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
[Prefix Cache] Use LoRA name for consistent KV-cache block hashing (#27211)
Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>
Sage committed
Oct 22, 2025 at 21:13 UTC
1651003c35d263fbc7d87c2e75b9cca5d590eb27
2 files changed
+25
-6
tests/v1/core/test_kv_cache_utils.py
+19
@@ -8,6 +8,7 @@ import torch
8
9
import vllm.v1.core.kv_cache_utils as kv_cache_utils
10
from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
11
+from vllm.lora.request import LoRARequest
12
from vllm.multimodal.inputs import (
13
MultiModalFeatureSpec,
14
MultiModalKwargsItem,
@@ -449,6 +450,24 @@ def test_generate_block_hash_extra_keys_cache_salt():
450
assert next_mm_idx == 1
451
452
453
+def test_generate_block_hash_extra_keys_lora():
454
+ request = make_request(
455
+ request_id="0",
456
+ prompt_token_ids=[_ for _ in range(6)],
457
+ )
458
+
459
+ request.lora_request = LoRARequest(
460
+ lora_name="test_lora_adapter", lora_int_id=1, lora_path="/path/to/lora"
461
+ )
462
+
463
+ extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0)
464
+ assert extra_keys == ("test_lora_adapter",)
465
+
466
+ request.lora_request = None
467
+ extra_keys, _ = generate_block_hash_extra_keys(request, 0, 3, 0)
468
+ assert extra_keys is None
469
+
470
+
471
@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
472
def test_hash_block_tokens(hash_fn):
473
parent_block_hash = BlockHash(b"123")
vllm/v1/core/kv_cache_utils.py
+6
-6
@@ -373,7 +373,7 @@ def need_extra_keys(request: Request) -> bool:
373
"""
374
375
# Multimodal requests need to include the MM hash.
376
- # LoRA requests need to include the LoRA ID.
376
+ # LoRA requests need to include the LoRA name.
377
# Request with provided cache salt need to include the salt.
378
return (
379
bool(request.mm_features)
@@ -446,26 +446,26 @@ def _gen_mm_extra_hash_keys(
446
return extra_keys, curr_mm_idx
447
448
449
-def _gen_lora_extra_hash_keys(request: Request) -> list[int]:
449
+def _gen_lora_extra_hash_keys(request: Request) -> list[str]:
450
"""Generate extra keys related to LoRA for block hash computation.
451
452
Args:
453
request: The request object.
454
455
Returns:
456
- Return LoRA id of the request if it is a LoRA request. Return empty
456
+ Return LoRA name of the request if it is a LoRA request. Return empty
457
list otherwise.
458
"""
459
if not request.lora_request:
460
return []
461
- return [request.lora_request.lora_int_id]
461
+ return [request.lora_request.lora_name]
462
463
464
def generate_block_hash_extra_keys(
465
request: Request, start_token_idx: int, end_token_idx: int, start_mm_idx: int
466
) -> tuple[tuple[Any, ...] | None, int]:
467
"""Generate extra keys for the block hash. The extra keys can come from
468
- the multi-modal inputs and request specific metadata (e.g., LoRA ID).
468
+ the multi-modal inputs and request specific metadata (e.g., LoRA name).
469
470
Args:
471
request: The request object.
@@ -480,7 +480,7 @@ def generate_block_hash_extra_keys(
480
mm_extra_keys, new_start_mm_idx = _gen_mm_extra_hash_keys(
481
request, start_token_idx, end_token_idx, start_mm_idx
482
)
483
- lora_extra_keys: list[int] = _gen_lora_extra_hash_keys(request)
483
+ lora_extra_keys: list[str] = _gen_lora_extra_hash_keys(request)
484
cache_salt_keys: list[str] = (
485
[request.cache_salt] if (start_token_idx == 0 and request.cache_salt) else []
486
)