注册用户内存

功能介绍

将用户申请的内存注册给llm_datadist，注册后，不同节点之间注册的用户内存可以实现通信访问。

涉及的接口：

• cache_manager.register_cache：将用户的连续cache内存地址注册给llm_datadist。
• cache_manager.register_blocks_cache：将用户的block的cache内存地址注册给llm_datadist
• llm_datadist.query_register_mem_status：查询内存注册状态。内存实际注册是发生在link阶段，作为异步操作在注册，所以需要提供查询接口，查询内存注册是否实际完成。

使用场景

用户不想用llm_datadist分配和管理kv cache的内存，想直接使用框架中已分配好的kv cache内存。

常用于PA场景下将框架中已申请的kv cache内存注册给llm_datadist，通过llm_datadist的kv cache传输接口实现不同节点之间的kv cache传输。

功能示例

正常的业务流程，此处以PyTorch框架为例，给出PyTorch的tensor注册给llm_datadist示例。接口参数可参考《LLM DataDist接口参考（Python）》。

1. 在llm_datadist初始化时使能cache_manager的option

   1 2 3 4 5 6

   # 初始llm_datadist llm_datadist = LLMDataDist(role, cluster_id) llm_config = LLMConfig() llm_config.device_id = device_id llm_config.enable_cache_manager =True# 需要使能cache_manager的option llm_datadist.init(options)

2. 将框架中已经申请好的kv cache的内存注册给llm_datadist。

   1 2 3 4 5 6 7 8

   # 将torch的kv tensor内存注册给llm_datadist cache_manager = llm_datadist.cache_manager cache_desc = CacheDesc(num_tensors=1, shape=[2, 1024*1024], data_type=DataType.DT_FLOAT) tensor = torch.ones(2, 1024*1024, dtype=torch.float).npu() addr =int(tensor.data_ptr()) # block_cache_key作为可选参数，在数据接收端注册的时候可以不用配置。根据实际业务自行选择register_blocks_cache或者register_cache cache = cache_manager.register_blocks_cache(cache_desc, [addr], BlocksCacheKey(1, 0)) # cache = cache_manager.register_cache(cache_desc, [addr])

3. 参考集群建链示例进行建链。
4. 查询内存注册状态，等待内存在不同节点间注册完成。

   1 2 3 4

   while True: ret = llm_datadist.query_register_mem_status("建链成功返回的comm_id") if ret == RegisterMemStatus.OK: break

5. P侧模型执行推理，D侧模型可以pull kv，再执行增量推理。
6. 参考集群建链示例进行断链以及llm_datadist资源释放。

异常处理

当内存注册失败或者超时时，就不能继续后续流程。需要检查用户内存的有效性、是否有重叠，确认内存地址没问题后重试。更多异常处理请参考《LLM DataDist接口参考（Python）》中的"LLMStatusCode"章节。