Encoder cache extension hooks (#48218)
Signed-off-by: hotTea <958436561@qq.com> Signed-off-by: hanxi-java <634498162@qq.com> Co-authored-by: hanxi-java <634498162@qq.com> Co-authored-by: 韩熙 <63780107+hanxi-java@users.noreply.github.com>
hotTea committed
Jul 24, 2026 at 17:52 UTC
833483f3578a2c4766977d0fe431bd32a01ce250
8 files changed
+106
-15
tests/v1/worker/test_gpu_model_runner_mm_gather.py
+1
@@ -46,6 +46,7 @@ def _gather(features, cached, *, num_scheduled, shift, num_computed=0):
46
input_batch=SimpleNamespace(req_ids=["req0"]),
47
requests={"req0": req_state},
48
encoder_cache=encoder_cache,
49
+ _get_encoder_output_from_cache=lambda mm_hash: encoder_cache.get(mm_hash),
50
is_multimodal_pruning_enabled=False,
51
uses_mrope=False,
52
)
vllm/config/__init__.py
+3
@@ -11,6 +11,7 @@ from vllm.config.compilation import (
11
)
12
from vllm.config.device import DeviceConfig
13
from vllm.config.diffusion import DiffusionConfig
14
+from vllm.config.ec_manager_config import EncoderCacheManagerConfig
15
from vllm.config.ec_transfer import ECTransferConfig
16
from vllm.config.kernel import KernelConfig
17
from vllm.config.kv_events import KVEventsConfig
@@ -75,6 +76,8 @@ __all__ = [
76
"DeviceConfig",
77
# From vllm.config.diffusion
78
"DiffusionConfig",
79
+ # From vllm.config.ec_manager_config
80
+ "EncoderCacheManagerConfig",
81
# From vllm.config.ec_transfer
82
"ECTransferConfig",
83
# From vllm.config.kernel
vllm/config/ec_manager_config.py
new
+29
@@ -0,0 +1,29 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC
5
+
6
+from vllm.config.utils import config
7
+from vllm.utils.import_utils import resolve_obj_by_qualname
8
+
9
+
10
+@config
11
+class EncoderCacheManagerConfig:
12
+ encoder_cache_manager_cls: str | None = None
13
+ """The qualified class name of the custom encoder cache manager.
14
+ """
15
+
16
+ def get_encoder_cache_manager_obj(self):
17
+ cls_path = self.encoder_cache_manager_cls
18
+ if cls_path is None:
19
+ return None
20
+ return resolve_obj_by_qualname(cls_path)
21
+
22
+
23
+class EncoderCacheManagerMetadata(ABC): # noqa: B024
24
+ """
25
+ Abstract Metadata used to communicate between the
26
+ Scheduler EncoderCacheManager and Worker EncoderCacheManager.
27
+ """
28
+
29
+ pass
vllm/config/vllm.py
+5
@@ -32,6 +32,7 @@ from .cache import CacheConfig
32
from .compilation import CompilationConfig, CompilationMode, CUDAGraphMode
33
from .device import DeviceConfig
34
from .diffusion import DiffusionConfig
35
+from .ec_manager_config import EncoderCacheManagerConfig
36
from .ec_transfer import ECTransferConfig
37
from .kernel import KernelConfig
38
from .kv_events import KVEventsConfig
@@ -367,6 +368,10 @@ class VllmConfig:
368
"""The configurations for event publishing."""
369
ec_transfer_config: ECTransferConfig | None = None
370
"""The configurations for distributed EC cache transfer."""
371
+ ec_manager_config: EncoderCacheManagerConfig = Field(
372
+ default_factory=EncoderCacheManagerConfig
373
+ )
374
+ """The configurations for custom encoder cache manager."""
375
reasoning_config: ReasoningConfig | None = None
376
"""The configurations for reasoning model."""
377
# some opaque config, only used to provide additional information
vllm/v1/core/encoder_cache_manager.py
+4
@@ -5,6 +5,7 @@ from collections import OrderedDict
5
from collections.abc import Mapping
6
from typing import TYPE_CHECKING
7
8
+from vllm.config.ec_manager_config import EncoderCacheManagerMetadata
9
from vllm.logger import init_logger
10
from vllm.v1.request import Request
11
@@ -265,6 +266,9 @@ class EncoderCacheManager:
266
self.freed = []
267
return freed
268
269
+ def get_manager_metadata(self) -> EncoderCacheManagerMetadata | None:
270
+ return None
271
+
272
273
def compute_mm_encoder_budget(
274
scheduler_config: "SchedulerConfig",
vllm/v1/core/sched/output.py
+4
-1
@@ -5,6 +5,8 @@ from dataclasses import dataclass
5
from functools import cached_property
6
from typing import TYPE_CHECKING
7
8
+from vllm.config.ec_manager_config import EncoderCacheManagerMetadata
9
+
10
if TYPE_CHECKING:
11
import numpy as np
12
import numpy.typing as npt
@@ -246,7 +248,8 @@ class SchedulerOutput:
248
249
# EC Cache Connector metadata
250
ec_connector_metadata: ECConnectorMetadata | None = None
249
-
251
+ # EC Cache Manager metadata
252
+ ec_manager_metadata: EncoderCacheManagerMetadata | None = None
253
# Block IDs freshly allocated from the pool during this scheduling step.
254
# The worker zeros the corresponding GPU memory before the blocks are used,
255
# preventing stale NaN/data from corrupting attention or SSM computation.
vllm/v1/core/sched/scheduler.py
+10
-6
@@ -227,12 +227,15 @@ class Scheduler(SchedulerInterface):
227
mm_budget.encoder_compute_budget if mm_budget else 0
228
)
229
encoder_cache_size = mm_budget.encoder_cache_size if mm_budget else 0
230
- self.encoder_cache_manager = (
231
- EncoderDecoderCacheManager(cache_size=encoder_cache_size)
232
- if self.is_encoder_decoder
233
- else EncoderCacheManager(cache_size=encoder_cache_size)
234
- )
235
-
230
+ manager_cls_obj = vllm_config.ec_manager_config.get_encoder_cache_manager_obj()
231
+ if manager_cls_obj is not None:
232
+ self.encoder_cache_manager = manager_cls_obj(cache_size=encoder_cache_size)
233
+ else:
234
+ self.encoder_cache_manager = (
235
+ EncoderDecoderCacheManager(cache_size=encoder_cache_size)
236
+ if self.is_encoder_decoder
237
+ else EncoderCacheManager(cache_size=encoder_cache_size)
238
+ )
239
speculative_config = vllm_config.speculative_config
240
self.use_eagle = False
241
self.num_spec_tokens = vllm_config.num_speculative_tokens
@@ -1151,6 +1154,7 @@ class Scheduler(SchedulerInterface):
1154
new_block_ids_to_zero=self._get_new_block_ids_to_zero(),
1155
kv_cache_block_copies=pending_kv_cache_block_copies,
1156
num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
1157
+ ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(),
1158
)
1159
1160
# NOTE(Kuntai): this function is designed for multiple purposes:
vllm/v1/worker/gpu_model_runner.py
+50
-8
@@ -37,6 +37,7 @@ from vllm.config import (
37
update_config,
38
)
39
from vllm.config.cache import CacheConfig
40
+from vllm.config.ec_manager_config import EncoderCacheManagerMetadata
41
from vllm.config.model import PROCESSED_LOGPROBS_MODES
42
from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
43
from vllm.distributed.eplb.eplb_state import EplbState
@@ -1164,6 +1165,22 @@ class GPUModelRunner(
1165
self.async_output_copy_stream = stream
1166
return stream
1167
1168
+ def _on_request_state_removed(
1169
+ self,
1170
+ req_id: str,
1171
+ req_state: CachedRequestState | None,
1172
+ ) -> None:
1173
+ """Hook for platform runners to clean request-scoped side caches."""
1174
+ del req_id, req_state
1175
+
1176
+ def _process_encoder_cache_scheduler_output(
1177
+ self,
1178
+ scheduler_output: "SchedulerOutput",
1179
+ ) -> None:
1180
+ """Apply scheduler-side encoder cache lifecycle updates."""
1181
+ for mm_hash in scheduler_output.free_encoder_mm_hashes:
1182
+ self.encoder_cache.pop(mm_hash, None)
1183
+
1184
def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None:
1185
"""Update the cached states and the persistent batch with the scheduler
1186
output.
@@ -1176,7 +1193,8 @@ class GPUModelRunner(
1193
"""
1194
# Remove finished requests from the cached states.
1195
for req_id in scheduler_output.finished_req_ids:
1179
- self.requests.pop(req_id, None)
1196
+ req_state = self.requests.pop(req_id, None)
1197
+ self._on_request_state_removed(req_id, req_state)
1198
self.num_prompt_logprobs.pop(req_id, None)
1199
self.late_interaction_runner.on_requests_finished(
1200
scheduler_output.finished_req_ids
@@ -1202,8 +1220,7 @@ class GPUModelRunner(
1220
)
1221
1222
# Free the cached encoder outputs.
1205
- for mm_hash in scheduler_output.free_encoder_mm_hashes:
1206
- self.encoder_cache.pop(mm_hash, None)
1223
+ self._process_encoder_cache_scheduler_output(scheduler_output)
1224
1225
# Remove the unscheduled requests from the persistent batch.
1226
# NOTE(woosuk): The unscheduled requests are either preempted requests
@@ -2958,6 +2975,18 @@ class GPUModelRunner(
2975
2976
return mm_hashes, mm_kwargs, mm_lora_refs
2977
2978
+ def _cache_encoder_output(
2979
+ self,
2980
+ mm_hash: str,
2981
+ output: torch.Tensor,
2982
+ ec_manager_metadata: "EncoderCacheManagerMetadata | None",
2983
+ free_encoder_mm_hashes: list[str],
2984
+ ) -> None:
2985
+ """Store an encoder output for later multimodal embedding gather."""
2986
+ del ec_manager_metadata, free_encoder_mm_hashes
2987
+ self.encoder_cache[mm_hash] = output
2988
+ self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash)
2989
+
2990
def _execute_mm_encoder(
2991
self, scheduler_output: "SchedulerOutput"
2992
) -> list[torch.Tensor]:
@@ -2983,8 +3012,12 @@ class GPUModelRunner(
3012
pe_tensor = mm_kwargs[i][1]["embedding"].data
3013
assert isinstance(pe_tensor, torch.Tensor)
3014
2986
- self.encoder_cache[mm_hashes[i]] = pe_tensor.to(self.device)
2987
- self.maybe_save_ec_to_connector(self.encoder_cache, mm_hashes[i])
3015
+ self._cache_encoder_output(
3016
+ mm_hashes[i],
3017
+ pe_tensor.to(self.device),
3018
+ scheduler_output.ec_manager_metadata,
3019
+ scheduler_output.free_encoder_mm_hashes,
3020
+ )
3021
# Filter out `prompt_embeds` items from mm_kwargs/mm_hashes/mm_lora_refs
3022
# since they don't require further encoder processing.
3023
mm_hashes = [h for i, h in enumerate(mm_hashes) if i not in pe_indices]
@@ -3161,12 +3194,21 @@ class GPUModelRunner(
3194
3195
# Cache the encoder outputs by mm_hash
3196
for mm_hash, output in zip(mm_hashes, encoder_outputs):
3164
- self.encoder_cache[mm_hash] = output
3197
+ self._cache_encoder_output(
3198
+ mm_hash,
3199
+ output,
3200
+ scheduler_output.ec_manager_metadata,
3201
+ scheduler_output.free_encoder_mm_hashes,
3202
+ )
3203
logger.debug("Finish execute for mm hash %s", mm_hash)
3166
- self.maybe_save_ec_to_connector(self.encoder_cache, mm_hash)
3204
3205
return encoder_outputs
3206
3207
+ def _get_encoder_output_from_cache(self, mm_hash: str) -> torch.Tensor | None:
3208
+ """Return a cached encoder output for multimodal
3209
+ embedding gather."""
3210
+ return self.encoder_cache.get(mm_hash, None)
3211
+
3212
def _gather_mm_embeddings(
3213
self,
3214
scheduler_output: "SchedulerOutput",
@@ -3220,7 +3262,7 @@ class GPUModelRunner(
3262
continue
3263
3264
mm_hash = mm_feature.identifier
3223
- encoder_output = self.encoder_cache.get(mm_hash, None)
3265
+ encoder_output = self._get_encoder_output_from_cache(mm_hash)
3266
if encoder_output is None:
3267
# A feature starting at/after the processed boundary is only
3268
# reached via the drafter's +1 look-ahead and might not be