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

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

[CI] Fix mypy for `vllm/distributed` (#26593)

Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Wentao Ye committed Oct 13, 2025 at 16:02 UTC 314285d4f2134111345b4c4cbb583a925db9a91d

14 files changed +122 -65

tools/pre_commit/mypy.py

+1 -1

@@ -26,6 +26,7 @@ import regex as re

26 FILES = [

27 "vllm/*.py",

28 "vllm/assets",

29 + "vllm/distributed",

30 "vllm/entrypoints",

31 "vllm/inputs",

32 "vllm/logging_utils",

@@ -42,7 +43,6 @@ SEPARATE_GROUPS = [

43 "tests",

44 "vllm/attention",

45 "vllm/compilation",

45 - "vllm/distributed",

46 "vllm/engine",

47 "vllm/executor",

48 "vllm/inputs",

vllm/config/kv_transfer.py

+1 -1

@@ -27,7 +27,7 @@ class KVTransferConfig:

27 engine_id: str | None = None

28 """The engine id for KV transfers."""

30 - kv_buffer_device: str | None = "cuda"

30 + kv_buffer_device: str = "cuda"

31 """The device used by kv connector to buffer the KV cache. Choices are

32 'cuda' and 'cpu'."""

vllm/distributed/device_communicators/all2all.py

+24 -12

@@ -15,9 +15,11 @@ from vllm.utils.flashinfer import has_flashinfer_all2all

15 from .base_device_communicator import All2AllManagerBase, Cache

17 if has_flashinfer_all2all():

18 - from flashinfer.comm import Mapping

19 - from flashinfer.comm.mnnvl import MnnvlConfig

20 - from flashinfer.comm.trtllm_alltoall import MnnvlMoe

18 + from flashinfer.comm import Mapping # type: ignore[import-not-found]

19 + from flashinfer.comm.mnnvl import MnnvlConfig # type: ignore[import-not-found]

20 + from flashinfer.comm.trtllm_alltoall import (

21 + MnnvlMoe, # type: ignore[import-not-found]

22 + )

24 logger = init_logger(__name__)

@@ -65,6 +67,7 @@ class NaiveAll2AllManager(All2AllManagerBase):

67 ) -> tuple[torch.Tensor, torch.Tensor]:

68 sp_size = self.tp_group.world_size if is_sequence_parallel else 1

69 dp_metadata = get_forward_context().dp_metadata

70 + assert dp_metadata is not None

71 cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)

73 hidden_states = self.naive_multicast(

@@ -81,6 +84,7 @@ class NaiveAll2AllManager(All2AllManagerBase):

84 ep_rank = self.rank if is_sequence_parallel else self.dp_rank

86 dp_metadata = get_forward_context().dp_metadata

87 + assert dp_metadata is not None

88 sp_size = self.tp_group.world_size if is_sequence_parallel else 1

89 cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)

@@ -113,7 +117,10 @@ class AgRsAll2AllManager(All2AllManagerBase):

117 """

118 Gather hidden_states and router_logits from all dp ranks.

119 """

116 - sizes = get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()

120 + dp_metadata = get_forward_context().dp_metadata

121 + assert dp_metadata is not None

122 + sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

123 + assert sizes is not None

124

125 dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()

126 assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

@@ -130,7 +137,10 @@ class AgRsAll2AllManager(All2AllManagerBase):

137 """

138 Reduce-scatter hidden_states across all dp ranks.

139 """

133 - sizes = get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()

140 + dp_metadata = get_forward_context().dp_metadata

141 + assert dp_metadata is not None

142 + sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

143 + assert sizes is not None

144

145 dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()

146 hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)

@@ -155,7 +165,7 @@ class PPLXAll2AllManager(All2AllManagerBase):

165 if self.internode:

166 # inter-node communication needs nvshmem,

167 # intra-node communication uses p2p mapping directly

