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

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, 23:45 UTC

Revert "[V1] Exception Handling when Loading KV Cache from Remote Store" (#21778)

Signed-off-by: KuntaiDu <kuntai@uchicago.edu>

Kuntai Du committed Jul 28, 2025 at 13:15 UTC b18b417fbf36b6e2747420d7ce86a685d54b4ff7

10 files changed +5 -229

tests/v1/kv_connector/kv_load_exception_handling/random_drop_connector.py deleted

-120

@@ -1,120 +0,0 @@

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

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

3 -import logging

4 -import random

5 -from dataclasses import dataclass

6 -from typing import TYPE_CHECKING

7 -

8 -import torch

9 -

10 -from vllm.config import VllmConfig

11 -from vllm.distributed.kv_transfer.kv_connector.v1.base import (

12 - KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)

13 -from vllm.v1.core.sched.output import SchedulerOutput

14 -

15 -if TYPE_CHECKING:

16 - from vllm.attention.backends.abstract import AttentionMetadata

17 - from vllm.v1.core.kv_cache_manager import KVCacheBlocks

18 - from vllm.v1.request import Request

19 -

20 -logger = logging.getLogger()

21 -logging.basicConfig(level=logging.INFO)

22 -

23 -

24 -@dataclass

25 -class RandomDropConnectorMetadata(KVConnectorMetadata):

26 - req_meta: dict[str, list[int]]

27 -

28 -

29 -class RandomDropConnector(KVConnectorBase_V1):

30 - """

31 - A connector designed for fault tolerance testing by randomly dropping

32 - kv data during the process of loading or receiving KV cache.

33 -

34 - This class simulates real-world scenarios where requests or data

35 - might be lost or timeout, allowing developers to test and validate the

36 - system's ability to handle such failures.

37 -

38 - Attributes:

39 - finished_recving_kv_req_ids (set[str]): A set of request IDs that

40 - have completed receiving KV cache data.

41 - finished_loading_dict (dict[str, int]): A dictionary that tracks

42 - the actual number of tokens loaded from the remote KV store

43 - for each completed request. The keys are request IDs, and

44 - the values are the corresponding token counts.

45 - """

46 -

47 - def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):

48 - super().__init__(vllm_config=vllm_config, role=role)

49 -

50 - self.failure_request: list[str] = []

51 - self._reqs_need_recv: dict[str, list[int]] = {}

52 - self._finish_load: dict[str, int] = {}

53 -

54 - self.chunk_size = 256

55 -

56 - ############################################################

57 - # Scheduler Side Methods

58 - ############################################################

59 -

60 - def get_num_new_matched_tokens(

61 - self, request: "Request",

62 - num_computed_tokens: int) -> tuple[int, bool]:

63 - if request.request_id in self.failure_request:

64 - self.failure_request.remove(request.request_id)

65 - return 0, False

66 - num_external_hit_tokens = request.num_prompt_tokens - 1

67 - logger.info(

68 - "request %s num_prompt_tokens %d num_external_hit_tokens %d",

69 - request.request_id, request.num_prompt_tokens,

70 - num_external_hit_tokens)

71 - return num_external_hit_tokens, True

72 -

73 - def update_state_after_alloc(self, request: "Request",

74 - blocks: "KVCacheBlocks",

75 - num_external_tokens: int):

76 - if num_external_tokens > 0:

77 - self._reqs_need_recv[

78 - request.

79 - request_id] = request.prompt_token_ids[:num_external_tokens]

80 -

81 - def build_connector_meta(

82 - self,

83 - scheduler_output: SchedulerOutput,

84 - ) -> KVConnectorMetadata:

85 - req_meta = self._reqs_need_recv.copy()

86 - self._reqs_need_recv.clear()

87 - return RandomDropConnectorMetadata(req_meta)

88 -

89 - def add_failure_request(self, request: "Request"):

90 - self.failure_request.append(request.request_id)

91 -

92 - def start_load_kv(self, forward_context, **kwargs) -> None:

93 - for request_id, hit_tokens in self._get_connector_metadata(

94 - ).req_meta.items():

95 - num_actual_load_tokens = self.load_kv(request_id, hit_tokens)

96 - logger.info("request %s hit_tokens %d num_actual_load_tokens %d",

97 - request_id, len(hit_tokens), num_actual_load_tokens)

