@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 77073c77bc

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 15, 12:15 UTC

[Core] Prevent side-channel attacks via cache salting (#17045)

Signed-off-by: Marko Rosenmueller <5467316+dr75@users.noreply.github.com>

Marko Rosenmueller committed Apr 30, 2025 at 14:27 UTC 77073c77bc2006eb80ea6d5128f076f5e6c6f54f

18 files changed +324 -122

docs/source/design/v1/prefix_caching.md

+19 -1

@@ -16,7 +16,7 @@ In the example above, the KV cache in the first block can be uniquely identified

17 * Parent hash value: The hash value of the parent hash block.

18 * Block tokens: A tuple of tokens in this block. The reason to include the exact tokens is to reduce potential hash value collision.

19 -* Extra hashes: Other values required to make this block unique, such as LoRA IDs and multi-modality input hashes (see the example below).

19 +* Extra hashes: Other values required to make this block unique, such as LoRA IDs, multi-modality input hashes (see the example below), and cache salts to isolate caches in multi-tenant environments.

21 > **Note 1:** We only cache full blocks.

@@ -76,6 +76,24 @@ Block 3

77 In the rest of this document, we first introduce the data structure used for prefix caching in vLLM v1, followed by the prefix caching workflow of major KV cache operators (e.g., allocate, append, free, eviction). Finally, we use an example to illustrate the end to end prefix caching workflow.

79 +**Cache Isolation for Security**

80 +To improve privacy in shared environments, vLLM supports isolating prefix cache reuse through optional per-request salting. By including a `cache_salt` in the request, this value is injected into the hash of the first block, ensuring that only requests with the same salt can reuse cached KV blocks. This prevents timing-based attacks where an adversary could infer cached content by observing latency differences. This offers protection without compromising performance.

81 +

82 +```json

83 +{

84 + "messages": [

85 + {"role": "system", "content": "You are a helpful assistant."},

86 + {"role": "user", "content": "Here is a document with details about the world series: ..."},

87 + {"role": "user", "content": "Who won the world series in 2020?"}

88 + ],

89 + "cache_salt": "Z3V2bmV3aGxza3ZubGFoZ3Zud3V3ZWZ2bmd0b3V2bnZmc2xpZ3RoZ2x2aQ=="

90 +}

91 +```

92 +

93 +With this setup, cache sharing is limited to users or requests that explicitly agree on a common salt, enabling cache reuse within a trust group while isolating others.

94 +

95 +> **Note:** Cache isolation is not supported in engine V0.

96 +

97 ## Data Structure

99 The prefix caching in vLLM v1 is implemented in the KV cache manager. The basic building block is the “Block” data class (simplified):

tests/entrypoints/openai/test_serving_chat.py

+40

@@ -272,3 +272,43 @@ def test_serving_chat_could_load_correct_generation_config():

272

273 assert mock_engine.generate.call_args.args[1].temperature == 0.0

274 assert mock_engine.generate.call_args.args[1].repetition_penalty == 1.05

275 +

276 +

277 +def test_serving_chat_did_set_correct_cache_salt():

278 + mock_model_config = MockModelConfig()

279 +

280 + mock_engine = MagicMock(spec=MQLLMEngineClient)

281 + mock_engine.get_tokenizer.return_value = get_tokenizer(MODEL_NAME)

282 + mock_engine.errored = False

283 +

284 + # Initialize the serving chat

285 + models = OpenAIServingModels(engine_client=mock_engine,

286 + base_model_paths=BASE_MODEL_PATHS,

287 + model_config=mock_model_config)

288 + serving_chat = OpenAIServingChat(mock_engine,

289 + mock_model_config,

290 + models,

291 + response_role="assistant",

292 + chat_template=CHAT_TEMPLATE,

293 + chat_template_content_format="auto",

294 + request_logger=None)

295 +

296 + # Test cache_salt

297 + req = ChatCompletionRequest(

298 + model=MODEL_NAME,

299 + messages=[{

300 + "role": "user",

301 + "content": "what is 1+1?"

302 + }],

303 + )

304 +

305 + # By default cache_salt in the engine prompt is not set

306 + with suppress(Exception):

307 + asyncio.run(serving_chat.create_chat_completion(req))

308 + assert "cache_salt" not in mock_engine.generate.call_args.args[0]

309 +

310 + # Test with certain cache_salt

311 + req.cache_salt = "test_salt"

312 + with suppress(Exception):

313 + asyncio.run(serving_chat.create_chat_completion(req))

314 + assert mock_engine.generate.call_args.args[0]["cache_salt"] == "test_salt"

tests/tokenization/test_detokenize.py

+10 -2

@@ -60,8 +60,16 @@ def _run_incremental_decode(tokenizer,

60 skip_special_tokens=skip_special_tokens,

61 spaces_between_special_tokens=spaces_between_special_tokens,

62 )

63 - request = EngineCoreRequest("", prompt_token_ids, None, None, None, params,

64 - None, 0.0, None)

63 + request = EngineCoreRequest("",

64 + prompt_token_ids,

65 + None,

66 + None,

67 + None,

68 + params,

69 + None,

70 + 0.0,

71 + None,

72 + cache_salt=None)

74 if fast is None:

75 detokenizer = IncrementalDetokenizer.from_new_request(

tests/v1/core/test_kv_cache_utils.py

+42 -1

@@ -29,7 +29,8 @@ from vllm.v1.request import Request

29 def make_request(request_id,

30 prompt_token_ids,

31 mm_positions=None,

32 - mm_hashes=None):

32 + mm_hashes=None,

33 + cache_salt=None):

34 if mm_positions is None:

35 multi_modal_inputs = None

36 else:

@@ -45,6 +46,7 @@ def make_request(request_id,

46 eos_token_id=100,

47 arrival_time=0,

48 lora_request=None,

49 + cache_salt=cache_salt,

50 )

@@ -213,6 +215,45 @@ def test_generate_block_hash_extra_keys_no_mm_inputs():

215 assert next_mm_idx == 0

216

217

218 +def test_generate_block_hash_extra_keys_cache_salt():

219 + request = make_request(

220 + request_id=0,

221 + prompt_token_ids=[_ for _ in range(6)],

222 + mm_positions=None,

223 + mm_hashes=None,

224 + cache_salt="salt",

225 + )

226 +

227 + # salt is added for the first token

228 + extra_keys, _ = generate_block_hash_extra_keys(request, 0, 1, 0)

229 + assert extra_keys == ('salt', )

230 + extra_keys, _ = generate_block_hash_extra_keys(request, 0, 10, 0)

231 + assert extra_keys == ('salt', )

232 +

233 + # no salt added for other tokens

234 + extra_keys, _ = generate_block_hash_extra_keys(request, 1, 2, 0)

235 + assert extra_keys is None

236 + extra_keys, _ = generate_block_hash_extra_keys(request, 6, 10, 0)

237 + assert extra_keys is None

238 +

239 + # works together with other extra keys

240 + request_mm = make_request(

241 + request_id=0,

242 + prompt_token_ids=[_ for _ in range(20)],

243 + mm_positions=[

244 + PlaceholderRange(offset=0, length=5),

245 + ],

246 + mm_hashes=["hash1"],

247 + cache_salt="salt",

248 + )

249 +

250 + # Test with no extra keys

251 + extra_keys, next_mm_idx = generate_block_hash_extra_keys(

252 + request_mm, 0, 5, 0)

253 + assert extra_keys == ("hash1", "salt")

254 + assert next_mm_idx == 1

255 +

256 +

257 @pytest.mark.parametrize("hash_fn", [sha256, hash])

258 def test_hash_block_tokens(hash_fn):

259 parent_block_hash = 123

tests/v1/core/test_prefix_caching.py

+63 -1

@@ -21,7 +21,8 @@ def make_request(request_id,

21 prompt_token_ids,

22 mm_positions=None,

23 mm_hashes=None,

24 - prompt_logprobs: Optional[int] = None):

24 + prompt_logprobs: Optional[int] = None,

25 + cache_salt: Optional[str] = None):