158 - from pplx_kernels.nvshmem import (

168 + from pplx_kernels.nvshmem import ( # type: ignore[import-not-found]

169 nvshmem_alloc_empty_unique_id,

170 nvshmem_get_unique_id,

171 nvshmem_init,

@@ -182,7 +192,7 @@ class PPLXAll2AllManager(All2AllManagerBase):

192 self.handle_cache = Cache()

193

194 def get_handle(self, kwargs):

185 - import pplx_kernels as pplx

195 + import pplx_kernels as pplx # type: ignore[import-not-found]

196

197 return self.handle_cache.get_or_create(

198 kwargs,

@@ -208,7 +218,9 @@ class PPLXAll2AllManager(All2AllManagerBase):

218 handle.destroy()

219

220 if self.internode:

211 - from pplx_kernels.nvshmem import nvshmem_finalize

221 + from pplx_kernels.nvshmem import (

222 + nvshmem_finalize, # type: ignore[import-not-found]

223 + )

224

225 logger.debug("PPLX NVSHMEM finalize")

226 nvshmem_finalize()

@@ -288,7 +300,7 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):

300 "args are computed in the Manager itself."

301 )

302

291 - import deep_ep

303 + import deep_ep # type: ignore[import-not-found]

304

305 buffer_kwargs = self._make_all2all_kwargs()

306 logger.debug("DeepEP all2all args %s", buffer_kwargs)

@@ -298,7 +310,7 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):

310 return handle

311

312 def set_num_sms(self, num_sms: int):

301 - import deep_ep

313 + import deep_ep # type: ignore[import-not-found]

314

315 # Right now the buffers are sized for only what the kernels were

316 # created with. So we can only reduce the number of SMS used

@@ -332,7 +344,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

344 num_global_experts: Number of experts in the model.

345 num_local_experts: Number of experts in an EP rank.

346 """

335 - import deep_ep

347 + import deep_ep # type: ignore[import-not-found]

348

349 # Defaults for internode and intranode are taken from DeepEP tests.

350 num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024

@@ -358,7 +370,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

370 The kwargs for DeepEPLLAll2AllManager is dictated by

371 _make_all2all_kwargs.

372 """

361 - import deep_ep

373 + import deep_ep # type: ignore[import-not-found]

374

375 buffer_kwargs = self._make_all2all_kwargs(**kwargs)

376 logger.debug("DeepEP all2all args %s", buffer_kwargs)

vllm/distributed/device_communicators/custom_all_reduce.py

+12 -7

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

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

4 from contextlib import contextmanager

5 +from typing import cast

7 import torch

8 import torch.distributed as dist

@@ -118,15 +119,18 @@ class CustomAllreduce:

119 # now `device` is a `torch.device` object

120 assert isinstance(device, torch.device)

121 self.device = device

121 - device_capability = current_platform.get_device_capability().as_version_str()

122 + device_capability = current_platform.get_device_capability()

123 if (

124 current_platform.is_cuda()

125 and symm_mem_enabled

125 - and device_capability in CUSTOM_ALL_REDUCE_MAX_SIZES

126 + and device_capability is not None

127 ):

127 - max_size = min(

128 - CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability][world_size], max_size

129 - )

128 + device_capability_str = device_capability.as_version_str()

129 + if device_capability_str in CUSTOM_ALL_REDUCE_MAX_SIZES:

130 + max_size = min(

131 + CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],

132 + max_size,

133 + )

134 cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES

135 if cuda_visible_devices:

136 device_ids = list(map(int, cuda_visible_devices.split(",")))

@@ -213,6 +217,7 @@ class CustomAllreduce:

217 # We cannot directly use `dist.all_gather_object` here

218 # because it is incompatible with `gloo` backend under inference mode.

219 # see https://github.com/pytorch/pytorch/issues/126032 for details.

220 + all_data: list[list[list[int] | None]]

221 all_data = [[None, None] for _ in range(dist.get_world_size(group=self.group))]

222 all_data[self.rank] = [handle, offset]

223 ranks = sorted(dist.get_process_group_ranks(group=self.group))

@@ -221,8 +226,8 @@ class CustomAllreduce:

226 all_data[i], src=rank, group=self.group, device="cpu"

227 )

228 # Unpack list of tuples to tuple of lists.

224 - handles = [d[0] for d in all_data] # type: ignore

225 - offsets = [d[1] for d in all_data] # type: ignore

229 + handles = cast(list[list[int]], [d[0] for d in all_data])

230 + offsets = cast(list[list[int]], [d[1] for d in all_data])

231 ops.register_graph_buffers(self._ptr, handles, offsets)

232

233 def should_custom_ar(self, inp: torch.Tensor):

vllm/distributed/device_communicators/symm_mem.py

+8 -3

@@ -52,9 +52,14 @@ class SymmMemCommunicator:

52 self.device = device

53 self.group = group

54 self.world_size = dist.get_world_size(self.group)

55 - self.device_capability = (

56 - current_platform.get_device_capability().as_version_str()

57 - )

55 + capability = current_platform.get_device_capability()

56 + if capability is None:

57 + logger.warning(

58 + "SymmMemCommunicator: device capability is unknown, "

59 + "communicator is not available."

60 + )

61 + return

62 + self.device_capability = capability.as_version_str()

63 if self.device_capability not in SYMM_MEM_ALL_REDUCE_MAX_SIZES:

64 logger.warning(

65 "SymmMemCommunicator: Device capability %s not supported, "

vllm/distributed/kv_transfer/kv_connector/factory.py

+12 -2

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

4 import importlib

5 from collections.abc import Callable

6 -from typing import TYPE_CHECKING

6 +from typing import TYPE_CHECKING, cast

8 import vllm.envs as envs

9 from vllm.distributed.kv_transfer.kv_connector.base import (

@@ -48,6 +48,8 @@ class KVConnectorFactory:

48 )

50 kv_transfer_config = config.kv_transfer_config

51 + if kv_transfer_config is None:

52 + raise ValueError("kv_transfer_config must be set to create a connector")

53 connector_cls = cls.get_connector_class(kv_transfer_config)

54 logger.info(

55 "Creating v1 connector with name: %s and engine_id: %s",

@@ -70,6 +72,8 @@ class KVConnectorFactory:

72 ) -> type[KVConnectorBaseType]:

73 """Get the connector class by name."""

74 connector_name = kv_transfer_config.kv_connector

75 + if connector_name is None:

76 + raise ValueError("Connector name is not set in KVTransferConfig")

77 if connector_name in cls._registry:

78 connector_cls = cls._registry[connector_name]()

79 else:

@@ -77,7 +81,13 @@ class KVConnectorFactory:

81 if connector_module_path is None:

82 raise ValueError(f"Unsupported connector type: {connector_name}")

83 connector_module = importlib.import_module(connector_module_path)

80 - connector_cls = getattr(connector_module, connector_name)

84 + try:

85 + connector_cls = getattr(connector_module, connector_name)

86 + except AttributeError as e:

87 + raise AttributeError(

88 + f"Class {connector_name} not found in {connector_module_path}"

89 + ) from e

90 + connector_cls = cast(type[KVConnectorBaseType], connector_cls)

91 return connector_cls

vllm/distributed/kv_transfer/kv_connector/utils.py

+7 -7

@@ -151,21 +151,21 @@ class KVOutputAggregator:

151 aggregated_kv_connector_stats = None

152 invalid_block_ids = set[int]()

153 for model_runner_output in outputs:

154 - output = model_runner_output.kv_connector_output

155 - if not output:

154 + kv_output = model_runner_output.kv_connector_output

155 + if not kv_output:

156 continue

157 update_finished_set(

158 - output.finished_sending, self._send_remaining_count, finished_sending

158 + kv_output.finished_sending, self._send_remaining_count, finished_sending

159 )

160 update_finished_set(

161 - output.finished_recving, self._recv_remaining_count, finished_recving

161 + kv_output.finished_recving, self._recv_remaining_count, finished_recving

162 )

163

164 # Aggregate kv_connector_stats from all workers.

165 if aggregated_kv_connector_stats is None:

166 # Use the first worker's kv_connector_stats as accumulator.

167 - aggregated_kv_connector_stats = output.kv_connector_stats

168 - elif kv_connector_stats := output.kv_connector_stats:

167 + aggregated_kv_connector_stats = kv_output.kv_connector_stats

168 + elif kv_connector_stats := kv_output.kv_connector_stats:

169 if aggregated_kv_connector_stats is None:

170 aggregated_kv_connector_stats = kv_connector_stats

171 else:

@@ -176,7 +176,7 @@ class KVOutputAggregator:

176 aggregated_kv_connector_stats.aggregate(kv_connector_stats)

177 )

178

179 - invalid_block_ids |= output.invalid_block_ids

179 + invalid_block_ids |= kv_output.invalid_block_ids

180

181 # select output of the worker specified by output_rank

182 output = outputs[output_rank]

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -95,6 +95,10 @@ class KVConnectorBase_V1(ABC):

95 )

96 self._connector_metadata: KVConnectorMetadata | None = None