98 - self._finish_load[request_id] = num_actual_load_tokens

99 -

100 - def wait_for_layer_load(self, layer_name: str) -> None:

101 - pass

102 -

103 - def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor,

104 - attn_metadata: "AttentionMetadata", **kwargs) -> None:

105 - pass

106 -

107 - def wait_for_save(self):

108 - pass

109 -

110 - def load_kv(self, request_id, hit_tokens):

111 - num_actual_load_tokens = random.randint(0, len(hit_tokens))

112 - return num_actual_load_tokens

113 -

114 - def get_finished_loading(self) -> dict[str, int]:

115 - if not self._finish_load:

116 - return {}

117 - finished_loading = self._finish_load.copy()

118 - self._finish_load.clear()

119 -

120 - return finished_loading

tests/v1/kv_connector/kv_load_exception_handling/test.sh deleted

-16

@@ -1,16 +0,0 @@

1 -#!/bin/bash

2 -

3 -SCRIPT_DIR=$(dirname "$(readlink -f "$0")")

4 -export PYTHONPATH=$PYTHONPATH:$SCRIPT_DIR

5 -

6 -vllm serve DeepSeek-V2-Lite-Chat \

7 ---trust-remote-code \

8 ---served-model-name vllm_cpu_offload \

9 ---max-model-len 32768 \

10 ---no-enable-prefix-caching \

11 ---max-seq-len-to-capture 10000 \

12 ---max-num-seqs 64 \

13 ---gpu-memory-utilization 0.9 \

14 ---host 0.0.0.0 \

15 --tp 2 \

16 ---kv-transfer-config '{"kv_connector":"RandomDropConnector","kv_role":"kv_both","kv_connector_module_path":"random_drop_connector"}'

vllm/distributed/kv_transfer/kv_connector/utils.py

+1 -15

@@ -139,27 +139,13 @@ class KVOutputAggregator:

139 finished_set.add(req_id)

140 del remaining_count_dict[req_id]

141

142 - def update_finished_load_dict(worker_finished_loading_dict: dict[str,

143 - int],

144 - finished_loading_dict: dict[str, int]):

145 - for req_id, num_actual_load_tokens in (worker_finished_loading_dict

146 - or {}).items():

147 - if req_id in finished_loading_dict:

148 - finished_loading_dict[req_id] = min(

149 - finished_loading_dict[req_id], num_actual_load_tokens)

150 - else:

151 - finished_loading_dict[req_id] = num_actual_load_tokens

152 -

142 finished_sending = set[str]()

143 finished_recving = set[str]()

155 - finished_loading_dict: dict[str, int] = {}

144 for output in outputs:

145 update_finished_set(output.finished_sending,

146 self._send_remaining_count, finished_sending)

147 update_finished_set(output.finished_recving,

148 self._recv_remaining_count, finished_recving)

161 - update_finished_load_dict(output.finished_loading_dict,

162 - finished_loading_dict)

149

150 # select output of the worker specified by output_rank

151 output = outputs[output_rank]

@@ -171,7 +157,7 @@ class KVOutputAggregator:

157 # send/recv

158 output.finished_sending = finished_sending if finished_sending else None

159 output.finished_recving = finished_recving if finished_recving else None

174 - output.finished_loading_dict = finished_loading_dict or None

160 +

161 return output

162