26 if mm_positions is None:

27 multi_modal_inputs = None

28 else:

@@ -38,6 +39,7 @@ def make_request(request_id,

39 eos_token_id=100,

40 arrival_time=0,

41 lora_request=None,

42 + cache_salt=cache_salt,

43 )

@@ -603,6 +605,66 @@ def test_mm_prefix_caching():

605 assert num_computed_tokens == 3 * 16

606

607

608 +def test_cache_key_salting():

609 + """

610 + This tests that cache salts are applied during hashing and the cache

611 + is separated cache as expected.

612 + """

613 + block_size = 16

614 + manager = KVCacheManager(

615 + make_kv_cache_config(block_size, 11),

616 + max_model_len=8192,

617 + enable_caching=True,

618 + )

619 +

620 + # 3 complete blocks and an incomplete block with 11 tokens.

621 + common_token_ids = [i for i in range(3) for _ in range(block_size)]

622 + token_ids = common_token_ids + [3] * 11

623 + req0 = make_request("0", token_ids, cache_salt="salt1")

624 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)

625 +

626 + # Completed block should have hashes with extra keys.

627 + assert not computed_blocks

628 + assert num_computed_tokens == 0

629 + block_hashes = manager.req_to_block_hashes[req0.request_id]

630 + assert len(block_hashes) == 3

