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异常。
约束说明
无
