KV Cache管理

功能介绍

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

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

使用场景

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

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

本示例分别介绍PA场景和非PA场景下接口的使用，主要涉及KV Cache的申请、释放、传输。

如下将根据业务角色给出伪代码示例，接口参数可参考LLM-DataDist接口列表。

1. P侧和D侧根据链路管理章节的示例完成LLMDataDist的初始化和建链操作。
2. 在P侧给每个请求申请对应大小的KV Cache内存，以torch为例，将KV Cache转化为torch tensor，进行全量模型推理。

   1 2 3 4 5 6 7 8 9 10 11 12 13 14 15

   import torchair import torch import torch_npu # 从初始化完的llm_datadist中获取kv_cache_manager kv_cache_manager = llm_datadist.kv_cache_manager # 根据模型中KV Cache的shape以及总个数创建CacheDesc cache_desc = CacheDesc(num_tensors=4, shape=[4, 4, 8], data_type=DataType.DT_FLOAT16) # 根据初始化llm_datadist时的cluster_id创建对应请求的cache_key，当P侧是多batch模型时，需要创建batch数量的cache_key batch_num = 8 kv_cache = kv_cache_manager.allocate_cache(cache_desc, [CacheKey(prompt_cluster_id=0, req_id=i, model_id=0) for i in range(batch_num)]) # 将申请好的KV Cache转换为torch tensor kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0] kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs) # 将转换的kv_tensors传给模型推理计算产生KV Cache，将模型输出传输给增量推理模型作为输入

3. 在D侧申请一份用于模型执行的KV Cache内存。

   1 2 3 4 5 6

   # 从初始化完的llm_datadist中获取kv_cache_manager kv_cache_manager = llm_datadist.kv_cache_manager # 根据模型中KV Cache的shape以及总个数创建CacheDesc cache_desc = CacheDesc(num_tensors=4, shape=[4, 4, 8], data_type=DataType.DT_FLOAT16) # 调用allocate_cache接口申请对应请求的KV Cache内存 kv_cache = kv_cache_manager.allocate_cache(cache_desc)

