KV Cache管理

功能介绍

在LLM-DataDist初始化时会预申请一块指定大小的内存池，由OPTION_BUF_POOL_CFG配置项决定其大小，后续的KV Cache的内存申请及释放都在预申请的内存上进行，相比每次执行时申请一块cache内存，可以节省耗时。

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

接口名称	功能
AllocateCache	分配Cache。
DeallocateCache	释放Cache。
PullKvCache	从远端节点拉取Cache到本地Cache，仅当角色为Decoder时可调用。
PullKvBlocks	PA场景下通过block列表的方式，从远端节点拉取Cache到本地Cache，仅当角色为Decoder时可调用。
CopyKvCache	拷贝KV Cache。支持D2D，D2H的拷贝。当期望PullKvCache和其他使用Cache的操作流水时，可以额外申请一块中转Cache。当其他流程在使用Cache时，可以先将下一次的Cache pull到中转Cache，待其他流程使用完Cache后，拷贝到指定的位置，从而通过pipeline流水将PullKvCache的耗时隐藏，减少总耗时。公共前缀场景在新请求推理前，可以将公共前缀拷贝到新的内存中与当前请求的kv合并推理。
CopyKvBlocks	PA场景下，通过block列表的方式拷贝KV Cache。支持D2D，D2H，H2D的拷贝。 D2D场景主要是针对当多个回答需要共用相同block，block没填满时，新增的token需要拷贝到新的block上继续迭代。 H2D和D2H的拷贝主要用于对应block_index上Cache内存的换入换出。

使用场景

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

功能示例（非PA场景）

本示例主要涉及KV Cache的申请、释放、传输。如下将根据业务角色给出伪代码示例，接口参数请参考LLM-DataDist接口列表。

P侧和D侧根据建链章节的示例完成LLM-DataDist的初始化和建链操作。

在P/D侧给每个请求申请对应大小的KV Cache内存，若失败，则需要释放对应的资源。

        
         
           
           
             void OnError(LlmDataDist &llmDataDist, Cache &cache)
{
    if (cache.cache_id > 0) {
        (void) llmDataDist.DeallocateCache(cache.cache_id);
    }
    llmDataDist.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};
Cache cache{};
auto ret = llmDataDist.AllocateCache(kv_cache_desc, cache);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] AllocateCache failed, ret = %u\n", ret);
    OnError(llmDataDist, cache);
    return -1;
}

            

          

        
       

将Cache从P侧传输到D侧，由D侧调用PullKvCache触发。

        
             // D侧
// 等待P侧写完cache，pull拉取cache
CacheIndex cacheIndex{PROMPT_CLUSTER_ID, 1, 0};
cacheIndex.batch_index = 0;
ret = llmDataDist.PullKvCache(cacheIndex, cache, 0);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PullKvCache failed, ret = %u\n", ret);
    return -1;
}
// 进行增量推理
// P侧
// 进行全量推理写cache，通知D侧可以pull

P/D侧根据业务中Cache的使用时机自行释放对应请求的KV Cache内存。

        
             ret = llmDataDist.DeallocateCache(cache.cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] DeallocateCache failed, ret = %u\n", ret);
} else {
    printf("[INFO] DeallocateCache success\n");
}

业务退出时，P侧和D侧根据断链章节的示例进行断链和调用finalize接口释放资源。

功能示例（PA场景）

本示例主要涉及KV Cache的申请、释放、传输。如下将根据业务角色给出伪代码示例，接口参数请参考LLM-DataDist接口列表。

P侧和D侧根据集群建链的示例完成LLM-DataDist的初始化和建链操作。

在P侧和D侧模型的每层按照计算好的num_block数量调用AllocateCache申请KV Cache。不同请求对创建的num_block大小的KV Cache进行复用，上层框架自行管理，业务结束后释放申请的内存。

       
        
          
          
            void OnError(LlmDataDist &llmDataDist, Cache &cache)
{
    if (cache.cache_id > 0) {
        (void) llmDataDist.DeallocateCache(cache.cache_id);
    }
    llmDataDist.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};
Cache cache{};
auto ret = llmDataDist.AllocateCache(kv_cache_desc, cache);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] AllocateCache failed, ret = %u\n", ret);
    OnError(llmDataDist, cache);
    return -1;
}

           

         

       
      

P侧有新请求进来后，会给每个请求分配好对应的block_index，这块是推理框架提供的功能，模型推理完之后，该请求对应的KV Cache就在对应的block_index所在的内存上，将模型输出和请求对应的block_table传输给D侧推理模型作为输入。

D侧有新请求进来后，也会给每个请求分配好对应的block_index，然后调用pull_blocks接口，根据P侧的block_index和D侧的block_index的对应关系，将KV Cache传输到指定位置。

       
            // D侧
// 等待P侧写完cache
std::vector<uint64_t> promptBlocks {1,2,3,6,5,4,7}; // 由P侧传过来
std::vector<uint64_t> decoderBlocks {1,2,3,6,5,4,7}; // 由推理框架自行分配
auto ret = llmDataDist.PullKvBlocks(cacheIndex, cache, promptBlocks, decoderBlocks);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] PullKvBlocks failed, ret = %u\n", ret);
    return -1;
}
// 进行增量推理
// P侧
// 进行全量推理写cache，通知D侧可以pull

P/D侧根据业务中Cache的使用时机自行释放对应请求的KV Cache内存。

       
            ret = llmDataDist.DeallocateCache(cache.cache_id);
if (ret != LLM_SUCCESS) {
    printf("[ERROR] DeallocateCache failed, ret = %u\n", ret);
} else {
    printf("[INFO] DeallocateCache success\n");
}

业务退出时，P侧和D侧根据断链章节的示例进行断链和资源释放。

异常处理

遇到LLM_DEVICE_OUT_OF_MEMORY，表示device申请KV Cache内存失败。需要检查初始化时设置的OPTION_BUF_POOL_CFG大小以及申请KV Cache的大小，查看是否有请求KV Cache拉取之后没有释放内存。
遇到LLM_KV_CACHE_NOT_EXIST，表示对端KV Cache不存在，需要检查对端进程是否异常或者对应KV Cache的请求有没有推理完成。该错误不影响其他请求流程，确认流程后可以重试。
遇到LLM_TIMEOUT，表示pull kv超时，说明链路出现问题，需要重新断链建链再尝试。
遇到LLM_NOT_YET_LINK，说明与远端cluster没有建链。

更多异常处理请参考错误码。

父主题： 功能介绍