register_cache

函数功能

非PagedAttention场景下，调用此接口注册一个自行申请的内存。

函数原型

register_cache(cache_desc: CacheDesc, addrs: List[int], cache_keys: Union[Tuple[CacheKey], List[CacheKey]] = (), remote_accessible: Optional[bool] = None) -> Cache:

参数说明

参数名称	数据类型	取值说明
cache_desc	CacheDesc	Cache的描述信息。
addrs	List[int]	Cache的地址。register_cache中的地址个数与register_blocks_cache的地址个数之和不超过256。
cache_keys	Union[Tuple[CacheKey], List[CacheKey]]	Cache的索引。
remote_accessible	Optional[bool]	指定当前注册内存是否与对端进行内存描述符交换，预留参数，暂不生效。

调用示例

from llm_datadist import *
import torch
...
cache_mgr = llm_datadist.cache_manager
data_addrs = []
for i in range(80):
    tensor = torch.ones(4, 1024 * 1024, dtype=torch.float32).npu()
    tensor_addr = tensor.data_ptr()
    data_addrs.append(tensor_addr)
cache_desc = CacheDesc(80, [4, 1024 * 1024], DataType.DT_FLOAT, Placement.DEVICE)
cache = cache_mgr.register_cache(cache_desc, data_addrs)
print('register cache suc:', cache.cache_id, flush=True)

返回值

正常情况下返回注册的cache。

传入数据类型错误，src和dst不匹配情况下会抛出TypeError或ValueError异常。

传入参数为None，会抛出AttributeError异常。

约束说明

Device内存需要先注册再进行link，host内存不约束顺序。

注册内存地址需自行保证不重复。

如果通过HCCS进行数据传输，且内存不是通过aclrtMalloc申请，则地址需要按照页大小对齐(如果确定页大小是多少，推荐2MB对齐)，否则可能导致link失败。

父主题： CacheManager