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

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

[V0 Deprecation] Remove V0 core (#25321)

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

Woosuk Kwon committed Sep 20, 2025 at 19:58 UTC c99db8c8ddd50c0d30c14057012bd27b42d32b2d

29 files changed +24 -9020

.buildkite/test-pipeline.yaml

-3

@@ -148,7 +148,6 @@ steps:

148 num_gpus: 4

149 source_file_dependencies:

150 - vllm/distributed/

151 - - vllm/core/

151 - tests/distributed/test_utils

152 - tests/distributed/test_pynccl

153 - tests/distributed/test_events

@@ -867,8 +866,6 @@ steps:

866 - tests/distributed/

867 - vllm/compilation

868 - vllm/worker/worker_base.py

870 - - vllm/worker/worker.py

871 - - vllm/worker/model_runner.py

869 - entrypoints/llm/test_collective_rpc.py

870 - tests/v1/test_async_llm_dp.py

871 - tests/v1/test_external_lb_dp.py

.github/CODEOWNERS

-2

@@ -4,10 +4,8 @@

4 # This lists cover the "core" components of vLLM that require careful review

5 /vllm/attention @LucasWilkinson

6 /vllm/attention/backends/abstract.py @WoosukKwon @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill

7 -/vllm/core @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill

7 /vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill @22quinn

8 /vllm/worker/worker_base.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill @22quinn

10 -/vllm/worker/worker.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill

9 /vllm/model_executor/layers/fused_moe @mgoin

10 /vllm/model_executor/layers/sampler.py @zhuohan123 @youkaichao @alexm-redhat @comaniac @njhill @NickLucche

11 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256

pyproject.toml

-2

@@ -70,7 +70,6 @@ line-length = 80

70 "vllm/_version.py" = ["ALL"]

71 # Python 3.8 typing - skip V0 code

72 "vllm/attention/**/*.py" = ["UP006", "UP035"]

73 -"vllm/core/**/*.py" = ["UP006", "UP035"]

73 "vllm/engine/**/*.py" = ["UP006", "UP035"]

74 "vllm/executor/**/*.py" = ["UP006", "UP035"]

75 "vllm/worker/**/*.py" = ["UP006", "UP035"]

@@ -117,7 +116,6 @@ files = [

116 "vllm/*.py",

117 "vllm/assets",

118 "vllm/entrypoints",

120 - "vllm/core",

119 "vllm/inputs",

120 "vllm/logging_utils",

121 "vllm/multimodal",

vllm/attention/backends/differential_flash_attn.py

+4 -8

@@ -4,7 +4,7 @@

4 from collections import defaultdict

5 from dataclasses import dataclass

6 from itertools import accumulate

7 -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type

7 +from typing import Any, Dict, List, Optional, Tuple, Type

9 import torch

10 from einops import rearrange

@@ -34,9 +34,6 @@ from vllm.utils import async_tensor_h2d, make_tensor_with_pad

34 from vllm.vllm_flash_attn import (flash_attn_varlen_func,

35 flash_attn_with_kvcache)

37 -if TYPE_CHECKING:

38 - from vllm.worker.model_runner import ModelInputForGPUBuilder

39 -

37 logger = init_logger(__name__)

@@ -329,7 +326,7 @@ class DifferentialFlashAttentionMetadata(AttentionMetadata):

326 class DifferentialFlashAttentionMetadataBuilder(

327 AttentionMetadataBuilder[DifferentialFlashAttentionMetadata]):

328

332 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):

329 + def __init__(self, input_builder):

330 self.input_builder = input_builder

331 self.runner = input_builder.runner

332 self.sliding_window = input_builder.sliding_window

@@ -350,9 +347,8 @@ class DifferentialFlashAttentionMetadataBuilder(

347 self.num_decode_tokens = 0

348 self.has_prefix_cache_hit = False

349

353 - def _add_seq_group(

354 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

355 - chunked_prefill_enabled: bool, prefix_cache_hit: bool):

350 + def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

351 + prefix_cache_hit: bool):

352 """Add a sequence group to the metadata. Specifically update/append

353 1. context length.

354 2. block table.

vllm/attention/backends/dual_chunk_flash_attn.py

+3 -7

@@ -4,7 +4,7 @@

4 """

5 import math

6 from dataclasses import dataclass

7 -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type

7 +from typing import Any, Dict, List, Optional, Tuple, Type

9 import torch

10 import torch.distributed

@@ -22,9 +22,6 @@ from vllm.utils import async_tensor_h2d

22 from vllm.vllm_flash_attn import (flash_attn_varlen_func,

23 flash_attn_with_kvcache, sparse_attn_func)

25 -if TYPE_CHECKING:

26 - from vllm.worker.model_runner import ModelInputForGPUBuilder

27 -

25 logger = init_logger(__name__)

@@ -224,9 +221,8 @@ class DualChunkFlashAttentionMetadataBuilder(FlashAttentionMetadataBuilder):

221 super().prepare()

222 self.orig_seq_lens: List[int] = []

223

227 - def _add_seq_group(

228 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

229 - chunked_prefill_enabled: bool, prefix_cache_hit: bool):

224 + def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

225 + prefix_cache_hit: bool):

226 super()._add_seq_group(inter_data, chunked_prefill_enabled,

227 prefix_cache_hit)

228 for prompt_len, seq_len in zip(inter_data.prompt_lens,

vllm/attention/backends/flash_attn.py

+4 -8

@@ -4,7 +4,7 @@

4 from collections import defaultdict

5 from dataclasses import dataclass

6 from itertools import accumulate

7 -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type

7 +from typing import Dict, List, Optional, Tuple, Type

9 import torch

@@ -31,9 +31,6 @@ from vllm.utils import async_tensor_h2d, make_tensor_with_pad

31 from vllm.vllm_flash_attn import (flash_attn_varlen_func,

32 flash_attn_with_kvcache)

34 -if TYPE_CHECKING:

35 - from vllm.worker.model_runner import ModelInputForGPUBuilder

36 -

34 logger = init_logger(__name__)

@@ -312,7 +309,7 @@ class FlashAttentionMetadata(AttentionMetadata):

309 class FlashAttentionMetadataBuilder(

310 AttentionMetadataBuilder[FlashAttentionMetadata]):

311

315 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):

312 + def __init__(self, input_builder):

313 self.input_builder = input_builder

314 self.runner = input_builder.runner

315 self.sliding_window = input_builder.sliding_window

@@ -332,9 +329,8 @@ class FlashAttentionMetadataBuilder(

329 self.num_decode_tokens = 0

330 self.has_prefix_cache_hit = False

331

335 - def _add_seq_group(

336 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

337 - chunked_prefill_enabled: bool, prefix_cache_hit: bool):

332 + def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

333 + prefix_cache_hit: bool):

334 """Add a sequence group to the metadata. Specifically update/append

335 1. context length.

336 2. block table.

vllm/attention/backends/mla/common.py

+4 -9

@@ -193,8 +193,7 @@ from collections import defaultdict

193 from contextlib import contextmanager

194 from dataclasses import dataclass

195 from itertools import accumulate

196 -from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Tuple,

197 - Type, TypeVar)

196 +from typing import Any, Dict, Generic, List, Optional, Tuple, Type, TypeVar

197

198 import torch

199

@@ -233,9 +232,6 @@ except ImportError:

232 except ImportError:

233 flash_attn_varlen_func = None

234

236 -if TYPE_CHECKING:

237 - from vllm.worker.model_runner import ModelInputForGPUBuilder

238 -

235 is_hip = current_platform.is_rocm()

236

237

@@ -638,7 +634,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

634 """

635 BLOCK_TABLE_EXTENDER: list[list[int]] = []

636

641 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):

637 + def __init__(self, input_builder):

638 self.input_builder = input_builder

639 self.runner = input_builder.runner

640 self.sliding_window = input_builder.sliding_window

@@ -668,9 +664,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[T], Generic[T]):

664 self.num_decode_tokens = 0

665 self.has_prefix_cache_hit = False

666

671 - def _add_seq_group(

672 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

673 - chunked_prefill_enabled: bool, prefix_cache_hit: bool):

667 + def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool,

668 + prefix_cache_hit: bool):

669 """Add a sequence group to the metadata. Specifically update/append

670 1. context length.

671 2. block table.

vllm/attention/backends/placeholder_attn.py

+3 -8

@@ -4,7 +4,7 @@

4 from collections import defaultdict

5 from dataclasses import dataclass

6 from itertools import accumulate

7 -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type

7 +from typing import Dict, List, Optional, Tuple, Type

9 import torch

@@ -13,9 +13,6 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

13 AttentionMetadataBuilder)

14 from vllm.attention.backends.utils import CommonAttentionState

15 from vllm.multimodal import MultiModalPlaceholderMap

16 -

17 -if TYPE_CHECKING:

18 - from vllm.worker.model_runner import (ModelInputForGPUBuilder)

16 from vllm.utils import async_tensor_h2d

18 # Placeholder attention backend for models like Mamba and pooling models that

@@ -204,7 +201,7 @@ class PlaceholderAttentionMetadata(AttentionMetadata):

201 class PlaceholderAttentionMetadataBuilder(

202 AttentionMetadataBuilder[PlaceholderAttentionMetadata]):

203

207 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):

204 + def __init__(self, input_builder):

205

206 self.input_builder = input_builder

207 self.runner = input_builder.runner

@@ -220,9 +217,7 @@ class PlaceholderAttentionMetadataBuilder(

217 self.num_prefill_tokens = 0

218 self.num_decode_tokens = 0

219

223 - def _add_seq_group(

224 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

225 - chunked_prefill_enabled: bool):

220 + def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool):

221 """Add a sequence group to the metadata. Specifically update/append

222 1. context length.

223 """

vllm/attention/backends/rocm_aiter_mla.py

+2 -5

@@ -3,7 +3,7 @@

4 from contextlib import contextmanager

5 from dataclasses import dataclass

6 -from typing import TYPE_CHECKING, Optional, Type, Union

6 +from typing import Optional, Type, Union

8 import torch

@@ -19,9 +19,6 @@ from vllm.attention.backends.utils import (compute_slot_mapping,

19 from vllm.attention.ops.rocm_aiter_mla import (aiter_mla_decode_fwd,

20 get_aiter_mla_metadata)

22 -if TYPE_CHECKING:

23 - from vllm.worker.model_runner import ModelInputForGPUBuilder

24 -

23 def is_aiter_mla_enabled() -> bool:

24 return envs.VLLM_ROCM_USE_AITER \

@@ -110,7 +107,7 @@ class AiterMLAMetadata(MLACommonMetadata):

107 class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

108 BLOCK_TABLE_EXTENDER: list[list[int]] = [[]]

109

113 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):

110 + def __init__(self, input_builder):

111 super().__init__(input_builder)

112 assert self.block_size == 1, "AITER MLA requires only block size 1."

113

vllm/attention/backends/utils.py

+2 -7

@@ -35,9 +35,6 @@ PAD_SLOT_ID = -1

35 # if we have at least this many elements. Could be tuned further.

36 _COMPUTE_SLOT_MAPPING_NUMPY_NUMEL = 256

38 -if TYPE_CHECKING:

39 - from vllm.worker.model_runner import ModelInputForGPUBuilder

40 -

39 def is_block_tables_empty(block_tables: Union[None, Dict]):

40 """

@@ -129,7 +126,7 @@ class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]):

126

127 _metadata_cls: Type[TAttentionMetadata]

128

132 - def __init__(self, input_builder: "ModelInputForGPUBuilder"):

129 + def __init__(self, input_builder):

130 self.input_builder = input_builder

131 self.runner = input_builder.runner

132

@@ -149,9 +146,7 @@ class CommonMetadataBuilder(AttentionMetadataBuilder[TAttentionMetadata]):

146 self.num_prefill_tokens = 0

147 self.num_decode_tokens = 0

148

152 - def _add_seq_group(

153 - self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",

154 - chunked_prefill_enabled: bool):

149 + def _add_seq_group(self, inter_data, chunked_prefill_enabled: bool):

150 is_prompt = inter_data.is_prompt

151 block_tables = inter_data.block_tables

152

vllm/core/__init__.py

vllm/core/block/__init__.py

vllm/core/block/block_table.py deleted

-399

@@ -1,399 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import math

5 -from typing import List, Optional

6 -

7 -from vllm.core.block.common import BlockList

8 -from vllm.core.block.interfaces import Block, DeviceAwareBlockAllocator

9 -from vllm.utils import Device, cdiv, chunk_list

10 -

11 -

12 -class BlockTable:

13 - """A class to manage blocks for a specific sequence.

14 -

15 - The BlockTable maps a sequence of tokens to a list of blocks, where each

16 - block represents a contiguous memory allocation for a portion of the

17 - sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is

18 - responsible for allocating and freeing memory for the blocks.

19 -

20 - Args:

21 - block_size (int): The maximum number of tokens that can be stored in a

22 - single block.

23 - block_allocator (DeviceAwareBlockAllocator): The block allocator used to

24 - manage memory for the blocks.

25 - _blocks (Optional[List[Block]], optional): An optional list of existing

26 - blocks to initialize the BlockTable with. If not provided, an empty

27 - BlockTable is created.

28 - max_block_sliding_window (Optional[int], optional): The number of

29 - blocks to keep around for each sequence. If None, all blocks

30 - are kept (eg., when sliding window is not used).

31 - It should at least fit the sliding window size of the model.

32 -

33 - Attributes:

34 - _block_size (int): The maximum number of tokens that can be stored in a

35 - single block.

36 - _allocator (DeviceAwareBlockAllocator): The block allocator used to

37 - manage memory for the blocks.

38 - _blocks (Optional[List[Block]]): The list of blocks managed by this

39 - BlockTable.

40 - _num_full_slots (int): The number of tokens currently stored in the

41 - blocks.

42 - """

43 -

44 - def __init__(

45 - self,

46 - block_size: int,

47 - block_allocator: DeviceAwareBlockAllocator,

48 - _blocks: Optional[List[Block]] = None,

49 - max_block_sliding_window: Optional[int] = None,

50 - ):

51 - self._block_size = block_size

52 - self._allocator = block_allocator

53 - if _blocks is None:

54 - _blocks = []

55 - self._blocks: BlockList = BlockList(_blocks)

56 -

57 - self._max_block_sliding_window = max_block_sliding_window

58 - self._num_full_slots = self._get_num_token_ids()

59 -

60 - @staticmethod

61 - def get_num_required_blocks(token_ids: List[int],

62 - block_size: int,

63 - num_lookahead_slots: int = 0) -> int:

64 - """Calculates the minimum number of blocks required to store a given

65 - sequence of token IDs along with any look-ahead slots that may be

66 - required (like in multi-step + chunked-prefill).

67 -

68 - This assumes worst-case scenario, where every block requires a new

69 - allocation (e.g. ignoring prefix caching).

70 -

71 - Args:

72 - token_ids (List[int]): The sequence of token IDs to be stored.

73 - block_size (int): The maximum number of tokens that can be stored in

74 - a single block.

75 - num_lookahead_slots (int): look-ahead slots that the sequence may

76 - require.

77 -

78 - Returns:

79 - int: The minimum number of blocks required to store the given

80 - sequence of token IDs along with any required look-ahead slots.

81 - """

82 - return cdiv(len(token_ids) + num_lookahead_slots, block_size)

83 -

84 - def allocate(self,

85 - token_ids: List[int],

86 - device: Device = Device.GPU,

87 - extra_hash: Optional[int] = None) -> None:

88 - """Allocates memory blocks for storing the given sequence of token IDs.

89 -

90 - This method allocates the required number of blocks to store the given

91 - sequence of token IDs.

92 -

93 - Args:

94 - token_ids (List[int]): The sequence of token IDs to be stored.

95 - device (Device, optional): The device on which the blocks should be

96 - allocated. Defaults to Device.GPU.

97 - extra_hash (Optional[int]): The hash value of additional

98 - factors, such as adapters, that influence the block hash

99 - in the prefixcaching block.

100 - """

101 - assert not self._is_allocated

102 - assert token_ids

103 - blocks = self._allocate_blocks_for_token_ids(prev_block=None,

104 - token_ids=token_ids,

105 - device=device,

106 - extra_hash=extra_hash)

107 - self.update(blocks)

108 - self._num_full_slots = len(token_ids)

109 -

110 - def update(self, blocks: List[Block]) -> None:

111 - """Resets the table to the newly provided blocks

112 - (with their corresponding block ids)

113 - """

114 - self._blocks.update(blocks)

115 -

116 - def append_token_ids(self,

117 - token_ids: List[int],

118 - num_lookahead_slots: int = 0,

119 - num_computed_slots: Optional[int] = None,

120 - extra_hash: Optional[int] = None) -> None:

121 - """Appends a sequence of token IDs to the existing blocks in the

122 - BlockTable.

123 -

124 - This method appends the given sequence of token IDs to the existing

125 - blocks in the BlockTable. If there is not enough space in the existing

126 - blocks, new blocks are allocated using the `ensure_num_empty_slots`

127 - method to accommodate the additional tokens.

128 -

129 - The token IDs are divided into chunks of size `block_size` (except for

130 - the first chunk, which may be smaller), and each chunk is appended to a

131 - separate block.

132 -

133 - Args:

134 - token_ids (List[int]): The sequence of token IDs to be appended.

135 - num_computed_slots (Optional[int]): The number of KV cache slots

136 - that are already filled (computed).

137 - When sliding window is enabled, this is used to compute how many

138 - blocks to drop at the front of the sequence.

139 - Without sliding window, None can be passed.

140 - Without chunked prefill, it should be the same as

141 - _num_full_slots.

142 - extra_hash (Optional[int]): The hash value of additional

143 - factors such as adapters that influence the block, apart

144 - from the token_ids.

145 - """

146 - assert self._is_allocated, "no blocks have been allocated"

147 - assert len(self._blocks) > 0

148 -

149 - # Drop blocks that are no longer needed due to sliding window

150 - if self._max_block_sliding_window is not None:

151 - null_block = self._allocator.allocate_or_get_null_block()

152 - assert num_computed_slots is not None

153 - end_block_idx = (num_computed_slots //

154 - self._block_size) - self._max_block_sliding_window

155 - for idx in range(0, end_block_idx):

156 - b = self._blocks[idx]

157 - if b is not null_block:

158 - self._allocator.free(b)

159 - self._blocks[idx] = null_block

160 -

161 - # Ensure there are enough empty slots for the new tokens plus

162 - # lookahead slots

163 - self.ensure_num_empty_slots(num_empty_slots=len(token_ids) +

164 - num_lookahead_slots,

165 - extra_hash=extra_hash)

166 -

167 - # Update the blocks with the new tokens

168 - first_block_idx = self._num_full_slots // self._block_size

169 - token_blocks = self._chunk_token_blocks_for_append(token_ids)

170 -

171 - for i, token_block in enumerate(token_blocks):

172 - self._blocks.append_token_ids(first_block_idx + i, token_block)

173 -

174 - self._num_full_slots += len(token_ids)

175 -

176 - def ensure_num_empty_slots(self,

177 - num_empty_slots: int,

178 - extra_hash: Optional[int] = None) -> None:

179 - """Ensures that the BlockTable has at least the specified number of

180 - empty slots available.

181 -

182 - This method checks if the BlockTable has enough empty slots (i.e.,

183 - available space) to accommodate the requested number of tokens. If not,

184 - it allocates additional blocks on the GPU to ensure that the required

185 - number of empty slots is available.

186 -

187 - Args:

188 - num_empty_slots (int): The minimum number of empty slots required.

189 - extra_hash (Optional[int]): The hash value of additional

190 - factors such as adapters that influence the block, apart

191 - from the token_ids.

192 - """

193 - # Currently the block table only supports

194 - # appending tokens to GPU blocks.

195 - device = Device.GPU

196 - assert self._is_allocated

197 -

198 - if self._num_empty_slots >= num_empty_slots:

199 - return

200 -

201 - slots_to_allocate = num_empty_slots - self._num_empty_slots

202 - blocks_to_allocate = cdiv(slots_to_allocate, self._block_size)

203 -

204 - for _ in range(blocks_to_allocate):

205 - assert len(self._blocks) > 0

206 - self._blocks.append(

207 - self._allocator.allocate_mutable_block(

208 - prev_block=self._blocks[-1],

209 - device=device,

210 - extra_hash=extra_hash))

211 -

212 - def fork(self) -> "BlockTable":

213 - """Creates a new BlockTable instance with a copy of the blocks from the

214 - current instance.

215 -

216 - This method creates a new BlockTable instance with the same block size,

217 - block allocator, and a copy of the blocks from the current instance. The

218 - new BlockTable has its own independent set of blocks, but shares the

219 - same underlying memory allocation with the original BlockTable.

220 -

221 - Returns:

222 - BlockTable: A new BlockTable instance with a copy of the blocks from

223 - the current instance.

224 - """

225 - assert self._is_allocated

226 - assert len(self._blocks) > 0

227 - forked_blocks = self._allocator.fork(self._blocks[-1])

228 - return BlockTable(

229 - block_size=self._block_size,

230 - block_allocator=self._allocator,

231 - _blocks=forked_blocks,

232 - max_block_sliding_window=self._max_block_sliding_window,

233 - )

234 -

235 - def free(self) -> None:

236 - """Frees the memory occupied by the blocks in the BlockTable.

237 -

238 - This method iterates over all the blocks in the `_blocks` list and calls

239 - the `free` method of the `_allocator` object to release the memory

240 - occupied by each block. After freeing all the blocks, the `_blocks` list

241 - is set to `None`.

242 - """

243 - for block in self.blocks:

244 - self._allocator.free(block)

245 - self._blocks.reset()

246 -

247 - @property

248 - def physical_block_ids(self) -> List[int]:

249 - """Returns a list of physical block indices for the blocks in the

250 - BlockTable.

251 -

252 - This property returns a list of integers, where each integer represents

253 - the physical block index of a corresponding block in the `_blocks` list.

254 - The physical block index is a unique identifier for the memory location

255 - occupied by the block.

256 -

257 - Returns:

258 - List[int]: A list of physical block indices for the blocks in the

259 - BlockTable.

260 - """

261 - return self._blocks.ids()

262 -

263 - def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]:

264 - """Get the number of "unseen" tokens in the sequence.

265 -

266 - Unseen tokens are tokens in the sequence corresponding to this block

267 - table, but are not yet appended to this block table.

268 -

269 - Args:

270 - sequence_token_ids (List[int]): The list of token ids in the

271 - sequence.

272 -

273 - Returns:

274 - List[int]: The postfix of sequence_token_ids that has not yet been

275 - appended to the block table.

276 - """

277 -

278 - # Since the block table is append-only, the unseen token ids are the

279 - # ones after the appended ones.

280 - return sequence_token_ids[self.num_full_slots:]

281 -

282 - def _allocate_blocks_for_token_ids(

283 - self,

284 - prev_block: Optional[Block],

285 - token_ids: List[int],

286 - device: Device,

287 - extra_hash: Optional[int] = None) -> List[Block]:

288 - blocks: List[Block] = []

289 -

290 - block_token_ids = []

291 - tail_token_ids = []

292 - for cur_token_ids in chunk_list(token_ids, self._block_size):

293 - if len(cur_token_ids) == self._block_size:

294 - block_token_ids.append(cur_token_ids)

295 - else:

296 - tail_token_ids.append(cur_token_ids)

297 -

298 - if block_token_ids:

299 - blocks.extend(

300 - self._allocator.allocate_immutable_blocks(

301 - prev_block,

302 - block_token_ids=block_token_ids,

303 - device=device,

304 - extra_hash=extra_hash))

305 - prev_block = blocks[-1]

306 -

307 - if tail_token_ids:

308 - assert len(tail_token_ids) == 1

309 - cur_token_ids = tail_token_ids[0]

310 -

311 - block = self._allocator.allocate_mutable_block(

312 - prev_block=prev_block, device=device, extra_hash=extra_hash)

313 - block.append_token_ids(cur_token_ids)

314 -

315 - blocks.append(block)

316 -

317 - return blocks

318 -

319 - def _get_all_token_ids(self) -> List[int]:

320 - # NOTE: This function is O(seq_len); use sparingly.

321 - token_ids: List[int] = []

322 -

323 - if not self._is_allocated:

324 - return token_ids

325 -

326 - for block in self.blocks:

327 - token_ids.extend(block.token_ids)

328 -

329 - return token_ids

330 -

331 - def _get_num_token_ids(self) -> int:

332 - res = 0

333 - for block in self.blocks:

334 - res += len(block.token_ids)

335 -

336 - return res

337 -

338 - @property

339 - def _is_allocated(self) -> bool:

340 - return len(self._blocks) > 0

341 -

342 - @property

343 - def blocks(self) -> List[Block]:

344 - return self._blocks.list()

345 -

346 - @property

347 - def _num_empty_slots(self) -> int:

348 - assert self._is_allocated

349 - return len(self._blocks) * self._block_size - self._num_full_slots

350 -

351 - @property

352 - def num_full_slots(self) -> int:

353 - """Returns the total number of tokens currently stored in the

354 - BlockTable.

355 -

356 - Returns:

357 - int: The total number of tokens currently stored in the BlockTable.

358 - """

359 - return self._num_full_slots

360 -

361 - def get_num_blocks_touched_by_append_slots(

362 - self, token_ids: List[int], num_lookahead_slots: int) -> int:

363 - """Determine how many blocks will be "touched" by appending the token

364 - ids.

365 -

366 - This is required for the scheduler to determine whether a sequence can

367 - continue generation, or if it must be preempted.

368 - """

369 - # Math below is equivalent to:

370 - # all_token_ids = token_ids + [-1] * num_lookahead_slots

371 - # token_blocks = self._chunk_token_blocks_for_append(all_token_ids)

372 - # return len(token_blocks)

373 -

374 - num_token_ids = len(token_ids) + num_lookahead_slots

375 - first_chunk_size = self._block_size - (self._num_full_slots %

376 - self._block_size)

377 - num_token_blocks = (1 + math.ceil(

378 - (num_token_ids - first_chunk_size) / self._block_size))

379 - return num_token_blocks

380 -

381 - def _chunk_token_blocks_for_append(

382 - self, token_ids: List[int]) -> List[List[int]]:

383 - """Split the token ids into block-sized chunks so they can be easily

384 - appended to blocks. The first such "token block" may have less token ids

385 - than the block size, since the last allocated block may be partially

386 - full.

387 -

388 - If no token ids are provided, then no chunks are returned.

389 - """

390 -

391 - if not token_ids:

392 - return []

393 -

394 - first_chunk_size = self._block_size - (self._num_full_slots %

395 - self._block_size)

396 - token_blocks = [token_ids[:first_chunk_size]]

397 - token_blocks.extend(

398 - chunk_list(token_ids[first_chunk_size:], self._block_size))

399 - return token_blocks

vllm/core/block/common.py deleted

-371

@@ -1,371 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from collections import deque

5 -from dataclasses import dataclass

6 -from typing import Deque, Dict, Iterable, List, Optional, Protocol, Tuple

7 -

8 -from vllm.core.block.interfaces import Block, BlockAllocator

9 -

10 -BlockId = int

11 -RefCount = int

12 -

13 -

14 -class RefCounterProtocol(Protocol):

15 -

16 - def incr(self, block_id: BlockId) -> RefCount:

17 - raise NotImplementedError

18 -

19 - def decr(self, block_id: BlockId) -> RefCount:

20 - raise NotImplementedError

21 -

22 - def get(self, block_id: BlockId) -> RefCount:

23 - raise NotImplementedError

24 -

25 -

26 -class RefCounter(RefCounterProtocol):

27 - """A class for managing reference counts for a set of block indices.

28 -

29 - The RefCounter class maintains a dictionary that maps block indices to their

30 - corresponding reference counts. It provides methods to increment, decrement,

31 - and retrieve the reference count for a given block index.

32 -

33 - Args:

34 - all_block_indices (Iterable[BlockId]): An iterable of block indices

35 - to initialize the reference counter with.

36 - """

37 -

38 - def __init__(self, all_block_indices: Iterable[BlockId]):

39 - deduped = set(all_block_indices)

40 - self._refcounts: Dict[BlockId, RefCount] = {

41 - index: 0

42 - for index in deduped

43 - }

44 -

45 - def incr(self, block_id: BlockId) -> RefCount:

46 - assert block_id in self._refcounts

47 - pre_incr_refcount = self._refcounts[block_id]

48 -

49 - assert pre_incr_refcount >= 0

50 -

51 - post_incr_refcount = pre_incr_refcount + 1

52 - self._refcounts[block_id] = post_incr_refcount

53 - return post_incr_refcount

54 -

55 - def decr(self, block_id: BlockId) -> RefCount:

56 - assert block_id in self._refcounts

57 - refcount = self._refcounts[block_id]

58 -

59 - assert refcount > 0

60 - refcount -= 1

61 -

62 - self._refcounts[block_id] = refcount

63 -

64 - return refcount

65 -

66 - def get(self, block_id: BlockId) -> RefCount:

67 - assert block_id in self._refcounts

68 - return self._refcounts[block_id]

69 -

70 - def as_readonly(self) -> "ReadOnlyRefCounter":

71 - return ReadOnlyRefCounter(self)

72 -

73 -

74 -class ReadOnlyRefCounter(RefCounterProtocol):

75 - """A read-only view of the RefCounter class.

76 -

77 - The ReadOnlyRefCounter class provides a read-only interface to access the

78 - reference counts maintained by a RefCounter instance. It does not allow

79 - modifications to the reference counts.

80 -

81 - Args:

82 - refcounter (RefCounter): The RefCounter instance to create a read-only

83 - view for.

84 - """

85 -

86 - def __init__(self, refcounter: RefCounter):

87 - self._refcounter = refcounter

88 -

89 - def incr(self, block_id: BlockId) -> RefCount:

90 - raise ValueError("Incr not allowed")

91 -

92 - def decr(self, block_id: BlockId) -> RefCount:

93 - raise ValueError("Decr not allowed")

94 -

95 - def get(self, block_id: BlockId) -> RefCount:

96 - return self._refcounter.get(block_id)

97 -

98 -

99 -class CopyOnWriteTracker:

100 - """A class for tracking and managing copy-on-write operations for blocks.

101 -

102 - The CopyOnWriteTracker class maintains a mapping of source block indices to

103 - their corresponding copy-on-write destination block indices. It works in

104 - conjunction with a RefCounter.

105 -

106 - Args:

107 - refcounter (RefCounter): The reference counter used to track block

108 - reference counts.

109 - """

110 -

111 - def __init__(self, refcounter: RefCounterProtocol):

112 - self._copy_on_writes: List[Tuple[BlockId, BlockId]] = []

113 - self._refcounter = refcounter

114 -

115 - def is_appendable(self, block: Block) -> bool:

116 - """Checks if the block is shared or not. If shared, then it cannot

117 - be appended and needs to be duplicated via copy-on-write

118 - """

119 - block_id = block.block_id

120 - if block_id is None:

121 - return True

122 -

123 - refcount = self._refcounter.get(block_id)

124 - return refcount <= 1

125 -

126 - def record_cow(self, src_block_id: Optional[BlockId],

127 - trg_block_id: Optional[BlockId]) -> None:

128 - """Records a copy-on-write operation from source to target block id

129 - Args:

130 - src_block_id (BlockId): The source block id from which to copy

131 - the data

132 - trg_block_id (BlockId): The target block id to which the data

133 - is copied

134 - """

135 - assert src_block_id is not None

136 - assert trg_block_id is not None

137 - self._copy_on_writes.append((src_block_id, trg_block_id))

138 -

139 - def clear_cows(self) -> List[Tuple[BlockId, BlockId]]:

140 - """Clears the copy-on-write tracking information and returns the current

141 - state.

142 -

143 - This method returns a list mapping source block indices to

144 - destination block indices for the current copy-on-write operations.

145 - It then clears the internal tracking information.

146 -

147 - Returns:

148 - List[Tuple[BlockId, BlockId]]: A list mapping source

149 - block indices to destination block indices for the

150 - current copy-on-write operations.

151 - """

152 - cows = self._copy_on_writes

153 - self._copy_on_writes = []

154 - return cows

155 -

156 -

157 -class BlockPool:

158 - """Used to pre-allocate block objects, in order to avoid excessive python

159 - object allocations/deallocations.

160 - The pool starts from "pool_size" objects and will increase to more objects

161 - if necessary

162 -

163 - Note that multiple block objects may point to the same physical block id,

164 - which is why this pool is needed, so that it will be easier to support

165 - prefix caching and more complicated sharing of physical blocks.

166 - """

167 -

168 - def __init__(self, block_size: int, create_block: Block.Factory,

169 - allocator: BlockAllocator, pool_size: int):

170 - self._block_size = block_size

171 - self._create_block = create_block

172 - self._allocator = allocator

173 - self._pool_size = pool_size

174 - assert self._pool_size >= 0

175 -

176 - self._free_ids: Deque[int] = deque(range(self._pool_size))

177 - self._pool = []

178 - for i in range(self._pool_size):

179 - self._pool.append(

180 - self._create_block(prev_block=None,

181 - token_ids=[],

182 - block_size=self._block_size,

183 - allocator=self._allocator,

184 - block_id=None,

185 - extra_hash=None))

186 -

187 - def increase_pool(self):

188 - """Doubles the internal pool size

189 - """

190 - cur_pool_size = self._pool_size

191 - new_pool_size = cur_pool_size * 2

192 - self._pool_size = new_pool_size

193 -

194 - self._free_ids += deque(range(cur_pool_size, new_pool_size))

195 -

196 - for i in range(cur_pool_size, new_pool_size):

197 - self._pool.append(

198 - self._create_block(prev_block=None,

199 - token_ids=[],

200 - block_size=self._block_size,

201 - allocator=self._allocator,

202 - block_id=None,

203 - extra_hash=None))

204 -

205 - def init_block(self,

206 - prev_block: Optional[Block],

207 - token_ids: List[int],

208 - block_size: int,

209 - physical_block_id: Optional[int],

210 - extra_hash: Optional[int] = None) -> Block:

211 - if len(self._free_ids) == 0:

212 - self.increase_pool()

213 - assert len(self._free_ids) > 0

214 -

215 - pool_id = self._free_ids.popleft()

216 -

217 - block = self._pool[pool_id]

218 - block.__init__( # type: ignore[misc]

219 - prev_block=prev_block,

220 - token_ids=token_ids,

221 - block_size=block_size,

222 - allocator=block._allocator, # type: ignore[attr-defined]

223 - block_id=physical_block_id,

224 - extra_hash=extra_hash)

225 - block.pool_id = pool_id # type: ignore[attr-defined]

226 - return block

227 -

228 - def free_block(self, block: Block) -> None:

229 - self._free_ids.appendleft(block.pool_id) # type: ignore[attr-defined]

230 -

231 -

232 -class BlockList:

233 - """This class is an optimization to allow fast-access to physical

234 - block ids. It maintains a block id list that is updated with the

235 - block list and this avoids the need to reconstruct the block id

236 - list on every iteration of the block manager

237 - """

238 -

239 - def __init__(self, blocks: List[Block]):

240 - self._blocks: List[Block] = []

241 - self._block_ids: List[int] = []

242 -

243 - self.update(blocks)

244 -

245 - def _add_block_id(self, block_id: Optional[BlockId]) -> None:

246 - assert block_id is not None

247 - self._block_ids.append(block_id)

248 -

249 - def _update_block_id(self, block_index: int,

250 - new_block_id: Optional[BlockId]) -> None:

251 - assert new_block_id is not None

252 - self._block_ids[block_index] = new_block_id

253 -

254 - def update(self, blocks: List[Block]):

255 - self._blocks = blocks

256 -

257 - # Cache block ids for fast query

258 - self._block_ids = []

259 - for block in self._blocks:

260 - self._add_block_id(block.block_id)

261 -

262 - def append_token_ids(self, block_index: int, token_ids: List[int]) -> None:

263 - block = self._blocks[block_index]

264 - prev_block_id = block.block_id

265 -

266 - block.append_token_ids(token_ids)

267 -

268 - # CoW or promotion may update the internal block_id

269 - if prev_block_id != block.block_id:

270 - self._update_block_id(block_index, block.block_id)

271 -

272 - def append(self, new_block: Block):

273 - self._blocks.append(new_block)

274 - self._add_block_id(new_block.block_id)

275 -

276 - def __len__(self) -> int:

277 - return len(self._blocks)

278 -

279 - def __getitem__(self, block_index: int) -> Block:

280 - return self._blocks[block_index]

281 -

282 - def __setitem__(self, block_index: int, new_block: Block) -> None:

283 - self._blocks[block_index] = new_block

284 - self._update_block_id(block_index, new_block.block_id)

285 -

286 - def reset(self):

287 - self._blocks = []

288 - self._block_ids = []

289 -

290 - def list(self) -> List[Block]:

291 - return self._blocks

292 -

293 - def ids(self) -> List[int]:

294 - return self._block_ids

295 -

296 -

297 -@dataclass

298 -class CacheMetricData:

299 - """A utility dataclass to maintain cache metric.

300 - To avoid overflow, we maintain the hit rate in block granularity, so that

301 - we can maintain a single hit rate for n_completed_block x block_size,

302 - and calculate the real time hit rate by the following:

303 - BS = The number of queries per block.

304 - nB = The number of completed blocks.

305 - HR = hit rate of (nB x BS) queries.

306 - Q = current number of queries (< BS).

307 - H = current number of hits (< BS).

308 - hit rate = ((HR x nB) + (H / Q) x (Q / BS)) / (nB + Q / BS)

309 - """

310 - num_completed_blocks: int = 0

311 - completed_block_cache_hit_rate: float = 0.0

312 - num_incompleted_block_queries: int = 0

313 - num_incompleted_block_hit: int = 0

314 - block_size: int = 1000

315 -

316 - def query(self, hit: bool):

317 - self.num_incompleted_block_queries += 1

318 - self.num_incompleted_block_hit += 1 if hit else 0

319 -

320 - # When a block is completed, update the cache hit rate

321 - # and reset the incomplete numbers.

322 - if self.num_incompleted_block_queries == self.block_size:

323 - hit_rate = (self.num_incompleted_block_hit /

324 - self.num_incompleted_block_queries)

325 - self.completed_block_cache_hit_rate = (

326 - self.completed_block_cache_hit_rate * self.num_completed_blocks

327 - + hit_rate) / (self.num_completed_blocks + 1)

328 - self.num_incompleted_block_queries = 0

329 - self.num_incompleted_block_hit = 0

330 - self.num_completed_blocks += 1

331 -

332 - def get_hit_rate(self):

333 - incomplete_ratio = self.num_incompleted_block_queries / self.block_size

334 - total_blocks = self.num_completed_blocks + incomplete_ratio

335 - if total_blocks == 0:

336 - return 0.0

337 -

338 - completed_block_hit, incompleted_block_hit = 0.0, 0.0

339 - if self.num_completed_blocks > 0:

340 - completed_block_hit = (self.completed_block_cache_hit_rate *

341 - self.num_completed_blocks)

342 - if self.num_incompleted_block_queries > 0:

343 - incompleted_hit_rate = (self.num_incompleted_block_hit /

344 - self.num_incompleted_block_queries)

345 - incompleted_block_hit = (incompleted_hit_rate * incomplete_ratio)

346 - return (completed_block_hit + incompleted_block_hit) / total_blocks

347 -

348 -

349 -def get_all_blocks_recursively(last_block: Block) -> List[Block]:

350 - """Retrieves all the blocks in a sequence starting from the last block.

351 -

352 - This function recursively traverses the sequence of blocks in reverse order,

353 - starting from the given last block, and returns a list of all the blocks in

354 - the sequence.

355 -

356 - Args:

357 - last_block (Block): The last block in the sequence.

358 -

359 - Returns:

360 - List[Block]: A list of all the blocks in the sequence, in the order they

361 - appear.

362 - """

363 -

364 - def recurse(block: Block, lst: List[Block]) -> None:

365 - if block.prev_block is not None:

366 - recurse(block.prev_block, lst)

367 - lst.append(block)

368 -

369 - all_blocks: List[Block] = []

370 - recurse(last_block, all_blocks)

371 - return all_blocks

vllm/core/block/cpu_gpu_block_allocator.py deleted

-439

@@ -1,439 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from typing import Dict, FrozenSet, List, Optional, Tuple

5 -

6 -from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId,

7 - DeviceAwareBlockAllocator)

8 -from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator

9 -from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator

10 -from vllm.utils import Device

11 -

12 -

13 -class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):

14 - """A block allocator that can allocate blocks on both CPU and GPU memory.

15 -

16 - This class implements the `DeviceAwareBlockAllocator` interface and provides

17 - functionality for allocating and managing blocks of memory on both CPU and

18 - GPU devices.

19 -

20 - The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU

21 - blocks, and allows for allocation, deallocation, forking, and swapping of

22 - blocks across these memory pools.

23 - """

24 -

25 - @staticmethod

26 - def create(

27 - allocator_type: str,

28 - num_gpu_blocks: int,

29 - num_cpu_blocks: int,

30 - block_size: int,

31 - ) -> DeviceAwareBlockAllocator:

32 - """Creates a CpuGpuBlockAllocator instance with the specified

33 - configuration.

34 -

35 - This static method creates and returns a CpuGpuBlockAllocator instance

36 - based on the provided parameters. It initializes the CPU and GPU block

37 - allocators with the specified number of blocks, block size, and

38 - allocator type.

39 -

40 - Args:

41 - allocator_type (str): The type of block allocator to use for CPU

42 - and GPU blocks. Currently supported values are "naive" and

43 - "prefix_caching".

44 - num_gpu_blocks (int): The number of blocks to allocate for GPU

45 - memory.

46 - num_cpu_blocks (int): The number of blocks to allocate for CPU

47 - memory.

48 - block_size (int): The size of each block in number of tokens.

49 -

50 - Returns:

51 - DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the

52 - specified configuration.

53 -

54 - Notes:

55 - - The block IDs are assigned contiguously, with GPU block IDs coming

56 - before CPU block IDs.

57 - """

58 - reserved_blocks = 0

59 - block_ids = list(

60 - range(reserved_blocks, num_gpu_blocks + num_cpu_blocks))

61 - num_gpu_blocks -= reserved_blocks

62 - gpu_block_ids = block_ids[:num_gpu_blocks]

63 - cpu_block_ids = block_ids[num_gpu_blocks:]

64 -

65 - if allocator_type == "naive":

66 - gpu_allocator: BlockAllocator = NaiveBlockAllocator(

67 - create_block=NaiveBlock, # type: ignore

68 - num_blocks=num_gpu_blocks,

69 - block_size=block_size,

70 - block_ids=gpu_block_ids,

71 - )

72 -

73 - cpu_allocator: BlockAllocator = NaiveBlockAllocator(

74 - create_block=NaiveBlock, # type: ignore

75 - num_blocks=num_cpu_blocks,

76 - block_size=block_size,

77 - block_ids=cpu_block_ids,

78 - )

79 - elif allocator_type == "prefix_caching":

80 - gpu_allocator = PrefixCachingBlockAllocator(

81 - num_blocks=num_gpu_blocks,

82 - block_size=block_size,

83 - block_ids=gpu_block_ids,

84 - )

85 -

86 - cpu_allocator = PrefixCachingBlockAllocator(

87 - num_blocks=num_cpu_blocks,

88 - block_size=block_size,

89 - block_ids=cpu_block_ids,

90 - )

91 - else:

92 - raise ValueError(f"Unknown allocator type {allocator_type=}")

93 -

94 - return CpuGpuBlockAllocator(

95 - cpu_block_allocator=cpu_allocator,

96 - gpu_block_allocator=gpu_allocator,

97 - )

98 -

99 - def __init__(self, cpu_block_allocator: BlockAllocator,

100 - gpu_block_allocator: BlockAllocator):

101 - assert not (

102 - cpu_block_allocator.all_block_ids

103 - & gpu_block_allocator.all_block_ids

104 - ), "cpu and gpu block allocators can't have intersection of block ids"

105 -

106 - self._allocators = {

107 - Device.CPU: cpu_block_allocator,

108 - Device.GPU: gpu_block_allocator,

109 - }

110 -

111 - self._swap_mapping: Dict[int, int] = {}

112 - self._null_block: Optional[Block] = None

113 -

114 - self._block_ids_to_allocator: Dict[int, BlockAllocator] = {}

115 - for _, allocator in self._allocators.items():

116 - for block_id in allocator.all_block_ids:

117 - self._block_ids_to_allocator[block_id] = allocator

118 -

119 - def allocate_or_get_null_block(self) -> Block:

120 - if self._null_block is None:

121 - self._null_block = NullBlock(

122 - self.allocate_mutable_block(None, Device.GPU))

123 - return self._null_block

124 -

125 - def allocate_mutable_block(self,

126 - prev_block: Optional[Block],

127 - device: Device,

128 - extra_hash: Optional[int] = None) -> Block:

129 - """Allocates a new mutable block on the specified device.

130 -

131 - Args:

132 - prev_block (Optional[Block]): The previous block to in the sequence.

133 - Used for prefix hashing.

134 - device (Device): The device on which to allocate the new block.

135 - extra_hash (Optional[int]): The hash value of additional

136 - factors, such as adapters, that influence the block hash

137 - in the prefix caching block.

138 -

139 - Returns:

140 - Block: The newly allocated mutable block.

141 - """

142 - return self._allocators[device].allocate_mutable_block(

143 - prev_block, extra_hash=extra_hash)

144 -

145 - def allocate_immutable_blocks(

146 - self,

147 - prev_block: Optional[Block],

148 - block_token_ids: List[List[int]],

149 - device: Device,

150 - extra_hash: Optional[int] = None) -> List[Block]:

151 - """Allocates a new group of immutable blocks with the provided block

152 - token IDs on the specified device.

153 -

154 - Args:

155 - prev_block (Optional[Block]): The previous block in the sequence.

156 - Used for prefix hashing.

157 - block_token_ids (List[int]): The list of block token IDs to be

158 - stored in the new blocks.

159 - device (Device): The device on which to allocate the new block.

160 - extra_hash (Optional[int]): The hash value of additional

161 - factors, such as adapters, that influence the block hash

162 - in the prefix caching block.

163 -

164 - Returns:

165 - List[Block]: The newly allocated list of immutable blocks

166 - containing the provided block token IDs.

167 - """

168 - return self._allocators[device].allocate_immutable_blocks(

169 - prev_block, block_token_ids, extra_hash=extra_hash)

170 -

171 - def allocate_immutable_block(self,

172 - prev_block: Optional[Block],

173 - token_ids: List[int],

174 - device: Device,

175 - extra_hash: Optional[int] = None) -> Block:

176 - """Allocates a new immutable block with the provided token IDs on the

177 - specified device.

178 -

179 - Args:

180 - prev_block (Optional[Block]): The previous block in the sequence.

181 - Used for prefix hashing.

182 - token_ids (List[int]): The list of token IDs to be stored in the new

183 - block.

184 - device (Device): The device on which to allocate the new block.

185 - extra_hash (Optional[int]): The hash value of additional

186 - factors, such as adapters, that influence the block hash

187 - in the prefix caching block.

188 -

189 - Returns:

190 - Block: The newly allocated immutable block containing the provided

191 - token IDs.

192 - """

193 - return self._allocators[device].allocate_immutable_block(

194 - prev_block, token_ids, extra_hash=extra_hash)

195 -

196 - def free(self, block: Block) -> None:

197 - """Frees the memory occupied by the given block.

198 -

199 - Args:

200 - block (Block): The block to be freed.

201 - """

202 - # Null block should never be freed

203 - if isinstance(block, NullBlock):

204 - return

205 - block_id = block.block_id

206 - assert block_id is not None

207 - allocator = self._block_ids_to_allocator[block_id]

208 - allocator.free(block)

209 -

210 - def fork(self, last_block: Block) -> List[Block]:

211 - """Creates a new sequence of blocks that shares the same underlying

212 - memory as the original sequence.

213 -

214 - Args:

215 - last_block (Block): The last block in the original sequence.

216 -

217 - Returns:

218 - List[Block]: A new list of blocks that shares the same memory as the

219 - original sequence.

220 - """

221 - # do not attempt to fork the null block

222 - assert not isinstance(last_block, NullBlock)

223 - block_id = last_block.block_id

224 - assert block_id is not None

225 - allocator = self._block_ids_to_allocator[block_id]

226 - return allocator.fork(last_block)

227 -

228 - def get_num_free_blocks(self, device: Device) -> int:

229 - """Returns the number of free blocks available on the specified device.

230 -

231 - Args:

232 - device (Device): The device for which to query the number of free

233 - blocks. AssertionError is raised if None is passed.

234 -

235 - Returns:

236 - int: The number of free blocks available on the specified device.

237 - """

238 - return self._allocators[device].get_num_free_blocks()

239 -

240 - def get_num_total_blocks(self, device: Device) -> int:

241 - return self._allocators[device].get_num_total_blocks()

242 -

243 - def get_physical_block_id(self, device: Device, absolute_id: int) -> int:

244 - """Returns the zero-offset block id on certain device given the

245 - absolute block id.

246 -

247 - Args:

248 - device (Device): The device for which to query relative block id.

249 - absolute_id (int): The absolute block id for the block in

250 - whole allocator.

251 -

252 - Returns:

253 - int: The zero-offset block id on certain device.

254 - """

255 - return self._allocators[device].get_physical_block_id(absolute_id)

256 -

257 - def swap(self, blocks: List[Block], src_device: Device,

258 - dst_device: Device) -> Dict[int, int]:

259 - """Execute the swap for the given blocks from source_device

260 - on to dest_device, save the current swap mapping and append

261 - them to the accumulated `self._swap_mapping` for each

262 - scheduling move.

263 -

264 - Args:

265 - blocks: List of blocks to be swapped.

266 - src_device (Device): Device to swap the 'blocks' from.

267 - dst_device (Device): Device to swap the 'blocks' to.

268 -

269 - Returns:

270 - Dict[int, int]: Swap mapping from source_device

271 - on to dest_device.

272 - """

273 - src_block_ids = [block.block_id for block in blocks]

274 - self._allocators[src_device].swap_out(blocks)

275 - self._allocators[dst_device].swap_in(blocks)

276 - dst_block_ids = [block.block_id for block in blocks]

277 -

278 - current_swap_mapping: Dict[int, int] = {}

279 - for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids):

280 - if src_block_id is not None and dst_block_id is not None:

281 - self._swap_mapping[src_block_id] = dst_block_id

282 - current_swap_mapping[src_block_id] = dst_block_id

283 - return current_swap_mapping

284 -

285 - def get_num_full_blocks_touched(self, blocks: List[Block],

286 - device: Device) -> int:

287 - """Returns the number of full blocks that will be touched by

288 - swapping in/out the given blocks on to the 'device'.

289 -

290 - Args:

291 - blocks: List of blocks to be swapped.

292 - device (Device): Device to swap the 'blocks' on.

293 -

294 - Returns:

295 - int: the number of full blocks that will be touched by

296 - swapping in/out the given blocks on to the 'device'.

297 - Non full blocks are ignored when deciding the number

298 - of blocks to touch.

299 - """

300 - return self._allocators[device].get_num_full_blocks_touched(blocks)

301 -

302 - def clear_copy_on_writes(self) -> List[Tuple[int, int]]:

303 - """Clears the copy-on-write (CoW) state and returns the mapping of

304 - source to destination block IDs.

305 -

306 - Returns:

307 - List[Tuple[int, int]]: A list mapping source block IDs to

308 - destination block IDs.

309 - """

310 - # CoW only supported on GPU

311 - device = Device.GPU

312 - return self._allocators[device].clear_copy_on_writes()

313 -

314 - def mark_blocks_as_accessed(self, block_ids: List[int],

315 - now: float) -> None:

316 - """Mark blocks as accessed, only use for prefix caching."""

317 - # Prefix caching only supported on GPU.

318 - device = Device.GPU

319 - return self._allocators[device].mark_blocks_as_accessed(block_ids, now)

320 -

321 - def mark_blocks_as_computed(self, block_ids: List[int]) -> None:

322 - """Mark blocks as accessed, only use for prefix caching."""

323 - # Prefix caching only supported on GPU.

324 - device = Device.GPU

325 - return self._allocators[device].mark_blocks_as_computed(block_ids)

326 -

327 - def get_common_computed_block_ids(

328 - self, computed_seq_block_ids: List[List[int]]) -> List[int]:

329 - # Prefix caching only supported on GPU.

330 - device = Device.GPU

331 - return self._allocators[device].get_common_computed_block_ids(

332 - computed_seq_block_ids)

333 -

334 - @property

335 - def all_block_ids(self) -> FrozenSet[int]:

336 - return frozenset(self._block_ids_to_allocator.keys())

337 -

338 - def get_prefix_cache_hit_rate(self, device: Device) -> float:

339 - """Prefix cache hit rate. -1 means not supported or disabled."""

340 - assert device in self._allocators

341 - return self._allocators[device].get_prefix_cache_hit_rate()

342 -

343 - def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

344 - """Reset prefix cache for specified or all devices."""

345 - if device:

346 - return self._allocators[device].reset_prefix_cache()

347 - success = True

348 - for allocator in self._allocators.values():

349 - success = success and allocator.reset_prefix_cache()

350 - return success

351 -

352 - def get_and_reset_swaps(self) -> List[Tuple[int, int]]:

353 - """Returns and clears the mapping of source to destination block IDs.

354 - Will be called after every swapping operations for now, and after every

355 - schedule when BlockManagerV2 become default. Currently not useful.

356 -

357 - Returns:

358 - List[Tuple[int, int]]: A mapping of source to destination block IDs.

359 - """

360 - mapping = self._swap_mapping.copy()

361 - self._swap_mapping.clear()

362 - return list(mapping.items())

363 -

364 - def find_cached_blocks_prefix(

365 - self,

366 - block_hashes: List[int],

367 - device: Device = Device.GPU,

368 - ) -> List[int]:

369 - return self._allocators[device].find_cached_blocks_prefix(block_hashes)

370 -

371 -

372 -class NullBlock(Block):

373 - """

374 - Null blocks are used as a placeholders for KV cache blocks that have

375 - been dropped due to sliding window.

376 - This implementation just wraps an ordinary block and prevents it from

377 - being modified. It also allows for testing if a block is NullBlock

378 - via isinstance().

379 - """

380 -

381 - def __init__(self, proxy: Block):

382 - super().__init__()

383 - self._proxy = proxy

384 -

385 - def append_token_ids(self, token_ids: List[BlockId]):

386 - raise ValueError("null block should not be modified")

387 -

388 - @property

389 - def block_id(self):

390 - return self._proxy.block_id

391 -

392 - @block_id.setter

393 - def block_id(self, value: Optional[BlockId]):

394 - raise ValueError("null block should not be modified")

395 -

396 - @property

397 - def token_ids(self) -> List[BlockId]:

398 - return self._proxy.token_ids

399 -

400 - @property

401 - def num_tokens_total(self) -> int:

402 - raise NotImplementedError(

403 - "num_tokens_total is not used for null block")

404 -

405 - @property

406 - def num_empty_slots(self) -> BlockId:

407 - return self._proxy.num_empty_slots

408 -

409 - @property

410 - def is_full(self):

411 - return self._proxy.is_full

412 -

413 - @property

414 - def prev_block(self):

415 - return self._proxy.prev_block

416 -

417 - @property

418 - def extra_hash(self):

419 - return None

420 -

421 - @property

422 - def computed(self):

423 - return self._proxy.computed

424 -

425 - @computed.setter

426 - def computed(self, value):

427 - self._proxy.computed = value

428 -

429 - @property

430 - def last_accessed(self) -> float:

431 - return self._proxy.last_accessed

432 -

433 - @last_accessed.setter

434 - def last_accessed(self, last_accessed_ts: float):

435 - self._proxy.last_accessed = last_accessed_ts

436 -

437 - @property

438 - def content_hash(self):

439 - return self._proxy.content_hash

vllm/core/block/interfaces.py deleted

-319

@@ -1,319 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from abc import ABC, abstractmethod

5 -from typing import Dict, FrozenSet, List, Optional, Protocol, Tuple

6 -

7 -from vllm.utils import Device

8 -

9 -BlockId = int

10 -

11 -

12 -class Block(ABC):

13 -

14 - @abstractmethod

15 - def append_token_ids(self, token_ids: List[int]) -> None:

16 - pass

17 -

18 - @property

19 - @abstractmethod

20 - def block_id(self) -> Optional[int]:

21 - pass

22 -

23 - @block_id.setter

24 - @abstractmethod

25 - def block_id(self, value: Optional[int]) -> None:

26 - """NOTE: Do not use this API outside Block."""

27 - self._block_id = value

28 -

29 - @property

30 - @abstractmethod

31 - def token_ids(self) -> List[int]:

32 - pass

33 -

34 - @property

35 - @abstractmethod

36 - def num_tokens_total(self) -> int:

37 - """The number of tokens till the current block (inclusive)

38 - """

39 - pass

40 -

41 - @property

42 - @abstractmethod

43 - def num_empty_slots(self) -> int:

44 - pass

45 -

46 - @property

47 - @abstractmethod

48 - def is_full(self) -> bool:

49 - pass

50 -

51 - @property

52 - @abstractmethod

53 - def prev_block(self) -> Optional["Block"]:

54 - pass

55 -

56 - @property

57 - @abstractmethod

58 - def extra_hash(self) -> Optional[int]:

59 - return None

60 -

61 - @property

62 - @abstractmethod

63 - def computed(self) -> bool:

64 - raise NotImplementedError

65 -

66 - @computed.setter

67 - @abstractmethod

68 - def computed(self, value) -> bool:

69 - """Should be only used by PrefixCacingAllocator"""

70 - raise NotImplementedError

71 -

72 - @property

73 - @abstractmethod

74 - def last_accessed(self) -> float:

75 - raise NotImplementedError

76 -

77 - @last_accessed.setter

78 - @abstractmethod

79 - def last_accessed(self, last_accessed_ts: float):

80 - raise NotImplementedError

81 -

82 - class Factory(Protocol):

83 -

84 - @abstractmethod

85 - def __call__(

86 - self,

87 - prev_block: Optional["Block"],

88 - token_ids: List[int],

89 - block_size: int,

90 - allocator: "BlockAllocator",

91 - block_id: Optional[int] = None,

92 - computed: bool = False,

93 - extra_hash: Optional[int] = None,

94 - ) -> "Block":

95 - pass

96 -

97 - @property

98 - @abstractmethod

99 - def content_hash(self) -> Optional[int]:

100 - """Return the content-based hash of the current block, or None if it is

101 - not yet defined or not supported.

102 -

103 - For the content-based hash to be defined, the current block must be

104 - full.

105 - """

106 - return None

107 -

108 -

109 -class BlockAllocator(ABC):

110 -

111 - @abstractmethod

112 - def allocate_mutable_block(self, prev_block: Optional[Block],

113 - extra_hash: Optional[int]) -> Block:

114 - pass

115 -

116 - @abstractmethod

117 - def allocate_immutable_block(self, prev_block: Optional[Block],

118 - token_ids: List[int],

119 - extra_hash: Optional[int]) -> Block:

120 - pass

121 -

122 - @abstractmethod

123 - def allocate_immutable_blocks(self, prev_block: Optional[Block],

124 - block_token_ids: List[List[int]],

125 - extra_hash: Optional[int]) -> List[Block]:

126 - pass

127 -

128 - @abstractmethod

129 - def free(self, block: Block) -> None:

130 - pass

131 -

132 - @abstractmethod

133 - def fork(self, last_block: Block) -> List[Block]:

134 - pass

135 -

136 - @abstractmethod

137 - def get_num_total_blocks(self) -> int:

138 - pass

139 -

140 - @abstractmethod

141 - def get_num_free_blocks(self) -> int:

142 - pass

143 -

144 - @abstractmethod

145 - def get_physical_block_id(self, absolute_id: int) -> int:

146 - pass

147 -

148 - @abstractmethod

149 - def swap_out(self, blocks: List[Block]) -> None:

150 - pass

151 -

152 - @abstractmethod

153 - def swap_in(self, blocks: List[Block]) -> None:

154 - pass

155 -

156 - @property

157 - @abstractmethod

158 - def all_block_ids(self) -> FrozenSet[int]:

159 - pass

160 -

161 - @abstractmethod

162 - def clear_copy_on_writes(self) -> List[Tuple[int, int]]:

163 - pass

164 -

165 - @abstractmethod

166 - def mark_blocks_as_accessed(self, block_ids: List[int],

167 - now: float) -> None:

168 - pass

169 -

170 - @abstractmethod

171 - def mark_blocks_as_computed(self, block_ids: List[int]) -> None:

172 - pass

173 -

174 - @abstractmethod

175 - def get_common_computed_block_ids(

176 - self, computed_seq_block_ids: List[List[int]]) -> List[int]:

177 - pass

178 -

179 - @abstractmethod

180 - def cow_block_if_not_appendable(self, block: Block) -> BlockId:

181 - """NOTE: This should not be used besides Block"""

182 - pass

183 -

184 - @abstractmethod

185 - def promote_to_immutable_block(self, block: Block) -> BlockId:

186 - """NOTE: This should not be used besides Block"""

187 - pass

188 -

189 - @abstractmethod

190 - def get_num_full_blocks_touched(self, blocks: List[Block]) -> int:

191 - pass

192 -

193 - @abstractmethod

194 - def get_prefix_cache_hit_rate(self) -> float:

195 - """Prefix cache hit rate. -1 means not supported or disabled."""

196 - pass

197 -

198 - @abstractmethod

199 - def reset_prefix_cache(self) -> bool:

200 - """Reset prefix cache."""

201 - pass

202 -

203 - class NoFreeBlocksError(ValueError):

204 - pass

205 -

206 - @abstractmethod

207 - def find_cached_blocks_prefix(

208 - self,

209 - block_hashes: List[int],

210 - ) -> List[int]:

211 - pass

212 -

213 -

214 -class DeviceAwareBlockAllocator(ABC):

215 -

216 - @abstractmethod

217 - def allocate_mutable_block(self,

218 - prev_block: Optional[Block],

219 - device: Device,

220 - extra_hash: Optional[int] = None) -> Block:

221 - pass

222 -

223 - @abstractmethod

224 - def allocate_immutable_block(self,

225 - prev_block: Optional[Block],

226 - token_ids: List[int],

227 - device: Device,

228 - extra_hash: Optional[int] = None) -> Block:

229 - pass

230 -

231 - @abstractmethod

232 - def allocate_immutable_blocks(

233 - self,

234 - prev_block: Optional[Block],

235 - block_token_ids: List[List[int]],

236 - device: Device,

237 - extra_hash: Optional[int] = None,

238 - ) -> List[Block]:

239 - pass

240 -

241 - @abstractmethod

242 - def get_num_free_blocks(self, device: Device) -> int:

243 - pass

244 -

245 - @abstractmethod

246 - def get_num_total_blocks(self, device: Device) -> int:

247 - pass

248 -

249 - @abstractmethod

250 - def free(self, block: Block) -> None:

251 - pass

252 -

253 - @abstractmethod

254 - def fork(self, last_block: Block) -> List[Block]:

255 - pass

256 -

257 - @property

258 - @abstractmethod

259 - def all_block_ids(self) -> FrozenSet[int]:

260 - pass

261 -

262 - @abstractmethod

263 - def clear_copy_on_writes(self) -> List[Tuple[int, int]]:

264 - pass

265 -

266 - @abstractmethod

267 - def mark_blocks_as_accessed(self, block_ids: List[int],

268 - now: float) -> None:

269 - pass

270 -

271 - @abstractmethod

272 - def mark_blocks_as_computed(self, block_ids: List[int]) -> None:

273 - pass

274 -

275 - @abstractmethod

276 - def get_common_computed_block_ids(

277 - self, computed_seq_block_ids: List[List[int]]) -> List[int]:

278 - pass

279 -

280 - @abstractmethod

281 - def get_num_full_blocks_touched(self, blocks: List[Block],

282 - device: Device) -> int:

283 - pass

284 -

285 - @abstractmethod

286 - def swap(self, blocks: List[Block], src_device: Device,

287 - dst_device: Device) -> Dict[int, int]:

288 - pass

289 -

290 - @abstractmethod

291 - def get_physical_block_id(self, device: Device, absolute_id: int) -> int:

292 - pass

293 -

294 - @abstractmethod

295 - def allocate_or_get_null_block(self) -> Block:

296 - """

297 - Null blocks are used as a placeholders for KV cache blocks that have

298 - been dropped due to sliding window.

299 - There is at most one null block per allocator.

300 - """

301 - pass

302 -

303 - @abstractmethod

304 - def get_prefix_cache_hit_rate(self, device: Device) -> float:

305 - """Prefix cache hit rate. -1 means not supported or disabled."""

306 - pass

307 -

308 - @abstractmethod

309 - def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

310 - """Reset prefix cache."""

311 - pass

312 -

313 - @abstractmethod

314 - def find_cached_blocks_prefix(

315 - self,

316 - block_hashes: List[int],

317 - device: Device = Device.GPU,

318 - ) -> List[int]:

319 - pass

vllm/core/block/naive_block.py deleted

-466

@@ -1,466 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from collections import deque

5 -from typing import Deque, FrozenSet, Iterable, List, Optional, Tuple, Union

6 -

7 -from vllm.core.block.common import (BlockPool, CopyOnWriteTracker, RefCounter,

8 - get_all_blocks_recursively)

9 -from vllm.core.block.interfaces import Block, BlockAllocator, BlockId, Device

10 -

11 -Refcount = int

12 -

13 -

14 -class NaiveBlockAllocator(BlockAllocator):

15 - """A simple block allocator that manages blocks of memory without prefix

16 - caching.

17 -

18 - Args:

19 - create_block (Block.Factory): A factory function for creating new

20 - blocks. This is used when a NaiveBlockAllocator is composed within

21 - a prefix caching allocator -- the naive block allocator must

22 - construct prefix caching blocks (but shouldn't know anything else

23 - about them).

24 - num_blocks (int): The total number of blocks to manage.

25 - block_size (int): The size of each block in tokens.

26 - block_ids (Optional[Iterable[int]], optional): An optional iterable of

27 - block IDs. If not provided, block IDs will be assigned sequentially

28 - from 0 to num_blocks - 1.

29 - """

30 -

31 - def __init__(

32 - self,

33 - create_block: Block.Factory,

34 - num_blocks: int,

35 - block_size: int,

36 - block_ids: Optional[Iterable[int]] = None,

37 - block_pool: Optional[BlockPool] = None,

38 - ):

39 - if block_ids is None:

40 - block_ids = range(num_blocks)

41 -

42 - self._free_block_indices: Deque[BlockId] = deque(block_ids)

43 - self._all_block_indices = frozenset(block_ids)

44 - assert len(self._all_block_indices) == num_blocks

45 -

46 - self._refcounter = RefCounter(

47 - all_block_indices=self._free_block_indices)

48 - self._block_size = block_size

49 -

50 - self._cow_tracker = CopyOnWriteTracker(

51 - refcounter=self._refcounter.as_readonly())

52 -

53 - if block_pool is None:

54 - extra_factor = 4

55 - # Pre-allocate "num_blocks * extra_factor" block objects.

56 - # The "* extra_factor" is a buffer to allow more block objects

57 - # than physical blocks

58 - self._block_pool = BlockPool(self._block_size, create_block, self,

59 - num_blocks * extra_factor)

60 - else:

61 - # In this case, the block pool is provided by the caller,

62 - # which means that there is most likely a need to share

63 - # a block pool between allocators

64 - self._block_pool = block_pool

65 -

66 - def allocate_immutable_block(self,

67 - prev_block: Optional[Block],

68 - token_ids: List[int],

69 - extra_hash: Optional[int] = None,

70 - device: Optional[Device] = None) -> Block:

71 - """Allocates a new immutable block with the given token IDs, linked to

72 - the previous block.

73 -

74 - Args:

75 - prev_block (Optional[Block]): The previous block in the sequence. If

76 - None, then the block to be allocated is the first block in the

77 - sequence.

78 - token_ids (List[int]): The token IDs to be stored in the new block.

79 -

80 - Returns:

81 - Block: The newly allocated immutable block.

82 - """

83 - assert device is None

84 - block = self.allocate_mutable_block(prev_block=prev_block)

85 - block.append_token_ids(token_ids)

86 - return block

87 -

88 - def allocate_immutable_blocks(

89 - self,

90 - prev_block: Optional[Block],

91 - block_token_ids: List[List[int]],

92 - extra_hash: Optional[int] = None,

93 - device: Optional[Device] = None) -> List[Block]:

94 - assert device is None

95 - num_blocks = len(block_token_ids)

96 -

97 - block_ids = []

98 - for i in range(num_blocks):

99 - block_ids.append(self._allocate_block_id())

100 -

101 - blocks = []

102 - for i in range(num_blocks):

103 - prev_block = self._block_pool.init_block(

104 - prev_block=prev_block,

105 - token_ids=block_token_ids[i],

106 - block_size=self._block_size,

107 - physical_block_id=block_ids[i])

108 - blocks.append(prev_block)

109 -

110 - return blocks

111 -

112 - def allocate_mutable_block(self,

113 - prev_block: Optional[Block],

114 - extra_hash: Optional[int] = None,

115 - device: Optional[Device] = None) -> Block:

116 - """Allocates a new mutable block, linked to the previous block.

117 -

118 - Args:

119 - prev_block (Optional[Block]): The previous block in the sequence. If

120 - None, then the block to be allocated is the first block in the

121 - sequence.

122 -

123 - Returns:

124 - Block: The newly allocated mutable block.

125 - """

126 - assert device is None

127 - block_id = self._allocate_block_id()

128 - block = self._block_pool.init_block(prev_block=prev_block,

129 - token_ids=[],

130 - block_size=self._block_size,

131 - physical_block_id=block_id)

132 - return block

133 -

134 - def _allocate_block_id(self) -> BlockId:

135 - if not self._free_block_indices:

136 - raise BlockAllocator.NoFreeBlocksError()

137 -

138 - block_id = self._free_block_indices.popleft()

139 - self._refcounter.incr(block_id)

140 - return block_id

141 -

142 - def _free_block_id(self, block: Union[Block, BlockId]) -> None:

143 - if isinstance(block, Block):

144 - block_id = block.block_id

145 - block.block_id = None

146 - else:

147 - block_id = block

148 - assert block_id is not None

149 -

150 - refcount = self._refcounter.decr(block_id)

151 - if refcount == 0:

152 - self._free_block_indices.appendleft(block_id)

153 -

154 - def free(self, block: Block, keep_block_object: bool = False) -> None:

155 - # Release the physical block id

156 - self._free_block_id(block)

157 -

158 - # Release the block object

159 - if not keep_block_object:

160 - self._block_pool.free_block(block)

161 -

162 - def free_block_id(self, block_id: BlockId) -> None:

163 - self._free_block_id(block_id)

164 -

165 - def fork(self, last_block: Block) -> List[Block]:

166 - """Creates a new sequence of blocks that shares the same underlying

167 - memory as the original sequence.

168 -

169 - Args:

170 - last_block (Block): The last block in the original sequence.

171 -

172 - Returns:

173 - List[Block]: The new sequence of blocks that shares the same memory

174 - as the original sequence.

175 - """

176 - source_blocks = get_all_blocks_recursively(last_block)

177 -

178 - forked_blocks: List[Block] = []

179 - prev_block = None

180 - for block in source_blocks:

181 -

182 - # Increment refcount for each block.

183 - assert block.block_id is not None

184 - refcount = self._refcounter.incr(block.block_id)

185 - assert refcount != 1, "can't fork freed block"

186 -

187 - forked_block = self._block_pool.init_block(

188 - prev_block=prev_block,

189 - token_ids=block.token_ids,

190 - block_size=self._block_size,

191 - physical_block_id=block.block_id)

192 -

193 - forked_blocks.append(forked_block)

194 - prev_block = forked_blocks[-1]

195 -

196 - return forked_blocks

197 -

198 - def get_num_free_blocks(self) -> int:

199 - return len(self._free_block_indices)

200 -

201 - def get_num_total_blocks(self) -> int:

202 - return len(self._all_block_indices)

203 -

204 - def get_physical_block_id(self, absolute_id: int) -> int:

205 - """Returns the zero-offset block id on certain block allocator

206 - given the absolute block id.

207 -

208 - Args:

209 - absolute_id (int): The absolute block id for the block

210 - in whole allocator.

211 -

212 - Returns:

213 - int: The zero-offset block id on certain device.

214 - """

215 - return sorted(self._all_block_indices).index(absolute_id)

216 -

217 - @property

218 - def refcounter(self):

219 - return self._refcounter

220 -

221 - @property

222 - def all_block_ids(self) -> FrozenSet[int]:

223 - return self._all_block_indices

224 -

225 - def cow_block_if_not_appendable(self, block: Block) -> BlockId:

226 - """Performs a copy-on-write operation on the given block if it is not

227 - appendable.

228 -

229 - Args:

230 - block (Block): The block to check for copy-on-write.

231 -

232 - Returns:

233 - BlockId: The block index of the new block if a copy-on-write

234 - operation was performed, or the original block index if

235 - no copy-on-write was necessary.

236 - """

237 - src_block_id = block.block_id

238 - assert src_block_id is not None

239 -

240 - if self._cow_tracker.is_appendable(block):

241 - return src_block_id

242 -

243 - self._free_block_id(block)

244 - trg_block_id = self._allocate_block_id()

245 -

246 - self._cow_tracker.record_cow(src_block_id, trg_block_id)

247 -

248 - return trg_block_id

249 -

250 - def clear_copy_on_writes(self) -> List[Tuple[BlockId, BlockId]]:

251 - """Returns the copy-on-write source->destination mapping and clears it.

252 -

253 - Returns:

254 - List[Tuple[BlockId, BlockId]]: A list mapping source

255 - block indices to destination block indices.

256 - """

257 - return self._cow_tracker.clear_cows()

258 -

259 - def mark_blocks_as_accessed(self, block_ids: List[int],

260 - now: float) -> None:

261 - """Mark blocks as accessed, used in prefix caching.

262 -

263 - Since the naive allocator does not implement prefix caching, we do

264 - nothing.

265 - """

266 - pass

267 -

268 - def mark_blocks_as_computed(self, block_ids: List[int]) -> None:

269 - """Mark blocks as computed, used in prefix caching.

270 -

271 - Since the naive allocator does not implement prefix caching, we do

272 - nothing.

273 - """

274 - pass

275 -

276 - def get_common_computed_block_ids(

277 - self, computed_seq_block_ids: List[List[int]]) -> List[int]:

278 - """Determine blocks that can be skipped in prefill.

279 -

280 - Since the naive allocator does not support prefix caching, always return

281 - an empty list.

282 - """

283 - return []

284 -

285 - def promote_to_immutable_block(self, block: Block) -> BlockId:

286 - raise NotImplementedError("There is no promotion for naive blocks")

287 -

288 - def get_num_full_blocks_touched(self, blocks: List[Block]) -> int:

289 - """Returns the number of full blocks that will be touched by

290 - swapping in/out.

291 -

292 - Args:

293 - blocks: List of blocks to be swapped.

294 - Returns:

295 - int: the number of full blocks that will be touched by

296 - swapping in/out the given blocks. Non full blocks are ignored

297 - when deciding the number of blocks to touch.

298 - """

299 - # NOTE: for naive block, we use set to eliminate common blocks among

300 - # seqs, also we compare the empty slots in the mutable blocks with

301 - # lookahead slots to get the number of unique new block that are

302 - # needed.

303 - old_block_set = set()

304 - for block in blocks:

305 - if block.is_full:

306 - old_block_set.add(block)

307 - return len(old_block_set)

308 -

309 - def swap_out(self, blocks: List[Block]) -> None:

310 - for block in blocks:

311 - self._free_block_id(block)

312 -

313 - def swap_in(self, blocks: List[Block]) -> None:

314 - for block in blocks:

315 - # Here we allocate either immutable or mutable block and then

316 - # extract its block_id. Note that the block object is released

317 - # and the block_id is assigned to "block" to allow reusing the

318 - # existing "block" object

319 - if block.is_full:

320 - tmp_block = self.allocate_immutable_block(

321 - prev_block=block.prev_block, token_ids=block.token_ids)

322 - else:

323 - tmp_block = self.allocate_mutable_block(

324 - prev_block=block.prev_block)

325 - tmp_block.append_token_ids(block.token_ids)

326 -

327 - block_id = tmp_block.block_id

328 - tmp_block.block_id = None

329 - self._block_pool.free_block(tmp_block)

330 -

331 - block.block_id = block_id # Assign block_id

332 -

333 - def get_prefix_cache_hit_rate(self) -> float:

334 - return -1

335 -

336 - def reset_prefix_cache(self) -> bool:

337 - """No prefix cache for naive block allocator."""

338 - return True

339 -

340 - def find_cached_blocks_prefix(self, block_hashes: List[int]) -> List[int]:

341 - # Not applicable for naive block allocator.

342 - return []

343 -

344 -

345 -class NaiveBlock(Block):

346 - """An implementation of the Block class that does not support prefix

347 - caching.

348 -

349 - The NaiveBlock class represents a block of token IDs with a fixed size. It

350 - provides methods for appending token IDs to the block and manages copy-on

351 - -write operations when necessary.

352 -

353 - Args:

354 - prev_block (Block): The previous block in the sequence.

355 - token_ids (List[int]): The initial token IDs to be stored in the block.

356 - block_size (int): The maximum number of token IDs that can be stored in

357 - the block.

358 - allocator (BlockAllocator): The block allocator associated with this

359 - block.

360 - block_id (Optional[int], optional): The physical block index

361 - of this block. Defaults to None, which means no allocation has been

362 - made.

363 - _cow_target (Optional[Block], optional): The copy-on-write target block.

364 - If not provided, it defaults to self.

365 - """

366 -

367 - def __init__(self,

368 - prev_block: Optional[Block],

369 - token_ids: List[int],

370 - block_size: int,

371 - allocator: BlockAllocator,

372 - block_id: Optional[int] = None,

373 - _cow_target: Optional[Block] = None,

374 - extra_hash: Optional[int] = None):

375 - self._token_ids: List[int] = []

376 - self._block_size = block_size

377 - self._prev_block = prev_block

378 - self._block_id = block_id

379 - self._allocator = allocator

380 - self._cow_target = _cow_target if _cow_target is not None else self

381 -

382 - self._append_token_ids_no_cow(token_ids)

383 -

384 - def append_token_ids(self, token_ids: List[int]) -> None:

385 - """Appends the given token IDs to the block and performs a

386 - copy-on-write if necessary.

387 -

388 - Args:

389 - token_ids (Optional[List[int]]): The token IDs to be appended

390 - to the block.

391 - """

392 - self._append_token_ids_no_cow(token_ids)

393 -

394 - if self._block_id is not None:

395 - self._block_id = (self._allocator.cow_block_if_not_appendable(

396 - self._cow_target))

397 -

398 - def _append_token_ids_no_cow(self, token_ids: List[int]) -> None:

399 - """Appends the given token IDs to the block

400 -

401 - Args:

402 - token_ids (List[int]): The token IDs to be appended to the block.

403 - """

404 - if len(token_ids) == 0:

405 - return

406 -

407 - assert len(token_ids) <= self.num_empty_slots

408 -

409 - self._token_ids.extend(token_ids)

410 -

411 - @property

412 - def computed(self) -> bool:

413 - raise NotImplementedError

414 -

415 - @computed.setter

416 - def computed(self, value) -> None:

417 - raise NotImplementedError

418 -

419 - @property

420 - def last_accessed(self) -> float:

421 - raise NotImplementedError

422 -

423 - @last_accessed.setter

424 - def last_accessed(self, last_accessed_ts: float):

425 - raise NotImplementedError

426 -

427 - @property

428 - def block_id(self) -> Optional[int]:

429 - return self._block_id

430 -

431 - @block_id.setter

432 - def block_id(self, value: Optional[int]) -> None:

433 - self._block_id = value

434 -

435 - @property

436 - def is_full(self) -> bool:

437 - return self.num_empty_slots == 0

438 -

439 - @property

440 - def num_empty_slots(self) -> int:

441 - return self._block_size - len(self.token_ids)

442 -

443 - @property

444 - def token_ids(self) -> List[int]:

445 - return self._token_ids

446 -

447 - @property

448 - def num_tokens_total(self) -> int:

449 - raise NotImplementedError(

450 - "num_tokens_total is not used for naive block")

451 -

452 - @property

453 - def block_size(self) -> int:

454 - return self._block_size

455 -

456 - @property

457 - def prev_block(self) -> Optional["Block"]:

458 - return self._prev_block

459 -

460 - @property

461 - def extra_hash(self):

462 - return None

463 -

464 - @property

465 - def content_hash(self) -> Optional[int]:

466 - return None

vllm/core/block/prefix_caching_block.py deleted

-1135

@@ -1,1135 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""Token blocks."""

4 -import sys

5 -from bisect import bisect_left

6 -from os.path import commonprefix

7 -from typing import (Callable, Dict, FrozenSet, Iterable, List, Optional, Set,

8 - Tuple)

9 -

10 -from vllm.core.block.common import (CacheMetricData, CopyOnWriteTracker,

11 - get_all_blocks_recursively)

12 -from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId, Device,

13 - DeviceAwareBlockAllocator)

14 -from vllm.core.block.naive_block import (BlockPool, NaiveBlock,

15 - NaiveBlockAllocator)

16 -from vllm.core.evictor import EvictionPolicy, Evictor, make_evictor

17 -from vllm.logger import init_logger

18 -from vllm.sequence import Sequence

19 -

20 -PrefixHash = int

21 -

22 -# By default, we init our block access time as _DEFAULT_LAST_ACCESSED_TIME

23 -# so that if we find one block is still hold _DEFAULT_LAST_ACCESSED_TIME,

24 -# then we know this block hasn't been accessed yet.

25 -_DEFAULT_LAST_ACCESSED_TIME = -1

26 -

27 -logger = init_logger(__name__)

28 -

29 -

30 -class BlockTracker:

31 - """Used to track the status of a block inside the prefix caching allocator

32 - """

33 - __slots__ = ("active", "last_accessed", "computed")

34 -

35 - def reset(self):

36 - self.last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME

37 - self.computed: bool = False

38 -

39 - def __init__(self):

40 - self.active: bool = False

41 - self.reset()

42 -

43 - def enable(self):

44 - assert not self.active

45 - self.active = True

46 - self.reset()

47 -

48 - def disable(self):

49 - assert self.active

50 - self.active = False

51 - self.reset()

52 -

53 -

54 -class PrefixCachingBlockAllocator(BlockAllocator):

55 - """A block allocator that implements prefix caching.

56 -

57 - The PrefixCachingBlockAllocator maintains a cache of blocks based on their

58 - content hash. It reuses blocks with the same content hash to avoid redundant

59 - memory allocation. The allocator also supports copy-on-write operations.

60 -

61 - Args:

62 - num_blocks (int): The total number of blocks to manage.

63 - block_size (int): The size of each block in tokens.

64 - block_ids (Optional[Iterable[int]], optional): An optional iterable of

65 - block IDs. If not provided, block IDs will be assigned sequentially

66 - from 0 to num_blocks - 1.

67 - """

68 -

69 - # Note that we use 'None' as a string here instead of None because

70 - # as of Python 3.12, hash(None) returns a constant predictable value.

71 - # This could possibly make it easier to find and exploit hash

72 - # collisions. 'None' as a string will be hashed differently per process,

73 - # but consistently within the same process. This is the same as the

74 - # behavior of None prior to Python 3.12.

75 - _none_hash: int = hash('None')

76 -

77 - # Implements Block.Factory.

78 - def __init__(

79 - self,

80 - num_blocks: int,

81 - block_size: int,

82 - block_ids: Optional[Iterable[int]] = None,

83 - eviction_policy: EvictionPolicy = EvictionPolicy.LRU,

84 - ):

85 - if block_ids is None:

86 - block_ids = range(num_blocks)

87 -

88 - self._block_size = block_size

89 -

90 - # A mapping of prefix hash to block index. All blocks which have a

91 - # prefix hash will be in this dict, even if they have refcount 0.

92 - self._cached_blocks: Dict[PrefixHash, BlockId] = {}

93 -

94 - # A list of immutable block IDs that have been touched by scheduler

95 - # and should be marked as computed after an entire batch of sequences

96 - # are scheduled.

97 - self._touched_blocks: Set[BlockId] = set()

98 -

99 - # Used to track status of each physical block id

100 - self._block_tracker: Dict[BlockId, BlockTracker] = {}

101 - for block_id in block_ids:

102 - self._block_tracker[block_id] = BlockTracker()

103 -

104 - # Pre-allocate "num_blocks * extra_factor" block objects.

105 - # The "* extra_factor" is a buffer to allow more block objects

106 - # than physical blocks

107 - extra_factor = 4

108 - self._block_pool = BlockPool(self._block_size, self._create_block,

109 - self, num_blocks * extra_factor)

110 -

111 - # An allocator for blocks that do not have prefix hashes.

112 - self._hashless_allocator = NaiveBlockAllocator(

113 - create_block=self._create_block, # type: ignore

114 - num_blocks=num_blocks,

115 - block_size=block_size,

116 - block_ids=block_ids,

117 - block_pool=self._block_pool, # Share block pool here

118 - )

119 -

120 - # Evitor used to maintain how we want to handle those computed blocks

121 - # if we find memory pressure is high.

122 - self.eviction_policy = eviction_policy

123 - self.evictor: Evictor = make_evictor(self.eviction_policy)

124 -

125 - # We share the refcounter between allocators. This allows us to promote

126 - # blocks originally allocated in the hashless allocator to immutable

127 - # blocks.

128 - self._refcounter = self._hashless_allocator.refcounter

129 -

130 - self._cow_tracker = CopyOnWriteTracker(

131 - refcounter=self._refcounter.as_readonly())

132 -

133 - self.metric_data = CacheMetricData()

134 -

135 - def _create_block(

136 - self,

137 - prev_block: Optional[Block],

138 - token_ids: List[int],

139 - block_size: int,

140 - allocator: BlockAllocator,

141 - block_id: Optional[int] = None,

142 - computed: bool = False,

143 - extra_hash: Optional[int] = None,

144 - ) -> Block:

145 - # Bind block to self.

146 - allocator = self

147 -

148 - return PrefixCachingBlock(

149 - prev_block=prev_block,

150 - token_ids=token_ids,

151 - block_size=block_size,

152 - block_id=block_id,

153 - allocator=allocator,

154 - computed=computed,

155 - extra_hash=extra_hash,

156 - )

157 -

158 - def allocate_immutable_block(self,

159 - prev_block: Optional[Block],

160 - token_ids: List[int],

161 - extra_hash: Optional[int] = None,

162 - device: Optional[Device] = None) -> Block:

163 - """Allocates an immutable block with the given token IDs, reusing cached

164 - blocks if possible.

165 -

166 - Args:

167 - prev_block (Optional[Block]): The previous block in the sequence.

168 - token_ids (List[int]): The token IDs to be stored in the block.

169 -

170 - Returns:

171 - Block: The allocated immutable block.

172 - """

173 - assert device is None

174 - assert_prefix_caching_block_or_none(prev_block)

175 -

176 - # First, try to create a block that points to cached data

177 - block = self._block_pool.init_block(prev_block=prev_block,

178 - token_ids=token_ids,

179 - block_size=self._block_size,

180 - physical_block_id=None,

181 - extra_hash=extra_hash)

182 - assert block.content_hash is not None

183 -

184 - cached_block_id = self._cached_blocks.get(block.content_hash, None)

185 - if cached_block_id is not None:

186 - self.metric_data.query(hit=True)

187 - block.block_id = cached_block_id

188 - self._incr_refcount_cached_block(block)

189 - return block

190 - self.metric_data.query(hit=False)

191 - self._block_pool.free_block(block)

192 -

193 - # No cached block => Allocate a new block

194 - block = self.allocate_mutable_block(prev_block, extra_hash=extra_hash)

195 - block.append_token_ids(token_ids)

196 - return block

197 -

198 - def allocate_immutable_blocks(

199 - self,

200 - prev_block: Optional[Block],

201 - block_token_ids: List[List[int]],

202 - extra_hash: Optional[int] = None,

203 - device: Optional[Device] = None) -> List[Block]:

204 - blocks = []

205 - for token_ids in block_token_ids:

206 - prev_block = self.allocate_immutable_block(prev_block=prev_block,

207 - token_ids=token_ids,

208 - device=device,

209 - extra_hash=extra_hash)

210 - blocks.append(prev_block)

211 - return blocks

212 -

213 - def allocate_mutable_block(self,

214 - prev_block: Optional[Block],

215 - extra_hash: Optional[int] = None,

216 - device: Optional[Device] = None) -> Block:

217 - """Allocates a mutable block. If there are no free blocks, this will

218 - evict unused cached blocks.

219 -

220 - Args:

221 - prev_block (Block): The previous block in the sequence.

222 - None is not allowed unlike it is super class.

223 -

224 - Returns:

225 - Block: The allocated mutable block.

226 - """

227 - assert device is None

228 - assert_prefix_caching_block_or_none(prev_block)

229 -

230 - block_id = self._allocate_block_id()

231 - block = self._block_pool.init_block(prev_block=prev_block,

232 - token_ids=[],

233 - block_size=self._block_size,

234 - physical_block_id=block_id,

235 - extra_hash=extra_hash)

236 - assert not block.computed

237 - assert block.content_hash is None

238 - return block

239 -

240 - def _incr_refcount_cached_block(self, block: Block) -> None:

241 - # Set this block to be "computed" since it is pointing to a

242 - # cached block id (which was already computed)

243 - block.computed = True

244 -

245 - block_id = block.block_id

246 - assert block_id is not None

247 -

248 - refcount = self._refcounter.incr(block_id)

249 - if refcount == 1:

250 - # In case a cached block was evicted, restore its tracking

251 - if block_id in self.evictor:

252 - self.evictor.remove(block_id)

253 -

254 - self._track_block_id(block_id, computed=True)

255 -

256 - def _decr_refcount_cached_block(self, block: Block) -> None:

257 - # Ensure this is immutable/cached block

258 - assert block.content_hash is not None

259 -

260 - block_id = block.block_id

261 - assert block_id is not None

262 -

263 - refcount = self._refcounter.decr(block_id)

264 - if refcount > 0:

265 - block.block_id = None

266 - return

267 - else:

268 - assert refcount == 0

269 -

270 - # No longer used

271 - assert block.content_hash in self._cached_blocks

272 -

273 - # Add the cached block to the evictor

274 - # (This keeps the cached block around so it can be reused)

275 - self.evictor.add(block_id, block.content_hash, block.num_tokens_total,

276 - self._block_tracker[block_id].last_accessed)

277 -

278 - # Stop tracking the block

279 - self._untrack_block_id(block_id)

280 -

281 - block.block_id = None

282 -

283 - def _decr_refcount_hashless_block(self, block: Block) -> None:

284 - block_id = block.block_id

285 - assert block_id is not None

286 -

287 - # We may have a fork case where block is shared,

288 - # in which case, we cannot remove it from tracking

289 - refcount = self._refcounter.get(block_id)

290 - if refcount == 1:

291 - self._untrack_block_id(block_id)

292 -

293 - # Decrement refcount of the block_id, but do not free the block object

294 - # itself (will be handled by the caller)

295 - self._hashless_allocator.free(block, keep_block_object=True)

296 -

297 - def _allocate_block_id(self) -> BlockId:

298 - """First tries to allocate a block id from the hashless allocator,

299 - and if there are no blocks, then tries to evict an unused cached block.

300 - """

301 - hashless_block_id = self._maybe_allocate_hashless_block_id()

302 - if hashless_block_id is not None:

303 - return hashless_block_id

304 -

305 - evicted_block_id = self._maybe_allocate_evicted_block_id()

306 - if evicted_block_id is not None:

307 - return evicted_block_id

308 -

309 - # No block available in hashless allocator, nor in unused cache blocks.

310 - raise BlockAllocator.NoFreeBlocksError()

311 -

312 - def _maybe_allocate_hashless_block_id(self) -> Optional[BlockId]:

313 - try:

314 - # Allocate mutable block and extract its block_id

315 - block = self._hashless_allocator.allocate_mutable_block(

316 - prev_block=None)

317 - block_id = block.block_id

318 - self._block_pool.free_block(block)

319 -

320 - self._track_block_id(block_id, computed=False)

321 - return block_id

322 - except BlockAllocator.NoFreeBlocksError:

323 - return None

324 -

325 - def _maybe_allocate_evicted_block_id(self) -> Optional[BlockId]:

326 - if self.evictor.num_blocks == 0:

327 - return None

328 -

329 - # Here we get an evicted block, which is only added

330 - # into evictor if its ref counter is 0

331 - # and since its content would be changed, we need

332 - # to remove it from _cached_blocks's tracking list

333 - block_id, content_hash_to_evict = self.evictor.evict()

334 -

335 - # Sanity checks

336 - assert content_hash_to_evict in self._cached_blocks

337 - _block_id = self._cached_blocks[content_hash_to_evict]

338 - assert self._refcounter.get(_block_id) == 0

339 - assert _block_id == block_id

340 -

341 - self._cached_blocks.pop(content_hash_to_evict)

342 -

343 - self._refcounter.incr(block_id)

344 - self._track_block_id(block_id, computed=False)

345 -

346 - return block_id

347 -

348 - def _free_block_id(self, block: Block) -> None:

349 - """Decrements the refcount of the block. The block may be in two

350 - possible states: (1) immutable/cached or (2) mutable/hashless.

351 - In the first case, the refcount is decremented directly and the block

352 - may be possibly added to the evictor. In other case, hashless

353 - allocator free(..) with keep_block_object=True is called to only free

354 - the block id (since the block object may be reused by the caller)

355 - """

356 - block_id = block.block_id

357 - assert block_id is not None, "Freeing unallocated block is undefined"

358 -

359 - if block.content_hash is not None:

360 - # Immutable: This type of block is always cached, and we want to

361 - # keep it in the evictor for future reuse

362 - self._decr_refcount_cached_block(block)

363 - else:

364 - # Mutable: This type of block is not cached, so we release it

365 - # directly to the hashless allocator

366 - self._decr_refcount_hashless_block(block)

367 -

368 - assert block.block_id is None

369 -

370 - def free(self, block: Block, keep_block_object: bool = False) -> None:

371 - """Release the block (look at free_block_id(..) docs)

372 - """

373 - # Release the physical block index

374 - self._free_block_id(block)

375 -

376 - # Release the block object to the pool

377 - if not keep_block_object:

378 - self._block_pool.free_block(block)

379 -

380 - def fork(self, last_block: Block) -> List[Block]:

381 - """Creates a new sequence of blocks that shares the same underlying

382 - memory as the original sequence.

383 -

384 - Args:

385 - last_block (Block): The last block in the original sequence.

386 -

387 - Returns:

388 - List[Block]: The new sequence of blocks that shares the same memory

389 - as the original sequence.

390 - """

391 - source_blocks = get_all_blocks_recursively(last_block)

392 -

393 - forked_blocks: List[Block] = []

394 - prev_block = None

395 - for block in source_blocks:

396 - block_id = block.block_id

397 - assert block_id is not None

398 -

399 - refcount = self._refcounter.incr(block_id)

400 - assert refcount != 1, "can't fork free'd block_id = {}".format(

401 - block_id)

402 -

403 - forked_block = self._block_pool.init_block(

404 - prev_block=prev_block,

405 - token_ids=block.token_ids,

406 - block_size=self._block_size,

407 - physical_block_id=block_id,

408 - extra_hash=block.extra_hash)

409 -

410 - forked_blocks.append(forked_block)

411 - prev_block = forked_blocks[-1]

412 -

413 - return forked_blocks

414 -

415 - def get_num_free_blocks(self, device: Optional[Device] = None) -> int:

416 - assert device is None

417 - # The number of free blocks is the number of hashless free blocks

418 - # plus the number of blocks evictor could free from its list.

419 - return self._hashless_allocator.get_num_free_blocks(

420 - ) + self.evictor.num_blocks

421 -

422 - def get_num_total_blocks(self) -> int:

423 - return self._hashless_allocator.get_num_total_blocks()

424 -

425 - def get_physical_block_id(self, absolute_id: int) -> int:

426 - """Returns the zero-offset block id on certain block allocator

427 - given the absolute block id.

428 -

429 - Args:

430 - absolute_id (int): The absolute block id for the block

431 - in whole allocator.

432 -

433 - Returns:

434 - int: The rzero-offset block id on certain device.

435 - """

436 - return sorted(self.all_block_ids).index(absolute_id)

437 -

438 - @property

439 - def all_block_ids(self) -> FrozenSet[int]:

440 - return self._hashless_allocator.all_block_ids

441 -

442 - def get_prefix_cache_hit_rate(self) -> float:

443 - return self.metric_data.get_hit_rate()

444 -

445 - def reset_prefix_cache(self) -> bool:

446 - """Reset prefix cache. This function may be used in RLHF

447 - flows to invalid prefix caching after the weights are updated,

448 - or used for resetting prefix caching status for benchmarking.

449 -

450 - Returns:

451 - bool: True if the prefix cache is successfully reset,

452 - False otherwise.

453 - """

454 - num_used_blocks = (self.get_num_total_blocks() -

455 - self.get_num_free_blocks())

456 - if num_used_blocks > 0:

457 - logger.warning(

458 - "Failed to reset prefix cache because some "

459 - "blocks (%d) are not freed yet", num_used_blocks)

460 - return False

461 -

462 - # Free all blocks in the evictor.

463 - while (block_id :=

464 - self._maybe_allocate_evicted_block_id()) is not None:

465 - self._hashless_allocator.free_block_id(block_id)

466 -

467 - # Should not have any cached blocks because all blocks are evicted.

468 - assert not self._cached_blocks

469 -

470 - # Reset the evictor.

471 - self.evictor = make_evictor(self.eviction_policy)

472 -

473 - # Reset the block tracker.

474 - for block_id in self._block_tracker:

475 - self._block_tracker[block_id] = BlockTracker()

476 -

477 - # Reset the metrics.

478 - self.metric_data = CacheMetricData()

479 -

480 - logger.info("Successfully reset prefix cache")

481 - return True

482 -

483 - def is_block_cached(self, block: Block) -> bool:

484 - assert block.content_hash is not None

485 - return block.content_hash in self._cached_blocks

486 -

487 - def promote_to_immutable_block(self, block: Block) -> BlockId:

488 - """Once a mutable block is full, it can be promoted to an immutable

489 - block. This means that its content can be referenced by future blocks

490 - having the same prefix.

491 -

492 - Note that if we already have a cached block with the same content, we

493 - will replace the newly-promoted block's mapping with the existing cached

494 - block id.

495 -

496 - Args:

497 - block: The mutable block to be promoted.

498 -

499 - Returns:

500 - BlockId: Either the original block index, or the block index of

501 - the previously cached block matching the same content.

502 - """

503 - # Ensure block can be promoted

504 - assert block.content_hash is not None

505 - assert block.block_id is not None

506 - assert self._refcounter.get(block.block_id) > 0

507 -

508 - if block.content_hash not in self._cached_blocks:

509 - # No cached content hash => Set this block as cached.

510 - # Note that this block cannot be marked as computed yet

511 - # because other sequences in the same batch cannot reuse

512 - # this block.

513 - self._cached_blocks[block.content_hash] = block.block_id

514 - # Mark this block as touched so that it can be marked as

515 - # computed after the entire batch of sequences are scheduled.

516 - self._touched_blocks.add(block.block_id)

517 - return block.block_id

518 -

519 - # Reuse the cached content hash

520 - self._decr_refcount_hashless_block(block)

521 - block.block_id = self._cached_blocks[block.content_hash]

522 -

523 - # Increment refcount of the cached block and (possibly) restore

524 - # it from the evictor.

525 - # Note that in this case, the block is marked as computed

526 - self._incr_refcount_cached_block(block)

527 -

528 - return block.block_id

529 -

530 - def cow_block_if_not_appendable(self, block: Block) -> BlockId:

531 - """Performs a copy-on-write operation on the given block if it is not

532 - appendable.

533 -

534 - Args:

535 - block (Block): The block to check for copy-on-write.

536 -

537 - Returns:

538 - BlockId: The block index of the new block if a copy-on-write

539 - operation was performed, or the original block index if

540 - no copy-on-write was necessary.

541 - """

542 - src_block_id = block.block_id

543 - assert src_block_id is not None

544 -

545 - if self._cow_tracker.is_appendable(block):

546 - return src_block_id

547 -

548 - self._free_block_id(block)

549 - trg_block_id = self._allocate_block_id()

550 -

551 - self._cow_tracker.record_cow(src_block_id, trg_block_id)

552 -

553 - return trg_block_id

554 -

555 - def clear_copy_on_writes(self) -> List[Tuple[BlockId, BlockId]]:

556 - """Returns the copy-on-write source->destination mapping and clears it.

557 -

558 - Returns:

559 - List[Tuple[BlockId, BlockId]]: A list mapping source

560 - block indices to destination block indices.

561 - """

562 - return self._cow_tracker.clear_cows()

563 -

564 - def mark_blocks_as_accessed(self, block_ids: List[int],

565 - now: float) -> None:

566 - """Mark blocks as accessed, used in prefix caching.

567 -

568 - If the block is added into evictor, we need to update corresponding

569 - info in evictor's metadata.

570 - """

571 -

572 - for block_id in block_ids:

573 - if self._block_tracker[block_id].active:

574 - self._block_tracker[block_id].last_accessed = now

575 - elif block_id in self.evictor:

576 - self.evictor.update(block_id, now)

577 - else:

578 - raise ValueError(

579 - "Mark block as accessed which is not belonged to GPU")

580 -

581 - def mark_blocks_as_computed(self, block_ids: List[int]) -> None:

582 - # Mark all touched blocks as computed.

583 - for block_id in self._touched_blocks:

584 - self._block_tracker[block_id].computed = True

585 - self._touched_blocks.clear()

586 -

587 - def _track_block_id(self, block_id: Optional[BlockId],

588 - computed: bool) -> None:

589 - assert block_id is not None

590 - self._block_tracker[block_id].enable()

591 - self._block_tracker[block_id].computed = computed

592 -

593 - def _untrack_block_id(self, block_id: Optional[BlockId]) -> None:

594 - assert block_id is not None

595 - self._block_tracker[block_id].disable()

596 -

597 - def block_is_computed(self, block_id: int) -> bool:

598 - if self._block_tracker[block_id].active:

599 - return self._block_tracker[block_id].computed

600 - else:

601 - return block_id in self.evictor

602 -

603 - def get_common_computed_block_ids(

604 - self, computed_seq_block_ids: List[List[int]]) -> List[int]:

605 - """Return the block ids that are common for a given sequence group.

606 -

607 - Only those blocks that are immutable and already be marked

608 - compyted would be taken consideration.

609 - """

610 -

611 - # NOTE We exclude the last block to avoid the case where the entire

612 - # prompt is cached. This would cause erroneous behavior in model

613 - # runner.

614 -

615 - # It returns a list of int although type annotation says list of string.

616 - if len(computed_seq_block_ids) == 1:

617 - return computed_seq_block_ids[0]

618 -

619 - return commonprefix([

620 - ids for ids in computed_seq_block_ids # type: ignore

621 - if ids

622 - ])

623 -

624 - def get_num_full_blocks_touched(self, blocks: List[Block]) -> int:

625 - """Returns the number of full blocks that will be touched by

626 - swapping in/out.

627 -

628 - Args:

629 - blocks: List of blocks to be swapped.

630 - Returns:

631 - int: the number of full blocks that will be touched by

632 - swapping in/out the given blocks. Non full blocks are ignored

633 - when deciding the number of blocks to touch.

634 - """

635 - num_touched_blocks: int = 0

636 - for block in blocks:

637 - # If the block has a match in the cache and the cached

638 - # block is not referenced, then we still count it as a

639 - # touched block

640 - if block.is_full and (not self.is_block_cached(block) or \

641 - (block.content_hash is not None and \

642 - self._cached_blocks[block.content_hash] in \

643 - self.evictor)):

644 - num_touched_blocks += 1

645 - return num_touched_blocks

646 -

647 - def swap_out(self, blocks: List[Block]) -> None:

648 - """Execute the swap out actions. Basically just free the

649 - given blocks.

650 -

651 - Args:

652 - blocks: List of blocks to be swapped out.

653 - """

654 - for block in blocks:

655 - self._free_block_id(block)

656 -

657 - def swap_in(self, blocks: List[Block]) -> None:

658 - """Execute the swap in actions. Change the block id from

659 - old allocator to current allocator for each block to finish

660 - the block table update.

661 -

662 - Args:

663 - blocks: List of blocks to be swapped in.

664 - """

665 - for block in blocks:

666 - # Here we allocate either immutable or mutable block and then

667 - # extract its block_id. Note that the block object is released

668 - # and the block_id is assigned to "block" to allow reusing the

669 - # existing "block" object

670 - if block.is_full:

671 - tmp_block = self.allocate_immutable_block(

672 - prev_block=block.prev_block,

673 - token_ids=block.token_ids,

674 - extra_hash=block.extra_hash)

675 - else:

676 - tmp_block = self.allocate_mutable_block(

677 - prev_block=block.prev_block, extra_hash=block.extra_hash)

678 - tmp_block.append_token_ids(block.token_ids)

679 -

680 - block_id = tmp_block.block_id

681 - self._block_pool.free_block(tmp_block)

682 -

683 - block.block_id = block_id # Assign block_id

684 -

685 - def find_cached_blocks_prefix(self, block_hashes: List[int]) -> List[int]:

686 - """

687 - Given a list of block hashes, return the prefix of the block hashes that

688 - are all cached.

689 -

690 - Since a block's block hash includes the hashes of all previous blocks,

691 - and we only allocate/deallocate blocks in the entire sequence, so if a

692 - block is cached, then all previous blocks are also cached. With this

693 - property, we can use binary search to find the prefix of cached blocks.

694 -

695 - Args:

696 - block_hashes (List[int]): The list of block hashes.

697 -

698 - Returns:

699 - List[int]: The prefix of the `block_hashes` that are cached.

700 - """

701 -

702 - def _block_is_cached(block_hash: PrefixHash) -> bool:

703 - if block_hash not in self._cached_blocks:

704 - return False

705 -

706 - cached_block_id = self._cached_blocks[block_hash]

707 - # We only consider the blocks that are marked as computed.

708 - return self.block_is_computed(cached_block_id)

709 -

710 - def _bisect_left(a, x, key: Callable[[PrefixHash], bool]) -> int:

711 -

712 - # python <= 3.10 don't have the key argument

713 - if sys.version_info < (3, 10):

714 - a = [key(e) for e in a]

715 - return bisect_left(a, x)

716 - else:

717 - return bisect_left(a, x, key=key)

718 -

719 - # Look for the first block that's not cached, and returns the prefix

720 - # i.e. blocks that are cached.

721 - idx = _bisect_left(block_hashes,

722 - True,

723 - key=lambda x: not _block_is_cached(x))

724 - return block_hashes[:idx]

725 -

726 -

727 -class PrefixCachingBlock(Block):

728 - """A block implementation that supports prefix caching.

729 -

730 - The PrefixCachingBlock class represents a block of token IDs with prefix

731 - caching capabilities. It wraps a NaiveBlock internally and provides

732 - additional functionality for content hashing and promoting immutable blocks

733 - with the prefix caching allocator.

734 -

735 - Args:

736 - prev_block (Optional[PrefixCachingBlock]): The previous block in the

737 - sequence.

738 - token_ids (List[int]): The initial token IDs to be stored in the block.

739 - block_size (int): The maximum number of token IDs that can be stored in

740 - the block.

741 - allocator (BlockAllocator): The prefix

742 - caching block allocator associated with this block.

743 - block_id (Optional[int], optional): The physical block index

744 - of this block. Defaults to None.

745 - extra_hash (Optional[int]): The hash value of additional factors

746 - such as adapters that influence the block, apart from the token_ids.

747 - """

748 -

749 - # Note that we use 'None' as a string here instead of None because

750 - # as of Python 3.12, hash(None) returns a constant predictable value.

751 - # This could possibly make it easier to find and exploit hash

752 - # collisions. 'None' as a string will be hashed differently per process,

753 - # but consistently within the same process. This is the same as the

754 - # behavior of None prior to Python 3.12.

755 - _none_hash: int = hash('None')

756 -

757 - def __init__(

758 - self,

759 - prev_block: Optional[Block],

760 - token_ids: List[int],

761 - block_size: int,

762 - allocator: BlockAllocator,

763 - block_id: Optional[int] = None,

764 - computed: bool = False,

765 - extra_hash: Optional[int] = None,

766 - ):

767 - assert isinstance(allocator, PrefixCachingBlockAllocator), (

768 - "Currently this class is only tested with "

769 - "PrefixCachingBlockAllocator. Got instead allocator = {}".format(

770 - allocator))

771 - assert_prefix_caching_block_or_none(prev_block)

772 -

773 - self._prev_block = prev_block

774 - self._cached_content_hash: Optional[int] = None

775 - self._cached_num_tokens_total: int = 0

776 - self._allocator = allocator

777 - self._last_accessed: float = _DEFAULT_LAST_ACCESSED_TIME

778 - self._computed = computed

779 - self._extra_hash = extra_hash

780 -

781 - # On the first time, we create the block object, and next we only

782 - # reinitialize it

783 - if hasattr(self, "_block"):

784 - self._block.__init__( # type: ignore[has-type]

785 - prev_block=prev_block,

786 - token_ids=token_ids,

787 - block_size=block_size,

788 - block_id=block_id,

789 - allocator=self._allocator)

790 - else:

791 - self._block = NaiveBlock(prev_block=prev_block,

792 - token_ids=token_ids,

793 - block_size=block_size,

794 - block_id=block_id,

795 - allocator=self._allocator)

796 -

797 - self._update_num_tokens_total()

798 -

799 - def _update_num_tokens_total(self):

800 - """Incrementally computes the number of tokens that there is

801 - till the current block (included)

802 - """

803 - res = 0

804 -

805 - # Add all previous blocks

806 - if self._prev_block is not None:

807 - res += self._prev_block.num_tokens_total

808 -

809 - # Add current block

810 - res += len(self.token_ids)

811 -

812 - self._cached_num_tokens_total = res

813 -

814 - @property

815 - def computed(self) -> bool:

816 - return self._computed

817 -

818 - @computed.setter

819 - def computed(self, value) -> None:

820 - self._computed = value

821 -

822 - @property

823 - def last_accessed(self) -> float:

824 - return self._last_accessed

825 -

826 - @last_accessed.setter

827 - def last_accessed(self, last_accessed_ts: float):

828 - self._last_accessed = last_accessed_ts

829 -

830 - def append_token_ids(self, token_ids: List[int]) -> None:

831 - """Appends the given token IDs to the block and registers the block as

832 - immutable if the block becomes full.

833 -

834 - Args:

835 - token_ids (List[int]): The token IDs to be appended to the block.

836 - """

837 - # Ensure this is mutable block (not promoted)

838 - assert self.content_hash is None

839 - assert not self.computed

840 -

841 - if len(token_ids) == 0:

842 - return

843 -

844 - # Ensure there are input tokens

845 - assert token_ids, "Got token_ids = {}".format(token_ids)

846 -

847 - # Naive block handles CoW.

848 - self._block.append_token_ids(token_ids)

849 - self._update_num_tokens_total()

850 -

851 - # If the content hash is present, then the block can be made immutable.

852 - # Register ourselves with the allocator, potentially replacing the

853 - # physical block index.

854 - if self.content_hash is not None:

855 - self.block_id = self._allocator.promote_to_immutable_block(self)

856 -

857 - @property

858 - def block_id(self) -> Optional[int]:

859 - return self._block.block_id

860 -

861 - @block_id.setter

862 - def block_id(self, value) -> None:

863 - self._block.block_id = value

864 -

865 - @property

866 - def is_full(self) -> bool:

867 - return self._block.is_full

868 -

869 - @property

870 - def num_empty_slots(self) -> int:

871 - return self._block.num_empty_slots

872 -

873 - @property

874 - def num_tokens_total(self) -> int:

875 - return self._cached_num_tokens_total

876 -

877 - @property

878 - def block_size(self) -> int:

879 - return self._block.block_size

880 -

881 - @property

882 - def token_ids(self) -> List[int]:

883 - return self._block.token_ids

884 -

885 - @property

886 - def prev_block(self) -> Optional[Block]:

887 - return self._prev_block

888 -

889 - @property

890 - def extra_hash(self) -> Optional[int]:

891 - return self._extra_hash

892 -

893 - @property

894 - def content_hash(self) -> Optional[int]:

895 - """Return the content-based hash of the current block, or None if it is

896 - not yet defined.

897 -

898 - For the content-based hash to be defined, the current block must be

899 - full.

900 - """

901 - # If the hash is already computed, return it.

902 - if self._cached_content_hash is not None:

903 - return self._cached_content_hash

904 -

905 - # We cannot compute a hash for the current block because it is not full.

906 - if not self.is_full:

907 - return None

908 -

909 - is_first_block = self._prev_block is None

910 - prev_block_hash = (

911 - self._none_hash if is_first_block else

912 - self._prev_block.content_hash # type: ignore

913 - )

914 -

915 - # Previous block exists but does not yet have a hash.

916 - # Return no hash in this case.

917 - if prev_block_hash == self._none_hash and not is_first_block:

918 - return None

919 -

920 - self._cached_content_hash = PrefixCachingBlock.hash_block_tokens(

921 - is_first_block,

922 - prev_block_hash,

923 - cur_block_token_ids=self.token_ids,

924 - extra_hash=self._extra_hash)

925 - return self._cached_content_hash

926 -

927 - @classmethod

928 - def hash_block_tokens(cls,

929 - is_first_block: bool,

930 - prev_block_hash: Optional[int],

931 - cur_block_token_ids: List[int],

932 - extra_hash: Optional[int] = None) -> int:

933 - """Computes a hash value corresponding to the contents of a block and

934 - the contents of the preceding block(s). The hash value is used for

935 - prefix caching.

936 -

937 - Parameters:

938 - - is_first_block (bool): A flag indicating if the block is the first in

939 - the sequence.

940 - - prev_block_hash (Optional[int]): The hash of the previous block. None

941 - if this is the first block.

942 - - cur_block_token_ids (List[int]): A list of token ids in the current

943 - block. The current block is assumed to be full.

944 - - extra_hash (Optional[int]): The hash value of additional factors

945 - such as adapters that influence the block, apart from the token_ids.

946 -

947 - Returns:

948 - - int: The computed hash value for the block.

949 - """

950 - if is_first_block and prev_block_hash is None:

951 - prev_block_hash = cls._none_hash

952 - return hash((is_first_block, prev_block_hash, *cur_block_token_ids,

953 - extra_hash))

954 -

955 -

956 -class ComputedBlocksTracker:

957 - """

958 - Tracks the computed blocks for each sequence.

959 -

960 - Internally, it maintains a map from sequence id to the list of block hashes

961 - for the sequence. We cache the hashes of the full blocks for each sequence,

962 - and make sure the hash is calculated in the same way as the allocator.

963 - When a sequence is being decoded, we also update the sequence's hash

964 - accordingly and incrementally.

965 -

966 - From the sequence hash, with prefix caching enabled, we could also calculate

967 - the number of cached tokens for the sequence by looking up the number of

968 - cached block hashes in the allocator.

969 - """

970 -

971 - # Note that we use 'None' as a string here instead of None because

972 - # as of Python 3.12, hash(None) returns a constant predictable value.

973 - # This could possibly make it easier to find and exploit hash

974 - # collisions. 'None' as a string will be hashed differently per process,

975 - # but consistently within the same process. This is the same as the

976 - # behavior of None prior to Python 3.12.

977 - _none_hash: int = hash('None')

978 -

979 - def __init__(

980 - self,

981 - allocator: DeviceAwareBlockAllocator,

982 - block_size: int,

983 - enable_caching: bool,

984 - ):

985 - self._allocator = allocator

986 - self._block_size = block_size

987 - self._enable_caching = enable_caching

988 -

989 - # A map from seq_id to the list of block hashes for the

990 - # sequence. This is so that we don't have to recompute the block hashes

991 - # for the sequence when we need to check if the sequence is cached.

992 - # Note a block that's not full will not have its hash calculated and

993 - # recorded.

994 - self._seq_id_to_blocks_hashes: Dict[int, List[int]] = {}

995 -

996 - # A map from seq_id to the number of tokens that are cached for the

997 - # sequence.

998 - # We need this so that a sequence in continuous prefill doesn't

999 - # accidentally see its cached token count change. See comments in

1000 - # `get_num_cached_tokens` for more details.

1001 - self._seq_id_to_num_tokens_computed: Dict[int, int] = {}

1002 -

1003 - def _update_seq_hashes(self, seq: Sequence) -> None:

1004 - """Incrementally update the sequence's block hashes and record them."""

1005 - assert self._enable_caching

1006 -

1007 - block_hashes_recorded = self._seq_id_to_blocks_hashes.get(

1008 - seq.seq_id, [])

1009 - cur_num_blocks_recorded = len(block_hashes_recorded)

1010 - token_ids = seq.get_token_ids()

1011 - assert len(token_ids) >= cur_num_blocks_recorded * self._block_size, (

1012 - f"The sequence has {len(token_ids)} tokens, but"

1013 - f" already recorded {cur_num_blocks_recorded} blocks. "

1014 - "This should not happen since we assume blocks are "

1015 - "only appended other than recomputation. When the sequence is "

1016 - "recomputed, we should have removed the info of the old blocks.")

1017 - # Update the computed block hashes for the sequence. Since only full

1018 - # blocks are considered as "computed", we take floor here.

1019 - num_computed_blocks = len(token_ids) // self._block_size

1020 -

1021 - # We need to know the hash of the previous block to compute the hash of

1022 - # the current block so that blocks could be uniquely identified across

1023 - # sequences of prefixes.

1024 - prev_block_hash = (self._none_hash if cur_num_blocks_recorded == 0 else

1025 - block_hashes_recorded[-1])

1026 - # Only update the computed block hashes for the new blocks

1027 - for i in range(cur_num_blocks_recorded, num_computed_blocks):

1028 - assert len(token_ids) >= (i + 1) * self._block_size

1029 - block_token_ids = token_ids[i * self._block_size:(i + 1) *

1030 - self._block_size]

1031 -

1032 - # NOTE: If there are any factors affecting the block besides

1033 - # token_ids, they should be added as input to extra_hash.

1034 - extra_hash = seq.extra_hash()

1035 -

1036 - # This has to be kept in sync with the allocator's hash

1037 - # calculation.

1038 - block_hash = PrefixCachingBlock.hash_block_tokens(

1039 - is_first_block=prev_block_hash == self._none_hash,

1040 - prev_block_hash=prev_block_hash,

1041 - cur_block_token_ids=block_token_ids,

1042 - extra_hash=extra_hash,

1043 - )

1044 - block_hashes_recorded.append(block_hash)

1045 - prev_block_hash = block_hash

1046 -

1047 - self._seq_id_to_blocks_hashes[seq.seq_id] = block_hashes_recorded

1048 -

1049 - def get_num_cached_tokens(self, seq: Sequence) -> int:

1050 - if not self._enable_caching:

1051 - return 0

1052 -

1053 - # We always try to update the sequence hashes on the fly.

1054 - # This is to ensure that we don't miss any cached tokens for the

1055 - # sequence during decode.

1056 - # This routine should only update hash for any new blocks too.

1057 - self._update_seq_hashes(seq)

1058 -

1059 - num_computed_tokens_prev = self._seq_id_to_num_tokens_computed.get(

1060 - seq.seq_id, None)

1061 -

1062 - # TODO(rickyx): This hack could be removed once we mark blocks as

1063 - # computed correctly with chunked prefills.

1064 - if num_computed_tokens_prev is not None and seq.is_prefill():

1065 - # For a sequence that is still in prefill, we don't

1066 - # recompute the number of cached tokens.

1067 - # This also handles correctly chunked prefill since currently

1068 - # we mark blocks as computed even if the sequence is still partially

1069 - # prefilled. So a continuously prefilled sequence should not

1070 - # see its cached token count change while running.

1071 - return num_computed_tokens_prev

1072 -

1073 - block_hashes = self._seq_id_to_blocks_hashes[seq.seq_id]

1074 -

1075 - # This is O(logN), where N is the number of blocks.

1076 - num_cached_blocks = len(

1077 - self._allocator.find_cached_blocks_prefix(block_hashes))

1078 - num_cached_tokens = num_cached_blocks * self._block_size

1079 - self._seq_id_to_num_tokens_computed[seq.seq_id] = num_cached_tokens

1080 - return num_cached_tokens

1081 -

1082 - def remove_seq(self, seq_id: int) -> None:

1083 - """Stop tracking the sequence."""

1084 - if not self._enable_caching:

1085 - return

1086 - assert seq_id in self._seq_id_to_blocks_hashes

1087 - del self._seq_id_to_blocks_hashes[seq_id]

1088 -

1089 - assert seq_id in self._seq_id_to_num_tokens_computed

1090 - del self._seq_id_to_num_tokens_computed[seq_id]

1091 -

1092 -

1093 -class LastAccessBlocksTracker:

1094 - """Manages the last access time of the tracked sequences, in order to allow

1095 - an efficient update of allocator's block last access times

1096 - """

1097 -

1098 - def __init__(self, allocator):

1099 - self._allocator = allocator

1100 - self._seq_last_access: Dict[int, Optional[float]] = {}

1101 -

1102 - def add_seq(self, seq_id: int) -> None:

1103 - """Start tracking seq_id

1104 - """

1105 - assert seq_id not in self._seq_last_access

1106 - self._seq_last_access[seq_id] = None

1107 -

1108 - def remove_seq(self, seq_id: int) -> None:

1109 - """Stop tracking seq_id

1110 - """

1111 - assert seq_id in self._seq_last_access

1112 - del self._seq_last_access[seq_id]

1113 -

1114 - def update_last_access(self, seq_id: int, time: float) -> None:

1115 - assert seq_id in self._seq_last_access

1116 - self._seq_last_access[seq_id] = time

1117 -

1118 - def update_seq_blocks_last_access(self, seq_id: int,

1119 - block_ids: List[int]) -> None:

1120 - assert seq_id in self._seq_last_access

1121 -

1122 - ts = self._seq_last_access[seq_id]

1123 -

1124 - if ts is None:

1125 - # No last access was recorded, no need to update.

1126 - return

1127 -

1128 - self._allocator.mark_blocks_as_accessed(block_ids, ts)

1129 -

1130 -

1131 -def assert_prefix_caching_block_or_none(block: Optional[Block]):

1132 - if block is None:

1133 - return

1134 - assert isinstance(block,

1135 - PrefixCachingBlock), "Got block = {}".format(block)

vllm/core/block/utils.py deleted

-28

@@ -1,28 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""Block manager utils."""

4 -from vllm.sequence import SequenceGroup

5 -from vllm.utils import (STR_NOT_IMPL_ENC_DEC_PREFIX_CACHE,

6 - STR_NOT_IMPL_ENC_DEC_SWA)

7 -

8 -

9 -def check_no_caching_or_swa_for_blockmgr_encdec(

10 - block_mgr, seq_group: SequenceGroup) -> None:

11 - '''

12 - Enforce that prefix caching & sliding-window attention (SWA)

13 - are currently unsupported *specifically* for encoder/decoder models.

14 -

15 - Raises NotImplementedError if unsupported scenario is detected.

16 -

17 - Arguments:

18 -

19 - * block_mgr: BlockSpaceManager instance

20 - * seq_group: SequenceGroup passed to block_mgr

21 - '''

22 -

23 - if seq_group.is_encoder_decoder():

24 - if block_mgr.max_block_sliding_window is not None:

25 - raise NotImplementedError(STR_NOT_IMPL_ENC_DEC_SWA)

26 -

27 - if block_mgr.enable_caching:

28 - raise NotImplementedError(STR_NOT_IMPL_ENC_DEC_PREFIX_CACHE)

vllm/core/block_manager.py deleted

-523

@@ -1,523 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""A block manager that manages token blocks."""

4 -from typing import Dict, List, Optional

5 -from typing import Sequence as GenericSequence

6 -from typing import Tuple

7 -

8 -from vllm.core.block.block_table import BlockTable

9 -from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator

10 -from vllm.core.block.interfaces import Block

11 -from vllm.core.block.prefix_caching_block import (ComputedBlocksTracker,

12 - LastAccessBlocksTracker)

13 -from vllm.core.block.utils import check_no_caching_or_swa_for_blockmgr_encdec

14 -from vllm.core.interfaces import AllocStatus, BlockSpaceManager

15 -from vllm.sequence import Sequence, SequenceGroup, SequenceStatus

16 -from vllm.utils import Device

17 -

18 -SeqId = int

19 -EncoderSeqId = str

20 -

21 -

22 -class SelfAttnBlockSpaceManager(BlockSpaceManager):

23 - """BlockSpaceManager which manages the allocation of KV cache.

24 -

25 - It owns responsibility for allocation, swapping, allocating memory for

26 - autoregressively-generated tokens, and other advanced features such as

27 - prefix caching, forking/copy-on-write, and sliding-window memory allocation.

28 -

29 - This class implements the design described in

30 - https://github.com/vllm-project/vllm/pull/3492.

31 -

32 - Lookahead slots

33 - The block manager has the notion of a "lookahead slot". These are slots

34 - in the KV cache that are allocated for a sequence. Unlike the other

35 - allocated slots, the content of these slots is undefined -- the worker

36 - may use the memory allocations in any way.

37 -

38 - In practice, a worker could use these lookahead slots to run multiple

39 - forward passes for a single scheduler invocation. Each successive

40 - forward pass would write KV activations to the corresponding lookahead

41 - slot. This allows low inter-token latency use-cases, where the overhead

42 - of continuous batching scheduling is amortized over >1 generated tokens.

43 -

44 - Speculative decoding uses lookahead slots to store KV activations of

45 - proposal tokens.

46 -

47 - See https://github.com/vllm-project/vllm/pull/3250 for more information

48 - on lookahead scheduling.

49 -

50 - Args:

51 - block_size (int): The size of each memory block.

52 - num_gpu_blocks (int): The number of memory blocks allocated on GPU.

53 - num_cpu_blocks (int): The number of memory blocks allocated on CPU.

54 - watermark (float, optional): The threshold used for memory swapping.

55 - Defaults to 0.01.

56 - sliding_window (Optional[int], optional): The size of the sliding

57 - window. Defaults to None.

58 - enable_caching (bool, optional): Flag indicating whether caching is

59 - enabled. Defaults to False.

60 - """

61 -

62 - def __init__(

63 - self,

64 - block_size: int,

65 - num_gpu_blocks: int,

66 - num_cpu_blocks: int,

67 - watermark: float = 0.01,

68 - sliding_window: Optional[int] = None,

69 - enable_caching: bool = False,

70 - ) -> None:

71 - self.block_size = block_size

72 - self.num_total_gpu_blocks = num_gpu_blocks

73 - self.num_total_cpu_blocks = num_cpu_blocks

74 -

75 - self.sliding_window = sliding_window

76 - # max_block_sliding_window is the max number of blocks that need to be

77 - # allocated

78 - self.max_block_sliding_window = None

79 - if sliding_window is not None:

80 - # +1 here because // rounds down

81 - num_blocks = sliding_window // block_size + 1

82 - # +1 here because the last block may not be full,

83 - # and so the sequence stretches one more block at the beginning

84 - # For example, if sliding_window is 3 and block_size is 4,

85 - # we may need 2 blocks when the second block only holds 1 token.

86 - self.max_block_sliding_window = num_blocks + 1

87 -

88 - self.watermark = watermark

89 - assert watermark >= 0.0

90 -

91 - self.enable_caching = enable_caching

92 -

93 - self.watermark_blocks = int(watermark * num_gpu_blocks)

94 -

95 - self.block_allocator = CpuGpuBlockAllocator.create(

96 - allocator_type="prefix_caching" if enable_caching else "naive",

97 - num_gpu_blocks=num_gpu_blocks,

98 - num_cpu_blocks=num_cpu_blocks,

99 - block_size=block_size,

100 - )

101 -

102 - self.block_tables: Dict[SeqId, BlockTable] = {}

103 - self.cross_block_tables: Dict[EncoderSeqId, BlockTable] = {}

104 -

105 - self._computed_blocks_tracker = ComputedBlocksTracker(

106 - self.block_allocator, self.block_size, self.enable_caching)

107 - self._last_access_blocks_tracker = LastAccessBlocksTracker(

108 - self.block_allocator)

109 -

110 - def can_allocate(self,

111 - seq_group: SequenceGroup,

112 - num_lookahead_slots: int = 0) -> AllocStatus:

113 - # FIXME(woosuk): Here we assume that all sequences in the group share

114 - # the same prompt. This may not be true for preempted sequences.

115 -

116 - check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)

117 -

118 - seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]

119 - num_required_blocks = BlockTable.get_num_required_blocks(

120 - seq.get_token_ids(),

121 - block_size=self.block_size,

122 - num_lookahead_slots=num_lookahead_slots,

123 - )

124 -

125 - if seq_group.is_encoder_decoder():

126 - encoder_seq = seq_group.get_encoder_seq()

127 - assert encoder_seq is not None

128 - num_required_blocks += BlockTable.get_num_required_blocks(

129 - encoder_seq.get_token_ids(),

130 - block_size=self.block_size,

131 - )

132 -

133 - if self.max_block_sliding_window is not None:

134 - num_required_blocks = min(num_required_blocks,

135 - self.max_block_sliding_window)

136 -

137 - num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(

138 - device=Device.GPU)

139 -

140 - # Use watermark to avoid frequent cache eviction.

141 - if (self.num_total_gpu_blocks - num_required_blocks

142 - < self.watermark_blocks):

143 - return AllocStatus.NEVER

144 - if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:

145 - return AllocStatus.OK

146 - else:

147 - return AllocStatus.LATER

148 -

149 - def _allocate_sequence(self, seq: Sequence) -> BlockTable:

150 - block_table = BlockTable(

151 - block_size=self.block_size,

152 - block_allocator=self.block_allocator,

153 - max_block_sliding_window=self.max_block_sliding_window,

154 - )

155 - if seq.get_token_ids():

156 - # NOTE: If there are any factors affecting the block besides

157 - # token_ids, they should be added as input to extra_hash.

158 - extra_hash = seq.extra_hash()

159 -

160 - # Add blocks to the block table only if the sequence is non empty.

161 - block_table.allocate(token_ids=seq.get_token_ids(),

162 - extra_hash=extra_hash)

163 -

164 - return block_table

165 -

166 - def allocate(self, seq_group: SequenceGroup) -> None:

167 -

168 - # Allocate self-attention block tables for decoder sequences

169 - waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)

170 - assert not (set(seq.seq_id for seq in waiting_seqs)

171 - & self.block_tables.keys()), "block table already exists"

172 -

173 - # NOTE: Here we assume that all sequences in the group have the same

174 - # prompt.

175 - seq = waiting_seqs[0]

176 - block_table: BlockTable = self._allocate_sequence(seq)

177 - self.block_tables[seq.seq_id] = block_table

178 -

179 - # Track seq

180 - self._last_access_blocks_tracker.add_seq(seq.seq_id)

181 -

182 - # Assign the block table for each sequence.

183 - for seq in waiting_seqs[1:]:

184 - self.block_tables[seq.seq_id] = block_table.fork()

185 -

186 - # Track seq

187 - self._last_access_blocks_tracker.add_seq(seq.seq_id)

188 -

189 - # Allocate cross-attention block table for encoder sequence

190 - #

191 - # NOTE: Here we assume that all sequences in the group have the same

192 - # encoder prompt.

193 - request_id = seq_group.request_id

194 -

195 - assert (request_id

196 - not in self.cross_block_tables), \

197 - "block table already exists"

198 -

199 - check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)

200 -

201 - if seq_group.is_encoder_decoder():

202 - encoder_seq = seq_group.get_encoder_seq()

203 - assert encoder_seq is not None

204 - block_table = self._allocate_sequence(encoder_seq)

205 - self.cross_block_tables[request_id] = block_table

206 -

207 - def can_append_slots(self, seq_group: SequenceGroup,

208 - num_lookahead_slots: int) -> bool:

209 - """Determine if there is enough space in the GPU KV cache to continue

210 - generation of the specified sequence group.

211 -

212 - We use a worst-case heuristic: assume each touched block will require a

213 - new allocation (either via CoW or new block). We can append slots if the

214 - number of touched blocks is less than the number of free blocks.

215 -

216 - "Lookahead slots" are slots that are allocated in addition to the slots

217 - for known tokens. The contents of the lookahead slots are not defined.

218 - This is used by speculative decoding when speculating future tokens.

219 - """

220 -

221 - num_touched_blocks = 0

222 - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):

223 - block_table = self.block_tables[seq.seq_id]

224 -

225 - num_touched_blocks += (

226 - block_table.get_num_blocks_touched_by_append_slots(

227 - token_ids=block_table.get_unseen_token_ids(

228 - seq.get_token_ids()),

229 - num_lookahead_slots=num_lookahead_slots,

230 - ))

231 -

232 - num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(

233 - Device.GPU)

234 - return num_touched_blocks <= num_free_gpu_blocks

235 -

236 - def append_slots(

237 - self,

238 - seq: Sequence,

239 - num_lookahead_slots: int,

240 - ) -> List[Tuple[int, int]]:

241 -

242 - block_table = self.block_tables[seq.seq_id]

243 -

244 - block_table.append_token_ids(

245 - token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()),

246 - num_lookahead_slots=num_lookahead_slots,

247 - num_computed_slots=seq.data.get_num_computed_tokens(),

248 - extra_hash=seq.extra_hash(),

249 - )

250 - # Return any new copy-on-writes.

251 - new_cows = self.block_allocator.clear_copy_on_writes()

252 - return new_cows

253 -

254 - def free(self, seq: Sequence) -> None:

255 - seq_id = seq.seq_id

256 -

257 - if seq_id not in self.block_tables:

258 - # Already freed or haven't been scheduled yet.

259 - return

260 -

261 - # Update seq block ids with the latest access time

262 - self._last_access_blocks_tracker.update_seq_blocks_last_access(

263 - seq_id, self.block_tables[seq.seq_id].physical_block_ids)

264 -

265 - # Untrack seq

266 - self._last_access_blocks_tracker.remove_seq(seq_id)

267 - self._computed_blocks_tracker.remove_seq(seq_id)

268 -

269 - # Free table/blocks

270 - self.block_tables[seq_id].free()

271 - del self.block_tables[seq_id]

272 -

273 - def remove_seq_from_computed_blocks_tracker(self, seq: Sequence) -> None:

274 - seq_id = seq.seq_id

275 - self._computed_blocks_tracker.remove_seq(seq_id)

276 -

277 - def free_cross(self, seq_group: SequenceGroup) -> None:

278 - request_id = seq_group.request_id

279 - if request_id not in self.cross_block_tables:

280 - # Already freed or hasn't been scheduled yet.

281 - return

282 - self.cross_block_tables[request_id].free()

283 - del self.cross_block_tables[request_id]

284 -

285 - def get_block_table(self, seq: Sequence) -> List[int]:

286 - block_ids = self.block_tables[seq.seq_id].physical_block_ids

287 - return block_ids # type: ignore

288 -

289 - def get_cross_block_table(self, seq_group: SequenceGroup) -> List[int]:

290 - request_id = seq_group.request_id

291 - assert request_id in self.cross_block_tables

292 - block_ids = self.cross_block_tables[request_id].physical_block_ids

293 - assert all(b is not None for b in block_ids)

294 - return block_ids # type: ignore

295 -

296 - def access_all_blocks_in_seq(self, seq: Sequence, now: float):

297 - if self.enable_caching:

298 - # Record the latest access time for the sequence. The actual update

299 - # of the block ids is deferred to the sequence free(..) call, since

300 - # only during freeing of block ids, the blocks are actually added to

301 - # the evictor (which is when the most updated time is required)

302 - # (This avoids expensive calls to mark_blocks_as_accessed(..))

303 - self._last_access_blocks_tracker.update_last_access(

304 - seq.seq_id, now)

305 -

306 - def mark_blocks_as_computed(self, seq_group: SequenceGroup,

307 - token_chunk_size: int):

308 - # If prefix caching is enabled, mark immutable blocks as computed

309 - # right after they have been scheduled (for prefill). This assumes

310 - # the scheduler is synchronous so blocks are actually computed when

311 - # scheduling the next batch.

312 - self.block_allocator.mark_blocks_as_computed([])

313 -

314 - def get_common_computed_block_ids(

315 - self, seqs: List[Sequence]) -> GenericSequence[int]:

316 - """Determine which blocks for which we skip prefill.

317 -

318 - With prefix caching we can skip prefill for previously-generated blocks.

319 - Currently, the attention implementation only supports skipping cached

320 - blocks if they are a contiguous prefix of cached blocks.

321 -

322 - This method determines which blocks can be safely skipped for all

323 - sequences in the sequence group.

324 - """

325 - computed_seq_block_ids = []

326 - for seq in seqs:

327 - all_blocks = self.block_tables[seq.seq_id].physical_block_ids

328 - num_cached_tokens = (

329 - self._computed_blocks_tracker.get_num_cached_tokens(seq))

330 - assert num_cached_tokens % self.block_size == 0

331 - num_cached_blocks = num_cached_tokens // self.block_size

332 - computed_block_ids = all_blocks[:num_cached_blocks]

333 - computed_seq_block_ids.append(computed_block_ids)

334 -

335 - # NOTE(sang): This assumes seq_block_ids doesn't contain any None.

336 - return self.block_allocator.get_common_computed_block_ids(

337 - computed_seq_block_ids) # type: ignore

338 -

339 - def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:

340 - if parent_seq.seq_id not in self.block_tables:

341 - # Parent sequence has either been freed or never existed.

342 - return

343 - src_block_table = self.block_tables[parent_seq.seq_id]

344 - self.block_tables[child_seq.seq_id] = src_block_table.fork()

345 -

346 - # Track child seq

347 - self._last_access_blocks_tracker.add_seq(child_seq.seq_id)

348 -

349 - def can_swap_in(self, seq_group: SequenceGroup,

350 - num_lookahead_slots: int) -> AllocStatus:

351 - """Returns the AllocStatus for the given sequence_group

352 - with num_lookahead_slots.

353 -

354 - Args:

355 - seq_group (SequenceGroup): The sequence group to swap in.

356 - num_lookahead_slots (int): Number of lookahead slots used in

357 - speculative decoding, default to 0.

358 -

359 - Returns:

360 - AllocStatus: The AllocStatus for the given sequence group.

361 - """

362 - return self._can_swap(seq_group, Device.GPU, SequenceStatus.SWAPPED,

363 - num_lookahead_slots)

364 -

365 - def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:

366 - """Returns the block id mapping (from CPU to GPU) generated by

367 - swapping in the given seq_group with num_lookahead_slots.

368 -

369 - Args:

370 - seq_group (SequenceGroup): The sequence group to swap in.

371 -

372 - Returns:

373 - List[Tuple[int, int]]: The mapping of swapping block from CPU

374 - to GPU.

375 - """

376 - physical_block_id_mapping = []

377 - for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):

378 - blocks = self.block_tables[seq.seq_id].blocks

379 - if len(blocks) == 0:

380 - continue

381 -

382 - seq_swap_mapping = self.block_allocator.swap(blocks=blocks,

383 - src_device=Device.CPU,

384 - dst_device=Device.GPU)

385 -

386 - # Refresh the block ids of the table (post-swap)

387 - self.block_tables[seq.seq_id].update(blocks)

388 -

389 - seq_physical_block_id_mapping = {

390 - self.block_allocator.get_physical_block_id(

391 - Device.CPU, cpu_block_id):

392 - self.block_allocator.get_physical_block_id(

393 - Device.GPU, gpu_block_id)

394 - for cpu_block_id, gpu_block_id in seq_swap_mapping.items()

395 - }

396 -

397 - physical_block_id_mapping.extend(

398 - list(seq_physical_block_id_mapping.items()))

399 -

400 - return physical_block_id_mapping

401 -

402 - def can_swap_out(self, seq_group: SequenceGroup) -> bool:

403 - """Returns whether we can swap out the given sequence_group

404 - with num_lookahead_slots.

405 -

406 - Args:

407 - seq_group (SequenceGroup): The sequence group to swap out.

408 -

409 - Returns:

410 - bool: Whether it's possible to swap out current sequence group.

411 - """

412 - alloc_status = self._can_swap(seq_group, Device.CPU,

413 - SequenceStatus.RUNNING)

414 - return alloc_status == AllocStatus.OK

415 -

416 - def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:

417 - """Returns the block id mapping (from GPU to CPU) generated by

418 - swapping out the given sequence_group with num_lookahead_slots.

419 -

420 - Args:

421 - seq_group (SequenceGroup): The sequence group to swap out.

422 -

423 - Returns:

424 - List[Tuple[int, int]]: The mapping of swapping block from

425 - GPU to CPU.

426 - """

427 - physical_block_id_mapping = []

428 - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):

429 - blocks = self.block_tables[seq.seq_id].blocks

430 - if len(blocks) == 0:

431 - continue

432 -

433 - seq_swap_mapping = self.block_allocator.swap(blocks=blocks,

434 - src_device=Device.GPU,

435 - dst_device=Device.CPU)

436 -

437 - # Refresh the block ids of the table (post-swap)

438 - self.block_tables[seq.seq_id].update(blocks)

439 -

440 - seq_physical_block_id_mapping = {

441 - self.block_allocator.get_physical_block_id(

442 - Device.GPU, gpu_block_id):

443 - self.block_allocator.get_physical_block_id(

444 - Device.CPU, cpu_block_id)

445 - for gpu_block_id, cpu_block_id in seq_swap_mapping.items()

446 - }

447 -

448 - physical_block_id_mapping.extend(

449 - list(seq_physical_block_id_mapping.items()))

450 -

451 - return physical_block_id_mapping

452 -

453 - def get_num_free_gpu_blocks(self) -> int:

454 - return self.block_allocator.get_num_free_blocks(Device.GPU)

455 -

456 - def get_num_free_cpu_blocks(self) -> int:

457 - return self.block_allocator.get_num_free_blocks(Device.CPU)

458 -

459 - def get_prefix_cache_hit_rate(self, device: Device) -> float:

460 - return self.block_allocator.get_prefix_cache_hit_rate(device)

461 -

462 - def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

463 - return self.block_allocator.reset_prefix_cache(device)

464 -

465 - def _can_swap(self,

466 - seq_group: SequenceGroup,

467 - device: Device,

468 - status: SequenceStatus,

469 - num_lookahead_slots: int = 0) -> AllocStatus:

470 - """Returns the AllocStatus for swapping in/out the given sequence_group

471 - on to the 'device'.

472 -

473 - Args:

474 - seq_group (SequenceGroup): The sequence group to swap in/out.

475 - device (Device): device to swap the 'seq_group' on.

476 - status (SequenceStatus): The status of sequence which is needed

477 - for action. RUNNING for swap out and SWAPPED for swap in

478 - num_lookahead_slots (int): Number of lookahead slots used in

479 - speculative decoding, default to 0.

480 -

481 - Returns:

482 - AllocStatus: The AllocStatus for swapping in/out the given

483 - sequence_group on to the 'device'.

484 - """

485 - # First determine the number of blocks that will be touched by this

486 - # swap. Then verify if there are available blocks in the device

487 - # to perform the swap.

488 - num_blocks_touched = 0

489 - blocks: List[Block] = []

490 - for seq in seq_group.get_seqs(status=status):

491 - block_table = self.block_tables[seq.seq_id]

492 - if block_table.blocks is not None:

493 - # Compute the number blocks to touch for the tokens to be

494 - # appended. This does NOT include the full blocks that need

495 - # to be touched for the swap.

496 - num_blocks_touched += \

497 - block_table.get_num_blocks_touched_by_append_slots(

498 - block_table.get_unseen_token_ids(seq.get_token_ids()),

499 - num_lookahead_slots=num_lookahead_slots)

500 - blocks.extend(block_table.blocks)

501 - # Compute the number of full blocks to touch and add it to the

502 - # existing count of blocks to touch.

503 - num_blocks_touched += self.block_allocator.get_num_full_blocks_touched(

504 - blocks, device=device)

505 -

506 - watermark_blocks = 0

507 - if device == Device.GPU:

508 - watermark_blocks = self.watermark_blocks

509 -

510 - if self.block_allocator.get_num_total_blocks(

511 - device) < num_blocks_touched:

512 - return AllocStatus.NEVER

513 - elif self.block_allocator.get_num_free_blocks(

514 - device) - num_blocks_touched >= watermark_blocks:

515 - return AllocStatus.OK

516 - else:

517 - return AllocStatus.LATER

518 -

519 - def get_num_cached_tokens(self, seq: Sequence) -> int:

520 - """Get the number of tokens in blocks that are already computed and

521 - cached in the block manager for the sequence.

522 - """

523 - return self._computed_blocks_tracker.get_num_cached_tokens(seq)

vllm/core/evictor.py deleted

-157

@@ -1,157 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import enum

5 -import heapq

6 -from abc import ABC, abstractmethod

7 -from typing import Dict, List, Tuple

8 -

9 -

10 -class EvictionPolicy(enum.Enum):

11 - """Enum for eviction policy used by make_evictor to instantiate the correct

12 - Evictor subclass.

13 - """

14 - LRU = enum.auto()

15 -

16 -

17 -class Evictor(ABC):

18 - """The Evictor subclasses should be used by the BlockAllocator class to

19 - handle eviction of freed Blocks.

20 - """

21 -

22 - @abstractmethod

23 - def __init__(self):

24 - pass

25 -

26 - @abstractmethod

27 - def __contains__(self, block_id: int) -> bool:

28 - pass

29 -

30 - @abstractmethod

31 - def evict(self) -> Tuple[int, int]:

32 - """Runs the eviction algorithm and returns the evicted block's

33 - content hash along with physical block id along with physical block id

34 - """

35 - pass

36 -

37 - @abstractmethod

38 - def add(self, block_id: int, content_hash: int, num_hashed_tokens: int,

39 - last_accessed: float):

40 - """Adds block to the evictor, making it a candidate for eviction"""

41 - pass

42 -

43 - @abstractmethod

44 - def update(self, block_id: int, last_accessed: float):

45 - """Update corresponding block's access time in metadata"""

46 - pass

47 -

48 - @abstractmethod

49 - def remove(self, block_id: int):

50 - """Remove a given block id from the cache."""

51 - pass

52 -

53 - @property

54 - @abstractmethod

55 - def num_blocks(self) -> int:

56 - pass

57 -

58 -

59 -class BlockMetaData:

60 - """Data structure for storing key data describe cached block, so that

61 - evictor could use to make its decision which one to choose for eviction

62 -

63 - Here we use physical block id as the dict key, as there maybe several

64 - blocks with the same content hash, but their physical id is unique.

65 - """

66 -

67 - def __init__(self, content_hash: int, num_hashed_tokens: int,

68 - last_accessed: float):

69 - self.content_hash = content_hash

70 - self.num_hashed_tokens = num_hashed_tokens

71 - self.last_accessed = last_accessed

72 -

73 -

74 -class LRUEvictor(Evictor):

75 - """Evicts in a least-recently-used order using the last_accessed timestamp

76 - that's recorded in the Block. If there are multiple blocks with

77 - the same last_accessed time, then the one with the largest num_hashed_tokens

78 - will be evicted. If two blocks each have the lowest last_accessed time and

79 - highest num_hashed_tokens value, then one will be chosen arbitrarily

80 - """

81 -

82 - # CLEANUP_THRESHOLD determines the maximum allowable size of the priority

83 - # queue relative to the free table size. When this threshold is exceeded,

84 - # a cleanup operation is triggered to reduce memory usage.

85 - CLEANUP_THRESHOLD = 50

86 -

87 - def __init__(self):

88 - self.free_table: Dict[int, BlockMetaData] = {}

89 - self.priority_queue = []

90 -

91 - def __contains__(self, block_id: int) -> bool:

92 - return block_id in self.free_table

93 -

94 - def evict(self) -> Tuple[int, int]:

95 - if len(self.free_table) == 0:

96 - raise ValueError("No usable cache memory left")

97 -

98 - while self.priority_queue:

99 - # We do not remove outdated entries from the priority queue at the

100 - # time of updating the last_accessed timestamp. Instead, outdated

101 - # entries are filtered out here during eviction. Outdated entries

102 - # would either not in the free table, or have older last accessed

103 - # time.

104 - last_accessed, _, block_id, content_hash = heapq.heappop(

105 - self.priority_queue)

106 - if (block_id in self.free_table and

107 - self.free_table[block_id].last_accessed == last_accessed):

108 - self.free_table.pop(block_id)

109 - return block_id, content_hash

110 -

111 - raise ValueError("No usable cache memory left")

112 -

113 - def add(self, block_id: int, content_hash: int, num_hashed_tokens: int,

114 - last_accessed: float):

115 - self.free_table[block_id] = BlockMetaData(content_hash,

116 - num_hashed_tokens,

117 - last_accessed)

118 - heapq.heappush(

119 - self.priority_queue,

120 - (last_accessed, -num_hashed_tokens, block_id, content_hash))

121 - self._cleanup_if_necessary()

122 -

123 - def update(self, block_id: int, last_accessed: float):

124 - self.free_table[block_id].last_accessed = last_accessed

125 -

126 - def _cleanup_if_necessary(self):

127 - if len(self.priority_queue) > LRUEvictor.CLEANUP_THRESHOLD * len(

128 - self.free_table):

129 - self._cleanup()

130 -

131 - def _cleanup(self):

132 - new_priority_queue: List[Tuple[float, int, int, int]] = []

133 -

134 - for block_id, block in self.free_table.items():

135 - new_priority_queue.append(

136 - (block.last_accessed, -block.num_hashed_tokens, block_id,

137 - block.content_hash))

138 - heapq.heapify(new_priority_queue)

139 -

140 - self.priority_queue = new_priority_queue

141 -

142 - def remove(self, block_id: int):

143 - if block_id not in self.free_table:

144 - raise ValueError(

145 - "Attempting to remove block that's not in the evictor")

146 - self.free_table.pop(block_id)

147 -

148 - @property

149 - def num_blocks(self) -> int:

150 - return len(self.free_table)

151 -

152 -

153 -def make_evictor(eviction_policy: EvictionPolicy) -> Evictor:

154 - if eviction_policy == EvictionPolicy.LRU:

155 - return LRUEvictor()

156 - else:

157 - raise ValueError(f"Unknown cache eviction policy: {eviction_policy}")

vllm/core/interfaces.py deleted

-139

@@ -1,139 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import enum

5 -from abc import ABC, abstractmethod

6 -from typing import List, Optional

7 -from typing import Sequence as GenericSequence

8 -from typing import Tuple

9 -

10 -from vllm.sequence import Sequence, SequenceGroup

11 -from vllm.utils import Device

12 -

13 -

14 -class AllocStatus(enum.Enum):

15 - """Result for BlockSpaceManager.can_allocate

16 -

17 - 1. Ok: seq_group can be allocated now.

18 - 2. Later: seq_group cannot be allocated.

19 - The capacity of allocator is larger than seq_group required.

20 - 3. Never: seq_group can never be allocated.

21 - The seq_group is too large to allocated in GPU.

22 - """

23 - OK = enum.auto()

24 - LATER = enum.auto()

25 - NEVER = enum.auto()

26 -

27 -

28 -class BlockSpaceManager(ABC):

29 -

30 - @staticmethod

31 - def get_block_space_manager_class(version: str):

32 - version = version.lower()

33 -

34 - if version == "selfattn":

35 - from vllm.core.block_manager import SelfAttnBlockSpaceManager

36 - return SelfAttnBlockSpaceManager

37 -

38 - if version == "placeholder":

39 - from vllm.core.placeholder_block_space_manager import (

40 - PlaceholderBlockSpaceManager)

41 - return PlaceholderBlockSpaceManager

42 -

43 - raise ValueError(f"Unknown version {version=}")

44 -

45 - @abstractmethod

46 - def can_allocate(self,

47 - seq_group: SequenceGroup,

48 - num_lookahead_slots: int = 0) -> AllocStatus:

49 - pass

50 -

51 - @abstractmethod

52 - def allocate(self, seq_group: SequenceGroup) -> None:

53 - pass

54 -

55 - @abstractmethod

56 - def can_append_slots(self, seq_group: SequenceGroup,

57 - num_lookahead_slots: int) -> bool:

58 - pass

59 -

60 - @abstractmethod

61 - def append_slots(

62 - self,

63 - seq: Sequence,

64 - num_lookahead_slots: int,

65 - ) -> List[Tuple[int, int]]:

66 - pass

67 -

68 - @abstractmethod

69 - def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:

70 - pass

71 -

72 - @abstractmethod

73 - def can_swap_in(self, seq_group: SequenceGroup,

74 - num_lookahead_slots: int) -> AllocStatus:

75 - pass

76 -

77 - @abstractmethod

78 - def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:

79 - pass

80 -

81 - @abstractmethod

82 - def can_swap_out(self, seq_group: SequenceGroup) -> bool:

83 - pass

84 -

85 - @abstractmethod

86 - def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:

87 - pass

88 -

89 - @abstractmethod

90 - def free(self, seq: Sequence) -> None:

91 - pass

92 -

93 - @abstractmethod

94 - def get_block_table(self, seq: Sequence) -> List[int]:

95 - pass

96 -

97 - @abstractmethod

98 - def get_num_free_gpu_blocks(self) -> int:

99 - pass

100 -

101 - @abstractmethod

102 - def get_num_free_cpu_blocks(self) -> int:

103 - pass

104 -

105 - @abstractmethod

106 - def access_all_blocks_in_seq(

107 - self,

108 - seq: Sequence,

109 - access_time: float,

110 - ) -> None:

111 - pass

112 -

113 - @abstractmethod

114 - def get_common_computed_block_ids(

115 - self, seqs: List[Sequence]) -> GenericSequence[int]:

116 - pass

117 -

118 - @abstractmethod

119 - def mark_blocks_as_computed(self, seq_group: SequenceGroup,

120 - token_chunk_size: int):

121 - pass

122 -

123 - @abstractmethod

124 - def get_prefix_cache_hit_rate(self, device: Device) -> float:

125 - """Prefix cache hit rate. -1 means not supported or disabled."""

126 - pass

127 -

128 - @abstractmethod

129 - def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

130 - """Reset prefix cache for specified or all devices."""

131 - pass

132 -

133 - @abstractmethod

134 - def get_num_cached_tokens(self, seq: Sequence) -> int:

135 - pass

136 -

137 - @abstractmethod

138 - def remove_seq_from_computed_blocks_tracker(self, seq: Sequence) -> None:

139 - pass

\ No newline at end of file

vllm/core/placeholder_block_space_manager.py deleted

-103

@@ -1,103 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -from typing import List, Optional, Tuple

5 -

6 -from vllm.core.interfaces import AllocStatus, BlockSpaceManager

7 -from vllm.sequence import Sequence, SequenceGroup

8 -from vllm.utils import Device

9 -

10 -

11 -class PlaceholderBlockSpaceManager(BlockSpaceManager):

12 - """A version of BlockSpaceManager for use in environments

13 - where block management is not required.

14 - For example: pooling models or attention-free models like Mamba.

15 -

16 - This class provides the same interface as BlockSpaceManager, but its

17 - methods perform no actions or return simple values like True in specific

18 - actions. It's designed to be used in scenarios where the overhead of

19 - block management is unnecessary, such as in an embedding environment.

20 - """

21 -

22 - def __init__(

23 - self,

24 - **kwargs,

25 - ) -> None:

26 - pass

27 -

28 - def can_allocate(self,

29 - seq_group: SequenceGroup,

30 - num_lookahead_slots: int = 0) -> AllocStatus:

31 - # Always return OK for dummy purposes

32 - return AllocStatus.OK

33 -

34 - def allocate(self, seq_group: SequenceGroup) -> None:

35 - # No actual allocation logic needed

36 - pass

37 -

38 - def can_append_slots(self, seq_group: SequenceGroup,

39 - num_lookahead_slots: int) -> bool:

40 - return True

41 -

42 - def append_slots(

43 - self,

44 - seq: Sequence,

45 - num_lookahead_slots: int,

46 - ) -> List[Tuple[int, int]]:

47 - return []

48 -

49 - def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:

50 - pass

51 -

52 - def can_swap_in(self, seq_group: SequenceGroup,

53 - num_lookahead_slots: int) -> AllocStatus:

54 - return AllocStatus.OK

55 -

56 - def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:

57 - return None # type: ignore

58 -

59 - def can_swap_out(self, seq_group: SequenceGroup) -> bool:

60 - return True

61 -

62 - def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:

63 - return None # type: ignore

64 -

65 - def free(self, seq: Sequence) -> None:

66 - # No operation on free

67 - return

68 -

69 - def get_block_table(self, seq: Sequence) -> List[int]:

70 - return None # type: ignore

71 -

72 - def get_num_free_gpu_blocks(self) -> int:

73 - return 1

74 -

75 - def get_num_free_cpu_blocks(self) -> int:

76 - return 1

77 -

78 - def access_all_blocks_in_seq(

79 - self,

80 - seq: Sequence,

81 - access_time: float,

82 - ) -> None:

83 - pass

84 -

85 - def get_common_computed_block_ids(self,

86 - seq_group: List[Sequence]) -> List[int]:

87 - return []

88 -

89 - def mark_blocks_as_computed(self, seq_group: SequenceGroup,

90 - token_chunk_size: int):

91 - pass

92 -

93 - def get_prefix_cache_hit_rate(self, device: Device) -> float:

94 - return -1

95 -

96 - def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

97 - return True

98 -

99 - def get_num_cached_tokens(self, seq: Sequence) -> int:

100 - return 0

101 -

102 - def remove_seq_from_computed_blocks_tracker(self, seq: Sequence) -> None:

103 - return

vllm/core/scheduler.py deleted

-2028

@@ -1,2028 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import enum

5 -import os

6 -import random

7 -import time

8 -from collections import deque

9 -from dataclasses import dataclass, field

10 -from typing import Callable, Deque, Dict, Iterable, List, Optional

11 -from typing import Sequence as GenericSequence

12 -from typing import Set, Tuple, Union

13 -

14 -from vllm.config import CacheConfig, SchedulerConfig

15 -from vllm.config.lora import LoRAConfig

16 -from vllm.core.interfaces import AllocStatus, BlockSpaceManager

17 -from vllm.logger import init_logger

18 -from vllm.lora.request import LoRARequest

19 -from vllm.sequence import (Sequence, SequenceData, SequenceGroup,

20 - SequenceGroupBase, SequenceGroupMetadata,

21 - SequenceGroupMetadataDelta, SequenceStage,

22 - SequenceStatus)

23 -from vllm.utils import Device, PyObjectCache

24 -

25 -logger = init_logger(__name__)

26 -

27 -# Test-only. If configured, decode is preempted with

28 -# ARTIFICIAL_PREEMPTION_PROB% probability.

29 -ENABLE_ARTIFICIAL_PREEMPT = bool(

30 - os.getenv("VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT", False)) # noqa

31 -ARTIFICIAL_PREEMPTION_PROB = 0.5

32 -ARTIFICIAL_PREEMPTION_MAX_CNT = 500

33 -

34 -

35 -class PreemptionMode(enum.Enum):

36 - """Preemption modes.

37 -

38 - 1. Swapping: Swap out the blocks of the preempted sequences to CPU memory

39 - and swap them back in when the sequences are resumed.

40 - 2. Recomputation: Discard the blocks of the preempted sequences and

41 - recompute them when the sequences are resumed, treating the sequences as

42 - new prompts.

43 - """

44 -

45 - SWAP = enum.auto()

46 - RECOMPUTE = enum.auto()

47 -

48 -

49 -@dataclass

50 -class SchedulingBudget:

51 - """The available slots for scheduling.

52 -

53 - TODO(sang): Right now, the budget is request_id-aware meaning it can ignore

54 - budget update from the same request_id. It is because in normal scheduling

55 - path, we update RUNNING num_seqs ahead of time, meaning it could be

56 - updated more than once when scheduling RUNNING requests. Since this won't

57 - happen if we only have chunked prefill scheduling, we can remove this

58 - feature from the API when chunked prefill is enabled by default.

59 - """

60 -

61 - token_budget: int

62 - max_num_seqs: int

63 - _request_ids_num_batched_tokens: Set[str] = field(default_factory=set)

64 - _request_ids_num_curr_seqs: Set[str] = field(default_factory=set)

65 - # Number of cached tokens in the batch.

66 - _num_cached_tokens: int = 0

67 - # Number of actual non-cached tokens in the batch.

68 - _num_batched_tokens: int = 0

69 - _num_curr_seqs: int = 0

70 -

71 - def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int):

72 - # We allow num_new_tokens to be 0 when the entire sequence has

73 - # been cached.

74 - assert num_new_tokens >= 0

75 - assert num_new_seqs != 0

76 - return (self.num_batched_tokens + num_new_tokens <= self.token_budget

77 - and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs)

78 -

79 - def remaining_token_budget(self):

80 - return self.token_budget - self.num_batched_tokens

81 -

82 - def add_num_batched_tokens(self,

83 - req_id: str,

84 - num_batched_tokens: int,

85 - num_cached_tokens: int = 0):

86 - if req_id in self._request_ids_num_batched_tokens:

87 - return

88 - assert num_cached_tokens >= 0

89 - assert num_batched_tokens >= 0

90 -

91 - self._request_ids_num_batched_tokens.add(req_id)

92 - self._num_batched_tokens += num_batched_tokens

93 - self._num_cached_tokens += num_cached_tokens

94 -

95 - def subtract_num_batched_tokens(self, req_id: str,

96 - num_batched_tokens: int):

97 - if req_id in self._request_ids_num_batched_tokens:

98 - self._request_ids_num_batched_tokens.remove(req_id)

99 - self._num_batched_tokens -= num_batched_tokens

100 -

101 - def add_num_seqs(self, req_id: str, num_curr_seqs: int):

102 - if req_id in self._request_ids_num_curr_seqs:

103 - return

104 -

105 - self._request_ids_num_curr_seqs.add(req_id)

106 - self._num_curr_seqs += num_curr_seqs

107 -

108 - def subtract_num_seqs(self, req_id: str, num_curr_seqs: int):

109 - if req_id in self._request_ids_num_curr_seqs:

110 - self._request_ids_num_curr_seqs.remove(req_id)

111 - self._num_curr_seqs -= num_curr_seqs

112 -

113 - @property

114 - def num_batched_tokens(self):

115 - return self._num_batched_tokens

116 -

117 - @property

118 - def num_curr_seqs(self):

119 - return self._num_curr_seqs

120 -

121 - @property

122 - def num_cached_tokens(self):

123 - return self._num_cached_tokens

124 -

125 -

126 -@dataclass

127 -class ScheduledSequenceGroup:

128 - # A sequence group that's scheduled.

129 - seq_group: SequenceGroup

130 - # The total chunk size (number of tokens) to process for next iteration.

131 - # 1 for decoding. Same as prompt tokens for prefill, but if prefill is

132 - # chunked, it can be smaller than that.

133 - token_chunk_size: int

134 -

135 -

136 -@dataclass

137 -class SchedulerOutputs:

138 - """The scheduling decision made from a scheduler."""

139 -

140 - # Scheduled sequence groups.

141 - scheduled_seq_groups: GenericSequence[ScheduledSequenceGroup]

142 - # Number of prefill groups scheduled.

143 - num_prefill_groups: int

144 - # Total number of batched tokens.

145 - num_batched_tokens: int

146 - # Blocks to swap in. List of CPU -> GPU block number.

147 - blocks_to_swap_in: List[Tuple[int, int]]

148 - # Blocks to swap out. List of GPU -> CPU block number.

149 - blocks_to_swap_out: List[Tuple[int, int]]

150 - # Blocks to copy. Source to dest block.

151 - blocks_to_copy: List[Tuple[int, int]]

152 - # Sequence groups that are going to be ignored.

153 - ignored_seq_groups: List[SequenceGroup]

154 - # The number of slots for lookahead decoding.

155 - num_lookahead_slots: int

156 - # The number of requests in the running queue

157 - running_queue_size: int

158 - preempted: int

159 -

160 - def __post_init__(self):

161 - # Swap in and swap out should never happen at the same time.

162 - assert not (self.blocks_to_swap_in and self.blocks_to_swap_out)

163 -

164 - self.num_loras: int = len(self.lora_requests)

165 - if self.num_loras > 0:

166 - self._sort_by_lora_ids()

167 -

168 - def is_empty(self) -> bool:

169 - # NOTE: We do not consider the ignored sequence groups.

170 - return (not self.scheduled_seq_groups and not self.blocks_to_swap_in

171 - and not self.blocks_to_swap_out and not self.blocks_to_copy)

172 -

173 - def _sort_by_lora_ids(self):

174 - assert 0 <= self.num_prefill_groups <= len(self.scheduled_seq_groups)

175 -

176 - def key_fn(group: ScheduledSequenceGroup):

177 - key = (group.seq_group.lora_int_id, group.seq_group.request_id)

178 - if 0 < self.num_prefill_groups < len(self.scheduled_seq_groups):

179 - # Sort sequence groups so that all prefills come before all

180 - # decodes as required by chunked prefill.

181 - return (not group.seq_group.is_prefill(), *key)

182 - return key

183 -

184 - self.scheduled_seq_groups = sorted(self.scheduled_seq_groups,

185 - key=key_fn)

186 -

187 - @property

188 - def lora_requests(self) -> Set[LoRARequest]:

189 - return {

190 - g.seq_group.lora_request

191 - for g in self.scheduled_seq_groups

192 - if g.seq_group.lora_request is not None

193 - }

194 -

195 -

196 -@dataclass

197 -class SchedulerRunningOutputs:

198 - """The requests that are scheduled from a running queue.

199 -

200 - Could contain prefill (prefill that's chunked) or decodes. If there's not

201 - enough memory, it can be preempted (for recompute) or swapped out.

202 - """

203 -

204 - # Selected sequences that are running and in a decoding phase.

205 - decode_seq_groups: List[ScheduledSequenceGroup]

206 - # Selected sequences that are running and in a prefill phase.

207 - # I.e., it means the prefill has been chunked.

208 - prefill_seq_groups: List[ScheduledSequenceGroup]

209 - # The preempted sequences.

210 - preempted: List[SequenceGroup]

211 - # Sequences that are swapped out.

212 - swapped_out: List[SequenceGroup]

213 - # The blocks to swap out.

214 - blocks_to_swap_out: List[Tuple[int, int]]

215 - # The blocks to copy.

216 - blocks_to_copy: List[Tuple[int, int]]

217 - # The number of slots for lookahead decoding.

218 - num_lookahead_slots: int

219 -

220 - # Optimization for fast-access to seq_group lists

221 - decode_seq_groups_list: List[SequenceGroup]

222 - prefill_seq_groups_list: List[SequenceGroup]

223 -

224 - @classmethod

225 - def create_empty(cls) -> "SchedulerRunningOutputs":

226 - return SchedulerRunningOutputs(

227 - decode_seq_groups=[],

228 - prefill_seq_groups=[],

229 - preempted=[],

230 - swapped_out=[],

231 - blocks_to_swap_out=[],

232 - blocks_to_copy=[],

233 - num_lookahead_slots=0,

234 - decode_seq_groups_list=[],

235 - prefill_seq_groups_list=[],

236 - )

237 -

238 -

239 -@dataclass

240 -class SchedulerSwappedInOutputs:

241 - """The requests that are scheduled from a swap queue.

242 -

243 - Could contain prefill (prefill that's chunked) or decodes.

244 - """

245 -

246 - # Selected sequences that are going to be swapped in and is in a

247 - # decoding phase.

248 - decode_seq_groups: List[ScheduledSequenceGroup]

249 - # Selected sequences that are going to be swapped in and in a prefill

250 - # phase. I.e., it means the prefill has been chunked.

251 - prefill_seq_groups: List[ScheduledSequenceGroup]

252 - # The blocks to swap in.

253 - blocks_to_swap_in: List[Tuple[int, int]]

254 - # The blocks to copy.

255 - blocks_to_copy: List[Tuple[int, int]]

256 - # The number of slots for lookahead decoding.

257 - num_lookahead_slots: int

258 - # Infeasible sequence groups.

259 - infeasible_seq_groups: List[SequenceGroup]

260 -

261 - @classmethod

262 - def create_empty(cls) -> "SchedulerSwappedInOutputs":

263 - return SchedulerSwappedInOutputs(

264 - decode_seq_groups=[],

265 - prefill_seq_groups=[],

266 - blocks_to_swap_in=[],

267 - blocks_to_copy=[],

268 - num_lookahead_slots=0,

269 - infeasible_seq_groups=[],

270 - )

271 -

272 -

273 -@dataclass

274 -class SchedulerPrefillOutputs:

275 - """The requests that are scheduled from a waiting queue.

276 -

277 - Could contain a fresh prefill requests or preempted requests that need

278 - to be recomputed from scratch.

279 - """

280 -

281 - # Selected sequences for prefill.

282 - seq_groups: List[ScheduledSequenceGroup]

283 - # Ignored sequence groups.

284 - ignored_seq_groups: List[SequenceGroup]

285 - num_lookahead_slots: int

286 -

287 - @classmethod

288 - def create_empty(cls) -> "SchedulerPrefillOutputs":

289 - return SchedulerPrefillOutputs(

290 - seq_groups=[],

291 - ignored_seq_groups=[],

292 - num_lookahead_slots=0,

293 - )

294 -

295 -

296 -def seq_group_metadata_builder():

297 - return SequenceGroupMetadata(request_id="",

298 - is_prompt=False,

299 - seq_data={},

300 - sampling_params=None,

301 - block_tables={})

302 -

303 -

304 -def scheduler_running_outputs_builder():

305 - return SchedulerRunningOutputs(decode_seq_groups=[],

306 - prefill_seq_groups=[],

307 - preempted=[],

308 - swapped_out=[],

309 - blocks_to_swap_out=[],

310 - blocks_to_copy=[],

311 - num_lookahead_slots=0,

312 - prefill_seq_groups_list=[],

313 - decode_seq_groups_list=[])

314 -

315 -

316 -def scheduled_seq_group_builder():

317 - return ScheduledSequenceGroup(SequenceGroup.__new__(SequenceGroup),

318 - token_chunk_size=0)

319 - # return ScheduledSequenceGroup(seq_group=None, token_chunk_size=0)

320 -

321 -

322 -@dataclass

323 -class PartialPrefillMetadata:

324 - """Holds information about the partial prefills that are currently running

325 - during a single iteration of the Scheduler.

326 - When chunked prefill is enabled, we allow a certain number of seqs to be

327 - partially prefilled during each iteration. Having multiple partial prefills

328 - in flight allows us to minimize TTFT and avoid decode starvation in cases

329 - where a single sequence group with a very large prompt blocks the queue for

330 - too many iterations.

331 - The number of long prefill requests is limited so that smaller

332 - requests may jump the queue in front of them and get to the decode

333 - phase faster.

334 - """

335 -

336 - # A minimum bound on the total number of prefills to be scheduled during

337 - # this iteration

338 - schedulable_prefills: int

339 -

340 - # The number of long prefill requests currently running

341 - long_prefills: int

342 -

343 - scheduler_config: SchedulerConfig

344 -

345 - def can_schedule(self, seq_group: SequenceGroup) -> bool:

346 - """When concurrent partial prefills are enabled,

347 - we limit the number of long requests and only accept

348 - shorter requests from the queue while running them

349 - concurrently"""

350 - return not (seq_group.first_seq.get_num_new_tokens()

351 - > self.scheduler_config.long_prefill_token_threshold

352 - and self.long_prefills

353 - >= self.scheduler_config.max_long_partial_prefills

354 - and self.scheduler_config.max_num_partial_prefills > 1)

355 -

356 - def maybe_increment_partial_prefills(self,

357 - seq_group: SequenceGroup) -> None:

358 - # When a new prefill is scheduled, we need to know if it is a

359 - # long request

360 - if (seq_group.first_seq.get_num_new_tokens()

361 - > self.scheduler_config.long_prefill_token_threshold):

362 - self.long_prefills += 1

363 -

364 - @classmethod

365 - def from_queues(

366 - cls,

367 - running: Deque[SequenceGroup],

368 - waiting: Deque[SequenceGroup],

369 - scheduler_config: SchedulerConfig,

370 - ) -> "PartialPrefillMetadata":

371 - """Create a PartialPrefillMetadata object from the current state of

372 - the scheduler's queues.

373 - This accounts for the currently running prefill requests, and peeks into

374 - the waiting queue to see if there are more prefills to potentially be

375 - scheduled during this iteration."""

376 - prefills = 0

377 - long_prefills = 0

378 -

379 - waiting_long_prefills = 0

380 -

381 - for sg in running:

382 - if sg.first_seq.data.stage == SequenceStage.PREFILL:

383 - prefills += 1

384 - if (sg.first_seq.get_num_new_tokens()

385 - > scheduler_config.long_prefill_token_threshold):

386 - long_prefills += 1

387 -

388 - for sg in waiting:

389 - # Don't bother looping through the rest of the queue if we know

390 - # there are already at

391 - # least max_partial_prefills requests to fill

392 - if prefills >= scheduler_config.max_num_partial_prefills:

393 - break

394 -

395 - # Don't count long requests from the waiting queue if we aren't

396 - # going to schedule them anyway

397 - if (sg.first_seq.get_num_new_tokens()

398 - > scheduler_config.long_prefill_token_threshold):

399 - if (long_prefills + waiting_long_prefills

400 - >= scheduler_config.max_long_partial_prefills):

401 - continue

402 - waiting_long_prefills += 1

403 - prefills += 1

404 -

405 - # NB: long_prefills and waiting_long_prefills are tracked separately.

406 - # We don't account for the waiting requests here because we need to use

407 - # this metadata to track how many have actually been scheduled.

408 - return PartialPrefillMetadata(

409 - schedulable_prefills=min(

410 - prefills, scheduler_config.max_num_partial_prefills),

411 - long_prefills=long_prefills,

412 - scheduler_config=scheduler_config,

413 - )

414 -

415 -

416 -class Scheduler:

417 -

418 - def __init__(

419 - self,

420 - scheduler_config: SchedulerConfig,

421 - cache_config: CacheConfig,

422 - lora_config: Optional[LoRAConfig],

423 - pipeline_parallel_size: int = 1,

424 - output_proc_callback: Optional[Callable] = None,

425 - ) -> None:

426 - self.scheduler_config = scheduler_config

427 - self.cache_config = cache_config

428 - # Note for LoRA scheduling: the current policy is extremely

429 - # simple and NOT fair. It can lead to starvation of some

430 - # LoRAs. This should be improved in the future.

431 - self.lora_config = lora_config

432 -

433 - version = "selfattn"

434 - if (self.scheduler_config.runner_type == "pooling"

435 - or self.cache_config.is_attention_free):

436 - version = "placeholder"

437 -

438 - BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class(

439 - version)

440 -

441 - num_gpu_blocks = cache_config.num_gpu_blocks

442 - if num_gpu_blocks:

443 - num_gpu_blocks //= pipeline_parallel_size

444 -

445 - num_cpu_blocks = cache_config.num_cpu_blocks

446 - if num_cpu_blocks:

447 - num_cpu_blocks //= pipeline_parallel_size

448 -

449 - # Create the block space manager.

450 - self.block_manager = BlockSpaceManagerImpl(

451 - block_size=self.cache_config.block_size,

452 - num_gpu_blocks=num_gpu_blocks,

453 - num_cpu_blocks=num_cpu_blocks,

454 - sliding_window=self.cache_config.sliding_window,

455 - enable_caching=self.cache_config.enable_prefix_caching,

456 - )

457 -

458 - # Sequence groups in the WAITING state.

459 - # Contain new prefill or preempted requests.

460 - self.waiting: Deque[SequenceGroup] = deque()

461 - # Sequence groups in the RUNNING state.

462 - # Contain decode requests.

463 - self.running: Deque[SequenceGroup] = deque()

464 - # Sequence groups in the SWAPPED state.

465 - # Contain decode requests that are swapped out.

466 - self.swapped: Deque[SequenceGroup] = deque()

467 - # Sequence groups finished requests ids since last step iteration.

468 - # It lets the model know that any state associated with these requests

469 - # can and must be released after the current step.

470 - # This is used to evict the finished requests from the Mamba cache.

471 - self._finished_requests_ids: List[str] = list()

472 - # Time at previous scheduling step

473 - self.prev_time = 0.0

474 - # Did we schedule a prompt at previous step?

475 - self.prev_prompt = False

476 - # Latency of the last prompt step

477 - self.last_prompt_latency = 0.0

478 - # preemption mode, RECOMPUTE or SWAP

479 - self.user_specified_preemption_mode = scheduler_config.preemption_mode

480 -

481 - # The following field is test-only. It is used to inject artificial

482 - # preemption.

483 - self.enable_artificial_preemption = ENABLE_ARTIFICIAL_PREEMPT

484 - self.artificial_preempt_cnt = (ARTIFICIAL_PREEMPTION_MAX_CNT

485 - if self.enable_artificial_preemption

486 - else 0)

487 - self.num_cumulative_preemption: int = 0

488 -

489 - # Used to cache python objects

490 - self._seq_group_metadata_cache: List[PyObjectCache] = []

491 - self._scheduler_running_outputs_cache: List[PyObjectCache] = []

492 - self._scheduled_seq_group_cache: List[PyObjectCache] = []

493 -

494 - # For async output processing, we need to swap cache buffers between

495 - # iterations. I.e. since the output processing is lagged one step,

496 - # we cannot reuse the cached objects immediately when the schedule()

497 - # is called again, but only when schedule() is called the second time.

498 - self.output_proc_callback = output_proc_callback

499 - self.use_async_output_proc = self.output_proc_callback is not None

500 - self.num_cache_iters = 2 if self.use_async_output_proc else 1

501 -

502 - self.cache_id = 0

503 - for i in range(self.num_cache_iters):

504 - self._seq_group_metadata_cache.append(

505 - PyObjectCache(seq_group_metadata_builder))

506 - self._scheduler_running_outputs_cache.append(

507 - PyObjectCache(scheduler_running_outputs_builder))

508 - self._scheduled_seq_group_cache.append(

509 - PyObjectCache(scheduled_seq_group_builder))

510 -

511 - # For async postprocessor, the extra decode run cannot be done

512 - # when the request reaches max_model_len. In this case, the request

513 - # will be stopped during schedule() call and added to this stop list

514 - # for processing and deallocation by the free_finished_seq_groups()

515 - self._async_stopped: List[SequenceGroup] = []

516 -

517 - # List with the chunk sizes to hand out to each sequence depending

518 - # on how many partial prefills are running. This is slightly faster than

519 - # running an integer division every time a prefill is scheduled.

520 - # This splits the budget evenly among all prefills.

521 - self.partial_prefill_budget_lookup_list = [0] * (

522 - self.scheduler_config.max_num_partial_prefills + 1)

523 - self.partial_prefill_budget_lookup_list[0] = (

524 - scheduler_config.max_num_batched_tokens)

525 - for i in range(1, self.scheduler_config.max_num_partial_prefills + 1):

526 - self.partial_prefill_budget_lookup_list[i] = (

527 - scheduler_config.max_num_batched_tokens // i)

528 -

529 - @property

530 - def next_cache_id(self):

531 - return (self.cache_id + 1) % self.num_cache_iters

532 -

533 - @property

534 - def lora_enabled(self) -> bool:

535 - return bool(self.lora_config)

536 -

537 - @property

538 - def num_decoding_tokens_per_seq(self) -> int:

539 - """The number of new tokens."""

540 - return 1

541 -

542 - def add_seq_group(self, seq_group: SequenceGroup) -> None:

543 - # Add sequence groups to the waiting queue.

544 - self.waiting.append(seq_group)

545 -

546 - def _add_seq_group_to_running(self, seq_group: SequenceGroup) -> None:

547 - # Add sequence groups to the running queue.

548 - # Only for testing purposes.

549 - self.running.append(seq_group)

550 -

551 - def _add_seq_group_to_swapped(self, seq_group: SequenceGroup) -> None:

552 - # Add sequence groups to the swapped queue.

553 - # Only for testing purposes.

554 - self.swapped.append(seq_group)

555 -

556 - def abort_seq_group(

557 - self,

558 - request_id: Union[str, Iterable[str]],

559 - seq_id_to_seq_group: Optional[Dict[str, SequenceGroupBase]] = None,

560 - ) -> None:

561 - """Aborts a sequence group with the given ID.

562 -

563 - Check if the sequence group with the given ID

564 - is present in any of the state queue.

565 - If present, remove the sequence group from the state queue.

566 - Also, if any of the sequences in the sequence group is not finished,

567 - free the sequence with status `FINISHED_ABORTED`.

568 - Otherwise, do nothing.

569 -

570 - Args:

571 - request_id: The ID(s) of the sequence group to abort.

572 - seq_id_to_seq_group: helper for groups with n>1

573 - """

574 - if isinstance(request_id, str):

575 - request_id = (request_id, )

576 - request_ids = set(request_id)

577 - seq_id_to_seq_group = seq_id_to_seq_group or {}

578 - for state_queue in [self.waiting, self.running, self.swapped]:

579 - aborted_groups: List[SequenceGroup] = []

580 - for seq_group in state_queue:

581 - # When n>1, seq_group.request_id looks like

582 - # foo_parallel_sample_0, while request_ids is just foo, and we

583 - # should resolve it as real_request_id to match.

584 - if seq_group.request_id in seq_id_to_seq_group:

585 - real_request_id = seq_id_to_seq_group[

586 - seq_group.request_id].group_id

587 - else:

588 - real_request_id = seq_group.request_id

589 - if real_request_id in request_ids:

590 - # Appending aborted group into pending list.

591 - aborted_groups.append(seq_group)

592 - # We can't remove real_request_id in request_ids here,

593 - # because there may be other seq groups sharing the same

594 - # real_request_id

595 - for aborted_group in aborted_groups:

596 - # Remove the sequence group from the state queue.

597 - state_queue.remove(aborted_group)

598 - # Remove the aborted request from the Mamba cache.

599 - self._finished_requests_ids.append(aborted_group.request_id)

600 - for seq in aborted_group.get_seqs():

601 - if seq.is_finished():

602 - continue

603 - seq.status = SequenceStatus.FINISHED_ABORTED

604 - self.free_seq(seq)

605 - if aborted_group.request_id in seq_id_to_seq_group:

606 - del seq_id_to_seq_group[aborted_group.request_id]

607 -

608 - self._free_seq_group_cross_attn_blocks(aborted_group)

609 -

610 - def _free_seq_group_cross_attn_blocks(

611 - self,

612 - seq_group: SequenceGroup,

613 - ) -> None:

614 - """

615 - Free a sequence group from a cross-attention block table.

616 - Has no effect on decoder-only models.

617 - """

618 - if seq_group.is_encoder_decoder():

619 - self.block_manager.free_cross(seq_group)

620 -

621 - def has_unfinished_seqs(self) -> bool:

622 - return (len(self.waiting) != 0 or len(self.running) != 0

623 - or len(self.swapped) != 0)

624 -

625 - def get_prefix_cache_hit_rate(self, device: Device) -> float:

626 - return self.block_manager.get_prefix_cache_hit_rate(device)

627 -

628 - def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:

629 - return self.block_manager.reset_prefix_cache(device)

630 -

631 - def get_num_unfinished_seq_groups(self) -> int:

632 - return len(self.waiting) + len(self.running) + len(self.swapped)

633 -

634 - def get_and_reset_finished_requests_ids(self) -> List[str]:

635 - """Flushes the list of request ids of previously finished seq_groups."""

636 - finished_requests_ids = self._finished_requests_ids

637 - self._finished_requests_ids = list()

638 - return finished_requests_ids

639 -

640 - def _schedule_running(

641 - self,

642 - budget: SchedulingBudget,

643 - curr_loras: Optional[Set[int]],

644 - enable_chunking: bool = False,

645 - partial_prefill_metadata: Optional[PartialPrefillMetadata] = None,

646 - ) -> SchedulerRunningOutputs:

647 - """Schedule sequence groups that are running.

648 -

649 - Running queue should include decode and chunked prefill requests.

650 -

651 - Args:

652 - budget: The scheduling budget. The argument is in-place updated

653 - when any decodes are preempted.

654 - curr_loras: Currently batched lora request ids. The argument is

655 - in-place updated when any decodes are preempted.

656 - enable_chunking: If True, seq group can be chunked and only a

657 - chunked number of tokens are scheduled if

658 - `budget.num_batched_tokens` has not enough capacity to schedule

659 - all tokens.

660 - partial_prefill_metadata: information about the partial prefills

661 - that are currently running

662 -

663 - Returns:

664 - SchedulerRunningOutputs.

665 - """

666 - ret: SchedulerRunningOutputs = self._scheduler_running_outputs_cache[

667 - self.cache_id].get_object()

668 - ret.blocks_to_swap_out.clear()

669 - ret.blocks_to_copy.clear()

670 - ret.decode_seq_groups.clear()

671 - ret.prefill_seq_groups.clear()

672 - ret.preempted.clear()

673 - ret.swapped_out.clear()

674 -

675 - ret.num_lookahead_slots = self._get_num_lookahead_slots(

676 - is_prefill=False, enable_chunking=enable_chunking)

677 -

678 - ret.decode_seq_groups_list.clear()

679 - ret.prefill_seq_groups_list.clear()

680 -

681 - # Blocks that need to be swapped or copied before model execution.

682 - blocks_to_swap_out: List[Tuple[int, int]] = ret.blocks_to_swap_out

683 - blocks_to_copy: List[Tuple[int, int]] = ret.blocks_to_copy

684 -

685 - decode_seq_groups: List[ScheduledSequenceGroup] = ret.decode_seq_groups

686 - prefill_seq_groups: List[

687 - ScheduledSequenceGroup] = ret.prefill_seq_groups

688 - preempted: List[SequenceGroup] = ret.preempted

689 - swapped_out: List[SequenceGroup] = ret.swapped_out

690 -

691 - running_queue = self.running

692 - assert len(self._async_stopped) == 0

693 - while running_queue:

694 - seq_group = running_queue[0]

695 - # We discard the cached tokens info here because we don't need it

696 - # for running sequence:

697 - # 1. If a sequence is running with chunked prefill, the cached

698 - # tokens info was already used for the first prefill.

699 - # 2. If a sequence is running with non-chunked prefill, then

700 - # there it's a decoding sequence, and the cached tokens info is

701 - # irrelevant.

702 - num_uncached_new_tokens, _ = \

703 - self._get_num_new_uncached_and_cached_tokens(

704 - seq_group,

705 - SequenceStatus.RUNNING,

706 - enable_chunking,

707 - budget,

708 - partial_prefill_metadata,

709 - )

710 -

711 - num_running_tokens = num_uncached_new_tokens

712 - if num_running_tokens == 0:

713 - # No budget => Stop

714 - break

715 -

716 - running_queue.popleft()

717 -

718 - # With async postprocessor, an extra decode run is done

719 - # to process the final tokens. The check below avoids this extra

720 - # decode run when the model max len is reached, in order to avoid

721 - # a memory overflow.

722 - if (self.use_async_output_proc and seq_group.seqs[0].get_len()

723 - > self.scheduler_config.max_model_len):

724 - self._async_stopped.append(seq_group)

725 - continue

726 -

727 - # NOTE(woosuk): Preemption happens only when there is no available

728 - # slot to keep all the sequence groups in the RUNNING state.

729 - while not self._can_append_slots(seq_group, enable_chunking):

730 - budget.subtract_num_batched_tokens(seq_group.request_id,

731 - num_running_tokens)

732 - num_running_seqs = seq_group.get_max_num_running_seqs()

733 - budget.subtract_num_seqs(seq_group.request_id,

734 - num_running_seqs)

735 -

736 - if (curr_loras is not None and seq_group.lora_int_id > 0

737 - and seq_group.lora_int_id in curr_loras):

738 - curr_loras.remove(seq_group.lora_int_id)

739 -

740 - # Determine victim sequence

741 - cont_loop = True

742 - if running_queue:

743 - # Preempt the lowest-priority sequence group.

744 - victim_seq_group = running_queue.pop()

745 - else:

746 - # No other sequence group can be preempted.

747 - # Preempt the current sequence group.

748 - # Note: This is also where we stop this loop

749 - # (since there is nothing else to preempt)

750 - victim_seq_group = seq_group

751 - cont_loop = False

752 -

753 - # With async postprocessor, before preempting a sequence

754 - # we need to ensure it has no pending async postprocessor

755 - do_preempt = True

756 - if self.use_async_output_proc:

757 - assert self.output_proc_callback is not None

758 - self.output_proc_callback(

759 - request_id=victim_seq_group.request_id)

760 -

761 - # It may be that the async pending "victim_seq_group"

762 - # becomes finished, in which case we simply free it.

763 - if victim_seq_group.is_finished():

764 - self._free_finished_seq_group(victim_seq_group)

765 - do_preempt = False

766 -

767 - # Do preemption

768 - if do_preempt:

769 - preempted_mode = self._preempt(victim_seq_group,

770 - blocks_to_swap_out)

771 - if preempted_mode == PreemptionMode.RECOMPUTE:

772 - preempted.append(victim_seq_group)

773 - else:

774 - swapped_out.append(victim_seq_group)

775 -

776 - if not cont_loop:

777 - break

778 - else:

779 - self._append_slots(seq_group, blocks_to_copy, enable_chunking)

780 - is_prefill = seq_group.is_prefill()

781 -

782 - scheduled_seq_group: ScheduledSequenceGroup = (

783 - self._scheduled_seq_group_cache[

784 - self.cache_id].get_object())

785 - scheduled_seq_group.seq_group = seq_group

786 - if is_prefill:

787 - scheduled_seq_group.token_chunk_size = num_running_tokens

788 - prefill_seq_groups.append(scheduled_seq_group)

789 - ret.prefill_seq_groups_list.append(seq_group)

790 - else:

791 - scheduled_seq_group.token_chunk_size = 1

792 - decode_seq_groups.append(scheduled_seq_group)

793 - ret.decode_seq_groups_list.append(seq_group)

794 -

795 - budget.add_num_batched_tokens(seq_group.request_id,

796 - num_running_tokens)

797 - # OPTIMIZATION: Note that get_max_num_running_seqs is

798 - # expensive. For the default scheduling chase where

799 - # enable_chunking is False, num_seqs are updated before running

800 - # this method, so we don't have to update it again here.

801 - if enable_chunking:

802 - num_running_seqs = seq_group.get_max_num_running_seqs()

803 - budget.add_num_seqs(seq_group.request_id, num_running_seqs)

804 - if curr_loras is not None and seq_group.lora_int_id > 0:

805 - curr_loras.add(seq_group.lora_int_id)

806 -

807 - self._scheduler_running_outputs_cache[self.next_cache_id].reset()

808 - self._scheduled_seq_group_cache[self.next_cache_id].reset()

809 -

810 - return ret

811 -

812 - def _schedule_swapped(

813 - self,

814 - budget: SchedulingBudget,

815 - curr_loras: Optional[Set[int]],

816 - enable_chunking: bool = False,

817 - ) -> SchedulerSwappedInOutputs:

818 - """Schedule sequence groups that are swapped out.

819 -

820 - It schedules swapped requests as long as it fits `budget` and

821 - curr_loras <= max_lora from the scheduling config. The input arguments

822 - `budget` and `curr_loras` are updated based on scheduled seq_groups.

823 -

824 - Args:

825 - budget: The scheduling budget. The argument is in-place updated

826 - when any requests are swapped in.

827 - curr_loras: Currently batched lora request ids. The argument is

828 - in-place updated when any requests are swapped in.

829 - enable_chunking: If True, seq group can be chunked and only a

830 - chunked number of tokens are scheduled if

831 - `budget.num_batched_tokens` has not enough capacity to schedule

832 - all tokens.

833 -

834 - Returns:

835 - SchedulerSwappedInOutputs.

836 - """

837 - # Blocks that need to be swapped or copied before model execution.

838 - blocks_to_swap_in: List[Tuple[int, int]] = []

839 - blocks_to_copy: List[Tuple[int, int]] = []

840 - decode_seq_groups: List[ScheduledSequenceGroup] = []

841 - prefill_seq_groups: List[ScheduledSequenceGroup] = []

842 - infeasible_seq_groups: List[SequenceGroup] = []

843 -

844 - swapped_queue = self.swapped

845 -

846 - leftover_swapped: Deque[SequenceGroup] = deque()

847 - while swapped_queue:

848 - seq_group = swapped_queue[0]

849 -

850 - # If the sequence group cannot be swapped in, stop.

851 - is_prefill = seq_group.is_prefill()

852 - alloc_status = self.block_manager.can_swap_in(

853 - seq_group,

854 - self._get_num_lookahead_slots(is_prefill, enable_chunking))

855 - if alloc_status == AllocStatus.LATER:

856 - break

857 - elif alloc_status == AllocStatus.NEVER:

858 - logger.warning(

859 - "Failing the request %s because there's not enough kv "

860 - "cache blocks to run the entire sequence.",

861 - seq_group.request_id,

862 - )

863 - for seq in seq_group.get_seqs():

864 - seq.status = SequenceStatus.FINISHED_IGNORED

865 - infeasible_seq_groups.append(seq_group)

866 - swapped_queue.popleft()

867 - continue

868 -

869 - lora_int_id = 0

870 - if self.lora_enabled:

871 - lora_int_id = seq_group.lora_int_id

872 - assert curr_loras is not None

873 - assert self.lora_config is not None

874 - if (lora_int_id > 0 and (lora_int_id not in curr_loras)

875 - and len(curr_loras) >= self.lora_config.max_loras):

876 - # We don't have a space for another LoRA, so

877 - # we ignore this request for now.

878 - leftover_swapped.appendleft(seq_group)

879 - swapped_queue.popleft()

880 - continue

881 -

882 - # The total number of sequences in the RUNNING state should not

883 - # exceed the maximum number of sequences.

884 - num_new_seqs = seq_group.get_max_num_running_seqs()

885 - num_new_tokens_uncached, num_new_tokens_cached = (

886 - self._get_num_new_uncached_and_cached_tokens(

887 - seq_group, SequenceStatus.SWAPPED, enable_chunking,

888 - budget))

889 -

890 - if num_new_tokens_uncached == 0 or not budget.can_schedule(

891 - num_new_tokens=num_new_tokens_uncached,

892 - num_new_seqs=num_new_seqs,

893 - ):

894 - self.remove_seq_from_computed_blocks_tracker(

895 - seq_group, SequenceStatus.SWAPPED)

896 - break

897 -

898 - if lora_int_id > 0 and curr_loras is not None:

899 - curr_loras.add(lora_int_id)

900 - swapped_queue.popleft()

901 - self._swap_in(seq_group, blocks_to_swap_in)

902 - self._append_slots(seq_group, blocks_to_copy, enable_chunking)

903 - if is_prefill:

904 - prefill_seq_groups.append(

905 - ScheduledSequenceGroup(

906 - seq_group,

907 - token_chunk_size=num_new_tokens_uncached +

908 - num_new_tokens_cached,

909 - ))

910 - else:

911 - decode_seq_groups.append(

912 - ScheduledSequenceGroup(seq_group, token_chunk_size=1))

913 - budget.add_num_batched_tokens(

914 - seq_group.request_id,

915 - num_batched_tokens=num_new_tokens_uncached,

916 - num_cached_tokens=num_new_tokens_cached,

917 - )

918 - budget.add_num_seqs(seq_group.request_id, num_new_seqs)

919 -

920 - swapped_queue.extendleft(leftover_swapped)

921 -

922 - return SchedulerSwappedInOutputs(

923 - decode_seq_groups=decode_seq_groups,

924 - prefill_seq_groups=prefill_seq_groups,

925 - blocks_to_swap_in=blocks_to_swap_in,

926 - blocks_to_copy=blocks_to_copy,

927 - num_lookahead_slots=self._get_num_lookahead_slots(

928 - is_prefill=False, enable_chunking=enable_chunking),

929 - infeasible_seq_groups=infeasible_seq_groups,

930 - )

931 -

932 - def _get_prompt_limit(self, seq_group: SequenceGroup) -> int:

933 - if self.scheduler_config.chunked_prefill_enabled:

934 - prompt_limit = self.scheduler_config.max_model_len

935 - else:

936 - prompt_limit = min(

937 - self.scheduler_config.max_model_len,

938 - self.scheduler_config.max_num_batched_tokens,

939 - )

940 -

941 - # Model is fine tuned with long context. Return the fine tuned max_len.

942 - if seq_group.lora_request and seq_group.lora_request.long_lora_max_len:

943 - assert prompt_limit <= seq_group.lora_request.long_lora_max_len

944 - return seq_group.lora_request.long_lora_max_len

945 - else:

946 - return prompt_limit

947 -

948 - def _get_priority(self,

949 - seq_group: SequenceGroup) -> Tuple[Optional[int], float]:

950 - """Get the priority of the sequence group.

951 - Highest preference to user-defined priority, followed by arrival time.

952 - Args:

953 - seq_group: The sequence group input.

954 - Returns:

955 - The priority of the sequence group.

956 - """

957 - return seq_group.priority, seq_group.arrival_time

958 -

959 - def _schedule_priority_preemption(

960 - self,

961 - budget: SchedulingBudget,

962 - ) -> int:

963 - """Sorts waiting and running queue. Also, force preempt requests

964 - from the running queue if their priority is lower.

965 - Priority-based preemption is used with the priority policy.

966 - Args:

967 - budget: The scheduling budget. The argument is in-place updated

968 - when any requests are scheduled.

969 - Returns:

970 - A count of priority-based preemptions.

971 - """

972 -

973 - waiting_queue = self.waiting

974 -

975 - running_queue = deque(sorted(self.running, key=self._get_priority))

976 -

977 - blocks_to_swap_out: List[Tuple[int, int]] = []

978 - force_preemption_count = 0

979 -

980 - if waiting_queue:

981 - seq_group = waiting_queue.popleft()

982 - num_new_seqs = seq_group.get_max_num_running_seqs()

983 - num_new_tokens_uncached, _ = \

984 - self._get_num_new_uncached_and_cached_tokens(

985 - seq_group, SequenceStatus.WAITING, False, budget)

986 -

987 - # Only preempt if priority inversion exists

988 - while running_queue and self._get_priority(

989 - running_queue[-1]) > self._get_priority(seq_group):

990 - # Only preempt if waiting sequence cannot be allocated

991 - can_allocate = self.block_manager.can_allocate(seq_group)

992 - if (num_new_tokens_uncached > 0

993 - and can_allocate == AllocStatus.OK

994 - and budget.can_schedule(

995 - num_new_tokens=num_new_tokens_uncached,

996 - num_new_seqs=num_new_seqs,

997 - )):

998 - break

999 -

1000 - # Adjust budget to remove the victim sequence group

1001 - vseq_group = running_queue.pop()

1002 - num_running_tokens_uncached, _ = (

1003 - self._get_num_new_uncached_and_cached_tokens(

1004 - vseq_group, SequenceStatus.RUNNING, False, budget))

1005 - budget.subtract_num_batched_tokens(

1006 - vseq_group.request_id, num_running_tokens_uncached)

1007 - num_running_seqs = vseq_group.get_max_num_running_seqs()

1008 - budget.subtract_num_seqs(vseq_group.request_id,

1009 - num_running_seqs)

1010 -

1011 - # Preempt out the victim sequence group

1012 - self._preempt(vseq_group, blocks_to_swap_out)

1013 - waiting_queue.appendleft(vseq_group)

1014 - force_preemption_count += 1

1015 - # Put the sequence back into the waiting queue

1016 - waiting_queue.appendleft(seq_group)

1017 -

1018 - self.remove_seq_from_computed_blocks_tracker(

1019 - seq_group, SequenceStatus.WAITING)

1020 -

1021 - waiting_queue = deque(sorted(waiting_queue, key=self._get_priority))

1022 -

1023 - self.waiting = waiting_queue

1024 - self.running = running_queue

1025 - return force_preemption_count

1026 -

1027 - def _schedule_prefills(

1028 - self,

1029 - budget: SchedulingBudget,

1030 - curr_loras: Optional[Set[int]],

1031 - enable_chunking: bool = False,

1032 - partial_prefill_metadata: Optional[PartialPrefillMetadata] = None,

1033 - ) -> SchedulerPrefillOutputs:

1034 - """Schedule sequence groups that are in prefill stage.

1035 -

1036 - Note that the current scheduler treats PREEMPTED_FOR_RECOMPUTE

1037 - as a new prefill (that starts from beginning -> most recently generated

1038 - tokens).

1039 -

1040 - It schedules waiting requests as long as it fits `budget` and

1041 - curr_loras <= max_lora from the scheduling config. The input arguments

1042 - `budget` and `curr_loras` are updated based on scheduled seq_groups.

1043 -

1044 - Args:

1045 - budget: The scheduling budget. The argument is in-place updated

1046 - when any requests are scheduled.

1047 - curr_loras: Currently batched lora request ids. The argument is

1048 - in-place updated when any requests are scheduled.

1049 - enable_chunking: If True, seq group can be chunked and only a

1050 - chunked number of tokens are scheduled if

1051 - `budget.num_batched_tokens` has not enough capacity to schedule

1052 - all tokens.

1053 - partial_prefill_metadata: information about the partial prefills

1054 - that are currently running

1055 -

1056 - Returns:

1057 - SchedulerPrefillOutputs.

1058 - """

1059 - if budget.remaining_token_budget() == 0:

1060 - # Do nothing: Can't add any more prefill anyway

1061 - return SchedulerPrefillOutputs(

1062 - seq_groups=[],

1063 - ignored_seq_groups=[],

1064 - num_lookahead_slots=self._get_num_lookahead_slots(

1065 - is_prefill=True, enable_chunking=enable_chunking),

1066 - )

1067 - ignored_seq_groups: List[SequenceGroup] = []

1068 - seq_groups: List[ScheduledSequenceGroup] = []

1069 - using_prompt_embeds: bool = False

1070 -

1071 - waiting_queue = self.waiting

1072 -

1073 - leftover_waiting_sequences: Deque[SequenceGroup] = deque()

1074 - while self._passed_delay(time.time()) and waiting_queue:

1075 - seq_group = waiting_queue[0]

1076 -

1077 - waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)

1078 - assert len(waiting_seqs) == 1, (

1079 - "Waiting sequence group should have only one prompt "

1080 - "sequence.")

1081 - if (partial_prefill_metadata is not None

1082 - and not partial_prefill_metadata.can_schedule(seq_group)):

1083 - leftover_waiting_sequences.appendleft(seq_group)

1084 - waiting_queue.popleft()

1085 - continue

1086 - num_new_tokens_uncached, num_new_tokens_cached = (

1087 - self._get_num_new_uncached_and_cached_tokens(

1088 - seq_group,

1089 - SequenceStatus.WAITING,

1090 - enable_chunking,

1091 - budget,

1092 - partial_prefill_metadata=partial_prefill_metadata,

1093 - ))

1094 - num_new_tokens = num_new_tokens_uncached + num_new_tokens_cached

1095 -

1096 - if not enable_chunking:

1097 - num_prompt_tokens = waiting_seqs[0].get_len()

1098 - assert num_new_tokens == num_prompt_tokens

1099 -

1100 - prompt_limit = self._get_prompt_limit(seq_group)

1101 - if num_new_tokens > prompt_limit:

1102 - logger.warning(

1103 - "Input prompt (%d tokens) is too long"

1104 - " and exceeds limit of %d",

1105 - num_new_tokens,

1106 - prompt_limit,

1107 - )

1108 - for seq in waiting_seqs:

1109 - seq.status = SequenceStatus.FINISHED_IGNORED

1110 - self.remove_seq_from_computed_blocks_tracker(

1111 - seq_group, SequenceStatus.FINISHED_IGNORED)

1112 - ignored_seq_groups.append(seq_group)

1113 - waiting_queue.popleft()

1114 - continue

1115 -

1116 - num_lookahead_slots: int = 0

1117 -

1118 - # If the sequence group cannot be allocated, stop.

1119 - can_allocate = self.block_manager.can_allocate(

1120 - seq_group, num_lookahead_slots=num_lookahead_slots)

1121 - if can_allocate == AllocStatus.LATER:

1122 - self.remove_seq_from_computed_blocks_tracker(

1123 - seq_group, SequenceStatus.WAITING)

1124 - break

1125 - elif can_allocate == AllocStatus.NEVER:

1126 - logger.warning(

1127 - "Input prompt (%d tokens) + lookahead slots (%d) is "

1128 - "too long and exceeds the capacity of block_manager",

1129 - num_new_tokens,

1130 - num_lookahead_slots,

1131 - )

1132 - for seq in waiting_seqs:

1133 - seq.status = SequenceStatus.FINISHED_IGNORED

1134 - self.remove_seq_from_computed_blocks_tracker(

1135 - seq_group, SequenceStatus.FINISHED_IGNORED)

1136 - ignored_seq_groups.append(seq_group)

1137 - waiting_queue.popleft()

1138 - continue

1139 -

1140 - # We cannot mix sequence groups that use prompt embeds and

1141 - # those that do not.

1142 - if len(seq_groups) == 0:

1143 - using_prompt_embeds = seq_group.uses_prompt_embeds()

1144 - if using_prompt_embeds != seq_group.uses_prompt_embeds():

1145 - self.remove_seq_from_computed_blocks_tracker(

1146 - seq_group, SequenceStatus.WAITING)

1147 - leftover_waiting_sequences.appendleft(seq_group)

1148 - waiting_queue.popleft()

1149 - continue

1150 -

1151 - lora_int_id = 0

1152 - if self.lora_enabled:

1153 - lora_int_id = seq_group.lora_int_id

1154 - assert curr_loras is not None

1155 - assert self.lora_config is not None

1156 - if (self.lora_enabled and lora_int_id > 0

1157 - and lora_int_id not in curr_loras

1158 - and len(curr_loras) >= self.lora_config.max_loras):

1159 - # We don't have a space for another LoRA, so

1160 - # we ignore this request for now.

1161 - self.remove_seq_from_computed_blocks_tracker(

1162 - seq_group, SequenceStatus.WAITING)

1163 - leftover_waiting_sequences.appendleft(seq_group)

1164 - waiting_queue.popleft()

1165 - continue

1166 -

1167 - if (budget.num_batched_tokens

1168 - >= self.scheduler_config.max_num_batched_tokens):

1169 - # We've reached the budget limit - since there might be

1170 - # continuous prefills in the running queue, we should break

1171 - # to avoid scheduling any new prefills.

1172 - self.remove_seq_from_computed_blocks_tracker(

1173 - seq_group, SequenceStatus.WAITING)

1174 - break

1175 -

1176 - num_new_seqs = seq_group.get_max_num_running_seqs()

1177 - if num_new_tokens_uncached == 0 or not budget.can_schedule(

1178 - num_new_tokens=num_new_tokens_uncached,

1179 - num_new_seqs=num_new_seqs,

1180 - ):

1181 - self.remove_seq_from_computed_blocks_tracker(

1182 - seq_group, SequenceStatus.WAITING)

1183 - break

1184 -

1185 - # Can schedule this request.

1186 - if curr_loras is not None and lora_int_id > 0:

1187 - curr_loras.add(lora_int_id)

1188 - waiting_queue.popleft()

1189 - self._allocate_and_set_running(seq_group)

1190 -

1191 - if partial_prefill_metadata is not None:

1192 - partial_prefill_metadata.maybe_increment_partial_prefills(

1193 - seq_group)

1194 -

1195 - seq_groups.append(

1196 - ScheduledSequenceGroup(seq_group=seq_group,

1197 - token_chunk_size=num_new_tokens))

1198 - budget.add_num_batched_tokens(

1199 - seq_group.request_id,

1200 - num_batched_tokens=num_new_tokens_uncached,

1201 - num_cached_tokens=num_new_tokens_cached,

1202 - )

1203 - budget.add_num_seqs(seq_group.request_id, num_new_seqs)

1204 -

1205 - # Queue requests that couldn't be scheduled.

1206 - waiting_queue.extendleft(leftover_waiting_sequences)

1207 - if len(seq_groups) > 0:

1208 - self.prev_prompt = True

1209 -

1210 - return SchedulerPrefillOutputs(

1211 - seq_groups=seq_groups,

1212 - ignored_seq_groups=ignored_seq_groups,

1213 - num_lookahead_slots=self._get_num_lookahead_slots(

1214 - is_prefill=True, enable_chunking=enable_chunking),

1215 - )

1216 -

1217 - def _schedule_default(self) -> SchedulerOutputs:

1218 - """Schedule queued requests.

1219 -

1220 - The current policy is designed to optimize the throughput. First,

1221 - it batches as many prefill requests as possible. And it schedules

1222 - decodes. If there's a pressure on GPU memory, decode requests can

1223 - be swapped or preempted.

1224 - """

1225 - # Include running requests to the budget.

1226 - budget = SchedulingBudget(

1227 - token_budget=self.scheduler_config.max_num_batched_tokens,

1228 - max_num_seqs=self.scheduler_config.max_num_seqs,

1229 - )

1230 - # Make sure we include num running seqs before scheduling prefill,

1231 - # so that we don't schedule beyond max_num_seqs for prefill.

1232 - for seq_group in self.running:

1233 - budget.add_num_seqs(seq_group.request_id,

1234 - seq_group.get_max_num_running_seqs())

1235 - curr_loras = (set(

1236 - seq_group.lora_int_id for seq_group in self.running

1237 - if seq_group.lora_int_id > 0) if self.lora_enabled else None)

1238 -

1239 - prefills = SchedulerPrefillOutputs.create_empty()

1240 - running_scheduled = SchedulerRunningOutputs.create_empty()

1241 - swapped_in = SchedulerSwappedInOutputs.create_empty()

1242 -

1243 - # If any requests are swapped, prioritized swapped requests.

1244 - if not self.swapped:

1245 - prefills = self._schedule_prefills(budget,

1246 - curr_loras,

1247 - enable_chunking=False)

1248 -

1249 - if len(prefills.seq_groups

1250 - ) == 0 and self.scheduler_config.policy == "priority":

1251 - self._schedule_priority_preemption(budget)

1252 -

1253 - # Don't schedule decodes if prefills are scheduled.

1254 - # NOTE: If `_schedule_prefills` doesn't enable chunking, self.running

1255 - # only contains decode requests, not chunked prefills.

1256 - if len(prefills.seq_groups) == 0:

1257 - running_scheduled = self._schedule_running(budget,

1258 - curr_loras,

1259 - enable_chunking=False)

1260 -

1261 - # If any sequence group is preempted, do not swap in any sequence

1262 - # group. because it means there's no slot for new running requests.

1263 - if (len(running_scheduled.preempted) +

1264 - len(running_scheduled.swapped_out) == 0):

1265 - swapped_in = \

1266 - self._schedule_swapped(budget, curr_loras)

1267 -

1268 - assert (budget.num_batched_tokens

1269 - <= self.scheduler_config.max_num_batched_tokens)

1270 - assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs

1271 -

1272 - # Update waiting requests.

1273 - self.waiting.extendleft(running_scheduled.preempted)

1274 - # Update new running requests.

1275 - if len(prefills.seq_groups) > 0:

1276 - self.running.extend([s.seq_group for s in prefills.seq_groups])

1277 -

1278 - self.running.extend(running_scheduled.decode_seq_groups_list)

1279 -

1280 - if len(swapped_in.decode_seq_groups) > 0:

1281 - self.running.extend(

1282 - [s.seq_group for s in swapped_in.decode_seq_groups])

1283 -

1284 - # Update swapped requests.

1285 - self.swapped.extend(running_scheduled.swapped_out)

1286 - preempted = len(running_scheduled.preempted) + len(

1287 - running_scheduled.swapped_out)

1288 -

1289 - # There should be no prefill from running queue because this policy

1290 - # doesn't allow chunked prefills.

1291 - assert len(running_scheduled.prefill_seq_groups) == 0

1292 - assert len(swapped_in.prefill_seq_groups) == 0

1293 -

1294 - # Merge lists

1295 - num_prefill_groups = len(prefills.seq_groups)

1296 - ignored_seq_groups_for_embeds = list[SequenceGroup]()

1297 - if num_prefill_groups > 0:

1298 - scheduled_seq_groups = prefills.seq_groups

1299 - scheduled_seq_groups.extend(running_scheduled.decode_seq_groups)

1300 - ignored_seq_groups_for_embeds.clear()

1301 - else:

1302 - scheduled_seq_groups = running_scheduled.decode_seq_groups

1303 - if len(scheduled_seq_groups) > 0:

1304 - using_prompt_embeds = scheduled_seq_groups[

1305 - 0].seq_group.uses_prompt_embeds()

1306 - ignored_seq_groups_for_embeds.clear()

1307 - indices_ignored = list[int]()

1308 - for i, schedule_seq_group in enumerate(scheduled_seq_groups):

1309 - if using_prompt_embeds !=\

1310 - schedule_seq_group.seq_group.uses_prompt_embeds():

1311 - ignored_seq_groups_for_embeds.append(

1312 - schedule_seq_group.seq_group)

1313 - indices_ignored.append(i)

1314 - if len(ignored_seq_groups_for_embeds) > 0:

1315 - scheduled_seq_groups = [

1316 - group for i, group in enumerate(scheduled_seq_groups)

1317 - if i not in indices_ignored

1318 - ]

1319 - else:

1320 - ignored_seq_groups_for_embeds.clear()

1321 -

1322 - scheduled_seq_groups.extend(swapped_in.decode_seq_groups)

1323 -

1324 - blocks_to_copy = running_scheduled.blocks_to_copy

1325 - blocks_to_copy.extend(swapped_in.blocks_to_copy)

1326 -

1327 - ignored_seq_groups = prefills.ignored_seq_groups

1328 - ignored_seq_groups.extend(ignored_seq_groups_for_embeds)

1329 - ignored_seq_groups.extend(swapped_in.infeasible_seq_groups)

1330 -

1331 - return SchedulerOutputs(

1332 - scheduled_seq_groups=scheduled_seq_groups,

1333 - num_prefill_groups=num_prefill_groups,

1334 - num_batched_tokens=budget.num_batched_tokens +

1335 - budget.num_cached_tokens,

1336 - blocks_to_swap_in=swapped_in.blocks_to_swap_in,

1337 - blocks_to_swap_out=running_scheduled.blocks_to_swap_out,

1338 - blocks_to_copy=blocks_to_copy,

1339 - ignored_seq_groups=ignored_seq_groups,

1340 - num_lookahead_slots=running_scheduled.num_lookahead_slots,

1341 - running_queue_size=len(self.running),

1342 - preempted=preempted,

1343 - )

1344 -

1345 - def _schedule_chunked_prefill(self) -> SchedulerOutputs:

1346 - """Schedule queued requests.

1347 -

1348 - Chunked prefill allows to chunk prefill requests, batch them together

1349 - with decode requests. This policy 1. schedule as many decoding requests

1350 - as possible. 2. schedule chunked prefill requests that are not

1351 - finished. 3. schedule swapped request. 4. schedule new prefill

1352 - requests.

1353 -

1354 - The policy can sustain the high GPU utilization because it can put

1355 - prefill and decodes requests to the same batch, while it improves

1356 - inter token latency because decodes requests don't need to be blocked

1357 - by prefill requests.

1358 - """

1359 - budget = SchedulingBudget(

1360 - token_budget=self.scheduler_config.max_num_batched_tokens,

1361 - max_num_seqs=self.scheduler_config.max_num_seqs,

1362 - )

1363 - curr_loras: Set[int] = set()

1364 -

1365 - prefills = SchedulerPrefillOutputs.create_empty()

1366 - swapped_in = SchedulerSwappedInOutputs.create_empty()

1367 -

1368 - # Create partial prefill metadata

1369 - partial_prefill_metadata = PartialPrefillMetadata.from_queues(

1370 - running=self.running,

1371 - waiting=self.waiting,

1372 - scheduler_config=self.scheduler_config,

1373 - )

1374 -

1375 - # Decoding should be always scheduled first by fcfs.

1376 - running_scheduled = self._schedule_running(

1377 - budget,

1378 - curr_loras,

1379 - enable_chunking=True,

1380 - partial_prefill_metadata=partial_prefill_metadata,

1381 - )

1382 -

1383 - # Schedule swapped out requests.

1384 - # If preemption happens, it means we don't have space for swap-in.

1385 - if len(running_scheduled.preempted) + len(

1386 - running_scheduled.swapped_out) == 0:

1387 - swapped_in = self._schedule_swapped(budget, curr_loras)

1388 -

1389 - prefills = self._schedule_prefills(

1390 - budget,

1391 - curr_loras,

1392 - enable_chunking=True,

1393 - partial_prefill_metadata=partial_prefill_metadata,

1394 - )

1395 -

1396 - assert (budget.num_batched_tokens

1397 - <= self.scheduler_config.max_num_batched_tokens)

1398 - assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs

1399 -

1400 - # Update waiting requests.

1401 - self.waiting.extendleft(running_scheduled.preempted)

1402 -

1403 - # Update new running requests.

1404 - # By default, vLLM scheduler prioritizes prefills.

1405 - # Once chunked prefill is enabled,

1406 - # the policy is changed to prioritize decode requests.

1407 - self.running.extend(

1408 - [s.seq_group for s in swapped_in.decode_seq_groups])

1409 - self.running.extend(

1410 - [s.seq_group for s in swapped_in.prefill_seq_groups])

1411 - self.running.extend(

1412 - [s.seq_group for s in running_scheduled.decode_seq_groups])

1413 - # Because multiple prefills may be running concurrently, we need to

1414 - # make sure that prefills which are scheduled to finish are listed

1415 - # before those that won't. This is so that on the next scheduling

1416 - # iteration when they have transitioned to the decode stage, they are

1417 - # properly prioritized over sequences that are still in the prefill

1418 - # stage.

1419 - self.running.extend(

1420 - self._order_finishing_prefills_first(

1421 - running_scheduled.prefill_seq_groups))

1422 - self.running.extend([s.seq_group for s in prefills.seq_groups])

1423 -

1424 - # Update swapped requests.

1425 - self.swapped.extend(running_scheduled.swapped_out)

1426 - # Put prefills first due to Attention backend ordering assumption.

1427 - scheduled_seq_groups = (prefills.seq_groups +

1428 - running_scheduled.prefill_seq_groups +

1429 - swapped_in.prefill_seq_groups +

1430 - running_scheduled.decode_seq_groups +

1431 - swapped_in.decode_seq_groups)

1432 - num_prefill_groups = (len(prefills.seq_groups) +

1433 - len(swapped_in.prefill_seq_groups) +

1434 - len(running_scheduled.prefill_seq_groups))

1435 - return SchedulerOutputs(

1436 - scheduled_seq_groups=scheduled_seq_groups,

1437 - num_prefill_groups=num_prefill_groups,

1438 - num_batched_tokens=budget.num_batched_tokens +

1439 - budget.num_cached_tokens,

1440 - blocks_to_swap_in=swapped_in.blocks_to_swap_in,

1441 - blocks_to_swap_out=running_scheduled.blocks_to_swap_out,

1442 - blocks_to_copy=running_scheduled.blocks_to_copy +

1443 - swapped_in.blocks_to_copy,

1444 - ignored_seq_groups=prefills.ignored_seq_groups +

1445 - swapped_in.infeasible_seq_groups,

1446 - num_lookahead_slots=0,

1447 - running_queue_size=len(self.running),

1448 - preempted=(len(running_scheduled.preempted) +

1449 - len(running_scheduled.swapped_out)),

1450 - )

1451 -

1452 - def _order_finishing_prefills_first(

1453 - self, scheduled_prefill_seqs: List[ScheduledSequenceGroup]

1454 - ) -> List[SequenceGroup]:

1455 - """Returns a list of prefilling SequenceGroups where sequences that are

1456 - scheduled to finish prefilling are listed first"""

1457 - finishing = [

1458 - s.seq_group for s in scheduled_prefill_seqs

1459 - if s.seq_group.get_num_uncomputed_tokens() == s.token_chunk_size

1460 - ]

1461 - not_finishing = [

1462 - s.seq_group for s in scheduled_prefill_seqs

1463 - if s.seq_group.get_num_uncomputed_tokens() != s.token_chunk_size

1464 - ]

1465 - return finishing + not_finishing

1466 -

1467 - def _schedule(self) -> SchedulerOutputs:

1468 - """Schedule queued requests."""

1469 - if self.scheduler_config.chunked_prefill_enabled:

1470 - return self._schedule_chunked_prefill()

1471 - else:

1472 - return self._schedule_default()

1473 -

1474 - def _can_append_slots(self, seq_group: SequenceGroup,

1475 - enable_chunking: bool) -> bool:

1476 - """Determine whether or not we have enough space in the KV cache to

1477 - continue generation of the sequence group.

1478 - """

1479 - # It is True only for testing case to trigger artificial preemption.

1480 - if (self.enable_artificial_preemption

1481 - and random.uniform(0, 1) < ARTIFICIAL_PREEMPTION_PROB

1482 - and self.artificial_preempt_cnt > 0):

1483 - self.artificial_preempt_cnt -= 1

1484 - return False

1485 -

1486 - is_prefill = seq_group.is_prefill()

1487 - num_lookahead_slots = self._get_num_lookahead_slots(

1488 - is_prefill, enable_chunking)

1489 -

1490 - return self.block_manager.can_append_slots(

1491 - seq_group=seq_group, num_lookahead_slots=num_lookahead_slots)

1492 -

1493 - def _allow_async_output_proc(self, seq_group: SequenceGroup) -> bool:

1494 - # async_output_proc is allowed only when we have a single sequence

1495 - # in the sequence group

1496 - no_single_seq = seq_group.sampling_params is None or (

1497 - seq_group.sampling_params.n == 1)

1498 - return no_single_seq

1499 -

1500 - def schedule(

1501 - self

1502 - ) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs, bool]:

1503 - # Schedule sequence groups.

1504 - # This function call changes the internal states of the scheduler

1505 - # such as self.running, self.swapped, and self.waiting.

1506 - scheduler_start_time = time.perf_counter()

1507 -

1508 - scheduler_outputs: SchedulerOutputs = self._schedule()

1509 - now = time.time()

1510 -

1511 - if not self.cache_config.enable_prefix_caching:

1512 - common_computed_block_nums = []

1513 -

1514 - allow_async_output_proc: bool = self.use_async_output_proc

1515 -

1516 - # Create input data structures.

1517 - seq_group_metadata_list: List[SequenceGroupMetadata] = []

1518 - for i, scheduled_seq_group in enumerate(

1519 - scheduler_outputs.scheduled_seq_groups):

1520 - seq_group = scheduled_seq_group.seq_group

1521 - token_chunk_size = scheduled_seq_group.token_chunk_size

1522 - seq_group.maybe_set_first_scheduled_time(now)

1523 -

1524 - seq_group_metadata = self._seq_group_metadata_cache[

1525 - self.cache_id].get_object()

1526 - seq_group_metadata.seq_data.clear()

1527 - seq_group_metadata.block_tables.clear()

1528 -

1529 - # seq_id -> SequenceData

1530 - seq_data: Dict[int, SequenceData] = {}

1531 - # seq_id -> physical block numbers

1532 - block_tables: Dict[int, List[int]] = {}

1533 -

1534 - if seq_group.is_encoder_decoder():

1535 - # Encoder associated with SequenceGroup

1536 - encoder_seq = seq_group.get_encoder_seq()

1537 - assert encoder_seq is not None

1538 - encoder_seq_data = encoder_seq.data

1539 - # Block table for cross-attention

1540 - # Also managed at SequenceGroup level

1541 - cross_block_table = self.block_manager.get_cross_block_table(

1542 - seq_group)

1543 - else:

1544 - encoder_seq_data = None

1545 - cross_block_table = None

1546 -

1547 - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):

1548 - seq_id = seq.seq_id

1549 - seq_data[seq_id] = seq.data

1550 - block_tables[seq_id] = self.block_manager.get_block_table(seq)

1551 - self.block_manager.access_all_blocks_in_seq(seq, now)

1552 -

1553 - if self.cache_config.enable_prefix_caching:

1554 - common_computed_block_nums = (

1555 - self.block_manager.get_common_computed_block_ids(

1556 - seq_group.get_seqs(status=SequenceStatus.RUNNING)))

1557 -

1558 - do_sample = True

1559 - is_prompt = seq_group.is_prefill()

1560 - # We should send the metadata to workers when the first prefill

1561 - # is sent. Subsequent requests could be chunked prefill or decode.

1562 - is_first_prefill = False

1563 - if is_prompt:

1564 - seqs = seq_group.get_seqs()

1565 - # Prefill has only 1 sequence.

1566 - assert len(seqs) == 1

1567 - num_computed_tokens = seqs[0].data.get_num_computed_tokens()

1568 - is_first_prefill = num_computed_tokens == 0

1569 - # In the next iteration, all prompt tokens are not computed.

1570 - # It means the prefill is chunked, and we don't need sampling.

1571 - # NOTE: We use get_len instead of get_prompt_len because when

1572 - # a sequence is preempted, prefill includes previous generated

1573 - # output tokens.

1574 - if (token_chunk_size + num_computed_tokens

1575 - < seqs[0].data.get_len()):

1576 - do_sample = False

1577 -

1578 - # It assumes the scheduled_seq_groups is ordered by

1579 - # prefill < decoding.

1580 - if is_first_prefill or not self.scheduler_config.send_delta_data:

1581 - seq_group_metadata = SequenceGroupMetadata(

1582 - request_id=seq_group.request_id,

1583 - is_prompt=is_prompt,

1584 - seq_data=seq_data,

1585 - sampling_params=seq_group.sampling_params,

1586 - block_tables=block_tables,

1587 - do_sample=do_sample,

1588 - pooling_params=seq_group.pooling_params,

1589 - token_chunk_size=token_chunk_size,

1590 - lora_request=seq_group.lora_request,

1591 - computed_block_nums=common_computed_block_nums,

1592 - encoder_seq_data=encoder_seq_data,

1593 - cross_block_table=cross_block_table,

1594 - state=seq_group.state,

1595 - # `multi_modal_data` will only be present for the 1st comm

1596 - # between engine and worker.

1597 - # the subsequent comms can still use delta, but

1598 - # `multi_modal_data` will be None.

1599 - multi_modal_data=(seq_group.multi_modal_data

1600 - if scheduler_outputs.num_prefill_groups

1601 - > 0 else None),

1602 - multi_modal_placeholders=(

1603 - seq_group.multi_modal_placeholders

1604 - if scheduler_outputs.num_prefill_groups > 0 else None),

1605 - )

1606 - else:

1607 - # When SPMD mode is enabled, we only send delta data except for

1608 - # the first request to reduce serialization cost.

1609 - seq_data_delta = {}

1610 - for id, data in seq_data.items():

1611 - seq_data_delta[id] = data.get_delta_and_reset()

1612 - seq_group_metadata = SequenceGroupMetadataDelta(

1613 - seq_data_delta,

1614 - seq_group.request_id,

1615 - block_tables,

1616 - is_prompt,

1617 - do_sample=do_sample,

1618 - token_chunk_size=token_chunk_size,

1619 - computed_block_nums=common_computed_block_nums,

1620 - )

1621 - seq_group_metadata_list.append(seq_group_metadata)

1622 -

1623 - if allow_async_output_proc:

1624 - allow_async_output_proc = self._allow_async_output_proc(

1625 - seq_group)

1626 -

1627 - # Now that the batch has been created, we can assume all blocks in the

1628 - # batch will have been computed before the next scheduling invocation.

1629 - # This is because the engine assumes that a failure in model execution

1630 - # will crash the vLLM instance / will not retry.

1631 - for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups:

1632 - self.block_manager.mark_blocks_as_computed(

1633 - scheduled_seq_group.seq_group,

1634 - scheduled_seq_group.token_chunk_size)

1635 -

1636 - self._seq_group_metadata_cache[self.next_cache_id].reset()

1637 -

1638 - scheduler_time = time.perf_counter() - scheduler_start_time

1639 - # Add this to scheduler time to all the sequences that are currently

1640 - # running. This will help estimate if the scheduler is a significant

1641 - # component in the e2e latency.

1642 - for seq_group in self.running:

1643 - if seq_group is not None and seq_group.metrics is not None:

1644 - if seq_group.metrics.scheduler_time is not None:

1645 - seq_group.metrics.scheduler_time += scheduler_time

1646 - else:

1647 - seq_group.metrics.scheduler_time = scheduler_time

1648 -

1649 - # Move to next cache (if exists)

1650 - self.cache_id = self.next_cache_id

1651 -

1652 - # Return results

1653 - return (seq_group_metadata_list, scheduler_outputs,

1654 - allow_async_output_proc)

1655 -

1656 - def fork_seq(self, parent_seq: Sequence, child_seq: Sequence) -> None:

1657 - self.block_manager.fork(parent_seq, child_seq)

1658 -

1659 - def free_seq(self, seq: Sequence) -> None:

1660 - """Free a sequence from a block table."""

1661 - self.block_manager.free(seq)

1662 -

1663 - def remove_seq_from_computed_blocks_tracker(

1664 - self, seq_group: SequenceGroup,

1665 - status: Optional[SequenceStatus]) -> None:

1666 - seqs = seq_group.get_seqs(status=status)

1667 - for seq in seqs:

1668 - self._remove_seq_from_computed_blocks_tracker(seq)

1669 -

1670 - def _remove_seq_from_computed_blocks_tracker(self, seq: Sequence) -> None:

1671 - """

1672 - Free a sequence computed blocks tracker _seq_id_to_blocks_hashes

1673 - and _seq_id_to_num_tokens_computed.

1674 - """

1675 - self.block_manager.remove_seq_from_computed_blocks_tracker(seq)

1676 -

1677 - def _free_finished_seqs(self, seq_group: SequenceGroup) -> None:

1678 - """Free finished seqs in a sequence group."""

1679 - for seq in seq_group.get_seqs():

1680 - if seq.is_finished():

1681 - self.free_seq(seq)

1682 -

1683 - def _free_finished_seq_group(self, seq_group: SequenceGroup) -> None:

1684 - if seq_group.is_finished():

1685 - # Free cross-attention block table, if it exists

1686 - self._free_seq_group_cross_attn_blocks(seq_group)

1687 -

1688 - # Add the finished requests to the finished requests list.

1689 - # This list will be used to update the Mamba cache in the

1690 - # next step.

1691 - self._finished_requests_ids.append(seq_group.request_id)

1692 -

1693 - # Free finished seqs

1694 - self._free_finished_seqs(seq_group)

1695 -

1696 - def free_finished_seq_groups(self) -> None:

1697 - remaining: Deque[SequenceGroup] = deque()

1698 - for seq_group in self.running:

1699 - self._free_finished_seq_group(seq_group)

1700 - if not seq_group.is_finished():

1701 - remaining.append(seq_group)

1702 -

1703 - self.running = remaining

1704 -

1705 - # Handle async stopped sequence groups

1706 - # (ones that reached max model len)

1707 - if self._async_stopped:

1708 - for seq_group in self._async_stopped:

1709 - self._free_seq_group_cross_attn_blocks(seq_group)

1710 - self._finished_requests_ids.append(seq_group.request_id)

1711 -

1712 - # Free finished seqs

1713 - self._free_finished_seqs(seq_group)

1714 -

1715 - self._async_stopped.clear()

1716 -

1717 - def _allocate_and_set_running(self, seq_group: SequenceGroup) -> None:

1718 - self.block_manager.allocate(seq_group)

1719 - for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):

1720 - seq.status = SequenceStatus.RUNNING

1721 -

1722 - def _append_slots(

1723 - self,

1724 - seq_group: SequenceGroup,

1725 - blocks_to_copy: List[Tuple[int, int]],

1726 - enable_chunking: bool = False,

1727 - ) -> None:

1728 - """Appends new slots to the sequences in the given sequence group.

1729 -

1730 - Args:

1731 - seq_group (SequenceGroup): The sequence group containing the

1732 - sequences to append slots to.

1733 - blocks_to_copy (List[Tuple[int, int]]): A list of tuple of two

1734 - ints, the first int is the source block index, and the second

1735 - int is the destination block index. This list is updated with

1736 - the new source and destination block indices for the appended

1737 - slots.

1738 - enable_chunking (bool): True if chunked prefill is enabled.

1739 - """

1740 - is_prefill: bool = seq_group.is_prefill()

1741 - num_lookahead_slots: int = self._get_num_lookahead_slots(

1742 - is_prefill, enable_chunking)

1743 -

1744 - seq_status: Optional[SequenceStatus] = SequenceStatus.RUNNING

1745 - for seq in seq_group.get_seqs(status=seq_status):

1746 - cows = self.block_manager.append_slots(seq, num_lookahead_slots)

1747 - if len(cows) > 0:

1748 - blocks_to_copy.extend(cows)

1749 -

1750 - def _preempt(self, seq_group: SequenceGroup,

1751 - blocks_to_swap_out: List[Tuple[int, int]]) -> PreemptionMode:

1752 - # If preemption mode is not specified, we determine the mode as follows:

1753 - # We use recomputation by default since it incurs lower overhead than

1754 - # swapping. However, when the sequence group has multiple sequences

1755 - # (e.g., beam search), recomputation is not currently supported. In

1756 - # such a case, we use swapping instead.

1757 - # FIXME(woosuk): This makes our scheduling policy a bit bizarre.

1758 - # As swapped sequences are prioritized over waiting sequences,

1759 - # sequence groups with multiple sequences are implicitly prioritized

1760 - # over sequence groups with a single sequence.

1761 - # TODO(woosuk): Support recomputation for sequence groups with multiple

1762 - # sequences. This may require a more sophisticated CUDA kernel.

1763 - if self.user_specified_preemption_mode is None:

1764 - if seq_group.get_max_num_running_seqs() == 1:

1765 - preemption_mode = PreemptionMode.RECOMPUTE

1766 - else:

1767 - preemption_mode = PreemptionMode.SWAP

1768 -

1769 - elif self.user_specified_preemption_mode == "swap":

1770 - preemption_mode = PreemptionMode.SWAP

1771 - else:

1772 - preemption_mode = PreemptionMode.RECOMPUTE

1773 -

1774 - if self.num_cumulative_preemption % 50 == 0:

1775 - logger.warning(

1776 - "Sequence group %s is preempted by %s mode because there is "

1777 - "not enough KV cache space. This can affect the end-to-end "

1778 - "performance. Increase gpu_memory_utilization or "

1779 - "tensor_parallel_size to provide more KV cache memory. "

1780 - "total_num_cumulative_preemption=%d",

1781 - seq_group.request_id,

1782 - preemption_mode,

1783 - self.num_cumulative_preemption + 1,

1784 - )

1785 - self.num_cumulative_preemption += 1

1786 -

1787 - if preemption_mode == PreemptionMode.RECOMPUTE:

1788 - self._preempt_by_recompute(seq_group)

1789 - elif preemption_mode == PreemptionMode.SWAP:

1790 - self._preempt_by_swap(seq_group, blocks_to_swap_out)

1791 - else:

1792 - raise AssertionError("Invalid preemption mode.")

1793 - return preemption_mode

1794 -

1795 - def _preempt_by_recompute(

1796 - self,

1797 - seq_group: SequenceGroup,

1798 - ) -> None:

1799 - seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)

1800 - assert len(seqs) == 1

1801 - for seq in seqs:

1802 - seq.status = SequenceStatus.WAITING

1803 - self.free_seq(seq)

1804 - seq.reset_state_for_recompute()

1805 - self._free_seq_group_cross_attn_blocks(seq_group)

1806 -

1807 - def _preempt_by_swap(

1808 - self,

1809 - seq_group: SequenceGroup,

1810 - blocks_to_swap_out: List[Tuple[int, int]],

1811 - ) -> None:

1812 - self._swap_out(seq_group, blocks_to_swap_out)

1813 -

1814 - def _swap_in(

1815 - self,

1816 - seq_group: SequenceGroup,

1817 - blocks_to_swap_in: List[Tuple[int, int]],

1818 - ) -> None:

1819 - mapping = self.block_manager.swap_in(seq_group)

1820 - blocks_to_swap_in.extend(mapping)

1821 - for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):

1822 - seq.status = SequenceStatus.RUNNING

1823 -

1824 - def _swap_out(

1825 - self,

1826 - seq_group: SequenceGroup,

1827 - blocks_to_swap_out: List[Tuple[int, int]],

1828 - ) -> None:

1829 - if not self.block_manager.can_swap_out(seq_group):

1830 - # FIXME(woosuk): Abort the sequence group instead of aborting the

1831 - # entire engine.

1832 - raise RuntimeError(

1833 - "Aborted due to the lack of CPU swap space. Please increase "

1834 - "the swap space to avoid this error.")

1835 - mapping = self.block_manager.swap_out(seq_group)

1836 - blocks_to_swap_out.extend(mapping)

1837 - for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):

1838 - seq.status = SequenceStatus.SWAPPED

1839 -

1840 - def _passed_delay(self, now: float) -> bool:

1841 - if self.prev_prompt:

1842 - self.last_prompt_latency = now - self.prev_time

1843 - self.prev_time, self.prev_prompt = now, False

1844 - # Delay scheduling prompts to let waiting queue fill up

1845 - if self.scheduler_config.delay_factor > 0 and self.waiting:

1846 - earliest_arrival_time = min(

1847 - [e.metrics.arrival_time for e in self.waiting])

1848 - passed_delay = ((now - earliest_arrival_time)

1849 - > (self.scheduler_config.delay_factor *

1850 - self.last_prompt_latency) or not self.running)

1851 - else:

1852 - passed_delay = True

1853 - return passed_delay

1854 -

1855 - def _get_num_lookahead_slots(self, is_prefill: bool,

1856 - enable_chunking: bool) -> int:

1857 - """The number of slots to allocate per sequence per step, beyond known

1858 - token ids. Speculative decoding uses these slots to store KV activations

1859 - of tokens which may or may not be accepted.

1860 - """

1861 - return 0

1862 -

1863 - def _get_num_new_uncached_and_cached_tokens(

1864 - self,

1865 - seq_group: SequenceGroup,

1866 - status: SequenceStatus,

1867 - enable_chunking: bool,

1868 - budget: SchedulingBudget,

1869 - partial_prefill_metadata: Optional[PartialPrefillMetadata] = None,

1870 - ) -> Tuple[int, int]:

1871 - """

1872 - Returns the number of new uncached and cached tokens to schedule for a

1873 - given sequence group that's in a given `status`.

1874 -

1875 - The API could chunk the number of tokens to compute based on `budget`

1876 - if `enable_chunking` is True. If a sequence group has multiple

1877 - sequences (e.g., running beam search), it means it is in decoding

1878 - phase, so chunking doesn't happen.

1879 -

1880 - Returns (0, 0) if the new token cannot be computed due to token budget.

1881 -

1882 - The cached tokens's blocks are already computed, and the attention

1883 - backend will reuse the cached blocks rather than recomputing them. So

1884 - the scheduler could schedule these cached tokens "for free".

1885 -

1886 - Args:

1887 - seq_group: The sequence group to get the number of new tokens to

1888 - schedule.

1889 - status: The status of the sequences to get the number of new tokens

1890 - to schedule.

1891 - enable_chunking: Whether to chunk the number of tokens to compute.

1892 - budget: The budget to chunk the number of tokens to compute.

1893 - partial_prefill_metadata: information about the partial prefills

1894 - that are currently running

1895 -

1896 -

1897 - Returns:

1898 - A tuple of two ints. The first int is the number of new uncached

1899 - tokens to schedule. The second int is the number of cached tokens.

1900 - If no more new tokens can be scheduled, returns (0, 0).

1901 - """

1902 - num_cached_new_tokens = 0

1903 - num_uncached_new_tokens = 0

1904 -

1905 - seqs = seq_group.get_seqs(status=status)

1906 - # Compute the number of new uncached and cached tokens for

1907 - # each sequence.

1908 - for seq in seqs:

1909 - if not seq.is_prefill():

1910 - # Decode sequences should always just have 1 uncached token

1911 - # TODO(rickyx): Actually is this still correct for multi-step?

1912 - num_uncached_new_tokens += 1

1913 - continue

1914 -

1915 - num_computed_tokens_seq = seq.get_num_computed_tokens()

1916 - all_num_new_tokens_seq = seq.get_len() - num_computed_tokens_seq

1917 - if not self.cache_config.enable_prefix_caching:

1918 - # If prefix caching is not enabled, all new tokens are uncached.

1919 - num_uncached_new_tokens += all_num_new_tokens_seq

1920 - continue

1921 -

1922 - # NOTE: the cache token might be currently in a block that's in an

1923 - # evictor meaning that it's not yet allocated. However, we don't

1924 - # exclude such tokens in the cache count because it will be

1925 - # guaranteed to be allocated later if the sequence can be allocated.

1926 - num_cached_tokens_seq = self.block_manager.get_num_cached_tokens(

1927 - seq)

1928 -

1929 - # Sanity check.

1930 - if num_cached_tokens_seq < num_computed_tokens_seq:

1931 - # This should only happen with chunked prefill, and

1932 - # the seq is still in prefill. The `num_cached_tokens_seq`

1933 - # is the value we calculated on scheduling the first prefill.

1934 - # For subsequent continuous prefill steps, we cached the

1935 - # number of cache tokens for the sequence so the cached token

1936 - # count could be less than the number of computed tokens.

1937 - # See comments on `ComputedBlocksTracker` for more details.

1938 - assert (

1939 - seq.is_prefill() and seq.status == SequenceStatus.RUNNING

1940 - and self.scheduler_config.chunked_prefill_enabled

1941 - ), ("Number of cached tokens should not be less than the "

1942 - "number of computed tokens for a sequence that's still "

1943 - f"in prefill. But there are {num_cached_tokens_seq} cached "

1944 - f"tokens and {num_computed_tokens_seq} computed tokens "

1945 - f"for sequence {seq.seq_id}.")

1946 -

1947 - num_cached_new_tokens_seq = max(

1948 - 0, num_cached_tokens_seq - num_computed_tokens_seq)

1949 - num_uncached_new_tokens_seq = (all_num_new_tokens_seq -

1950 - num_cached_new_tokens_seq)

1951 -

1952 - num_uncached_new_tokens += num_uncached_new_tokens_seq

1953 - num_cached_new_tokens += num_cached_new_tokens_seq

1954 -

1955 - if num_uncached_new_tokens == 0 and num_cached_new_tokens > 0:

1956 - # For a fully cached hit sequence, we actually need to recompute the

1957 - # last token. So we need at least 1 uncached token to schedule.

1958 - # See ModelRunner._compute_for_prefix_cache_hit for more details.

1959 - num_uncached_new_tokens = 1

1960 - num_cached_new_tokens -= 1

1961 -

1962 - if enable_chunking and len(seqs) == 1:

1963 - # Chunk if a running request cannot fit in the given budget.

1964 - # If number of seq > 1, it means it is doing beam search

1965 - # in a decode phase. Do not chunk.

1966 - num_uncached_new_tokens = self._chunk_new_tokens_to_schedule(

1967 - self.scheduler_config,

1968 - self.cache_config,

1969 - budget,

1970 - self._get_prompt_limit(seq_group),

1971 - num_uncached_new_tokens,

1972 - self.partial_prefill_budget_lookup_list,

1973 - partial_prefill_metadata,

1974 - )

1975 -

1976 - return num_uncached_new_tokens, num_cached_new_tokens

1977 -

1978 - @staticmethod

1979 - def _chunk_new_tokens_to_schedule(

1980 - scheduler_config: SchedulerConfig,

1981 - cache_config: CacheConfig,

1982 - budget: SchedulingBudget,

1983 - prompt_limit: int,

1984 - num_new_tokens: int,

1985 - partial_prefill_budget_lookup_list: List[int],

1986 - partial_prefill_metadata: Optional[PartialPrefillMetadata] = None,

1987 - ) -> int:

1988 - """

1989 - Chunks the number of new tokens to schedule based on the budget when

1990 - chunked prefill is enabled.

1991 -

1992 - Args:

1993 - scheduler_config: The scheduler config.

1994 - cache_config: The cache config.

1995 - budget: The budget to chunk the number of tokens to compute.

1996 - prompt_limit: The maximum number of tokens allowed in a prompt.

1997 - num_new_tokens: The number of new tokens to schedule.

1998 -

1999 - Returns:

2000 - The number of new tokens to schedule after chunking.

2001 - """

2002 - remaining_token_budget = budget.remaining_token_budget()

2003 -

2004 - # Get the number of tokens to allocate to this prefill slot

2005 - prefill_slot_budget = (

2006 - remaining_token_budget if partial_prefill_metadata is None else

2007 - partial_prefill_budget_lookup_list[

2008 - partial_prefill_metadata.schedulable_prefills])

2009 -

2010 - if cache_config.enable_prefix_caching:

2011 - # When prefix caching is enabled and we're partially prefilling

2012 - # a sequence, we always allocate a number of new tokens that is

2013 - # divisible by the block size to avoid partial block matching.

2014 - block_size = cache_config.block_size

2015 - # Don't exceed either the total budget or slot budget.

2016 - # Take min of those and get the next lowest multiple of the

2017 - # block size:

2018 - remaining_token_budget = (

2019 - min(remaining_token_budget, prefill_slot_budget) //

2020 - block_size) * block_size

2021 - # NB: In the case where num_new_tokens < budget, we are

2022 - # finishing prefill for this sequence, so we do not need to

2023 - # allocate a full block.

2024 -

2025 - num_new_tokens = min(num_new_tokens, remaining_token_budget,

2026 - prefill_slot_budget)

2027 -

2028 - return num_new_tokens

vllm/engine/protocol.py

+1 -7

@@ -7,13 +7,11 @@ from typing import Any, AsyncGenerator, Iterable, Mapping, Optional, Union

8 from vllm.beam_search import BeamSearchSequence, create_sort_beams_key_function

9 from vllm.config import ModelConfig, VllmConfig

10 -from vllm.core.scheduler import SchedulerOutputs

10 from vllm.inputs.data import PromptType, TokensPrompt

11 from vllm.inputs.parse import is_explicit_encoder_decoder_prompt

12 from vllm.inputs.preprocess import InputPreprocessor

13 from vllm.logger import init_logger

14 from vllm.lora.request import LoRARequest

16 -from vllm.model_executor.layers.sampler import SamplerOutput

15 from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput

16 from vllm.plugins.io_processors.interface import IOProcessor

17 from vllm.pooling_params import PoolingParams

@@ -266,11 +264,7 @@ class EngineClient(ABC):

264 ...

265

266 @abstractmethod

269 - async def do_log_stats(

270 - self,

271 - scheduler_outputs: Optional[SchedulerOutputs] = None,

272 - model_output: Optional[list[SamplerOutput]] = None,

273 - ) -> None:

267 + async def do_log_stats(self) -> None:

268 ...

269

270 @abstractmethod

vllm/v1/engine/async_llm.py

+1 -5

@@ -601,11 +601,7 @@ class AsyncLLM(EngineClient):

601 async def is_tracing_enabled(self) -> bool:

602 return self.observability_config.otlp_traces_endpoint is not None

603

604 - async def do_log_stats(

605 - self,

606 - scheduler_outputs=None,

607 - model_output=None,

608 - ) -> None:

604 + async def do_log_stats(self) -> None:

605 if self.logger_manager:

606 self.logger_manager.log()

607

vllm/worker/cache_engine.py deleted

-145

@@ -1,145 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""CacheEngine class for managing the KV cache."""

4 -from typing import List

5 -

6 -import torch

7 -

8 -from vllm.attention import get_attn_backend

9 -from vllm.config import CacheConfig, DeviceConfig, ModelConfig, ParallelConfig

10 -from vllm.logger import init_logger

11 -from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, LayerBlockType,

12 - get_dtype_size, is_pin_memory_available)

13 -

14 -logger = init_logger(__name__)

15 -

16 -

17 -class CacheEngine:

18 - """Manages the KV cache.

19 -

20 - This class is responsible for initializing and managing the GPU and CPU KV

21 - caches. It also provides methods for performing KV cache operations, such

22 - as swapping and copying.

23 - """

24 -

25 - def __init__(

26 - self,

27 - cache_config: CacheConfig,

28 - model_config: ModelConfig,

29 - parallel_config: ParallelConfig,

30 - device_config: DeviceConfig,

31 - ) -> None:

32 - self.cache_config = cache_config

33 - self.model_config = model_config

34 - self.parallel_config = parallel_config

35 - self.device_config = device_config

36 -

37 - self.head_size = model_config.get_head_size()

38 - # Models like Jamba, have mixed typed layers, E.g Mamba

39 - self.num_attention_layers = model_config.get_num_layers_by_block_type(

40 - parallel_config, LayerBlockType.attention)

41 - self.num_kv_heads = model_config.get_num_kv_heads(parallel_config)

42 -

43 - self.block_size = cache_config.block_size

44 - self.num_gpu_blocks = cache_config.num_gpu_blocks

45 - if self.num_gpu_blocks:

46 - self.num_gpu_blocks //= parallel_config.pipeline_parallel_size

47 - self.num_cpu_blocks = cache_config.num_cpu_blocks

48 - if self.num_cpu_blocks:

49 - self.num_cpu_blocks //= parallel_config.pipeline_parallel_size

50 -

51 - if cache_config.cache_dtype == "auto":

52 - self.dtype = model_config.dtype

53 - else:

54 - self.dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]

55 -

56 - # Get attention backend.

57 - self.attn_backend = get_attn_backend(self.head_size,

58 - model_config.dtype,

59 - cache_config.cache_dtype,

60 - self.block_size,

61 - model_config.is_attention_free,

62 - use_mla=model_config.use_mla)

63 -

64 - # Initialize the cache.

65 - self.gpu_cache = self._allocate_kv_cache(

66 - self.num_gpu_blocks, self.device_config.device_type)

67 - self.cpu_cache = self._allocate_kv_cache(self.num_cpu_blocks, "cpu")

68 -

69 - def _allocate_kv_cache(

70 - self,

71 - num_blocks: int,

72 - device: str,

73 - ) -> List[torch.Tensor]:

74 - """Allocates KV cache on the specified device."""

75 - kv_cache_generic_shape = self.attn_backend.get_kv_cache_shape(

76 - num_blocks, self.block_size, self.num_kv_heads, self.head_size)

77 - pin_memory = is_pin_memory_available() if device == "cpu" else False

78 - kv_cache: List[torch.Tensor] = []

79 - try:

80 - kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order(

81 - )

82 - except (AttributeError, NotImplementedError):

83 - kv_cache_stride_order = tuple(range(len(kv_cache_generic_shape)))

84 -

85 - # The allocation respects the backend-defined stride order to ensure

86 - # the semantic remains consistent for each backend. We first obtain the

87 - # generic kv cache shape and then permute it according to the stride

88 - # order which could result in a non-contiguous tensor.

89 - kv_cache_allocation_shape = tuple(kv_cache_generic_shape[i]

90 - for i in kv_cache_stride_order)

91 -

92 - for _ in range(self.num_attention_layers):

93 - # null block in CpuGpuBlockAllocator requires at least that

94 - # block to be zeroed-out.

95 - # We zero-out everything for simplicity.

96 - layer_kv_cache = torch.zeros(

97 - kv_cache_allocation_shape,

98 - dtype=self.dtype,

99 - pin_memory=pin_memory,

100 - device=device).permute(*kv_cache_stride_order)

101 -

102 - # view back to (TOTAL_PAGES, PAGE_SIZE, entry_shape...) for cases

103 - # when entry_shape is higher than 1D

104 - kv_cache.append(layer_kv_cache)

105 - return kv_cache

106 -

107 - def swap_in(self, src_to_dst: torch.Tensor) -> None:

108 - for i in range(self.num_attention_layers):

109 - self.attn_backend.swap_blocks(self.cpu_cache[i], self.gpu_cache[i],

110 - src_to_dst)

111 -

112 - def swap_out(self, src_to_dst: torch.Tensor) -> None:

113 - for i in range(self.num_attention_layers):

114 - self.attn_backend.swap_blocks(self.gpu_cache[i], self.cpu_cache[i],

115 - src_to_dst)

116 -

117 - def copy(self, src_to_dsts: torch.Tensor) -> None:

118 - self.attn_backend.copy_blocks(self.gpu_cache, src_to_dsts)

119 -

120 - @staticmethod

121 - def get_cache_block_size(

122 - cache_config: CacheConfig,

123 - model_config: ModelConfig,

124 - parallel_config: ParallelConfig,

125 - ) -> int:

126 - head_size = model_config.get_head_size()

127 - num_heads = model_config.get_num_kv_heads(parallel_config)

128 - num_attention_layers = model_config.get_num_layers_by_block_type(

129 - parallel_config, LayerBlockType.attention)

130 -

131 - if cache_config.cache_dtype == "auto":

132 - dtype = model_config.dtype

133 - else:

134 - dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]

135 -

136 - key_cache_entry = num_heads * head_size

137 -

138 - # For MLA there is no value cache, since the latent vector

139 - # is joint keys and values.

140 - value_cache_entry = key_cache_entry if not model_config.use_mla else 0

141 - total = num_attention_layers * cache_config.block_size * \

142 - (key_cache_entry + value_cache_entry)

143 -

144 - dtype_size = get_dtype_size(dtype)

145 - return dtype_size * total

vllm/worker/model_runner.py deleted

-2031

@@ -1,2031 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -

4 -import dataclasses

5 -import gc

6 -import inspect

7 -import itertools

8 -import time

9 -import weakref

10 -from contextlib import contextmanager

11 -from dataclasses import dataclass

12 -from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set,

13 - Tuple, Type, TypeVar, Union)

14 -

15 -import numpy as np

16 -import torch

17 -import torch.distributed

18 -import torch.nn as nn

19 -from tqdm.auto import tqdm

20 -

21 -from vllm.attention import AttentionMetadata, get_attn_backend

22 -from vllm.attention.backends.abstract import AttentionState

23 -from vllm.attention.backends.utils import CommonAttentionState

24 -from vllm.compilation.counter import compilation_counter

25 -from vllm.config import CompilationLevel, VllmConfig

26 -from vllm.core.scheduler import SchedulerOutputs

27 -from vllm.distributed import broadcast_tensor_dict, get_pp_group

28 -from vllm.distributed.kv_transfer import get_kv_transfer_group

29 -from vllm.distributed.parallel_state import (get_tensor_model_parallel_rank,

30 - graph_capture)

31 -from vllm.forward_context import get_forward_context, set_forward_context

32 -from vllm.inputs import INPUT_REGISTRY, InputRegistry

33 -from vllm.logger import init_logger

34 -from vllm.lora.layers import LoRAMapping

35 -from vllm.lora.request import LoRARequest

36 -from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager

37 -from vllm.model_executor import SamplingMetadata, SamplingMetadataCache

38 -from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding

39 -from vllm.model_executor.layers.sampler import (Sampler, SamplerOutput,

40 - get_sampler)

41 -from vllm.model_executor.model_loader import get_model

42 -from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

43 -from vllm.model_executor.models import (supports_lora, supports_mrope,

44 - supports_multimodal)

45 -from vllm.model_executor.models.utils import set_cpu_offload_max_bytes

46 -from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,

47 - MultiModalKwargs, MultiModalPlaceholderMap,

48 - MultiModalRegistry)

49 -from vllm.sampling_params import SamplingParams

50 -from vllm.sequence import IntermediateTensors, SequenceGroupMetadata

51 -from vllm.utils import (DeviceMemoryProfiler, GiB_bytes, PyObjectCache,

52 - async_tensor_h2d, flatten_2d_lists,

53 - is_pin_memory_available, supports_dynamo,

54 - weak_ref_tensor)

55 -from vllm.worker.model_runner_base import (

56 - InputProcessingError, ModelRunnerBase, ModelRunnerInputBase,

57 - ModelRunnerInputBuilderBase, _add_attn_metadata_broadcastable_dict,

58 - _add_sampling_metadata_broadcastable_dict,

59 - _init_attn_metadata_from_tensor_dict,

60 - _init_sampling_metadata_from_tensor_dict)

61 -

62 -if TYPE_CHECKING:

63 - from vllm.attention.backends.abstract import AttentionBackend

64 -

65 -logger = init_logger(__name__)

66 -

67 -LORA_WARMUP_RANK = 8

68 -

69 -_NUM_WARMUP_ITERS = 2

70 -

71 -TModelInputForGPU = TypeVar('TModelInputForGPU', bound="ModelInputForGPU")

72 -

73 -# For now, bump up cache limits for recompilations during CUDA graph warmups.

74 -torch._dynamo.config.cache_size_limit = 128

75 -torch._dynamo.config.accumulated_cache_size_limit = 128

76 -

77 -

78 -@dataclass(frozen=True)

79 -class ModelInputForGPU(ModelRunnerInputBase):

80 - """

81 - This base class contains metadata needed for the base model forward pass

82 - but not metadata for possible additional steps, e.g., sampling. Model

83 - runners that run additional steps should subclass this method to add

84 - additional fields.

85 - """

86 - input_tokens: Optional[torch.Tensor] = None

87 - inputs_embeds: Optional[torch.Tensor] = None

88 - input_positions: Optional[torch.Tensor] = None

89 - seq_lens: Optional[List[int]] = None

90 - query_lens: Optional[List[int]] = None

91 - lora_mapping: Optional["LoRAMapping"] = None

92 - lora_requests: Optional[Set[LoRARequest]] = None

93 - attn_metadata: Optional["AttentionMetadata"] = None

94 - multi_modal_kwargs: Optional[BatchedTensorInputs] = None

95 - request_ids_to_seq_ids: Optional[Dict[str, List[int]]] = None

96 - finished_requests_ids: Optional[List[str]] = None

97 - virtual_engine: int = 0

98 - async_callback: Optional[Callable] = None

99 - scheduler_outputs: Optional[SchedulerOutputs] = None

100 - previous_hidden_states: Optional[torch.Tensor] = None

101 -

102 - def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:

103 - tensor_dict = {

104 - "input_tokens": self.input_tokens,

105 - "inputs_embeds": self.inputs_embeds,

106 - "input_positions": self.input_positions,

107 - "lora_requests": self.lora_requests,

108 - "lora_mapping": self.lora_mapping,

109 - "multi_modal_kwargs": self.multi_modal_kwargs,

110 - "virtual_engine": self.virtual_engine,

111 - "request_ids_to_seq_ids": self.request_ids_to_seq_ids,

112 - "finished_requests_ids": self.finished_requests_ids,

113 - }

114 - _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)

115 - return tensor_dict

116 -

117 - @classmethod

118 - def from_broadcasted_tensor_dict(

119 - cls: Type[TModelInputForGPU],

120 - tensor_dict: Dict[str, Any],

121 - attn_backend: Optional["AttentionBackend"] = None,

122 - ) -> TModelInputForGPU:

123 - if attn_backend is not None:

124 - tensor_dict = _init_attn_metadata_from_tensor_dict(

125 - attn_backend, tensor_dict)

126 - return cls(**tensor_dict)

127 -

128 - # Exclude `async_callback` to be able to pickle this object

129 - def __getstate__(self):

130 - state = self.__dict__.copy()

131 - del state["async_callback"]

132 - return state

133 -

134 - # TODO: What happens when we depickle this object?

135 - # How can we update this callback to properly pass it to the engine?

136 - def __setstate__(self, state):

137 - self.__dict__.update(state)

138 - self.__dict__.update({'async_callback': None})

139 -

140 -

141 -@dataclass(frozen=True)

142 -class ModelInputForGPUWithSamplingMetadata(ModelInputForGPU):

143 - """

144 - Used by the ModelRunner.

145 - """

146 - sampling_metadata: Optional["SamplingMetadata"] = None

147 - # Used for speculative decoding. We do not broadcast it because it is only

148 - # used by the driver worker.

149 - is_prompt: Optional[bool] = None

150 -

151 - def as_broadcastable_tensor_dict(self) -> Dict[str, Any]:

152 - tensor_dict = {

153 - "input_tokens": self.input_tokens,

154 - "inputs_embeds": self.inputs_embeds,

155 - "input_positions": self.input_positions,

156 - "lora_requests": self.lora_requests,

157 - "lora_mapping": self.lora_mapping,

158 - "multi_modal_kwargs": self.multi_modal_kwargs,

159 - "virtual_engine": self.virtual_engine,

160 - "request_ids_to_seq_ids": self.request_ids_to_seq_ids,

161 - "finished_requests_ids": self.finished_requests_ids,

162 - }

163 - _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)

164 - _add_sampling_metadata_broadcastable_dict(tensor_dict,

165 - self.sampling_metadata)

166 - return tensor_dict

167 -

168 - @classmethod

169 - def from_broadcasted_tensor_dict(

170 - cls,

171 - tensor_dict: Dict[str, Any],

172 - attn_backend: Optional["AttentionBackend"] = None,

173 - ) -> "ModelInputForGPUWithSamplingMetadata":

174 - tensor_dict = _init_sampling_metadata_from_tensor_dict(tensor_dict)

175 - if attn_backend is not None:

176 - tensor_dict = _init_attn_metadata_from_tensor_dict(

177 - attn_backend, tensor_dict)

178 - return cls(**tensor_dict)

179 -

180 -

181 -class ModelInputForGPUBuilder(ModelRunnerInputBuilderBase[ModelInputForGPU]):

182 - """Build ModelInputForGPU from SequenceGroupMetadata."""

183 -

184 - # Note: ideally we would be using a dataclass(kw_only=True)

185 - # here, so that this can be subclassed easily,

186 - # but kw_only is not supported in python<3.10.

187 - class InterDataForSeqGroup:

188 - """Intermediate data for the current sequence group."""

189 -

190 - def simple_reinit(self):

191 - self.input_tokens[0].clear() # type: ignore

192 - self.inputs_embeds = None # type: ignore

193 - self.input_positions[0].clear() # type: ignore

194 - self.mrope_input_positions = None # type: ignore

195 - self.seq_lens[0] = 0 # type: ignore

196 - self.orig_seq_lens[0] = 0 # type: ignore

197 - self.prompt_lens[0] = 0 # type: ignore

198 - self.query_lens[0] = 0 # type: ignore

199 - self.context_lens[0] = 0 # type: ignore

200 - self.curr_sliding_window_blocks[0] = 0 # type: ignore

201 - self.lora_index_mapping.clear() # type: ignore

202 - self.lora_prompt_mapping.clear() # type: ignore

203 - self.lora_requests.clear() # type: ignore

204 -

205 - def __init__(

206 - self,

207 - *,

208 - # From sequence group metadata.

209 - request_id: str,

210 - seq_ids: List[int],

211 - is_prompt: bool,

212 - block_tables: Optional[Dict[int, List[int]]],

213 - computed_block_nums: List[int],

214 - n_seqs: int = 0,

215 -

216 - # Input tokens and positions.

217 - input_tokens: Optional[List[List[int]]] = None,

218 - inputs_embeds: Optional[torch.Tensor] = None,

219 - input_positions: Optional[List[List[int]]] = None,

220 - mrope_input_positions: Optional[List[List[List[int]]]] = None,

221 -

222 - # The sequence length (may be capped to the sliding window).

223 - seq_lens: Optional[List[int]] = None,

224 - # The original sequence length (before applying sliding window).

225 - # This is used to compute slot mapping.

226 - orig_seq_lens: Optional[List[int]] = None,

227 - # This is used in the dual-chunk flash attention backend.

228 - prompt_lens: Optional[List[int]] = None,

229 - # The query length.

230 - query_lens: Optional[List[int]] = None,

231 - # The number of tokens that are already computed.

232 - context_lens: Optional[List[int]] = None,

233 - # The current sliding window block.

234 - curr_sliding_window_blocks: Optional[List[int]] = None,

235 -

236 - # LoRA inputs.

237 - lora_index_mapping: Optional[List[List[int]]] = None,

238 - lora_prompt_mapping: Optional[List[List[int]]] = None,

239 - lora_requests: Optional[Set[LoRARequest]] = None,

240 -

241 - # Multi-modal inputs.

242 - multi_modal_kwargs: Optional[MultiModalKwargs] = None,

243 - multi_modal_placeholder_maps: Optional[Dict[

244 - str, MultiModalPlaceholderMap]] = None,

245 -

246 - # Whether the prefix cache is hit (prefill only).

247 - prefix_cache_hit: bool = False,

248 - reinit: bool = False,

249 - reinit_use_defaults: bool = False,

250 - encoder_seq_len: int = 0,

251 - ):

252 - if reinit:

253 - assert len(self.seq_ids) == len(seq_ids) # type: ignore

254 - for i, seq_id in enumerate(seq_ids):

255 - self.seq_ids[i] = seq_id # type: ignore

256 - else:

257 - self.seq_ids = seq_ids

258 -

259 - self.request_id = request_id

260 - self.is_prompt = is_prompt

261 - self.block_tables = block_tables

262 - self.computed_block_nums = computed_block_nums

263 - self.n_seqs = n_seqs

264 - self.encoder_seq_len = encoder_seq_len

265 -

266 - if reinit:

267 - if len(self.seq_ids) == 1 and reinit_use_defaults:

268 - self.simple_reinit()

269 - else:

270 - if input_tokens:

271 - self.input_tokens = input_tokens

272 - else:

273 - for seq_id in range(len(self.seq_ids)):

274 - self.input_tokens[seq_id].clear()

275 -

276 - self.inputs_embeds = inputs_embeds

277 -

278 - if input_positions:

279 - self.input_positions = input_positions

280 - else:

281 - for seq_id in range(len(self.seq_ids)):

282 - self.input_positions[seq_id].clear()

283 -

284 - self.mrope_input_positions = None

285 -

286 - if seq_lens:

287 - self.seq_lens = seq_lens

288 - else:

289 - for seq_id in range(len(self.seq_ids)):

290 - self.seq_lens[seq_id] = 0

291 -

292 - if orig_seq_lens:

293 - self.orig_seq_lens = orig_seq_lens

294 - else:

295 - for seq_id in range(len(self.seq_ids)):

296 - self.orig_seq_lens[seq_id] = 0

297 -

298 - if prompt_lens:

299 - self.prompt_lens = prompt_lens

300 - else:

301 - for seq_id in range(len(self.seq_ids)):

302 - self.prompt_lens[seq_id] = 0

303 -

304 - if query_lens:

305 - self.query_lens = query_lens

306 - else:

307 - for seq_id in range(len(self.seq_ids)):

308 - self.query_lens[seq_id] = 0

309 -

310 - if context_lens:

311 - self.context_lens = context_lens

312 - else:

313 - for seq_id in range(len(self.seq_ids)):

314 - self.context_lens[seq_id] = 0

315 -

316 - if curr_sliding_window_blocks:

317 - self.curr_sliding_window_blocks = \

318 - curr_sliding_window_blocks

319 - else:

320 - for seq_id in range(len(self.seq_ids)):

321 - self.curr_sliding_window_blocks[seq_id] = 0

322 -

323 - if lora_index_mapping:

324 - self.lora_index_mapping = lora_index_mapping

325 - else:

326 - self.lora_index_mapping.clear()

327 -

328 - if lora_prompt_mapping:

329 - self.lora_prompt_mapping = lora_prompt_mapping

330 - else:

331 - self.lora_prompt_mapping.clear()

332 -

333 - if lora_requests:

334 - self.lora_requests = lora_requests

335 - else:

336 - self.lora_requests.clear()

337 -

338 - else:

339 - self.input_tokens = input_tokens or []

340 - self.inputs_embeds = inputs_embeds

341 - self.input_positions = input_positions or []

342 - self.mrope_input_positions = mrope_input_positions or None

343 - self.seq_lens = seq_lens or []

344 - self.orig_seq_lens = orig_seq_lens or []

345 - self.prompt_lens = prompt_lens or []

346 - self.query_lens = query_lens or []

347 - self.context_lens = context_lens or []

348 - self.curr_sliding_window_blocks = \

349 - curr_sliding_window_blocks or []

350 -

351 - self.lora_index_mapping = lora_index_mapping or []

352 - self.lora_prompt_mapping = lora_prompt_mapping or []

353 - self.lora_requests = lora_requests or set()

354 -

355 - self.multi_modal_kwargs = multi_modal_kwargs

356 - self.multi_modal_placeholder_maps = multi_modal_placeholder_maps

357 - self.prefix_cache_hit = prefix_cache_hit

358 -

359 - self.n_seqs = len(self.seq_ids)

360 -

361 - if not reinit:

362 - self.__post_init__()

363 -

364 - def __post_init__(self):

365 - self.n_seqs = len(self.seq_ids)

366 -

367 - self.input_tokens = [[] for _ in range(self.n_seqs)]

368 - self.input_positions = [[] for _ in range(self.n_seqs)]

369 - self.mrope_input_positions = None

370 - self.seq_lens = [0] * self.n_seqs

371 - self.orig_seq_lens = [0] * self.n_seqs

372 - self.prompt_lens = [0] * self.n_seqs

373 - self.query_lens = [0] * self.n_seqs

374 - self.context_lens = [0] * self.n_seqs

375 - self.curr_sliding_window_blocks = [0] * self.n_seqs

376 -

377 - self.lora_index_mapping = []

378 - self.lora_prompt_mapping = []

379 -

380 - def __repr__(self) -> str:

381 - return (f"InterDataForSeqGroup("

382 - f"request_id={self.request_id}, "

383 - f"seq_ids={self.seq_ids}, "

384 - f"is_prompt={self.is_prompt}, "

385 - f"block_tables={self.block_tables}, "

386 - f"computed_block_nums={self.computed_block_nums}, "

387 - f"n_seqs={self.n_seqs}, "

388 - f"input_tokens={self.input_tokens}, "

389 - f"inputs_embeds.shape="

390 - f"{getattr(self.inputs_embeds, 'shape', None)}, "

391 - f"input_positions={self.input_positions}, "

392 - f"mrope_input_positions={self.mrope_input_positions}, "

393 - f"seq_lens={self.seq_lens}, "

394 - f"orig_seq_lens={self.orig_seq_lens}, "

395 - f"query_lens={self.query_lens}, "

396 - f"context_lens={self.context_lens}, "

397 - f"multi_modal_kwargs={self.multi_modal_kwargs}")

398 -

399 - def gen_inter_data_builder(self, num_seqs: int):

400 - return lambda: ModelInputForGPUBuilder.InterDataForSeqGroup(

401 - request_id="",

402 - seq_ids=[0] * num_seqs,

403 - is_prompt=True,

404 - block_tables=None,

405 - computed_block_nums=[])

406 -

407 - def init_cached_inter_data(self, *args, **kwargs):

408 - assert len(args) == 0

409 - assert "seq_ids" in kwargs

410 - seq_ids = kwargs["seq_ids"]

411 - num_seqs = len(seq_ids)

412 -

413 - # The inter-data cache is per model_runner

414 - inter_data_cache = self.runner.inter_data_cache

415 - if num_seqs not in inter_data_cache:

416 - inter_data_cache[num_seqs] = PyObjectCache(

417 - self.gen_inter_data_builder(num_seqs))

418 -

419 - obj = inter_data_cache[num_seqs].get_object()

420 - obj.__init__(*args, **kwargs)

421 - return obj

422 -

423 - def reset_cached_inter_data(self):

424 - for cache in self.runner.inter_data_cache.values():

425 - cache.reset()

426 -

427 - def __init__(self,

428 - runner: "GPUModelRunnerBase",

429 - finished_requests_ids: Optional[List[str]] = None):

430 - super().__init__()

431 - # Compute functions for each sequence in a sequence group.

432 - # WARNING: The order of the functions matters!

433 - self.per_seq_compute_fns = [

434 - self._compute_lens,

435 - self._compute_for_prefix_cache_hit,

436 - self._compute_for_sliding_window,

437 - self._compute_lora_input,

438 - ]

439 - # Compute functions for each sequence group.

440 - # WARNING: The order of the functions matters!

441 - self.per_seq_group_compute_fns = [

442 - self._compute_multi_modal_input,

443 - ]

444 -

445 - self.runner = runner

446 - self.model_input_cls = self.runner._model_input_cls

447 - self.attn_backend = self.runner.attn_backend

448 - self.scheduler_config = self.runner.scheduler_config

449 - self.sliding_window = self.runner.sliding_window

450 - self.block_size = self.runner.block_size

451 - self.enable_lora = self.runner.lora_config is not None

452 -

453 - # Attention metadata inputs.

454 - if self.attn_backend is not None:

455 - # spec decode (e.g. Medusa) does not have atten backend

456 - self.attn_metadata_builder = self.attn_backend.get_builder_cls()(

457 - weakref.proxy(self))

458 -

459 - # Engine/Model configurations.

460 - self.chunked_prefill_enabled = (

461 - self.scheduler_config is not None

462 - and self.scheduler_config.chunked_prefill_enabled)

463 - if self.sliding_window is not None:

464 - self.sliding_window_blocks = (

465 - self.sliding_window + self.block_size - 1) // self.block_size

466 - self.block_aligned_sliding_window = \

467 - self.sliding_window_blocks * self.block_size

468 -

469 - def prepare(self,

470 - finished_requests_ids: Optional[List[str]] = None) -> None:

471 - self.finished_requests_ids = finished_requests_ids

472 -

473 - # if the current batch is decode-only.

474 - # will be set to False if there is any non-decode request.

475 - self.decode_only = True

476 -

477 - # Intermediate data (data in CPU before going to GPU) for

478 - # the current sequence group.

479 - self.inter_data_list: List[

480 - ModelInputForGPUBuilder.InterDataForSeqGroup] = []

481 -

482 - self.attn_metadata_builder.prepare()

483 -

484 - def _compute_lens(self, inter_data: InterDataForSeqGroup, seq_idx: int,

485 - seq_group_metadata: SequenceGroupMetadata):

486 - """Compute context length, sequence length and tokens

487 - for the given sequence data.

488 - """

489 - seq_data = seq_group_metadata.seq_data[inter_data.seq_ids[seq_idx]]

490 - token_chunk_size = seq_group_metadata.token_chunk_size

491 -

492 - # Compute context length (the number of tokens that are

493 - # already computed) and sequence length (total number of tokens).

494 -

495 - seq_len = seq_data.get_len()

496 - if inter_data.is_prompt:

497 - context_len = seq_data.get_num_computed_tokens()

498 - seq_len = min(seq_len, context_len + token_chunk_size)

499 - elif self.runner.model_config.is_encoder_decoder:

500 - context_len = seq_len - 1

501 - else:

502 - context_len = seq_data.get_num_computed_tokens()

503 -

504 - # Compute tokens.

505 - if seq_data.prompt_embeds is None:

506 - tokens = seq_data.get_token_ids()[context_len:seq_len]

507 - prompt_embeds = None

508 - else:

509 - tokens = [0] * (seq_len - context_len)

510 - prompt_embeds = seq_data.get_token_embeddings(

511 - )[context_len:seq_len]

512 -

513 - inter_data.seq_lens[seq_idx] = seq_len

514 - inter_data.orig_seq_lens[seq_idx] = seq_len

515 - inter_data.prompt_lens[seq_idx] = seq_data.get_prompt_len()

516 - inter_data.context_lens[seq_idx] = context_len

517 - inter_data.input_tokens[seq_idx].extend(tokens)

518 - inter_data.inputs_embeds = prompt_embeds

519 - inter_data.input_positions[seq_idx].extend(range(context_len, seq_len))

520 - inter_data.query_lens[seq_idx] = seq_len - context_len

521 -

522 - if seq_data.mrope_position_delta is not None:

523 - if inter_data.mrope_input_positions is None:

524 - inter_data.mrope_input_positions = [None] * inter_data.n_seqs

525 -

526 - inter_data.mrope_input_positions[

527 - seq_idx] = MRotaryEmbedding.get_next_input_positions(

528 - seq_data.mrope_position_delta,

529 - context_len,

530 - seq_len,

531 - )

532 -

533 - def _compute_for_prefix_cache_hit(

534 - self, inter_data: InterDataForSeqGroup, seq_idx: int,

535 - seq_group_metadata: SequenceGroupMetadata):

536 - """Check if hit prefix cache (i.e., some blocks are already computed).

537 - If hit, update input tokens and positions to only compute the

538 - remaining blocks.

539 - """

540 - computed_block_nums = inter_data.computed_block_nums

541 -

542 - # Note that prefix caching does not support sliding window.

543 - prefix_cache_hit = (computed_block_nums is not None

544 - and len(computed_block_nums) > 0

545 - and self.sliding_window is None

546 - and inter_data.is_prompt)

547 - inter_data.prefix_cache_hit = prefix_cache_hit

548 -

549 - if not prefix_cache_hit:

550 - return

551 -

552 - assert computed_block_nums is not None

553 - # The cache hit prompt tokens in this sequence. Note that

554 - # this may be larger than the sequence length if chunked

555 - # prefill is enabled.

556 - prefix_cache_len = len(computed_block_nums) * self.block_size

557 - seq_group_metadata.seq_data[inter_data.seq_ids[

558 - seq_idx]].update_num_cached_tokens(prefix_cache_len)

559 -

560 - # The number of so far computed prompt tokens in this sequence.

561 - context_len = inter_data.context_lens[seq_idx]

562 - # The total number of prompt tokens in this sequence.

563 - # When chunked prefill is enabled, this is the token number of

564 - # computed chunks + current chunk.

565 - seq_len = inter_data.seq_lens[seq_idx]

566 - if prefix_cache_len <= context_len:

567 - # We already passed the cache hit region,

568 - # so do normal computation.

569 - pass

570 - elif context_len < prefix_cache_len < seq_len:

571 - # Partial hit. Compute the missing part.

572 - uncomputed_start = prefix_cache_len - context_len

573 - inter_data.input_tokens[seq_idx] = inter_data.input_tokens[

574 - seq_idx][uncomputed_start:]

575 - inter_data.input_positions[seq_idx] = inter_data.input_positions[

576 - seq_idx][uncomputed_start:]

577 - context_len = prefix_cache_len

578 -

579 - inter_data.context_lens[seq_idx] = context_len

580 - inter_data.query_lens[

581 - seq_idx] = inter_data.seq_lens[seq_idx] - context_len

582 - elif seq_len <= prefix_cache_len:

583 - # Full hit. Only compute the last token to avoid

584 - # erroneous behavior. FIXME: Ideally we should directly

585 - # mark all tokens as computed in the scheduler and do not

586 - # schedule this sequence, so this case should not happen.

587 - inter_data.input_tokens[seq_idx] = inter_data.input_tokens[

588 - seq_idx][-1:]

589 - inter_data.input_positions[seq_idx] = inter_data.input_positions[

590 - seq_idx][-1:]

591 - inter_data.query_lens[seq_idx] = 1

592 - inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1

593 -

594 - def _compute_for_sliding_window(self, inter_data: InterDataForSeqGroup,

595 - seq_idx: int,

596 - seq_group_metadata: SequenceGroupMetadata):

597 - """Update seq_len and curr_sliding_window_block for the given

598 - sequence data (only required by decoding) if sliding window is enabled.

599 - """

600 - curr_sliding_window_block = 0

601 - sliding_seq_len = inter_data.seq_lens[seq_idx]

602 - if not inter_data.is_prompt and self.sliding_window is not None:

603 - # TODO(sang): This is a hack to make sliding window work with

604 - # paged attn. We can remove it if we make paged attn kernel

605 - # to properly handle slinding window attn.

606 - curr_sliding_window_block = self.sliding_window_blocks

607 - # number of elements in last block

608 - suff_len = inter_data.seq_lens[seq_idx] % self.block_size

609 - sliding_seq_len = min(inter_data.seq_lens[seq_idx],

610 - self.block_aligned_sliding_window + suff_len)

611 - if suff_len > 0:

612 - curr_sliding_window_block += 1

613 -

614 - inter_data.curr_sliding_window_blocks[

615 - seq_idx] = curr_sliding_window_block

616 - inter_data.seq_lens[seq_idx] = sliding_seq_len

617 -

618 - def _compute_lora_input(self, inter_data: InterDataForSeqGroup,

619 - seq_idx: int,

620 - seq_group_metadata: SequenceGroupMetadata):

621 - """If LoRA is enabled, compute LoRA index and prompt mapping."""

622 - if not self.enable_lora:

623 - return

624 -

625 - lora_id = seq_group_metadata.lora_int_id

626 - if lora_id > 0:

627 - inter_data.lora_requests.add(seq_group_metadata.lora_request)

628 - query_len = inter_data.query_lens[seq_idx]

629 - inter_data.lora_index_mapping.append([lora_id] * query_len)

630 - sampling_params = seq_group_metadata.sampling_params

631 - if sampling_params and sampling_params.prompt_logprobs is not None:

632 - inter_data.lora_prompt_mapping.append([lora_id] * query_len)

633 - elif not self.chunked_prefill_enabled or seq_group_metadata.do_sample:

634 - inter_data.lora_prompt_mapping.append([lora_id])

635 - else:

636 - inter_data.lora_prompt_mapping.append([])

637 -

638 - def _compute_multi_modal_input(self, inter_data: InterDataForSeqGroup,

639 - seq_group_metadata: SequenceGroupMetadata):

640 - """If multi-modal data is given, add it to the input."""

641 - # NOTE: mm_kwargs only includes the subset of multi-modal items that

642 - # intersect with the current prefill positions.

643 - positions = inter_data.input_positions[0]

644 - mm_kwargs, placeholder_maps = MultiModalPlaceholderMap.from_seq_group(

645 - seq_group_metadata,

646 - range(positions[0], positions[0] + len(positions)))

647 -

648 - # M-RoPE requires mrope_positions even for plain text; return early

649 - # when mm_kwargs is empty only if inter_data.is_prompt is False.

650 - if not mm_kwargs and not inter_data.is_prompt:

651 - return

652 -

653 - inter_data.multi_modal_kwargs = mm_kwargs

654 - inter_data.multi_modal_placeholder_maps = placeholder_maps

655 -

656 - # special processing for mrope position deltas.

657 - if self.runner.model_config.uses_mrope:

658 - image_grid_thw = mm_kwargs.get("image_grid_thw", None)

659 - video_grid_thw = mm_kwargs.get("video_grid_thw", None)

660 - audio_feature_lengths = mm_kwargs.get("audio_feature_lengths",

661 - None)

662 -

663 - second_per_grid_ts = mm_kwargs.get("second_per_grid_ts", None)

664 - use_audio_in_video = mm_kwargs.get("use_audio_in_video", False)

665 - hf_config = self.runner.model_config.hf_config

666 -

667 - inter_data.mrope_input_positions = [None] * inter_data.n_seqs

668 - for seq_idx in range(inter_data.n_seqs):

669 - seq_data = seq_group_metadata.seq_data[

670 - inter_data.seq_ids[seq_idx]]

671 - token_ids = seq_data.get_token_ids()

672 -

673 - if supports_mrope(self.runner.model):

674 - mrope_input_positions, mrope_position_delta = \

675 - self.runner.model.get_mrope_input_positions(

676 - token_ids,

677 - hf_config=hf_config,

678 - image_grid_thw=image_grid_thw,

679 - video_grid_thw=video_grid_thw,

680 - second_per_grid_ts=second_per_grid_ts,

681 - context_len=inter_data.context_lens[seq_idx],

682 - seq_len=inter_data.seq_lens[seq_idx],

683 - audio_feature_lengths=audio_feature_lengths,

684 - use_audio_in_video=use_audio_in_video,

685 - )

686 - mrope_input_positions = mrope_input_positions.tolist()

687 - else:

688 - mrope_input_positions, mrope_position_delta = \

689 - MRotaryEmbedding.get_input_positions(

690 - token_ids,

691 - hf_config=hf_config,

692 - image_grid_thw=image_grid_thw,

693 - video_grid_thw=video_grid_thw,

694 - second_per_grid_ts=second_per_grid_ts,

695 - context_len=inter_data.context_lens[seq_idx],

696 - seq_len=inter_data.seq_lens[seq_idx],

697 - audio_feature_lengths=audio_feature_lengths,

698 - use_audio_in_video=use_audio_in_video,

699 - )

700 -

701 - seq_data.mrope_position_delta = mrope_position_delta

702 - inter_data.mrope_input_positions[

703 - seq_idx] = mrope_input_positions

704 -

705 - def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):

706 - """Add a sequence group to the builder."""

707 - seq_ids = seq_group_metadata.seq_data.keys()

708 - n_seqs = len(seq_ids)

709 - is_prompt = seq_group_metadata.is_prompt

710 -

711 - if is_prompt:

712 - assert n_seqs == 1

713 - self.decode_only = False

714 -

715 - encoder_seq_len = 0

716 -

717 - if self.runner.model_config.is_encoder_decoder:

718 - encoder_seq_len = seq_group_metadata.encoder_seq_data.get_len()

719 -

720 - inter_data = self.init_cached_inter_data(

721 - request_id=seq_group_metadata.request_id,

722 - seq_ids=seq_ids,

723 - is_prompt=is_prompt,

724 - block_tables=seq_group_metadata.block_tables,

725 - computed_block_nums=seq_group_metadata.computed_block_nums,

726 - reinit=True,

727 - reinit_use_defaults=True,

728 - encoder_seq_len=encoder_seq_len)

729 -

730 - self.inter_data_list.append(inter_data)

731 -

732 - for seq_idx in range(n_seqs):

733 - for per_seq_fn in self.per_seq_compute_fns:

734 - per_seq_fn(inter_data, seq_idx, seq_group_metadata)

735 - for per_seq_group_fn in self.per_seq_group_compute_fns:

736 - per_seq_group_fn(inter_data, seq_group_metadata)

737 -

738 - def _use_captured_graph(self,

739 - batch_size: int,

740 - decode_only: bool,

741 - max_decode_seq_len: int,

742 - max_encoder_seq_len: int = 0) -> bool:

743 - return (decode_only and not self.runner.model_config.enforce_eager

744 - and max_decode_seq_len <= self.runner.max_seq_len_to_capture

745 - and max_encoder_seq_len <= self.runner.max_seq_len_to_capture

746 - and batch_size <= self.runner.max_batchsize_to_capture)

747 -

748 - def _get_cuda_graph_pad_size(self,

749 - num_seqs: int,

750 - max_decode_seq_len: int,

751 - max_encoder_seq_len: int = 0) -> int:

752 - """

753 - Determine the number of padding sequences required for running in

754 - CUDA graph mode. Returns -1 if CUDA graphs cannot be used.

755 -

756 - In the multi-step + chunked-prefill case, only the first step

757 - has Prefills (if any). The rest of the steps are guaranteed to be all

758 - decodes. In this case, we set up the padding as if all the sequences

759 - are decodes so we may run all steps except the first step in CUDA graph

760 - mode.

761 -

762 - Args:

763 - num_seqs (int): Number of sequences scheduled to run.

764 - max_decode_seq_len (int): Greatest of all the decode sequence

765 - lengths. Used only in checking the viablility of using

766 - CUDA graphs.

767 - max_encoder_seq_len (int, optional): Greatest of all the encode

768 - sequence lengths. Defaults to 0. Used only in checking the

769 - viability of using CUDA graphs.

770 - Returns:

771 - int: Returns the determined number of padding sequences. If

772 - CUDA graphs is not viable, returns -1.

773 - """

774 - decode_only = self.decode_only

775 - if not decode_only:

776 - # Early exit so we can treat num_seqs as the batch_size below.

777 - return -1

778 -

779 - # batch_size out of this function refers to the number of input

780 - # tokens being scheduled. This conflation of num_seqs as batch_size

781 - # is valid as this is a decode-only case.

782 - batch_size = num_seqs

783 - if not self._use_captured_graph(batch_size, decode_only,

784 - max_decode_seq_len,

785 - max_encoder_seq_len):

786 - return -1

787 -

788 - graph_batch_size = self.runner.vllm_config.pad_for_cudagraph(

789 - batch_size)

790 - assert graph_batch_size >= batch_size

791 - return graph_batch_size - batch_size

792 -

793 - def build(self) -> ModelInputForGPU:

794 - """Finalize the builder intermediate data and

795 - create on-device tensors.

796 - """

797 - # Combine and flatten intermediate data.

798 - input_tokens = list[int]()

799 - inputs_embeds_list = list[torch.Tensor]()

800 - for inter_data in self.inter_data_list:

801 - for cur_input_tokens in inter_data.input_tokens:

802 - input_tokens.extend(cur_input_tokens)

803 - if inter_data.inputs_embeds is not None:

804 - inputs_embeds_list.append(

805 - inter_data.inputs_embeds.to(

806 - dtype=self.runner.model_config.dtype,

807 - device=self.runner.device))

808 - inputs_embeds: Optional[torch.Tensor]

809 - if len(inputs_embeds_list) == 0:

810 - inputs_embeds = None

811 - else:

812 - inputs_embeds = torch.cat(inputs_embeds_list, dim=0).to(

813 - dtype=self.runner.model_config.dtype,

814 - device=self.runner.device)

815 - assert len(inputs_embeds) == len(input_tokens)

816 -

817 - if not input_tokens and inputs_embeds is None:

818 - # This may happen when all prefill requests hit

819 - # prefix caching and there is no decode request.

820 - return self.model_input_cls()

821 -

822 - mrope_input_positions: Optional[List[List[int]]] = None

823 - if any(inter_data.mrope_input_positions is not None

824 - for inter_data in self.inter_data_list):

825 - mrope_input_positions = [[] for _ in range(3)]

826 - for idx in range(3):

827 - for inter_data in self.inter_data_list:

828 - msections = inter_data.mrope_input_positions

829 - if msections is None:

830 - for _seq_input_positions in inter_data.input_positions:

831 - mrope_input_positions[idx].extend(

832 - _seq_input_positions)

833 - else:

834 - for _seq_mrope_input_positions in msections:

835 - mrope_input_positions[idx].extend(

836 - _seq_mrope_input_positions[idx])

837 - input_positions = None

838 - else:

839 - input_positions = []

840 - for inter_data in self.inter_data_list:

841 - for cur_input_positions in inter_data.input_positions:

842 - input_positions.extend(cur_input_positions)

843 -

844 - seq_lens = []

845 - query_lens = []

846 - max_decode_seq_len = 0

847 - max_encoder_seq_len = 0

848 - for inter_data in self.inter_data_list:

849 - seq_lens.extend(inter_data.seq_lens)

850 - query_lens.extend(inter_data.query_lens)

851 - if not inter_data.is_prompt:

852 - max_decode_seq_len = max(max_decode_seq_len,

853 - max(inter_data.seq_lens))

854 - if self.runner.model_config.is_encoder_decoder:

855 - max_encoder_seq_len = max(max_encoder_seq_len,

856 - inter_data.encoder_seq_len)

857 -

858 - # Mapping from request IDs to sequence IDs. Used for Jamba models

859 - # that manages the cache by itself.

860 - request_ids_to_seq_ids = {

861 - data.request_id: data.seq_ids

862 - for data in self.inter_data_list

863 - }

864 -

865 - cuda_graph_pad_size = self._get_cuda_graph_pad_size(

866 - num_seqs=len(seq_lens),

867 - max_decode_seq_len=max_decode_seq_len,

868 - max_encoder_seq_len=max_encoder_seq_len)

869 -

870 - batch_size = len(input_tokens)

871 - if cuda_graph_pad_size != -1:

872 - # If cuda graph can be used, pad tensors accordingly.

873 - # See `capture_model` API for more details.

874 - # vLLM uses cuda graph only for decoding requests.

875 - batch_size += cuda_graph_pad_size

876 -

877 - # Tokens and positions.

878 - if cuda_graph_pad_size:

879 - input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size))

880 - assert self.runner.device is not None

881 - input_tokens_tensor = async_tensor_h2d(input_tokens, torch.long,

882 - self.runner.device,

883 - self.runner.pin_memory)

884 -

885 - if mrope_input_positions is not None:

886 - for idx in range(3):

887 - mrope_input_positions[idx].extend(

888 - itertools.repeat(0, cuda_graph_pad_size))

889 - input_positions_tensor = async_tensor_h2d(mrope_input_positions,

890 - torch.long,

891 - self.runner.device,

892 - self.runner.pin_memory)

893 - else:

894 - input_positions.extend(itertools.repeat(0, cuda_graph_pad_size))

895 - input_positions_tensor = async_tensor_h2d(input_positions,

896 - torch.long,

897 - self.runner.device,

898 - self.runner.pin_memory)

899 - # Sequence and query lengths.

900 - if cuda_graph_pad_size:

901 - seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size))

902 -

903 - # Attention metadata.

904 - attn_metadata = self.attn_metadata_builder.build(

905 - seq_lens, query_lens, cuda_graph_pad_size, batch_size)

906 -

907 - # LoRA data.

908 - lora_requests = set()

909 - lora_mapping = None

910 - if self.enable_lora:

911 - lora_requests = set(r for data in self.inter_data_list

912 - for r in data.lora_requests)

913 - lora_index_mapping = flatten_2d_lists([

914 - flatten_2d_lists(inter_data.lora_index_mapping)

915 - for inter_data in self.inter_data_list

916 - ])

917 - if cuda_graph_pad_size:

918 - lora_index_mapping.extend(

919 - itertools.repeat(0, cuda_graph_pad_size))

920 - lora_prompt_mapping = flatten_2d_lists([

921 - flatten_2d_lists(inter_data.lora_prompt_mapping)

922 - for inter_data in self.inter_data_list

923 - ])

924 -

925 - lora_mapping = LoRAMapping(

926 - **dict(index_mapping=lora_index_mapping,

927 - prompt_mapping=lora_prompt_mapping,

928 - is_prefill=not self.decode_only))

929 -

930 - # Multi-modal data.

931 - multi_modal_kwargs_list = [

932 - data.multi_modal_kwargs for data in self.inter_data_list

933 - if data.multi_modal_kwargs is not None

934 - ]

935 - multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)

936 -

937 - return self.model_input_cls(

938 - input_tokens=input_tokens_tensor,

939 - inputs_embeds=inputs_embeds,

940 - input_positions=input_positions_tensor,

941 - attn_metadata=attn_metadata,

942 - seq_lens=seq_lens,

943 - query_lens=query_lens,

944 - lora_mapping=lora_mapping,

945 - lora_requests=lora_requests,

946 - multi_modal_kwargs=multi_modal_kwargs,

947 - request_ids_to_seq_ids=request_ids_to_seq_ids,

948 - finished_requests_ids=self.finished_requests_ids)

949 -

950 -

951 -class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):

952 - """

953 - Helper class for shared methods between GPU model runners.

954 - """

955 - _model_input_cls: Type[TModelInputForGPU]

956 - _builder_cls: Type[ModelInputForGPUBuilder]

957 - builder: ModelInputForGPUBuilder

958 -

959 - def __init__(

960 - self,

961 - vllm_config: VllmConfig,

962 - kv_cache_dtype: Optional[str] = "auto",

963 - is_driver_worker: bool = False,

964 - return_hidden_states: bool = False,

965 - input_registry: InputRegistry = INPUT_REGISTRY,

966 - mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,

967 - ):

968 -

969 - ModelRunnerBase.__init__(self, vllm_config)

970 - model_config = self.model_config

971 - cache_config = self.cache_config

972 -

973 - self.is_driver_worker = is_driver_worker

974 - self.return_hidden_states = return_hidden_states

975 -

976 - self.device = self.device_config.device

977 - self.pin_memory = is_pin_memory_available()

978 -

979 - self.kv_cache_dtype = kv_cache_dtype

980 - self.sliding_window = model_config.get_sliding_window()

981 - self.block_size = cache_config.block_size

982 - self.max_seq_len_to_capture = self.model_config.max_seq_len_to_capture

983 - self.max_batchsize_to_capture = \

984 - self.vllm_config.compilation_config.max_capture_size

985 -

986 - #

987 - self.graph_runners: List[Dict[Tuple[int, bool], CUDAGraphRunner]] = [

988 - {} for _ in range(self.parallel_config.pipeline_parallel_size)

989 - ]

990 - self.graph_memory_pool: Optional[Tuple[

991 - int, int]] = None # Set during graph capture.

992 -

993 - self.has_inner_state = model_config.has_inner_state

994 -

995 - self.in_profile_run = False

996 -

997 - # When using CUDA graph, the input block tables must be padded to

998 - # max_seq_len_to_capture. However, creating the block table in

999 - # Python can be expensive. To optimize this, we cache the block table

1000 - # in numpy and only copy the actual input content at every iteration.

1001 - # The shape of the cached block table will be

1002 - # (max batch size to capture, max seq len to capture / block size).

1003 - self.graph_block_tables = np.zeros(

1004 - (self.max_batchsize_to_capture, self.get_max_block_per_batch()),

1005 - dtype=np.int32)

1006 -

1007 - self.cross_layer_shared_graph_block_tables = np.zeros(

1008 - (self.max_batchsize_to_capture, self.get_max_block_per_batch()),

1009 - dtype=np.int32)

1010 -

1011 - # Attention-free but stateful models like Mamba need a placeholder attn

1012 - # backend, as the attention metadata is needed to manage internal state.

1013 - # However we must bypass attention selection altogether for some models

1014 - # used for speculative decoding to avoid a divide-by-zero in

1015 - # model_config.get_head_size()

1016 - num_attn_heads = self.model_config.get_num_attention_heads(

1017 - self.parallel_config)

1018 - needs_attn_backend = (num_attn_heads != 0

1019 - or self.model_config.is_attention_free)

1020 -

1021 - self.attn_backend = get_attn_backend(

1022 - self.model_config.get_head_size(),

1023 - self.model_config.dtype,

1024 - self.kv_cache_dtype,

1025 - self.block_size,

1026 - self.model_config.is_attention_free,

1027 - use_mla=self.model_config.use_mla,

1028 - ) if needs_attn_backend else None

1029 - if self.attn_backend:

1030 - self.attn_state = self.attn_backend.get_state_cls()(

1031 - weakref.proxy(self))

1032 - else:

1033 - self.attn_state = CommonAttentionState(weakref.proxy(self))

1034 -

1035 - # Multi-modal data support

1036 - self.input_registry = input_registry

1037 - self.mm_registry = mm_registry

1038 -

1039 - # Lazy initialization

1040 - self.model: nn.Module # Set after load_model

1041 - # Set after load_model.

1042 - self.lora_manager: Optional[LRUCacheWorkerLoRAManager] = None

1043 - self.sampler = get_sampler()

1044 -

1045 - set_cpu_offload_max_bytes(

1046 - int(self.cache_config.cpu_offload_gb * 1024**3))

1047 -

1048 - # Used to cache python objects

1049 - self.inter_data_cache: Dict[int, PyObjectCache] = {}

1050 -

1051 - # Using the PythonizationCache in Pipeline-Parallel clobbers the

1052 - # SequenceGroupToSample object. In Pipeline-Parallel, we have

1053 - # more than 1 Scheduler, resulting in a potential back-to-back

1054 - # prepare_model_inputs() call. This clobbers the cached

1055 - # SequenceGroupToSample objects, as we reset the cache during

1056 - # every prepare_model_inputs() call.

1057 - self.sampling_metadata_cache: SamplingMetadataCache = \

1058 - SamplingMetadataCache() \

1059 - if self.parallel_config.pipeline_parallel_size == 1 else None

1060 -

1061 - if hasattr(self, "_builder_cls"):

1062 - # multi-step model runner does not have `_builder_cls`

1063 - self.builder = self._builder_cls(weakref.proxy(self))

1064 -

1065 - def load_model(self) -> None:

1066 - logger.info("Starting to load model %s...", self.model_config.model)

1067 - with DeviceMemoryProfiler(self.device) as m:

1068 - time_before_load = time.perf_counter()

1069 - self.model = get_model(vllm_config=self.vllm_config)

1070 - if self.lora_config:

1071 - assert supports_lora(

1072 - self.model

1073 - ), f"{self.model.__class__.__name__} does not support LoRA yet."

1074 -

1075 - if supports_multimodal(self.model):

1076 - logger.warning(

1077 - "Regarding multimodal models, vLLM currently "

1078 - "only supports adding LoRA to language model.")

1079 -

1080 - self.lora_manager = LRUCacheWorkerLoRAManager(

1081 - self.vllm_config,

1082 - self.device,

1083 - self.model.embedding_modules,

1084 - self.model.embedding_padding_modules,

1085 - )

1086 -

1087 - self.model = self.lora_manager.create_lora_manager(self.model)

1088 - time_after_load = time.perf_counter()

1089 -

1090 - self.model_memory_usage = m.consumed_memory

1091 - logger.info("Model loading took %.4f GiB and %.6f seconds",

1092 - self.model_memory_usage / GiB_bytes,

1093 - time_after_load - time_before_load)

1094 -

1095 -

1096 - if self.vllm_config.compilation_config.level ==\

1097 - CompilationLevel.DYNAMO_AS_IS and supports_dynamo():

1098 - backend = self.vllm_config.compilation_config.init_backend(

1099 - self.vllm_config)

1100 - compilation_counter.dynamo_as_is_count += 1

1101 - self.model = torch.compile(self.model,

1102 - fullgraph=True,

1103 - backend=backend)

1104 -

1105 - def get_model(self) -> nn.Module:

1106 - return self.model

1107 -

1108 - def save_sharded_state(

1109 - self,

1110 - path: str,

1111 - pattern: Optional[str] = None,

1112 - max_size: Optional[int] = None,

1113 - ) -> None:

1114 - from vllm.model_executor.model_loader import ShardedStateLoader

1115 - ShardedStateLoader.save_model(

1116 - self.model,

1117 - path,

1118 - pattern=pattern,

1119 - max_size=max_size,

1120 - )

1121 -

1122 - def save_tensorized_model(

1123 - self,

1124 - tensorizer_config: TensorizerConfig,

1125 - ) -> None:

1126 - from vllm.model_executor.model_loader import TensorizerLoader

1127 - TensorizerLoader.save_model(

1128 - self.model,

1129 - tensorizer_config=tensorizer_config,

1130 - model_config=self.model_config,

1131 - )

1132 -

1133 - def get_max_block_per_batch(self) -> int:

1134 - block_size = self.block_size

1135 - return (self.max_seq_len_to_capture + block_size - 1) // block_size

1136 -

1137 - def _prepare_model_input_tensors(

1138 - self,

1139 - seq_group_metadata_list: List[SequenceGroupMetadata],

1140 - finished_requests_ids: Optional[List[str]] = None

1141 - ) -> TModelInputForGPU:

1142 - """Helper method to prepare the model input based on a given sequence

1143 - group. Prepares metadata needed for the base model forward pass but not

1144 - metadata for possible additional steps, e.g., sampling.

1145 -

1146 - The API assumes seq_group_metadata_list is sorted by prefill -> decode.

1147 -

1148 - The result tensors and data structure also batches input in prefill

1149 - -> decode order. For example,

1150 -

1151 - - input_tokens[:num_prefill_tokens] contains prefill tokens.

1152 - - input_tokens[num_prefill_tokens:] contains decode tokens.

1153 -

1154 - If cuda graph is required, this API automatically pads inputs.

1155 - """

1156 - self.builder.prepare(finished_requests_ids)

1157 - for seq_group_metadata in seq_group_metadata_list:

1158 - try:

1159 - self.builder.add_seq_group(seq_group_metadata)

1160 - except Exception as e:

1161 - # Raise an exception that tracks the ID of the bad request

1162 - raise InputProcessingError(seq_group_metadata.request_id,

1163 - str(e)) from e

1164 -

1165 - self.builder.reset_cached_inter_data()

1166 -

1167 - return self.builder.build() # type: ignore

1168 -

1169 - @contextmanager

1170 - def set_in_profile_run(self):

1171 - self.in_profile_run = True

1172 - try:

1173 - yield

1174 - finally:

1175 - self.in_profile_run = False

1176 -

1177 - @torch.inference_mode()

1178 - def profile_run(self) -> None:

1179 - max_num_batched_tokens = \

1180 - self.scheduler_config.max_num_batched_tokens

1181 - max_num_seqs = self.scheduler_config.max_num_seqs

1182 - self._dummy_run(max_num_batched_tokens, max_num_seqs)

1183 -

1184 - def _add_dummy_loras(self, num_loras: int) -> list[LoRARequest]:

1185 - assert num_loras > 0

1186 - assert self.lora_manager is not None

1187 -

1188 - dummy_lora_requests: list[LoRARequest] = []

1189 - with self.lora_manager.dummy_lora_cache():

1190 - for idx in range(num_loras):

1191 - lora_id = idx + 1

1192 - dummy_lora_request = LoRARequest(

1193 - lora_name=f"warmup_{lora_id}",

1194 - lora_int_id=lora_id,

1195 - lora_path="/not/a/real/path",

1196 - )

1197 - self.lora_manager.add_dummy_lora(dummy_lora_request,

1198 - rank=LORA_WARMUP_RANK)

1199 - dummy_lora_requests.append(dummy_lora_request)

1200 - return dummy_lora_requests

1201 -

1202 - def _remove_dummy_loras(self):

1203 - # Remove dummy loras.

1204 - assert self.lora_manager is not None

1205 - self.remove_all_loras()

1206 -

1207 - def _dummy_run(self,

1208 - max_num_batched_tokens: int,

1209 - max_num_seqs: int = 1) -> None:

1210 - with self.set_in_profile_run():

1211 - # Enable top-k sampling to reflect the accurate memory usage.

1212 - sampling_params = \

1213 - SamplingParams(top_p=0.99, top_k=self.vocab_size - 1)

1214 -

1215 - # This represents the maximum number of different requests

1216 - # that will have unique loras, and therefore the max amount of

1217 - # memory consumption. Create dummy lora request copies from the

1218 - # lora request passed in, which contains a lora from the lora

1219 - # warmup path.

1220 - dummy_lora_requests: List[LoRARequest] = []

1221 - dummy_lora_requests_per_seq: List[LoRARequest] = []

1222 - if self.lora_config:

1223 - dummy_lora_requests = self._add_dummy_loras(

1224 - self.lora_config.max_loras)

1225 - assert len(dummy_lora_requests) == self.lora_config.max_loras

1226 - dummy_lora_requests_per_seq = [

1227 - dummy_lora_requests[idx % len(dummy_lora_requests)]

1228 - for idx in range(max_num_seqs)

1229 - ]

1230 -

1231 - # Profile memory usage with max_num_sequences sequences and the

1232 - # total number of tokens equal to max_num_batched_tokens.

1233 - seqs: List[SequenceGroupMetadata] = []

1234 - # Additional GPU memory may be needed for multi-modal encoding,

1235 - # which needs to be accounted for when calculating the GPU blocks

1236 - # for vLLM blocker manager.

1237 - # To exercise the worst scenario for GPU memory consumption,

1238 - # the number of seqs (batch_size) is chosen to maximize the number

1239 - # of images processed.

1240 -

1241 - max_mm_tokens = self.mm_registry.get_max_multimodal_tokens(

1242 - self.model_config)

1243 - if max_mm_tokens > 0:

1244 - max_num_seqs_orig = max_num_seqs

1245 - max_num_seqs = min(max_num_seqs,

1246 - max_num_batched_tokens // max_mm_tokens)

1247 - if max_num_seqs < 1:

1248 - expr = (f"min({max_num_seqs_orig}, "

1249 - f"{max_num_batched_tokens} // {max_mm_tokens})")

1250 - logger.warning(

1251 - "Computed max_num_seqs (%s) to be less than 1. "

1252 - "Setting it to the minimum value of 1.", expr)

1253 - max_num_seqs = 1

1254 -

1255 - batch_size = 0

1256 - for group_id in range(max_num_seqs):

1257 - seq_len = (max_num_batched_tokens // max_num_seqs +

1258 - (group_id < max_num_batched_tokens % max_num_seqs))

1259 - batch_size += seq_len

1260 -

1261 - dummy_data = self.input_registry \

1262 - .dummy_data_for_profiling(self.model_config,

1263 - seq_len,

1264 - self.mm_registry)

1265 -

1266 - seq = SequenceGroupMetadata(

1267 - request_id=str(group_id),

1268 - is_prompt=True,

1269 - seq_data={group_id: dummy_data.seq_data},

1270 - sampling_params=sampling_params,

1271 - block_tables=None,

1272 - lora_request=dummy_lora_requests_per_seq[group_id]

1273 - if dummy_lora_requests_per_seq else None,

1274 - multi_modal_data=dummy_data.multi_modal_data,

1275 - multi_modal_placeholders=dummy_data.

1276 - multi_modal_placeholders,

1277 - )

1278 - seqs.append(seq)

1279 -

1280 - # Run the model with the dummy inputs.

1281 - num_layers = self.model_config.get_num_layers(self.parallel_config)

1282 - # use an empty tensor instead of `None`` to force Dynamo to pass

1283 - # it by reference, rather by specializing on the value ``None``.

1284 - # the `dtype` argument does not matter, and we use `float32` as

1285 - # a placeholder (it has wide hardware support).

1286 - # it is important to create tensors inside the loop, rather than

1287 - # multiplying the list, to avoid Dynamo from treating them as

1288 - # tensor aliasing.

1289 - kv_caches = [

1290 - torch.tensor([], dtype=torch.float32, device=self.device)

1291 - for _ in range(num_layers)

1292 - ]

1293 - finished_requests_ids = [seq.request_id for seq in seqs]

1294 - model_input = self.prepare_model_input(

1295 - seqs, finished_requests_ids=finished_requests_ids)

1296 - intermediate_tensors = None

1297 - if not get_pp_group().is_first_rank:

1298 - intermediate_tensors = \

1299 - self.model.make_empty_intermediate_tensors(

1300 - batch_size=batch_size,

1301 - dtype=self.model_config.dtype,

1302 - device=self.device)

1303 -

1304 - # Disable KV Scale Calculation for dummy data during profile run

1305 - if model_input.attn_metadata is not None:

1306 - model_input.attn_metadata.enable_kv_scales_calculation = False

1307 -

1308 - self.execute_model(model_input, kv_caches, intermediate_tensors)

1309 - torch.cuda.synchronize()

1310 - if self.lora_config:

1311 - self._remove_dummy_loras()

1312 -

1313 - return

1314 -

1315 - def remove_all_loras(self):

1316 - if not self.lora_manager:

1317 - raise RuntimeError("LoRA is not enabled.")

1318 - self.lora_manager.remove_all_adapters()

1319 -

1320 - def set_active_loras(self, lora_requests: Set[LoRARequest],

1321 - lora_mapping: LoRAMapping) -> None:

1322 - if not self.lora_manager:

1323 - raise RuntimeError("LoRA is not enabled.")

1324 - self.lora_manager.set_active_adapters(lora_requests, lora_mapping)

1325 -

1326 - def add_lora(self, lora_request: LoRARequest) -> bool:

1327 - if not self.lora_manager:

1328 - raise RuntimeError("LoRA is not enabled.")

1329 - return self.lora_manager.add_adapter(lora_request)

1330 -

1331 - def remove_lora(self, lora_id: int) -> bool:

1332 - if not self.lora_manager:

1333 - raise RuntimeError("LoRA is not enabled.")

1334 - return self.lora_manager.remove_adapter(lora_id)

1335 -

1336 - def pin_lora(self, lora_id: int) -> bool:

1337 - if not self.lora_manager:

1338 - raise RuntimeError("LoRA is not enabled.")

1339 - return self.lora_manager.pin_adapter(lora_id)

1340 -

1341 - def list_loras(self) -> Set[int]:

1342 - if not self.lora_manager:

1343 - raise RuntimeError("LoRA is not enabled.")

1344 - return self.lora_manager.list_adapters()

1345 -

1346 - @torch.inference_mode()

1347 - def capture_model(self, kv_caches: List[List[torch.Tensor]]) -> int:

1348 - """Cuda graph capture a model and return cudagraph memory

1349 - consumption in bytes.

1350 -

1351 - Note that CUDA graph's performance gain is negligible if number

1352 - of batched tokens are larger than 200. And since CUDA graph

1353 - requires fixed sized tensors, supporting large/variable batch

1354 - size requires high GPU memory overhead. Thus, vLLM only captures

1355 - decoding requests. Mixed batch (chunked prefill + decoding) or

1356 - prefill requests are not captured.

1357 -

1358 - Since it is used for decoding-only, it assumes there's only 1 token

1359 - per sequence in the batch.

1360 - """

1361 - assert not self.model_config.enforce_eager

1362 - logger.info("Capturing cudagraphs for decoding. This may lead to "

1363 - "unexpected consequences if the model is not static. To "

1364 - "run the model in eager mode, set 'enforce_eager=True' or "

1365 - "use '--enforce-eager' in the CLI. "

1366 - "If out-of-memory error occurs during cudagraph capture,"

1367 - " consider decreasing `gpu_memory_utilization` or "

1368 - "switching to eager mode. You can also reduce the "

1369 - "`max_num_seqs` as needed to decrease memory usage.")

1370 - start_time = time.perf_counter()

1371 - start_free_gpu_memory = torch.cuda.mem_get_info()[0]

1372 -

1373 - # Prepare dummy inputs. These will be reused for all batch sizes.

1374 - max_batch_size = self.max_batchsize_to_capture

1375 - input_tokens = torch.zeros(max_batch_size,

1376 - dtype=torch.long,

1377 - device=self.device)

1378 - input_positions = torch.zeros(max_batch_size,

1379 - dtype=torch.long,

1380 - device=self.device)

1381 - inputs_embeds = torch.zeros(

1382 - (max_batch_size, self.model_config.get_hidden_size()),

1383 - dtype=self.model_config.dtype,

1384 - device=self.device)

1385 - if self.model_config.uses_mrope:

1386 - input_positions = torch.tile(input_positions,

1387 - (3, 1)).cuda(device=self.device)

1388 - # Prepare dummy previous_hidden_states only if needed by the model.

1389 - # This is used by draft models such as EAGLE.

1390 - previous_hidden_states = None

1391 - if "previous_hidden_states" in inspect.signature(

1392 - self.model.forward).parameters:

1393 - previous_hidden_states = torch.empty(

1394 - [max_batch_size,

1395 - self.model_config.get_hidden_size()],

1396 - dtype=self.model_config.dtype,

1397 - device=self.device)

1398 -

1399 - intermediate_inputs = None

1400 - if not get_pp_group().is_first_rank:

1401 - intermediate_inputs = self.model.make_empty_intermediate_tensors(

1402 - batch_size=max_batch_size,

1403 - dtype=self.model_config.dtype,

1404 - device=self.device)

1405 -

1406 - dummy_lora_id: Optional[int] = None

1407 - dummy_lora_request: LoRARequest = []

1408 - if self.lora_config:

1409 - # The goal is to capture the LoRA kernels in cuda graphs.

1410 - # for this purpose, as single dummy lora is sufficient.

1411 - dummy_lora_requests = self._add_dummy_loras(num_loras=1)

1412 - assert len(dummy_lora_requests) == 1

1413 - dummy_lora_request = dummy_lora_requests[0]

1414 - dummy_lora_id = dummy_lora_request.lora_int_id

1415 -

1416 - with self.attn_state.graph_capture(max_batch_size), graph_capture(

1417 - self.device) as graph_capture_context:

1418 - # NOTE: Capturing the largest batch size first may help reduce the

1419 - # memory usage of CUDA graph.

1420 - for virtual_engine in range(

1421 - self.parallel_config.pipeline_parallel_size):

1422 - # We need to not only iterate over batch sizes, but also whether

1423 - # to use inputs_embeds or not, hence we use the cartesian

1424 - # product.

1425 - cudagraph_capture_sizes = self.vllm_config.compilation_config\

1426 - .cudagraph_capture_sizes

1427 - cudagraph_inputs_embeds = ((

1428 - True, False) if self.model_config.enable_prompt_embeds else

1429 - (False, ))

1430 - compilation_cases = itertools.product(

1431 - cudagraph_capture_sizes,

1432 - cudagraph_inputs_embeds,

1433 - )

1434 - # Only rank 0 should print progress bar during capture

1435 - if get_tensor_model_parallel_rank() == 0:

1436 - compilation_cases = tqdm(

1437 - list(compilation_cases),

1438 - disable=not self.load_config.use_tqdm_on_load,

1439 - desc="Capturing CUDA graph shapes")

1440 - for batch_size, use_inputs_embeds in compilation_cases:

1441 - attn_metadata = (

1442 - self.attn_state.graph_capture_get_metadata_for_batch(

1443 - batch_size,

1444 - is_encoder_decoder_model=self.model_config.

1445 - is_encoder_decoder))

1446 - # Disable KV Scale Calculation for graph capture

1447 - attn_metadata.enable_kv_scales_calculation = False

1448 - if self.lora_config:

1449 - lora_mapping = LoRAMapping(

1450 - **dict(index_mapping=[dummy_lora_id] * batch_size,

1451 - prompt_mapping=[dummy_lora_id] * batch_size,

1452 - is_prefill=False))

1453 - self.set_active_loras(set([dummy_lora_request]),

1454 - lora_mapping)

1455 -

1456 - graph_runner = CUDAGraphRunner(

1457 - self.model, self.attn_backend.get_name(),

1458 - self.attn_state.graph_clone(batch_size),

1459 - self.model_config.is_encoder_decoder)

1460 -

1461 - capture_inputs = {

1462 - "input_ids":

1463 - input_tokens[:batch_size],

1464 - "inputs_embeds":

1465 - inputs_embeds[:batch_size]

1466 - if use_inputs_embeds else None,

1467 - "positions":

1468 - input_positions[..., :batch_size],

1469 - "intermediate_inputs":

1470 - intermediate_inputs[:batch_size]

1471 - if intermediate_inputs is not None else None,

1472 - "kv_caches":

1473 - kv_caches[virtual_engine],

1474 - "attn_metadata":

1475 - attn_metadata,

1476 - "memory_pool":

1477 - self.graph_memory_pool,

1478 - "stream":

1479 - graph_capture_context.stream

1480 - }

1481 - if previous_hidden_states is not None:

1482 - capture_inputs[

1483 - "previous_hidden_states"] = previous_hidden_states[:

1484 - batch_size]

1485 -

1486 - if self.has_inner_state:

1487 - # Only used by Mamba-based models CUDA graph atm (Jamba)

1488 - capture_inputs.update({

1489 - "seqlen_agnostic_capture_inputs":

1490 - self.model.get_seqlen_agnostic_capture_inputs(

1491 - batch_size)

1492 - })

1493 - if self.model_config.is_encoder_decoder:

1494 - # add the additional inputs to capture for

1495 - # encoder-decoder models.

1496 - self._update_inputs_to_capture_for_enc_dec_model(

1497 - capture_inputs)

1498 -

1499 - with set_forward_context(attn_metadata, self.vllm_config,

1500 - virtual_engine):

1501 - graph_runner.capture(**capture_inputs)

1502 - self.graph_memory_pool = graph_runner.graph.pool()

1503 - self.graph_runners[virtual_engine][(

1504 - batch_size, use_inputs_embeds)] = graph_runner

1505 -

1506 - if self.lora_config:

1507 - self._remove_dummy_loras()

1508 -

1509 - end_time = time.perf_counter()

1510 - end_free_gpu_memory = torch.cuda.mem_get_info()[0]

1511 - elapsed_time = end_time - start_time

1512 - cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory

1513 - # This usually takes < 10 seconds.

1514 - logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",

1515 - elapsed_time, cuda_graph_size / GiB_bytes)

1516 - return cuda_graph_size

1517 -

1518 - def _update_inputs_to_capture_for_enc_dec_model(self,

1519 - capture_inputs: Dict[str,

1520 - Any]):

1521 - """

1522 - Updates the set of input tensors needed for CUDA graph capture in an

1523 - encoder-decoder model.

1524 -

1525 - This method modifies the provided `capture_inputs` dictionary by

1526 - adding tensors specific to encoder-decoder specific models that

1527 - need to be captured for CUDA Graph replay.

1528 - """

1529 - # During the decode phase encoder_input_ids and encoder_positions are

1530 - # unset. Do the same thing for graph capture.

1531 - capture_inputs["encoder_input_ids"] = torch.tensor([],

1532 - dtype=torch.long,

1533 - device=self.device)

1534 - capture_inputs["encoder_positions"] = torch.tensor([],

1535 - dtype=torch.long,

1536 - device=self.device)

1537 -

1538 - @property

1539 - def vocab_size(self) -> int:

1540 - return self.model_config.get_vocab_size()

1541 -

1542 -

1543 -class ModelRunner(GPUModelRunnerBase[ModelInputForGPUWithSamplingMetadata]):

1544 - """

1545 - GPU model runner with sampling step.

1546 - """

1547 - _model_input_cls: Type[ModelInputForGPUWithSamplingMetadata] = (

1548 - ModelInputForGPUWithSamplingMetadata)

1549 - _builder_cls: Type[ModelInputForGPUBuilder] = ModelInputForGPUBuilder

1550 -

1551 - def make_model_input_from_broadcasted_tensor_dict(

1552 - self,

1553 - tensor_dict: Dict[str, Any],

1554 - ) -> ModelInputForGPUWithSamplingMetadata:

1555 - model_input = \

1556 - ModelInputForGPUWithSamplingMetadata.from_broadcasted_tensor_dict(

1557 - tensor_dict,

1558 - attn_backend=self.attn_backend,

1559 - )

1560 - return model_input

1561 -

1562 - def prepare_model_input(

1563 - self,

1564 - seq_group_metadata_list: List[SequenceGroupMetadata],

1565 - virtual_engine: int = 0,

1566 - finished_requests_ids: Optional[List[str]] = None,

1567 - ) -> ModelInputForGPUWithSamplingMetadata:

1568 - """Prepare the model input based on a given sequence group, including

1569 - metadata for the sampling step.

1570 -

1571 - The API assumes seq_group_metadata_list is sorted by prefill -> decode.

1572 -

1573 - The result tensors and data structure also batches input in prefill

1574 - -> decode order. For example,

1575 -

1576 - - input_tokens[:num_prefill_tokens] contains prefill tokens.

1577 - - input_tokens[num_prefill_tokens:] contains decode tokens.

1578 -

1579 - If cuda graph is required, this API automatically pads inputs.

1580 - """

1581 - model_input = self._prepare_model_input_tensors(

1582 - seq_group_metadata_list, finished_requests_ids)

1583 - if get_pp_group().is_last_rank:

1584 - # Sampling metadata is only required for the final pp group

1585 - generators = self.get_generators(finished_requests_ids)

1586 - sampling_metadata = SamplingMetadata.prepare(

1587 - seq_group_metadata_list, model_input.seq_lens,

1588 - model_input.query_lens, self.device, self.pin_memory,

1589 - generators, self.sampling_metadata_cache)

1590 - else:

1591 - sampling_metadata = None

1592 - is_prompt = (seq_group_metadata_list[0].is_prompt

1593 - if seq_group_metadata_list else None)

1594 - return dataclasses.replace(model_input,

1595 - sampling_metadata=sampling_metadata,

1596 - is_prompt=is_prompt,

1597 - virtual_engine=virtual_engine)

1598 -

1599 - @torch.inference_mode()

1600 - def execute_model(

1601 - self,

1602 - model_input: ModelInputForGPUWithSamplingMetadata,

1603 - kv_caches: List[torch.Tensor],

1604 - intermediate_tensors: Optional[IntermediateTensors] = None,

1605 - num_steps: int = 1,

1606 - **kwargs,

1607 - ) -> Optional[Union[List[SamplerOutput], IntermediateTensors]]:

1608 - if num_steps > 1:

1609 - raise ValueError("num_steps > 1 is not supported in ModelRunner")

1610 -

1611 - if self.lora_config:

1612 - assert model_input.lora_requests is not None

1613 - assert model_input.lora_mapping is not None

1614 - self.set_active_loras(model_input.lora_requests,

1615 - model_input.lora_mapping)

1616 -

1617 - self.attn_state.begin_forward(model_input)

1618 -

1619 - # Currently cuda graph is only supported by the decode phase.

1620 - assert model_input.attn_metadata is not None

1621 - prefill_meta = model_input.attn_metadata.prefill_metadata

1622 - decode_meta = model_input.attn_metadata.decode_metadata

1623 - # TODO(andoorve): We can remove this once all

1624 - # virtual engines share the same kv cache.

1625 - virtual_engine = model_input.virtual_engine

1626 - previous_hidden_states = kwargs.get("previous_hidden_states")

1627 - if prefill_meta is None and decode_meta.use_cuda_graph:

1628 - assert model_input.input_tokens is not None

1629 - graph_batch_size = model_input.input_tokens.shape[0]

1630 - use_inputs_embeds = model_input.inputs_embeds is not None

1631 - model_executable = self.graph_runners[virtual_engine][(

1632 - graph_batch_size, use_inputs_embeds)]

1633 - if previous_hidden_states is not None:

1634 - previous_hidden_states = torch.cat([

1635 - previous_hidden_states,

1636 - torch.empty([

1637 - graph_batch_size - previous_hidden_states.shape[0],

1638 - *previous_hidden_states.shape[1:]

1639 - ],

1640 - dtype=previous_hidden_states.dtype,

1641 - device=previous_hidden_states.device)

1642 - ])

1643 - else:

1644 - model_executable = self.model

1645 -

1646 - # Receive KV cache in distributed KV cache transfer setting

1647 - # In disagg prefill setting, it will also recv hidden states and bypass

1648 - # model forwarding

1649 - # In KV cache database setting, it will change the model input so that

1650 - # we can skip prefilling on tokens that successfully received KV caches

1651 - # NOTE: The receive operation is blocking

1652 - bypass_model_exec = False

1653 - if self.need_recv_kv(model_input, kv_caches):

1654 - hidden_or_intermediate_states, bypass_model_exec, model_input = \

1655 - get_kv_transfer_group().recv_kv_caches_and_hidden_states(

1656 - # model is used to know which layer the current worker

1657 - # is working on, so that we can receive KV for only those

1658 - # layers.

1659 - model_executable,

1660 - model_input,

1661 - kv_caches=kv_caches

1662 - )

1663 -

1664 - multi_modal_kwargs = model_input.multi_modal_kwargs or {}

1665 - seqlen_agnostic_kwargs = {

1666 - "finished_requests_ids": model_input.finished_requests_ids,

1667 - "request_ids_to_seq_ids": model_input.request_ids_to_seq_ids,

1668 - } if self.has_inner_state else {}

1669 - model_kwargs = {}

1670 - if previous_hidden_states is not None:

1671 - model_kwargs["previous_hidden_states"] = previous_hidden_states

1672 - if (self.observability_config is not None

1673 - and self.observability_config.collect_model_forward_time):

1674 - model_forward_start = torch.cuda.Event(enable_timing=True)

1675 - model_forward_end = torch.cuda.Event(enable_timing=True)

1676 - model_forward_start.record()

1677 -

1678 - if not bypass_model_exec:

1679 - with set_forward_context(model_input.attn_metadata,

1680 - self.vllm_config, virtual_engine):

1681 - hidden_or_intermediate_states = model_executable(

1682 - input_ids=model_input.input_tokens,

1683 - inputs_embeds=model_input.inputs_embeds,

1684 - positions=model_input.input_positions,

1685 - intermediate_tensors=intermediate_tensors,

1686 - **MultiModalKwargs.as_kwargs(

1687 - multi_modal_kwargs,

1688 - device=self.device,

1689 - ),

1690 - **seqlen_agnostic_kwargs,

1691 - **model_kwargs,

1692 - )

1693 -

1694 - if (self.observability_config is not None

1695 - and self.observability_config.collect_model_forward_time):

1696 - model_forward_end.record()

1697 -

1698 - # Sending KV cache in distributed KV cache transfer setting

1699 - # NOTE: the send operation is non-blocking

1700 - if self.need_send_kv(model_input, kv_caches):

1701 - get_kv_transfer_group().send_kv_caches_and_hidden_states(

1702 - # model_executable is used to know which layer the current

1703 - # worker is working on, so that we can send KV for only those

1704 - # layers.

1705 - model_executable,

1706 - model_input,

1707 - kv_caches,

1708 - hidden_or_intermediate_states,

1709 - )

1710 -

1711 - # Compute the logits in the last pipeline stage.

1712 - if not get_pp_group().is_last_rank:

1713 - if (self.is_driver_worker

1714 - and hidden_or_intermediate_states is not None

1715 - and isinstance(hidden_or_intermediate_states,

1716 - IntermediateTensors)

1717 - and self.observability_config is not None

1718 - and self.observability_config.collect_model_forward_time):

1719 - model_forward_end.synchronize()

1720 - model_forward_time = model_forward_start.elapsed_time(

1721 - model_forward_end)

1722 - orig_model_forward_time = 0.0

1723 - if intermediate_tensors is not None:

1724 - orig_model_forward_time = intermediate_tensors.tensors.get(

1725 - "model_forward_time", torch.tensor(0.0)).item()

1726 - hidden_or_intermediate_states.tensors["model_forward_time"] = (

1727 - torch.tensor(model_forward_time + orig_model_forward_time))

1728 - return hidden_or_intermediate_states

1729 -

1730 - logits = self.model.compute_logits(hidden_or_intermediate_states,

1731 - model_input.sampling_metadata)

1732 -

1733 - if self.is_driver_worker:

1734 - if model_input.async_callback is not None:

1735 - model_input.async_callback()

1736 -

1737 - # Sample the next token.

1738 - assert isinstance(self.sampler, Sampler)

1739 - orig_include_gpu_probs = self.sampler.include_gpu_probs_tensor

1740 - if model_input.inputs_embeds is not None:

1741 - self.sampler.include_gpu_probs_tensor = True

1742 -

1743 - output: SamplerOutput = self.sampler(

1744 - logits=logits,

1745 - sampling_metadata=model_input.sampling_metadata,

1746 - )

1747 - if (self.observability_config is not None

1748 - and self.observability_config.collect_model_forward_time

1749 - and output is not None):

1750 - model_forward_end.synchronize()

1751 - model_forward_time = model_forward_start.elapsed_time(

1752 - model_forward_end)

1753 - orig_model_forward_time = 0.0

1754 - if intermediate_tensors is not None:

1755 - orig_model_forward_time = intermediate_tensors.tensors.get(

1756 - "model_forward_time", torch.tensor(0.0)).item()

1757 - # If there are multiple workers, we are still tracking the

1758 - # latency from the start time of the driver worker to the end

1759 - # time of the driver worker. The model forward time will then

1760 - # end up covering the communication time as well.

1761 - output.model_forward_time = (orig_model_forward_time +

1762 - model_forward_time)

1763 -

1764 - if model_input.inputs_embeds is not None:

1765 - if self.is_driver_worker:

1766 - sampled_token_ids = []

1767 - valid_outputs = []

1768 - for sequence_group_output in output.outputs:

1769 - if len(sequence_group_output.samples) == 0:

1770 - continue

1771 - assert len(sequence_group_output.samples) == 1

1772 - valid_outputs.append(sequence_group_output)

1773 - sampled_token_ids.append(

1774 - sequence_group_output.samples[0].output_token)

1775 - sampled_token_ids = torch.tensor(sampled_token_ids).to(

1776 - self.device)

1777 - sampled_token_ids = broadcast_tensor_dict(

1778 - {"sampled_token_ids":

1779 - sampled_token_ids})["sampled_token_ids"]

1780 - else:

1781 - sampled_token_ids = broadcast_tensor_dict(

1782 - )["sampled_token_ids"]

1783 - if len(sampled_token_ids) > 0:

1784 - sampled_token_embeds = \

1785 - self.model.get_input_embeddings(sampled_token_ids)

1786 - if self.is_driver_worker:

1787 - self.sampler.include_gpu_probs_tensor = \

1788 - orig_include_gpu_probs

1789 - for i, sequence_group_output in enumerate(valid_outputs):

1790 - sequence_group_output.samples[0].output_embed = \

1791 - sampled_token_embeds[i]

1792 -

1793 - if not self.is_driver_worker:

1794 - return []

1795 -

1796 - if self.return_hidden_states:

1797 - # we only need to pass hidden states of most recent token

1798 - assert model_input.sampling_metadata is not None

1799 - indices = model_input.sampling_metadata.selected_token_indices

1800 - if model_input.is_prompt:

1801 - hidden_states = hidden_or_intermediate_states.index_select(

1802 - 0, indices)

1803 - output.prefill_hidden_states = hidden_or_intermediate_states

1804 - elif decode_meta.use_cuda_graph:

1805 - hidden_states = hidden_or_intermediate_states[:len(indices)]

1806 - else:

1807 - hidden_states = hidden_or_intermediate_states

1808 -

1809 - output.hidden_states = hidden_states

1810 -

1811 - return [output]

1812 -

1813 - def need_recv_kv(self, model_input: ModelInputForGPUWithSamplingMetadata,

1814 - kv_caches: List[torch.Tensor]) -> bool:

1815 - """Check if we need to receive kv-cache from the other worker.

1816 - We need to receive KV when

1817 - 1. current vLLM instance is KV cache consumer/decode vLLM instance

1818 - 2. this batch is not a profiling run

1819 - 3. this batch is a prefill run

1820 -

1821 - Args:

1822 - model_input: input to the model executable

1823 - kv_caches: vLLM's paged memory

1824 - """

1825 -

1826 - if self.vllm_config.kv_transfer_config is None:

1827 - return False

1828 -

1829 - if model_input.attn_metadata is None:

1830 - raise ValueError("model_input.attn_metadata cannot be None")

1831 -

1832 - prefill_meta = model_input.attn_metadata.prefill_metadata

1833 -

1834 - # check if the current run is profiling

1835 - is_profile_run = (kv_caches[0].numel() == 0)

1836 - # check if the current run is prefill

1837 - is_prefill_run = prefill_meta is not None

1838 -

1839 - return self.vllm_config.kv_transfer_config.is_kv_consumer and (

1840 - not is_profile_run) and is_prefill_run

1841 -

1842 - def need_send_kv(self, model_input: ModelInputForGPUWithSamplingMetadata,

1843 - kv_caches: List[torch.Tensor]) -> bool:

1844 - """Check if we need to send kv-cache to the other worker.

1845 - We need to send KV when

1846 - 1. current vLLM instance is KV cache producer/prefill vLLM instance

1847 - 2. this batch is not a profiling run

1848 - 3. this batch is a prefill run

1849 -

1850 - Args:

1851 - model_input: input to the model executable

1852 - kv_caches: vLLM's paged memory

1853 - """

1854 -

1855 - if self.vllm_config.kv_transfer_config is None:

1856 - return False

1857 -

1858 - if model_input.attn_metadata is None:

1859 - raise ValueError("model_input.attn_metadata cannot be None")

1860 -

1861 - prefill_meta = model_input.attn_metadata.prefill_metadata

1862 -

1863 - # check if the current run is profiling

1864 - is_profile_run = (kv_caches[0].numel() == 0)

1865 - # check if the current run is prefill

1866 - is_prefill_run = prefill_meta is not None

1867 -

1868 - return self.vllm_config.kv_transfer_config.is_kv_producer and (

1869 - not is_profile_run) and is_prefill_run

1870 -

1871 -

1872 -# NOTE: this is nn.Module so the profiler can properly capture/group

1873 -# kernels calls made within the graph

1874 -class CUDAGraphRunner(nn.Module):

1875 -

1876 - def __init__(self, model: nn.Module, backend_name: str,

1877 - attn_state: AttentionState, is_encoder_decoder_model: bool):

1878 - super().__init__()

1879 - self.model = model

1880 - self.backend_name = backend_name

1881 - self.attn_state = attn_state

1882 -

1883 - self.input_buffers: Dict[str, torch.Tensor] = {}

1884 - self.output_buffers: Dict[str, torch.Tensor] = {}

1885 -

1886 - self._graph: Optional[torch.cuda.CUDAGraph] = None

1887 - self._is_encoder_decoder_model = is_encoder_decoder_model

1888 -

1889 - @property

1890 - def graph(self):

1891 - assert self._graph is not None

1892 - return self._graph

1893 -

1894 - def capture(

1895 - self,

1896 - input_ids: torch.Tensor,

1897 - inputs_embeds: Optional[torch.Tensor],

1898 - positions: torch.Tensor,

1899 - intermediate_inputs: Optional[IntermediateTensors],

1900 - kv_caches: List[torch.Tensor],

1901 - attn_metadata: AttentionMetadata,

1902 - memory_pool: Optional[Tuple[int, int]],

1903 - stream: torch.cuda.Stream,

1904 - **kwargs,

1905 - ):

1906 - assert self._graph is None

1907 - # Run the model a few times without capturing the graph.

1908 - # This is to make sure that the captured graph does not include the

1909 - # kernel launches for initial benchmarking (e.g., Triton autotune).

1910 - # Note one iteration is not enough for torch.compile

1911 - for _ in range(_NUM_WARMUP_ITERS):

1912 - self.model(

1913 - input_ids=input_ids,

1914 - inputs_embeds=inputs_embeds,

1915 - positions=positions,

1916 - intermediate_tensors=intermediate_inputs,

1917 - **kwargs,

1918 - )

1919 - # Wait for the warm up operations to finish before proceeding with

1920 - # Graph Capture.

1921 - torch.cuda.synchronize()

1922 - # Capture the graph.

1923 - self._graph = torch.cuda.CUDAGraph()

1924 - with torch.cuda.graph(self._graph, pool=memory_pool, stream=stream):

1925 - output_hidden_or_intermediate_states = self.model(

1926 - input_ids=input_ids,

1927 - **({

1928 - "inputs_embeds": inputs_embeds,

1929 - } if inputs_embeds is not None else {}),

1930 - positions=positions,

1931 - intermediate_tensors=intermediate_inputs,

1932 - **kwargs,

1933 - )

1934 -

1935 - if isinstance(output_hidden_or_intermediate_states, torch.Tensor):

1936 - hidden_or_intermediate_states = weak_ref_tensor(

1937 - output_hidden_or_intermediate_states)

1938 - elif isinstance(output_hidden_or_intermediate_states,

1939 - IntermediateTensors):

1940 - hidden_or_intermediate_states = IntermediateTensors(

1941 - tensors={

1942 - key: weak_ref_tensor(value)

1943 - for key, value in

1944 - output_hidden_or_intermediate_states.tensors.items()

1945 - })

1946 -

1947 - del output_hidden_or_intermediate_states

1948 - # make sure `output_hidden_or_intermediate_states` is deleted

1949 - # in the graph's memory pool

1950 - gc.collect()

1951 - torch.cuda.synchronize()

1952 -

1953 - # Save the input and output buffers.

1954 - self.input_buffers = {

1955 - "input_ids":

1956 - input_ids,

1957 - **({

1958 - "inputs_embeds": inputs_embeds,

1959 - } if inputs_embeds is not None else {}),

1960 - "positions":

1961 - positions,

1962 - "kv_caches":

1963 - kv_caches,

1964 - **self.attn_state.get_graph_input_buffers(

1965 - attn_metadata, self._is_encoder_decoder_model),

1966 - **kwargs,

1967 - }

1968 - if intermediate_inputs is not None:

1969 - self.input_buffers.update(intermediate_inputs.tensors)

1970 - if get_pp_group().is_last_rank:

1971 - self.output_buffers = {

1972 - "hidden_states": hidden_or_intermediate_states

1973 - }

1974 - else:

1975 - self.output_buffers = hidden_or_intermediate_states

1976 -

1977 - def forward(

1978 - self,

1979 - input_ids: torch.Tensor,

1980 - inputs_embeds: Optional[torch.Tensor],

1981 - positions: torch.Tensor,

1982 - intermediate_tensors: Optional[IntermediateTensors],

1983 - **kwargs,

1984 - ) -> torch.Tensor:

1985 - attn_metadata: AttentionMetadata = get_forward_context().attn_metadata

1986 -

1987 - # Copy the input tensors to the input buffers.

1988 - self.input_buffers["input_ids"].copy_(input_ids, non_blocking=True)

1989 - if positions is not None:

1990 - # in some case like MLA, it will reuse positions in metadata

1991 - # but truncate them to the original size

1992 - # so the shape is not padded, we need to copy partial only

1993 - self.input_buffers["positions"][:positions.shape[0]].copy_(

1994 - positions, non_blocking=True)

1995 - if inputs_embeds is not None:

1996 - self.input_buffers["inputs_embeds"][:inputs_embeds.shape[0]].copy_(

1997 - inputs_embeds, non_blocking=True)

1998 -

1999 - if self.backend_name != "NO_ATTENTION":

2000 - self.input_buffers["slot_mapping"].copy_(

2001 - attn_metadata.slot_mapping, non_blocking=True)

2002 -

2003 - self.attn_state.prepare_graph_input_buffers(

2004 - self.input_buffers, attn_metadata, self._is_encoder_decoder_model)

2005 -

2006 - if "seqlen_agnostic_capture_inputs" in self.input_buffers:

2007 - self.model.copy_inputs_before_cuda_graphs(self.input_buffers,

2008 - **kwargs)

2009 -

2010 - if "previous_hidden_states" in self.input_buffers:

2011 - self.input_buffers["previous_hidden_states"].copy_(

2012 - kwargs["previous_hidden_states"], non_blocking=True)

2013 -

2014 - if intermediate_tensors is not None:

2015 - for key in intermediate_tensors.tensors:

2016 - if key != "model_execute_time" and key != "model_forward_time":

2017 - self.input_buffers[key].copy_(intermediate_tensors[key],

2018 - non_blocking=True)

2019 - if self._is_encoder_decoder_model:

2020 - self.input_buffers["encoder_input_ids"].copy_(

2021 - kwargs['encoder_input_ids'], non_blocking=True)

2022 - self.input_buffers["encoder_positions"].copy_(

2023 - kwargs['encoder_positions'], non_blocking=True)

2024 -

2025 - # Run the graph.

2026 - self.graph.replay()

2027 - # Return the output tensor.

2028 - if get_pp_group().is_last_rank:

2029 - return self.output_buffers["hidden_states"]

2030 -

2031 - return self.output_buffers

vllm/worker/worker.py deleted

-666

@@ -1,666 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

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

3 -"""A GPU worker class."""

4 -import gc

5 -import os

6 -from contextlib import nullcontext

7 -from typing import Dict, List, Optional, Set, Tuple, Type, Union

8 -

9 -import torch

10 -import torch.distributed

11 -

12 -import vllm.envs as envs

13 -from vllm.attention.layer import Attention

14 -from vllm.config import VllmConfig, get_layers_from_vllm_config

15 -from vllm.device_allocator.cumem import CuMemAllocator

16 -from vllm.distributed import (ensure_model_parallel_initialized,

17 - init_distributed_environment,

18 - set_custom_all_reduce)

19 -from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized

20 -from vllm.logger import init_logger

21 -from vllm.lora.request import LoRARequest

22 -from vllm.model_executor import set_random_seed

23 -from vllm.model_executor.layers.sampler import SamplerOutput

24 -from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

25 -from vllm.platforms import current_platform

26 -from vllm.sequence import (ExecuteModelRequest, IntermediateTensors,

27 - SequenceGroupMetadata, SequenceGroupMetadataDelta)

28 -from vllm.utils import (GiB_bytes, MemorySnapshot, bind_kv_cache,

29 - memory_profiling)

30 -from vllm.worker.cache_engine import CacheEngine

31 -from vllm.worker.model_runner import GPUModelRunnerBase, ModelRunner

32 -from vllm.worker.worker_base import (LocalOrDistributedWorkerBase, WorkerBase,

33 - WorkerInput)

34 -

35 -logger = init_logger(__name__)

36 -

37 -

38 -class Worker(LocalOrDistributedWorkerBase):

39 - """A worker class that executes (a partition of) the model on a GPU.

40 -

41 - Each worker is associated with a single GPU. The worker is responsible for

42 - maintaining the KV cache and executing the model on the GPU. In case of

43 - distributed inference, each worker is assigned a partition of the model.

44 - """

45 -

46 - def __init__(

47 - self,

48 - vllm_config: VllmConfig,

49 - local_rank: int,

50 - rank: int,

51 - distributed_init_method: str,

52 - is_driver_worker: bool = False,

53 - model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None,

54 - ) -> None:

55 - WorkerBase.__init__(self, vllm_config)

56 - self.parallel_config.rank = rank

57 - self.local_rank = local_rank

58 - self.rank = rank

59 - self.distributed_init_method = distributed_init_method

60 - self.is_driver_worker = is_driver_worker

61 - if self.model_config.trust_remote_code:

62 - # note: lazy import to avoid importing torch before initializing

63 - from vllm.utils import init_cached_hf_modules

64 - init_cached_hf_modules()

65 -

66 - # Return hidden states from target model if the draft model is an

67 - # mlp_speculator

68 - speculative_config = self.speculative_config

69 - model_config = self.model_config

70 - speculative_args = {} if speculative_config is None \

71 - or (speculative_config.draft_model_config.hf_config.model_type ==

72 - model_config.hf_config.model_type) \

73 - or (speculative_config.draft_model_config.hf_config.model_type

74 - not in ("medusa",

75 - "mlp_speculator",

76 - "eagle",

77 - "deepseek_mtp",

78 - "glm4_moe_mtp",

79 - "mimo_mtp",

80 - "ernie_mtp",

81 - "qwen3_next_mtp")) \

82 - else {"return_hidden_states": True}

83 -

84 - self.model_runner: GPUModelRunnerBase = ModelRunner(

85 - vllm_config=self.vllm_config,

86 - kv_cache_dtype=self.cache_config.cache_dtype,

87 - is_driver_worker=is_driver_worker,

88 - **speculative_args,

89 - )

90 - if model_runner_cls is not None:

91 - self.model_runner = model_runner_cls(self.model_runner)

92 -

93 - # Uninitialized cache engine. Will be initialized by

94 - # initialize_cache.

95 - self.cache_engine: List[CacheEngine]

96 - self.gpu_cache: Optional[List[List[torch.Tensor]]] = None

97 - self._seq_group_metadata_cache: Dict[str, SequenceGroupMetadata] = {}

98 -

99 - # Buffers saved before sleep

100 - self._sleep_saved_buffers: Dict[str, torch.Tensor] = {}

101 -

102 - # Torch profiler. Enabled and configured through env vars:

103 - # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace

104 - if envs.VLLM_TORCH_PROFILER_DIR:

105 - torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR

106 - logger.info("Profiling enabled. Traces will be saved to: %s",

107 - torch_profiler_trace_dir)

108 - self.profiler = torch.profiler.profile(

109 - activities=[

110 - torch.profiler.ProfilerActivity.CPU,

111 - torch.profiler.ProfilerActivity.CUDA,

112 - ],

113 - with_stack=True,

114 - on_trace_ready=torch.profiler.tensorboard_trace_handler(

115 - torch_profiler_trace_dir, use_gzip=True))

116 - else:

117 - self.profiler = None

118 -

119 - def start_profile(self):

120 - if self.profiler is None:

121 - raise RuntimeError("Profiler is not enabled.")

122 - self.profiler.start()

123 -

124 - def stop_profile(self):

125 - if self.profiler is None:

126 - raise RuntimeError("Profiler is not enabled.")

127 - self.profiler.stop()

128 - # only print profiler results on rank 0

129 - if self.local_rank == 0:

130 - print(self.profiler.key_averages().table(

131 - sort_by="self_cuda_time_total"))

132 -

133 - def sleep(self, level: int = 1) -> None:

134 - free_bytes_before_sleep = torch.cuda.mem_get_info()[0]

135 -

136 - # Save the buffers before level 2 sleep

137 - if level == 2:

138 - model = self.model_runner.model

139 - self._sleep_saved_buffers = {

140 - name: buffer.cpu().clone()

141 - for name, buffer in model.named_buffers()

142 - }

143 -

144 - allocator = CuMemAllocator.get_instance()

145 - allocator.sleep(offload_tags=("weights", ) if level == 1 else tuple())

146 - free_bytes_after_sleep, total = torch.cuda.mem_get_info()

147 - freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep

148 - used_bytes = total - free_bytes_after_sleep

149 - assert freed_bytes >= 0, "Memory usage increased after sleeping."

150 - logger.info(

151 - "Sleep mode freed %.2f GiB memory, "

152 - "%.2f GiB memory is still in use.", freed_bytes / GiB_bytes,

153 - used_bytes / GiB_bytes)

154 -

155 - def wake_up(self, tags: Optional[list[str]] = None) -> None:

156 - allocator = CuMemAllocator.get_instance()

157 - allocator.wake_up(tags=tags)

158 -

159 - # Restore the buffers after level 2 sleep

160 - if len(self._sleep_saved_buffers):

161 - model = self.model_runner.model

162 - for name, buffer in model.named_buffers():

163 - if name in self._sleep_saved_buffers:

164 - buffer.data.copy_(self._sleep_saved_buffers[name].data)

165 - self._sleep_saved_buffers = {}

166 -

167 - def init_device(self) -> None:

168 - if self.device_config.device.type == "cuda":

169 - # torch.distributed.all_reduce does not free the input tensor until

170 - # the synchronization point. This causes the memory usage to grow

171 - # as the number of all_reduce calls increases. This env var disables

172 - # this behavior.

173 - # Related issue:

174 - # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573

175 - os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"

176 -

177 - # This env var set by Ray causes exceptions with graph building.

178 - os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)

179 - self.device = torch.device(f"cuda:{self.local_rank}")

180 - torch.cuda.set_device(self.device)

181 -

182 - _check_if_gpu_supports_dtype(self.model_config.dtype)

183 - gc.collect()

184 - torch.cuda.empty_cache()

185 - torch.cuda.reset_peak_memory_stats()

186 - self.baseline_snapshot = MemorySnapshot()

187 - else:

188 - raise RuntimeError(

189 - f"Not support device type: {self.device_config.device}")

190 - # Initialize the distributed environment.

191 - init_worker_distributed_environment(self.vllm_config, self.rank,

192 - self.distributed_init_method,

193 - self.local_rank)

194 - # Set random seed.

195 - set_random_seed(self.model_config.seed)

196 -

197 - def load_model(self):

198 - if self.vllm_config.model_config.enable_sleep_mode:

199 - allocator = CuMemAllocator.get_instance()

200 - assert allocator.get_current_usage() == 0, (

201 - "Sleep mode can only be "

202 - "used for one instance per process.")

203 - context = allocator.use_memory_pool(tag="weights")

204 - else:

205 - context = nullcontext()

206 - with context:

207 - self.model_runner.load_model()

208 -

209 - def save_sharded_state(

210 - self,

211 - path: str,

212 - pattern: Optional[str] = None,

213 - max_size: Optional[int] = None,

214 - ) -> None:

215 - self.model_runner.save_sharded_state(

216 - path,

217 - pattern=pattern,

218 - max_size=max_size,

219 - )

220 -

221 - def save_tensorized_model(

222 - self,

223 - tensorizer_config: TensorizerConfig,

224 - ) -> None:

225 - self.model_runner.save_tensorized_model(

226 - tensorizer_config=tensorizer_config, )

227 -

228 - @torch.inference_mode()

229 - def determine_available_kv_cache_memory(self,

230 - total_gpu_memory: int) -> float:

231 - if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:

232 - # still need a profile run which compiles the model for

233 - # max_num_batched_tokens

234 - self.model_runner.profile_run()

235 -

236 - GiB = lambda b: b / GiB_bytes

237 - msg = (

238 - f"Initial free memory "

239 - f"{GiB(self.baseline_snapshot.free_memory):.2f} "

240 - f"GiB, reserved {GiB(kv_cache_memory_bytes):.2f}GiB memory for "

241 - "KV Cache as specified by kv_cache_memory_bytes config and "

242 - "skipped memory profiling. This does does not respect the "

243 - "gpu_memory_utilization config. Only use kv_cache_memory_bytes "

244 - "config when you want manual control of KV cache memory "

245 - "size. If OOM'ed, check the difference of initial free "

246 - "memory between the current run and the previous run "

247 - "where kv_cache_memory_bytes is suggested and update it "

248 - "correspondingly.")

249 - logger.info(msg)

250 - return self.cache_config.kv_cache_memory_bytes

251 -

252 - # Execute a forward pass with dummy inputs to profile the memory usage

253 - # of the model.

254 - with memory_profiling(

255 - self.baseline_snapshot,

256 - weights_memory=self.model_runner.model_memory_usage) as result:

257 - self.model_runner.profile_run()

258 -

259 - self.non_torch_memory = result.non_torch_increase

260 - self.peak_activation_memory = result.torch_peak_increase

261 -

262 - self._assert_memory_footprint_increased_during_profiling()

263 -

264 - self.requested_memory = total_gpu_memory * \

265 - self.cache_config.gpu_memory_utilization

266 -

267 - self.available_kv_cache_memory = (self.requested_memory -

268 - result.non_kv_cache_memory)

269 -

270 - msg = (f"Memory profiling takes {result.profile_time:.2f} seconds\n"

271 - "the current vLLM instance can use "

272 - "total_gpu_memory "

273 - f"({(total_gpu_memory / GiB_bytes):.2f}GiB)"

274 - " x gpu_memory_utilization "

275 - f"({self.cache_config.gpu_memory_utilization:.2f})"

276 - f" = {(self.requested_memory / GiB_bytes):.2f}GiB\n"

277 - "model weights take "

278 - f"{(result.weights_memory / GiB_bytes):.2f}GiB;"

279 - " non_torch_memory takes "

280 - f"{(result.non_torch_increase / GiB_bytes):.2f}GiB;"

281 - " PyTorch activation peak memory takes "

282 - f"{(result.torch_peak_increase / GiB_bytes):.2f}GiB;"

283 - " the rest of the memory reserved for KV Cache is "

284 - f"{(self.available_kv_cache_memory / GiB_bytes):.2f}GiB.")

285 -

286 - logger.info(msg)

287 - return self.available_kv_cache_memory

288 -

289 - @torch.inference_mode()

290 - def determine_num_available_blocks(self) -> Tuple[int, int]:

291 - """Profiles the peak memory usage of the model to determine how many

292 - KV blocks may be allocated without OOMs.

293 -

294 - The engine will first conduct a profiling of the existing memory usage.

295 - Then, it calculates the maximum possible number of GPU and CPU blocks

296 - that can be allocated with the remaining free memory.

297 -

298 - Tip:

299 - You may limit the usage of GPU memory

300 - by adjusting the `gpu_memory_utilization` parameter.

301 - """

302 - # Profile the memory usage of the model and get the maximum number of

303 - # cache blocks that can be allocated with the remaining free memory.

304 - torch.cuda.empty_cache()

305 - torch.cuda.reset_peak_memory_stats()

306 -

307 - free_memory_pre_profile, total_gpu_memory = torch.cuda.mem_get_info()

308 - available_kv_cache_memory = self.determine_available_kv_cache_memory(

309 - total_gpu_memory)

310 -

311 - # Calculate the number of blocks that can be allocated with the

312 - # profiled peak memory.

313 - cache_block_size = self.get_cache_block_size_bytes()

314 - if cache_block_size == 0:

315 - num_gpu_blocks = 0

316 - num_cpu_blocks = 0

317 - else:

318 - num_gpu_blocks = int(available_kv_cache_memory // cache_block_size)

319 - num_cpu_blocks = int(self.cache_config.swap_space_bytes //

320 - cache_block_size)

321 - num_gpu_blocks = max(num_gpu_blocks, 0)

322 - num_cpu_blocks = max(num_cpu_blocks, 0)

323 -

324 - # Final cleanup

325 - gc.collect()

326 -

327 - return num_gpu_blocks, num_cpu_blocks

328 -

329 - def _assert_memory_footprint_increased_during_profiling(self):

330 - # NOTE(woosuk): Here we assume that the other processes using the same

331 - # GPU did not change their memory usage during the profiling.

332 - free_gpu_memory, total = torch.cuda.mem_get_info()

333 - cuda_memory = total - free_gpu_memory

334 - assert self.baseline_snapshot.cuda_memory < cuda_memory, (

335 - "Error in memory profiling. "

336 - f"Initial used memory {self.baseline_snapshot.cuda_memory}, "

337 - f"currently used memory {cuda_memory}. "

338 - f"This happens when the GPU memory was "

339 - "not properly cleaned up before initializing the vLLM instance.")

340 -

341 - def initialize_cache(self, num_gpu_blocks: int,

342 - num_cpu_blocks: int) -> None:

343 - """Allocate GPU and CPU KV cache with the specified number of blocks.

344 -

345 - This also warms up the model, which may record CUDA graphs.

346 - """

347 - raise_if_cache_size_invalid(

348 - num_gpu_blocks, self.cache_config.block_size,

349 - self.cache_config.is_attention_free,

350 - self.model_config.max_model_len,

351 - self.parallel_config.pipeline_parallel_size)

352 -

353 - self.cache_config.num_gpu_blocks = num_gpu_blocks

354 - self.cache_config.num_cpu_blocks = num_cpu_blocks

355 -

356 - if self.vllm_config.model_config.enable_sleep_mode:

357 - allocator = CuMemAllocator.get_instance()

358 - context = allocator.use_memory_pool(tag="kv_cache")

359 - else:

360 - context = nullcontext()

361 - with context:

362 - self._init_cache_engine()

363 - self._warm_up_model()

364 -

365 - def _init_cache_engine(self):

366 - assert self.cache_config.num_gpu_blocks is not None

367 - self.cache_engine = [

368 - CacheEngine(self.cache_config, self.model_config,

369 - self.parallel_config, self.device_config)

370 - for _ in range(self.parallel_config.pipeline_parallel_size)

371 - ]

372 - self.gpu_cache = [

373 - self.cache_engine[ve].gpu_cache

374 - for ve in range(self.parallel_config.pipeline_parallel_size)

375 - ]

376 -

377 - # Layer pairings for cross-layer KV sharing.

378 - # If an Attention layer `layer_name` is in the keys of this dict, it

379 - # means this layer will perform attention using the keys and values

380 - # from the KV cache of `shared_kv_cache_layers[layer_name]`.

381 - shared_kv_cache_layers: dict[str, str] = {}

382 -

383 - attn_layers = get_layers_from_vllm_config(self.vllm_config, Attention)

384 -

385 - for layer_name, attn_module in attn_layers.items():

386 - if (kv_tgt_layer :=

387 - attn_module.kv_sharing_target_layer_name) is not None:

388 - # The layer doesn't need its own KV cache and will use that of

389 - # the target layer. We skip creating a KVCacheSpec for it, so

390 - # that KV cache management logic will act as this layer does

391 - # not exist, and doesn't allocate KV cache for the layer. This

392 - # enables the memory saving of cross-layer kv sharing, allowing

393 - # a given amount of memory to accommodate longer context lengths

394 - # or enable more requests to be processed simultaneously.

395 - shared_kv_cache_layers[layer_name] = kv_tgt_layer

396 -

397 - bind_kv_cache(self.compilation_config.static_forward_context,

398 - self.gpu_cache, shared_kv_cache_layers)

399 -

400 - def _warm_up_model(self) -> None:

401 - # warm up sizes that are not in cudagraph capture sizes,

402 - # but users still want to compile for better performance,

403 - # e.g. for the max-num-batched token size in chunked prefill.

404 - warmup_sizes = self.vllm_config.compilation_config.compile_sizes.copy()

405 - if not self.model_config.enforce_eager:

406 - warmup_sizes = [

407 - x for x in warmup_sizes if x not in

408 - self.vllm_config.compilation_config.cudagraph_capture_sizes

409 - ]

410 - for size in sorted(warmup_sizes, reverse=True):

411 - logger.info("Compile and warming up model for size %d", size)

412 - self.model_runner._dummy_run(size)

413 -

414 - cuda_graph_memory_bytes = 0

415 - if not self.model_config.enforce_eager:

416 - cuda_graph_memory_bytes = self.model_runner.capture_model(

417 - self.gpu_cache)

418 -

419 - if (self.cache_config.kv_cache_memory_bytes is None

420 - and hasattr(self, "peak_activation_memory")):

421 - # Suggests optimal kv cache memory size if we rely on

422 - # memory_profiling to guess the kv cache memory size which

423 - # provides peak_activation_memory and a few other memory

424 - # consumption. `memory_profiling` does not consider

425 - # CUDAGraph memory size and may not utilize all gpu memory.

426 - # Users may want fine-grained control to specify kv cache

427 - # memory size.

428 - GiB = lambda b: round(b / GiB_bytes, 2)

429 - non_kv_cache_memory = (self.model_runner.model_memory_usage +

430 - self.peak_activation_memory +

431 - self.non_torch_memory +

432 - cuda_graph_memory_bytes)

433 -

434 - # empirically observed that the memory profiling may

435 - # slightly underestimate the memory consumption.

436 - # So leave a small buffer (=150MiB) to avoid OOM.

437 - redundancy_buffer_memory = 150 * (1 << 20)

438 - kv_cache_memory_bytes_to_gpu_limit = (

439 - self.baseline_snapshot.free_memory - non_kv_cache_memory -

440 - redundancy_buffer_memory)

441 - kv_cache_memory_bytes_to_requested_limit = (

442 - int(self.requested_memory) - non_kv_cache_memory -

443 - redundancy_buffer_memory)

444 -

445 - msg = (

446 - f"Free memory on device "

447 - f"({GiB(self.baseline_snapshot.free_memory)}/"

448 - f"{GiB(self.baseline_snapshot.total_memory)} GiB) on startup. "

449 - f"Desired GPU memory utilization is "

450 - f"({self.cache_config.gpu_memory_utilization}, "

451 - f"{GiB(self.requested_memory)} GiB). "

452 - f"Actual usage is {GiB(self.model_runner.model_memory_usage)} "

453 - f"GiB for weight, {GiB(self.peak_activation_memory)} GiB "

454 - f"for peak activation, {GiB(self.non_torch_memory)} GiB "

455 - f"for non-torch memory, and {GiB(cuda_graph_memory_bytes)} "

456 - f"GiB for CUDAGraph memory. Replace gpu_memory_utilization "

457 - f"config with `--kv-cache-memory="

458 - f"{kv_cache_memory_bytes_to_requested_limit}` to fit into "

459 - f"requested memory, or `--kv-cache-memory="

460 - f"{kv_cache_memory_bytes_to_gpu_limit}` to fully "

461 - f"utilize gpu memory. Current kv cache memory in use is "

462 - f"{int(self.available_kv_cache_memory)} bytes.")

463 - logger.info(msg)

464 -

465 - # Reset the seed to ensure that the random state is not affected by

466 - # the model initialization and profiling.

467 - set_random_seed(self.model_config.seed)

468 -

469 - @property

470 - def do_metadata_broadcast(self) -> bool:

471 - return self.parallel_config.tensor_parallel_size > 1

472 -

473 - @property

474 - def kv_cache(self) -> Optional[List[List[torch.Tensor]]]:

475 - return self.gpu_cache

476 -

477 - @torch.inference_mode()

478 - def prepare_worker_input(

479 - self, execute_model_req: ExecuteModelRequest) -> WorkerInput:

480 - virtual_engine = execute_model_req.virtual_engine

481 - num_steps = execute_model_req.num_steps

482 - num_seq_groups = len(execute_model_req.seq_group_metadata_list)

483 - # `blocks_to_swap_in` and `blocks_to_swap_out` are cpu tensors.

484 - # they contain parameters to launch cudamemcpyasync.

485 - blocks_to_swap_in = torch.tensor(execute_model_req.blocks_to_swap_in,

486 - device="cpu",

487 - dtype=torch.int64).view(-1, 2)

488 - blocks_to_swap_out = torch.tensor(execute_model_req.blocks_to_swap_out,

489 - device="cpu",

490 - dtype=torch.int64).view(-1, 2)

491 - # `blocks_to_copy` is a gpu tensor. The src and tgt of

492 - # blocks to copy are in the same device, and `blocks_to_copy`

493 - # can be used directly within cuda kernels.

494 - blocks_to_copy = torch.tensor(execute_model_req.blocks_to_copy,

495 - device=self.device,

496 - dtype=torch.int64).view(-1, 2)

497 -

498 - return WorkerInput(

499 - num_seq_groups=num_seq_groups,

500 - blocks_to_swap_in=blocks_to_swap_in,

501 - blocks_to_swap_out=blocks_to_swap_out,

502 - blocks_to_copy=blocks_to_copy,

503 - virtual_engine=virtual_engine,

504 - num_steps=num_steps,

505 - )

506 -

507 - @torch.inference_mode()

508 - def execute_worker(self, worker_input: WorkerInput) -> None:

509 - virtual_engine = worker_input.virtual_engine

510 - # Issue cache operations.

511 - if (worker_input.blocks_to_swap_in is not None

512 - and worker_input.blocks_to_swap_in.numel() > 0):

513 - self.cache_engine[virtual_engine].swap_in(

514 - worker_input.blocks_to_swap_in)

515 - if (worker_input.blocks_to_swap_out is not None

516 - and worker_input.blocks_to_swap_out.numel() > 0):

517 - self.cache_engine[virtual_engine].swap_out(

518 - worker_input.blocks_to_swap_out)

519 - if (worker_input.blocks_to_copy is not None

520 - and worker_input.blocks_to_copy.numel() > 0):

521 - self.cache_engine[virtual_engine].copy(worker_input.blocks_to_copy)

522 -

523 - def _get_cached_seq_group_metadata(

524 - self,

525 - seq_group_metadata_list: List[Union[SequenceGroupMetadata,

526 - SequenceGroupMetadataDelta]],

527 - finished_request_ids: List[str]) -> List[SequenceGroupMetadata]:

528 - """Return a list of cached Sequence Group Metadata after updating its

529 - state.

530 -

531 - It is used because scheduler only sends delta to workers to reduce

532 - the data payload size. The function also cleans up cache based on

533 - a given `finished_request_ids`.

534 - """

535 - new_seq_group_metadata_list = []

536 - for metadata_or_delta in seq_group_metadata_list:

537 - request_id = metadata_or_delta.request_id

538 - if request_id not in self._seq_group_metadata_cache:

539 - # The first prefill.

540 - assert isinstance(metadata_or_delta, SequenceGroupMetadata)

541 - self._seq_group_metadata_cache[request_id] = metadata_or_delta

542 - else:

543 - # The first prefill is already cached.

544 - if isinstance(metadata_or_delta, SequenceGroupMetadataDelta):

545 - self._seq_group_metadata_cache[request_id].apply_delta(

546 - metadata_or_delta)

547 - else:

548 - # If metadata snapshot is sent again, it is

549 - # preempted. Reset the cache because we need to start

550 - # from scratch.

551 - assert isinstance(metadata_or_delta, SequenceGroupMetadata)

552 - self._seq_group_metadata_cache[

553 - request_id] = metadata_or_delta

554 -

555 - new_seq_group_metadata_list.append(

556 - self._seq_group_metadata_cache[request_id])

557 -

558 - # Clean up finished ids

559 - for finished_id in finished_request_ids:

560 - del self._seq_group_metadata_cache[finished_id]

561 -

562 - return new_seq_group_metadata_list

563 -

564 - def _execute_model_spmd(

565 - self,

566 - execute_model_req: ExecuteModelRequest,

567 - intermediate_tensors: Optional[IntermediateTensors] = None,

568 - ) -> Optional[List[SamplerOutput]]:

569 - if execute_model_req is not None:

570 - new_seq_group_metadata_list = self._get_cached_seq_group_metadata(

571 - execute_model_req.seq_group_metadata_list,

572 - execute_model_req.finished_requests_ids)

573 -

574 - execute_model_req.seq_group_metadata_list = (

575 - new_seq_group_metadata_list)

576 - output = super()._execute_model_spmd(execute_model_req,

577 - intermediate_tensors)

578 - return output

579 -

580 - def add_lora(self, lora_request: LoRARequest) -> bool:

581 - return self.model_runner.add_lora(lora_request)

582 -

583 - def remove_lora(self, lora_id: int) -> bool:

584 - return self.model_runner.remove_lora(lora_id)

585 -

586 - def pin_lora(self, lora_id: int) -> bool:

587 - return self.model_runner.pin_lora(lora_id)

588 -

589 - def list_loras(self) -> Set[int]:

590 - return self.model_runner.list_loras()

591 -

592 - @property

593 - def max_model_len(self) -> int:

594 - return self.model_config.max_model_len

595 -

596 - @property

597 - def vocab_size(self) -> int:

598 - return self.model_runner.vocab_size

599 -

600 - def get_cache_block_size_bytes(self) -> int:

601 - """Get the size of the KV cache block size in bytes.

602 - """

603 - return CacheEngine.get_cache_block_size(self.cache_config,

604 - self.model_config,

605 - self.parallel_config)

606 -

607 -

608 -def init_worker_distributed_environment(

609 - vllm_config: VllmConfig,

610 - rank: int,

611 - distributed_init_method: Optional[str] = None,

612 - local_rank: int = -1,

613 -) -> None:

614 - """Initialize the distributed environment."""

615 - parallel_config = vllm_config.parallel_config

616 - set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)

617 -

618 - init_distributed_environment(parallel_config.world_size, rank,

619 - distributed_init_method, local_rank,

620 - current_platform.dist_backend)

621 - ensure_model_parallel_initialized(

622 - parallel_config.tensor_parallel_size,

623 - parallel_config.pipeline_parallel_size,

624 - parallel_config.decode_context_parallel_size)

625 -

626 - ensure_kv_transfer_initialized(vllm_config)

627 -

628 -

629 -def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):

630 - # Check if the GPU supports the dtype.

631 - if torch_dtype == torch.bfloat16: # noqa: SIM102

632 - if not current_platform.has_device_capability(80):

633 - capability = current_platform.get_device_capability()

634 - gpu_name = current_platform.get_device_name()

635 -

636 - if capability is None:

637 - compute_str = "does not have a compute capability"

638 - else:

639 - version_str = capability.as_version_str()

640 - compute_str = f"has compute capability {version_str}"

641 -

642 - raise ValueError(

643 - "Bfloat16 is only supported on GPUs with compute capability "

644 - f"of at least 8.0. Your {gpu_name} GPU {compute_str}. "

645 - "You can use float16 instead by explicitly setting the "

646 - "`dtype` flag in CLI, for example: --dtype=half.")

647 -

648 -

649 -def raise_if_cache_size_invalid(num_gpu_blocks, block_size, is_attention_free,

650 - max_model_len, pipeline_parallel_size) -> None:

651 - if is_attention_free and num_gpu_blocks != 0:

652 - raise ValueError("No memory should be allocated for the cache blocks "

653 - f"for an attention-free model, but {num_gpu_blocks} "

654 - "blocks are allocated.")

655 - if not is_attention_free and num_gpu_blocks <= 0:

656 - raise ValueError("No available memory for the cache blocks. "

657 - "Try increasing `gpu_memory_utilization` when "

658 - "initializing the engine.")

659 - max_seq_len = block_size * (num_gpu_blocks // pipeline_parallel_size)

660 - if not is_attention_free and max_model_len > max_seq_len:

661 - raise ValueError(

662 - f"The model's max seq len ({max_model_len}) "

663 - "is larger than the maximum number of tokens that can be "

664 - f"stored in KV cache ({max_seq_len}). Try increasing "

665 - "`gpu_memory_utilization` or decreasing `max_model_len` when "

666 - "initializing the engine.")