163 def async_aggregate(self,

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

-20

@@ -28,9 +28,6 @@ The class provides the following primitives:

29 get_finished() - called with ids of finished requests, returns

30 ids of requests that have completed async sending/recving.

31 - get_finished_loading() - called with scheduler outputs, returns

32 - a dictionary that the keys are request IDs and the values are

33 - the actual number of tokens loaded from the remote KV cache

31 """

33 import enum

@@ -222,23 +219,6 @@ class KVConnectorBase_V1(ABC):

219 """

220 return None, None

221

225 - def get_finished_loading(

226 - self, scheduler_output: "SchedulerOutput") -> dict[str, int]:

227 - """

228 - Retrieves the actual number of tokens loaded for requests that have

229 - completed the asynchronous loading process from the remote KV cache.

230 -

231 - This function is used by the scheduler process (via the Executors)

232 - to track the progress of requests and determine which requests have

233 - successfully finished loading their KV cache data.

234 -

235 - Returns:

236 - A dictionary where the keys are request IDs and the values are the

237 - corresponding number of tokens that have been successfully loaded

238 - for each request.

239 - """

240 - return {}

241 -

222 # ==============================

223 # Scheduler-side methods

224 # ==============================

vllm/sequence.py

-2

@@ -1167,8 +1167,6 @@ class IntermediateTensors:

1167 # [req_ids]

1168 finished_sending: Optional[set[str]] = None

1169 finished_recving: Optional[set[str]] = None

1170 - #req_id -> num_actual_load_tokens

1171 - finished_loading_dict: Optional[dict[str, int]] = None

1170

1171 def __init__(self, tensors):

1172 # manually define this function, so that

vllm/v1/core/sched/scheduler.py

-30

@@ -118,9 +118,6 @@ class Scheduler(SchedulerInterface):

118

119 # KV Connector: requests in process of async KV loading or recving

120 self.finished_recving_kv_req_ids: set[str] = set()

121 - # The keys are request IDs, and the values are corresponding token

122 - # count that have been successfully loaded from the remote KV store

123 - self.finished_loading_dict: dict[str, int] = {}

121

122 # Encoder-related.

123 # Calculate encoder cache size if applicable

@@ -1097,27 +1094,6 @@ class Scheduler(SchedulerInterface):

1094 (block_ids, ) = self.kv_cache_manager.get_block_ids(request.request_id)

1095 return self.connector.request_finished(request, block_ids)

1096

1100 - def _update_actual_load_token_num_from_remote_kv(self,

1101 - request: Request) -> bool:

1102 -

1103 - num_actual_load_tokens = self.finished_loading_dict.pop(

1104 - request.request_id)

1105 - num_computed_tokens = num_actual_load_tokens

1106 - assert self.connector is not None

1107 - if num_actual_load_tokens <= 0 and hasattr(self.connector,

1108 - "add_failure_request"):

1109 - self.connector.add_failure_request(request)

1110 - return True

1111 -

1112 - if num_actual_load_tokens == request.num_tokens:

1113 - num_computed_tokens -= 1

1114 -

1115 - self.kv_cache_manager.cache_blocks(request, num_computed_tokens)

1116 -

1117 - # Update the request state for scheduling.

1118 - request.num_computed_tokens = num_computed_tokens

1119 - return True

1120 -

1097 def _update_waiting_for_remote_kv(self, request: Request) -> bool:

1098 """

1099 KV Connector: check if the request_id is finished_recving.

@@ -1131,9 +1107,6 @@ class Scheduler(SchedulerInterface):

1107 WAITING_FOR_REMOTE_KV.

1108 """

1109 assert self.connector is not None

1134 - if request.request_id in self.finished_loading_dict:

1135 - return self._update_actual_load_token_num_from_remote_kv(request)

1136 -

1110 if request.request_id not in self.finished_recving_kv_req_ids:

1111 return False

1112

@@ -1172,6 +1145,3 @@ class Scheduler(SchedulerInterface):

1145 for req_id in (model_runner_output.finished_sending or ()):

1146 logger.debug("Finished sending KV transfer for request %s", req_id)

1147 self._free_blocks(self.requests[req_id])

1175 - if model_runner_output.finished_loading_dict:

1176 - self.finished_loading_dict.update(

1177 - model_runner_output.finished_loading_dict)

vllm/v1/outputs.py

-3

@@ -107,8 +107,6 @@ class ModelRunnerOutput:

107 # [req_ids]

108 finished_sending: Optional[set[str]] = None

109 finished_recving: Optional[set[str]] = None

110 - # req_id -> actual_load_token from connector

111 - finished_loading_dict: Optional[dict[str, int]] = None

110

111 # req_id -> num_nans_in_logits

112 num_nans_in_logits: Optional[dict[str, int]] = None

@@ -123,5 +121,4 @@ EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[],

121 pooler_output=[],

122 finished_sending=None,

123 finished_recving=None,

126 - finished_loading_dict=None,

124 num_nans_in_logits=None)

vllm/v1/worker/gpu_model_runner.py

+2 -8

@@ -1375,7 +1375,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1375 num_scheduled_tokens_np: np.ndarray,

1376 finished_sending: Optional[set[str]],

1377 finished_recving: Optional[set[str]],

1378 - finished_loading_dict: Optional[dict[str, int]],

1378 ) -> ModelRunnerOutput:

1379 assert self.input_batch.num_reqs ==\

1380 len(self.input_batch.pooling_params), \

@@ -1412,7 +1411,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1411 pooler_output=pooler_output,

1412 finished_sending=finished_sending,

1413 finished_recving=finished_recving,

1415 - finished_loading_dict=finished_loading_dict,

1414 )

1415

1416 @torch.inference_mode()

@@ -1532,7 +1530,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1530 self.maybe_wait_for_kv_save()

1531 finished_sending, finished_recving = (

1532 self.get_finished_kv_transfers(scheduler_output))

1535 - finished_loading_dict = self.get_finished_loading(scheduler_output)

1533

1534 if self.use_aux_hidden_state_outputs:

1535 hidden_states, aux_hidden_states = model_output

@@ -1550,11 +1547,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1547 if not get_pp_group().is_last_rank:

1548 # For mid-pipeline stages, return the hidden states.

1549 if not broadcast_pp_output:

1553 - if (finished_sending or finished_recving

1554 - or finished_loading_dict):

1550 + if finished_sending or finished_recving:

1551 hidden_states.finished_sending = finished_sending

1552 hidden_states.finished_recving = finished_recving

1557 - hidden_states.finished_loading_dict = finished_loading_dict

1553 return hidden_states

1554 assert isinstance(hidden_states, IntermediateTensors)

1555 get_pp_group().send_tensor_dict(hidden_states.tensors,

@@ -1564,7 +1559,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1559 if self.input_batch.pooling_params:

1560 return self._pool(hidden_states, num_scheduled_tokens,

1561 num_scheduled_tokens_np, finished_sending,

1567 - finished_recving, finished_loading_dict)

1562 + finished_recving)

1563

1564 sample_hidden_states = hidden_states[logits_indices]

1565 logits = self.model.compute_logits(sample_hidden_states, None)

@@ -1716,7 +1711,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1711 pooler_output=[],

1712 finished_sending=finished_sending,

1713 finished_recving=finished_recving,

1719 - finished_loading_dict=finished_loading_dict,

1714 num_nans_in_logits=num_nans_in_logits,

1715 )

1716

vllm/v1/worker/gpu_worker.py

+1 -3

@@ -359,12 +359,10 @@ class Worker(WorkerBase):

359 # In case of PP with kv transfer, we need to pass through the

360 # finished_sending and finished_recving buffers.

361 new_output = EMPTY_MODEL_RUNNER_OUTPUT

362 - if (output.finished_sending or output.finished_recving

363 - or output.finished_loading_dict):

362 + if output.finished_sending or output.finished_recving:

363 new_output = copy.copy(new_output)

364 new_output.finished_sending = output.finished_sending

365 new_output.finished_recving = output.finished_recving

367 - new_output.finished_loading_dict = output.finished_loading_dict

366 output = new_output

367

368 assert isinstance(output, ModelRunnerOutput)

vllm/v1/worker/kv_connector_model_runner_mixin.py

+1 -12

@@ -53,14 +53,6 @@ class KVConnectorModelRunnerMixin:

53 scheduler_output.finished_req_ids)

54 return None, None

56 - @staticmethod

57 - def get_finished_loading(

58 - scheduler_output: "SchedulerOutput", ) -> dict[str, int]:

59 - if has_kv_transfer_group():

60 - return get_kv_transfer_group().get_finished_loading(

61 - scheduler_output)

62 - return {}

63 -

56 def kv_connector_no_forward(self, scheduler_output: "SchedulerOutput",

57 vllm_config: VllmConfig) -> ModelRunnerOutput:

58 # KV send/recv even if no work to do.

@@ -68,14 +60,11 @@ class KVConnectorModelRunnerMixin:

60 self.maybe_setup_kv_connector(scheduler_output)

61 finished_sending, finished_recving = (

62 self.get_finished_kv_transfers(scheduler_output))

71 - finished_loading_dict = self.get_finished_loading(scheduler_output)

73 - if (not finished_sending and not finished_recving

74 - and not finished_loading_dict):

64 + if not finished_sending and not finished_recving:

65 return EMPTY_MODEL_RUNNER_OUTPUT

67 output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)

68 output.finished_sending = finished_sending

69 output.finished_recving = finished_recving

80 - output.finished_loading_dict = finished_loading_dict

70 return output