4. 将KV Cache从Prompt传输到Decode，有以下两种方式：
   • 在Decode调用pull_cache接口拉取对应请求的KV Cache到申请的内存中。

     1 2 3

     # 创建和P侧申请cache时相同的cache_key，用于拉取对应的KV cache cache_key = CacheKey(prompt_cluster_id=0, req_id=1, model_id=0) kv_cache_manager.pull_cache(cache_key, kv_cache, batch_index=1) # 拉到batch index为1的位置上

   • 在Prompt调用transfer_cache_async接口将数据传输至Decode。

     1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22

     from llm_datadist import LayerSynchronizer, TransferConfig class LayerSynchronizerImpl(LayerSynchronizer): def __init__(self, events): self._events = events def synchronize_layer(self, layer_index: int, timeout_in_millis: Optional[int]) -> bool: self._events[layer_index].wait() return True events = [torch.npu.Event() for _ in range(cache_desc.num_tensors // 2)] # 执行模型，模型在各层计算完成后调用events[layer_index].record()记录完成状态 # 模型执行由用户实现 # user_model.Predict(kv_tensors, events) # 模型下发完成后，调用transfer_cache_async传输数据，此处需要填写Decode已申请的KV Cache各层tensor的内存地址 transfer_config = TransferConfig(DECODER_CLUSTER_ID, decoder_kv_cache_addrs) cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(events), [transfer_config]) # 同步等待传输结果 cache_task.synchronize()

5. 以torch为例，将KV Cache转化为torch tensor，进行增量模型推理。

   1 2 3 4 5 6 7 8 9 10 11 12

   # 将申请好的KV Cache转换为框架中的KV Cache类型，不同框架中都需要提供根据KV Cache地址创建对应类型的KV Cache的接口。此处以PyTorch为例 # 转换操作和pull操作顺序不分先后 kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0] kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs) # 将转换后的tensor拆分为框架需要的kv配对方式，可以自定义组合kv mid = len(kv_tensors) // 2 k_tensors = kv_tensors[: mid] v_tensors = kv_tensors[mid:] kv_cache_tensors = list(zip(k_tensors, v_tensors)) # 将转换的kv_tensors传给模型进行迭代推理 # 等待请求增量推理完成

6. 根据业务中cache的使用时机自行释放对应请求的KV Cache内存。

   1 2 3 4

   # P侧：当batch中所有请求的KV Cache都被拉走后，调用deallocate_cache才会真正释放内存。如果D侧未拉走KV Cache，则还需要调用remove_cache_key。 kv_cache_manager.remove_cache_key(cache_key_0) kv_cache_manager.remove_cache_key(cache_key_2) kv_cache_manager.deallocate_cache(kv_cache)

   1 2	# D侧：由于申请时不需要cache_key，所以释放时只需要调用deallocate接口。 kv_cache_manager.deallocate_cache(kv_cache)

7. 业务退出时，P侧和D侧根据链路管理章节的示例进行断链和finalize。

   1	llm_datadist.finalize()

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

本示例分别介绍PA场景和非PA场景下接口的使用，主要涉及KV Cache的申请、释放、传输。

如下将根据业务角色给出伪代码示例，接口参数可参考LLM-DataDist接口列表。

1. P侧和D侧根据链路管理的示例完成LLMDataDist的初始化和建链操作。
2. 在P侧和D侧模型的每层按照计算好的num_block数量调用AllocateCache申请KV Cache。Blocks Cache场景下，不同请求对创建的num_block大小的KV Cache进行复用，上层框架自行管理，业务结束后释放申请的内存。

   1 2 3 4 5 6 7 8 9 10 11 12 13 14 15

   # P侧 # 从初始化完的llm_datadist中获取kv_cache_manager kv_cache_manager = llm_datadist.kv_cache_manager # 根据模型中KV Cache的shape以及总个数创建CacheDesc num_blocks = 10 block_mem_size = 128 cache_desc = CacheDesc(num_tensors=4, shape=[10, 128], data_type=DataType.DT_FLOAT16) # 根据初始化llm_datadist时的cluster_id创建对应请求的BlocksCacheKey cache_key = BlocksCacheKey(prompt_cluster_id=0, model_id=0) # 调用allocate_blocks_cache接口申请KV Cache内存 kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc, cache_key) # 将申请好的KV Cache转换为框架中的KV Cache类型，不同框架中需要提供根据KV Cache的地址创建对应类型的KV Cache的接口。此处以PyTorch为例 kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0] kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs)

   1 2 3 4 5 6 7 8 9

   # D侧 kv_cache_manager = llm_datadist.kv_cache_manager num_blocks = 10 block_mem_size = 128 cache_desc = CacheDesc(num_tensors=4, shape=[10, 128], data_type=DataType.DT_FLOAT16) kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc) kv_tensor_addrs = kv_cache.per_device_tensor_addrs[0] kv_tensors = torchair.llm_datadist.create_npu_tensors(kv_cache.cache_desc.shape, torch.float16, kv_tensor_addrs)

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传输到指定位置。此时有两种方式：
   • 在Decode调用pull_blocks接口拉取KV Cache。

     1 2 3 4

     # D侧根据P侧传过来的信息，添加新请求，并申请对应的block_table # D侧根据传过来请求的src_block_table和新申请的dst_block_table拉取kv到对应block cache_key = BlocksCacheKey(prompt_cluster_id=0, model_id=0) # P侧allocate_blocks_cache时的入参 kv_cache_manager.pull_blocks(cache_key, cache, [0, 1], [2, 3]) # 将P侧0, 1 block位置上的数据拉到D侧2, 3 block位置上

   • 在Prompt调用transfer_cache_async接口时将数据传输至Decode。

     1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21

     # 实现LayerSynchronizerImpl，通过torch Event获取各层计算结束状态，本例中通过Event机制实现 class LayerSynchronizerImpl(LayerSynchronizer): def __init__(self, events): self._events = events def synchronize_layer(self, layer_index: int, timeout_in_millis: Optional[int]) -> bool: self._events[layer_index].wait() return True events = [torch.npu.Event() for _ in range(cache_desc.num_tensors // 2)] # 执行模型, 模型在各层计算完成后调用events[layer_index].record()记录完成状态 # 该函数由用户实现 user_model.Predict(kv_cache_tensors, events) # 模型下发完成后，调用transfer_cache_async传输数据，此处需要填写Decode申请出的KV Cache各层tensor的内存地址 transfer_config = TransferConfig(DECODER_CLUSTER_ID, decoder_kv_cache_addrs) cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(events), [transfer_config], [0, 1], [2, 3]) # 同步等待传输结果 cache_task.synchronize()

5. 业务结束后P侧和D侧调用deallocate_cache释放已申请的KV Cache内存。

   1 2 3

   # 等待D侧拉取完对应请求的KV Cache # 根据业务中cache的使用时机自行释放对应请求的KV Cache。PA场景无需释放对应的cache_key kv_cache_manager.deallocate_cache(kv_cache)

6. 业务退出时，P侧和D侧根据集群断链的示例进行断链和llm_datadist的finalize。

异常处理

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

更多异常处理请参考LLMStatusCode。