631 + assert block_hashes[0].extra_keys == ("salt1", )

632 + assert block_hashes[1].extra_keys is None

633 + assert block_hashes[2].extra_keys is None

634 +

635 + blocks = manager.allocate_slots(req0, 59, computed_blocks)

636 + assert [b.block_id for b in blocks] == [1, 2, 3, 4]

637 + req0.num_computed_tokens = 59

638 +

639 + # Append slots without allocating a new block.

640 + for _ in range(5):

641 + req0.append_output_token_ids(8)

642 + new_blocks = manager.allocate_slots(req0, 5)

643 + assert new_blocks is not None and len(new_blocks) == 0

644 +

645 + # Now one more block that should not have extra keys.

646 + assert len(block_hashes) == 4

647 + assert block_hashes[3].extra_keys is None

648 +

649 + # Test cache hit with a new request that has the same salt.

650 + token_ids = common_token_ids + [4] * 11

651 + req1 = make_request("1", token_ids, cache_salt="salt1")

652 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)

653 + # Should match only a prefix of 3 blocks.

654 + assert len(computed_blocks) == 3

655 + assert num_computed_tokens == 3 * block_size

656 +

657 + # Test cache miss with same content but different salt.

658 + token_ids = common_token_ids + [4] * 11

659 + req2 = make_request("2", token_ids, cache_salt="salt2")

660 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)

661 + assert len(computed_blocks) == 0

662 + assert num_computed_tokens == 0

663 + block_hashes = manager.req_to_block_hashes[req2.request_id]

664 + assert len(block_hashes) == 3

665 + assert block_hashes[0].extra_keys == ("salt2", )

666 +

667 +

668 def test_prefill_not_enough_free_blocks_with_computed_blocks():

