1 摘要

大模型推理的Prefill和Decode阶段往往存在计算和内存受限问题，资源分配不均，导致成本居高不下，这些趋势和挑战，都驱动着推理系统的重构，以更高效的资源调度与并行解码，满足推理商业落地的经济性要求。在昇腾AI异构计算架构CANN最新推出的8.0版本中，基于LLM P-D分离部署方案设计并发布LLM-DataDist组件，API简单易用，可以低开发成本被MindIE-LLM、vLLM等大模型推理框架集成，并在近一年支撑了多个大模型业务的规模化商用。

2 P-D分离部署方案介绍

ChatGPT的推出意味着交互式人工智能逐渐成为商业化成熟产品，同时也进一步推动了底层技术大型语言模型LLM的研究和进步。现有主要使用的LLM有GPT系列、LLAMA系列、GLM系列，其模型基础架构都采用Transformer架构，并以Decoder-Only为主。

该类Transformer-Based-Decoder-Only的LLM在推理预测时，采用自回归生成（auto-aggressive generative）模式，即每个Token生成需要经过LLM模型的前向推理过程，即完成N次Transformer Layer层计算，意味着包含M个Token的语句完整生成需要经过M次完整LLM前向推理过程。

作为交互式AI技术的典型应用，ChatGPT对人机交互的响应速度有着极高的要求，从人类阅读速度出发，LLM的推理速度至少需要达到50ms/token。为了充分利用昇腾硬件强大的并行计算优势，推理系统在服务和任务调度层面实施了一系列工程优化措施，旨在通过提高数据密度来最大化并行计算的优势，例如增大Batch Size、扩展Sequence长度等。降低LLM推理时延、提升LLM推理集群的吞吐和算力利用率，成为相关AI应用大规模产品化变现的迫切要求。

在Transformer推理过程中利用KV Cache技术可降低Decoding阶段的计算量，目前已成为LLM推理系统的必选技术。

采用KV Cache的LLM推理过程通常分为预填充（Prefill）和解码（Decode）两个阶段。

• Prefill阶段：将用户请求的prompt传入大模型，进行计算，中间结果写入KV Cache并推出第1个token。随着Prompt Sequence Length长度线性增长，对首Token时延（TTFT）指标有要求的业务通常采用单batch方式执行LLM推理，该阶段属于计算密集型操作。
• Decode阶段：将请求的前1个Token传入大模型，从显存读取前文产生的KV Cache再进行计算。采用KV Cache技术后，单Token计算量低，而主要瓶颈在于搬运参数量，故通常需要采用多batch方式提升利用率，该阶段属于访存密集型操作。

图2-1 基于KV Cache的LLM推理过程

从实现上来看，整个LLM推理过程由Prefill和多轮迭代的Decode完成，如下图所示，想要在Decode阶段实现Continuous batching的前提是，每个被调度的Request需要空闲算力完成Prefill计算，按现有的部署模式，当Prefill和Decode部署在一起时，Prefill和Decode会交叉执行，从而导致增量Token时延（TBT）无法得到有效保障。

例如下图，当request5或request6到来时，系统可能会优先执行request5或request6的Prefill，此时request2/3/4的响应时延会受到一定影响，从而导致TBT不稳定。

图2-2 LLM响应时延分析

在实际的深度学习模型部署中，考虑到Prefill和Decode两阶段的计算/通信特征的差异特点，为了提升性能和资源利用效率，通过P-D分离部署方案将Prefill和Decode分别部署在不同规格和架构的集群中，并且配合服务层的任务调度，在满足TTFT和TBT指标范围内，结合Continuous batching机制尽可能提高Decode阶段的batch并发数，在提供更好用户体验的前提下，提升算力利用率。

图2-3 P-D分离部署方案

基于该方案，结合下图可以看到Prefill和Decode的执行互不影响，系统能够提供给用户一个稳定的TBT。

图2-4 基于P-D分离部署方案的LLM响应时延分析

基于以上收益，自2023年11月以来，业界对P-D分离的研究和应用持续火热，昇腾异构计算架构CANN针对LLM P-D分离部署方案设计LLM-DataDist组件，这一组件在支撑多个客户场景规模化落地方面发挥了重要作用，并经过多轮迭代优化后，作为CANN 8.0的关键特性之一正式发布。

3 LLM-DataDist整体架构

LLM-DataDist作为大模型分布式集群和数据管理组件，提供KV Cache传输及链路管理能力，并通过简易的API开放给用户，支持快速构建布式集群链路、动态扩缩等特性。同时利用昇腾集群多样化通信链路（HCCS/RoCE）可实现跨实例和集群的高效KV Cache传输，并支持与主流LLM推理框架如MindIE-LLM、vLLM等的集成。

图3-1 LLM-DataDist整体架构

从完整的P-D分离系统角度来看，LLM-DataDist通过link-mgr和cache-mgr支撑起了基于昇腾AI硬件的P-D分离系统的构建：

图3-2 基于CANN的P-D分离系统示意图

从一个简单的PD分离系统角度看来，基于LLM-DataDist提供的cache-mgr和link-mgr，上层可以构建如下能力：

