@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 83f478bb19
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

[KVConnector] Migrate the LMCache integration code to be vLLM native (#25542)

Signed-off-by: ApostaC <yihua98@uchicago.edu>

Yihua Cheng committed Oct 24, 2025 at 17:23 UTC 83f478bb19489b41e9d208b47b4bb5a95ac171ac
4 files changed +1637 -2
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+18 -2
@@ -3,7 +3,9 @@
3 from typing import TYPE_CHECKING, Any
4
5 import torch
6 -from lmcache.integration.vllm.vllm_v1_adapter import LMCacheConnectorV1Impl
6 +from lmcache.integration.vllm.vllm_v1_adapter import (
7 + LMCacheConnectorV1Impl as LMCacheConnectorLatestImpl,
8 +)
9
10 from vllm.config import VllmConfig
11 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -11,6 +13,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
13 KVConnectorMetadata,
14 KVConnectorRole,
15 )
16 +from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_integration import (
17 + vllm_v1_adapter as _adapter,
18 +)
19 from vllm.logger import init_logger
20 from vllm.v1.core.sched.output import SchedulerOutput
21
@@ -26,7 +31,18 @@ logger = init_logger(__name__)
31 class LMCacheConnectorV1(KVConnectorBase_V1):
32 def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
33 super().__init__(vllm_config=vllm_config, role=role)
29 - self._lmcache_engine = LMCacheConnectorV1Impl(vllm_config, role, self)
34 + assert vllm_config.kv_transfer_config is not None
35 + use_native = vllm_config.kv_transfer_config.get_from_extra_config(
36 + "use_native", False
37 + )
38 + if use_native:
39 + logger.info("Initializing native LMCache connector")
40 + cls = _adapter.LMCacheConnectorV1Impl
41 + else:
42 + logger.info("Initializing latest dev LMCache connector")
43 + cls = LMCacheConnectorLatestImpl
44 +
45 + self._lmcache_engine = cls(vllm_config, role, self)
46
47 # ==============================
48 # Worker-side methods
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/utils.py new
+221
@@ -0,0 +1,221 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Standard
4 +import os
5 +import threading
6 +from typing import TYPE_CHECKING, Union
7 +
8 +import torch
9 +from lmcache.config import LMCacheEngineConfig as Config
10 +from lmcache.logging import init_logger
11 +from lmcache.v1.config import LMCacheEngineConfig as V1Config
12 +
13 +if TYPE_CHECKING:
14 + from vllm.config import ModelConfig
15 + from vllm.multimodal.inputs import PlaceholderRange
16 + from vllm.v1.core.sched.output import NewRequestData
17 + from vllm.v1.request import Request
18 +
19 +logger = init_logger(__name__)
20 +ENGINE_NAME = "vllm-instance"
21 +
22 +# Thread-safe singleton storage
23 +_config_instance: Config | V1Config | None = None
24 +_config_lock = threading.Lock()
25 +
26 +
27 +def is_false(value: str) -> bool:
28 + """Check if the given string value is equivalent to 'false'."""
29 + return value.lower() in ("false", "0", "no", "n", "off")
30 +
31 +
32 +def lmcache_get_or_create_config() -> Config | V1Config:
33 + """Get the LMCache configuration from the environment variable
34 + `LMCACHE_CONFIG_FILE`. If the environment variable is not set, this
35 + function will return the default configuration.
36 +
37 + This function is thread-safe and implements singleton pattern,
38 + ensuring the configuration is loaded only once.
39 + """
40 + global _config_instance
41 +
42 + # Double-checked locking for thread-safe singleton
43 + if _config_instance is None:
44 + with _config_lock:
45 + if _config_instance is None: # Check again within lock
46 + if is_false(os.getenv("LMCACHE_USE_EXPERIMENTAL", "True")):
47 + logger.warning(
48 + "Detected LMCACHE_USE_EXPERIMENTAL is set to False. "
49 + "Using legacy configuration is deprecated and will "
50 + "be remove soon! Please set LMCACHE_USE_EXPERIMENTAL "
51 + "to True."
52 + )
53 + LMCacheEngineConfig = Config # type: ignore[assignment]
54 + else:
55 + LMCacheEngineConfig = V1Config # type: ignore[assignment]
56 +
57 + if "LMCACHE_CONFIG_FILE" not in os.environ:
58 + logger.warning(
59 + "No LMCache configuration file is set. Trying to read"
60 + " configurations from the environment variables."
61 + )
62 + logger.warning(
63 + "You can set the configuration file through "
64 + "the environment variable: LMCACHE_CONFIG_FILE"
65 + )
66 + _config_instance = LMCacheEngineConfig.from_env()
67 + else:
68 + config_file = os.environ["LMCACHE_CONFIG_FILE"]
69 + logger.info("Loading LMCache config file %s", config_file)
70 + _config_instance = LMCacheEngineConfig.from_file(config_file)
71 + # Update config from environment variables
72 + _config_instance.update_config_from_env()
73 + return _config_instance
74 +
75 +
76 +def hex_hash_to_int16(s: str) -> int:
77 + """
78 + Convert a hex hash string to a 16-bit integer.
79 + """
80 + return int(s, 16) & 0xFFFF
81 +
82 +
83 +def apply_mm_hashes_to_token_ids(
84 + token_ids: torch.Tensor,
85 + mm_hashes: list[str],
86 + mm_positions: list["PlaceholderRange"],
87 +) -> torch.Tensor:
88 + """
89 + Overwrite token_ids in-place for multimodal placeholders using
90 + efficient slice assignments.
91 + """
92 + n = token_ids.size(0)
93 + for hash_str, placeholder in zip(mm_hashes, mm_positions):
94 + start, length = placeholder.offset, placeholder.length
95 + if start >= n:
96 + continue
97 + end = min(start + length, n)
98 + token_ids[start:end] = hex_hash_to_int16(hash_str)
99 + return token_ids
100 +
101 +
102 +def mla_enabled(model_config: "ModelConfig") -> bool:
103 + return (
104 + hasattr(model_config, "use_mla")
105 + and isinstance(model_config.use_mla, bool)
106 + and model_config.use_mla
107 + )
108 +
109 +
110 +def create_lmcache_metadata(
111 + vllm_config=None, model_config=None, parallel_config=None, cache_config=None
112 +):
113 + """
114 + Create LMCacheEngineMetadata from vLLM configuration.
115 +
116 + This function extracts common metadata creation logic that was duplicated
117 + across multiple files.
118 +
119 + Args:
120 + vllm_config (VllmConfig): vLLM configuration object containing model,
121 + parallel, and cache configs (alternative to
122 + individual config parameters)
123 + model_config (ModelConfig): Model configuration (alternative to
124 + vllm_config)
125 + parallel_config (ParallelConfig): Parallel configuration (alternative
126 + to vllm_config)
127 + cache_config (CacheConfig): Cache configuration (alternative to
128 + vllm_config)
129 + """
130 + # Third Party
131 + # First Party
132 + from lmcache.config import LMCacheEngineMetadata
133 +
134 + from vllm.utils import get_kv_cache_torch_dtype
135 +
136 + config = lmcache_get_or_create_config()
137 + # Support both vllm_config object and individual config parameters
138 + if vllm_config is not None:
139 + model_cfg = vllm_config.model_config
140 + parallel_cfg = vllm_config.parallel_config
141 + cache_cfg = vllm_config.cache_config
142 + else:
143 + if model_config is None or parallel_config is None or cache_config is None:
144 + raise ValueError(
145 + "Either vllm_config must be provided, or all of "
146 + "model_config, parallel_config, and cache_config must be provided."
147 + )
148 + model_cfg = model_config
149 + parallel_cfg = parallel_config
150 + cache_cfg = cache_config
151 +
152 + # Get KV cache dtype
153 + kv_dtype = get_kv_cache_torch_dtype(cache_cfg.cache_dtype, model_cfg.dtype)
154 +
155 + # Check if MLA is enabled
156 + use_mla = mla_enabled(model_cfg)
157 +
158 + # Construct KV shape (for memory pool)
159 + num_layer = model_cfg.get_num_layers(parallel_cfg)
160 + chunk_size = config.chunk_size
161 + num_kv_head = model_cfg.get_num_kv_heads(parallel_cfg)
162 + head_size = model_cfg.get_head_size()
163 + kv_shape = (num_layer, 1 if use_mla else 2, chunk_size, num_kv_head, head_size)
164 +
165 + # Create metadata
166 + metadata = LMCacheEngineMetadata(
167 + model_cfg.model,
168 + parallel_cfg.world_size,
169 + parallel_cfg.rank,
170 + "vllm",
171 + kv_dtype,
172 + kv_shape,
173 + use_mla,
174 + )
175 +
176 + return metadata, config
177 +
178 +
179 +def extract_mm_features(
180 + request: Union["Request", "NewRequestData"], modify: bool = False
181 +) -> tuple[list[str], list["PlaceholderRange"]]:
182 + """
183 + Normalize multimodal information from a Request into parallel lists.
184 +
185 + This helper reads either:
186 + 1) `request.mm_features` (objects each exposing `.identifier` and
187 + `.mm_position`), or
188 + 2) legacy fields `request.mm_hashes` and `request.mm_positions`.
189 +
190 + It returns two equally sized lists: the multimodal hash identifiers and
191 + their corresponding positions. If the request contains no multimodal info,
192 + it returns `([], [])`.
193 +
194 + Args:
195 + request (Request): The source object.
196 + modify (bool):
197 + Controls copy semantics for the legacy-path return values.
198 + - If True and legacy fields are used, shallow-copies are returned so
199 + the caller can mutate the lists without affecting `request`.
200 + - If False, the original legacy sequences are returned as-is
201 + (zero-copy); treat them as read-only.
202 +
203 + Returns:
204 + tuple[list[str], list[PlaceholderRange]]: (`mm_hashes`, `mm_positions`).
205 + May be `([], [])` when no multimodal data is present.
206 + """
207 + if getattr(request, "mm_features", None):
208 + mm_hashes, mm_positions = zip(
209 + *((f.identifier, f.mm_position) for f in request.mm_features)
210 + )
211 + return (list(mm_hashes), list(mm_positions))
212 + elif getattr(request, "mm_hashes", None):
213 + if modify:
214 + return (
215 + request.mm_hashes.copy(), # type: ignore
216 + request.mm_positions.copy(), # type: ignore
217 + )
218 + else:
219 + return (request.mm_hashes, request.mm_positions) # type: ignore
220 + else:
221 + return ([], [])
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py new
+1396
@@ -0,0 +1,1396 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# Standard
4 +import os
5 +import uuid
6 +from collections.abc import Generator
7 +from dataclasses import dataclass, field
8 +from typing import TYPE_CHECKING, Any, Optional
9 +
10 +import torch
11 +from lmcache import utils
12 +from lmcache.config import LMCacheEngineMetadata
13 +from lmcache.logging import init_logger
14 +from lmcache.observability import LMCStatsMonitor
15 +from lmcache.utils import _lmcache_nvtx_annotate
16 +from lmcache.v1.cache_engine import LMCacheEngine, LMCacheEngineBuilder
17 +from lmcache.v1.compute.blend import LMCBlenderBuilder
18 +from lmcache.v1.config import LMCacheEngineConfig, _validate_and_set_config_value
19 +from lmcache.v1.gpu_connector import (
20 + VLLMBufferLayerwiseGPUConnector,
21 + VLLMPagedMemGPUConnectorV2,
22 + VLLMPagedMemLayerwiseGPUConnector,
23 +)
24 +from lmcache.v1.internal_api_server.api_server import InternalAPIServer
25 +from lmcache.v1.lookup_client import LookupClientFactory
26 +from lmcache.v1.lookup_client.lmcache_async_lookup_client import (
27 + LMCacheAsyncLookupServer,
28 +)
29 +from lmcache.v1.offload_server.zmq_server import ZMQOffloadServer
30 +from lmcache.v1.plugin.plugin_launcher import PluginLauncher
31 +
32 +from vllm.config import VllmConfig
33 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
34 + KVConnectorBase_V1,
35 + KVConnectorMetadata,
36 + KVConnectorRole,
37 +)
38 +from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_integration.utils import (
39 + ENGINE_NAME,
40 + apply_mm_hashes_to_token_ids,
41 + extract_mm_features,
42 + lmcache_get_or_create_config,
43 + mla_enabled,
44 +)
45 +from vllm.distributed.parallel_state import get_tensor_model_parallel_rank, get_tp_group
46 +from vllm.sampling_params import SamplingParams
47 +from vllm.utils import cdiv, get_kv_cache_torch_dtype
48 +from vllm.v1.core.sched.output import SchedulerOutput
49 +from vllm.version import __version__ as VLLM_VERSION
50 +
51 +if TYPE_CHECKING:
52 + from vllm.attention.backends.abstract import AttentionMetadata
53 + from vllm.forward_context import ForwardContext
54 + from vllm.multimodal.inputs import PlaceholderRange
55 + from vllm.v1.core.kv_cache_manager import KVCacheManager
56 + from vllm.v1.core.sched.output import NewRequestData
57 + from vllm.v1.request import Request
58 +
59 +logger = init_logger(__name__)
60 +
61 +
62 +@dataclass
63 +class LoadSpec:
64 + # Number of tokens cached in vLLM
65 + vllm_cached_tokens: int
66 + # Number of tokens that are cached in LMCache
67 + lmcache_cached_tokens: int
68 + # Whether the scheduler allow us to load the tokens
69 + can_load: bool
70 +
71 +
72 +@dataclass
73 +class SaveSpec:
74 + # Skip already saved tokens
75 + skip_leading_tokens: int
76 + # Whether the scheduler allow us to save the tokens
77 + can_save: bool
78 +
79 +
80 +@dataclass
81 +class DisaggSpec:
82 + req_id: str
83 + receiver_id: str
84 + receiver_host: str
85 + receiver_init_port: int
86 + receiver_alloc_port: int
87 + is_last_prefill: bool = False
88 + num_transferred_tokens: int = 0
89 +
90 +
91 +tmp_disagg_tracker: dict[str, DisaggSpec] = {}
92 +
93 +
94 +def extract_request_configs(sampling_params: SamplingParams) -> dict | None:
95 + request_configs = None
96 + if (
97 + sampling_params.extra_args is not None
98 + and "kv_transfer_params" in sampling_params.extra_args
99 + ):
100 + kv_transfer_params = sampling_params.extra_args.get("kv_transfer_params")
101 + if kv_transfer_params is None:
102 + return None
103 + assert isinstance(kv_transfer_params, dict)
104 + for k, v in kv_transfer_params.items():
105 + if k.startswith("lmcache."):
106 + if request_configs is None:
107 + request_configs = {}
108 + request_configs[k] = v
109 + return request_configs
110 +
111 +
112 +@dataclass
113 +class RequestTracker:
114 + # Request id
115 + req_id: str
116 +
117 + # Total prompt token length
118 + prompt_len: int
119 +
120 + # The token ids that has been scheduled so far
121 + token_ids: list[int]
122 +
123 + # The block ids that has been allocated so far
124 + # NOTE: allocated blocks could be more than the number of tokens
125 + allocated_block_ids: list[int]
126 +
127 + # The number of tokens that has been saved
128 + num_saved_tokens: int = 0
129 +
130 + # Disagg spec for the request
131 + disagg_spec: DisaggSpec | None = None
132 +
133 + # Multimodal hashes and positions
134 + mm_hashes: list[str] | None = None
135 + mm_positions: list["PlaceholderRange"] | None = None
136 +
137 + # The configs of the request, includes tags and other configs
138 + request_configs: dict | None = None
139 +
140 + # Whether the request is in decode phase
141 + is_decode_phase = False
142 +
143 + # Whether the request cache should be saved
144 + skip_save: bool = False
145 +
146 + @_lmcache_nvtx_annotate
147 + @staticmethod
148 + def from_new_request(
149 + lmcache_config: LMCacheEngineConfig,
150 + new_request: "NewRequestData",
151 + num_tokens_to_compute: int,
152 + lmcache_cached_tokens: int,
153 + skip_save: bool,
154 + ) -> "RequestTracker":
155 + """Create the request tracker from a new request.
156 +
157 + Args:
158 + lmcache_config (LMCacheEngineConfig): the LMCache engine config.
159 + new_request (NewRequestData): the new request data.
160 + num_tokens_to_compute (int): the number of tokens that will
161 + be 'computed', including the `num_computed_tokens` (vLLM's
162 + local cache hit) and new tokens that will be scheduled.
163 + lmcache_cached_tokens (int): the number of tokens that are
164 + cached in LMCache.
165 + skip_save (bool): whether the request cache should be saved
166 + """
167 + # vLLM 0.9.0 update: request.block_ids changed from list[int] to
168 + # list[list[int]]
169 + # Need to check the type of request.block_ids
170 +
171 + unfolded_block_ids = []
172 +
173 + if not isinstance(new_request.block_ids[0], list):
174 + unfolded_block_ids = new_request.block_ids.copy()
175 + else:
176 + # According to the vLLM code
177 + # (https://github.com/vllm-project/vllm/blob/main/vllm/v1/core/
178 + # sched/scheduler.py#L943),
179 + # only one KVCacheGroup is supported in connector for now.
180 + unfolded_block_ids = new_request.block_ids[0].copy()
181 +
182 + # NOTE: Initialized in `update_state_after_alloc`
183 + disagg_spec = tmp_disagg_tracker.pop(new_request.req_id, None)
184 +
185 + if new_request.sampling_params:
186 + request_configs = extract_request_configs(new_request.sampling_params)
187 + else:
188 + request_configs = None
189 +
190 + mm_hashes, mm_positions = extract_mm_features(new_request, modify=True)
191 +
192 + assert new_request.prompt_token_ids is not None
193 + return RequestTracker(
194 + req_id=new_request.req_id,
195 + prompt_len=len(new_request.prompt_token_ids),
196 + token_ids=new_request.prompt_token_ids[:num_tokens_to_compute].copy(),
197 + allocated_block_ids=unfolded_block_ids,
198 + num_saved_tokens=lmcache_cached_tokens,
199 + disagg_spec=disagg_spec,
200 + mm_hashes=mm_hashes,
201 + mm_positions=mm_positions,
202 + skip_save=skip_save,
203 + request_configs=request_configs,
204 + )
205 +
206 + def update(
207 + self,
208 + new_token_ids: list[int],
209 + new_block_ids: tuple[list[int], ...] | None | list[int],
210 + ) -> None:
211 + """Update the request tracker when a running request is
212 + scheduled again
213 + """
214 +
215 + self.token_ids.extend(new_token_ids)
216 +
217 + if new_block_ids is None:
218 + # https://github.com/vllm-project/vllm/commit/
219 + # b029de9902aa3ac58806c8c17776c7074175b6db
220 + new_block_ids = []
221 + elif len(new_block_ids) == 0:
222 + new_block_ids = []
223 + elif isinstance(new_block_ids, tuple):
224 + new_block_ids = new_block_ids[0]
225 + elif isinstance(new_block_ids, list):
226 + pass
227 + else:
228 + raise ValueError(f"Unsupported new_block_ids type {type(new_block_ids)}")
229 + self.allocated_block_ids.extend(new_block_ids)
230 +
231 + # When a request is scheduled again, and the number of new tokens
232 + # is 1 (excluding chunked prefill), the request is in decode phase.
233 + if len(new_token_ids) == 1:
234 + self.is_decode_phase = True
235 +
236 +
237 +@dataclass
238 +class ReqMeta:
239 + # Request id
240 + req_id: str
241 + # Request tokens
242 + token_ids: list[int] # torch.Tensor
243 + # Slot mapping
244 + slot_mapping: torch.Tensor
245 +
246 + # Whether is last prefill or not
247 + is_last_prefill: bool = False
248 +
249 + # Skip save or not
250 + save_spec: SaveSpec | None = None
251 + # load_spec
252 + load_spec: LoadSpec | None = None
253 + # disagg spec
254 + disagg_spec: DisaggSpec | None = None
255 + # the configs of the request
256 + request_configs: dict | None = None
257 +
258 + @staticmethod
259 + def from_request_tracker(
260 + tracker: RequestTracker,
261 + block_size: int,
262 + lmcache_chunk_size: int = 256,
263 + load_spec: LoadSpec | None = None,
264 + discard_partial_chunks: bool = True,
265 + save_decode_cache: bool = False,
266 + ) -> Optional["ReqMeta"]:
267 + """Create the request metadata from a request tracker.
268 +
269 + Args:
270 + tracker (RequestTracker): the request tracker.
271 + block_size (int): the block size in vLLM.
272 + lmcache_chunk_size (int): the chunk size for LMCache.
273 + load_spec (Optional[LoadSpec]): the load spec for KV cache loading.
274 + discard_partial_chunks (bool): whether to discard partial chunks.
275 + save_decode_cache (bool): whether to save the cache in decode phase.
276 +
277 + Returns:
278 + the request metadata if we need to perform load/save
279 + operations, None otherwise.
280 + """
281 + input_token_ids = tracker.token_ids
282 + input_token_len = len(input_token_ids)
283 +
284 + is_last_prefill = False
285 + if input_token_len == tracker.prompt_len:
286 + is_last_prefill = True
287 +
288 + # For save operation: do not save if the following condition is met
289 + # 1. has already been saved before (num_saved_tokens > 0)
290 + # 2. number of unsaved tokens is not reached the chunk boundary
291 + # 3. if save_decode_cache is False and it is in decode phase
292 +
293 + skip_leading_tokens = tracker.num_saved_tokens
294 + chunk_boundary = (
295 + cdiv(tracker.num_saved_tokens + 1, lmcache_chunk_size) * lmcache_chunk_size
296 + )
297 +
298 + # NOTE(vladnosiv): for disagg, you cannot skip saving, as saving is a
299 + # trqansfer. Check if request_configs has lmcache.skip_save set to True
300 + request_skip = (tracker.request_configs or {}).get("lmcache.skip_save", False)
301 +
302 + skip_save = tracker.disagg_spec is None and (
303 + tracker.skip_save
304 + or (tracker.num_saved_tokens > 0 and input_token_len < chunk_boundary)
305 + or (tracker.is_decode_phase and not save_decode_cache)
306 + or request_skip
307 + )
308 +
309 + if skip_save and load_spec is None:
310 + return None
311 +
312 + # Calculate number of tokens to save based on discard_partial_chunks
313 + # setting
314 +
315 + # NOTE(vladnosiv): for the input_token_len chunk prefill,
316 + # we are required to discard partial chunks,
317 + # as new tokens will be added in the next iteration.
318 + num_tokens_to_save = (
319 + (input_token_len // lmcache_chunk_size * lmcache_chunk_size)
320 + if not is_last_prefill or discard_partial_chunks
321 + else input_token_len
322 + )
323 +
324 + # If we need to save, update the number of saved tokens
325 + if not skip_save:
326 + tracker.num_saved_tokens = num_tokens_to_save
327 + save_spec = SaveSpec(skip_leading_tokens, not skip_save)
328 +
329 + # Calculate the token ids and slot mappings for load and save
330 + token_ids = input_token_ids[:num_tokens_to_save]
331 +
332 + # If the request has multimodal hashes, apply them to the token ids
333 + if tracker.mm_hashes:
334 + token_ids_tensor = torch.tensor(token_ids)
335 + assert tracker.mm_positions is not None, (
336 + "tracker got mm_hashes but no mm_positions"
337 + )
338 + apply_mm_hashes_to_token_ids(
339 + token_ids_tensor, tracker.mm_hashes, tracker.mm_positions
340 + )
341 + token_ids = token_ids_tensor.tolist()
342 +
343 + num_blocks = len(tracker.allocated_block_ids)
344 +
345 + if len(token_ids) > num_blocks * block_size:
346 + logger.error(
347 + "The number of tokens is more than the number of blocks."
348 + "Something might be wrong in scheduling logic!"
349 + )
350 + logger.error(
351 + "Num tokens: %d, num blocks: %d, block size: %d",
352 + len(token_ids),
353 + num_blocks,
354 + block_size,
355 + )
356 +
357 + block_ids = torch.tensor(tracker.allocated_block_ids, dtype=torch.long)
358 + block_offsets = torch.arange(0, block_size, dtype=torch.long)
359 + slot_mapping = (
360 + block_offsets.reshape((1, block_size))
361 + + block_ids.reshape((num_blocks, 1)) * block_size
362 + )
363 +
364 + slot_mapping = slot_mapping.flatten()[: len(token_ids)]
365 + assert slot_mapping.dtype == torch.long
366 +
367 + # For load operation: check whether the request is scheduled to load
368 + if load_spec is not None and load_spec.can_load:
369 + logger.debug(
370 + "Scheduled to load %d tokens for request %s",
371 + load_spec.lmcache_cached_tokens,
372 + tracker.req_id,
373 + )
374 + else:
375 + # Do not load if not in `can_load` state
376 + load_spec = None
377 +
378 + return ReqMeta(
379 + req_id=tracker.req_id,
380 + token_ids=token_ids,
381 + slot_mapping=slot_mapping,
382 + is_last_prefill=is_last_prefill,
383 + save_spec=save_spec,
384 + load_spec=load_spec,
385 + disagg_spec=tracker.disagg_spec,
386 + request_configs=tracker.request_configs,
387 + )
388 +
389 +
390 +def need_gpu_interm_buffer(lmcache_config: LMCacheEngineConfig):
391 + return lmcache_config.enable_pd
392 +
393 +
394 +def _calculate_mtp_layers(vllm_config, model_config):
395 + num_mtp_layers = 0
396 + if vllm_config is not None and vllm_config.speculative_config is not None:
397 + logger.info(
398 + "vllm_config.speculative_config: %s", vllm_config.speculative_config
399 + )
400 + # TODO(baoloongmao): Support other MTP methods
401 + if vllm_config.speculative_config.method == "deepseek_mtp":
402 + num_mtp_layers = getattr(
403 + model_config.hf_config, "num_nextn_predict_layers", 0
404 + )
405 + return num_mtp_layers
406 +
407 +
408 +def _init_lmcache_engine(
409 + lmcache_config: LMCacheEngineConfig,
410 + vllm_config: "VllmConfig",
411 +) -> LMCacheEngine:
412 + """Initialize the LMCache engine by the given model config and parallel
413 + config. This function will check the environment variable
414 + `LMCACHE_CONFIG_FILE` to load the configuration file. If that environment
415 + variable is not set, this function will return None.
416 +
417 + :param lmcache_config: The LMCache configuration.
418 + :type lmcache_config: LMCacheEngineConfig
419 + :param vllm_config: The vLLM configuration.
420 + :type vllm_config: VllmConfig
421 +
422 + :return: The initialized LMCache engine
423 + :rtype: LMCacheEngine
424 + """
425 + if curr_engine := LMCacheEngineBuilder.get(ENGINE_NAME):
426 + return curr_engine
427 +
428 + model_config = vllm_config.model_config
429 + parallel_config = vllm_config.parallel_config
430 + cache_config = vllm_config.cache_config
431 +
432 + assert isinstance(lmcache_config, LMCacheEngineConfig), (
433 + "LMCache v1 configuration is should be passed."
434 + )
435 +
436 + kv_dtype = get_kv_cache_torch_dtype(cache_config.cache_dtype, model_config.dtype)
437 +
438 + use_mla = mla_enabled(model_config)
439 + if use_mla and (
440 + lmcache_config.remote_serde != "naive"
441 + and lmcache_config.remote_serde is not None
442 + ):
443 + raise ValueError("MLA only works with naive serde mode..")
444 +
445 + # construct kv shape (for mem pool)
446 + num_layer = model_config.get_num_layers(parallel_config)
447 + num_mtp_layers = _calculate_mtp_layers(vllm_config, model_config)
448 + num_layer += num_mtp_layers
449 + chunk_size = lmcache_config.chunk_size
450 + num_kv_head = model_config.get_num_kv_heads(parallel_config)
451 + head_size = model_config.get_head_size()
452 + kv_shape = (num_layer, 1 if use_mla else 2, chunk_size, num_kv_head, head_size)
453 + logger.info(
454 + "use mla: %s, kv shape: %s, num_mtp_layers: %s",
455 + use_mla,
456 + kv_shape,
457 + num_mtp_layers,
458 + )
459 +
460 + # Change current device.
461 + num_gpus = torch.cuda.device_count()
462 + local_rank = parallel_config.rank % num_gpus
463 + torch.cuda.set_device(local_rank)
464 + device = torch.device(f"cuda:{local_rank}")
465 + metadata = LMCacheEngineMetadata(
466 + model_config.model,
467 + parallel_config.world_size,
468 + parallel_config.rank,
469 + "vllm",
470 + kv_dtype,
471 + kv_shape,
472 + use_mla,
473 + )
474 +
475 + use_gpu = need_gpu_interm_buffer(lmcache_config)
476 + vllm_gpu_connector: (
477 + VLLMBufferLayerwiseGPUConnector
478 + | VLLMPagedMemGPUConnectorV2
479 + | VLLMPagedMemLayerwiseGPUConnector
480 + )
481 +
482 + if use_mla and lmcache_config.use_layerwise:
483 + raise ValueError("layerwise MLA connector is not supported yet")
484 +
485 + # When use_mla is True, num_kv_head is 1
486 + hidden_dim_size = num_kv_head * head_size
487 + if lmcache_config.use_layerwise:
488 + if lmcache_config.enable_blending:
489 + # Use layerwise connector for blending
490 + vllm_gpu_connector = VLLMBufferLayerwiseGPUConnector(
491 + hidden_dim_size,
492 + num_layer,
493 + use_gpu=use_gpu,
494 + chunk_size=chunk_size,
495 + dtype=kv_dtype,
496 + device=device,
497 + )
498 + else:
499 + vllm_gpu_connector = VLLMPagedMemLayerwiseGPUConnector(
500 + hidden_dim_size,
501 + num_layer,
502 + use_gpu=use_gpu,
503 + chunk_size=chunk_size,
504 + dtype=kv_dtype,
505 + device=device,
506 + )
507 + else:
508 + vllm_gpu_connector = VLLMPagedMemGPUConnectorV2(
509 + hidden_dim_size,
510 + num_layer,
511 + use_gpu=use_gpu,
512 + chunk_size=chunk_size,
513 + dtype=kv_dtype,
514 + device=device,
515 + use_mla=use_mla,
516 + )
517 + tpg = get_tp_group()
518 + engine = LMCacheEngineBuilder.get_or_create(
519 + ENGINE_NAME,
520 + lmcache_config,
521 + metadata,
522 + vllm_gpu_connector,
523 + tpg.broadcast,
524 + tpg.broadcast_object,
525 + )
526 +
527 + return engine
528 +
529 +
530 +@dataclass
531 +class LMCacheConnectorMetadata(KVConnectorMetadata):
532 + requests: list[ReqMeta] = field(default_factory=list)
533 + lookup_requests_in_step: list[str] = field(default_factory=list)
534 +
535 + @_lmcache_nvtx_annotate
536 + def add_request(self, req_meta: ReqMeta) -> None:
537 + """Add a request to the metadata.
538 +
539 + Args:
540 + req_meta (ReqMeta): the request metadata.
541 + """
542 + self.requests.append(req_meta)
543 +
544 +
545 +class LMCacheConnectorV1Impl:
546 + def __init__(
547 + self,
548 + vllm_config: "VllmConfig",
549 + role: KVConnectorRole,
550 + parent: KVConnectorBase_V1,
551 + ):
552 + assert vllm_config.kv_transfer_config is not None
553 + self._parent = parent
554 + self._vllm_config = vllm_config
555 + self.kv_role = vllm_config.kv_transfer_config.kv_role
556 + self.worker_count = vllm_config.parallel_config.tensor_parallel_size
557 + config = lmcache_get_or_create_config()
558 + assert isinstance(config, LMCacheEngineConfig), (
559 + "LMCache v1 configuration is should be passed for vLLM v1."
560 + )
561 + # Put the leading with "lmcache." and matched configs from
562 + # vllm extra_config to the config
563 + kv_connector_extra_config = (
564 + vllm_config.kv_transfer_config.kv_connector_extra_config
565 + )
566 + if kv_connector_extra_config:
567 + for key, value in kv_connector_extra_config.items():
568 + if key.startswith("lmcache."):
569 + config_key = key[8:] # Remove "lmcache." prefix
570 + if _validate_and_set_config_value(config, config_key, value):
571 + logger.info(
572 + "Updated config %s from vLLM extra config: %s",
573 + config_key,
574 + value,
575 + )
576 +
577 + self.config = config
578 +
579 + self.async_loading = config.enable_async_loading
580 + self.layerwise_retrievers: list[Generator[torch.Tensor | None, None, None]] = []
581 + self._stats_monitor = LMCStatsMonitor.GetOrCreate()
582 + if role == KVConnectorRole.SCHEDULER:
583 + # Create lookup client using factory
584 + self.lookup_client = LookupClientFactory.create_lookup_client(
585 + vllm_config, config
586 + )
587 + self._unfinished_requests: dict[str, Request] = {}
588 + self._lookup_requests_in_step: list[str] = []
589 + self.lmcache_engine = None
590 + else:
591 + self.lmcache_engine = _init_lmcache_engine(
592 + config,
593 + vllm_config,
594 + )
595 +
596 + self.use_layerwise = config.use_layerwise
597 + self.enable_blending = config.enable_blending
598 +
599 + if self.enable_blending:
600 + self.blender = LMCBlenderBuilder.get_or_create(
601 + ENGINE_NAME,
602 + self.lmcache_engine,
603 + self.lmcache_engine.gpu_connector,
604 + config,
605 + )
606 +
607 + # Create lookup server using factory
608 + assert self.lmcache_engine is not None
609 + self.lookup_server = LookupClientFactory.create_lookup_server(
610 + self.lmcache_engine, vllm_config
611 + )
612 +
613 + self.offload_server = ZMQOffloadServer(
614 + self.lmcache_engine,
615 + vllm_config,
616 + get_tensor_model_parallel_rank(),
617 + )
618 +
619 + # In case of MLA, the lookup server is only created on worker 0
620 + if self.async_loading and self.lookup_server is not None:
621 + assert isinstance(self.lookup_server, LMCacheAsyncLookupServer)
622 + self.lmcache_engine.post_init(async_lookup_server=self.lookup_server)
623 +
624 + self.kv_caches: dict[str, torch.Tensor] = {}
625 +
626 + self._block_size = vllm_config.cache_config.block_size
627 +
628 + # request_id -> (vllm cached tokens, lmcache cached tokens)
629 + self.load_specs: dict[str, LoadSpec] = {}
630 +
631 + self.kv_cache_manager: KVCacheManager | None = None
632 +
633 + # request_id -> full_token_ids
634 + self._request_trackers: dict[str, RequestTracker] = {}
635 +
636 + # Whether to discard partial chunks
637 + self._discard_partial_chunks = (
638 + vllm_config.kv_transfer_config.get_from_extra_config(
639 + "discard_partial_chunks", False
640 + )
641 + or not config.save_unfull_chunk
642 + )
643 +
644 + self._lmcache_chunk_size = config.chunk_size
645 + self._save_decode_cache = config.save_decode_cache
646 +
647 + self.skip_last_n_tokens = vllm_config.kv_transfer_config.get_from_extra_config(
648 + "skip_last_n_tokens", 0
649 + )
650 +
651 + self.num_layers = vllm_config.model_config.get_num_layers(
652 + vllm_config.parallel_config
653 + )
654 + self.current_layer = 0
655 +
656 + self.force_skip_save = bool(os.environ.get("LMCACHE_FORCE_SKIP_SAVE", False))
657 +
658 + self._requests_priority: dict[str, int] = {}
659 +
660 + # TODO(baoloongmao): Internal api server & plugin framework support
661 + # dp > 1
662 + if (
663 + vllm_config.parallel_config.data_parallel_size_local == 1
664 + or vllm_config.parallel_config.data_parallel_rank_local == 0
665 + ):
666 + # Start internal API server if enabled
667 + # The enabled check is in the InternalAPIServer constructor
668 + self.api_server = InternalAPIServer(self)
669 + self.api_server.start()
670 + # Launch plugins
671 + self.plugin_launcher = PluginLauncher(
672 + self.config,
673 + role,
674 + self.worker_count,
675 + -1
676 + if self.lmcache_engine is None # scheduler side
677 + else self.lmcache_engine.metadata.worker_id,
678 + )
679 + self.plugin_launcher.launch_plugins()
680 + else:
681 + self.api_server = None # type: ignore[assignment]
682 + self.plugin_launcher = None # type: ignore[assignment]
683 + logger.info(
684 + "LMCache initialized for role %s with version %s, "
685 + "vllm version %s, lmcache cache_engine metadata: %s",
686 + role,
687 + utils.get_version(),
688 + VLLM_VERSION,
689 + getattr(self.lmcache_engine, "metadata", None),
690 + )
691 +
692 + def get_inference_info(self) -> dict:
693 + """Get inference information including vLLM config and related details.
694 +
695 + Returns:
696 + dict: Dictionary containing inference information
697 + """
698 + # Get vLLM config information
699 + vllm_config = self._vllm_config
700 +
701 + # Use vLLM config's string representation and add specific configs
702 + inference_info = {
703 + "vllm_version": VLLM_VERSION,
704 + "lmcache_version": utils.get_version(),
705 + "vllm_config": str(vllm_config),
706 + "model_config": {
707 + "model": getattr(vllm_config.model_config, "model", None),
708 + "dtype": str(getattr(vllm_config.model_config, "dtype", None)),
709 + "max_model_len": getattr(
710 + vllm_config.model_config, "max_model_len", None
711 + ),
712 + "vocab_size": getattr(vllm_config.model_config, "vocab_size", None),
713 + "num_layers": getattr(
714 + vllm_config.model_config, "get_num_layers", lambda _: None
715 + )(vllm_config.parallel_config),
716 + "num_attention_heads": getattr(
717 + vllm_config.model_config, "get_num_attention_heads", lambda _: None
718 + )(vllm_config.parallel_config),
719 + "num_kv_heads": getattr(
720 + vllm_config.model_config, "get_num_kv_heads", lambda _: None
721 + )(vllm_config.parallel_config),
722 + "head_size": getattr(
723 + vllm_config.model_config, "get_head_size", lambda: None
724 + )(),
725 + },
726 + "cache_config": {
727 + "block_size": getattr(vllm_config.cache_config, "block_size", None),
728 + "cache_dtype": str(
729 + getattr(vllm_config.cache_config, "cache_dtype", None)
730 + ),
731 + "gpu_memory_utilization": getattr(
732 + vllm_config.cache_config, "gpu_memory_utilization", None
733 + ),
734 + "swap_space": getattr(vllm_config.cache_config, "swap_space", None),
735 + "enable_prefix_caching": getattr(
736 + vllm_config.cache_config, "enable_prefix_caching", None
737 + ),
738 + },
739 + }
740 +
741 + return inference_info
742 +
743 + def get_inference_version(self) -> str:
744 + """Get vLLM version information.
745 +
746 + Returns:
747 + str: vLLM version string
748 + """
749 + return VLLM_VERSION
750 +
751 + @_lmcache_nvtx_annotate
752 + def _init_kv_caches_from_forward_context(self, forward_context: "ForwardContext"):
753 + for layer_name in forward_context.no_compile_layers:
754 + attn_layer = forward_context.no_compile_layers[layer_name]
755 + if not hasattr(attn_layer, "kv_cache"):
756 + logger.debug("The layer %s does not have kv_cache, skip it", layer_name)
757 + continue
758 +
759 + if layer_name not in self.kv_caches:
760 + self.kv_caches[layer_name] = attn_layer.kv_cache[
761 + forward_context.virtual_engine
762 + ]
763 +
764 + ####################
765 + # Worker side APIs
766 + ####################
767 +
768 + @_lmcache_nvtx_annotate
769 + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
770 + """Start loading the KV cache from the connector buffer to vLLM's
771 + paged KV buffer.
772 +
773 + Args:
774 + forward_context (ForwardContext): the forward context.
775 +
776 + Note:
777 + The number of elements in kv_caches and layer_names should be
778 + the same.
779 + """
780 + self.current_layer = 0
781 +
782 + if len(self.kv_caches) == 0:
783 + self._init_kv_caches_from_forward_context(forward_context)
784 +
785 + metadata = self._parent._get_connector_metadata()
786 + assert isinstance(metadata, LMCacheConnectorMetadata)
787 +
788 + assert len(self.kv_caches) > 0
789 + kvcaches = list(self.kv_caches.values())
790 +
791 + attn_metadata = forward_context.attn_metadata
792 + if attn_metadata is None:
793 + logger.debug("In connector.start_load_kv, but the attn_metadata is None")
794 + return
795 +
796 + assert self.lmcache_engine is not None
797 +
798 + self.lmcache_engine.post_init(kvcaches=kvcaches)
799 +
800 + self.layerwise_retrievers = []
801 +
802 + for idx, request in enumerate(metadata.requests):
803 + if request.load_spec is None:
804 + continue
805 + last_idx = idx
806 +
807 + for idx, request in enumerate(metadata.requests):
808 + if request.load_spec is None:
809 + continue
810 +
811 + tokens = request.token_ids
812 + # TODO: have a pre-allocated buffer to hold the slot_mappings
813 + slot_mapping = request.slot_mapping.cuda()
814 + assert len(tokens) == len(slot_mapping)
815 +
816 + self._stats_monitor.update_interval_vllm_hit_tokens(
817 + request.load_spec.vllm_cached_tokens
818 + )
819 + token_mask = torch.ones(len(tokens), dtype=torch.bool)
820 + masked_token_count = (
821 + request.load_spec.vllm_cached_tokens
822 + // self._lmcache_chunk_size
823 + * self._lmcache_chunk_size
824 + )
825 + token_mask[:masked_token_count] = False
826 +
827 + lmcache_cached_tokens = request.load_spec.lmcache_cached_tokens
828 + if self.use_layerwise:
829 + sync = idx == last_idx
830 + # NOTE(Jiayi): Perform blending before layerwise prefix caching
831 + if self.enable_blending:
832 + # TODO(Jiayi): Need to make prefix caching and blending
833 + # compatible
834 + self.blender.blend(
835 + tokens[:lmcache_cached_tokens],
836 + token_mask[:lmcache_cached_tokens],
837 + kvcaches=kvcaches,
838 + slot_mapping=slot_mapping[:lmcache_cached_tokens],
839 + )
840 + else:
841 + layerwise_retriever = self.lmcache_engine.retrieve_layer(
842 + tokens[:lmcache_cached_tokens],
843 + token_mask[:lmcache_cached_tokens],
844 + kvcaches=kvcaches,
845 + slot_mapping=slot_mapping[:lmcache_cached_tokens],
846 + sync=sync,
847 + )
848 + # NOTE: retrieve for two layers at the first layer
849 + next(layerwise_retriever)
850 + next(layerwise_retriever)
851 + self.layerwise_retrievers.append(layerwise_retriever)
852 + else:
853 + ret_token_mask = self.lmcache_engine.retrieve(
854 + tokens[:lmcache_cached_tokens],
855 + token_mask[:lmcache_cached_tokens],
856 + kvcaches=kvcaches,
857 + slot_mapping=slot_mapping[:lmcache_cached_tokens],
858 + request_configs=request.request_configs,
859 + req_id=request.req_id,
860 + )
861 +
862 + # Check the result
863 + num_retrieved_tokens = ret_token_mask.sum().item()
864 + num_expected_tokens = (
865 + lmcache_cached_tokens - request.load_spec.vllm_cached_tokens
866 + )
867 + if num_retrieved_tokens < num_expected_tokens:
868 + logger.error(
869 + "The number of retrieved tokens is less than the "
870 + "expected number of tokens! This should not happen!"
871 + )
872 + logger.error(
873 + "Num retrieved tokens: %d, num expected tokens: %d",
874 + num_retrieved_tokens,
875 + num_expected_tokens,
876 + )
877 +
878 + @_lmcache_nvtx_annotate
879 + def wait_for_layer_load(self, layer_name: str) -> None:
880 + """Blocking until the KV for a specific layer is loaded into vLLM's
881 + paged buffer.
882 +
883 + This interface will be useful for layer-by-layer pipelining.
884 +
885 + Args:
886 + layer_name: the name of that layer
887 + """
888 + if self.layerwise_retrievers:
889 + logger.debug("Waiting for layer %s to be loaded", self.current_layer)
890 +
891 + # Wait for the layer to be loaded
892 + for layerwise_retriever in self.layerwise_retrievers:
893 + ret_token_mask = next(layerwise_retriever)
894 +
895 + if self.current_layer == self.num_layers - 1:
896 + assert ret_token_mask is not None
897 + num_retrieved_tokens = ret_token_mask.sum().item()
898 + logger.info("Retrieved %s tokens", num_retrieved_tokens)
899 +
900 + return
901 +
902 + @_lmcache_nvtx_annotate
903 + def save_kv_layer(
904 + self,
905 + layer_name: str,
906 + kv_layer: torch.Tensor,
907 + attn_metadata: "AttentionMetadata",
908 + **kwargs,
909 + ) -> None:
910 + """Start saving the a layer of KV cache from vLLM's paged buffer
911 + to the connector.
912 +
913 + Args:
914 + layer_name (str): the name of the layer.
915 + kv_layer (torch.Tensor): the paged KV buffer of the current
916 + layer in vLLM.
917 + attn_metadata (AttentionMetadata): the attention metadata.
918 + """
919 + assert self.lmcache_engine is not None
920 +
921 + if not self.use_layerwise:
922 + return
923 +
924 + if self.kv_role == "kv_consumer":
925 + # Don't do save if the role is kv_consumer
926 + return
927 + if self._parent._connector_metadata is None:
928 + logger.warning(
929 + "In connector.save_kv_layer, but the connector metadata is None"
930 + )
931 + return
932 + connector_metadata = self._parent._get_connector_metadata()
933 + assert isinstance(connector_metadata, LMCacheConnectorMetadata)
934 +
935 + assert len(self.kv_caches) > 0
936 +
937 + kvcaches = list(self.kv_caches.values())
938 + if self.current_layer == 0:
939 + self.layerwise_storers = []
940 +
941 + is_first = True
942 +
943 + for idx, request in enumerate(connector_metadata.requests):
944 + save_spec = request.save_spec
945 + if save_spec is None or not save_spec.can_save:
946 + continue
947 +
948 + token_ids = request.token_ids
949 + assert isinstance(token_ids, list)
950 +
951 + slot_mapping = request.slot_mapping
952 + assert isinstance(slot_mapping, torch.Tensor)
953 + assert len(slot_mapping) == len(token_ids)
954 +
955 + # TODO: have a pre-allocated buffer to hold the slot_mappings
956 + slot_mapping = slot_mapping.cuda()
957 +
958 + if self.kv_role == "kv_producer":
959 + skip_leading_tokens = 0
960 + else:
961 + skip_leading_tokens = save_spec.skip_leading_tokens
962 +
963 + if skip_leading_tokens == len(token_ids):
964 + continue # skip this request
965 + # Align to lmcache chunk size
966 + skip_leading_tokens = (
967 + skip_leading_tokens
968 + // self._lmcache_chunk_size
969 + * self._lmcache_chunk_size
970 + )
971 +
972 + store_mask = torch.ones(len(token_ids), dtype=torch.bool)
973 + store_mask[:skip_leading_tokens] = False
974 +
975 + logger.info(
976 + "Storing KV cache for %d out of %d tokens "
977 + "(skip_leading_tokens=%d) for request %s",
978 + len(token_ids) - skip_leading_tokens,
979 + len(token_ids),
980 + skip_leading_tokens,
981 + request.req_id,
982 + )
983 +
984 + # TODO (Jiayi): need to make layerwise storing
985 + # compatible with disagg spec
986 + layerwise_storer = self.lmcache_engine.store_layer(
987 + token_ids,
988 + mask=store_mask,
989 + kvcaches=kvcaches,
990 + slot_mapping=slot_mapping,
991 + offset=skip_leading_tokens,
992 + sync=is_first,
993 + )
994 + self.layerwise_storers.append(layerwise_storer)
995 + if is_first:
996 + is_first = False
997 +
998 + for layerwise_storer in self.layerwise_storers:
999 + next(layerwise_storer)
1000 +
1001 + self.current_layer += 1
1002 +
1003 + @_lmcache_nvtx_annotate
1004 + def wait_for_save(self):
1005 + """Blocking until the KV cache is saved to the connector buffer."""
1006 +
1007 + connector_metadata = self._parent._get_connector_metadata()
1008 + assert isinstance(connector_metadata, LMCacheConnectorMetadata)
1009 +
1010 + self.lmcache_engine.lookup_unpin( # type: ignore
1011 + connector_metadata.lookup_requests_in_step
1012 + )
1013 +
1014 + if self.kv_role == "kv_consumer":
1015 + # Don't do save if the role is kv_consumer
1016 + return
1017 +
1018 + if self.use_layerwise:
1019 + for layerwise_storer in self.layerwise_storers:
1020 + next(layerwise_storer)
1021 + return
1022 +
1023 + assert len(self.kv_caches) > 0
1024 + kvcaches = list(self.kv_caches.values())
1025 +
1026 + assert self.lmcache_engine is not None
1027 +
1028 + for request in connector_metadata.requests:
1029 + save_spec = request.save_spec
1030 + if (
1031 + save_spec is None or not save_spec.can_save
1032 + ) and self.kv_role != "kv_producer":
1033 + continue
1034 +
1035 + token_ids = request.token_ids
1036 +
1037 + slot_mapping = request.slot_mapping
1038 + assert isinstance(slot_mapping, torch.Tensor)
1039 + assert len(slot_mapping) == len(token_ids)
1040 + assert save_spec is not None
1041 +
1042 + # TODO: have a pre-allocated buffer to hold the slot_mappings
1043 + slot_mapping = slot_mapping.cuda()
1044 +
1045 + skip_leading_tokens = save_spec.skip_leading_tokens
1046 + if self.kv_role == "kv_producer":
1047 + assert request.disagg_spec is not None
1048 + skip_leading_tokens = min(
1049 + skip_leading_tokens, request.disagg_spec.num_transferred_tokens
1050 + )
1051 +
1052 + if skip_leading_tokens == len(token_ids):
1053 + continue # skip this request
1054 + # Align to lmcache chunk size
1055 + skip_leading_tokens = (
1056 + skip_leading_tokens
1057 + // self._lmcache_chunk_size
1058 + * self._lmcache_chunk_size
1059 + )
1060 +
1061 + store_mask = torch.ones(len(token_ids), dtype=torch.bool)
1062 + store_mask[:skip_leading_tokens] = False
1063 +
1064 + logger.info(
1065 + "Storing KV cache for %d out of %d tokens "
1066 + "(skip_leading_tokens=%d) for request %s",
1067 + len(token_ids) - skip_leading_tokens,
1068 + len(token_ids),
1069 + skip_leading_tokens,
1070 + request.req_id,
1071 + )
1072 +
1073 + is_last_prefill = request.is_last_prefill
1074 + if is_last_prefill:
1075 + if request.disagg_spec:
1076 + request.disagg_spec.is_last_prefill = True
1077 + else:
1078 + token_len = len(token_ids)
1079 + aligned_token_len = (
1080 + token_len // self._lmcache_chunk_size * self._lmcache_chunk_size
1081 + )
1082 + token_ids = token_ids[:aligned_token_len]
1083 + store_mask = store_mask[:aligned_token_len]
1084 + slot_mapping = slot_mapping[:aligned_token_len]
1085 +
1086 + self.lmcache_engine.store(
1087 + token_ids,
1088 + mask=store_mask,
1089 + kvcaches=kvcaches,
1090 + slot_mapping=slot_mapping,
1091 + offset=skip_leading_tokens,
1092 + transfer_spec=request.disagg_spec,
1093 + request_configs=request.request_configs,
1094 + )
1095 +
1096 + # NOTE(Jiayi): We assume all tokens are saved
1097 + save_spec.skip_leading_tokens = len(token_ids)
1098 + if request.disagg_spec:
1099 + request.disagg_spec.num_transferred_tokens = len(token_ids)
1100 +
1101 + @_lmcache_nvtx_annotate
1102 + def get_finished(
1103 + self, finished_req_ids: set[str]
1104 + ) -> tuple[set[str] | None, set[str] | None]:
1105 + return None, None
1106 +
1107 + ###################
1108 + # Scheduler side APIs
1109 + ####################
1110 +
1111 + @_lmcache_nvtx_annotate
1112 + def get_num_new_matched_tokens(
1113 + self,
1114 + request: "Request",
1115 + num_computed_tokens: int,
1116 + ) -> int | None:
1117 + """
1118 + Check for external KV cache hit.
1119 +
1120 + Args:
1121 + request (Request): the request object.
1122 + num_computed_tokens (int): the number of locally
1123 + computed tokens for this request
1124 +
1125 + Returns:
1126 + the number of tokens that can be loaded from the
1127 + external KV cache beyond what is already computed.
1128 + """
1129 + if self.kv_role == "kv_producer" and not hasattr(
1130 + self.lookup_client, "supports_producer_reuse"
1131 + ):
1132 + return 0
1133 +
1134 + self._requests_priority[request.request_id] = request.priority
1135 +
1136 + token_ids = request.prompt_token_ids
1137 +
1138 + # If the request has multimodal hashes, apply them to the token ids
1139 + mm_hashes, mm_positions = extract_mm_features(request)
1140 + if mm_hashes and mm_positions:
1141 + # TODO(Jiayi): Optimize this
1142 + token_ids_tensor = torch.tensor(request.prompt_token_ids)
1143 + apply_mm_hashes_to_token_ids(token_ids_tensor, mm_hashes, mm_positions)
1144 + token_ids = token_ids_tensor.tolist()
1145 +
1146 + if request.sampling_params:
1147 + request_configs = extract_request_configs(request.sampling_params)
1148 + else:
1149 + request_configs = None
1150 +
1151 + if self.skip_last_n_tokens > 0:
1152 + assert token_ids is not None
1153 + token_ids = token_ids[: -self.skip_last_n_tokens]
1154 + lookup_id = request.request_id if self.async_loading else str(uuid.uuid4())
1155 +
1156 + self._lookup_requests_in_step.append(lookup_id)
1157 +
1158 + num_external_hit_tokens = self.lookup_client.lookup(
1159 + token_ids,
1160 + lookup_id=lookup_id,
1161 + request_configs=request_configs,
1162 + )
1163 +
1164 + if num_external_hit_tokens is None:
1165 + logger.info(
1166 + "Reqid: %s, Total tokens %d, LMCache hit tokens: None.",
1167 + request.request_id,
1168 + request.num_tokens,
1169 + )
1170 + return None
1171 +
1172 + # When prompt length is divisible by the block size and all
1173 + # blocks are cached, we need to recompute the last token.
1174 + # This will be removed in the future if vLLM's scheduler provides
1175 + # a better support for this case.
1176 + need_to_allocate = num_external_hit_tokens - num_computed_tokens
1177 +
1178 + # In, full-prompt-hit case, we need to recompute the last token
1179 + if num_external_hit_tokens == request.num_tokens:
1180 + need_to_allocate -= 1
1181 +
1182 + logger.info(
1183 + "Reqid: %s, Total tokens %d, LMCache hit tokens: %d, need to load: %d",
1184 + request.request_id,
1185 + request.num_tokens,
1186 + num_external_hit_tokens,
1187 + need_to_allocate,
1188 + )
1189 +
1190 + self.load_specs[request.request_id] = LoadSpec(
1191 + vllm_cached_tokens=num_computed_tokens,
1192 + lmcache_cached_tokens=num_external_hit_tokens,
1193 + can_load=False,
1194 + )
1195 +
1196 + if need_to_allocate <= 0:
1197 + return 0
1198 +
1199 + return need_to_allocate
1200 +
1201 + @_lmcache_nvtx_annotate
1202 + def update_state_after_alloc(self, request: "Request", num_external_tokens: int):
1203 + """
1204 + Update KVConnector state after temporary buffer alloc.
1205 +
1206 + For SharedStorageConnector, update _request_needs_load
1207 + if the CacheManager this allocated blocks for us.
1208 + """
1209 +
1210 + kv_transfer_params = (
1211 + request.kv_transfer_params
1212 + if hasattr(request, "kv_transfer_params")
1213 + else None
1214 + )
1215 +
1216 + if kv_transfer_params is not None and "disagg_spec" in kv_transfer_params:
1217 + req_disagg_spec = kv_transfer_params["disagg_spec"]
1218 +
1219 + receiver_id = req_disagg_spec["receiver_host"] + str(
1220 + req_disagg_spec["receiver_init_port"]
1221 + )
1222 +
1223 + disagg_spec = DisaggSpec(
1224 + req_id=req_disagg_spec["req_id"],
1225 + receiver_id=receiver_id,
1226 + receiver_host=req_disagg_spec["receiver_host"],
1227 + receiver_init_port=req_disagg_spec["receiver_init_port"],
1228 + receiver_alloc_port=req_disagg_spec["receiver_alloc_port"],
1229 + )
1230 +
1231 + tmp_disagg_tracker[request.request_id] = disagg_spec
1232 + self._unfinished_requests[request.request_id] = request
1233 +
1234 + if request.request_id not in self.load_specs:
1235 + # No KV tokens from external KV cache, return
1236 + return
1237 +
1238 + if num_external_tokens == 0:
1239 + # No need to load anything
1240 + self.load_specs[request.request_id].can_load = False
1241 + return
1242 +
1243 + # Only check for non-prompt-hit case
1244 + if (
1245 + self.load_specs[request.request_id].lmcache_cached_tokens
1246 + != request.num_tokens
1247 + ):
1248 + assert (
1249 + num_external_tokens > 0
1250 + and num_external_tokens
1251 + == self.load_specs[request.request_id].lmcache_cached_tokens
1252 + - self.load_specs[request.request_id].vllm_cached_tokens
1253 + ), (
1254 + f"Mismatch in number of tokens: {num_external_tokens} vs "
1255 + f"{self.load_specs[request.request_id].lmcache_cached_tokens} -"
1256 + f" {self.load_specs[request.request_id].vllm_cached_tokens}"
1257 + f" for request {request.request_id}"
1258 + )
1259 +
1260 + self.load_specs[request.request_id].can_load = True
1261 +
1262 + @_lmcache_nvtx_annotate
1263 + def build_connector_meta(
1264 + self, scheduler_output: SchedulerOutput
1265 + ) -> KVConnectorMetadata:
1266 + """Attach the connector metadata to the request object.
1267 +
1268 + This function should NOT modify other fields in the scheduler_output
1269 + except the `kv_connector_metadata` field.
1270 + Also, calling this function will reset the state of the connector.
1271 +
1272 + Args:
1273 + scheduler_output (SchedulerOutput): the scheduler output object.
1274 + """
1275 +
1276 + force_skip_save = self.kv_role == "kv_consumer" or self.force_skip_save
1277 +
1278 + meta = LMCacheConnectorMetadata()
1279 +
1280 + # set and update lookup requests for unpin
1281 + meta.lookup_requests_in_step = self._lookup_requests_in_step
1282 + self._lookup_requests_in_step = []
1283 +
1284 + for finished_req_id in scheduler_output.finished_req_ids:
1285 + self._request_trackers.pop(finished_req_id, None)
1286 + self._unfinished_requests.pop(finished_req_id, None)
1287 +
1288 + for request in scheduler_output.scheduled_new_reqs:
1289 + # Right now, we only load KV for new requests
1290 + load_spec = self.load_specs.pop(request.req_id, None)
1291 + num_tokens_to_compute = (
1292 + request.num_computed_tokens
1293 + + scheduler_output.num_scheduled_tokens[request.req_id]
1294 + )
1295 + lmcache_cached_tokens = 0
1296 + if load_spec is not None:
1297 + lmcache_cached_tokens = load_spec.lmcache_cached_tokens
1298 + request_priority = self._requests_priority.pop(request.req_id, 0)
1299 +
1300 + skip_save = force_skip_save or (
1301 + self.config.priority_limit is not None
1302 + and request_priority > self.config.priority_limit
1303 + )
1304 +
1305 + request_tracker = RequestTracker.from_new_request(
1306 + self.config,
1307 + request,
1308 + num_tokens_to_compute,
1309 + lmcache_cached_tokens,
1310 + skip_save,
1311 + )
1312 + self._request_trackers[request.req_id] = request_tracker
1313 +
1314 + req_meta = ReqMeta.from_request_tracker(
1315 + request_tracker,
1316 + self._block_size,
1317 + self._lmcache_chunk_size,
1318 + load_spec=load_spec,
1319 + discard_partial_chunks=self._discard_partial_chunks,
1320 + save_decode_cache=self._save_decode_cache,
1321 + )
1322 + if req_meta is not None:
1323 + meta.add_request(req_meta)
1324 +
1325 + cached_reqs = scheduler_output.scheduled_cached_reqs
1326 +
1327 + # NOTE: For backward compatibility with vllm version < 0.9.2,
1328 + # In the latest vllm version, the type of scheduled_cached_reqs has
1329 + # changed from list to object `CachedRequestData`
1330 + if isinstance(cached_reqs, list):
1331 + for i, req in enumerate(cached_reqs):
1332 + request_tracker = self._request_trackers[req.req_id]
1333 + request_tracker.update(req.new_token_ids, req.new_block_ids)
1334 +
1335 + req_meta = ReqMeta.from_request_tracker(
1336 + request_tracker,
1337 + self._block_size,
1338 + self._lmcache_chunk_size,
1339 + load_spec=None,
1340 + discard_partial_chunks=self._discard_partial_chunks,
1341 + )
1342 + if req_meta is not None:
1343 + meta.add_request(req_meta)
1344 + return meta
1345 +
1346 + for i, req_id in enumerate(cached_reqs.req_ids):
1347 + request_tracker = self._request_trackers[req_id]
1348 + num_new_tokens = scheduler_output.num_scheduled_tokens[req_id]
1349 + if cached_request := self._unfinished_requests.get(req_id):
1350 + num_current_tokens = len(request_tracker.token_ids)
1351 + new_token_ids = cached_request.all_token_ids[
1352 + num_current_tokens : num_current_tokens + num_new_tokens
1353 + ]
1354 + else:
1355 + raise ValueError(
1356 + f"Request {req_id} is not in _unfinished_requests, "
1357 + f"but it is scheduled to be cached"
1358 + )
1359 + new_block_ids = cached_reqs.new_block_ids[i]
1360 +
1361 + request_tracker.update(new_token_ids, new_block_ids)
1362 +
1363 + req_meta = ReqMeta.from_request_tracker(
1364 + request_tracker,
1365 + self._block_size,
1366 + self._lmcache_chunk_size,
1367 + load_spec=None,
1368 + discard_partial_chunks=self._discard_partial_chunks,
1369 + save_decode_cache=self._save_decode_cache,
1370 + )
1371 + if req_meta is not None:
1372 + meta.add_request(req_meta)
1373 +
1374 + return meta
1375 +
1376 + @_lmcache_nvtx_annotate
1377 + def request_finished(
1378 + self,
1379 + request: "Request",
1380 + block_ids: list[int],
1381 + ) -> tuple[bool, dict[str, Any] | None]:
1382 + params = (
1383 + request.kv_transfer_params
1384 + if hasattr(request, "kv_transfer_params")
1385 + else None
1386 + )
1387 + return_params = None
1388 +
1389 + # NOTE: Used to stream back the first token
1390 + # for disagg prefill
1391 + if params is not None and "ret_first_tok" in params:
1392 + return_params = {
1393 + "first_tok": request._output_token_ids[0],
1394 + }
1395 +
1396 + return False, return_params