lmcache--lmcache
wehub-resource-sync
58 行
2.3 KiB
Python
58 行
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Blend V2 protocol definitions.
|
|
|
|
This module defines the V2 variants of blend lookup/retrieve protocols that use
|
|
CBMatchResult instead of plain (start, end) integer tuples. CBMatchResult carries
|
|
old/cur position ranges plus the pre-computed chunk hash, enabling direct storage
|
|
key lookup without re-hashing on the server side.
|
|
"""
|
|
|
|
# First Party
|
|
from lmcache.v1.multiprocess.custom_types import CBMatchResult, IPCCacheServerKey
|
|
from lmcache.v1.multiprocess.protocols.base import HandlerType, ProtocolDefinition
|
|
|
|
# Define request names for this protocol group
|
|
REQUEST_NAMES = [
|
|
"CB_LOOKUP_PRE_COMPUTED_V2",
|
|
"CB_RETRIEVE_PRE_COMPUTED_V2",
|
|
]
|
|
|
|
|
|
def get_protocol_definitions() -> dict[str, ProtocolDefinition]:
|
|
"""
|
|
Returns protocol definitions for blend V2 lookup/retrieve operations.
|
|
|
|
Returns:
|
|
Dictionary mapping request names to their protocol definitions
|
|
"""
|
|
return {
|
|
# Lookup pre-computed chunks (V2)
|
|
# Payload:
|
|
# - key: IPCCacheServerKey - The key containing the token ids
|
|
# Returns: List of CBMatchResult with match positions and chunk hashes
|
|
"CB_LOOKUP_PRE_COMPUTED_V2": ProtocolDefinition(
|
|
payload_classes=[IPCCacheServerKey],
|
|
response_class=list[CBMatchResult],
|
|
handler_type=HandlerType.BLOCKING,
|
|
),
|
|
# Retrieve pre-computed chunks (V2)
|
|
# Payload:
|
|
# - key: IPCCacheServerKey - The key containing the token ids
|
|
# - cb_match_result: list[CBMatchResult] - Match results returned by
|
|
# CB_LOOKUP_PRE_COMPUTED_V2, with per-chunk hashes
|
|
# and query positions
|
|
# - offset: int - The starting offset in the CB KV cache buffer
|
|
# - instance_id: int - Unique identifier for the vLLM instance
|
|
# - event_ipc_handle: bytes - IPC handle for event notification when the
|
|
# retrieval is complete
|
|
# Returns:
|
|
# - IPC handle bytes
|
|
# - boolean flag indicating if the retrieval is successful
|
|
"CB_RETRIEVE_PRE_COMPUTED_V2": ProtocolDefinition(
|
|
payload_classes=[IPCCacheServerKey, list[CBMatchResult], int, int, bytes],
|
|
response_class=tuple[bytes, bool],
|
|
handler_type=HandlerType.BLOCKING,
|
|
),
|
|
}
|