CacheDesc

Function Description

Constructs a CacheDesc object, which is usually used as the parameter type in allocate_cache of CacheManager.

Prototype

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

Parameters

Parameter	Data Type	Description
num_tensors	int	Number of tensors in the cache. When the cache is operated, all tensors perform the same operation.
shape	Union[Tuple[int], List[int]]	Tensor shape.
data_type	DataType	Data type of a tensor.
placement	Placement	Type of the device where the cache is located. The default value is Placement.DEVICE.
batch_dim_index	int	Dimension of the batch size in the shape. The default value is 0, indicating dimension 0.
seq_len_dim_index	int	Dimension of seq_len in the shape. The default value is -1, indicating that the dimension is not configured.
kv_tensor_format	str	Cache format. By default, it is not configured.

Example

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

Returns

In normal cases, a CacheDesc instance is returned.

If the input data type is incorrect, the TypeError or ValueError exception is reported.

Restrictions

None

Parent topic: API Reference