669 """

670 This is a unit test that tests the correctness of the allocate_slots

tests/v1/engine/test_engine_core.py

@@ -40,6 +40,7 @@ def make_request() -> EngineCoreRequest:

40 eos_token_id=None,

41 arrival_time=time.time(),

42 lora_request=None,

43 + cache_salt=None,

44 )

tests/v1/engine/test_engine_core_client.py

@@ -43,6 +43,7 @@ def make_request(params: SamplingParams) -> EngineCoreRequest:

43 eos_token_id=None,

44 arrival_time=time.time(),

45 lora_request=None,

46 + cache_salt=None,

47 )

tests/v1/engine/test_output_processor.py

+6 -1

@@ -57,6 +57,7 @@ def test_incremental_detokenization(request_output_kind: RequestOutputKind,

57 mm_placeholders=None,

58 eos_token_id=None,

59 lora_request=None,

60 + cache_salt=None,

61 sampling_params=SamplingParams(

62 skip_special_tokens=False,

63 spaces_between_special_tokens=False,

@@ -403,6 +404,7 @@ def test_logprobs_processor(request_output_kind: RequestOutputKind,

404 mm_placeholders=None,

405 eos_token_id=None,

406 lora_request=None,

407 + cache_salt=None,

408 sampling_params=SamplingParams(

409 skip_special_tokens=False,

410 spaces_between_special_tokens=False,

@@ -503,7 +505,7 @@ def test_stop_token(include_stop_str_in_output: bool,

505 reason should be "stop" (i.e. first control token causes stop

506 and is represented in output text)

507

506 - * else, the detokenized string should be

508 + * else, the detokenized string should be

509 <token><token>...<token> and the finish reason should be "stop"

510 (i.e. first control token causes stop but is not represented

511 in output text.)

@@ -565,6 +567,7 @@ def test_stop_token(include_stop_str_in_output: bool,

567 mm_placeholders=None,

568 eos_token_id=eos_token_id,

569 lora_request=None,

570 + cache_salt=None,

571 sampling_params=SamplingParams(

572 skip_special_tokens=False,

573 spaces_between_special_tokens=False,

@@ -661,6 +664,7 @@ def test_stop_string(include_stop_str_in_output: bool,

664 mm_placeholders=None,

665 eos_token_id=None,

666 lora_request=None,

667 + cache_salt=None,

668 sampling_params=SamplingParams(

669 skip_special_tokens=False,

670 spaces_between_special_tokens=False,

@@ -774,6 +778,7 @@ def test_iteration_stats(dummy_test_vectors):

778 mm_placeholders=None,

779 eos_token_id=None,

780 lora_request=None,

781 + cache_salt=None,

782 sampling_params=SamplingParams(),

783 ) for idx, prompt_tokens in enumerate(dummy_test_vectors.prompt_tokens)

784 ]

vllm/entrypoints/openai/protocol.py

+28 -4

@@ -14,6 +14,7 @@ from pydantic import (BaseModel, ConfigDict, Field, TypeAdapter,

14 ValidationInfo, field_validator, model_validator)

15 from typing_extensions import TypeAlias

17 +from vllm import envs

18 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam

19 from vllm.logger import init_logger

20 from vllm.pooling_params import PoolingParams

@@ -408,6 +409,15 @@ class ChatCompletionRequest(OpenAIBaseModel):

409 "If specified with 'logprobs', tokens are represented "

410 " as strings of the form 'token_id:{token_id}' so that tokens "

411 "that are not JSON-encodable can be identified."))

412 + cache_salt: Optional[str] = Field(

413 + default=None,

414 + description=(

415 + "If specified, the prefix cache will be salted with the provided "

416 + "string to prevent an attacker to guess prompts in multi-user "

417 + "environments. The salt should be random, protected from "

418 + "access by 3rd parties, and long enough to be "

419 + "unpredictable (e.g., 43 characters base64-encoded, corresponding "

420 + "to 256 bit). Not supported by vLLM engine V0."))

421

422 # doc: end-chat-completion-extra-params

423

@@ -726,6 +736,20 @@ class ChatCompletionRequest(OpenAIBaseModel):

736 "`add_generation_prompt` to True.")

737 return data

738

739 + @model_validator(mode="before")

740 + @classmethod

741 + def check_cache_salt_support(cls, data):

742 + if data.get("cache_salt") is not None:

743 + if not envs.VLLM_USE_V1:

744 + raise ValueError(

745 + "Parameter 'cache_salt' is not supported with "

746 + "this instance of vLLM, which uses engine V0.")

747 + if not isinstance(data["cache_salt"],

748 + str) or not data["cache_salt"]:

749 + raise ValueError("Parameter 'cache_salt' must be a "

750 + "non-empty string if provided.")

751 + return data

752 +

753

754 class CompletionRequest(OpenAIBaseModel):

755 # Ordered by official OpenAI API documentation

@@ -1622,9 +1646,9 @@ class TranscriptionRequest(OpenAIBaseModel):

1646

1647 # doc: begin-transcription-extra-params

1648 stream: Optional[bool] = False

1625 - """Custom field not present in the original OpenAI definition. When set,

1649 + """Custom field not present in the original OpenAI definition. When set,

1650 it will enable output to be streamed in a similar fashion as the Chat

1627 - Completion endpoint.

1651 + Completion endpoint.

1652 """

1653 # Flattened stream option to simplify form data.

1654 stream_include_usage: Optional[bool] = False

@@ -1642,7 +1666,7 @@ class TranscriptionRequest(OpenAIBaseModel):

1666 """

1667

1668 top_p: Optional[float] = None

1645 - """Enables nucleus (top-p) sampling, where tokens are selected from the

