@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f4f4e7ef27
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][P/D] Deprecate v0 `KVConnectorBase` code (1/2) (#21785)

Signed-off-by: Linkun Chen <github@lkchen.net>

lkchen committed Aug 4, 2025 at 19:11 UTC f4f4e7ef273645192fac837718b3fdcf073c597a
13 files changed +31 -1040
.buildkite/test-pipeline.yaml
-1
@@ -749,7 +749,6 @@ steps:
749 # this test fails consistently.
750 # TODO: investigate and fix
751 - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
752 - - VLLM_USE_V1=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s kv_transfer/test_disagg.py
752 - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
753 - pytest -v -s models/multimodal/generation/test_maverick.py
754
tests/kv_transfer/test_disagg.py deleted
-120
@@ -1,120 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import os
5 -import subprocess
6 -import sys
7 -import time
8 -from subprocess import Popen
9 -
10 -import pytest
11 -import requests
12 -import torch
13 -
14 -
15 -# Fixture to set up environment variables and teardown servers after tests
16 -@pytest.fixture(scope="module", autouse=True)
17 -def setup_servers():
18 - if torch.cuda.device_count() < 2:
19 - pytest.skip("Skipping test: fewer than 2 GPUs available")
20 -
21 - # Set up environment variables
22 - VLLM_HOST_IP = subprocess.check_output("hostname -I | awk '{print $1}'",
23 - shell=True).decode().strip()
24 - os.environ["VLLM_HOST_IP"] = VLLM_HOST_IP
25 -
26 - # Start prefill instance
27 - prefill_cmd = [
28 - sys.executable,
29 - "-m",
30 - "vllm.entrypoints.openai.api_server",
31 - "--model",
32 - "meta-llama/Llama-3.2-1B-Instruct",
33 - "--port",
34 - "8100",
35 - "--gpu-memory-utilization",
36 - "0.5",
37 - "--max-model-len",
38 - "1000",
39 - "--kv-transfer-config",
40 - '{"kv_connector":"PyNcclConnector","kv_role":"kv_producer",'\
41 - '"kv_rank":0,"kv_parallel_size":2}',
42 - ]
43 - prefill_env = os.environ.copy()
44 - prefill_env["CUDA_VISIBLE_DEVICES"] = "0"
45 - prefill_proc = Popen(prefill_cmd, env=prefill_env)
46 -
47 - # Start decode instance
48 - decode_cmd = [
49 - sys.executable,
50 - "-m",
51 - "vllm.entrypoints.openai.api_server",
52 - "--model",
53 - "meta-llama/Llama-3.2-1B-Instruct",
54 - "--port",
55 - "8200",
56 - "--gpu-memory-utilization",
57 - "0.5",
58 - "--max-model-len",
59 - "1000",
60 - "--kv-transfer-config",
61 - '{"kv_connector":"PyNcclConnector","kv_role":"kv_consumer",'\
62 - '"kv_rank":1,"kv_parallel_size":2}',
63 - ]
64 - decode_env = os.environ.copy()
65 - decode_env["CUDA_VISIBLE_DEVICES"] = "1"
66 - decode_proc = Popen(decode_cmd, env=decode_env)
67 -
68 - # Wait for servers to be ready
69 - assert wait_for_server(8100), "Prefill server did not start in time"
70 - assert wait_for_server(8200), "Decode server did not start in time"
71 -
72 - # Yield to the test function and handle teardown after tests
73 - yield
74 -
75 - # Cleanup: kill the processes
76 - prefill_proc.terminate()
77 - decode_proc.terminate()
78 -
79 - # Additional cleanup if needed
80 - prefill_proc.wait()
81 - decode_proc.wait()
82 -
83 -
84 -# Helper function to wait for server
85 -def wait_for_server(port, timeout=240):
86 - start_time = time.time()
87 - while time.time() - start_time < timeout:
88 - try:
89 - response = requests.get(f"http://localhost:{port}/v1/completions")
90 - if response.status_code in [200, 405]:
91 - return True
92 - except requests.ConnectionError:
93 - time.sleep(1)
94 - return False
95 -
96 -
97 -# Test function to send curl requests and validate responses
98 -@pytest.mark.parametrize("prompt", ["San Francisco is a", "Santa Clara is a"])
99 -def test_disaggregated_prefilling(prompt):
100 - # Send to prefill
101 - response = requests.post("http://localhost:8100/v1/completions",
102 - headers={"Content-Type": "application/json"},
103 - json={
104 - "model": "meta-llama/Llama-3.2-1B-Instruct",
105 - "prompt": prompt,
106 - "max_tokens": 1,
107 - "temperature": 0
108 - })
109 - assert response.status_code == 200
110 -
111 - # Send to decode
112 - response = requests.post("http://localhost:8200/v1/completions",
113 - headers={"Content-Type": "application/json"},
114 - json={
115 - "model": "meta-llama/Llama-3.2-1B-Instruct",
116 - "prompt": prompt,
117 - "max_tokens": 10,
118 - "temperature": 0
119 - })
120 - assert response.status_code == 200
vllm/distributed/kv_transfer/kv_connector/base.py
+4 -136
@@ -1,142 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -KVConnectorBase Class for Distributed KV Cache & Hidden State communication
5 -
6 -The class provides two primary abstract methods:
7 -1. send_kv_caches_and_hidden_states(): Send KV caches and hidden states
8 -2. recv_kv_caches_and_hidden_states(): Recv KV caches and hidden states
9 -"""
10 -
11 -from abc import ABC, abstractmethod
12 -from typing import TYPE_CHECKING, Optional, Union
13 -
14 -import torch
3 +"""Defines the base type for KV cache connectors."""
4
5 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
17 -from vllm.sequence import IntermediateTensors
18 -
19 -if TYPE_CHECKING:
20 - from vllm.config import VllmConfig
21 - from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
22 -
23 -
24 -class KVConnectorBase(ABC):
25 - """
26 - Abstract base class for a KV connector.
27 -
28 - The class provides two primary abstract methods:
29 - 1. send_kv_caches_and_hidden_states(): Send KV caches and hidden states
30 - 2. recv_kv_caches_and_hidden_states(): Recv KV caches and hidden states
31 - """
32 -
33 - @abstractmethod
34 - def __init__(
35 - self,
36 - rank: int,
37 - local_rank: int,
38 - config: "VllmConfig",
39 - ):
40 - raise NotImplementedError
41 -
42 - @abstractmethod
43 - def close(self) -> None:
44 - """Close the buffer and release resources.
45 -
46 - This method is responsible for cleaning up resources related to the
47 - connector when it is no longer needed.
48 -
49 - Raises:
50 - NotImplementedError: This method must be implemented in subclasses.
51 - """
52 - raise NotImplementedError
53 -
54 - @abstractmethod
55 - def send_kv_caches_and_hidden_states(
56 - self,
57 - model_executable: torch.nn.Module,
58 - model_input: "ModelInputForGPUWithSamplingMetadata",
59 - kv_caches: list[torch.Tensor],
60 - hidden_or_intermediate_states: Union[torch.Tensor,
61 - IntermediateTensors],
62 - ) -> None:
63 - """
64 - Send KV caches and hidden states to the connector.
65 -
66 - This method processes the input tokens, KV caches, and
67 - hidden/intermediate states for a given model and sends the data to the
68 - decode instance.
69 -
70 - Args:
71 - model_executable (torch.nn.Module): The model executable containing
72 - start and end layer information.
73 - model_input (ModelInputForGPUWithSamplingMetadata): The input
74 - metadata from vLLM.
75 - kv_caches (list[torch.Tensor]): List of KV caches (keys and values)
76 - for each layer.
77 - hidden_or_intermediate_states (Union[torch.Tensor,
78 - IntermediateTensors]):
79 - The hidden or intermediate states associated with the tokens.
80 -
81 - Returns:
82 - None
83 -
84 - """
85 -
86 - raise NotImplementedError
87 -
88 - @abstractmethod
89 - def recv_kv_caches_and_hidden_states(
90 - self, model_executable: torch.nn.Module,
91 - model_input: "ModelInputForGPUWithSamplingMetadata",
92 - kv_caches: list[torch.Tensor]
93 - ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,
94 - "ModelInputForGPUWithSamplingMetadata"]:
95 - """
96 - Receive KV caches and hidden states from the connector.
97 -
98 - This method attempts to retrieve KV caches and hidden states for input
99 - tokens. If all required KV caches and hidden states are received, it
100 - will bypass model input, else it will fall back to normal vLLM model
101 - forwarding.
102 -
103 - Args:
104 - model_executable (torch.nn.Module):
105 - The model executable from vLLM modelrunner.
106 - model_input (ModelInputForGPUWithSamplingMetadata):
107 - The model input from vLLM modelrunner.
108 - kv_caches (list[torch.Tensor]):
109 - List of KV caches for each layer.
110 -
111 - Returns:
112 - - hidden_or_intermediate_states (torch.Tensor or
113 - IntermediateTensors):
114 - Concatenated hidden states if all required data is retrieved,
115 - otherwise `None`.
116 - - bypass_model_exec (bool):
117 - Indicates whether the model execution can be skipped (True) or
118 - needs to be redone (False).
119 - - model_input (ModelInputForGPUWithSamplingMetadata):
120 - Optionally adjusted input metadata for re-execution when
121 - `bypass_model_exec=False`.
122 -
123 - """
124 -
125 - raise NotImplementedError
126 -
127 - @classmethod
128 - def get_required_kvcache_layout(
129 - cls, vllm_config: "VllmConfig") -> Optional[str]:
130 - """
131 - Get the required KV cache layout for this connector.
132 - Args:
133 - vllm_config (VllmConfig): the vllm config.
134 -
135 - Returns:
136 - str: the required KV cache layout. e.g. HND, or NHD.
137 - None if the connector does not require a specific layout.
138 - """
139 - return None
6
7 +KVConnectorBase = KVConnectorBase_V1
8 +KVConnectorBaseType = KVConnectorBase_V1
9
142 -KVConnectorBaseType = Union[KVConnectorBase, KVConnectorBase_V1]
10 +__all__ = ["KVConnectorBase", "KVConnectorBaseType"]
vllm/distributed/kv_transfer/kv_connector/factory.py
+13 -55
@@ -5,14 +5,10 @@ import importlib
5 from typing import TYPE_CHECKING, Callable
6
7 import vllm.envs as envs
8 -from vllm.config import KVTransferConfig
9 -from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType
10 -from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
11 - KVConnectorRole)
8 +from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
9 +from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
10 from vllm.logger import init_logger
11
14 -from .base import KVConnectorBase
15 -
12 if TYPE_CHECKING:
13 from vllm.config import VllmConfig
14
@@ -20,7 +16,7 @@ logger = init_logger(__name__)
16
17
18 class KVConnectorFactory:
23 - _registry: dict[str, Callable[[], type[KVConnectorBaseType]]] = {}
19 + _registry: dict[str, Callable[[], type[KVConnectorBase]]] = {}
20
21 @classmethod
22 def register_connector(cls, name: str, module_path: str,
@@ -29,28 +25,23 @@ class KVConnectorFactory:
25 if name in cls._registry:
26 raise ValueError(f"Connector '{name}' is already registered.")
27
32 - def loader() -> type[KVConnectorBaseType]:
28 + def loader() -> type[KVConnectorBase]:
29 module = importlib.import_module(module_path)
30 return getattr(module, class_name)
31
32 cls._registry[name] = loader
33
34 @classmethod
39 - def create_connector_v0(cls, rank: int, local_rank: int,
40 - config: "VllmConfig") -> KVConnectorBase:
41 - if envs.VLLM_USE_V1:
42 - raise ValueError("Attempting to initialize a V0 Connector, "
35 + def create_connector(
36 + cls,
37 + config: "VllmConfig",
38 + role: KVConnectorRole,
39 + ) -> KVConnectorBase:
40 + if not envs.VLLM_USE_V1:
41 + raise ValueError("Attempting to initialize a V1 Connector, "
42 f"but found {envs.VLLM_USE_V1=}")
43
45 - connector_cls = cls.get_connector_class(config.kv_transfer_config)
46 - assert issubclass(connector_cls, KVConnectorBase)
47 - return connector_cls(rank, local_rank, config)
48 -
49 - @classmethod
50 - def get_connector_class(
51 - cls, kv_transfer_config: "KVTransferConfig"
52 - ) -> type[KVConnectorBaseType]:
53 - """Get the connector class by name."""
44 + kv_transfer_config = config.kv_transfer_config
45 connector_name = kv_transfer_config.kv_connector
46 if connector_name in cls._registry:
47 connector_cls = cls._registry[connector_name]()
@@ -61,21 +52,7 @@ class KVConnectorFactory:
52 f"Unsupported connector type: {connector_name}")
53 connector_module = importlib.import_module(connector_module_path)
54 connector_cls = getattr(connector_module, connector_name)
64 - return connector_cls
65 -
66 - @classmethod
67 - def create_connector_v1(
68 - cls,
69 - config: "VllmConfig",
70 - role: KVConnectorRole,
71 - ) -> KVConnectorBase_V1:
72 - if not envs.VLLM_USE_V1:
73 - raise ValueError("Attempting to initialize a V1 Connector, "
74 - f"but found {envs.VLLM_USE_V1=}")
75 -
76 - kv_transfer_config = config.kv_transfer_config
77 - connector_cls = cls.get_connector_class(kv_transfer_config)
78 - assert issubclass(connector_cls, KVConnectorBase_V1)
55 + assert issubclass(connector_cls, KVConnectorBase)
56 logger.info("Creating v1 connector with name: %s and engine_id: %s",
57 connector_cls.__name__, kv_transfer_config.engine_id)
58 # NOTE(Kuntai): v1 connector is explicitly separated into two roles.
@@ -92,25 +69,6 @@ class KVConnectorFactory:
69 # Register various connectors here.
70 # The registration should not be done in each individual file, as we want to
71 # only load the files corresponding to the current connector.
95 -KVConnectorFactory.register_connector(
96 - "PyNcclConnector",
97 - "vllm.distributed.kv_transfer.kv_connector.simple_connector",
98 - "SimpleConnector")
99 -
100 -KVConnectorFactory.register_connector(
101 - "MooncakeConnector",
102 - "vllm.distributed.kv_transfer.kv_connector.simple_connector",
103 - "SimpleConnector")
104 -
105 -KVConnectorFactory.register_connector(
106 - "LMCacheConnector",
107 - "vllm.distributed.kv_transfer.kv_connector.lmcache_connector",
108 - "LMCacheConnector")
109 -
110 -KVConnectorFactory.register_connector(
111 - "MooncakeStoreConnector",
112 - "vllm.distributed.kv_transfer.kv_connector.mooncake_store_connector",
113 - "MooncakeStoreConnector")
72
73 KVConnectorFactory.register_connector(
74 "SharedStorageConnector",
vllm/distributed/kv_transfer/kv_connector/lmcache_connector.py deleted
-99
@@ -1,99 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -LMCache KV Cache Connector for Distributed Machine Learning Inference
5 -
6 -The LMCacheConnector can (1) transfer KV caches between prefill vLLM worker
7 -(KV cache producer) and decode vLLM worker (KV cache consumer) using LMCache;
8 -(2) offload and share KV caches.
9 -"""
10 -
11 -from typing import TYPE_CHECKING, Union
12 -
13 -import torch
14 -
15 -from vllm.config import VllmConfig
16 -from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
17 -from vllm.logger import init_logger
18 -from vllm.sequence import IntermediateTensors
19 -
20 -if TYPE_CHECKING:
21 - from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
22 -
23 -logger = init_logger(__name__)
24 -
25 -
26 -class LMCacheConnector(KVConnectorBase):
27 -
28 - def __init__(
29 - self,
30 - rank: int,
31 - local_rank: int,
32 - config: VllmConfig,
33 - ):
34 -
35 - self.transfer_config = config.kv_transfer_config
36 - self.vllm_config = config
37 -
38 - from lmcache.experimental.cache_engine import LMCacheEngineBuilder
39 - from lmcache.integration.vllm.utils import ENGINE_NAME
40 - from lmcache.integration.vllm.vllm_adapter import (
41 - RetrieveStatus, StoreStatus, init_lmcache_engine,
42 - lmcache_retrieve_kv, lmcache_should_retrieve, lmcache_should_store,
43 - lmcache_store_kv)
44 - logger.info("Initializing LMCacheConfig under kv_transfer_config %s",
45 - self.transfer_config)
46 -
47 - # TODO (Jiayi): Find model_config, parallel_config, and cache_config
48 - self.engine = init_lmcache_engine(config.model_config,
49 - config.parallel_config,
50 - config.cache_config)
51 - self.lmcache_engine_name = ENGINE_NAME
52 - self.lmcache_engine_builder = LMCacheEngineBuilder
53 -
54 - self.model_config = config.model_config
55 - self.parallel_config = config.parallel_config
56 - self.cache_config = config.cache_config
57 - self.lmcache_retrieve_kv = lmcache_retrieve_kv
58 - self.lmcache_store_kv = lmcache_store_kv
59 - self.lmcache_should_retrieve = lmcache_should_retrieve
60 - self.lmcache_should_store = lmcache_should_store
61 - self.store_status = StoreStatus
62 - self.retrieve_status = RetrieveStatus
63 -
64 - def recv_kv_caches_and_hidden_states(
65 - self, model_executable: torch.nn.Module,
66 - model_input: "ModelInputForGPUWithSamplingMetadata",
67 - kv_caches: list[torch.Tensor]
68 - ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,
69 - "ModelInputForGPUWithSamplingMetadata"]:
70 -
71 - retrieve_status = self.lmcache_should_retrieve(model_input)
72 - model_input, bypass_model_exec, hidden_or_intermediate_states =\
73 - self.lmcache_retrieve_kv(
74 - model_executable, model_input, self.cache_config, kv_caches,
75 - retrieve_status)
76 - return hidden_or_intermediate_states, bypass_model_exec, model_input
77 -
78 - def send_kv_caches_and_hidden_states(
79 - self,
80 - model_executable: torch.nn.Module,
81 - model_input: "ModelInputForGPUWithSamplingMetadata",
82 - kv_caches: list[torch.Tensor],
83 - hidden_or_intermediate_states: Union[torch.Tensor,
84 - IntermediateTensors],
85 - ) -> None:
86 -
87 - store_status = self.lmcache_should_store(model_input)
88 - self.lmcache_store_kv(
89 - self.model_config,
90 - self.parallel_config,
91 - self.cache_config,
92 - model_executable,
93 - model_input,
94 - kv_caches,
95 - store_status,
96 - )
97 -
98 - def close(self):
99 - self.lmcache_engine_builder.destroy(self.lmcache_engine_name)
vllm/distributed/kv_transfer/kv_connector/mooncake_store_connector.py deleted
-203
@@ -1,203 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -MooncakeStore Connector for Distributed Machine Learning Inference
5 -The MooncakeStoreConnector transfers KV caches between prefill vLLM workers
6 -(KV cache producer) and decode vLLM workers (KV cache consumer) using a
7 -database-style KVStore.
8 -"""
9 -import hashlib
10 -from typing import TYPE_CHECKING, Union
11 -
12 -import torch
13 -
14 -from vllm.config import VllmConfig
15 -from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
16 -from vllm.distributed.kv_transfer.kv_connector.utils import (
17 - model_aware_kv_ops_helper as kv_helper)
18 -from vllm.logger import init_logger
19 -from vllm.sequence import IntermediateTensors
20 -
21 -if TYPE_CHECKING:
22 - from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
23 -
24 -logger = init_logger(__name__)
25 -
26 -
27 -class MooncakeStoreConnector(KVConnectorBase):
28 -
29 - def __init__(
30 - self,
31 - rank: int,
32 - local_rank: int,
33 - config: VllmConfig,
34 - ):
35 - self.kv_transfer_config = config.kv_transfer_config
36 - self.kv_helper = kv_helper(config)
37 - self.local_tp_rank = local_rank
38 -
39 - # Init kv_store
40 - if self.kv_transfer_config.kv_connector == "MooncakeStoreConnector":
41 - # Check if MOONCAKE_CONFIG_PATH is set
42 - import os
43 - use_mooncake_store = os.getenv('MOONCAKE_CONFIG_PATH') is not None
44 -
45 - if not use_mooncake_store:
46 - raise ValueError(
47 - "To use MooncakeStoreConnector, you need to pass the ENV: "
48 - "'MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json'.")
49 - else:
50 - from vllm.distributed.kv_transfer.kv_lookup_buffer.mooncake_store import ( # noqa: E501
51 - MooncakeStore)
52 - logger.info(
53 - "Initializing KVStoreConnector under kv_transfer_config %s",
54 - self.kv_transfer_config)
55 - self.kv_store = MooncakeStore(config)
56 - else:
57 - logger.error("Can not find %s",
58 - self.kv_transfer_config.kv_connector)
59 -
60 - assert self.kv_store is not None
61 -
62 - def close(self) -> None:
63 - """Close the buffer and release resources.
64 - This method is responsible for cleaning up resources related to the
65 - connector when it is no longer needed.
66 - Raises:
67 - NotImplementedError: This method must be implemented in subclasses.
68 - """
69 - self.kv_store.close()
70 -
71 - def send_kv_caches_and_hidden_states(
72 - self,
73 - model_executable: torch.nn.Module,
74 - model_input: "ModelInputForGPUWithSamplingMetadata",
75 - kv_caches: list[torch.Tensor],
76 - hidden_or_intermediate_states: Union[torch.Tensor,
77 - IntermediateTensors],
78 - ) -> None:
79 - input_tokens_tensor = model_input.input_tokens
80 - seq_lens = model_input.attn_metadata.seq_lens
81 - slot_mapping_flat = model_input.attn_metadata.slot_mapping.flatten()
82 - start_layer = model_executable.model.start_layer
83 - end_layer = model_executable.model.end_layer
84 - num_heads, head_size = self.kv_helper.get_model_args(model_executable)
85 -
86 - for idx, slen in enumerate(seq_lens):
87 - start_pos = sum(seq_lens[:idx])
88 - end_pos = start_pos + slen
89 -
90 - current_tokens = input_tokens_tensor[start_pos:end_pos]
91 - store_key_prefix = self.tensor_hash(current_tokens)
92 - keys, values = [], []
93 -
94 - for layer_id in range(start_layer, end_layer):
95 - kv_cache = kv_caches[layer_id - start_layer]
96 - key_cache, value_cache = self.kv_helper.get_kv_from_cache(
97 - kv_cache, num_heads, head_size)
98 - current_slot_mapping = slot_mapping_flat[start_pos:end_pos]
99 -
100 - keys.append(key_cache[current_slot_mapping].unsqueeze(0))
101 - values.append(value_cache[current_slot_mapping].unsqueeze(0))
102 -
103 - keys = torch.cat(keys, dim=0)
104 - values = torch.cat(values, dim=0)
105 - kvcache_to_sent = torch.stack((keys, values), dim=0)
106 - store_kvcache_key = f"{store_key_prefix}_{self.local_tp_rank}"
107 - self.kv_store.put(store_kvcache_key, kvcache_to_sent)
108 -
109 - hidden_key = f"{store_key_prefix}_hidden_{self.local_tp_rank}"
110 - self.kv_store.put(hidden_key,
111 - hidden_or_intermediate_states[start_pos:end_pos])
112 -
113 - logger.debug("[rank%d]: KV send DONE.", torch.distributed.get_rank())
114 -
115 - def recv_kv_caches_and_hidden_states(
116 - self, model_executable: torch.nn.Module,
117 - model_input: "ModelInputForGPUWithSamplingMetadata",
118 - kv_caches: list[torch.Tensor]
119 - ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,
120 - "ModelInputForGPUWithSamplingMetadata"]:
121 - bypass_model_exec = True
122 - input_tokens_tensor = model_input.input_tokens
123 - seq_lens = model_input.attn_metadata.seq_lens
124 - num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
125 - slot_mapping = model_input.attn_metadata.slot_mapping.flatten()
126 - start_layer = model_executable.model.start_layer
127 - end_layer = model_executable.model.end_layer
128 - hidden_or_intermediate_states_for_one_req = []
129 -
130 - for idx, slen in enumerate(seq_lens):
131 - start_pos = sum(seq_lens[:idx])
132 - end_pos = start_pos + slen
133 -
134 - if start_pos >= num_prefill_tokens:
135 - # This can happen during inflight batching. See:
136 - # vllm/worker/model_runner.py::_prepare_model_input_tensors:
137 - # - input_tokens[:num_prefill_tokens] contains prefill tokens.
138 - # - input_tokens[num_prefill_tokens:] contains decode tokens.
139 - logger.warning("You should set --enable_chunked_prefill=False "
140 - "and --max_num_batched_tokens "
141 - "should be equal to max_seq_len_to_capture")
142 - bypass_model_exec = False
143 - assert start_pos == num_prefill_tokens
144 - break
145 -
146 - current_tokens = input_tokens_tensor[start_pos:end_pos]
147 -
148 - # get roi for current seq
149 - load_key_prefix = self.tensor_hash(current_tokens)
150 - load_kvcache_key = f"{load_key_prefix}_{self.local_tp_rank}"
151 - remote_kv = self.kv_store.get(load_kvcache_key)
152 - hidden_key = f"{load_key_prefix}_hidden_{self.local_tp_rank}"
153 - hidden = self.kv_store.get(hidden_key)
154 -
155 - if remote_kv is None or hidden is None:
156 - # didn't find any match.
157 - bypass_model_exec = False
158 - continue
159 -
160 - num_computed_tokens = current_tokens.shape[0]
161 -
162 - # update the end position based on how many tokens are cached.
163 - end_pos = start_pos + num_computed_tokens
164 -
165 - # call self.kv_store to get kv layer by layer
166 - for layer_id in range(start_layer, end_layer):
167 - layer = model_executable.model.layers[layer_id]
168 - # get kvcache object
169 - kv_cache = kv_caches[layer_id - start_layer]
170 -
171 - # get remote kvcache
172 - remote_k, remote_v = remote_kv[0][layer_id], remote_kv[1][
173 - layer_id]
174 -
175 - self.kv_helper.put_kv_to_cache(model_executable, remote_k,
176 - remote_v, layer, kv_cache,
177 - slot_mapping, start_pos,
178 - end_pos)
179 -
180 - hidden_or_intermediate_states_for_one_req.append(hidden)
181 -
182 - if not bypass_model_exec:
183 - logger.warning(
184 - "[rank%d]: Failed to receive all KVs and hidden "
185 - "states, redo model forwarding.", torch.distributed.get_rank())
186 - hidden_or_intermediate_states = None
187 -
188 - else:
189 - logger.debug(
190 - "[rank%d]: Successfully received all KVs and hidden "
191 - "states, skip model forwarding.", torch.distributed.get_rank())
192 - hidden_or_intermediate_states = torch.cat(
193 - hidden_or_intermediate_states_for_one_req, dim=0)
194 -
195 - return hidden_or_intermediate_states, bypass_model_exec, model_input
196 -
197 - @staticmethod
198 - def tensor_hash(tensor: torch.Tensor) -> int:
199 - """Calculate the hash value of the tensor."""
200 - tensor_bytes = tensor.clone().detach().cpu().numpy().tobytes()
201 - hash_object = hashlib.blake2b(tensor_bytes)
202 - hash_hex = hash_object.hexdigest()
203 - return int(hash_hex[:16], 16)
vllm/distributed/kv_transfer/kv_connector/simple_connector.py deleted
-329
@@ -1,329 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -Simple KV Cache Connector for Distributed Machine Learning Inference
5 -
6 -The SimpleConnector transfers KV caches between prefill vLLM worker (KV cache
7 -producer) and decode vLLM worker (KV cache consumer) using PyNcclPipe or
8 -MooncakePipe.
9 -
10 -But the logic can be extended to support other pipe and lookup buffer.
11 -"""
12 -from typing import TYPE_CHECKING, Optional, Union
13 -
14 -import torch
15 -
16 -from vllm.config import VllmConfig
17 -from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
18 -from vllm.distributed.kv_transfer.kv_connector.utils import (
19 - model_aware_kv_ops_helper as kv_helper)
20 -from vllm.distributed.kv_transfer.kv_lookup_buffer.simple_buffer import (
21 - SimpleBuffer)
22 -from vllm.logger import init_logger
23 -from vllm.sequence import IntermediateTensors
24 -
25 -if TYPE_CHECKING:
26 - from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
27 -
28 -logger = init_logger(__name__)
29 -
30 -
31 -class SimpleConnector(KVConnectorBase):
32 -
33 - def __init__(
34 - self,
35 - rank: int,
36 - local_rank: int,
37 - config: VllmConfig,
38 - ):
39 -
40 - self.config = config.kv_transfer_config
41 - self.kv_helper = kv_helper(config)
42 -
43 - if self.config.kv_connector == "PyNcclConnector":
44 - from vllm.distributed.kv_transfer.kv_pipe.pynccl_pipe import (
45 - PyNcclPipe)
46 - logger.info(
47 - "Initializing PyNcclConfig under kv_transfer_config %s",
48 - self.config)
49 - elif self.config.kv_connector == "MooncakeConnector":
50 - # Check if MOONCAKE_CONFIG_PATH is set
51 - import os
52 - use_mooncake_distributed_pipe = os.getenv(
53 - 'MOONCAKE_CONFIG_PATH') is not None
54 -
55 - if not use_mooncake_distributed_pipe:
56 - raise ValueError(
57 - "To use MooncakeConnector, you need to pass the ENV: "
58 - "'MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json'.")
59 - else:
60 - from vllm.distributed.kv_transfer.kv_pipe.mooncake_pipe import ( # noqa: E501
61 - MooncakePipe)
62 - logger.info(
63 - "Initializing MooncakeConfig under kv_transfer_config %s",
64 - self.config)
65 -
66 - self.lookup_buffer_size = self.config.kv_buffer_size
67 -
68 - self.producer_buffer: Optional[SimpleBuffer] = None
69 - self.consumer_buffer: Optional[SimpleBuffer] = None
70 -
71 - self.producer_data_pipe: Union[PyNcclPipe, MooncakePipe]
72 - self.consumer_data_pipe: Union[PyNcclPipe, MooncakePipe]
73 - self.producer_signal_pipe: Union[PyNcclPipe, MooncakePipe]
74 - self.consumer_signal_pipe: Union[PyNcclPipe, MooncakePipe]
75 -
76 - # 2 pipes for every rank in the world
77 - port_offset_base = 2 * rank
78 -
79 - # In disaggregated prefill, the prefill vLLM only uses send pipe
80 - # and the decode vLLM only uses recv pipe
81 - if self.config.is_kv_producer:
82 -
83 - if self.config.kv_connector == "PyNcclConnector":
84 - self.producer_data_pipe = PyNcclPipe(
85 - local_rank=local_rank,
86 - config=self.config,
87 - port_offset=port_offset_base,
88 - )
89 - self.producer_signal_pipe = PyNcclPipe(
90 - local_rank=local_rank,
91 - config=self.config,
92 - port_offset=port_offset_base + 1,
93 - device="cpu",
94 - )
95 - elif self.config.kv_connector == "MooncakeConnector":
96 - self.producer_data_pipe = MooncakePipe(
97 - local_rank=local_rank,
98 - config=self.config,
99 - )
100 - # We only need to initialize MooncakePipe once
101 - self.producer_signal_pipe = self.producer_data_pipe
102 -
103 - self.producer_buffer = SimpleBuffer(self.producer_signal_pipe,
104 - self.producer_data_pipe,
105 - self.config.kv_buffer_size)
106 -
107 - else:
108 -
109 - # the current vLLM instance is KV consumer, so it needs to connect
110 - # its recv pipe to the send pipe of KV producer
111 - if self.config.kv_connector == "PyNcclConnector":
112 - self.consumer_data_pipe = PyNcclPipe(
113 - local_rank=local_rank,
114 - config=self.config,
115 - port_offset=port_offset_base,
116 - )
117 - self.consumer_signal_pipe = PyNcclPipe(
118 - local_rank=local_rank,
119 - config=self.config,
120 - port_offset=port_offset_base + 1,
121 - device="cpu",
122 - )
123 - elif self.config.kv_connector == "MooncakeConnector":
124 - self.consumer_data_pipe = MooncakePipe(
125 - local_rank=local_rank,
126 - config=self.config,
127 - )
128 - self.consumer_signal_pipe = self.consumer_data_pipe
129 -
130 - self.consumer_buffer = SimpleBuffer(
131 - self.consumer_signal_pipe,
132 - self.consumer_data_pipe,
133 - self.config.kv_buffer_size,
134 - )
135 -
136 - def select(self, input_tokens: Optional[torch.Tensor],
137 - roi: Optional[torch.Tensor]) -> list[Optional[torch.Tensor]]:
138 -
139 - assert self.consumer_buffer is not None, "Please initialize the "\
140 - "consumer buffer before calling select."
141 - return self.consumer_buffer.drop_select(input_tokens, roi)
142 -
143 - def insert(self, input_tokens: torch.Tensor, roi: torch.Tensor,
144 - key: torch.Tensor, value: torch.Tensor,
145 - hidden: torch.Tensor) -> None:
146 -
147 - assert self.producer_buffer is not None, "Please initialize the "\
148 - "producer buffer before calling insert."
149 -
150 - self.producer_buffer.insert(input_tokens, roi, key, value, hidden)
151 -
152 - def send_kv_caches_and_hidden_states(
153 - self,
154 - model_executable: torch.nn.Module,
155 - model_input: "ModelInputForGPUWithSamplingMetadata",
156 - kv_caches: list[torch.Tensor],
157 - hidden_or_intermediate_states: Union[torch.Tensor,
158 - IntermediateTensors],
159 - ) -> None:
160 -
161 - input_tokens_tensor = model_input.input_tokens
162 - seq_lens = model_input.attn_metadata.seq_lens
163 - slot_mapping_flat = model_input.attn_metadata.slot_mapping.flatten()
164 - num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
165 - start_layer = model_executable.model.start_layer
166 - end_layer = model_executable.model.end_layer
167 - num_heads, head_size = self.kv_helper.get_model_args(model_executable)
168 -
169 - # query_lens contains new KV caches that are added to vLLM.
170 - # so we will send them to decode instance
171 - # FIXME(Kuntai): This assume that all requests are prefill.
172 - for idx, slen in enumerate(seq_lens):
173 - start_pos = sum(seq_lens[:idx])
174 - end_pos = start_pos + slen
175 -
176 - if start_pos >= num_prefill_tokens:
177 - # vllm/worker/model_runner.py::_prepare_model_input_tensors:
178 - # - input_tokens[:num_prefill_tokens] contains prefill tokens.
179 - # - input_tokens[num_prefill_tokens:] contains decode tokens.
180 - logger.warning("You have some decode requests while using "
181 - "SimpleConnector. Their KVCache won't be sent.")
182 - break
183 -
184 - current_tokens = input_tokens_tensor[start_pos:end_pos]
185 -
186 - keys, values = [], []
187 -
188 - for layer_id in range(start_layer, end_layer):
189 - kv_cache = kv_caches[layer_id - start_layer]
190 - key_cache, value_cache = self.kv_helper.get_kv_from_cache(
191 - kv_cache, num_heads, head_size)
192 -
193 - current_slot_mapping = slot_mapping_flat[start_pos:end_pos]
194 -
195 - keys.append(key_cache[current_slot_mapping].unsqueeze(0))
196 - values.append(value_cache[current_slot_mapping].unsqueeze(0))
197 -
198 - keys = torch.cat(keys, dim=0)
199 - values = torch.cat(values, dim=0)
200 -
201 - self.insert(current_tokens,
202 - torch.ones_like(current_tokens,
203 - dtype=bool), keys, values,
204 - hidden_or_intermediate_states[start_pos:end_pos])
205 -
206 - logger.debug("[rank%d]: KV send DONE.", torch.distributed.get_rank())
207 -
208 - def recv_kv_caches_and_hidden_states(
209 - self, model_executable: torch.nn.Module,
210 - model_input: "ModelInputForGPUWithSamplingMetadata",
211 - kv_caches: list[torch.Tensor]
212 - ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,
213 - "ModelInputForGPUWithSamplingMetadata"]:
214 -
215 - # When bypass_model_exec is set to False, it means that at least for one
216 - # request its corresponding KV cache or hidden state is missing.
217 - # In this case we need to do prefilling to recompute missing KV cache
218 - # and hidden states.
219 - bypass_model_exec = True
220 -
221 - input_tokens_tensor = model_input.input_tokens
222 - seq_lens = model_input.attn_metadata.seq_lens
223 - num_prefill_tokens = model_input.attn_metadata.num_prefill_tokens
224 - slot_mapping = model_input.attn_metadata.slot_mapping.flatten()
225 - start_layer = model_executable.model.start_layer
226 - end_layer = model_executable.model.end_layer
227 -
228 - hidden_or_intermediate_states_for_one_req = []
229 -
230 - input_tokens_list = []
231 - num_computed_tokens_list = []
232 - start_pos_list = []
233 -
234 - # enumerate different requests
235 - # FIXME(Kuntai): This impl assumes that all requests are prefill.
236 - for idx, slen in enumerate(seq_lens):
237 - start_pos = sum(seq_lens[:idx])
238 - end_pos = start_pos + slen
239 -
240 - if start_pos >= num_prefill_tokens:
241 - # This can happen during inflight batching. See:
242 - # vllm/worker/model_runner.py::_prepare_model_input_tensors:
243 - # - input_tokens[:num_prefill_tokens] contains prefill tokens.
244 - # - input_tokens[num_prefill_tokens:] contains decode tokens.
245 - logger.warning("You should set --enable_chunked_prefill=False "
246 - "and --max_num_batched_tokens "
247 - "should be equal to --max_seq_len_to_capture")
248 - bypass_model_exec = False
249 - assert start_pos == num_prefill_tokens
250 - break
251 -
252 - current_tokens = input_tokens_tensor[start_pos:end_pos]
253 - num_tokens = slen
254 -
255 - # collecting data for rebuilding the input
256 - input_tokens_list.append(current_tokens)
257 - start_pos_list.append(start_pos)
258 -
259 - ret = self.select(current_tokens,
260 - torch.ones_like(current_tokens, dtype=bool))
261 - if ret[0] is None:
262 - # didn't find any match.
263 - bypass_model_exec = False
264 - num_computed_tokens_list.append(0)
265 - continue
266 -
267 - roi: torch.Tensor = ret[1]
268 - keys: torch.Tensor = ret[2]
269 - values: torch.Tensor = ret[3]
270 - hidden: torch.Tensor = ret[4]
271 -
272 - num_computed_tokens = roi.shape[0]
273 - num_computed_tokens_list.append(num_computed_tokens)
274 -
275 - # check if both KV cache and the hidden states are received
276 - # If not, need to redo the forwarding to compute missing states
277 - if not all([(num_computed_tokens == num_tokens), hidden is not None
278 - ]):
279 - bypass_model_exec = False
280 -
281 - # update the end position based on how many tokens are cached.
282 - end_pos = start_pos + num_computed_tokens
283 -
284 - # put received KV caches into paged memory
285 - for cur_layer in range(start_layer, end_layer):
286 -
287 - layer_id = cur_layer - start_layer
288 - kv_cache = kv_caches[layer_id]
289 - layer = model_executable.model.layers[cur_layer]
290 -
291 - # get remote kvcache
292 - remote_k, remote_v = keys[layer_id], values[layer_id]
293 -
294 - self.kv_helper.put_kv_to_cache(model_executable, remote_k,
295 - remote_v, layer, kv_cache,
296 - slot_mapping, start_pos,
297 - end_pos)
298 -
299 - hidden_or_intermediate_states_for_one_req.append(hidden)
300 -
301 - if not bypass_model_exec:
302 - # Some of the KV cache is not retrieved
303 - # Here we will fall back to normal model forwarding
304 - # But optionally you can adjust model_input so that you only do
305 - # prefilling on those tokens that are missing KV caches.
306 - logger.warning(
307 - "[rank%d]: Failed to receive all KVs and hidden "
308 - "states, redo model forwarding.", torch.distributed.get_rank())
309 - hidden_or_intermediate_states = None
310 -
311 - else:
312 - logger.debug(
313 - "[rank%d]: Successfully received all KVs and hidden "
314 - "states, skip model forwarding.", torch.distributed.get_rank())
315 - hidden_or_intermediate_states = torch.cat(
316 - hidden_or_intermediate_states_for_one_req, dim=0)
317 -
318 - return hidden_or_intermediate_states, bypass_model_exec, model_input
319 -
320 - def close(self):
321 - self.producer_data_pipe.close()
322 - self.consumer_data_pipe.close()
323 - if self.config.kv_connector == "PyNcclConnector":
324 - self.producer_signal_pipe.close()
325 - self.consumer_signal_pipe.close()
326 - elif self.config.kv_connector == "MooncakeConnector":
327 - # MooncakePipe reuses data_pipe for signal_pipe, so we only have to
328 - # close the data_pipe.
329 - pass
vllm/distributed/kv_transfer/kv_connector/utils.py
+4 -5
@@ -13,8 +13,8 @@ import torch
13 import vllm.envs as envs
14 from vllm import _custom_ops as ops
15 from vllm.config import VllmConfig, get_current_vllm_config
16 -from vllm.distributed.kv_transfer.kv_connector.factory import (
17 - KVConnectorFactory)
16 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
17 + KVConnectorBase_V1)
18 from vllm.logger import init_logger
19 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
20
@@ -106,9 +106,8 @@ def get_kv_connector_cache_layout():
106 vllm_config = get_current_vllm_config()
107 kv_config = vllm_config.kv_transfer_config
108 if kv_config is not None:
109 - connector_cls = KVConnectorFactory.get_connector_class(kv_config)
110 - required_kvcache_layout = connector_cls.get_required_kvcache_layout(
111 - vllm_config)
109 + required_kvcache_layout = (
110 + KVConnectorBase_V1.get_required_kvcache_layout(vllm_config))
111 if required_kvcache_layout is not None:
112 return required_kvcache_layout
113 logger.info_once("Connectors do not specify a " \
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+4 -4
@@ -52,7 +52,7 @@ class MultiConnector(KVConnectorBase_V1):
52 temp_config.kv_transfer_config = KVTransferConfig(
53 **ktc, engine_id=engine_id)
54 self._connectors.append(
55 - KVConnectorFactory.create_connector_v1(temp_config, role))
55 + KVConnectorFactory.create_connector(temp_config, role))
56
57 # A mapping from request id to the index of the connector chosen to
58 # load the request from (if any).
@@ -223,9 +223,9 @@ class MultiConnector(KVConnectorBase_V1):
223 for ktc in ktcs:
224 kv_transfer_config = KVTransferConfig(**ktc)
225 temp_vllm_config.kv_transfer_config = kv_transfer_config
226 - required_kvcache_layout = KVConnectorFactory.get_connector_class(
227 - kv_transfer_config).get_required_kvcache_layout(
228 - temp_vllm_config)
226 + required_kvcache_layout = (
227 + KVConnectorBase_V1.get_required_kvcache_layout(
228 + temp_vllm_config))
229 if required_kvcache_layout is not None:
230 layouts.add(required_kvcache_layout)
231
vllm/distributed/kv_transfer/kv_connector_agent.py deleted
-77
@@ -1,77 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""A centralized entrypoint to perform distributed KV cache transfer.
4 -
5 -This implementation is a shim wrapper on two APIs exposed by `kv_connector`:
6 -1. `send_kv_caches_and_hidden_states`
7 -2. `recv_kv_caches_and_hidden_states
8 -"""
9 -from typing import TYPE_CHECKING, Union
10 -
11 -if TYPE_CHECKING:
12 - from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
13 - from vllm.config import VllmConfig
14 -
15 -import torch
16 -
17 -from vllm.distributed.kv_transfer.kv_connector.factory import (
18 - KVConnectorFactory)
19 -from vllm.logger import init_logger
20 -from vllm.sequence import IntermediateTensors
21 -
22 -logger = init_logger(__name__)
23 -
24 -
25 -class KVTransferAgent:
26 - """
27 - A class designated for distributed KV transfer
28 -
29 - Target use cases:
30 - 1. Disaggregated prefill
31 - 2. Remote KV cache storage
32 - """
33 -
34 - def __init__(
35 - self,
36 - rank: int,
37 - local_rank: int,
38 - config: "VllmConfig",
39 - ):
40 -
41 - self.config = config
42 -
43 - if config.kv_transfer_config is None:
44 - raise ValueError("KVTransferConfig is not set in the VllmConfig,"
45 - " cannot initialize KVConnector.")
46 -
47 - assert self.config.kv_transfer_config.is_kv_transfer_instance, "KV"\
48 - "TransferAgent should only be used when kv_connector is set."
49 -
50 - self.connector = KVConnectorFactory.create_connector_v0(
51 - rank, local_rank, config)
52 -
53 - def send_kv_caches_and_hidden_states(
54 - self,
55 - model_executable: torch.nn.Module,
56 - model_input: "ModelInputForGPUWithSamplingMetadata",
57 - kv_caches: list[torch.Tensor],
58 - hidden_or_intermediate_states: Union[torch.Tensor,
59 - IntermediateTensors],
60 - ) -> None:
61 -
62 - self.connector.send_kv_caches_and_hidden_states(
63 - model_executable, model_input, kv_caches,
64 - hidden_or_intermediate_states)
65 -
66 - def close(self) -> None:
67 - self.connector.close()
68 -
69 - def recv_kv_caches_and_hidden_states(
70 - self, model_executable: torch.nn.Module,
71 - model_input: "ModelInputForGPUWithSamplingMetadata",
72 - kv_caches: list[torch.Tensor]
73 - ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,
74 - "ModelInputForGPUWithSamplingMetadata"]:
75 -
76 - return self.connector.recv_kv_caches_and_hidden_states(
77 - model_executable, model_input, kv_caches)
vllm/distributed/kv_transfer/kv_transfer_state.py
+2 -7
@@ -8,7 +8,6 @@ from vllm.distributed.kv_transfer.kv_connector.factory import (
8 KVConnectorFactory)
9 from vllm.distributed.kv_transfer.kv_connector.v1 import (KVConnectorBase_V1,
10 KVConnectorRole)
11 -from vllm.distributed.parallel_state import get_world_group
11
12 if TYPE_CHECKING:
13 from vllm.config import VllmConfig
@@ -61,11 +60,7 @@ def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
60 if (vllm_config.kv_transfer_config.is_kv_transfer_instance
61 and _KV_CONNECTOR_AGENT is None):
62 if envs.VLLM_USE_V1:
64 - _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v1(
63 + _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector(
64 config=vllm_config, role=KVConnectorRole.WORKER)
65 else:
67 - _KV_CONNECTOR_AGENT = KVConnectorFactory.create_connector_v0(
68 - rank=get_world_group().rank,
69 - local_rank=get_world_group().local_rank,
70 - config=vllm_config,
71 - )
66 + raise ValueError("V0 is no longer supported")
vllm/v1/core/sched/scheduler.py
+1 -1
@@ -83,7 +83,7 @@ class Scheduler(SchedulerInterface):
83 assert len(self.kv_cache_config.kv_cache_groups) == 1, (
84 "Multiple KV cache groups are not currently supported "
85 "with KV connectors")
86 - self.connector = KVConnectorFactory.create_connector_v1(
86 + self.connector = KVConnectorFactory.create_connector(
87 config=self.vllm_config, role=KVConnectorRole.SCHEDULER)
88
89 self.kv_event_publisher = EventPublisherFactory.create(
vllm/v1/worker/kv_connector_model_runner_mixin.py
+3 -3
@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Optional
11 from vllm.config import VllmConfig
12 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
13 has_kv_transfer_group)
14 -from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
14 +from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
15 from vllm.forward_context import get_forward_context, set_forward_context
16 from vllm.logger import init_logger
17 from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput,
@@ -31,7 +31,7 @@ class KVConnectorModelRunnerMixin:
31 # Update KVConnector with the KVConnector metadata forward().
32 if has_kv_transfer_group():
33 kv_connector = get_kv_transfer_group()
34 - assert isinstance(kv_connector, KVConnectorBase_V1)
34 + assert isinstance(kv_connector, KVConnectorBase)
35 assert scheduler_output.kv_connector_metadata is not None
36 kv_connector.bind_connector_metadata(
37 scheduler_output.kv_connector_metadata)
@@ -93,7 +93,7 @@ class KVConnectorModelRunnerMixin:
93
94 # Update KVConnector with the KVConnector metadata forward().
95 kv_connector = get_kv_transfer_group()
96 - assert isinstance(kv_connector, KVConnectorBase_V1)
96 + assert isinstance(kv_connector, KVConnectorBase)
97 assert scheduler_output.kv_connector_metadata is not None
98 kv_connector.bind_connector_metadata(
99 scheduler_output.kv_connector_metadata)