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

[Metrics] Log multi-modal cache stats and fix reset (#26285)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>

Cyrus Leung committed Oct 10, 2025 at 16:45 UTC ad430a67cab89ddc6060cf493f730c291826eb9d
25 files changed +588 -237
tests/entrypoints/llm/test_mm_cache_stats.py new
+74
@@ -0,0 +1,74 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm import LLM
7 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
8 +from vllm.v1.metrics.reader import Counter, Metric
9 +
10 +from ..openai.test_vision import TEST_IMAGE_ASSETS
11 +
12 +
13 +def _make_messages(image_url: str) -> list[ChatCompletionMessageParam]:
14 + return [
15 + {
16 + "role": "user",
17 + "content": [
18 + {
19 + "type": "image_url",
20 + "image_url": {"url": image_url},
21 + },
22 + ],
23 + }
24 + ]
25 +
26 +
27 +def _get_counter_value(metrics: list[Metric], name: str):
28 + metric = next(m for m in metrics if m.name == name)
29 + assert isinstance(metric, Counter)
30 + return metric.value
31 +
32 +
33 +def _get_mm_cache_stats(metrics: list[Metric]):
34 + mm_cache_queries = _get_counter_value(metrics, "vllm:mm_cache_queries")
35 + mm_cache_hits = _get_counter_value(metrics, "vllm:mm_cache_hits")
36 +
37 + return mm_cache_queries, mm_cache_hits
38 +
39 +
40 +@pytest.mark.parametrize("image_urls", [TEST_IMAGE_ASSETS[:2]], indirect=True)
41 +@pytest.mark.parametrize("mm_processor_cache_type", ["lru", "shm"])
42 +def test_mm_cache_stats(
43 + num_gpus_available,
44 + image_urls,
45 + mm_processor_cache_type,
46 +):
47 + llm = LLM(
48 + model="llava-hf/llava-1.5-7b-hf",
49 + max_model_len=4096,
50 + max_num_seqs=5,
51 + enforce_eager=True,
52 + mm_processor_cache_type=mm_processor_cache_type,
53 + disable_log_stats=False,
54 + limit_mm_per_prompt={"image": 2},
55 + )
56 +
57 + llm.chat(_make_messages(image_urls[0]))
58 + assert _get_mm_cache_stats(llm.get_metrics()) == (1, 0)
59 +
60 + llm.chat(_make_messages(image_urls[1]))
61 + assert _get_mm_cache_stats(llm.get_metrics()) == (2, 0)
62 +
63 + llm.chat(_make_messages(image_urls[0]))
64 + assert _get_mm_cache_stats(llm.get_metrics()) == (3, 1)
65 +
66 + # NOTE: This only resets hit rate stats in CachingMetrics
67 + # The raw queries and hits counts remain unaffected
68 + llm.reset_mm_cache()
69 +
70 + llm.chat(_make_messages(image_urls[0]))
71 + assert _get_mm_cache_stats(llm.get_metrics()) == (4, 1)
72 +
73 + llm.chat(_make_messages(image_urls[1]))
74 + assert _get_mm_cache_stats(llm.get_metrics()) == (5, 1)
tests/entrypoints/openai/test_metrics.py
+115 -65
@@ -18,10 +18,18 @@ from vllm import version
18
19 from ...utils import RemoteOpenAIServer
20
21 -MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
21 +MODELS = {
22 + "text": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
23 + "multimodal": "HuggingFaceTB/SmolVLM-256M-Instruct",
24 +}
25 PREV_MINOR_VERSION = version._prev_minor_version()
26
27
28 +@pytest.fixture(scope="module", params=list(MODELS.keys()))
29 +def model_key(request):
30 + yield request.param
31 +
32 +
33 @pytest.fixture(scope="module")
34 def default_server_args():
35 return [
@@ -45,11 +53,12 @@ def default_server_args():
53 f"--show-hidden-metrics-for-version={PREV_MINOR_VERSION}",
54 ],
55 )
48 -def server(default_server_args, request):
56 +def server(model_key, default_server_args, request):
57 if request.param:
58 default_server_args.append(request.param)
59
52 - with RemoteOpenAIServer(MODEL_NAME, default_server_args) as remote_server:
60 + model_name = MODELS[model_key]
61 + with RemoteOpenAIServer(model_name, default_server_args) as remote_server:
62 yield remote_server
63
64
@@ -60,64 +69,70 @@ async def client(server):
69
70
71 _PROMPT = "Hello my name is Robert and I love magic"
63 -tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
64 -_TOKENIZED_PROMPT = tokenizer(_PROMPT)["input_ids"]
65 -
66 -_NUM_REQUESTS = 10
67 -_NUM_PROMPT_TOKENS_PER_REQUEST = len(_TOKENIZED_PROMPT)
68 -_NUM_GENERATION_TOKENS_PER_REQUEST = 10
69 -
70 -# {metric_family: [(suffix, expected_value)]}
71 -EXPECTED_VALUES = {
72 - "vllm:time_to_first_token_seconds": [("_count", _NUM_REQUESTS)],
73 - "vllm:time_per_output_token_seconds": [
74 - ("_count", _NUM_REQUESTS * (_NUM_GENERATION_TOKENS_PER_REQUEST - 1))
75 - ],
76 - "vllm:e2e_request_latency_seconds": [("_count", _NUM_REQUESTS)],
77 - "vllm:request_queue_time_seconds": [("_count", _NUM_REQUESTS)],
78 - "vllm:request_inference_time_seconds": [("_count", _NUM_REQUESTS)],
79 - "vllm:request_prefill_time_seconds": [("_count", _NUM_REQUESTS)],
80 - "vllm:request_decode_time_seconds": [("_count", _NUM_REQUESTS)],
81 - "vllm:request_prompt_tokens": [
82 - ("_sum", _NUM_REQUESTS * _NUM_PROMPT_TOKENS_PER_REQUEST),
83 - ("_count", _NUM_REQUESTS),
84 - ],
85 - "vllm:request_generation_tokens": [
86 - ("_sum", _NUM_REQUESTS * _NUM_GENERATION_TOKENS_PER_REQUEST),
87 - ("_count", _NUM_REQUESTS),
88 - ],
89 - "vllm:request_params_n": [("_count", _NUM_REQUESTS)],
90 - "vllm:request_params_max_tokens": [
91 - ("_sum", _NUM_REQUESTS * _NUM_GENERATION_TOKENS_PER_REQUEST),
92 - ("_count", _NUM_REQUESTS),
93 - ],
94 - "vllm:iteration_tokens_total": [
95 - (
96 - "_sum",
97 - _NUM_REQUESTS
98 - * (_NUM_PROMPT_TOKENS_PER_REQUEST + _NUM_GENERATION_TOKENS_PER_REQUEST),
99 - ),
100 - ("_count", _NUM_REQUESTS * _NUM_GENERATION_TOKENS_PER_REQUEST),
101 - ],
102 - "vllm:prompt_tokens": [("_total", _NUM_REQUESTS * _NUM_PROMPT_TOKENS_PER_REQUEST)],
103 - "vllm:generation_tokens": [
104 - ("_total", _NUM_REQUESTS * _NUM_PROMPT_TOKENS_PER_REQUEST)
105 - ],
106 - "vllm:request_success": [("_total", _NUM_REQUESTS)],
107 -}
72 +_IMAGE_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
73 +
74 +
75 +def _get_expected_values(num_requests: int, prompt_ids: list[int], max_tokens: int):
76 + num_prompt_tokens = len(prompt_ids)
77 +
78 + # {metric_family: [(suffix, expected_value)]}
79 + return {
80 + "vllm:time_to_first_token_seconds": [("_count", num_requests)],
81 + "vllm:time_per_output_token_seconds": [
82 + ("_count", num_requests * (max_tokens - 1))
83 + ],
84 + "vllm:e2e_request_latency_seconds": [("_count", num_requests)],
85 + "vllm:request_queue_time_seconds": [("_count", num_requests)],
86 + "vllm:request_inference_time_seconds": [("_count", num_requests)],
87 + "vllm:request_prefill_time_seconds": [("_count", num_requests)],
88 + "vllm:request_decode_time_seconds": [("_count", num_requests)],
89 + "vllm:request_prompt_tokens": [
90 + ("_sum", num_requests * num_prompt_tokens),
91 + ("_count", num_requests),
92 + ],
93 + "vllm:request_generation_tokens": [
94 + ("_sum", num_requests * max_tokens),
95 + ("_count", num_requests),
96 + ],
97 + "vllm:request_params_n": [("_count", num_requests)],
98 + "vllm:request_params_max_tokens": [
99 + ("_sum", num_requests * max_tokens),
100 + ("_count", num_requests),
101 + ],
102 + "vllm:iteration_tokens_total": [
103 + (
104 + "_sum",
105 + num_requests * (num_prompt_tokens + max_tokens),
106 + ),
107 + ("_count", num_requests * max_tokens),
108 + ],
109 + "vllm:prompt_tokens": [("_total", num_requests * num_prompt_tokens)],
110 + "vllm:generation_tokens": [("_total", num_requests * max_tokens)],
111 + "vllm:request_success": [("_total", num_requests)],
112 + }
113
114
115 @pytest.mark.asyncio
116 async def test_metrics_counts(
117 server: RemoteOpenAIServer,
118 client: openai.AsyncClient,
119 + model_key: str,
120 ):
115 - for _ in range(_NUM_REQUESTS):
121 + if model_key == "multimodal":
122 + pytest.skip("Unnecessary test")
123 +
124 + model_name = MODELS[model_key]
125 + tokenizer = AutoTokenizer.from_pretrained(model_name)
126 + prompt_ids = tokenizer.encode(_PROMPT)
127 + num_requests = 10
128 + max_tokens = 10
129 +
130 + for _ in range(num_requests):
131 # sending a request triggers the metrics to be logged.
132 await client.completions.create(
118 - model=MODEL_NAME,
119 - prompt=_TOKENIZED_PROMPT,
120 - max_tokens=_NUM_GENERATION_TOKENS_PER_REQUEST,
133 + model=model_name,
134 + prompt=prompt_ids,
135 + max_tokens=max_tokens,
136 )
137
138 response = requests.get(server.url_for("metrics"))
@@ -125,8 +140,9 @@ async def test_metrics_counts(
140 assert response.status_code == HTTPStatus.OK
141
142 # Loop over all expected metric_families
128 - for metric_family, suffix_values_list in EXPECTED_VALUES.items():
129 - if (metric_family not in EXPECTED_METRICS_V1) or (
143 + expected_values = _get_expected_values(num_requests, prompt_ids, max_tokens)
144 + for metric_family, suffix_values_list in expected_values.items():
145 + if metric_family not in EXPECTED_METRICS_V1 or (
146 not server.show_hidden_metrics
147 and metric_family in HIDDEN_DEPRECATED_METRICS
148 ):
@@ -217,6 +233,11 @@ EXPECTED_METRICS_V1 = [
233 "vllm:request_decode_time_seconds_count",
234 ]
235
236 +EXPECTED_METRICS_MM = [
237 + "vllm:mm_cache_queries",
238 + "vllm:mm_cache_hits",
239 +]
240 +
241 HIDDEN_DEPRECATED_METRICS: list[str] = [
242 "vllm:gpu_cache_usage_perc",
243 "vllm:gpu_prefix_cache_queries",
@@ -231,19 +252,43 @@ HIDDEN_DEPRECATED_METRICS: list[str] = [
252 async def test_metrics_exist(
253 server: RemoteOpenAIServer,
254 client: openai.AsyncClient,
255 + model_key: str,
256 ):
257 + model_name = MODELS[model_key]
258 +
259 # sending a request triggers the metrics to be logged.
236 - await client.completions.create(
237 - model=MODEL_NAME,
238 - prompt="Hello, my name is",
239 - max_tokens=5,
240 - temperature=0.0,
241 - )
260 + if model_key == "text":
261 + await client.completions.create(
262 + model=model_name,
263 + prompt="Hello, my name is",
264 + max_tokens=5,
265 + temperature=0.0,
266 + )
267 + else:
268 + await client.chat.completions.create(
269 + model=model_name,
270 + messages=[
271 + {
272 + "role": "user",
273 + "content": [
274 + {"type": "image_url", "image_url": {"url": _IMAGE_URL}},
275 + {"type": "text", "text": "What's in this image?"},
276 + ],
277 + }
278 + ],
279 + max_tokens=5,
280 + temperature=0.0,
281 + )
282
283 response = requests.get(server.url_for("metrics"))
284 assert response.status_code == HTTPStatus.OK
285
246 - for metric in EXPECTED_METRICS_V1:
286 + expected_metrics = EXPECTED_METRICS_V1
287 + if model_key == "multimodal":
288 + # NOTE: Don't use in-place assignment
289 + expected_metrics = expected_metrics + EXPECTED_METRICS_MM
290 +
291 + for metric in expected_metrics:
292 if metric in HIDDEN_DEPRECATED_METRICS and not server.show_hidden_metrics:
293 continue
294 assert metric in response.text
@@ -253,9 +298,14 @@ async def test_metrics_exist(
298 async def test_abort_metrics_reset(
299 server: RemoteOpenAIServer,
300 client: openai.AsyncClient,
301 + model_key: str,
302 ):
303 + model_name = MODELS[model_key]
304 + tokenizer = AutoTokenizer.from_pretrained(model_name)
305 + prompt_ids = tokenizer.encode(_PROMPT)
306 +
307 running_requests, waiting_requests, kv_cache_usage = _get_running_metrics_from_api(
258 - server
308 + server,
309 )
310
311 # Expect no running requests or kvcache usage
@@ -268,8 +318,8 @@ async def test_abort_metrics_reset(
318 for _ in range(3):
319 task = asyncio.create_task(
320 client.completions.create(
271 - model=MODEL_NAME,
272 - prompt=_TOKENIZED_PROMPT,
321 + model=model_name,
322 + prompt=prompt_ids,
323 max_tokens=100, # Long generation to give time to abort
324 temperature=0.0,
325 )
@@ -281,7 +331,7 @@ async def test_abort_metrics_reset(
331
332 # Check that we have running requests
333 running_requests, waiting_requests, kv_cache_usage = _get_running_metrics_from_api(
284 - server
334 + server,
335 )
336
337 # Expect running requests and kvcache usage
tests/v1/core/test_kv_cache_utils.py
+3 -4
@@ -20,7 +20,6 @@ from vllm.v1.core.kv_cache_utils import (
20 BlockHash,
21 FreeKVCacheBlockQueue,
22 KVCacheBlock,
23 - PrefixCachingMetrics,
23 estimate_max_model_len,
24 generate_block_hash_extra_keys,
25 generate_scheduler_kv_cache_config,
@@ -42,7 +41,7 @@ from vllm.v1.kv_cache_interface import (
41 SlidingWindowSpec,
42 UniformTypeKVCacheSpecs,
43 )
45 -from vllm.v1.metrics.stats import PrefixCacheStats
44 +from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
45 from vllm.v1.request import Request
46
47 pytestmark = pytest.mark.cpu_test
@@ -536,7 +535,7 @@ def test_metrics():
535 """
536 Test the prefix caching metrics.
537 """
539 - metrics = PrefixCachingMetrics(max_recent_requests=5)
538 + metrics = CachingMetrics(max_recent_requests=5)
539 assert metrics.hit_rate == 0.0
540
541 metrics.observe(_stats(1, 20, 9))
@@ -568,7 +567,7 @@ def test_metrics_empty_stats():
567 """
568 Test the prefix caching metrics with empty stats.
569 """
571 - metrics = PrefixCachingMetrics(max_recent_requests=5)
570 + metrics = CachingMetrics(max_recent_requests=5)
571 metrics.observe(_stats(0, 0, 0))
572 metrics.observe(_stats(1, 20, 9))
573 metrics.observe(_stats(0, 0, 0))
tests/v1/distributed/test_async_llm_dp.py
+2 -1
@@ -17,7 +17,7 @@ from vllm.sampling_params import RequestOutputKind
17 from vllm.v1.engine.async_llm import AsyncLLM
18 from vllm.v1.engine.core_client import DPAsyncMPClient
19 from vllm.v1.metrics.loggers import StatLoggerBase
20 -from vllm.v1.metrics.stats import IterationStats, SchedulerStats
20 +from vllm.v1.metrics.stats import IterationStats, MultiModalCacheStats, SchedulerStats
21
22 DP_SIZE = int(os.getenv("DP_SIZE", 2))
23
@@ -93,6 +93,7 @@ async def test_load(
93 self,
94 scheduler_stats: Optional[SchedulerStats],
95 iteration_stats: Optional[IterationStats],
96 + mm_cache_stats: Optional[MultiModalCacheStats] = None,
97 engine_idx: int = 0,
98 ):
99 if iteration_stats:
vllm/entrypoints/llm.py
+4
@@ -354,6 +354,10 @@ class LLM:
354 else:
355 self.llm_engine.tokenizer = get_cached_tokenizer(tokenizer)
356
357 + def reset_mm_cache(self) -> None:
358 + self.processor.clear_mm_cache()
359 + self.llm_engine.reset_mm_cache()
360 +
361 def get_default_sampling_params(self) -> SamplingParams:
362 if self.default_sampling_params is None:
363 self.default_sampling_params = self.model_config.get_diff_sampling_param()
vllm/entrypoints/openai/serving_engine.py
+4
@@ -274,6 +274,10 @@ class OpenAIServing:
274 self.model_config = self.models.model_config
275 self.max_model_len = self.model_config.max_model_len
276
277 + async def reset_mm_cache(self) -> None:
278 + self.processor.clear_mm_cache()
279 + await self.engine_client.reset_mm_cache()
280 +
281 async def beam_search(
282 self,
283 prompt: PromptType,
vllm/executor/executor_base.py
+4
@@ -169,6 +169,10 @@ class ExecutorBase(ABC):
169 assert s == sets[0], "All workers should have the same LORAs."
170 return sets[0]
171
172 + def reset_mm_cache(self) -> None:
173 + """Reset the multi-modal cache in each worker."""
174 + self.collective_rpc("reset_mm_cache")
175 +
176 def start_profile(self) -> None:
177 self.collective_rpc("start_profile")
178
vllm/executor/uniproc_executor.py
+2 -10
@@ -12,11 +12,8 @@ import torch.distributed as dist
12 import vllm.envs as envs
13 from vllm.executor.executor_base import ExecutorBase
14 from vllm.logger import init_logger
15 -from vllm.multimodal import MULTIMODAL_REGISTRY
16 -from vllm.multimodal.cache import worker_receiver_cache_from_config
15 from vllm.utils import get_distributed_init_method, get_ip, get_open_port, run_method
16 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
19 -from vllm.v1.executor.utils import get_and_update_mm_cache
17 from vllm.v1.outputs import AsyncModelRunnerOutput
18 from vllm.v1.worker.worker_base import WorkerWrapperBase
19
@@ -30,16 +27,13 @@ class UniProcExecutor(ExecutorBase):
27 """Initialize the worker and load the model."""
28 self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config, rpc_rank=0)
29 distributed_init_method, rank, local_rank = self._distributed_args()
33 - is_driver_worker = True
30 kwargs = dict(
31 vllm_config=self.vllm_config,
32 local_rank=local_rank,
33 rank=rank,
34 distributed_init_method=distributed_init_method,
39 - is_driver_worker=is_driver_worker,
40 - )
41 - self.mm_receiver_cache = worker_receiver_cache_from_config(
42 - self.vllm_config, MULTIMODAL_REGISTRY, Lock()
35 + is_driver_worker=True,
36 + shared_worker_lock=Lock(),
37 )
38
39 self.async_output_thread: Optional[ThreadPoolExecutor] = None
@@ -74,8 +68,6 @@ class UniProcExecutor(ExecutorBase):
68 ) -> list[Any]:
69 if kwargs is None:
70 kwargs = {}
77 - if self.mm_receiver_cache is not None and method == "execute_model":
78 - get_and_update_mm_cache(self.mm_receiver_cache, args)
71
72 if not non_block:
73 return [run_method(self.driver_worker, method, args, kwargs)]
vllm/inputs/preprocess.py
+39 -3
@@ -19,6 +19,7 @@ from vllm.multimodal.inputs import (
19 from vllm.multimodal.processing import BaseMultiModalProcessor
20 from vllm.transformers_utils.tokenizer import AnyTokenizer
21 from vllm.utils.jsontree import json_iter_leaves
22 +from vllm.v1.metrics.stats import MultiModalCacheStats
23
24 from .data import (
25 DecoderOnlyInputs,
@@ -56,6 +57,8 @@ class InputPreprocessor:
57 self.mm_registry = mm_registry
58 self.mm_processor_cache = mm_processor_cache
59
60 + self.mm_cache_stats = MultiModalCacheStats() if mm_processor_cache else None
61 +
62 def get_tokenizer(self) -> AnyTokenizer:
63 if self.tokenizer is None:
64 raise ValueError(
@@ -664,14 +667,13 @@ class InputPreprocessor:
667
668 return self._build_decoder_only_llm_inputs(prompt_comps)
669
667 - def preprocess(
670 + def _preprocess(
671 self,
672 prompt: PromptType,
673 tokenization_kwargs: Optional[dict[str, Any]] = None,
674 *,
675 mm_uuids: Optional[MultiModalUUIDDict] = None,
676 ) -> ProcessorInputs:
674 - """Preprocess the input prompt."""
677 if self.model_config.is_encoder_decoder:
678 # Encoder-decoder model requires special mapping of
679 # input prompts to encoder & decoder.
@@ -694,6 +696,40 @@ class InputPreprocessor:
696 mm_uuids=mm_uuids,
697 )
698
697 - def clear_cache(self) -> None:
699 + def preprocess(
700 + self,
701 + prompt: PromptType,
702 + tokenization_kwargs: Optional[dict[str, Any]] = None,
703 + *,
704 + mm_uuids: Optional[MultiModalUUIDDict] = None,
705 + ) -> ProcessorInputs:
706 + """Preprocess the input prompt."""
707 + res = self._preprocess(
708 + prompt,
709 + tokenization_kwargs,
710 + mm_uuids=mm_uuids,
711 + )
712 +
713 + if self.mm_processor_cache and self.mm_cache_stats is not None:
714 + delta = self.mm_processor_cache.make_stats(delta=True)
715 + self.mm_cache_stats.requests += 1
716 + self.mm_cache_stats.queries += delta.total
717 + self.mm_cache_stats.hits += delta.hits
718 +
719 + return res
720 +
721 + def stat_mm_cache(self) -> Optional[MultiModalCacheStats]:
722 + mm_cache_stats = self.mm_cache_stats
723 + if mm_cache_stats is None:
724 + return None
725 +
726 + self.mm_cache_stats = MultiModalCacheStats()
727 +
728 + return mm_cache_stats
729 +
730 + def clear_mm_cache(self) -> None:
731 if self.mm_processor_cache is not None:
732 self.mm_processor_cache.clear_cache()
733 +
734 + if self.mm_cache_stats is not None:
735 + self.mm_cache_stats.reset = True
vllm/multimodal/cache.py
+46 -1
@@ -18,7 +18,7 @@ from vllm.distributed.device_communicators.shm_object_storage import (
18 from vllm.envs import VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME
19 from vllm.logger import init_logger
20 from vllm.utils import GiB_bytes, MiB_bytes
21 -from vllm.utils.cache import LRUCache
21 +from vllm.utils.cache import CacheInfo, LRUCache
22 from vllm.utils.jsontree import json_count_leaves, json_map_leaves, json_reduce_leaves
23
24 from .inputs import (
@@ -302,6 +302,16 @@ class BaseMultiModalProcessorCache(
302 """
303 return [self.is_cached_item(mm_hash) for mm_hash in mm_hashes]
304
305 + @abstractmethod
306 + def make_stats(self, *, delta: bool = False) -> CacheInfo:
307 + """
308 + Get (and reset) the multi-modal cache stats.
309 +
310 + Returns:
311 + The current multi-modal caching stats.
312 + """
313 + raise NotImplementedError
314 +
315
316 class MultiModalProcessorOnlyCache(BaseMultiModalProcessorCache):
317 """
@@ -347,6 +357,10 @@ class MultiModalProcessorOnlyCache(BaseMultiModalProcessorCache):
357 def clear_cache(self) -> None:
358 self._cache.clear()
359
360 + @override
361 + def make_stats(self, *, delta: bool = False) -> CacheInfo:
362 + return self._cache.stat(delta=delta)
363 +
364
365 class MultiModalProcessorSenderCache(BaseMultiModalProcessorCache):
366 """
@@ -397,6 +411,10 @@ class MultiModalProcessorSenderCache(BaseMultiModalProcessorCache):
411 def clear_cache(self) -> None:
412 self._cache.clear()
413
414 + @override
415 + def make_stats(self, *, delta: bool = False) -> CacheInfo:
416 + return self._cache.stat(delta=delta)
417 +
418
419 class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
420 """
@@ -430,6 +448,20 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
448 # cache (prompt_updates, modality) for P0 only
449 self._p0_cache: dict[str, tuple[Sequence[ResolvedPromptUpdate], str]] = {}
450
451 + self._hits = 0
452 + self._total = 0
453 + self._last_info = CacheInfo(hits=0, total=0)
454 +
455 + def _stat(self, *, delta: bool = False) -> CacheInfo:
456 + info = CacheInfo(hits=self._hits, total=self._total)
457 +
458 + if delta:
459 + info_delta = info - self._last_info
460 + self._last_info = info
461 + info = info_delta
462 +
463 + return info
464 +
465 @override
466 def is_cached_item(self, mm_hash: str) -> bool:
467 return self._shm_cache.is_cached(mm_hash)
@@ -441,12 +473,17 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
473 mm_hash: str,
474 ) -> MultiModalProcessorCacheOutItem:
475 if self._shm_cache.is_cached(mm_hash):
476 + self._hits += 1
477 + self._total += 1
478 +
479 address, monotonic_id = self._shm_cache.get_cached(mm_hash)
480 prompt_updates, modality = self._p0_cache[mm_hash]
481 return self.address_as_item(address, monotonic_id, modality), prompt_updates
482
483 assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
484
485 + self._total += 1
486 +
487 try:
488 address, monotonic_id = self._shm_cache.put(mm_hash, mm_item[0])
489 # Try to remove dangling items if p0 cache is too large.
@@ -469,6 +506,14 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
506 self._shm_cache.clear()
507 self._p0_cache.clear()
508
509 + self._hits = 0
510 + self._total = 0
511 + self._last_info = CacheInfo(hits=0, total=0)
512 +
513 + @override
514 + def make_stats(self, *, delta: bool = False) -> CacheInfo:
515 + return self._stat(delta=delta)
516 +
517 def remove_dangling_items(self) -> None:
518 """Remove items that are no longer in the shared memory cache."""
519 cached_hashes = self._shm_cache.key_index.keys()
vllm/v1/core/kv_cache_utils.py
+1 -74
@@ -4,7 +4,7 @@
4
5 import copy
6 import os
7 -from collections import defaultdict, deque
7 +from collections import defaultdict
8 from collections.abc import Iterable, Sequence
9 from dataclasses import dataclass
10 from typing import Any, Callable, NewType, Optional, Union
@@ -23,7 +23,6 @@ from vllm.v1.kv_cache_interface import (
23 SlidingWindowSpec,
24 UniformTypeKVCacheSpecs,
25 )
26 -from vllm.v1.metrics.stats import PrefixCacheStats
26 from vllm.v1.request import Request
27
28 # BlockHash represents the hash of a single KV-cache block used for
@@ -101,78 +100,6 @@ def init_none_hash(hash_fn: Callable[[Any], bytes]):
100 NONE_HASH = BlockHash(hash_fn(hash_seed))
101
102
104 -class PrefixCachingMetrics:
105 - """Metrics for prefix caching with a hit rate of the max recent N requests.
106 -
107 - Args:
108 - max_recent_requests: The number of the max recent requests to aggregate.
109 - Defaults to 1000.
110 - """
111 -
112 - def __init__(self, max_recent_requests: int = 1000):
113 - self.max_recent_requests = max_recent_requests
114 - # The current aggregated values.
115 - self.aggregated_requests = 0
116 - self.aggregated_query_total = 0
117 - self.aggregated_query_hit = 0
118 - # A deque of (requests, queries, hits) for the most recent requests.
119 - self.query_queue: deque[tuple[int, int, int]] = deque()
120 -
121 - def observe(self, stats: PrefixCacheStats):
122 - """Observe the prefix caching for a set of requests.
123 -
124 - This function is called with information gathered when new requests
125 - are being scheduled and are looking for computed blocks.
126 -
127 - When there are more than `max_recent_requests` requests, the oldest set
128 - of requests are removed from the metrics.
129 -
130 - Args:
131 - stats: The prefix cache stats.
132 - """
133 - # reset_prefix_cache was invoked before the current update.
134 - # Reset the metrics before aggregating the current stats.
135 - if stats.reset:
136 - self.reset()
137 -
138 - # DO NOT appending empty stats to avoid helpful info get kicked out
139 - # due to sliding window.
140 - if stats.requests == 0:
141 - return
142 -
143 - # Update the metrics.
144 - self.query_queue.append((stats.requests, stats.queries, stats.hits))
145 - self.aggregated_requests += stats.requests
146 - self.aggregated_query_total += stats.queries
147 - self.aggregated_query_hit += stats.hits
148 -
149 - # Remove the oldest stats until number of requests does not exceed
150 - # the limit.
151 - # NOTE: We preserve the latest added stats regardless.
152 - while (
153 - len(self.query_queue) > 1
154 - and self.aggregated_requests > self.max_recent_requests
155 - ):
156 - old_requests, old_queries, old_hits = self.query_queue.popleft()
157 - self.aggregated_requests -= old_requests
158 - self.aggregated_query_total -= old_queries
159 - self.aggregated_query_hit -= old_hits
160 -
161 - def reset(self):
162 - """Reset the metrics."""
163 - self.aggregated_requests = 0
164 - self.aggregated_query_total = 0
165 - self.aggregated_query_hit = 0
166 - self.query_queue.clear()
167 -
168 - @property
169 - def hit_rate(self) -> float:
170 - """Calculate the hit rate for the past N requests."""
171 - if self.aggregated_query_total == 0:
172 - return 0.0
173 - return self.aggregated_query_hit / self.aggregated_query_total
174 -
175 -
103 @dataclass
104 class KVCacheBlock:
105 """KV-cache block metadata."""
vllm/v1/engine/async_llm.py
+3 -1
@@ -463,6 +463,7 @@ class AsyncLLM(EngineClient):
463 output_processor = self.output_processor
464 log_stats = self.log_stats
465 logger_manager = self.logger_manager
466 + processor = self.processor
467
468 async def output_handler():
469 try:
@@ -511,6 +512,7 @@ class AsyncLLM(EngineClient):
512 engine_idx=outputs.engine_index,
513 scheduler_stats=outputs.scheduler_stats,
514 iteration_stats=iteration_stats,
515 + mm_cache_stats=processor.stat_mm_cache(),
516 )
517 except Exception as e:
518 logger.exception("AsyncLLM output_handler failed.")
@@ -660,7 +662,7 @@ class AsyncLLM(EngineClient):
662 await asyncio.gather(*coros)
663
664 async def reset_mm_cache(self) -> None:
663 - self.processor.clear_cache()
665 + self.processor.clear_mm_cache()
666 await self.engine_core.reset_mm_cache_async()
667
668 async def reset_prefix_cache(self, device: Optional[Device] = None) -> None:
vllm/v1/engine/core.py
+5 -2
@@ -319,7 +319,7 @@ class EngineCore:
319 )
320 engine_core_outputs = self.scheduler.update_from_output(
321 scheduler_output, model_output
322 - ) # type: ignore
322 + )
323
324 return (engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0)
325
@@ -400,16 +400,19 @@ class EngineCore:
400
401 def reset_mm_cache(self):
402 # NOTE: Since this is mainly for debugging, we don't attempt to
403 - # re-sync the internal caches (P0 processor, P0 mirror, P1 mirror)
403 + # re-sync the internal caches (P0 sender, P1 receiver)
404 if self.scheduler.has_unfinished_requests():
405 logger.warning(
406 "Resetting the multi-modal cache when requests are "
407 "in progress may lead to desynced internal caches."
408 )
409
410 + # The cache either exists in EngineCore or WorkerWrapperBase
411 if self.mm_receiver_cache is not None:
412 self.mm_receiver_cache.clear_cache()
413
414 + self.model_executor.reset_mm_cache()
415 +
416 def reset_prefix_cache(self):
417 self.scheduler.reset_prefix_cache()
418
vllm/v1/engine/llm_engine.py
+3 -1
@@ -306,9 +306,11 @@ class LLMEngine:
306 # 4) Record stats
307 if self.logger_manager is not None:
308 assert outputs.scheduler_stats is not None
309 +
310 self.logger_manager.record(
311 scheduler_stats=outputs.scheduler_stats,
312 iteration_stats=iteration_stats,
313 + mm_cache_stats=self.processor.stat_mm_cache(),
314 )
315 self.do_log_stats_with_interval()
316
@@ -321,7 +323,7 @@ class LLMEngine:
323 self.engine_core.profile(False)
324
325 def reset_mm_cache(self):
324 - self.processor.clear_cache()
326 + self.processor.clear_mm_cache()
327 self.engine_core.reset_mm_cache()
328
329 def reset_prefix_cache(self, device: Optional[Device] = None):
vllm/v1/engine/processor.py
+6 -2
@@ -21,6 +21,7 @@ from vllm.sampling_params import SamplingParams
21 from vllm.transformers_utils.tokenizer import AnyTokenizer
22 from vllm.utils import length_from_prompt_token_ids_or_embeds
23 from vllm.v1.engine import EngineCoreRequest
24 +from vllm.v1.metrics.stats import MultiModalCacheStats
25 from vllm.v1.structured_output.backend_guidance import validate_guidance_grammar
26 from vllm.v1.structured_output.backend_lm_format_enforcer import (
27 validate_structured_output_request_lm_format_enforcer,
@@ -573,5 +574,8 @@ class Processor:
574 # check that chunked prefill does not truncate them
575 # max_batch_len = self.scheduler_config.max_num_batched_tokens
576
576 - def clear_cache(self) -> None:
577 - self.input_preprocessor.clear_cache()
577 + def stat_mm_cache(self) -> Optional[MultiModalCacheStats]:
578 + return self.input_preprocessor.stat_mm_cache()
579 +
580 + def clear_mm_cache(self) -> None:
581 + self.input_preprocessor.clear_mm_cache()
vllm/v1/executor/multiproc_executor.py
+2 -14
@@ -33,8 +33,6 @@ from vllm.distributed.parallel_state import (
33 get_tp_group,
34 )
35 from vllm.logger import init_logger
36 -from vllm.multimodal import MULTIMODAL_REGISTRY
37 -from vllm.multimodal.cache import worker_receiver_cache_from_config
36 from vllm.utils import (
37 _maybe_force_spawn,
38 decorate_logs,
@@ -46,7 +44,6 @@ from vllm.utils import (
44 )
45 from vllm.v1.core.sched.output import SchedulerOutput
46 from vllm.v1.executor.abstract import Executor, FailureCallback
49 -from vllm.v1.executor.utils import get_and_update_mm_cache
47 from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
48 from vllm.v1.worker.worker_base import WorkerWrapperBase
49
@@ -422,6 +419,7 @@ class WorkerProc:
419 "rank": rank,
420 "distributed_init_method": distributed_init_method,
421 "is_driver_worker": is_driver_worker,
422 + "shared_worker_lock": shared_worker_lock,
423 }
424 wrapper.init_worker(all_kwargs)
425 self.worker = wrapper
@@ -445,11 +443,6 @@ class WorkerProc:
443 )
444 self.async_output_copy_thread.start()
445
448 - # Initialize multimodal receiver cache if needed
449 - self.mm_receiver_cache = worker_receiver_cache_from_config(
450 - vllm_config, MULTIMODAL_REGISTRY, shared_worker_lock
451 - )
452 -
446 # Initialize device
447 self.worker.init_device()
448
@@ -692,12 +685,7 @@ class WorkerProc:
685 func = getattr(self.worker, method)
686 elif isinstance(method, bytes):
687 func = partial(cloudpickle.loads(method), self.worker)
695 - # retrieve from shm cache if available
696 - if (
697 - self.mm_receiver_cache is not None
698 - and func.__name__ == "execute_model"
699 - ):
700 - get_and_update_mm_cache(self.mm_receiver_cache, args)
688 +
689 output = func(*args, **kwargs)
690 except Exception as e:
691 # Notes have been introduced in python 3.11
vllm/v1/executor/utils.py deleted
-24
@@ -1,24 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from vllm.multimodal.cache import ShmObjectStoreReceiverCache
4 -from vllm.v1.core.sched.output import SchedulerOutput
5 -
6 -
7 -def get_and_update_mm_cache(
8 - receiver_cache: ShmObjectStoreReceiverCache,
9 - args: tuple[SchedulerOutput],
10 -) -> None:
11 - """
12 - For each MultiModalKwargsItem in SchedulerOutput, fetch from shared memory
13 - cache as needed.
14 -
15 - Args:
16 - receiver_cache: The receiver cache to update.
17 - args: According to the collective_rpc call of execute_model method in
18 - executor, args is a tuple of only one SchedulerOutput element.
19 - """
20 - scheduler_output = args[0]
21 - for request_data in scheduler_output.scheduled_new_reqs:
22 - request_data.mm_features = receiver_cache.get_and_update_features(
23 - request_data.mm_features
24 - )
vllm/v1/metrics/loggers.py
+82 -13
@@ -11,10 +11,14 @@ import prometheus_client
11 from vllm.config import SupportsMetricsInfo, VllmConfig
12 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorLogging
13 from vllm.logger import init_logger
14 -from vllm.v1.core.kv_cache_utils import PrefixCachingMetrics
14 from vllm.v1.engine import FinishReason
15 from vllm.v1.metrics.prometheus import unregister_vllm_metrics
17 -from vllm.v1.metrics.stats import IterationStats, SchedulerStats
16 +from vllm.v1.metrics.stats import (
17 + CachingMetrics,
18 + IterationStats,
19 + MultiModalCacheStats,
20 + SchedulerStats,
21 +)
22 from vllm.v1.spec_decode.metrics import SpecDecodingLogging, SpecDecodingProm
23
24 logger = init_logger(__name__)
@@ -38,6 +42,7 @@ class StatLoggerBase(ABC):
42 self,
43 scheduler_stats: Optional[SchedulerStats],
44 iteration_stats: Optional[IterationStats],
45 + mm_cache_stats: Optional[MultiModalCacheStats] = None,
46 engine_idx: int = 0,
47 ): ...
48
@@ -53,10 +58,15 @@ class LoggingStatLogger(StatLoggerBase):
58 self.engine_index = engine_index
59 self.vllm_config = vllm_config
60 self._reset(time.monotonic())
61 +
62 self.last_scheduler_stats = SchedulerStats()
57 - # Prefix cache metrics. This cannot be reset.
63 + self.last_mm_cache_stats: Optional[MultiModalCacheStats] = None
64 +
65 + # Caching metrics. This cannot be reset.
66 # TODO: Make the interval configurable.
59 - self.prefix_caching_metrics = PrefixCachingMetrics()
67 + self.prefix_caching_metrics = CachingMetrics()
68 + self.mm_caching_metrics = CachingMetrics()
69 +
70 self.spec_decoding_logging = SpecDecodingLogging()
71 kv_tranfer_config = self.vllm_config.kv_transfer_config
72 self.kv_connector_logging = KVConnectorLogging(kv_tranfer_config)
@@ -86,6 +96,7 @@ class LoggingStatLogger(StatLoggerBase):
96 self,
97 scheduler_stats: Optional[SchedulerStats],
98 iteration_stats: Optional[IterationStats],
99 + mm_cache_stats: Optional[MultiModalCacheStats] = None,
100 engine_idx: int = 0,
101 ):
102 """Log Stats to standard output."""
@@ -101,6 +112,11 @@ class LoggingStatLogger(StatLoggerBase):
112 self.kv_connector_logging.observe(kv_connector_stats)
113 self.last_scheduler_stats = scheduler_stats
114
115 + if mm_cache_stats:
116 + self.mm_caching_metrics.observe(mm_cache_stats)
117 +
118 + self.last_mm_cache_stats = mm_cache_stats
119 +
120 def log(self):
121 now = time.monotonic()
122 prompt_throughput = self._get_throughput(self.num_prompt_tokens, now)
@@ -125,21 +141,32 @@ class LoggingStatLogger(StatLoggerBase):
141 self.last_prompt_throughput = prompt_throughput
142
143 # Format and print output.
128 - log_fn(
129 - "Engine %03d: "
130 - "Avg prompt throughput: %.1f tokens/s, "
131 - "Avg generation throughput: %.1f tokens/s, "
132 - "Running: %d reqs, Waiting: %d reqs, "
133 - "GPU KV cache usage: %.1f%%, "
144 + log_parts = [
145 + "Avg prompt throughput: %.1f tokens/s",
146 + "Avg generation throughput: %.1f tokens/s",
147 + "Running: %d reqs",
148 + "Waiting: %d reqs",
149 + "GPU KV cache usage: %.1f%%",
150 "Prefix cache hit rate: %.1f%%",
135 - self.engine_index,
151 + ]
152 + log_args = [
153 prompt_throughput,
154 generation_throughput,
155 scheduler_stats.num_running_reqs,
156 scheduler_stats.num_waiting_reqs,
157 scheduler_stats.kv_cache_usage * 100,
158 self.prefix_caching_metrics.hit_rate * 100,
159 + ]
160 + if self.last_mm_cache_stats:
161 + log_parts.append("MM cache hit rate: %.1f%%")
162 + log_args.append(self.mm_caching_metrics.hit_rate * 100)
163 +
164 + log_fn(
165 + "Engine %03d: " + ", ".join(log_parts),
166 + self.engine_index,
167 + *log_args,
168 )
169 +
170 self.spec_decoding_logging.log(log_fn=log_fn)
171 self.kv_connector_logging.log(log_fn=log_fn)
172
@@ -288,6 +315,32 @@ class PrometheusStatLogger(StatLoggerBase):
315 counter_prefix_cache_hits, engine_indexes, model_name
316 )
317
318 + #
319 + # Multi-modal cache
320 + #
321 +
322 + counter_mm_cache_queries = self._counter_cls(
323 + name="vllm:mm_cache_queries",
324 + documentation=(
325 + "Multi-modal cache queries, in terms of number of queried items."
326 + ),
327 + labelnames=labelnames,
328 + )
329 + self.counter_mm_cache_queries = make_per_engine(
330 + counter_mm_cache_queries, engine_indexes, model_name
331 + )
332 +
333 + counter_mm_cache_hits = self._counter_cls(
334 + name="vllm:mm_cache_hits",
335 + documentation=(
336 + "Multi-modal cache hits, in terms of number of cached items."
337 + ),
338 + labelnames=labelnames,
339 + )
340 + self.counter_mm_cache_hits = make_per_engine(
341 + counter_mm_cache_hits, engine_indexes, model_name
342 + )
343 +
344 #
345 # Counters
346 #
@@ -657,6 +710,7 @@ class PrometheusStatLogger(StatLoggerBase):
710 self,
711 scheduler_stats: Optional[SchedulerStats],
712 iteration_stats: Optional[IterationStats],
713 + mm_cache_stats: Optional[MultiModalCacheStats] = None,
714 engine_idx: int = 0,
715 ):
716 """Log to prometheus."""
@@ -694,6 +748,10 @@ class PrometheusStatLogger(StatLoggerBase):
748 scheduler_stats.spec_decoding_stats, engine_idx
749 )
750
751 + if mm_cache_stats is not None:
752 + self.counter_mm_cache_queries[engine_idx].inc(mm_cache_stats.queries)
753 + self.counter_mm_cache_hits[engine_idx].inc(mm_cache_stats.hits)
754 +
755 if iteration_stats is None:
756 return
757
@@ -871,6 +929,7 @@ class StatLoggerManager:
929 self,
930 scheduler_stats: Optional[SchedulerStats],
931 iteration_stats: Optional[IterationStats],
932 + mm_cache_stats: Optional[MultiModalCacheStats] = None,
933 engine_idx: Optional[int] = None,
934 ):
935 if engine_idx is None:
@@ -878,9 +937,19 @@ class StatLoggerManager:
937
938 per_engine_loggers = self.per_engine_logger_dict[engine_idx]
939 for logger in per_engine_loggers:
881 - logger.record(scheduler_stats, iteration_stats, engine_idx)
940 + logger.record(
941 + scheduler_stats,
942 + iteration_stats,
943 + mm_cache_stats=mm_cache_stats,
944 + engine_idx=engine_idx,
945 + )
946
883 - self.prometheus_logger.record(scheduler_stats, iteration_stats, engine_idx)
947 + self.prometheus_logger.record(
948 + scheduler_stats,
949 + iteration_stats,
950 + mm_cache_stats=mm_cache_stats,
951 + engine_idx=engine_idx,
952 + )
953
954 def log(self):
955 for per_engine_loggers in self.per_engine_logger_dict.values():
vllm/v1/metrics/stats.py
+109 -10
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import time
5 +from collections import deque
6 from dataclasses import dataclass, field
7 from typing import TYPE_CHECKING, Any, Optional
8
@@ -13,24 +14,122 @@ if TYPE_CHECKING:
14
15
16 @dataclass
16 -class PrefixCacheStats:
17 - """Stores prefix cache hit statistics."""
17 +class BaseCacheStats:
18 + """Stores cache hit statistics."""
19
19 - # Whether reset_prefix_cache was invoked.
20 reset: bool = False
21 - # The number of new requests in this update.
21 + """Whether the cache was reset."""
22 +
23 requests: int = 0
23 - # The number of queries in these requests. Note that "queries" here
24 - # means the number of tokens that were queried from the cache.
24 + """The number of requests in this update."""
25 +
26 queries: int = 0
26 - # The number of hits in these requests.
27 + """The number of queries in these requests."""
28 +
29 hits: int = 0
28 - # The number of previously preempted requests in this update.
30 + """The number of hits in these requests."""
31 +
32 +
33 +class CachingMetrics:
34 + """Metrics for caching with a hit rate of the most recent N requests.
35 + Args:
36 + interval: The number of the most recent requests to aggregate.
37 + Defaults to 1000.
38 + """
39 +
40 + def __init__(self, max_recent_requests: int = 1000) -> None:
41 + super().__init__()
42 +
43 + self.max_recent_requests = max_recent_requests
44 + # The current aggregated values.
45 + self.aggregated_requests = 0
46 + self.aggregated_query_total = 0
47 + self.aggregated_query_hit = 0
48 +
49 + # A deque of (requests, queries, hits) for the most recent requests.
50 + self.query_queue = deque[tuple[int, int, int]]()
51 +
52 + def observe(self, stats: BaseCacheStats):
53 + """Observe the prefix caching for a set of requests.
54 +
55 + This function is called with information gathered when new requests
56 + are being scheduled and are looking for computed blocks.
57 +
58 + When there are more than `max_recent_requests` requests, the oldest set
59 + of requests are removed from the metrics.
60 +
61 + Args:
62 + stats: The prefix cache stats.
63 + """
64 + # reset_prefix_cache was invoked before the current update.
65 + # Reset the metrics before aggregating the current stats.
66 + if stats.reset:
67 + self.reset()
68 +
69 + # DO NOT appending empty stats to avoid helpful info get kicked out
70 + # due to sliding window.
71 + if stats.requests == 0:
72 + return
73 +
74 + # Update the metrics.
75 + self.query_queue.append((stats.requests, stats.queries, stats.hits))
76 + self.aggregated_requests += stats.requests
77 + self.aggregated_query_total += stats.queries
78 + self.aggregated_query_hit += stats.hits
79 +
80 + # Remove the oldest stats until number of requests does not exceed
81 + # the limit.
82 + # NOTE: We preserve the latest added stats regardless.
83 + while (
84 + len(self.query_queue) > 1
85 + and self.aggregated_requests > self.max_recent_requests
86 + ):
87 + old_requests, old_queries, old_hits = self.query_queue.popleft()
88 + self.aggregated_requests -= old_requests
89 + self.aggregated_query_total -= old_queries
90 + self.aggregated_query_hit -= old_hits
91 +
92 + def reset(self):
93 + """Reset the metrics."""
94 + self.aggregated_requests = 0
95 + self.aggregated_query_total = 0
96 + self.aggregated_query_hit = 0
97 + self.query_queue.clear()
98 +
99 + @property
100 + def hit_rate(self) -> float:
101 + """Calculate the hit rate for the past N requests."""
102 + if self.aggregated_query_total == 0:
103 + return 0.0
104 + return self.aggregated_query_hit / self.aggregated_query_total
105 +
106 +
107 +@dataclass
108 +class PrefixCacheStats(BaseCacheStats):
109 + """
110 + Stores prefix cache hit statistics.
111 + - `reset`: Whether `reset_prefix_cache` was invoked.
112 + - `queries`: Refers to the number of tokens that were queried.
113 + """
114 +
115 preempted_requests: int = 0
30 - # The `queries` number for preempted requests.
116 + """The number of previously preempted requests in this update."""
117 +
118 preempted_queries: int = 0
32 - # The `hits` number for preempted requests.
119 + """The `queries` number for preempted requests."""
120 +
121 preempted_hits: int = 0
122 + """The `hits` number for preempted requests."""
123 +
124 +
125 +@dataclass
126 +class MultiModalCacheStats(BaseCacheStats):
127 + """
128 + Stores multi-modal cache hit statistics.
129 + - `reset`: Whether `reset_mm_cache` was invoked.
130 + - `queries`: Refers to the number of multi-modal data items
131 + that were queried.
132 + """
133
134
135 @dataclass
vllm/v1/worker/gpu_model_runner.py
+4
@@ -508,6 +508,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
508 pin_memory=self.pin_memory,
509 )
510
511 + def reset_mm_cache(self) -> None:
512 + if self.mm_budget:
513 + self.mm_budget.reset_cache()
514 +
515 def _get_positions(self, num_tokens: Any):
516 if isinstance(num_tokens, int):
517 if self.uses_mrope:
vllm/v1/worker/gpu_worker.py
+3
@@ -442,6 +442,9 @@ class Worker(WorkerBase):
442 # the model initialization and profiling.
443 set_random_seed(self.model_config.seed)
444
445 + def reset_mm_cache(self) -> None:
446 + self.model_runner.reset_mm_cache()
447 +
448 def get_model(self) -> nn.Module:
449 return self.model_runner.get_model()
450
vllm/v1/worker/tpu_model_runner.py
+4
@@ -371,6 +371,10 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
371 else:
372 self.sample_from_logits_func = self.sample_from_logits
373
374 + def reset_mm_cache(self) -> None:
375 + if self.mm_budget:
376 + self.mm_budget.reset_cache()
377 +
378 def _update_num_xla_graphs(self, case_str):
379 check_comp = self.check_recompilation and not self.enforce_eager
380 if not check_comp:
vllm/v1/worker/tpu_worker.py
+3
@@ -293,6 +293,9 @@ class TPUWorker:
293 # the model initialization and profiling.
294 set_random_seed(self.model_config.seed)
295
296 + def reset_mm_cache(self) -> None:
297 + self.model_runner.reset_mm_cache()
298 +
299 def get_model(self) -> nn.Module:
300 return self.model_runner.get_model()
301
vllm/v1/worker/utils.py
+4
@@ -126,6 +126,10 @@ class MultiModalBudget:
126
127 return max_items_per_prompt, max_items_per_batch
128
129 + def reset_cache(self) -> None:
130 + if self.cache is not None:
131 + self.cache.clear_cache()
132 +
133
134 @dataclass
135 class AttentionGroup:
vllm/v1/worker/worker_base.py
+66 -12
@@ -4,7 +4,7 @@
4 from __future__ import annotations
5
6 import os
7 -from typing import Any, Callable, TypeVar, Union
7 +from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union
8
9 import torch
10 import torch.nn as nn
@@ -12,7 +12,8 @@ import torch.nn as nn
12 from vllm.config import VllmConfig, set_current_vllm_config
13 from vllm.logger import init_logger
14 from vllm.lora.request import LoRARequest
15 -from vllm.sequence import ExecuteModelRequest
15 +from vllm.multimodal import MULTIMODAL_REGISTRY
16 +from vllm.multimodal.cache import worker_receiver_cache_from_config
17 from vllm.utils import (
18 enable_trace_function_call_for_thread,
19 resolve_obj_by_qualname,
@@ -21,7 +22,10 @@ from vllm.utils import (
22 warn_for_unimplemented_methods,
23 )
24 from vllm.v1.kv_cache_interface import KVCacheSpec
24 -from vllm.v1.outputs import SamplerOutput
25 +
26 +if TYPE_CHECKING:
27 + from vllm.v1.core.sched.output import SchedulerOutput
28 + from vllm.v1.outputs import ModelRunnerOutput
29
30 logger = init_logger(__name__)
31
@@ -103,6 +107,11 @@ class WorkerBase:
107 """Initialize the KV cache with the given size in blocks."""
108 raise NotImplementedError
109
110 + def reset_mm_cache(self) -> None:
111 + reset_fn = getattr(self.model_runner, "reset_mm_cache", None)
112 + if callable(reset_fn):
113 + reset_fn()
114 +
115 def get_model(self) -> nn.Module:
116 raise NotImplementedError
117
@@ -114,9 +123,7 @@ class WorkerBase:
123 """Load model onto target device."""
124 raise NotImplementedError
125
117 - def execute_model(
118 - self, execute_model_req: ExecuteModelRequest | None = None
119 - ) -> list[SamplerOutput] | None:
126 + def execute_model(self, scheduler_output: SchedulerOutput) -> ModelRunnerOutput:
127 raise NotImplementedError
128
129 def start_worker_execution_loop(self) -> None:
@@ -125,11 +132,7 @@ class WorkerBase:
132 You can stop the loop by executing a driver worker with an empty output.
133 See `stop_remote_worker_execution_loop` for more details.
134 """
128 - with self.current_platform.inference_mode():
129 - while True:
130 - output = self.execute_model(execute_model_req=None)
131 - if output is None:
132 - return None
135 + raise NotImplementedError("Dead V0 code")
136
137 def determine_num_available_blocks(self) -> tuple[int, int]:
138 """Determine the number of available blocks for the GPU KV cache and
@@ -289,6 +292,28 @@ class WorkerWrapperBase:
292 worker_class,
293 extended_calls,
294 )
295 +
296 + shared_worker_lock = kwargs.pop("shared_worker_lock", None)
297 + if shared_worker_lock is None:
298 + msg = (
299 + "Missing `shared_worker_lock` argument from executor. "
300 + "This argument is needed for mm_processor_cache_type='shm'."
301 + )
302 +
303 + mm_config = self.vllm_config.model_config.multimodal_config
304 + if mm_config and mm_config.mm_processor_cache_type == "shm":
305 + raise ValueError(msg)
306 + else:
307 + logger.warning_once(msg)
308 +
309 + self.mm_receiver_cache = None
310 + else:
311 + self.mm_receiver_cache = worker_receiver_cache_from_config(
312 + self.vllm_config,
313 + MULTIMODAL_REGISTRY,
314 + shared_worker_lock,
315 + )
316 +
317 with set_current_vllm_config(self.vllm_config):
318 # To make vLLM config available during worker initialization
319 self.worker = worker_class(**kwargs)
@@ -323,5 +348,34 @@ class WorkerWrapperBase:
348 logger.exception(msg)
349 raise e
350
326 - def __getattr__(self, attr):
351 + def __getattr__(self, attr: str):
352 return getattr(self.worker, attr)
353 +
354 + def _apply_mm_cache(self, scheduler_output: SchedulerOutput) -> None:
355 + mm_cache = self.mm_receiver_cache
356 + if mm_cache is None:
357 + return
358 +
359 + for req_data in scheduler_output.scheduled_new_reqs:
360 + req_data.mm_features = mm_cache.get_and_update_features(
361 + req_data.mm_features
362 + )
363 +
364 + def execute_model(
365 + self,
366 + scheduler_output: SchedulerOutput,
367 + *args,
368 + **kwargs,
369 + ) -> ModelRunnerOutput:
370 + self._apply_mm_cache(scheduler_output)
371 +
372 + assert self.worker is not None
373 + return self.worker.execute_model(scheduler_output, *args, **kwargs)
374 +
375 + def reset_mm_cache(self) -> None:
376 + mm_receiver_cache = self.mm_receiver_cache
377 + if mm_receiver_cache is not None:
378 + mm_receiver_cache.clear_cache()
379 +
380 + assert self.worker is not None
381 + self.worker.reset_mm_cache()