KV Cache管理

功能介绍

在调用LLM-DataDist建链前可自行申请一块内存，并调用RegisterKvCache接口注册到LLM-DataDist，在断链后调用UnregisterKvCache接口进行解注册。

KV Cache管理涉及的主要接口及功能如下：

接口名称

功能

RegisterKvCache

注册本地Cache。

UnregisterKvCache

解注册本地Cache。

PullKvCache

从远端节点拉取Cache到本地Cache，与角色无关。

PullKvBlocks

在传输Blocks Cache场景下，通过配置block列表的方式，从远端节点拉取Cache到本地Cache，与角色无关。

PushKvCache

从本地节点推送Cache到远端节点，与角色无关。

PushKvBlocks

在传输Blocks Cache场景下，通过配置block列表的方式，从本地节点推送Cache到远端节点，与角色无关。

使用场景

主要用于分布式集群间的KV Cache传输和搬移。

功能示例（一般Cache传输场景）

本示例介绍一般Cache传输场景下接口的使用，主要涉及KV Cache的注册、解注册、传输。如下将根据业务角色给出伪代码示例。

  1. P侧和D侧根据建链章节的示例完成LLM-DataDist的初始化并申请内存并进行注册。
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
    void OnError(LlmDataDist &llm_data_dist, int64_t cache_id)
{
    if (cache_id > 0) {
        (void) llmDataDist.UnregisterKvCache(cache_id);
    }
    llm_data_dist.Finalize();
}
CacheDesc kv_cache_desc{};
kv_cache_desc.num_tensors = NUM_TENSORS;
kv_cache_desc.data_type = DT_INT32;
kv_cache_desc.shape = {8, 16};
std::vector<uint64_t> addrs;  // 按需填充
int64_t cache_id = -1;
auto ret = llm_data_dist.RegisterKvCache(kv_cache_desc, addrs, {}, cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] RegisterKvCache failed, ret = %u\n", ret);
    return -1;
}
  2. P侧和D侧根据建链章节的示例完成LLM-DataDist的建链操作
  1. P/D侧按需对KV Cache内存进行传输，若失败，则需要释放对应的资源，Cache传输有两种方式
    • 调用PullKvCache接口拉取KV Cache。
    1
2
3
4
5
6
7
8
9
    // P侧进行全量推理写入cache，通知D侧可以pull
// D侧拉取cache
CacheIndex cache_index{PROMPT_CLUSTER_ID, 1, 0};
ret = llm_data_dist.PullKvCache(cache_index, cache, 0);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PullKvCache failed, ret = %u\n", ret);
    return -1;
}
// 进行增量推理
    • 调用PushKvCache接口推送KV Cache。
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
    // P侧一层计算完可在传输线程立即推送，以实现将大部分传输时间和计算重叠。
CacheIndex dst_cache_index{DECODE_CLUSTER_ID, 1, 0};
KvCacheExtParam ext_param{};
ext_param.src_layer_range =  std::pair<int32_t, int32_t>(0, 0);
ext_param.dst_layer_range =  std::pair<int32_t, int32_t>(0, 0);
// 每层tensor数量，可根据实际模型修改
ext_param.tensor_num_per_layer = 2;
ret = llm_data_dist.PushKvCache(cache, dst_cache_index, 0, -1, ext_param);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PushKvCache failed, ret = %u\n", ret);
    return -1;
}
  2. 业务退出时，P侧和D侧根据断链章节的示例进行断链，调用UnregisterKvCache进行解注册并调用finalize接口释放资源。

功能示例（Blocks Cache传输场景）

本示例介绍Blocks Cache（将Cache使用块状形式管理）传输场景下接口的使用，主要涉及KV Cache的注册、传输。如下将根据业务角色给出伪代码示例。

  1. P侧和D侧根据集群建链的示例完成LLM-DataDist的初始化和建链操作。
  2. 在P侧和D侧模型的每层按照计算好的num_block数量申请KV Cache，并注册给LLM-DataDist。上层框架根据不同请求对创建的num_block大小的KV Cache进行复用，业务结束后释放申请的内存。
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
    void OnError(LlmDataDist &llm_data_dist, int64_t cache_id)
{
    if (cache_id > 0) {
        (void) llmDataDist.UnregisterKvCache(cache_id);
    }
    llm_data_dist.Finalize();
}
CacheDesc kv_cache_desc{};
kv_cache_desc.num_tensors = NUM_TENSORS;
kv_cache_desc.data_type = DT_INT32;
kv_cache_desc.shape = {8, 16};
std::vector<uint64_t> addrs;  // 按需填充
int64_t cache_id = -1;
auto ret = llm_data_dist.RegisterKvCache(kv_cache_desc, addrs, {}, cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] RegisterKvCache failed, ret = %u\n", ret);
    return -1;
}
  3. P侧有新请求进来后，为每个请求分配好对应的block_index，该功能由推理框架提供。模型推理完之后，该请求对应的KV Cache就在对应的block_index所在的内存上，将模型输出和请求对应的block_table传输给D侧推理模型作为输入。
  4. D侧有新请求进来后，为给每个请求分配好对应的block_index，然后调用pull_blocks接口，根据P侧的block_index和D侧的block_index的对应关系，将KV Cache传输到指定位置。
    • 在D侧调用PullKvBlocks接口拉取KV Cache。
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
    // P侧进行全量推理写入cache，通知D侧可以拉取cache
// D侧拉取cache
CacheIndex cache_index{PROMPT_CLUSTER_ID, 1, 0};
std::vector<uint64_t> prompt_blocks = {0, 1, 2, 3};
std::vector<uint64_t> decoder_blocks = {3, 2, 1, 0};
auto ret = llm_data_dist.PullKvBlocks(cache_index, cache, prompt_blocks, decoder_blocks);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PullKvBlocks failed, ret = %u\n", ret);
    return -1;
}
// 进行增量推理
    • 在P侧调用PushKvBlocks接口推送KV Cache。
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
    // P侧一层计算完可在传输线程立即推送，以实现将大部分传输时间和计算重叠。
CacheIndex dst_cache_index{DECODE_CLUSTER_ID, 1};
KvCacheExtParam ext_param{};
ext_param.src_layer_range =  std::pair<int32_t, int32_t>(0, 0);
ext_param.dst_layer_range =  std::pair<int32_t, int32_t>(0, 0);
// 每层tensor数量，可根据实际模型修改
ext_param.tensor_num_per_layer = 2;
std::vector<uint64_t> prompt_blocks = {0, 1, 2, 3};
std::vector<uint64_t> decoder_blocks = {3, 2, 1, 0};
ret = llm_data_dist.PushKvBlocks(cache, dst_cache_index, prompt_blocks, decoder_blocks, ext_param);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PushKvBlocks failed, ret = %u\n", ret);
    return -1;
}
  1. 业务退出时，P侧和D侧根据断链章节的示例进行断链，调用UnregisterKvCache进行解注册并资源释放。

异常处理

父主题： 功能介绍