97 self._vllm_config = vllm_config

98 + if vllm_config.kv_transfer_config is not None:

99 + self._kv_transfer_config = vllm_config.kv_transfer_config

100 + else:

101 + raise ValueError("kv_transfer_config must be set for KVConnectorBase_V1")

102 self._role = role

103

104 @property

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

+3 -4

@@ -86,13 +86,11 @@ class MultiConnector(KVConnectorBase_V1):

86 super().__init__(vllm_config=vllm_config, role=role)

87 self._connectors: list[KVConnectorBase_V1] = []

88 self._ktc_kv_transfer_config = []

89 - ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(

90 - "connectors"

91 - )

89 + ktcs = self._kv_transfer_config.kv_connector_extra_config.get("connectors")

90 assert ktcs is not None

91 for ktc in ktcs:

92 temp_config = copy.copy(vllm_config)

95 - engine_id = ktc.get("engine_id", vllm_config.kv_transfer_config.engine_id)

93 + engine_id = ktc.get("engine_id", self._kv_transfer_config.engine_id)

94 temp_config.kv_transfer_config = KVTransferConfig(

95 **ktc, engine_id=engine_id

96 )

@@ -296,6 +294,7 @@ class MultiConnector(KVConnectorBase_V1):

294 str: the required KV cache layout. e.g. HND, or NHD.

295 None if the connector does not require a specific layout.

