register_blocks_cache

产品支持情况

产品	是否支持
Atlas A3 训练系列产品/Atlas A3 推理系列产品	√
Atlas 800I A2 推理产品/A200I A2 Box 异构组件	√
Atlas A2 训练系列产品	x
Atlas 200I/500 A2 推理产品	x
Atlas 推理系列产品	x
Atlas 训练系列产品	x

函数功能

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

函数原型

register_blocks_cache(cache_desc: CacheDesc, addrs: List[int], blocks_cache_key: Optional[BlocksCacheKey] = None, remote_accessible: Optional[bool] = None) -> Cache:

参数说明

参数名称	数据类型	取值说明
cache_desc	CacheDesc	Cache的描述信息。
addrs	List[int]	Cache的地址。register_cache中的地址个数与register_blocks_cache的地址个数之和不超过240。
blocks_cache_key	Optional[BlocksCacheKey]	可选的BlocksCacheKey索引。
remote_accessible	Optional[bool]	指定当前注册内存是否能用来做网络传输，对于device内存，默认值是True，对于host内存，默认值是False。

调用示例

from llm_datadist import *
import torch
...
cache_manager = llm_datadist.cache_manager
cache_desc = CacheDesc(1, [100, 1024 * 1024], DataType.DT_FLOAT)
cache_key = BlocksCacheKey(prompt_cluster_id=cluster_id, model_id=0)
cache = cache_manager.register_blocks_cache(cache_desc, [addr], cache_key)

返回值

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

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

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

约束说明

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

父主题： CacheManager