链接复制成功!
同步加载KVCache
功能介绍
将推理请求需要的多个KVCache数据块从EMS缓存池读取KVCache到加速卡显存中,并立即返回结果。
接口约束
- 接口会休眠阻塞,建议根据业务实际情况设置超时时间。
- 仅支持华为昇腾加速卡显存拷贝。
- 只返回键列表中连续命中的个数。
方法定义
ContextCaching.load(option, key_list, value_list)
请求参数说明
参数名称 |
参数类型 |
是否必选 |
描述 |
---|---|---|---|
option |
必选 |
参数解释: ContextCaching访问内存池的KV操作选项。 约束限制: 不能为None。 取值范围: 无 默认取值: 无 |
|
key_list |
List[str] |
必选 |
参数解释: ContextCaching访问内存池的键名列表。 约束限制: 所有键名必须唯一。 取值范围: 单个key的长度小于128,且保证全局唯一。 默认取值: 无 |
value_list |
List[List[KvBufferWrapper]] |
必选 |
参数解释: ContextCaching访问内存池的值列表。 约束限制: 值列表的数目必须跟键列表中的数目相同,形成一一对应的键值对。 取值范围: 无 默认取值: 无。 |
参数名称 |
参数类型 |
是否必选 |
描述 |
---|---|---|---|
write_rcache |
bool |
可选 |
参数解释: 是否将本次写入保存为本地读缓存,默认值为True。 约束限制: 无 取值范围: True:将本次写入保存为本地读缓存 False:不将本次写入保存为本地读缓存 默认取值: True |
read_local_only |
bool |
可选 |
参数解释: 是否只读本地缓存,如果置为True,则不会从其他节点读取数据,只有读流程生效。 约束限制: 无 取值范围: True:只读本地缓存 False:优先读本地缓存,如果本地未命中,则从其他节点读取数据 默认取值: False |
timeout |
int |
可选 |
参数解释: 请求超时时间,单位为毫秒。 约束限制: 无 取值范围: 大于等于0。 默认取值: 5000 |
返回结果说明
类型 |
说明 |
---|---|
参数解释: 返回异步执行Future句柄。 取值范围: 无 |
参数名称 |
参数类型 |
描述 |
---|---|---|
success |
int |
参数解释: 请求的批量key读写连续成功的个数。 约束限制: 无 取值范围: 0~请求key的个数。 默认取值: 无 |
total |
int |
参数解释: 请求的批量key总个数。 约束限制: 无 取值范围: 请求列表中批量key的个数。 默认取值: 无 |
代码样例
下面示例描述了同步加载KVCache流程,同时对返回异常进行处理。
import os, torch, torch_npu
from ems import Ems, EmsConfig, EmsException, CcConfig, CcKvOption, KvBufferWrapper
# 初始化cc配置
cc_config = CcConfig(rank_id=8, device_id=0, model_id="llama2-13b")
# 初始化Ems
config = EmsConfig(cc_config=cc_config)
try:
Ems.init(config)
except EmsException as e:
print(f"exception: {e}.")
exit(1)
# 获取context caching对象
cc = Ems.get_cc()
if cc is None:
print("cc is None.")
exit(1)
# 设置save请求的超时时间
option = CcKvOption(timeout=5000)
# 组成键值列表
key_list = ["123", "66"]
# 仅支持npu设备的tensor
tensor1 = torch.ones(2, device="npu:1")
tensor2 = torch.ones(6, device="npu:1")
len1 = tensor1.numel() * tensor1.element_size()
len2 = tensor2.numel() * tensor2.element_size()
val_list = [[KvBufferWrapper(tensor1.data_ptr, len1)], [KvBufferWrapper(tensor2.data_ptr, len2)]]
# 可以根据不同异常,采取不同处理方式,例如超时错误可以重试。
try:
cc_result = cc.load(option, key_list, val_list)
except EmsException as e:
print(f"failed to load, {e}.")
exit(2)
print(f"succeed to load key num {cc_result.success}.")
print(cc_result)