Skip to content

GitLab

Unverified Commit 80dc76e1 authored by ykwd's avatar ykwd Committed by GitHub
Browse files

[Fix] HiCache Bugfix & Mooncake Error Handling Enhance (#8901) 


Co-authored-by: default avatarZhiqiang Xie <xiezhq@stanford.edu>
parent 9b08d975
Hide whitespace changes
Inline Side-by-side
Showing with 285 additions and 246 deletions
python/sglang/srt/managers/cache_controller.py
View file @ 80dc76e1
......@@ -27,7 +27,7 @@ if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool_host import HostKVCache
from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
logger = logging.getLogger(__name__)
......@@ -240,28 +240,38 @@ class HiCacheController:
self.io_backend = io_backend
self.enable_storage = False
self.is_mla = isinstance(self.mem_pool_host, MLATokenToKVPoolHost)
# todo: move backend initialization to storage backend module
if storage_backend is not None:
self.storage_backend_type = storage_backend
from sglang.srt.mem_cache.hicache_storage import HiCacheFile, get_hash_str
from sglang.srt.mem_cache.hicache_storage import get_hash_str
self.get_hash_str = get_hash_str
# Currently, AscendMLAPagedTokenToKVPool is the subclass of MLATokenToKVPool.
is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool)
# In MLA backend, only one rank needs to backup the KV cache
self.backup_skip = (
is_mla_backend
# todo: for load balancing, decide which rank to backup the KV cache by hash value
and get_tensor_model_parallel_rank() != 0
# todo: support other storage backends
and self.storage_backend_type in ["file", "mooncake"]
)
if storage_backend == "file":
self.storage_backend = HiCacheFile(is_mla=self.is_mla)
self.get_hash_str = get_hash_str
from sglang.srt.mem_cache.hicache_storage import HiCacheFile
self.storage_backend = HiCacheFile(is_mla_backend=is_mla_backend)
elif storage_backend == "nixl":
from sglang.srt.mem_cache.storage.nixl.hicache_nixl import HiCacheNixl
self.storage_backend = HiCacheNixl()
self.get_hash_str = get_hash_str
elif storage_backend == "mooncake":
from sglang.srt.mem_cache.storage.mooncake_store.mooncake_store import (
MooncakeStore,
get_hash_str_mooncake,
)
self.storage_backend = MooncakeStore(is_mla=self.is_mla)
self.get_hash_str = get_hash_str_mooncake
self.storage_backend = MooncakeStore(is_mla_backend=is_mla_backend)
self.storage_backend.register_buffer(self.mem_pool_host.kv_buffer)
assert self.mem_pool_host.layout == "page_first"
elif storage_backend == "hf3fs":
......@@ -281,7 +291,6 @@ class HiCacheController:
self.storage_backend = HiCacheHF3FS.from_env_config(
bytes_per_page, dtype
)
self.get_hash_str = get_hash_str
else:
raise NotImplementedError(
f"Unsupported storage backend: {storage_backend}"
......@@ -400,15 +409,6 @@ class HiCacheController:
self.prefetch_thread.start()
self.backup_thread.start()
@property
def backup_skip(self):
return (
self.is_mla
and get_tensor_model_parallel_rank() != 0
# todo: only support file and mooncake
and self.storage_backend_type in ["file", "mooncake"]
)
def write(
self,
device_indices: torch.Tensor,
......@@ -570,57 +570,91 @@ class HiCacheController:
operation.mark_done()
return operation.completed_tokens, operation.hash_value
def zerocopy_page_transfer(self, operation, batch_size=8):
# zero copy
def _3fs_zero_copy_page_get(self, operation, hash_values, host_indices):
hashes, dsts = self.mem_pool_host.get_buffer_with_hash(
operation.hash_value, operation.host_indices
hash_values, host_indices
)
page_data = self.storage_backend.batch_get(hashes, dsts)
if page_data:
operation.increment(self.page_size * len(hashes))
else:
logger.warning(
f"Prefetch operation {operation.request_id} failed to retrieve page {hashes}."
)
# zero copy
def _mooncake_page_get(self, operation, hash_values, host_indices):
key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta(
hash_values,
host_indices,
)
get_result = self.storage_backend.batch_get(
key_strs,
target_location=buffer_ptrs,
target_sizes=buffer_sizes,
)
for i in range(0, len(hashes), batch_size):
page_hashes = hashes[i : i + batch_size]
page_dsts = dsts[i : i + batch_size]
page_data = self.storage_backend.batch_get(page_hashes, page_dsts)
if page_data is None:
if get_result != len(hash_values):
logger.warning(
f"Prefetch operation {operation.request_id} failed or partially failed."
)
if get_result != 0:
operation.increment(get_result * self.page_size)
# non-zero copy
def _generic_page_get(self, operation, hash_values, host_indices):
# todo: zero copy
dummy_page_dst = [self.mem_pool_host.get_dummy_flat_data_page()] * len(
hash_values
)
page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst)
if page_data is None:
return
for i in range(len(hash_values)):
if page_data[i] is None:
logger.warning(
f"Prefetch operation {operation.request_id} failed to retrieve page {page_hashes}."
f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}."
)
break
completed_tokens = operation.completed_tokens
if operation.increment(self.page_size * len(page_hashes)):
for i in range(len(page_hashes)):
completed_tokens += self.page_size
else:
break
self.mem_pool_host.set_from_flat_data_page(
host_indices[operation.completed_tokens],
page_data[i],
)
if not operation.increment(self.page_size):
break # Operation terminated by controller
def _page_transfer(self, operation):
# Select the get function and batch size
if self.is_mooncake_backend():
get_func = self._mooncake_page_get
batch_size = 128
elif self.storage_backend_type == "hf3fs":
if self.mem_pool_host.layout == "page_first":
get_func = self._3fs_zero_copy_page_get
elif self.mem_pool_host.layout == "layer_first":
get_func = self._generic_page_get
batch_size = 128
else:
get_func = self._generic_page_get
batch_size = 8
def generic_page_transfer(self, operation, batch_size=8):
# Transfer batch by batch
for i in range(0, len(operation.hash_value), batch_size):
page_hashes = operation.hash_value[i : i + batch_size]
# todo: zero copy
dummy_page_dst = [
self.mem_pool_host.get_dummy_flat_data_page()
for _ in range(len(page_hashes))
batch_hashes = operation.hash_value[i : i + batch_size]
batch_host_indices = operation.host_indices[
i * self.page_size : (i + len(batch_hashes)) * self.page_size
]
page_data = self.storage_backend.batch_get(page_hashes, dummy_page_dst)
if page_data is None:
logger.warning(
f"Prefetch operation {operation.request_id} failed to retrieve page {page_hashes}."
)
break
completed_tokens = operation.completed_tokens
if operation.increment(self.page_size * len(page_hashes)):
for i in range(len(page_hashes)):
self.mem_pool_host.set_from_flat_data_page(
operation.host_indices[completed_tokens],
page_data[i],
)
completed_tokens += self.page_size
else:
break
def mooncake_page_transfer(self, operation):
key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta(
operation.hash_value, operation.host_indices
)
self.storage_backend.batch_get(key_strs, buffer_ptrs, buffer_sizes)
operation.increment(len(operation.hash_value) * self.page_size)
prev_completed_tokens = operation.completed_tokens
# Get one batch token, and update the completed_tokens if succeed
get_func(operation, batch_hashes, batch_host_indices)
# Check termination
if (
operation.completed_tokens
!= prev_completed_tokens + len(batch_hashes) * self.page_size
):
break # Some operations fail or operation terminated by controller
# release pre-allocated memory
self.mem_pool_host.free(operation.host_indices[operation.completed_tokens :])
def is_mooncake_backend(self):
return self.storage_backend_type == "mooncake"
......@@ -632,15 +666,7 @@ class HiCacheController:
while not self.stop_event.is_set():
try:
operation = self.prefetch_buffer.get(block=True, timeout=1)
if self.is_mooncake_backend():
self.mooncake_page_transfer(operation)
elif self.storage_backend_type == "hf3fs":
if self.mem_pool_host.layout == "page_first":
self.zerocopy_page_transfer(operation, batch_size=128)
elif self.mem_pool_host.layout == "layer_first":
self.generic_page_transfer(operation, batch_size=128)
else:
self.generic_page_transfer(operation)
self._page_transfer(operation)
if self.tp_world_size > 1:
# to ensure all TP workers release the host memory at the same time
......@@ -662,6 +688,27 @@ class HiCacheController:
# todo: more sophisticated rate limiting based on storage backend performance
return True
def _generic_storage_hit_query(self, operation) -> tuple[list[str], int]:
last_hash = operation.last_hash
tokens_to_fetch = operation.token_ids
storage_query_count = 0
remaining_tokens = len(tokens_to_fetch)
hash_value = []
while remaining_tokens >= self.page_size:
last_hash = self.get_hash_str(
tokens_to_fetch[
storage_query_count : storage_query_count + self.page_size
],
last_hash,
)
hash_value.append(last_hash)
storage_query_count += self.page_size
remaining_tokens -= self.page_size
# deferring to batch exists
hit_page_num = self.storage_backend.batch_exists(hash_value)
return hash_value[:hit_page_num], hit_page_num * self.page_size
def prefetch_thread_func(self):
"""
Manage prefetching operations from storage backend to host memory.
......@@ -675,38 +722,12 @@ class HiCacheController:
if operation is None:
continue
storage_hit_count = 0
if (
operation.host_indices is not None
) and self.prefetch_rate_limit_check():
last_hash = operation.last_hash
tokens_to_fetch = operation.token_ids
remaining_tokens = len(tokens_to_fetch)
hash_value = []
while remaining_tokens >= self.page_size:
last_hash = self.get_hash_str(
tokens_to_fetch[
storage_hit_count : storage_hit_count + self.page_size
],
last_hash,
)
# todo, more unified interface
if not self.is_mooncake_backend():
if not self.storage_backend.exists(last_hash):
break
hash_value.append(last_hash)
storage_hit_count += self.page_size
remaining_tokens -= self.page_size
if self.is_mooncake_backend():
# deferring to batch exists for mooncake store
exist_result = self.storage_backend.exists(hash_value)
storage_hit_count = (
sum(1 for v in exist_result.values() if v != 0)
* self.page_size
)
hash_value, storage_hit_count = self._generic_storage_hit_query(
operation
)
if self.tp_world_size > 1:
storage_hit_count_tensor = torch.tensor(
......@@ -755,59 +776,64 @@ class HiCacheController:
self.backup_queue.put(operation)
return operation.id
def zerocopy_page_backup(self, operation, batch_size=8):
hashes, dsts = self.mem_pool_host.get_buffer_with_hash(
operation.hash_value, operation.host_indices
# non-zero copy
def _generic_page_set(self, hash_values, host_indices) -> bool:
data = [
self.mem_pool_host.get_flat_data_page(host_indices[i * self.page_size])
for i in range(len(hash_values))
]
return self.storage_backend.batch_set(hash_values, data)
# zero copy
def _mooncake_page_set(self, hash_values, host_indices) -> bool:
key_strs, buffer_ptrs, buffer_sizes = self.mem_pool_host.get_buffer_meta(
hash_values,
host_indices,
)
for i in range(0, len(hashes), batch_size):
page_hashes = hashes[i : i + batch_size]
page_data = dsts[i : i + batch_size]
success = self.storage_backend.batch_set(page_hashes, page_data)
if not success:
logger.warning(f"Failed to write page {page_hashes} to storage.")
break
operation.completed_tokens += self.page_size * len(page_hashes)
success = self.storage_backend.batch_set(
key_strs,
target_location=buffer_ptrs,
target_sizes=buffer_sizes,
)
return success
def generic_page_backup(self, operation, batch_size=8):
# zero copy
def _3fs_zero_copy_page_set(self, hash_values, host_indices) -> bool:
hashes, dsts = self.mem_pool_host.get_buffer_with_hash(
hash_values, host_indices
)
return self.storage_backend.batch_set(hashes, dsts)
# Backup batch by batch
def _page_backup(self, operation):
# Select the set function and batch size
if self.is_mooncake_backend():
backup_set_func = self._mooncake_page_set
batch_size = 128
elif self.storage_backend_type == "hf3fs":
if self.mem_pool_host.layout == "page_first":
backup_set_func = self._3fs_zero_copy_page_set
elif self.mem_pool_host.layout == "layer_first":
backup_set_func = self._generic_page_set
batch_size = 128
else:
backup_set_func = self._generic_page_set
batch_size = 8
# Backup batch by batch
for i in range(0, len(operation.hash_value), batch_size):
page_hashes = operation.hash_value[i : i + batch_size]
page_data = [
self.mem_pool_host.get_flat_data_page(
operation.host_indices[j * self.page_size]
)
for j in range(i, i + len(page_hashes))
batch_hashes = operation.hash_value[i : i + batch_size]
batch_host_indices = operation.host_indices[
i * self.page_size : (i + len(batch_hashes)) * self.page_size
]
success = self.storage_backend.batch_set(page_hashes, page_data)
# Set one batch token, and record if success.
# todo: allow partial success
success = backup_set_func(batch_hashes, batch_host_indices)
if not success:
logger.warning(f"Failed to write page {page_hashes} to storage.")
break
operation.completed_tokens += self.page_size * len(page_hashes)
def mooncake_page_backup(self, operation):
if len(operation.hash_value):
exist_hashvalues = self.storage_backend.exists(operation.hash_value)
indices = operation.host_indices.tolist()
non_exist_keys = []
non_exist_indices = []
for i in range(len(operation.hash_value)):
if not exist_hashvalues[operation.hash_value[i]]:
non_exist_keys.append(operation.hash_value[i])
non_exist_indices.extend(
indices[i * self.page_size : (i + 1) * self.page_size]
)
if len(non_exist_keys) > 0:
key_strs, buffer_ptrs, buffer_sizes = (
self.mem_pool_host.get_buffer_meta(
non_exist_keys, non_exist_indices
)
)
# TODO: check the return value of batch set to see how many tokens are set successfully
self.storage_backend.batch_set(
key_strs,
target_location=buffer_ptrs,
target_sizes=buffer_sizes,
logger.warning(
f"Write page to storage: {len(batch_hashes)} pages failed."
)
operation.completed_tokens += len(operation.hash_value) * self.page_size
break
operation.completed_tokens += self.page_size * len(batch_hashes)
def backup_thread_func(self):
"""
......@@ -820,15 +846,7 @@ class HiCacheController:
continue
if not self.backup_skip:
if self.is_mooncake_backend():
self.mooncake_page_backup(operation)
elif self.storage_backend_type == "hf3fs":
if self.mem_pool_host.layout == "page_first":
self.zerocopy_page_backup(operation, batch_size=128)
elif self.mem_pool_host.layout == "layer_first":
self.generic_page_backup(operation, batch_size=128)
else:
self.generic_page_backup(operation)
self._page_backup(operation)
min_completed_tokens = operation.completed_tokens
else:
min_completed_tokens = len(operation.token_ids)
......
python/sglang/srt/mem_cache/hicache_storage.py
View file @ 80dc76e1
......@@ -60,7 +60,7 @@ class HiCacheStorage(ABC):
keys: List[str],
target_locations: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> List[torch.Tensor | None]:
) -> List[torch.Tensor | None] | int:
"""
Retrieve values for multiple keys.
Returns a list of tensors or None for each key.
......@@ -96,17 +96,28 @@ class HiCacheStorage(ABC):
pass
@abstractmethod
def exists(self, key: str) -> bool | dict:
def exists(self, key: str) -> bool:
"""
Check if the key exists in the storage.
Returns True if the key exists, False otherwise.
"""
pass
def batch_exists(self, keys: List[str]) -> int:
"""
Check if the keys exist in the storage.
return the number of consecutive existing keys from the start.
Can be overridden by subclasses for more efficient implementation.
"""
for i in range(len(keys)):
if not self.exists(keys[i]):
return i
return len(keys)
class HiCacheFile(HiCacheStorage):
def __init__(self, file_path: str = "/tmp/hicache", is_mla: bool = False):
def __init__(self, file_path: str = "/tmp/hicache", is_mla_backend: bool = False):
self.file_path = os.getenv("SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR", file_path)
if is_dp_attention_enabled():
tp_rank = get_attention_tp_rank()
......@@ -115,7 +126,9 @@ class HiCacheFile(HiCacheStorage):
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()
self.tp_suffix = f"_{tp_rank}_{tp_size}" if tp_size > 1 and not is_mla else ""
self.tp_suffix = (
f"_{tp_rank}_{tp_size}" if tp_size > 1 and not is_mla_backend else ""
)
if not os.path.exists(self.file_path) and tp_rank == 0:
os.makedirs(self.file_path)
logger.info(f"Created HiCacheFile storage directory at {self.file_path}")
......
python/sglang/srt/mem_cache/memory_pool_host.py
View file @ 80dc76e1
......@@ -465,6 +465,7 @@ class MHATokenToKVPoolHost(HostKVCache):
raise ValueError(f"Unsupported layout: {self.layout}")
def get_buffer_meta(self, keys, indices):
local_rank = get_tensor_model_parallel_rank()
ptr_list = []
key_list = []
kv_buffer_data_ptr = self.kv_buffer.data_ptr()
......@@ -488,8 +489,8 @@ class MHATokenToKVPoolHost(HostKVCache):
ptr_list.append(k_ptr)
ptr_list.append(v_ptr)
key_ = keys[index // self.page_size]
key_list.append(f"{key_}_{get_tensor_model_parallel_rank()}_k")
key_list.append(f"{key_}_{get_tensor_model_parallel_rank()}_v")
key_list.append(f"{key_}_{local_rank}_k")
key_list.append(f"{key_}_{local_rank}_v")
element_size = (
self.layer_num
* self.dtype.itemsize
......@@ -704,6 +705,7 @@ class MLATokenToKVPoolHost(HostKVCache):
raise ValueError(f"Unsupported layout: {self.layout}")
def get_buffer_meta(self, keys, indices):
local_rank = get_tensor_model_parallel_rank()
ptr_list = []
key_list = []
kv_buffer_data_ptr = self.kv_buffer.data_ptr()
......@@ -717,7 +719,7 @@ class MLATokenToKVPoolHost(HostKVCache):
)
ptr_list.append(k_ptr)
key_ = keys[index // self.page_size]
key_list.append(f"{key_}_k")
key_list.append(f"{key_}_{local_rank}_k")
element_size = (
self.layer_num
* self.dtype.itemsize
......
python/sglang/srt/mem_cache/storage/mooncake_store/README.md
View file @ 80dc76e1
......@@ -55,12 +55,11 @@ Launch Mooncake meta server:
python -m mooncake.http_metadata_server
```
Start the SGLang server with Mooncake enabled. Mooncake configuration can be provided via environment variables:
Start the SGLang server with Mooncake enabled. Mooncake configuration can be provided via environment variables. Note that, for optimal performance, the Mooncake backend currently supports only the `page_first` layout.
```bash
MOONCAKE_TE_META_DATA_SERVER="http://127.0.0.1:8080/metadata" \
MOONCAKE_GLOBAL_SEGMENT_SIZE=4294967296 \
MOONCAKE_LOCAL_BUFFER_SIZE=134217728 \
MOONCAKE_PROTOCOL="rdma" \
MOONCAKE_DEVICE="erdma_0,erdma_1" \
MOONCAKE_MASTER=127.0.0.1:50051 \
......
python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py
View file @ 80dc76e1
......@@ -13,21 +13,11 @@ from sglang.srt.distributed import get_tensor_model_parallel_rank
from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
DEFAULT_LOCAL_BUFFER_SIZE = 128 * 1024 * 1024 # 128 MB
DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB
logger = logging.getLogger(__name__)
def get_hash_str_mooncake(token_ids: List[int], prior_hash: str = None):
prefix_str = ""
if prior_hash:
prefix_str = hashlib.sha256(prior_hash.encode()).hexdigest()
current_token_ids_bytes = np.array(token_ids).tobytes()
current_hash_object = hashlib.sha256(current_token_ids_bytes)
current_hash_hex = current_hash_object.hexdigest()
return f"{prefix_str}_{int(current_hash_hex[:16], 16)}"
@dataclass
class MooncakeStoreConfig:
local_hostname: str
......@@ -54,9 +44,8 @@ class MooncakeStoreConfig:
global_segment_size=config.get(
"global_segment_size", DEFAULT_GLOBAL_SEGMENT_SIZE
),
local_buffer_size=config.get(
"local_buffer_size", DEFAULT_LOCAL_BUFFER_SIZE
),
# Zero copy interface does not need local buffer
local_buffer_size=DEFAULT_LOCAL_BUFFER_SIZE,
protocol=config.get("protocol", "tcp"),
device_name=config.get("device_name", "auto"),
master_server_address=config.get("master_server_address"),
......@@ -79,9 +68,8 @@ class MooncakeStoreConfig:
global_segment_size=int(
os.getenv("MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_SIZE)
),
local_buffer_size=int(
os.getenv("MOONCAKE_LOCAL_BUFFER_SIZE", DEFAULT_LOCAL_BUFFER_SIZE)
),
# Zero copy interface does not need local buffer
local_buffer_size=DEFAULT_LOCAL_BUFFER_SIZE,
protocol=os.getenv("MOONCAKE_PROTOCOL", "tcp"),
device_name=os.getenv("MOONCAKE_DEVICE", "auto"),
master_server_address=os.getenv("MOONCAKE_MASTER"),
......@@ -96,7 +84,15 @@ class MooncakeStoreConfig:
class MooncakeStore(HiCacheStorage):
def __init__(self, is_mla: bool = False):
def __init__(self, is_mla_backend: bool = False):
"""
Initialize MooncakeStore.
Args:
is_mla_backend: If the backend is MLA
"""
self.is_mla_backend = is_mla_backend
try:
from mooncake.store import MooncakeDistributedStore
except ImportError as e:
......@@ -126,7 +122,6 @@ class MooncakeStore(HiCacheStorage):
logger.info("Connect to Mooncake store successfully.")
self.warmup()
logger.info("Mooncake store warmup successfully.")
self.is_mla = is_mla
except ValueError as e:
logger.error("Configuration loading failed: %s", e)
......@@ -135,14 +130,14 @@ class MooncakeStore(HiCacheStorage):
logger.error("An error occurred while loading the configuration: %s", exc)
raise
self.local_rank = get_tensor_model_parallel_rank()
def warmup(self):
warmup_key = "sglang_mooncake_store_warmup_key" + uuid.uuid4().hex
# 10 MB
warmup_value = bytes(10 * 1024 * 1024)
self.store.put(warmup_key, warmup_value)
warmup_value = bytes(4 * 1024) # 4 KB
assert self.store.put(warmup_key, warmup_value) == 0
assert self.store.is_exist(warmup_key) == 1
self.store.get(warmup_key)
self.store.remove(warmup_key)
assert self.store.get(warmup_key) == warmup_value
def register_buffer(self, buffer: torch.Tensor) -> None:
try:
......@@ -162,78 +157,95 @@ class MooncakeStore(HiCacheStorage):
target_location: Optional[List[int]] = None,
target_sizes: Optional[List[int]] = None,
) -> bool:
assert len(key) == len(target_location) == len(target_sizes)
if len(key) == 0:
return
for i in range(len(key)):
if key[i] is None or target_location[i] is None or target_sizes[i] is None:
return
self._put_batch_zero_copy_impl(key, target_location, target_sizes)
return self.batch_set([key], [value], [target_location], [target_sizes])
def batch_set(
self,
keys: List[str],
value: Optional[Any] = None,
target_location: Optional[List[int]] = None,
target_sizes: Optional[List[int]] = None,
) -> bool:
assert len(keys) == len(target_location) == len(target_sizes)
if len(keys) == 0:
return
return False
for i in range(len(keys)):
if keys[i] is None or target_location[i] is None or target_sizes[i] is None:
return
return False
self._put_batch_zero_copy_impl(keys, target_location, target_sizes)
exist_result = self._batch_exist(keys)
set_keys = []
set_target_locations = []
set_target_sizes = []
set_indices = []
for i in range(len(keys)):
if exist_result[i] != 1:
set_keys.append(keys[i])
set_target_locations.append(target_location[i])
set_target_sizes.append(target_sizes[i])
set_indices.append(i)
# Only set non-existing keys to storage
put_result = self._put_batch_zero_copy_impl(
set_keys, set_target_locations, set_target_sizes
)
for i in range(len(set_indices)):
if put_result[i] == 0:
exist_result[set_indices[i]] = 1
success_count = 0
for i in range(len(keys)):
if exist_result[i] == 0:
break
success_count += 1
# TODO: return the number of consecutive successful operations from the start.
return success_count == len(keys)
def get(
self,
key,
target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> torch.Tensor | None:
assert len(key) == len(target_location) == len(target_sizes)
if len(key) == 0:
return
for i in range(len(key)):
if key[i] is None or target_location[i] is None or target_sizes[i] is None:
return
return self._get_batch_zero_copy_impl(key, target_location, target_sizes)
) -> bool:
return self.batch_get([key], [target_location], [target_sizes]) == 1
def batch_get(
self,
keys: List[str],
target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None,
) -> torch.Tensor | None:
) -> int:
assert len(keys) == len(target_location) == len(target_sizes)
if len(keys) == 0:
return
return 0
get_result = self._get_batch_zero_copy_impl(keys, target_location, target_sizes)
if self.is_mla_backend:
key_multiplier = 1
else:
key_multiplier = 2
for i in range(len(keys)):
if keys[i] is None or target_location[i] is None or target_sizes[i] is None:
return
return self._get_batch_zero_copy_impl(keys, target_location, target_sizes)
def exists(self, keys) -> bool | dict:
_keys = []
local_rank = get_tensor_model_parallel_rank()
for key in keys:
if key is None:
return None
if self.is_mla:
_keys.append(f"{key}_k")
else:
_keys.append(f"{key}_{local_rank}_k")
result = {k: v for k, v in zip(keys, self.store.batch_is_exist(_keys))}
return result
if get_result[i] < 0:
return i // key_multiplier
return len(keys) // key_multiplier
def exists(self, key) -> bool:
return self.batch_exists([key]) > 0
def batch_exists(self, keys) -> int:
if self.is_mla_backend:
query_keys = [f"{key}_k" for key in keys]
key_multiplier = 1
else:
query_keys = []
for key in keys:
query_keys.append(f"{key}_{self.local_rank}_k")
query_keys.append(f"{key}_{self.local_rank}_v")
key_multiplier = 2
exist_result = self._batch_exist(query_keys)
for i in range(len(query_keys)):
if exist_result[i] != 1:
return i // key_multiplier
return len(query_keys) // key_multiplier
def delete(self, key) -> None:
raise (NotImplementedError)
......@@ -248,18 +260,13 @@ class MooncakeStore(HiCacheStorage):
def _put_batch_zero_copy_impl(
self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int]
) -> None:
try:
self.store.batch_put_from(key_strs, buffer_ptrs, buffer_sizes)
except TypeError as err:
logger.error("Failed to put value to Mooncake Store: %s", err)
raise TypeError("Mooncake Store Put Type Error.") from err
) -> List[int]:
return self.store.batch_put_from(key_strs, buffer_ptrs, buffer_sizes)
def _get_batch_zero_copy_impl(
self, key_strs: List[str], buffer_ptrs: List[int], buffer_sizes: List[int]
) -> None:
try:
self.store.batch_get_into(key_strs, buffer_ptrs, buffer_sizes)
except TypeError as err:
logger.error("Failed to get value from Mooncake Store: %s", err)
raise TypeError("Mooncake Store Get Type Error.") from err
) -> List[int]:
return self.store.batch_get_into(key_strs, buffer_ptrs, buffer_sizes)
def _batch_exist(self, key_strs: List[str]) -> List[int]:
return self.store.batch_is_exist(key_strs)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment