CacheDesc

函数功能

构造CacheDesc。

函数原型

CacheDesc(num_tensors: int, shape: Union[Tuple[int], List[int]], data_type: DataType, placement: Placement, batch_dim_index: int=0, seq_len_dim_index: int=-1, kv_tensor_format: str=None)

参数说明

参数名称	数据类型	取值说明
num_tensors	int	cache中tensor的个数，操作cache时，所有tensor会做同样的操作。
shape	Union[Tuple[int], List[int]]	tensor的shape。
data_type	DataType	tensor的data type。
placement	Placement	表示cache所在的设备类型。默认值Placement.DEVICE。
batch_dim_index	int	表示shape中batch size所在维度。默认值0，表示在第0维。
seq_len_dim_index	int	表示shape中seq_len所在维度。默认值-1，表示未配置。
kv_tensor_format	str	表示cache的format。默认不配置，在check_capacity时使用。

调用示例

from llm_datadist import CacheDesc
cache_desc = CacheDesc(80, [4, 2048, 1, 128], DataType.DT_FLOAT16)

返回值

正常场景下无返回值。

传入数据类型错误情况下会抛出TypeError异常。

约束说明

无