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 13, 03:15 UTC
[Model] Allow users to control skip reading cache per request. (#28194)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Nov 16, 2025 at 16:04 UTC
a55b64635c272ff1f34d20593140faa1fcbe4580
5 files changed
+67
-8
tests/models/language/pooling/test_extract_hidden_states.py
+27
-2
@@ -11,7 +11,7 @@ from vllm import TokensPrompt
11
["Qwen/Qwen3-0.6B"],
12
)
13
@torch.inference_mode
14
-def test_embed_models(hf_runner, vllm_runner, model: str):
14
+def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
15
n_prompt_tokens = [55, 56, 57]
16
token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
17
@@ -21,7 +21,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
21
enforce_eager=True,
22
runner="pooling",
23
enable_chunked_prefill=False,
24
- enable_prefix_caching=False,
24
+ enable_prefix_caching=True,
25
) as vllm_model:
26
pooling_outputs = vllm_model.llm.encode(
27
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
@@ -30,4 +30,29 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
30
31
for n, output in zip(n_prompt_tokens, pooling_outputs):
32
assert len(output.prompt_token_ids) == n
33
+ assert len(output.outputs.data) == n
34
assert output.num_cached_tokens == 0
35
+
36
+ # test enable_prefix_caching plus all pooling
37
+ # we need to skip reading cache at this request by
38
+ # request.skip_reading_prefix_cache
39
+ pooling_outputs = vllm_model.llm.encode(
40
+ [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
41
+ pooling_task="token_embed",
42
+ )
43
+
44
+ for n, output in zip(n_prompt_tokens, pooling_outputs):
45
+ assert len(output.prompt_token_ids) == n
46
+ assert len(output.outputs.data) == n
47
+ assert output.num_cached_tokens == 0
48
+
49
+ # skip_reading_prefix_cache can still write to cache
50
+ # to accelerate following requests
51
+ pooling_outputs = vllm_model.llm.encode(
52
+ [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
53
+ pooling_task="embed",
54
+ )
55
+
56
+ for n, output in zip(n_prompt_tokens, pooling_outputs):
57
+ assert len(output.prompt_token_ids) == n
58
+ assert output.num_cached_tokens > 0
vllm/pooling_params.py
+12
@@ -57,6 +57,7 @@ class PoolingParams(
57
## Internal use only
58
task: PoolingTask | None = None
59
requires_token_ids: bool = False
60
+ skip_reading_prefix_cache: bool = None
61
extra_kwargs: dict[str, Any] | None = None
62
output_kind: RequestOutputKind = RequestOutputKind.FINAL_ONLY
63
@@ -93,6 +94,8 @@ class PoolingParams(
94
# plugin task uses io_processor.parse_request to verify inputs,
95
# skipping PoolingParams verify
96
if self.task == "plugin":
97
+ if self.skip_reading_prefix_cache is None:
98
+ self.skip_reading_prefix_cache = True
99
return
100
101
# NOTE: Task validation needs to done against the model instance,
@@ -122,6 +125,15 @@ class PoolingParams(
125
if getattr(self, k, None) is None:
126
setattr(self, k, getattr(pooler_config, k))
127
128
+ if self.skip_reading_prefix_cache is None:
129
+ # If prefix caching is enabled,
130
+ # the output of all pooling may less than n_prompt_tokens,
131
+ # we need to skip reading cache at this request.
132
+ if self.task in ["token_embed", "token_classify"]:
133
+ self.skip_reading_prefix_cache = True
134
+ else:
135
+ self.skip_reading_prefix_cache = False
136
+
137
self._verify_step_pooling(pooler_config, valid_parameters)
138
139
def _verify_step_pooling(
vllm/sampling_params.py
+8
@@ -254,6 +254,8 @@ class SamplingParams(
254
generated token can complete the sequence."""
255
_bad_words_token_ids: list[list[int]] | None = None
256
257
+ skip_reading_prefix_cache: bool = None
258
+
259
@staticmethod
260
def from_optional(
261
n: int | None = 1,
@@ -414,6 +416,12 @@ class SamplingParams(
416
self.structured_outputs = self.guided_decoding
417
self.guided_decoding = None
418
419
+ if self.skip_reading_prefix_cache is None:
420
+ # If prefix caching is enabled,
421
+ # the output of prompt logprobs may less than n_prompt_tokens,
422
+ # we need to skip reading cache at this request.
423
+ self.skip_reading_prefix_cache = self.prompt_logprobs is not None
424
+
425
def _verify_args(self) -> None:
426
if not isinstance(self.n, int):
427
raise ValueError(f"n must be an int, but is of type {type(self.n)}")
vllm/v1/core/kv_cache_manager.py
+5
-6
@@ -185,12 +185,11 @@ class KVCacheManager:
185
- A list of blocks that are computed for the request.
186
- The number of computed tokens.
187
"""
188
- # Prefix caching is disabled or
189
- # When the request requires prompt logprobs, we skip prefix caching.
190
- if not self.enable_caching or (
191
- request.sampling_params is not None
192
- and request.sampling_params.prompt_logprobs is not None
193
- ):
188
+ # We skip finding the prefix cache hit when prefix caching is
189
+ # disabled or the request is marked as skipping kv cache read
190
+ # (which happens when the request requires prompt logprobs
191
+ # or calls a pooling model with all pooling).
192
+ if not self.enable_caching or request.skip_reading_prefix_cache:
193
return self.empty_kv_cache_blocks, 0
194
195
# NOTE: When all tokens hit the cache, we must recompute the last token
vllm/v1/request.py
+15
@@ -127,6 +127,8 @@ class Request:
127
self.get_hash_new_full_blocks = partial(block_hasher, self)
128
self.block_hashes = self.get_hash_new_full_blocks()
129
130
+ self.skip_reading_prefix_cache = self.get_skip_reading_prefix_cache()
131
+
132
@classmethod
133
def from_engine_core_request(
134
cls,
@@ -180,6 +182,19 @@ class Request:
182
def num_output_tokens(self) -> int:
183
return len(self._output_token_ids)
184
185
+ def get_skip_reading_prefix_cache(self) -> bool:
186
+ if (
187
+ self.sampling_params is not None
188
+ and self.sampling_params.skip_reading_prefix_cache is not None
189
+ ):
190
+ return self.sampling_params.skip_reading_prefix_cache
191
+ elif (
192
+ self.pooling_params is not None
193
+ and self.pooling_params.skip_reading_prefix_cache is not None
194
+ ):
195
+ return self.pooling_params.skip_reading_prefix_cache
196
+ return False
197
+
198
def is_finished(self) -> bool:
199
return RequestStatus.is_finished(self.status)
200