296 """

297 + assert vllm_config.kv_transfer_config is not None

298 ktcs = vllm_config.kv_transfer_config.kv_connector_extra_config.get(

299 "connectors"

300 )

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+12 -6

@@ -297,6 +297,7 @@ class NixlConnectorScheduler:

297 + vllm_config.parallel_config.data_parallel_rank

298 * vllm_config.parallel_config.tensor_parallel_size

299 )

300 + assert vllm_config.kv_transfer_config is not None

301 self.use_host_buffer = vllm_config.kv_transfer_config.kv_buffer_device == "cpu"

302 logger.info("Initializing NIXL Scheduler %s", engine_id)

303

@@ -340,7 +341,8 @@ class NixlConnectorScheduler:

341

342 if params is not None and params.get("do_remote_prefill"):

343 # Remote prefill: get all prompt blocks from remote.

343 - count = len(request.prompt_token_ids) - num_computed_tokens

344 + token_ids = request.prompt_token_ids or []

345 + count = len(token_ids) - num_computed_tokens

346 if count > 0:

347 return count, True

348

@@ -521,6 +523,9 @@ class NixlConnectorWorker:

523 self.vllm_config = vllm_config

524 self.block_size = vllm_config.cache_config.block_size

525

526 + if vllm_config.kv_transfer_config is None:

527 + raise ValueError("kv_transfer_config must be set for NixlConnector")

528 +

529 self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config(

530 "backends", ["UCX"]

531 )

@@ -577,17 +582,18 @@ class NixlConnectorWorker:

582 self.use_host_buffer = self.kv_buffer_device == "cpu"

583 # support for oot platform which can't register nixl memory

584 # type based on kv_buffer_device

580 - self.nixl_memory_type = current_platform.get_nixl_memory_type()

581 - if self.nixl_memory_type is None:

585 + nixl_memory_type = current_platform.get_nixl_memory_type()

586 + if nixl_memory_type is None:

587 if self.kv_buffer_device == "cuda":

583 - self.nixl_memory_type = "VRAM"

588 + nixl_memory_type = "VRAM"

589 elif self.kv_buffer_device == "cpu":

585 - self.nixl_memory_type = "DRAM"

586 - if self.nixl_memory_type is None:

590 + nixl_memory_type = "DRAM"

591 + if nixl_memory_type is None:

592 raise RuntimeError(

593 f"{self.device_type} with {self.kv_buffer_device} kv_buffer "

594 "is not supported."

595 )

596 + self.nixl_memory_type = nixl_memory_type

597

598 # Note: host xfer buffer ops when use_host_buffer is True

599 self.copy_blocks: CopyBlocksOp | None = None

vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py

+11 -8

@@ -75,9 +75,8 @@ class P2pNcclConnector(KVConnectorBase_V1):

75 super().__init__(vllm_config=vllm_config, role=role)

76 self._block_size = vllm_config.cache_config.block_size

77 self._requests_need_load: dict[str, Any] = {}

78 - self.config = vllm_config.kv_transfer_config

79 - self.is_producer = self.config.is_kv_producer

80 - self.chunked_prefill: dict[str, Any] = {}

78 + self.is_producer = self._kv_transfer_config.is_kv_producer

79 + self.chunked_prefill: dict[str, tuple[list[int], list[int] | None]] = {}

81 self._rank = get_world_group().rank if role == KVConnectorRole.WORKER else 0

82 self._local_rank = (

@@ -87,7 +86,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

86 self.p2p_nccl_engine = (

87 P2pNcclEngine(

88 local_rank=self._local_rank,

90 - config=self.config,

89 + config=self._kv_transfer_config,

90 hostname="",

91 port_offset=self._rank,

92 )

@@ -346,7 +345,8 @@ class P2pNcclConnector(KVConnectorBase_V1):

345 if self.is_producer:

346 return 0, False

347

349 - num_external_tokens = len(request.prompt_token_ids) - 1 - num_computed_tokens

348 + prompt_token_ids = request.prompt_token_ids or []

349 + num_external_tokens = len(prompt_token_ids) - 1 - num_computed_tokens

350

351 if num_external_tokens < 0:

352 num_external_tokens = 0

@@ -387,7 +387,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

387 ]

388 num_tokens = num_scheduled_tokens + new_req.num_computed_tokens

389 # the request's prompt is chunked prefill

390 - if num_tokens < len(new_req.prompt_token_ids):

390 + if num_tokens < len(new_req.prompt_token_ids or []):

391 # 'CachedRequestData' has no attribute 'prompt_token_ids'

392 self.chunked_prefill[new_req.req_id] = (

393 new_req.block_ids[0],

@@ -397,7 +397,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

397 # the request's prompt is not chunked prefill

398 meta.add_request(

399 request_id=new_req.req_id,

400 - token_ids=new_req.prompt_token_ids,

400 + token_ids=new_req.prompt_token_ids or [],

401 block_ids=new_req.block_ids[0],

402 block_size=self._block_size,

403 )

@@ -405,7 +405,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

405 if new_req.req_id in self._requests_need_load:

406 meta.add_request(

407 request_id=new_req.req_id,

408 - token_ids=new_req.prompt_token_ids,

408 + token_ids=new_req.prompt_token_ids or [],

409 block_ids=new_req.block_ids[0],

410 block_size=self._block_size,

411 )

@@ -421,10 +421,12 @@ class P2pNcclConnector(KVConnectorBase_V1):

421 num_scheduled_tokens = (scheduler_output.num_scheduled_tokens)[req_id]

422 num_tokens = num_scheduled_tokens + num_computed_tokens

423 assert req_id in self.chunked_prefill

424 + assert new_block_ids is not None

425 block_ids = new_block_ids[0]

426 if not resumed_from_preemption:

427 block_ids = self.chunked_prefill[req_id][0] + block_ids

428 prompt_token_ids = self.chunked_prefill[req_id][1]

429 + assert prompt_token_ids is not None

430 # the request's prompt is chunked prefill again

431 if num_tokens < len(prompt_token_ids):

432 self.chunked_prefill[req_id] = (block_ids, prompt_token_ids)

@@ -450,6 +452,7 @@ class P2pNcclConnector(KVConnectorBase_V1):

452

453 # NOTE(rob): For resumed req, new_block_ids is all

454 # of the block_ids for the request.

455 + assert new_block_ids is not None

456 block_ids = new_block_ids[0]

457

458 meta.add_request(

vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py

+25 -14

@@ -90,11 +90,10 @@ class SharedStorageConnector(KVConnectorBase_V1):

90 super().__init__(vllm_config=vllm_config, role=role)

91 self._block_size = vllm_config.cache_config.block_size

92 self._requests_need_load: dict[str, Request] = {}

93 - transfer_config = vllm_config.kv_transfer_config

94 - self._storage_path = transfer_config.get_from_extra_config(

93 + self._storage_path = self._kv_transfer_config.get_from_extra_config(

94 "shared_storage_path", "/tmp"

95 )

97 - logger.info(vllm_config.kv_transfer_config)

96 + logger.info(self._kv_transfer_config)

97 logger.info("Shared storage path is %s", self._storage_path)

99 def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:

@@ -277,9 +276,8 @@ class SharedStorageConnector(KVConnectorBase_V1):

276

277 # Now, first num_tokens_to_check tokens are hit, we need to prepare

278 # the metadata for the worker connector to correctly load the KV

280 - num_tokens_to_check = align_to_block_size(

281 - len(request.prompt_token_ids) - 1, self._block_size

282 - )

279 + token_ids = request.prompt_token_ids or []

280 + num_tokens_to_check = align_to_block_size(len(token_ids) - 1, self._block_size)

281

282 return num_tokens_to_check - num_computed_tokens, False

283

@@ -311,13 +309,15 @@ class SharedStorageConnector(KVConnectorBase_V1):

309

310 total_need_load = 0

311 for new_req in scheduler_output.scheduled_new_reqs:

312 + token_ids = new_req.prompt_token_ids or []

313 + mm_hashes = [f.identifier for f in new_req.mm_features]

314 if new_req.req_id in self._requests_need_load:

315 meta.add_request(

316 - token_ids=new_req.prompt_token_ids,

316 + token_ids=token_ids,

317 block_ids=new_req.block_ids[0],

318 block_size=self._block_size,

319 is_store=False,

320 - mm_hashes=[f.identifier for f in new_req.mm_features],

320 + mm_hashes=mm_hashes,

321 )

322 total_need_load += 1

323 else:

@@ -325,13 +325,13 @@ class SharedStorageConnector(KVConnectorBase_V1):

325 # but a single request can have both store and load.

326 # NOTE(rob): for this debug implementation, we only cache

327 # the original prompt tokens.

328 - if not self._found_match_for_request(new_req):

328 + if not self._found_match_for_prompt(token_ids, mm_hashes):

329 meta.add_request(

330 - token_ids=new_req.prompt_token_ids,

330 + token_ids=token_ids,

331 block_ids=new_req.block_ids[0],

332 block_size=self._block_size,

333 is_store=True,

334 - mm_hashes=[f.identifier for f in new_req.mm_features],

334 + mm_hashes=mm_hashes,

335 )

336

337 cached_reqs = scheduler_output.scheduled_cached_reqs

@@ -355,6 +355,7 @@ class SharedStorageConnector(KVConnectorBase_V1):

355

356 # NOTE(rob): For resumed req, new_block_ids is all

357 # of the block_ids for the request.

358 + assert new_block_ids is not None

359 block_ids = new_block_ids[0]

360

361 meta.add_request(

@@ -379,12 +380,22 @@ class SharedStorageConnector(KVConnectorBase_V1):

380 request: "Request",

381 ) -> bool:

382 """Check if the cache is hit for the request."""

383 + return self._found_match_for_prompt(

384 + list(request.prompt_token_ids or []),

385 + [f.identifier for f in request.mm_features],

386 + )

387 +

388 + def _found_match_for_prompt(

389 + self,

390 + prompt_token_ids: list[int],

391 + mm_hashes: list[str],

392 + ) -> bool:

393 num_tokens_to_check = align_to_block_size(

383 - len(request.prompt_token_ids) - 1, self._block_size

394 + len(prompt_token_ids) - 1, self._block_size

395 )

396 foldername = self._generate_foldername_debug(

386 - torch.tensor(request.prompt_token_ids)[:num_tokens_to_check],

387 - [f.identifier for f in request.mm_features],

397 + torch.tensor(prompt_token_ids)[:num_tokens_to_check],

398 + mm_hashes,

399 create_folder=False,

400 )

401 return os.path.exists(foldername)

vllm/distributed/kv_transfer/kv_pipe/mooncake_pipe.py

@@ -236,6 +236,7 @@ class MooncakePipe(KVPipeBase):

236 self.config = config

237 self.local_rank = local_rank

238 self.kv_rank = self.config.kv_rank

239 + assert self.kv_rank is not None

240 if device is None:

241 self.device = self._select_device(self.config.kv_buffer_device)

242 else:

vllm/distributed/kv_transfer/kv_pipe/pynccl_pipe.py

@@ -53,6 +53,7 @@ class PyNcclPipe(KVPipeBase):

53 self.config = config

54 self.local_rank = local_rank

55 self.kv_rank = self.config.kv_rank

56 + assert self.kv_rank is not None

57 self.kv_parallel_size = self.config.kv_parallel_size

58 if device is None:

59 self.device = self._select_device(self.config.kv_buffer_device)