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 5, 17:15 UTC
[Refactor] Remove dead code in kv connector and model runner (#38383)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Mar 31, 2026 at 17:05 UTC
856589ed9aa7fbac86af8e86374e4f00ed601ecd
7 files changed
+16
-54
tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py
+2
-1
@@ -6,6 +6,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( #
6
)
7
from vllm.distributed.kv_transfer.kv_transfer_state import (
8
ensure_kv_transfer_initialized,
9
+ ensure_kv_transfer_shutdown,
10
get_kv_transfer_group,
11
)
12
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
@@ -57,4 +58,4 @@ def test_kv_connector_mixin_clears_metadata():
58
assert connector.call_record.get("clear_connector_metadata", 0) == 1
59
finally:
60
# Ensure we clean up the global connector between tests
60
- KVConnectorModelRunnerMixin.ensure_kv_transfer_shutdown()
61
+ ensure_kv_transfer_shutdown()
vllm/v1/worker/cpu_model_runner.py
-4
@@ -102,10 +102,6 @@ class CPUModelRunner(GPUModelRunner):
102
# so stale KV cache data never affects computation.
103
pass
104
105
- def get_dp_padding(self, num_tokens: int) -> tuple[int, torch.Tensor | None]:
106
- # Note: For CPU backend, dp padding is not required for now.
107
- return 0, None
108
-
105
106
@contextmanager
107
def _torch_cuda_wrapper():
vllm/v1/worker/ec_connector_model_runner_mixin.py
-8
@@ -34,14 +34,6 @@ class ECConnectorModelRunnerMixin:
34
connector = get_ec_transfer()
35
connector.save_caches(encoder_cache=encoder_cache, mm_hash=mm_hash)
36
37
- @staticmethod
38
- def get_finished_ec_transfers(
39
- scheduler_output: "SchedulerOutput",
40
- ) -> tuple[set[str] | None, set[str] | None]:
41
- if has_ec_transfer():
42
- return get_ec_transfer().get_finished(scheduler_output.finished_req_ids)
43
- return None, None
44
-
37
@staticmethod
38
def maybe_get_ec_connector_output(
39
scheduler_output: "SchedulerOutput",
vllm/v1/worker/gpu/kv_connector.py
-5
@@ -97,11 +97,6 @@ class ActiveKVConnector(KVConnector):
97
self.kv_connector.clear_connector_metadata()
98
return output
99
100
- def clear_metadata(self) -> None:
101
- """Clear the connector metadata. Call this after draft model runs."""
102
- if not self._disabled:
103
- self.kv_connector.clear_connector_metadata()
104
-
100
def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
101
if self._disabled:
102
return EMPTY_MODEL_RUNNER_OUTPUT
vllm/v1/worker/gpu_model_runner.py
-2
@@ -451,7 +451,6 @@ class GPUModelRunner(
451
# Model-related.
452
self.num_query_heads = model_config.get_num_attention_heads(parallel_config)
453
self.inputs_embeds_size = model_config.get_inputs_embeds_size()
454
- self.attention_chunk_size = model_config.attention_chunk_size
454
# Only relevant for models using ALiBi (e.g, MPT)
455
self.use_alibi = model_config.uses_alibi
456
@@ -594,7 +593,6 @@ class GPUModelRunner(
593
# NOTE(rob): num_prompt_logprobs only includes reqs
594
# that are currently in the prefill phase.
595
self.num_prompt_logprobs: dict[str, int] = {}
597
- self.comm_stream = torch.cuda.Stream()
596
597
# Input Batch
598
# NOTE(Chen): Ideally, we should initialize the input batch inside
vllm/v1/worker/kv_connector_model_runner_mixin.py
+1
-11
@@ -13,11 +13,7 @@ import torch
13
14
from vllm.config import VllmConfig
15
from vllm.config.cache import CacheDType
16
-from vllm.distributed.kv_transfer import (
17
- ensure_kv_transfer_shutdown,
18
- get_kv_transfer_group,
19
- has_kv_transfer_group,
20
-)
16
+from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
17
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
18
from vllm.forward_context import get_forward_context, set_forward_context
19
from vllm.logger import init_logger
@@ -38,12 +34,6 @@ logger = init_logger(__name__)
34
35
# Defined as a kv connector functionality mixin for ModelRunner (GPU, TPU)
36
class KVConnectorModelRunnerMixin:
41
- @staticmethod
42
- def ensure_kv_transfer_shutdown() -> None:
43
- # has_kv_transfer_group can be None during interpreter shutdown.
44
- if has_kv_transfer_group and has_kv_transfer_group(): # type: ignore[truthy-function]
45
- ensure_kv_transfer_shutdown()
46
-
37
@staticmethod
38
def kv_connector_no_forward(
39
scheduler_output: "SchedulerOutput", vllm_config: VllmConfig
vllm/v1/worker/xpu_model_runner.py
+13
-23
@@ -1,23 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from contextlib import contextmanager
4
-from typing import TYPE_CHECKING
4
5
import torch
6
7
from vllm.config import VllmConfig
9
-from vllm.logger import init_logger
8
from vllm.utils.torch_utils import supports_xpu_graph
9
from vllm.v1.worker.gpu.model_runner import (
10
GPUModelRunner as GPUModelRunnerV2,
11
)
12
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
13
16
-if TYPE_CHECKING:
17
- pass
18
-
19
-logger = init_logger(__name__)
20
-
14
15
class XPUModelRunner(GPUModelRunner):
16
"""A model runner for XPU devices."""
@@ -47,19 +40,16 @@ class XPUModelRunnerV2(GPUModelRunnerV2):
40
41
@contextmanager
42
def _torch_cuda_wrapper():
50
- try:
51
- # replace cuda APIs with xpu APIs, this should work by default
52
- torch.cuda.Stream = torch.xpu.Stream
53
- torch.cuda.default_stream = torch.xpu.current_stream
54
- torch.cuda.current_stream = torch.xpu.current_stream
55
- torch.cuda.stream = torch.xpu.stream
56
- torch.cuda.mem_get_info = torch.xpu.mem_get_info
57
- torch.cuda.Event = torch.Event
58
- torch.cuda.set_stream = torch.xpu.set_stream
59
- if supports_xpu_graph():
60
- torch.cuda.graph = torch.xpu.graph
61
- torch.cuda.CUDAGraph = torch.xpu.XPUGraph
62
- torch.cuda.graph_pool_handle = torch.xpu.graph_pool_handle
63
- yield
64
- finally:
65
- pass
43
+ # replace cuda APIs with xpu APIs, this should work by default
44
+ torch.cuda.Stream = torch.xpu.Stream
45
+ torch.cuda.default_stream = torch.xpu.current_stream
46
+ torch.cuda.current_stream = torch.xpu.current_stream
47
+ torch.cuda.stream = torch.xpu.stream
48
+ torch.cuda.mem_get_info = torch.xpu.mem_get_info
49
+ torch.cuda.Event = torch.Event
50
+ torch.cuda.set_stream = torch.xpu.set_stream
51
+ if supports_xpu_graph():
52
+ torch.cuda.graph = torch.xpu.graph
53
+ torch.cuda.CUDAGraph = torch.xpu.XPUGraph
54
+ torch.cuda.graph_pool_handle = torch.xpu.graph_pool_handle
55
+ yield