# SPDX-License-Identifier: Apache-2.0 """ vLLM compatibility notes: - PR#20511: Introduced kv_cache_utils.init_none_hash() https://github.com/vllm-project/vllm/pull/20511 - PR#23673: Renamed sha256_cbor_64bit to sha256_cbor https://github.com/vllm-project/vllm/pull/23673 - PR#27151: Moved hash functions to vllm.utils.hashing module https://github.com/vllm-project/vllm/pull/27151 TODO(baoloongmao): Move this to vllm_v1_adapter to decouple from vLLM """ # Standard from typing import Any, Iterable, List, Optional, Tuple, Union import abc import os # Third Party from transformers import AutoTokenizer import torch # First Party from lmcache.logging import init_logger from lmcache.utils import CacheEngineKey, _lmcache_nvtx_annotate from lmcache.v1.config import LMCacheEngineConfig from lmcache.v1.metadata import LMCacheMetadata logger = init_logger(__name__) NONE_HASH = 0 def _normalize_hash_to_int(hash_value: Union[int, bytes]) -> int: """Normalize hash outputs to LMCache's int chunk-hash representation. This function is triggered when vLLM's ``sha256_cbor`` hash function is used because it returns a 32-byte digest. vLLM's ``kv_cache_utils.init_none_hash`` can therefore initialize ``NONE_HASH`` as bytes, and direct hash calls for token chunks can also return bytes. LMCache stores chunk hashes in ``CacheEngineKey`` and serializes them with msgpack, so byte digests must be folded into uint64-compatible ints before they enter the prefix hash chain. This also keeps ``NONE_HASH`` and later prefix hashes using the same structural type for CBOR hashing. Args: hash_value: Hash output from vLLM or Python's builtin hash. Returns: The original int hash value, or the first eight bytes of a digest as a big-endian int. """ if isinstance(hash_value, bytes): return int.from_bytes(hash_value[:8], "big") return hash_value # Type alias for process_tokens return value # (start_index, end_index, cache_engine_key｜hash) ProcessTokensResult = Tuple[int, int, Union[CacheEngineKey, int]] class TokenDatabase(metaclass=abc.ABCMeta): """TokenDatabase is used to convert input tokens into list of cache engine keys. There are multiple ways to implement this: - ChunkedTokenDatabase: It processes tokens into chunks and convert each chunk into a cache engine key using prefix hash. - SegmentTokenDatabase: It processes tokens into segments based on special separators and convert each segment into a cache engine key. """ @abc.abstractmethod def __init__( self, config: Optional[LMCacheEngineConfig] = None, metadata: Optional[LMCacheMetadata] = None, ): global NONE_HASH hash_algorithm: str = ( config.pre_caching_hash_algorithm if config is not None else "builtin" ) # Get hash function with vLLM version compatibility self.hash_func = self._get_vllm_hash_func(hash_algorithm) # Initialize NONE_HASH (vLLM >= PR#20511) # NOTE: For centralized cache sharing, ensure PYTHONHASHSEED is # set consistently across all processes (e.g., export PYTHONHASHSEED=0). try: # Third Party from vllm.v1.core import kv_cache_utils if hasattr(kv_cache_utils, "init_none_hash"): kv_cache_utils.init_none_hash(self.hash_func) NONE_HASH = _normalize_hash_to_int(kv_cache_utils.NONE_HASH) logger.info( "Initialized NONE_HASH=%s from vLLM (>= PR#20511)", NONE_HASH ) else: NONE_HASH = 0 logger.info("Using default NONE_HASH=0 (vLLM < PR#20511)") except (ImportError, AttributeError): NONE_HASH = 0 logger.info("Using default NONE_HASH=0 (vLLM not available)") logger.info("Using hash algorithm: %s", hash_algorithm) self.metadata = metadata # Whether only the first rank should save cache. This flag is also used # to control the logical world_size embedded into CacheEngineKey. self.save_only_first_rank = False if config is not None and metadata is not None: # save_only_first_rank only works when use MLA, follow the same # semantics as LMCacheEngine and memory allocator. self.save_only_first_rank = ( config.get_extra_config_value("save_only_first_rank", metadata.use_mla) and metadata.use_mla ) def _get_vllm_hash_func(self, hash_algorithm: str): """Get hash function from vLLM with version compatibility. Tries multiple import paths to support different vLLM versions: - vllm.utils.hashing.get_hash_fn_by_name (>= PR#27151) - vllm.utils.get_hash_fn_by_name (< PR#27151) - Direct imports as fallback - sha256_cbor_64bit -> sha256_cbor rename (PR#23673) """ # Try get_hash_fn_by_name from both locations (PR#27151) for module_path in ["vllm.utils.hashing", "vllm.utils"]: try: module = __import__(module_path, fromlist=["get_hash_fn_by_name"]) get_hash_fn_by_name = module.get_hash_fn_by_name return self._try_get_hash( get_hash_fn_by_name, hash_algorithm, module_path ) except (ImportError, AttributeError, ValueError): continue # Try direct imports as fallback (for older vLLM versions) func_names = ( ["sha256_cbor", "sha256_cbor_64bit"] if hash_algorithm in ("sha256_cbor", "sha256_cbor_64bit") else [hash_algorithm] ) for module_path in ["vllm.utils.hashing", "vllm.utils"]: for func_name in func_names: try: module = __import__(module_path, fromlist=[func_name]) hash_func = getattr(module, func_name) logger.info( "Loaded '%s' from %s (direct import)", func_name, module_path ) return hash_func except (ImportError, AttributeError): continue # Fallback to builtin hash logger.warning( "Could not load '%s' from vLLM. Using builtin hash. " "This may cause inconsistencies in distributed caching.", hash_algorithm, ) # Check PYTHONHASHSEED when using builtin hash if os.getenv("PYTHONHASHSEED") is None: logger.warning( "Using builtin hash without PYTHONHASHSEED set. " "For production environments (non-testing scenarios), you MUST set " "PYTHONHASHSEED to ensure consistent hashing across processes. " "Example: export PYTHONHASHSEED=0" ) return hash def _try_get_hash(self, get_hash_fn_by_name, hash_algorithm: str, module_name: str): """Try to get hash function, handling sha256_cbor_64bit rename.""" # Handle sha256_cbor_64bit -> sha256_cbor rename (PR#23673) names_to_try = ( ["sha256_cbor", "sha256_cbor_64bit"] if hash_algorithm in ("sha256_cbor", "sha256_cbor_64bit") else [hash_algorithm] ) for name in names_to_try: try: hash_func = get_hash_fn_by_name(name) logger.info("Loaded '%s' from %s", name, module_name) return hash_func except ValueError: continue raise ValueError(f"Hash function '{hash_algorithm}' not found in {module_name}") @abc.abstractmethod def process_tokens( self, tokens: Optional[Union[torch.Tensor, List[int]]] = None, hashes: Optional[List[int]] = None, offsets: Optional[List[int]] = None, mask: Optional[torch.Tensor] = None, make_key: bool = True, request_configs: Optional[dict] = None, ) -> Iterable[ProcessTokensResult]: """Process the tokens and return the corresponding cache engine keys. :param Optional[Union[torch.Tensor, List[int]]] tokens: The tokens to process. :param Optional[List[int]] hashes: The hashes to process. If provided, it will be used instead of tokens to generate cache engine keys. :param Optional[List[int]] offsets: The number of tokens in each chunk. :param Optional[torch.Tensor] mask: The mask for the tokens. Should have the same length as tokens. And the mask should ALWAYS be like FFFFFTTTTTTT, where True means the tokens needs to be matched, and the Falses will ALWAYS be at the PREFIX of the tensor. :param bool make_key: Whether to make the cache engine key or not. If False, the hash value will be returned instead. :param Optional[dict] request_configs: The configs of the request. :returns: A iterable of tuples with three elements. The first element is the start index of the tokens for the key. The second element is the end index of the tokens for the key. The third element is the cache engine key (or hash) for the tokens. """ raise NotImplementedError def _make_key_by_hash( self, chunk_hash: int, request_configs: Optional[dict] = None ): assert self.metadata is not None # When save_only_first_rank is enabled (for MLA), we deliberately # collapse the CacheEngineKey.world_size to 1 so that cache keys # become world-size agnostic across compatible deployments. return CacheEngineKey( self.metadata.model_name, self.metadata.world_size if not self.save_only_first_rank else 1, self.metadata.worker_id, chunk_hash, self.metadata.kv_dtype, request_configs, ) def _canonicalize_hash_inputs( self, prefix_hash: Optional[int], tokens_tuple: Tuple[int, ...], extra_keys: Optional[List[Any]], ) -> Tuple[int, Tuple[int, ...], Tuple[Any, ...]]: """ Canonicalize hash inputs so that semantically identical requests produce structurally identical hash inputs across instances. - prefix_hash: int or NONE_HASH if None - tokens_tuple: tuple of token IDs - extra_keys: tuple of additional keys, empty if None """ return ( prefix_hash if prefix_hash is not None else NONE_HASH, tokens_tuple, tuple(extra_keys) if extra_keys is not None else (), ) def _hash_tokens( self, tokens: Union[torch.Tensor, List[int]], prefix_hash: Optional[int] = None, extra_keys: Optional[list[Any]] = None, ) -> int: if isinstance(tokens, torch.Tensor): tokens_tuple = tuple(tokens.cpu().tolist()) elif isinstance(tokens, list): tokens_tuple = tuple(tokens) else: raise ValueError(f"Unsupported tokens type: {type(tokens)}") # Ignore extra keys for now # Extra keys are for multi-modal inputs and # request specific metadata (e.g., LoRA ID). # Use default values for None to maintain a fixed tuple structure for hashing. # Use helper to canonicalize inputs to ensure consistent hashing # This replaces the logic that was causing inconsistency canon_prefix, canon_tokens, canon_extra = self._canonicalize_hash_inputs( prefix_hash, tokens_tuple, extra_keys ) return _normalize_hash_to_int( self.hash_func((canon_prefix, canon_tokens, canon_extra)) ) class ChunkedTokenDatabase(TokenDatabase): def __init__( self, config: Optional[LMCacheEngineConfig] = None, metadata: Optional[LMCacheMetadata] = None, ): super(ChunkedTokenDatabase, self).__init__(config, metadata) if config is not None: self.config = config self.chunk_size = config.chunk_size # Check for cross-process cache sharing setup if os.getenv("PYTHONHASHSEED") is None: if config.remote_url is not None: logger.warning( "Centralized cache sharing detected " "but PYTHONHASHSEED not set. " "For consistent caching, set: export PYTHONHASHSEED=0 " "before the engine starts." ) if config.enable_pd: logger.error( "P/D Disaggregation detected " "but PYTHONHASHSEED not set. " "For consistent caching, set: export PYTHONHASHSEED=0 " "before the engine starts. " "This will cause incorrect KV cache transfer." ) else: # Default values self.config = None self.chunk_size = 256 def _get_init_hash(self) -> int: return NONE_HASH def _chunk_tokens( self, tokens: Union[torch.Tensor, List[int]], ) -> Iterable[Union[torch.Tensor, List[int]]]: """ Chunk the tokens into chunks of size self.chunk_size. :param tokens: the input tokens, with shape [seq_len] device: the target device after chunking :return: a generator of chunks of tokens, each with shape [chunk_size] """ save_unfull_chunk = ( self.config.save_unfull_chunk if self.config is not None else True ) end = ( len(tokens) if save_unfull_chunk else (len(tokens) - len(tokens) % self.chunk_size) ) for i in range(0, end, self.chunk_size): yield tokens[i : i + self.chunk_size] def _prefix_hash( self, token_chunks: Iterable[Union[torch.Tensor, List[int]]], ) -> Iterable[int]: prefix_hash = self._get_init_hash() for token_chunk in token_chunks: prefix_hash = self._hash_tokens(token_chunk, prefix_hash) yield prefix_hash @_lmcache_nvtx_annotate def process_tokens( self, tokens: Optional[Union[torch.Tensor, List[int]]] = None, hashes: Optional[List[int]] = None, offsets: Optional[List[int]] = None, mask: Optional[torch.Tensor] = None, make_key: bool = True, request_configs: Optional[dict] = None, ) -> Iterable[ProcessTokensResult]: """Process the tokens/hashes and return the corresponding cache engine keys. :param Optional[Union[torch.Tensor, List[int]]] tokens: The tokens to process. :param Optional[List[int]] hashes: The hashes to process. If provided, it will be used instead of tokens to generate cache engine keys. :param Optional[List[int]] offsets: The number of tokens in each chunk. :param Optional[torch.Tensor] mask: The mask for the tokens. Should have the same length as tokens. And the mask should ALWAYS be like FFFFFTTTTTTT, where True means the tokens needs to be matched, and the Falses will ALWAYS be at the PREFIX of the tensor. :param bool make_key: Whether to make the cache engine key or not. If False, the hash value will be returned instead. :param Optional[dict] request_configs: The configs of the request. :returns: A iterable of tuples with three elements. The first element is the start index of the tokens for the key. The second element is the end index of the tokens for the key. The third element is the cache engine key (or hash) for the tokens. :raises: ValueError if the number of Falses in the mask is not a multiple of the chunk size. """ if mask is not None: num_falses = mask.numel() - mask.long().sum().item() else: num_falses = 0 if num_falses % self.chunk_size != 0: raise ValueError( "The number of Falses in the mask is not a multiple of the chunk size." ) if tokens is not None: total_len = len(tokens) token_chunks = self._chunk_tokens(tokens) prefix_hashes = self._prefix_hash(token_chunks) for chunk_id, hash_val in enumerate(prefix_hashes): start_idx = chunk_id * self.chunk_size end_idx = min(start_idx + self.chunk_size, total_len) if start_idx < num_falses: continue else: if make_key: yield ( start_idx, end_idx, self._make_key_by_hash(hash_val, request_configs), ) else: yield start_idx, end_idx, hash_val elif hashes is not None: assert offsets is not None, ( "If hashes are provided, offsets must also be provided." ) start_idx = 0 for hash_val, offset in zip(hashes, offsets, strict=False): end_idx = start_idx + offset if make_key: yield ( start_idx, end_idx, self._make_key_by_hash(hash_val, request_configs), ) else: yield start_idx, end_idx, hash_val start_idx = end_idx else: raise ValueError("Either tokens or hashes must be provided.") class SegmentTokenDatabase(TokenDatabase): """ Currently, we still use special separators to identify chunks. In the future, we might need to implement a fast substring match. """ def __init__(self, config: LMCacheEngineConfig, metadata: LMCacheMetadata): super(SegmentTokenDatabase, self).__init__(config, metadata) self.tokenizer = AutoTokenizer.from_pretrained(metadata.model_name) # TODO (Jiayi): figure out how to decide when # to use `1:` (whether there's a special starting token # in the beginning) self.sep_tokens = self.tokenizer.encode(config.blend_special_str)[1:] self.sep_tokens = torch.tensor(self.sep_tokens, device="cpu") self.sep_len = len(self.sep_tokens) def _fast_split_by_subtensor(self, tokens: torch.Tensor) -> Iterable[torch.Tensor]: """Match the `sep_tokens` with sliding windows""" if self.sep_len == 0 or len(tokens) < self.sep_len: yield tokens # Unfold into sliding windows # shape: (num_tokens-sep_len+1, sep_len) windows = tokens.unfold(0, self.sep_len, 1) # Compare each window with sep_tokens matches = ( (windows == self.sep_tokens).all(dim=1).nonzero(as_tuple=True)[0].tolist() ) # Split based on matches start = 0 for idx in matches: yield tokens[start:idx] start = idx + self.sep_len # yield last chunk yield tokens[start:] def process_tokens( self, tokens: Optional[Union[torch.Tensor, List[int]]] = None, hashes: Optional[List[int]] = None, offsets: Optional[List[int]] = None, mask: Optional[torch.Tensor] = None, make_key: bool = True, request_configs: Optional[dict] = None, ) -> Iterable[ProcessTokensResult]: """Process the tokens and return the corresponding cache engine keys. :param Union[torch.Tensor, List[int]] tokens: The tokens to process. :param Optional[List[int]] hashes: The hashes to process. If provided, it will be used instead of tokens to generate cache engine keys. :param Optional[List[int]] offsets: The number of tokens in each chunk. :param Optional[torch.Tensor] mask: The mask for the tokens. Should have the same length as tokens. And the mask should ALWAYS be like FFFFFTTTTTTT, where True means the tokens needs to be matched, and the Falses will ALWAYS be at the PREFIX of the tensor. :param bool make_key: Whether to make the cache engine key or not. If False, the hash value will be returned instead. :param Optional[dict] request_configs: The configs of the request. :returns: A iterable of tuples with three elements. The first element is the start index of the tokens for the key. The second element is the end index of the tokens for the key. The third element is the cache engine key for the tokens. """ if tokens is not None: if not isinstance(tokens, torch.Tensor): tokens = torch.tensor(tokens, dtype=torch.long, device="cpu") else: tokens = tokens.to(device="cpu", dtype=torch.long) if mask is not None: num_falses = mask.numel() - mask.long().sum().item() else: num_falses = 0 assert num_falses < len(tokens), ( "The number of Falses in the mask shouldn't " "be less than the length of tokens." ) token_chunks = self._fast_split_by_subtensor(tokens) start_idx = 0 for idx, token_chunk in enumerate(token_chunks): token_chunk_len = len(token_chunk) end_idx = start_idx + token_chunk_len if idx > 0: start_idx += self.sep_len end_idx += self.sep_len if start_idx >= num_falses: if make_key: yield ( start_idx, end_idx, self._make_key_by_hash( self._hash_tokens(token_chunk), request_configs ), ) else: yield start_idx, end_idx, self._hash_tokens(token_chunk) start_idx = end_idx elif hashes is not None: assert offsets is not None, ( "If hashes are provided, offsets must also be provided." ) start_idx = 0 for hash_val, offset in zip(hashes, offsets, strict=False): end_idx = start_idx + offset if make_key: yield ( start_idx, end_idx, self._make_key_by_hash(hash_val, request_configs), ) else: yield start_idx, end_idx, hash_val start_idx = end_idx else: raise ValueError("Either tokens or hashes must be provided.")