1669 + """Enables nucleus (top-p) sampling, where tokens are selected from the

1670 smallest possible set whose cumulative probability exceeds `p`.

1671 """

1672

@@ -1650,7 +1674,7 @@ class TranscriptionRequest(OpenAIBaseModel):

1674 """Limits sampling to the `k` most probable tokens at each step."""

1675

1676 min_p: Optional[float] = None

1653 - """Filters out tokens with a probability lower than `min_p`, ensuring a

1677 + """Filters out tokens with a probability lower than `min_p`, ensuring a

1678 minimum likelihood threshold during sampling.

1679 """

1680

vllm/entrypoints/openai/serving_engine.py

@@ -470,6 +470,9 @@ class OpenAIServing:

470 if request.mm_processor_kwargs is not None:

471 engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs

472

473 + if hasattr(request, "cache_salt") and request.cache_salt is not None:

474 + engine_prompt["cache_salt"] = request.cache_salt

475 +

476 return conversation, [request_prompt], [engine_prompt]

477

478 def _log_inputs(

vllm/inputs/data.py

+19 -1

@@ -28,6 +28,11 @@ class TextPrompt(TypedDict):

28 to pass the mm_processor_kwargs to each of them.

29 """

31 + cache_salt: NotRequired[str]

32 + """

33 + Optional cache salt to be used for prefix caching.

34 + """

35 +

37 class TokensPrompt(TypedDict):

38 """Schema for a tokenized prompt."""

@@ -52,6 +57,11 @@ class TokensPrompt(TypedDict):

57 to pass the mm_processor_kwargs to each of them.

58 """

60 + cache_salt: NotRequired[str]

61 + """

62 + Optional cache salt to be used for prefix caching.

63 + """

64 +

66 SingletonPrompt = Union[str, TextPrompt, TokensPrompt]

67 """

@@ -141,11 +151,17 @@ class TokenInputs(TypedDict):

151 The original prompt text corresponding to the token IDs, if available.

152 """

153

154 + cache_salt: NotRequired[str]

155 + """

156 + Optional cache salt to be used for prefix caching.

157 + """

158 +

159

160 def token_inputs(

161 prompt_token_ids: list[int],

162 token_type_ids: Optional[list[int]] = None,

163 prompt: Optional[str] = None,

164 + cache_salt: Optional[str] = None,

165 ) -> TokenInputs:

166 """Construct :class:`TokenInputs` from optional values."""

167 inputs = TokenInputs(type="token", prompt_token_ids=prompt_token_ids)

@@ -154,6 +170,8 @@ def token_inputs(

170 inputs["prompt"] = prompt

171 if token_type_ids is not None:

172 inputs["token_type_ids"] = token_type_ids

173 + if cache_salt is not None:

174 + inputs["cache_salt"] = cache_salt

175

176 return inputs

177

@@ -217,7 +235,7 @@ def zip_enc_dec_prompts(

235 """

236 Zip encoder and decoder prompts together into a list of

237 :class:`ExplicitEncoderDecoderPrompt` instances.

220 -

238 +

239 ``mm_processor_kwargs`` may also be provided; if a dict is passed, the same

240 dictionary will be used for every encoder/decoder prompt. If an iterable is

241 provided, it will be zipped with the encoder/decoder prompts.

vllm/inputs/preprocess.py

+72 -106

@@ -17,7 +17,8 @@ from vllm.transformers_utils.tokenizer_group import TokenizerGroup

18 from .data import (DecoderOnlyInputs, EncoderDecoderInputs, ProcessorInputs,

19 PromptType, SingletonInputs, SingletonPrompt, token_inputs)

20 -from .parse import is_explicit_encoder_decoder_prompt, parse_singleton_prompt

20 +from .parse import (ParsedStrPrompt, ParsedTextPrompt, ParsedTokensPrompt,

21 + is_explicit_encoder_decoder_prompt, parse_singleton_prompt)

23 logger = init_logger(__name__)

@@ -283,6 +284,29 @@ class InputPreprocessor:

284 return mm_processor.apply(prompt, mm_data, mm_processor_kwargs,

285 return_mm_hashes)

286