• 集群管理层：基于link-mgr来构建起P-D集群的管理，并可以在业务的忙闲时以及故障场景下完成Prefill或Decode节点的上下线。
• 服务化框架层：基于cache-mgr的能力来构建起P-D之间的KV Cache传输，可以结合常见的KV Cache优化手段，例如PA（Paged Attention）、公共前缀缓存等。

4 LLM-DataDist支持动态建链，灵活调整集群配比

根据业务流量和Input/output负载不同，基于LLM-DataDist的动态链路管理能力，可以灵活设置P-D节点的集群配比数，在Serving服务运行过程中支持时延无感知的节点动态扩缩容，也可以在运行过程中实时灵活切换P-D角色。

开放API支持P-D角色划分

为了方便角色的管理，抽象出LLMDataDist类来表示Prefill还是Decode，通过实例化参数role进行区分，LLMDataDist类主要的几个方法如下：

• init：初始化LLM-DataDist，初始化过程中可以设置相关的超时时间，所使用的device等信息。
• finalize：与初始化相对应，用于释放资源。
• switch_role：切换当前LLM-DataDist的角色，支持在不中断业务的情况下，无缝互换P-D角色。

开放API支持链路管理

NN模型执行时的HCCL集合通信接口是双边通信，即需要两边同时发起建链，而在P-D分离方案中，简化建链操作，由Client单侧发起建链。由于动态扩缩的部分往往是Decode侧，因此将P定义为Server端，D定义为Client端，建链过程实现由D向P发起建链的流程，接口能力上主要提供了如下两个API：

• link_clusters，发起方为Decode侧，参数为Prefill侧的卡ip和端口号，示例如下：

from llm_datadist import LLMDataDist, LLMRole, LLMStatusCode, LLMClusterInfo 
llm_datadist = LLMDataDist(LLMRole.DECODER, 0) 
cluster = LLMClusterInfo() 
cluster.remote_cluster_id = 1 
cluster.append_local_ip_info("1.1.1.1", 26000) 
cluster.append_remote_ip_info("1.1.1.1", 26000) 
ret, rets = llm_datadist.link_clusters([cluster], 5000) 
if ret != LLMStatusCode.LLM_SUCCESS: 
    raise Exception("link failed.") 
for cluster_i in range(len(rets)): 
    link_ret = rets[cluster_i] 
    if link_ret != LLMStatusCode.LLM_SUCCESS: 
        print(f"{cluster_i} link failed.")

• unlink_clusters，发起方为Decode侧，参数为Prefill侧的卡ip和端口号，示例如下：

from llm_datadist import LLMDataDist, LLMRole, LLMStatusCode, LLMClusterInfo 
llm_datadist = LLMDataDist(LLMRole.DECODER, 0) 
cluster = LLMClusterInfo() 
cluster.remote_cluster_id = 1 
cluster.append_local_ip_info("1.1.1.1", 26000) 
cluster.append_remote_ip_info("1.1.1.1", 26000) 
ret, rets = llm_datadist.unlink_clusters([cluster], 5000) 
if ret != LLMStatusCode.LLM_SUCCESS: 
    raise Exception("unlink failed.") 
for cluster_i in range(len(rets)): 
    unlink_ret = rets[cluster_i] 
    if unlink_ret != LLMStatusCode.LLM_SUCCESS: 
        print(f"{cluster_i} unlink failed.")

link-mgr主要业务应用场景举例：

• 集群可靠性，当有P或D节点故障时，在不影响整个集群可用性的前提下，下线对应机器。
• 业务闲忙时应对，可以动态调整P-D的配比，以最大化LLM集群的吞吐。

5 LLM-DataDist支持P-D间高效KV Cache传输

对于分离集群间的KV Cache传输，通过昇腾AI处理器自带的8 * 200Gb RoCE网卡，独立的KV Cache传输通信链路与P/D模型并行使用的HCCS链路完全物理隔离，使大多数场景传输额外时延控制在一个Token的生成开销内，做到KV Cache传输被计算时延所掩盖。

开放API支持KV Cache管理

cache-mgr主要提供如下API：

• allocate_cache：申请KV Cache，为方便后续的KV Cache传输，可以为对应的KV Cache设置一个key。
• deallocate_cache：释放KV Cache。
• remove_cache_key：基于key删除在Prefill节点上的一块KV Cache，一旦删除，这块KV Cache就不可以被拉取。
• pull_cache：由Decode节点发起，从Prefill节点拉取对应的KV Cache。
• copy_cache：将KV Cache进行copy，在上一次nn没执行完成时，可以先将下一次的KV Cache进行暂存，待nn执行结束后，copy到指定的位置，从而最大程度打满Decode的batch。

同时，针对PA（Paged Attention）场景提供单独一套API，在pull blocks时进行传输上的性能优化，利用Roce网卡的Recv Scatter能力进行加速：

• allocate_blocks_cache：PA场景下，KV blocks的申请接口。
• pull_blocks：PA场景下，KV blocks的拉取接口。
• copy_blocks：PA场景下，KV blocks的copy接口。
• swap_blocks：KV blocks的换入换出，主要面向用户需要自行管理KV Cache的场景。