minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

feat(hicache): Support passing prefix keys for l3 store. (#9045)

Browse Source 
Co-authored-by: pansicheng <sicheng.pan.chn@gmail.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
hzh0425
2025-10-10 00:22:05 -07:00
committed by GitHub
co-authored by pansicheng Zhiqiang Xie
parent d8467db727
commit ee3bd8a1c8
11 changed files with 107 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/cache_controller.py
+45 -13
View File
@@ -22,7 +22,10 @@ from typing import TYPE_CHECKING, List, NamedTuple, Optional, Set, Tuple
import torch
from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
from sglang.srt.mem_cache.hicache_storage import (
HiCacheStorageConfig,
HiCacheStorageExtraInfo,
)
if TYPE_CHECKING:
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
@@ -191,12 +194,14 @@ class StorageOperation:
token_ids: List[int],
last_hash: Optional[str] = None,
hash_value: Optional[List[str]] = None,
prefix_keys: Optional[List[str]] = None,
):
self.host_indices = host_indices
self.token_ids = token_ids
self.last_hash = last_hash
self.completed_tokens = 0
self.hash_value = hash_value if hash_value is not None else []
self.prefix_keys = prefix_keys
self.id = StorageOperation.counter
StorageOperation.counter += 1
@@ -212,6 +217,7 @@ class PrefetchOperation(StorageOperation):
host_indices: torch.Tensor,
token_ids: List[int],
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
):
self.request_id = request_id
@@ -219,7 +225,7 @@ class PrefetchOperation(StorageOperation):
self._terminated_flag = False
self.start_time = time.monotonic()
super().__init__(host_indices, token_ids, last_hash)
super().__init__(host_indices, token_ids, last_hash, prefix_keys=prefix_keys)
def increment(self, num_tokens: int):
with self._lock:
@@ -550,12 +556,13 @@ class HiCacheController:
host_indices: torch.Tensor,
new_input_tokens: List[int],
last_hash: Optional[str] = None,
prefix_keys: Optional[List[str]] = None,
) -> PrefetchOperation:
"""
Prefetch KV caches from storage backend to host memory.
"""
operation = PrefetchOperation(
request_id, host_indices, new_input_tokens, last_hash
request_id, host_indices, new_input_tokens, last_hash, prefix_keys
)
self.prefetch_queue.put(operation)
return operation
@@ -571,8 +578,12 @@ class HiCacheController:
for page in pages:
self.host_mem_release_queue.put(page)
def _page_get_zero_copy(self, operation, hash_values, host_indices):
results = self.storage_backend.batch_get_v1(hash_values, host_indices)
def _page_get_zero_copy(
self, operation, hash_values, host_indices, extra_info=None
):
results = self.storage_backend.batch_get_v1(
hash_values, host_indices, extra_info
)
inc = 0
for i in range(len(hash_values)):
if not results[i]:
@@ -584,7 +595,7 @@ class HiCacheController:
operation.increment(inc)
# todo: deprecate
def _generic_page_get(self, operation, hash_values, host_indices):
def _generic_page_get(self, operation, hash_values, host_indices, extra_info=None):
dummy_page_dst = [
self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values
]
@@ -608,6 +619,7 @@ class HiCacheController:
def _page_transfer(self, operation):
# Transfer batch by batch
prefix_keys = operation.prefix_keys
for i in range(0, len(operation.hash_value), self.storage_batch_size):
batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
batch_host_indices = operation.host_indices[
@@ -615,7 +627,8 @@ class HiCacheController:
]
prev_completed_tokens = operation.completed_tokens
# Get one batch token, and update the completed_tokens if succeed
self.page_get_func(operation, batch_hashes, batch_host_indices)
extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
self.page_get_func(operation, batch_hashes, batch_host_indices, extra_info)
# Check termination
if (
operation.completed_tokens
@@ -623,6 +636,10 @@ class HiCacheController:
):
operation.mark_terminate()
break # Some operations fail or operation terminated by controller
if prefix_keys and len(prefix_keys) > 0:
prefix_keys += batch_hashes
# release pre-allocated memory
self.append_host_mem_release(
operation.host_indices[operation.completed_tokens :]
@@ -656,6 +673,7 @@ class HiCacheController:
def _storage_hit_query(self, operation) -> tuple[list[str], int]:
last_hash = operation.last_hash
tokens_to_fetch = operation.token_ids
prefix_keys = operation.prefix_keys.copy() if operation.prefix_keys else None
storage_query_count = 0
hash_value = []
@@ -673,11 +691,15 @@ class HiCacheController:
batch_tokens[i : i + self.page_size], last_hash
)
batch_hashes.append(last_hash)
hit_page_num = self.storage_backend.batch_exists(batch_hashes)
extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info)
hash_value.extend(batch_hashes[:hit_page_num])
storage_query_count += hit_page_num * self.page_size
if hit_page_num < len(batch_hashes):
break
if prefix_keys and len(prefix_keys) > 0:
prefix_keys += batch_hashes
return hash_value, storage_query_count
def prefetch_thread_func(self):
@@ -734,28 +756,34 @@ class HiCacheController:
host_indices: torch.Tensor,
token_ids: List[int],
hash_value: Optional[List[str]] = None,
prefix_keys: Optional[List[str]] = None,
) -> int:
"""
Write KV caches from host memory to storage backend.
"""
operation = StorageOperation(host_indices, token_ids, hash_value=hash_value)
operation = StorageOperation(
host_indices, token_ids, hash_value=hash_value, prefix_keys=prefix_keys
)
self.backup_queue.put(operation)
return operation.id
# todo: deprecate
def _generic_page_set(self, hash_values, host_indices) -> bool:
def _generic_page_set(self, hash_values, host_indices, extra_info=None) -> bool:
data = [
self.mem_pool_host.get_data_page(host_indices[i * self.page_size])
for i in range(len(hash_values))
]
return self.storage_backend.batch_set(hash_values, data)
def _page_set_zero_copy(self, hash_values, host_indices) -> bool:
return all(self.storage_backend.batch_set_v1(hash_values, host_indices))
def _page_set_zero_copy(self, hash_values, host_indices, extra_info=None) -> bool:
return all(
self.storage_backend.batch_set_v1(hash_values, host_indices, extra_info)
)
# Backup batch by batch
def _page_backup(self, operation):
# Backup batch by batch
prefix_keys = operation.prefix_keys
for i in range(0, len(operation.hash_value), self.storage_batch_size):
batch_hashes = operation.hash_value[i : i + self.storage_batch_size]
batch_host_indices = operation.host_indices[
@@ -763,12 +791,16 @@ class HiCacheController:
]
# Set one batch token, and record if success.
# todo: allow partial success
success = self.page_set_func(batch_hashes, batch_host_indices)
extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys)
success = self.page_set_func(batch_hashes, batch_host_indices, extra_info)
if not success:
logger.warning(
f"Write page to storage: {len(batch_hashes)} pages failed."
)
break
if prefix_keys and len(prefix_keys) > 0:
prefix_keys += batch_hashes
operation.completed_tokens += self.page_size * len(batch_hashes)
def backup_thread_func(self):