287 + def _get_prompt_data(self, parsed_prompt: Union[ParsedStrPrompt,

288 + ParsedTextPrompt,

289 + ParsedTokensPrompt]):

290 + prompt_text = None

291 + prompt_token_ids = None

292 + token_type_ids = None

293 + cache_salt = None

294 +

295 + if parsed_prompt["type"] == "str":

296 + prompt_text = parsed_prompt["content"]

297 + else:

298 + cache_salt = parsed_prompt["content"].get("cache_salt")

299 + if parsed_prompt["type"] == "text":

300 + prompt_text = parsed_prompt["content"]["prompt"]

301 + elif parsed_prompt["type"] == "tokens":

302 + prompt_token_ids = parsed_prompt["content"].get(

303 + "prompt_token_ids")

304 + token_type_ids = parsed_prompt["content"].get("token_type_ids")

305 + else:

306 + assert_never(parsed_prompt)

307 +

308 + return prompt_text, prompt_token_ids, token_type_ids, cache_salt

309 +

310 def _prompt_to_llm_inputs(

311 self,

312 prompt: SingletonPrompt,

@@ -304,70 +328,36 @@ class InputPreprocessor:

328 * :class:`SingletonInputs` instance

329 """

330 parsed = parse_singleton_prompt(prompt)

307 -

308 - if parsed["type"] == "str":

309 - prompt_text = parsed["content"]

310 - prompt_token_ids = self._tokenize_prompt(

311 - prompt_text,

331 + prompt_text, prompt_token_ids, token_type_ids, cache_salt = \

332 + self._get_prompt_data(parsed)

333 +

334 + # If multimodal data is present, process and return immediately

335 + if parsed["type"] != "str" and parsed["content"].get(

336 + "multi_modal_data") is not None:

337 + inputs = self._process_multimodal(

338 + prompt_text if prompt_text is not None else prompt_token_ids,

339 + parsed["content"]["multi_modal_data"],

340 + parsed["content"].get("mm_processor_kwargs"),

341 lora_request=lora_request,

313 - tokenization_kwargs=tokenization_kwargs,

314 - )

315 -

316 - return token_inputs(

317 - prompt=prompt_text,

318 - prompt_token_ids=prompt_token_ids,

342 + return_mm_hashes=return_mm_hashes,

343 )

344 + if cache_salt is not None:

345 + inputs["cache_salt"] = cache_salt

346 + return inputs

347

321 - if parsed["type"] == "tokens":

322 - tokens_content = parsed["content"]

323 -

324 - prompt_token_ids = tokens_content["prompt_token_ids"]

325 - token_type_ids = tokens_content.get("token_type_ids")

326 - multi_modal_data = tokens_content.get("multi_modal_data")

327 - mm_processor_kwargs = tokens_content.get("mm_processor_kwargs")

328 -

329 - if multi_modal_data is not None:

330 - return self._process_multimodal(

331 - prompt_token_ids,

332 - multi_modal_data,

333 - mm_processor_kwargs,

334 - lora_request=lora_request,

335 - return_mm_hashes=return_mm_hashes,

336 - )

337 -

338 - return token_inputs(

339 - prompt_token_ids=prompt_token_ids,

340 - token_type_ids=token_type_ids,

341 - )

342 -

343 - if parsed["type"] == "text":

344 - text_content = parsed["content"]

345 -

346 - prompt_text = text_content["prompt"]

347 - multi_modal_data = text_content.get("multi_modal_data")

348 - mm_processor_kwargs = text_content.get("mm_processor_kwargs")

349 -

350 - if multi_modal_data is not None:

351 - return self._process_multimodal(

352 - prompt_text,

353 - multi_modal_data,

354 - mm_processor_kwargs,

355 - lora_request=lora_request,

356 - return_mm_hashes=return_mm_hashes,

357 - )

358 -

348 + if prompt_token_ids is None:

349 prompt_token_ids = self._tokenize_prompt(

350 prompt_text,

351 lora_request=lora_request,

352 tokenization_kwargs=tokenization_kwargs,

353 )

354

365 - return token_inputs(

366 - prompt=prompt_text,

367 - prompt_token_ids=prompt_token_ids,

368 - )

369 -

370 - assert_never(parsed)

355 + return token_inputs(

356 + prompt=prompt_text,

357 + prompt_token_ids=prompt_token_ids,

358 + token_type_ids=token_type_ids,

359 + cache_salt=cache_salt,

360 + )

361

362 async def _prompt_to_llm_inputs_async(

363 self,

@@ -379,64 +369,35 @@ class InputPreprocessor:

369 """Async version of :meth:`_extract_prompt_components`."""

370 parsed = parse_singleton_prompt(prompt)

371

382 - if parsed["type"] == "str":

383 - prompt_text = parsed["content"]

384 - prompt_token_ids = await self._tokenize_prompt_async(

385 - prompt_text,

386 - lora_request=lora_request,

387 - tokenization_kwargs=tokenization_kwargs,

388 - )

372 + prompt_text, prompt_token_ids, token_type_ids, cache_salt = \

373 + self._get_prompt_data(parsed)

374

390 - return token_inputs(

391 - prompt=prompt_text,

392 - prompt_token_ids=prompt_token_ids,

375 + if parsed["type"] != "str" and parsed["content"].get(

376 + "multi_modal_data") is not None:

377 + inputs = await self._process_multimodal_async(

378 + prompt_token_ids if prompt_text is None else prompt_text,

379 + parsed["content"]["multi_modal_data"],

380 + parsed["content"].get("mm_processor_kwargs"),

381 + lora_request=lora_request,

382 + return_mm_hashes=return_mm_hashes,

383 )

384 + if cache_salt is not None:

385 + inputs["cache_salt"] = cache_salt

386 + return inputs

387

395 - if parsed["type"] == "tokens":

396 - tokens_content = parsed["content"]

397 -

398 - prompt_token_ids = tokens_content["prompt_token_ids"]

399 - multi_modal_data = tokens_content.get("multi_modal_data")

400 - mm_processor_kwargs = tokens_content.get("mm_processor_kwargs")

401 -

402 - if multi_modal_data is not None:

403 - return await self._process_multimodal_async(

404 - prompt_token_ids,

405 - multi_modal_data,

406 - mm_processor_kwargs,

407 - lora_request=lora_request,

408 - return_mm_hashes=return_mm_hashes,

409 - )

410 -

411 - return token_inputs(prompt_token_ids=prompt_token_ids)

412 -

413 - if parsed["type"] == "text":

414 - text_content = parsed["content"]

415 -

416 - prompt_text = text_content["prompt"]

417 - multi_modal_data = text_content.get("multi_modal_data")

418 - mm_processor_kwargs = text_content.get("mm_processor_kwargs")

419 -

420 - if multi_modal_data is not None:

421 - return await self._process_multimodal_async(

422 - prompt_text,

423 - multi_modal_data,

424 - mm_processor_kwargs,

425 - lora_request=lora_request,

426 - return_mm_hashes=return_mm_hashes,

427 - )

428 -

388 + if prompt_token_ids is None:

389 prompt_token_ids = await self._tokenize_prompt_async(

390 prompt_text,

391 lora_request=lora_request,

392 + tokenization_kwargs=tokenization_kwargs,

393 )

394

434 - return token_inputs(

435 - prompt=prompt_text,

436 - prompt_token_ids=prompt_token_ids,

437 - )

438 -

439 - assert_never(parsed)

395 + return token_inputs(

396 + prompt=prompt_text,

397 + prompt_token_ids=prompt_token_ids,

398 + token_type_ids=token_type_ids,

399 + cache_salt=cache_salt,

400 + )

401

402 def _build_enc_dec_llm_inputs(

403 self,

@@ -516,6 +477,11 @@ class InputPreprocessor:

477 mm_hashes=inputs["mm_hashes"],

478 mm_placeholders=inputs["mm_placeholders"],

479 )

480 +

481 + cache_salt = inputs.get("cache_salt")

482 + if cache_salt is not None:

483 + decoder_inputs["cache_salt"] = cache_salt

484 +

485 elif inputs["type"] == "token":

486 # Text-only inputs

487 encoder_inputs = token_inputs(prompt="", prompt_token_ids=[])

vllm/multimodal/inputs.py

@@ -826,6 +826,11 @@ class MultiModalInputs(TypedDict):

826 :code:`prompt_token_ids`.

827 """

828

829 + cache_salt: NotRequired[str]

830 + """

831 + Optional cache salt to be used for prefix caching.

832 + """

833 +

834

835 class MultiModalEncDecInputs(MultiModalInputs):

836 """

vllm/multimodal/processing.py

+1 -1

@@ -1789,7 +1789,7 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):

1789 mm_data: MultiModalDataDict,

1790 ) -> Union[str, list[int]]:

1791 """

1792 - Create input prompt for the encoder. HF processor will be applied on

1792 + Create input prompt for the encoder. HF processor will be applied on

1793 this prompt during profiling and generation.

1794 """

1795 raise NotImplementedError

vllm/v1/core/kv_cache_utils.py

+9 -4

@@ -275,7 +275,10 @@ def need_extra_keys(request: Request) -> bool:

275

276 # Multimodal requests need to include the MM hash.

277 # LoRA requests need to include the LoRA ID.

278 - return bool(request.mm_positions) or (request.lora_request is not None)

278 + # Request with provided cache salt need to include the salt.

279 + return bool(request.mm_positions) or (request.lora_request

280 + is not None) or (request.cache_salt

281 + is not None)

282

283

284 def _gen_mm_extra_hash_keys(request: Request, start_token_idx: int,

@@ -380,8 +383,10 @@ def generate_block_hash_extra_keys(

383 mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys(

384 request, start_token_idx, end_token_idx, start_mm_idx)

385 lora_extra_keys: list[int] = _gen_lora_extra_hash_keys(request)

386 + cache_salt_keys: list[str] = [request.cache_salt] if (

387 + start_token_idx == 0 and request.cache_salt) else []

388

384 - extra_keys: list[Any] = lora_extra_keys + mm_extra_keys

389 + extra_keys: list[Any] = lora_extra_keys + mm_extra_keys + cache_salt_keys

390

391 if not extra_keys:

392 return None, new_start_mm_idx

@@ -657,10 +662,10 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,

662 def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):

663 """

664 Only models with one type of KV cache are supported yet. This function tries

660 - to convert the KV cache specs to one type if the model is a hybrid model

665 + to convert the KV cache specs to one type if the model is a hybrid model

666 with multiple type of KV cache. It will convert all SlidingWindowSpec to

667 FullAttentionSpec if both types are present.

663 -

668 +

669 Args:

670 kv_cache_spec: The kv cache spec of each attention layer in the model

671 """

vllm/v1/engine/__init__.py

@@ -57,6 +57,7 @@ class EngineCoreRequest(

57 eos_token_id: Optional[int]

58 arrival_time: float

59 lora_request: Optional[LoRARequest]

60 + cache_salt: Optional[str]

62 # Used in DP case to indicate which wave of requests this is expected to

63 # belong to, to cover a race condition where the request is sent before

vllm/v1/engine/processor.py

@@ -317,6 +317,7 @@ class Processor:

317 eos_token_id=eos_token_id,

318 arrival_time=arrival_time,

319 lora_request=lora_request,

320 + cache_salt=decoder_inputs.get("cache_salt"),

321 )

322

323 def _validate_model_inputs(self,

vllm/v1/request.py

@@ -29,6 +29,7 @@ class Request:

29 arrival_time: float,

30 lora_request: Optional["LoRARequest"] = None,

31 structured_output_request: Optional["StructuredOutputRequest"] = None,

32 + cache_salt: Optional[str] = None,

33 ) -> None:

34 self.request_id = request_id

35 self.sampling_params = sampling_params

@@ -51,6 +52,7 @@ class Request:

52 self._all_token_ids: list[int] = self.prompt_token_ids.copy()

53 self.spec_token_ids: list[int] = []

54 self.num_computed_tokens = 0

55 + self.cache_salt: Optional[str] = cache_salt

57 # Multi-modal related

58 self.mm_positions = multi_modal_placeholders or []

@@ -89,6 +91,7 @@ class Request:

91 lora_request=request.lora_request,

92 structured_output_request=StructuredOutputRequest(

93 sampling_params=request.sampling_params),

94 + cache_salt=request.cache_salt,

95 )

97 def append_output_token_ids(