from dataclasses import dataclass
from threading import Thread
from typing import Dict, List, Optional, Tuple, Union
from llm_datadist_v1 import llm_wrapper
from .data_type import DataType, python_dtype_2_dwrapper_dtype
from .llm_types import (
BlocksCacheKey,
CacheDesc,
CacheKey,
CacheKeyByIdAndIndex,
CacheTask,
KvCache,
LayerSynchronizer,
MemInfo,
Placement,
PushType,
TransferConfig,
TransferWithCacheKeyConfig,
)
from .status import (
LLMException,
LLMStatusCode,
code_2_status,
raise_if_false,
raise_if_true,
)
from .utils import log
from .utils.utils import (
check_isinstance,
check_list_uint64,
check_type,
check_uint32,
check_uint64,
)
_INVALID_ID = 2**64 - 1
_MAX_DISPLAYED_REQ_ID_COUNT = 8
_NUM_TENSORS_PER_LAYER = 2
def clone_cache_desc(cache_desc: CacheDesc) -> CacheDesc:
return CacheDesc(
cache_desc.num_tensors,
cache_desc.shape[:],
cache_desc.data_type,
cache_desc.placement,
cache_desc.batch_dim,
cache_desc.seq_len_dim_index,
cache_desc.kv_tensor_format,
)
def verify_cache_shape(shape: List[int]) -> None:
raise_if_false(len(shape) > 0, "scalar is not supported")
raise_if_false(0 not in shape, f"empty dimension is not supported, shape = {shape}")
for dim in shape:
if dim < -1:
raise LLMException(f"dim {dim} is not supported, shape = {shape}")
def to_data_type(data_type_str: str) -> DataType:
for dtype in DataType:
if str(dtype.value) == data_type_str:
return dtype
raise LLMException(
f"unsupported dtype: {data_type_str}",
status_code=LLMStatusCode.LLM_PARAM_INVALID,
)
def is_invalid_id(req_id_or_prefix_id: int) -> bool:
return (req_id_or_prefix_id < 0) or (req_id_or_prefix_id == _INVALID_ID)
def is_valid_id(req_id_or_prefix_id: int) -> bool:
return not is_invalid_id(req_id_or_prefix_id)
def calc_tensor_size(shape: List[int], data_type: DataType) -> int:
ge_data_type = python_dtype_2_dwrapper_dtype.get(data_type)
tensor_size = llm_wrapper.calc_tensor_size(shape, ge_data_type)
raise_if_false(
tensor_size >= 0,
"Failed to calc tensor size, shape = {0}, data_type = {1}",
shape,
data_type,
)
return tensor_size
class CacheDescParser(object):
@staticmethod
def parse_by_options(options: Dict[str, str]) -> CacheDesc:
raise_if_false("llm.RefInputShapes" in options, "llm.RefInputShapes not set")
kv_shape_str = options["llm.RefInputShapes"]
kv_shapes = kv_shape_str.split(";")
raise_if_false(
len(set(kv_shapes)) == 1,
f"zero or multiple different shapes is not supported, llm.RefInputShapes = {kv_shapes}",
)
raise_if_false("llm.RefInputDtypes" in options, "llm.RefInputDtypes not set")
kv_data_type_str = options["llm.RefInputDtypes"]
kv_data_types = kv_data_type_str.split(";")
raise_if_false(
len(set(kv_data_types)) == 1,
f"zero or multiple different data types is not supported, llm.RefInputDtypes = {kv_data_types}",
)
raise_if_false(
len(kv_shapes) == len(kv_data_types),
f"kv_shapes num ({len(kv_shapes)}) mismatches that of kv_data_type ({len(kv_data_types)})",
)
kv_shape = list((int(dim) for dim in kv_shapes[0].split(",")))
verify_cache_shape(kv_shape)
if "llm.RefInputSeqLenDimIndex" in options:
seq_len_dim_index = int(options["llm.RefInputSeqLenDimIndex"])
log.info(f"get seq_len_dim_index from option, value = {seq_len_dim_index}")
raise_if_false(
0 <= seq_len_dim_index < len(kv_shape),
f"seq_len_dim_index ({seq_len_dim_index}) out of range, kv_shape = {kv_shape}",
)
elif -1 in kv_shape[1:]:
raise_if_false(
kv_shape[1:].count(-1) == 1,
f"can only have one dynamic dim apart from batch_dim(0), "
f"but got {kv_shape[1:].count(-1)}, shape = {kv_shape}",
)
seq_len_dim_index = kv_shape[1:].index(-1) + 1
log.info(f"seq_len_dim_index inferred by shape, value = {seq_len_dim_index}")
else:
seq_len_dim_index = -1
log.info(
f"llm.RefInputSeqLenDimIndex not set and can not infer by shape,seq_len_dim_index = {seq_len_dim_index}"
)
kv_tensor_format = options.get("llm.kvTensorFormat", None)
kv_data_type = to_data_type(kv_data_types[0])
cache_desc = CacheDesc(
len(kv_shapes),
kv_shape,
kv_data_type,
Placement.DEVICE,
0,
seq_len_dim_index,
kv_tensor_format,
)
log.info(f"parse cache_desc from option, value = {cache_desc}")
return cache_desc
def build_npu_tensors(kv_cache: KvCache, device_index: int = 0) -> List[int]:
return llm_wrapper.build_npu_tensors(
kv_cache.cache_desc.shape,
kv_cache.cache_desc.data_type.value,
kv_cache.cache_desc.size,
kv_cache.per_device_tensor_addrs[device_index],
)
def pack_cache_desc(cache_desc: CacheDesc) -> Tuple[int, int, int, List[int], int, int]:
return (
cache_desc.num_tensors,
cache_desc.data_type.value,
cache_desc.seq_len_dim_index,
cache_desc.shape,
cache_desc.placement.value,
cache_desc._is_blocks,
)
def pack_cache_key(cache_key: CacheKey) -> Tuple[int, int, int, int, int, int, bool]:
return (
cache_key.prompt_cluster_id,
-1,
0,
cache_key.req_id,
cache_key.prefix_id,
cache_key.model_id,
False,
)
def pack_cache_key_by_id(
cache_key: CacheKeyByIdAndIndex,
) -> Tuple[int, int, int, int, int, int, bool]:
return (
cache_key.cluster_id,
cache_key.cache_id,
cache_key.batch_index,
_INVALID_ID,
_INVALID_ID,
0,
False,
)
def pack_block_cache_key(
cache_key: BlocksCacheKey,
) -> Tuple[int, int, int, int, int, int, bool]:
return (
cache_key.prompt_cluster_id,
-1,
0,
_INVALID_ID,
_INVALID_ID,
cache_key.model_id,
True,
)
def pack_mem_info(mem_info: MemInfo) -> Tuple[int, int, int]:
return (mem_info.mem_type.value, mem_info.addr, mem_info.size)
@dataclass
class TransferCacheParameters:
src_cache: "Union[Cache, KvCache]"
transfer_configs: Union[
List[Union[TransferConfig, TransferWithCacheKeyConfig]],
Tuple[Union[TransferConfig, TransferWithCacheKeyConfig]],
]
src_block_indices: Optional[List[int]] = None
dst_block_indices: Optional[List[int]] = None
dst_block_memory_size: Optional[int] = None
class TransferCacheJob:
task_id = 0
def __init__(
self,
params: TransferCacheParameters,
layer_synchronizer: LayerSynchronizer,
transfer_cache_func,
):
self._transfer_configs = params.transfer_configs
self._src_block_indices = params.src_block_indices
self._dst_block_indices = params.dst_block_indices
self._dst_block_memory_size = params.dst_block_memory_size
self._cache_id = params.src_cache.cache_id
self._cache_desc = params.src_cache.cache_desc
self._num_layers = 0
self._rets: Dict[int, LLMStatusCode] = {}
self._timeout_in_millis: Optional[int] = None
self._layer_synchronizer = layer_synchronizer
self._transfer_cache_func = transfer_cache_func
def init(self):
raise_if_false(
self._cache_desc.num_tensors % 2 == 0,
"cache_desc.num_tensors ({0}) is not even",
self._cache_desc.num_tensors,
)
self._num_layers = self._cache_desc.num_tensors // 2
for transfer_config in self._transfer_configs:
if transfer_config.src_layer_range is None:
transfer_config.src_layer_range = range(0, self._num_layers)
self.check_transfer_config(transfer_config)
def transfer_layers(self):
for layer_i, src_layer_index in enumerate(range(self._num_layers)):
to_transfer = [config for config in self._transfer_configs if src_layer_index in config.src_layer_range]
if not to_transfer:
log.info("src_layer %d sends to no destination", src_layer_index)
continue
if not self._layer_synchronizer.synchronize_layer(src_layer_index, self._timeout_in_millis):
log.error(f"Failed to synchronize layer {src_layer_index}")
for config in to_transfer:
self._rets[config.dst_cluster_id] = LLMStatusCode.LLM_PARAM_INVALID
return
for config in to_transfer:
dst_layer_index = (
src_layer_index if config.dst_layer_range is None else list(config.dst_layer_range)[layer_i]
)
ret = self.transfer_layer(src_layer_index, dst_layer_index, config)
if ret != LLMStatusCode.LLM_SUCCESS:
log.error(f"Failed to transfer layer {src_layer_index} to dst_cluster_id={config.dst_cluster_id}")
self._rets[config.dst_cluster_id] = ret
return
log.info(f"transfer layer {src_layer_index} to dst_cluster_id={config.dst_cluster_id} success")
if src_layer_index == config.src_layer_range.stop - 1:
self._rets[config.dst_cluster_id] = LLMStatusCode.LLM_SUCCESS
log.info(f"transfer all layers to dst_cluster_id={config.dst_cluster_id} finished")
log.info("transfer all layers finished")
def transfer_layer(
self,
src_layer_index: int,
dst_layer_idx,
transfer_config: Union[TransferConfig, TransferWithCacheKeyConfig],
) -> LLMStatusCode:
if check_type(transfer_config, TransferConfig):
dst_layer_index = src_layer_index - transfer_config.src_layer_range.start
dst_addrs = transfer_config.dst_addrs[
dst_layer_index * _NUM_TENSORS_PER_LAYER : dst_layer_index * _NUM_TENSORS_PER_LAYER
+ _NUM_TENSORS_PER_LAYER
]
transfer_config = (
self._cache_id,
transfer_config.src_batch_index,
src_layer_index,
dst_addrs,
transfer_config.dst_cluster_id,
0,
0,
PushType.NO_CACHE_KEY.value,
src_layer_index,
2,
)
else:
if check_type(transfer_config.cache_key, BlocksCacheKey):
transfer_config = (
self._cache_id,
0,
src_layer_index,
[],
transfer_config.cache_key.cluster_id,
transfer_config.cache_key.model_id,
0,
PushType.BLOCKS_CACHE_KEY.value,
dst_layer_idx,
2,
)
elif check_type(transfer_config.cache_key, CacheKeyByIdAndIndex):
transfer_config = (
self._cache_id,
transfer_config.src_batch_index,
src_layer_index,
[],
transfer_config.cache_key.cluster_id,
transfer_config.cache_key.cache_id,
transfer_config.cache_key.batch_index,
PushType.CACHE_KEY_BY_ID.value,
dst_layer_idx,
2,
)
block_config = (
self._dst_block_memory_size if self._dst_block_memory_size is not None else 0,
self._src_block_indices if self._src_block_indices is not None else [],
self._dst_block_indices if self._dst_block_indices is not None else [],
)
ret = self._transfer_cache_func(TransferCacheJob.task_id, transfer_config, block_config)
TransferCacheJob.task_id += 1
return code_2_status(ret)
def check_transfer_config(self, transfer_config: Union[TransferConfig, TransferWithCacheKeyConfig]):
if self._src_block_indices:
raise_if_false(
transfer_config.src_batch_index == 0,
"Invalid TransferConfig, src_batch_index ({0}) != 0 while src is blocks",
transfer_config.src_batch_index,
)
raise_if_false(
0 <= transfer_config.src_batch_index < self._cache_desc.batch_size,
"Invalid TransferConfig, src_batch_index ({0}) out of range: [0, {1})",
transfer_config.src_batch_index,
self._cache_desc.batch_size,
)
raise_if_false(
0 <= transfer_config.src_layer_range.start < transfer_config.src_layer_range.stop <= self._num_layers,
"src_layer_range: {0} out of range, src_layer_num = {1}",
transfer_config.src_layer_range,
self._num_layers,
)
num_tensors_to_transfer = (
transfer_config.src_layer_range.stop - transfer_config.src_layer_range.start
) * _NUM_TENSORS_PER_LAYER
if check_type(transfer_config, TransferConfig):
raise_if_false(
len(transfer_config.dst_addrs) == num_tensors_to_transfer,
"expect {0} dst_addrs, but len(dst_addrs) = {1}, range = {2}",
num_tensors_to_transfer,
len(transfer_config.dst_addrs),
transfer_config.src_layer_range,
)
def get_results(self) -> List[LLMStatusCode]:
rets = []
for config in self._transfer_configs:
ret = self._rets.get(config.dst_cluster_id)
rets.append(ret)
return rets
def num_transfer_configs(self):
return len(self._transfer_configs)
class TransferAsyncThread(Thread):
def __init__(self, transfer_job: TransferCacheJob, default_err_code=LLMStatusCode.LLM_TIMEOUT):
super().__init__()
self._transfer_job = transfer_job
self._rets = []
self._default_err_code = default_err_code
def run(self):
self._transfer_job.transfer_layers()
def get_results(self, timeout) -> List[LLMStatusCode]:
self.join(timeout)
if self.is_alive():
return [self._default_err_code] * self._transfer_job.num_transfer_configs()
return self._transfer_job.get_results()
def get(self, timeout) -> LLMStatusCode:
self.join(timeout)
if self.is_alive():
return self._default_err_code
rets = self._transfer_job.get_results()
for ret in rets:
if ret != LLMStatusCode.LLM_SUCCESS and ret is not None:
return ret
return LLMStatusCode.LLM_SUCCESS
def _check_block_indices(arg_name, arg_value):
if arg_value is not None:
check_isinstance(arg_name, arg_value, [list, tuple], int)
check_list_uint64(arg_name, arg_value)
def transfer_cache_async(
params: TransferCacheParameters,
layer_synchronizer: LayerSynchronizer,
transfer_cache_func,
default_error_code=LLMStatusCode.LLM_TIMEOUT,
enable_remote_cache=False,
) -> CacheTask:
_check_block_indices("dst_block_indices", params.dst_block_indices)
_check_block_indices("src_block_indices", params.src_block_indices)
if params.dst_block_memory_size is not None:
check_uint64("dst_block_memory_size", params.dst_block_memory_size)
if params.src_block_indices:
raise_if_false(params.dst_block_indices, "transfer from blocks to cache is not supported")
raise_if_false(
len(params.src_block_indices) == len(params.dst_block_indices),
"num_block_indices mismatches, src_num = {0}, dst_num = {1}",
len(params.src_block_indices),
len(params.dst_block_indices),
)
else:
raise_if_true(
(transfer_cache_func == llm_wrapper.transfer_cache) and (params.dst_block_indices is not None),
"transfer from cache to blocks is not supported",
)
if params.dst_block_indices:
raise_if_false(
params.dst_block_memory_size is not None,
"dst_block_memory_size must be set when transfer from cache to blocks",
)
check_isinstance("layer_synchronizer", layer_synchronizer, LayerSynchronizer, allow_none=False)
if not enable_remote_cache:
check_isinstance(
"transfer_configs",
params.transfer_configs,
[list, tuple],
TransferConfig,
allow_none=False,
)
else:
check_isinstance(
"transfer_configs",
params.transfer_configs,
[list, tuple],
TransferWithCacheKeyConfig,
"While enable_remote_cache_accessible is True, ",
allow_none=False,
)
raise_if_false(
params.dst_block_indices or params.dst_block_memory_size in (None, 0),
"dst_block_memory_size ({0}) is neither None nor 0 while dst is not blocks",
params.dst_block_memory_size,
)
transfer_job = TransferCacheJob(params, layer_synchronizer, transfer_cache_func)
transfer_job.init()
transfer_thread = TransferAsyncThread(transfer_job, default_error_code)
transfer_thread.start()
log.info("[transfer_cache_async] async task start")
cache_task = CacheTask(transfer_thread)
return cache_task
def layer_range_to_tensor_indices(src_layer_range: range, dst_layer_range: range, tensor_num_per_layer: int = 2):
check_isinstance("src_layer_range", src_layer_range, range)
check_isinstance("dst_layer_range", dst_layer_range, range)
raise_if_true(
(src_layer_range is not None) and (src_layer_range.step != 1),
"param check failed, src_layer_range step must be 1.",
)
raise_if_true(
(dst_layer_range is not None) and (dst_layer_range.step != 1),
"param check failed, dst_layer_range step must be 1.",
)
src_layer_indices = [] if src_layer_range is None else list(src_layer_range)
dst_layer_indices = [] if dst_layer_range is None else list(dst_layer_range)
one_layer_tensor_num = tensor_num_per_layer
src_tensor_indices = []
if len(src_layer_indices) != 0:
raise_if_true(
src_layer_indices[0] < 0,
"src_layer_range is invalid, the start value:{0} is < 0",
src_layer_indices[0],
)
check_uint32("src_layer_range", src_layer_indices[0])
check_uint32("src_layer_range", src_layer_indices[-1])
src_tensor_start_index = src_layer_indices[0] * one_layer_tensor_num
for i in range(len(src_layer_indices) * one_layer_tensor_num):
src_tensor_indices.append(src_tensor_start_index + i)
dst_tensor_indices = []
if len(dst_layer_indices) != 0:
raise_if_true(
dst_layer_indices[0] < 0,
"dst_layer_range is invalid, the start value:{0} is < 0",
dst_layer_indices[0],
)
check_uint32("dst_layer_range", dst_layer_indices[0])
check_uint32("dst_layer_range", dst_layer_indices[-1])
dst_tensor_start_index = dst_layer_indices[0] * one_layer_tensor_num
for i in range(len(dst_layer_indices) * one_layer_tensor_num):
dst_tensor_indices.append(dst_tensor_start_index + i)
return src_tensor_indices, dst_tensor_indices
def parse_listen_ip_info(listen_ip_info: str) -> (str, int):
check_isinstance("listen_ip_info", listen_ip_info, [str])
ip_and_port = listen_ip_info.split(":")
raise_if_false(
len(ip_and_port) == 2,
f'llm.listenIpInfo "{ip_and_port}" is invalid, format should be "ip:port"',
)
ip = ip_and_port[0]
port = int(ip_and_port[1])
return ip, port
