尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

HIXL llm_datadist Python API 之 Cache 对象:KV Cache 句柄的属性与创建全解析

HIXL llm_datadist Python API 之 Cache 对象:KV Cache 句柄的属性与创建全解析 HIXL llm_datadist Python API 之 Cache 对象KV Cache 句柄的属性与创建全解析【免费下载链接】hixlHIXLHuawei Xfer Library是一个灵活、高效的昇腾单边通信库面向集群场景提供简单、可靠、高效的点对点数据传输能力。项目地址: https://gitcode.com/cann/hixlCache 是 HIXL 项目中 LLM-DataDist 数据分发能力的核心 Python 数据对象它封装了分布在设备Device或主机Host内存上的一个 KV Cache含普通连续 Cache 与 PagedAttention Blocks Cache 两种形态是pull_cache、push_cache、transfer_cache_async等所有缓存传输操作的直接操作对象。本文基于 Cache.md 展开结合仓库内 Python 类型定义与调用样例完整讲解 Cache 的获取方式、三个只读属性cache_id、cache_desc、tensor_addrs与静态创建方法create_cpu_cache的参数、返回值与约束帮助你正确地在多节点推理集群中使用 Cache 句柄完成 KV 传输。Cache 在 LLM-DataDist 数据流中的角色在 LLM-DataDistllm_datadist的典型Prompt 写 KV / Decoder 读 KV场景中Cache 对象贯穿整个生命周期Prompt 侧通过CacheManager.allocate_cache或allocate_blocks_cache分配显存得到 Cache 对象必要时同时用 CacheKey/BlocksCacheKey 建立远端索引Decoder 侧同样分配本地 Cache然后调用pull_cache/pull_blocks将远端 Cache 数据拉到本地或调用push_cache/push_blocks把本地 Cache 推送到远端传输完成后通过CacheManager.deallocate_cache释放或调用unregister_cache解除对自行申请内存的注册。从源码结构看Cache类被设计为用户不应直接构造的句柄对象其定义位于 llm_types.py类注释明确写道Cache由 CacheManager.register_cache/allocate_cache/allocate_blocks_cache 创建用户不应直接构造。唯一的例外是静态方法create_cpu_cache它用于把用户在 Host 侧自行申请的 CPU 内存包装成 Cache 句柄。因此理解 Cache 之前建议先了解它的来源CacheManager见 CacheManager.md通过LLMDataDist.cache_manager属性获取实例生命周期与LLMDataDist初始化周期一致其构造细节见 LLMDataDist.md。产品支持情况Cache 对象在以下昇腾产品形态上受支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 推理系列产品 / Atlas A2 训练系列产品支持但仅限Atlas 800I A2 推理服务器与A200I A2 Box 异构组件两种形态。需要说明的是该支持范围与本 API 文档目录下其他数据类型如 CacheDesc.md、CacheKey.md保持一致属于整个 llm_datadist 模块的通用产品约束。Cache 构造函数一般不需要用户调用文档明确说明构造 Cache 通常不需要用户直接调用。Cache 对象由 CacheManager 的以下接口返回获取途径适用场景相关文档allocate_cache(cache_desc, cache_keys)非 PagedAttention 场景由库自动分配显存并构建普通连续 CacheCacheManager.mdallocate_blocks_cache(cache_desc, blocks_cache_key)PagedAttention 场景按 block 分配的块式 CacheCacheManager.mdregister_cache(cache_desc, addrs, cache_keys, remote_accessible)非 PagedAttention 场景注册用户自行申请的内存CacheManager.mdregister_blocks_cache(cache_desc, addrs, blocks_cache_key, remote_accessible)PagedAttention 场景注册用户自行申请的内存CacheManager.md从实现看allocate_cache在 cache_manager.py 中先做参数校验然后调用底层allocate_cache_v2获得(cache_id, tensor_addrs)最终构造Cache(cache_id, cache_desc, tensor_addrs, self, False, False)返回allocate_blocks_cache同样返回 Cache只是最后一个标记is_blocks_cacheTrue见 cache_manager.py。因此cache_id与tensor_addrs的实际来源都是底层引擎的内存分配结果。由于分配出的 Cache 会同时被cache_id与cache_keys若传入引用只有当这些引用都解除后资源才会真正释放cache_id引用通过deallocate_cache解除CacheKey 引用则通过 Decoder 侧pull_cache成功后自动解除或 Prompt 侧显式调用remove_cache_key解除。实际使用中应遵循谁分配、谁释放的准则见下方完整调用示例。cache_id获取 Cache 的全局标识函数功能获取 Cache 的 id。函数原型property cache_id() - int参数说明无。调用示例... kv_cache cache_manager.allocate_cache(cache_desc, cache_keys) print(kv_cache.cache_id)返回值正常情况返回类型为 Cache 的 idint。约束说明无。从源码看cache_id是 Cache 构造时的第一个参数底层为 64 位整数构造时通过check_int64(cache_id, cache_id)校验见 llm_types.py。它同时是unregister_cache(cache_id)、CacheKeyByIdAndIndex(cache_id...)等接口的关联标识特别地通过create_cpu_cache创建的 Cache其cache_id固定为 -1表示这是一个未纳入 CacheManager 索引的本地句柄。cache_desc获取 Cache 描述函数功能获取 Cache 描述即创建该 Cache 所用的CacheDesc。函数原型property cache_desc() - CacheDesc参数说明无。调用示例... kv_cache cache_manager.allocate_cache(cache_desc, cache_keys) print(kv_cache.cache_desc.num_tensors)返回值正常情况返回类型为 Cache 的 cache 描述CacheDesc实例。约束说明无。CacheDesc是描述 Cache 元信息的载体构造参数包括num_tensorstensor 个数、shapetensor 形状、data_type数据类型、placement设备类型默认Placement.DEVICE、batch_dim_indexbatch 维索引默认 0、seq_len_dim_indexseq_len 维索引默认 -1 表示未配置以及kv_tensor_format可选 format。典型构造如下from llm_datadist import CacheDesc, DataType cache_desc CacheDesc(80, [4, 2048, 1, 128], DataType.DT_FLOAT16)上述示例声明一个包含 80 个 FP16 tensor、形状为[4, 2048, 1, 128]的设备端 Cache。CacheDesc的完整字段说明与校验规则见 CacheDesc.md 及源码 llm_types.py其中size属性会调用llm_wrapper.calc_tensor_size计算单个 tensor 字节数用于传输大小推算。在仓库的 pull_cache 样例中实际创建的 cache_desc 为CacheDesc(num_tensors4, shape[2, 16 * 1024], data_typeDataType.DT_FLOAT16, placementPlacement.DEVICE)见 pull_cache_sample.py可见print(kv_cache.cache_desc.num_tensors)这类属性访问在调试与维测中十分常用。tensor_addrs获取 Cache 的地址列表函数功能获取 Cache 的地址。函数原型property tensor_addrs() - List[int]参数说明无。调用示例... kv_cache cache_manager.allocate_cache(cache_desc, cache_keys) print(kv_cache.tensor_addrs)返回值正常情况返回类型为 Cache 的地址List[int]每个元素对应一个 tensor 的起始地址。约束说明无。需要特别说明的是tensor_addrs是**内存地址整数**而非 Python 对象它由底层引擎在分配时返回长度为cache_desc.num_tensors。对于 Device 端 Cache这些地址是昇腾设备内存地址对于 CPU Cache见下文create_cpu_cache则是用户传入的 CPU 内存地址。该属性可用于与第三方框架如 PyTorch的地址互操作例如构造from_dlpack或as_strided视图但地址有效性需由用户保证。create_cpu_cache创建 CPU Cache函数功能创建一个 CPU cacheHost 侧内存的 Cache 句柄。与CacheManager分配的设备 Cache 不同CPU Cache 用于 D2H / H2D 等涉及主机内存的传输场景。函数原型create_cpu_cache(cache_desc: CacheDesc, addrs: List[int])参数说明参数名称数据类型取值说明cache_descCacheDesccache 的描述。addrsList[int]cpu cache 的地址。调用示例from llm_datadist import Cache cpu_cache Cache.create_cpu_cache(cpu_cache_desc, cpu_addrs) # cpu_addrs来自创建的cpu tensors返回值正常情况返回类型为 Cache 的 cpu_cache传入数据类型错误情况下会抛出TypeError或ValueError异常传入参数为 None会抛出AttributeError异常。约束说明无。从实现看create_cpu_cache是一个classmethod其内部做了三件事见 llm_types.py校验cache_desc必须是CacheDesc实例、addrs必须是元素为 int 的 list强制要求cache_desc.placement Placement.HOST否则抛出 cache_desc placement must be HOST 的校验错误强制要求cache_desc.num_tensors len(addrs)即地址数量必须与 tensor 个数一一对应。满足校验后返回cls(-1, cache_desc, addrs, None, False, True)即cache_id为 -1、cache_manager为 None、is_registered为 False、is_blocks_cache为 True。也就是说CPU Cache 天然按 blocks cache 语义管理且不归属任何 CacheManager因此不能直接用于deallocate_cache/unregister_cache等需要 CacheManager 上下文的接口它通常作为transfer_cache_async等传输接口的源或目标参与数据传输。Placement枚举的定义HOST 0、DEVICE 1见 llm_types.py完整说明见 Placement.md。结合样例的完整用法从分配到传输再到释放以下流程取自仓库真实样例 pull_cache_sample.py完整演示了 Cache 句柄在Prompt 写、Decoder 拉场景下的用法Decoder 侧拉取 KVcache_manager datadist.cache_manager cache_desc CacheDesc( num_tensors4, shape[2, 16 * 1024], data_typeDataType.DT_FLOAT16, placementPlacement.DEVICE, ) cache cache_manager.allocate_cache(cache_desc) dist.barrier() # cache ready cache_key_0 CacheKey(prompt_cluster_id1, req_id0, model_id0) cache_key_1 CacheKey(prompt_cluster_id1, req_id1, model_id0) cache_manager.pull_cache(cache_key_0, cache, batch_index0) cache_manager.pull_cache(cache_key_1, cache, batch_index1) dist.barrier() # pull_cache end cache_manager.deallocate_cache(cache) datadist.finalize()Prompt 侧分配并登记 KV 供拉取cache_desc CacheDesc( num_tensors4, shape[2, 16 * 1024], data_typeDataType.DT_FLOAT16, placementPlacement.DEVICE, ) cache_key_0 CacheKey(prompt_cluster_id1, req_id0, model_id0) cache_key_1 CacheKey(prompt_cluster_id1, req_id1, model_id0) cache cache_manager.allocate_cache(cache_desc, [cache_key_0, cache_key_1]) # ...等待 decoder pull_cache 完成... cache_manager.remove_cache_key(cache_key_0) cache_manager.remove_cache_key(cache_key_1) cache_manager.deallocate_cache(cache) datadist.finalize()这段代码中 Cache 对象被用作pull_cache的目标参数与deallocate_cache的释放参数其cache_id、cache_desc、tensor_addrs三个属性则在底层驱动整个传输调度。更丰富的样例push、blocks、异步分层传输等见 examples/python/README.md以及 examples/python/llm_datadist 目录下的push_cache_sample.py、pull_blocks_sample.py、push_blocks_sample.py、transfer_cache_async_sample.py、pull_from_cache_to_blocks.py等。常见错误与排查建议TypeError/ValueError多为create_cpu_cache传入的cache_desc类型错误、addrs元素非 int或cache_desc.num_tensors与地址数量不一致cache_desc placement must be HOSTcreate_cpu_cache要求描述 CPU 内存的CacheDesc.placement必须是Placement.HOST若误用默认的Placement.DEVICE会触发该校验CacheManager 失效调用LLMDataDist.finalize()后旧 CacheManager 不再可用通过旧实例操作会抛出LLM_ENGINE_FINALIZED需在重新init()后重新获取cache_manager与新 Cache 句柄地址失效tensor_addrs仅在对应内存未被释放前有效注册内存地址需自行保证不重复释放 Cache 后不要再引用其地址相关约束详见 CacheManager.md 中register_cache/register_blocks_cache的说明例如 Device 内存需先注册再建链、HCCS 传输要求首地址 2MB 对齐等。总结Cache 是 llm_datadist 所有 KV 传输操作的操作句柄本质是(cache_id, cache_desc, tensor_addrs)的三元组封装配合CacheManagerCacheManager.md、CacheDescCacheDesc.md、CacheKey/CacheKeyByIdAndIndexCacheKey.md、BlocksCacheKeyBlocksCacheKey.md等类型共同构成完整的 KV Cache 管理模型。掌握其三个只读属性与create_cpu_cache的使用边界即可正确地在 Prompt/Decoder 双角色集群中完成 KV 的分配、传输与释放。【免费下载链接】hixlHIXLHuawei Xfer Library是一个灵活、高效的昇腾单边通信库面向集群场景提供简单、可靠、高效的点对点数据传输能力。项目地址: https://gitcode.com/cann/hixl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表