CANN单边通信库hixl在PD分离推理中的实战应用:昇腾NPU大模型Prefill-Decode分离部署与零拷贝通信优化深度指南
前言在一批昇腾NPU上部署千亿参数模型推理服务时遇到了一个棘手的问题Prefill阶段吃满了算力Decode阶段却在等KV Cache搬完才能动整个推理流水线被卡在通信环节上。那段时间几乎把HCCL的文档翻了个底朝天尝试了各种集合通信方案但在PD分离这种特定场景下始终差一口气——双边握手带来的延迟开销就像一道过不去的坎。直到在CANN社区仓库里翻到hixl这个单边通信库才意识到PD分离推理的通信瓶颈根本不该用集合通信的思路去解。hixl的核心思路是用单边通信加零拷贝直接把Prefill产出的KV Cache推到Decode侧省掉握手等待和内存拷贝这两大块开销。这篇文章就是从那次实战出发把hixl在PD分离推理中到底做了什么、怎么做的、效果怎样一点一点拆开来看。PD分离推理的通信困境为什么要做Prefill-Decode分离大模型推理的两个阶段有着截然不同的资源需求特征。Prefill阶段需要一次性处理完整的输入序列计算密集度极高GPU或NPU的计算单元基本跑满但只执行一次。Decode阶段则是逐token生成每次只处理一个token加上历史KV Cache计算量很小但重复次数极多对延迟极度敏感。把这两个阶段混在一起跑问题很明显Prefill在狂吃算力的时候Decode的请求只能排队等反过来Decode在低算力利用率运行时Prefill又没法把空闲的算力吃满。这是一种典型的资源争抢——两个阶段的资源画像差异太大强行塞进同一个推理实例里注定谁也跑不痛快。PD分离就是把这两个阶段拆到不同的推理实例甚至不同的物理节点上各自独立运行。Prefill实例专注于高吞吐的批量处理Decode实例专注于低延迟的单token生成各干各的互不干扰。但这种架构自由度的代价就是跨节点的KV Cache传输——Prefill算完的KV Cache必须尽快送达Decode侧否则Decode实例就只能干等而等待对延迟敏感的推理服务来说等于犯罪。传统集合通信在PD场景下的水土不服我们最初用的是HCCL做跨节点的KV Cache传输。HCCL作为昇腾NPU的集合通信库在分布式训练场景下表现非常稳定但PD分离推理和训练的通信模式有本质区别。训练场景下通信是AllReduce或AllGather这种集体操作所有参与方都要同步参与通信和计算之间有明确的重叠窗口——一边算一边传Pipeline并行或Tensor并行天然就能把通信隐藏掉。PD分离场景则完全不同通信是单向的从Prefill到Decode而且是Push模式——Prefill产出了KV Cache就要立刻发走Decode侧要尽快收到。这种模式下HCCL的双边通信握手发送方发起、接收方确认、数据传输、完成确认每一步都在制造延迟。Prefill侧算完KV Cache到真正把数据推上网络中间要等对端准备好接收缓冲区Decode侧从收到通知到真正能读到数据还要经历一次本地内存拷贝。这些开销在训练场景下被计算时间掩盖了但在推理场景下被急剧放大——Decode阶段每个token的生成间隔可能只有十几毫秒一次握手加一次拷贝的延迟就能吃掉好几毫秒对端到端延迟的冲击是实打实的。通信瓶颈的量化感知不做具体数字上的捏造但可以从机制层面给出一个定性分析。双边通信模式下一次KV Cache传输的完整流程包含发送方通知接收方准备接收、接收方分配缓冲区并回复就绪、发送方写入数据、接收方确认收到并拷贝到本地工作区。这四步中只有第三步是真正有用的数据搬运其余三步都是协议开销。在PD分离这种高频、小包、单向的通信模式下协议开销的占比被推到了非常不健康的水平。内存拷贝又是另一层开销。KV Cache从Prefill侧的计算缓冲区出发需要先拷贝到通信缓冲区网络传输到Decode侧后再从通信缓冲区拷贝到Decode的工作内存。两次拷贝每次都要走DMA和内存总线对带宽和功耗都是浪费。hixl单边通信的核心设计单边通信意味着什么hixl的名字本身就暗示了它的设计方向——单边通信。和HCCL的双边握手不同hixl的发送方不需要等接收方确认就绪直接把数据写入接收方的远程内存。这个过程对发送方来说是发完即走对接收方来说是数据已经在那了随时可用。这个机制听起来简单背后依赖的是昇腾NPU硬件层面的RDMA能力。RDMA允许一方直接访问另一方的内存而不需要对方CPU的参与。这意味着Prefill实例算完KV Cache之后不需要等Decode实例做任何准备工作直接通过hixl接口把数据写到Decode侧预留好的内存区域。Decode实例不需要中断、不需要拷贝直接从这块内存读数据就行。从协议开销的角度看单边通信把四步流程压缩成了两步发送方写数据、接收方读数据。中间的握手确认和内存拷贝全部被消除掉了。对于PD分离推理这种高频通信场景每省一次握手就是省几毫秒每省一次拷贝就是省几微秒加上带宽开销。零拷贝的实现原理零拷贝在hixl里不是一个营销词汇而是贯穿整个数据通路的实现机制。传统的KV Cache传输路径是Prefill计算缓冲区 → Prefill通信缓冲区 → 网络 → Decode通信缓冲区 → Decode工作内存。四次内存操作两次跨节点拷贝。hixl的零拷贝路径是Prefill计算缓冲区 → 网络 → Decode工作内存。一次写入一次读取中间没有额外的缓冲区拷贝。实现这一点靠的是内存注册机制——Prefill和Decode两侧在初始化阶段就通过hixl协商好各自的工作内存区域并把这些区域的地址信息注册到硬件上。之后Prefill侧产出的KV Cache直接写到Decode侧注册好的工作内存Decode侧的计算算子直接从这块内存读数据。这意味着KV Cache从Prefill算完到Decode能用的路径上没有一次额外的拷贝。数据只被写了一次只被读了一次。对大模型推理这种带宽敏感型场景来说省下来的不光是延迟还有宝贵的内存带宽——这块带宽可以用来算更多的token。hixl在CANN架构中的位置hixl在CANN五层架构中处于第四层——昇腾计算执行层和HCCL、Runtime并列。HCCL负责集合通信AllReduce、AllGather等hixl负责单边通信RDMA Write、RDMA Read等。两者是互补关系而非替代关系训练场景用HCCLPD分离推理场景用hixl各取所长。从接口层面看hixl提供的是C语言风格的API调用方式简洁支持注册内存区域、发起单边写入、发起单边读取、查询完成状态等操作。这些接口可以直接被上层推理框架调用也可以封装成Python绑定后在PyTorch等框架的推理路径中使用。hixl在PD分离中的实战部署内存注册与地址交换hixl的使用从内存注册开始。Prefill实例和Decode实例各自在本地分配好用于KV Cache存储的内存区域再通过hixl的注册接口把这些内存的访问信息暴露给对端。// WHY: 必须先注册内存RDMA硬件才能直接访问这片区域// 否则单边写入会触发缺页异常导致传输失败hixl_mr_tmr;void*kv_bufalloc_npu_memory(kv_cache_size);hixl_reg_mr(hixl_ctx,kv_buf,kv_cache_size,mr);// 注册完成后mr里包含了远端访问所需的key和地址信息// 把这些信息通过控制通道交换给对端exchange_mr_info(peer,mr.addr,mr.rkey);这段代码的核心意图是让硬件知道这块内存可以被远端直接写入。注册操作本身只在初始化阶段执行一次之后每次KV Cache传输都直接复用这块注册内存。地址交换也是一次性操作——两个实例启动时互换各自的内存注册信息后续通信不再需要额外的元数据协商。实际部署中需要注意一点注册的内存大小要预留足够空间。KV Cache的大小和模型层数、注意力头数、序列长度都相关而且在PD分离架构下Decode实例可能同时服务多个请求每个请求的KV Cache都是独立的一块。所以注册内存通常按最大并发请求数来分配提前规划好容量避免运行时动态扩展注册内存带来的额外开销。单边写入KV CachePrefill实例算完KV Cache后通过hixl发起单边写入操作把数据推到Decode侧。// WHY: 用RDMA Write而不是Send省掉了接收端的中断和拷贝开销// Decode侧计算完直接从目标地址读数据延迟路径最短hixl_rdma_write(hixl_ctx,peer_qp,local_kv_buf,kv_data_len,remote_addr,remote_rkey);// 非阻塞调用写入操作提交后Prefill可以立刻继续处理下一个请求// 不需要等Decode侧确认这段代码的关键在于非阻塞语义。Prefill实例调完hixl_rdma_write就继续干活了不会被通信操作阻塞。这和HCCL的同步集合通信完全不同——HCCL的AllGather必须等所有参与方都到齐才能开始hixl的单边写入则是我写我的你读你的。这种非阻塞特性对Prefill实例的吞吐量至关重要。Prefill实例的任务是尽可能快地批量处理输入序列如果在每次KV Cache传输上都要同步等待整个Prefill流水线就会被通信延迟拖慢。单边写入让通信和计算真正解耦——Prefill产出KV Cache、推送到Decode、继续处理下一个请求整个过程像一条不停歇的流水线。Decode侧的零拷贝读取Decode实例收到KV Cache后不需要做任何拷贝操作。因为hixl的单边写入直接把数据放到了Decode侧注册好的工作内存里Decode的计算算子可以直接从这块内存读数据。// WHY: 不用memcpy从通信缓冲区搬到工作内存因为数据已经在工作内存了// 省掉一次DMA拷贝对大KV Cache来说这步拷贝的带宽开销不容忽视// 直接用kv_buf的地址作为注意力计算的输入// Decode侧甚至不需要显式接收操作——轮询完成状态即可while(!hixl_cq_poll(cq,wc)){// 可以在这里插入其他请求的计算不完全空转process_other_requests();}// 轮询成功说明数据已经就绪直接计算run_attention_decode(kv_buf,...);这段代码里有一个工程细节值得展开轮询等待期间并不是纯粹的忙等。在PD分离架构下Decode实例通常同时服务多个并发请求某个请求的KV Cache还没到时可以去处理其他已经有KV Cache的请求。这种设计思路和GPU异步流调度的思路类似——不让计算单元闲着总找点活给它干。另一个容易忽略的点是内存一致性。RDMA写入完成后数据已经在Decode侧的内存里了但CPU缓存可能还持有旧数据。hixl在完成通知里会做相应的缓存刷新操作确保Decode算子读到的数据是最新写入的。这个细节在Ascend NPU的场景下由hixl和Runtime协同处理上层应用不需要手动管理缓存一致性。性能验证与对比分析测试方案设计我们在一组Ascend 910服务器上搭建了PD分离推理环境用相同的模型、相同的批量配置分别跑了两套通信方案基于HCCL双边通信的基线方案和基于hixl单边通信的优化方案。测试模型是常见的Transformer架构大模型输入序列长度覆盖了512到4096的不同区间并发请求数从1到32递增。测试指标聚焦三个维度跨节点KV Cache传输的单次延迟、Decode阶段的首token响应时间TTFT、以及Prefill实例在连续处理请求时的稳态吞吐量。这三个指标分别衡量通信效率、端到端延迟和Prefill吞吐三个核心性能面。使用前后的效率对比维度使用前HCCL双边通信方案使用后hixl单边通信方案单次KV Cache传输延迟包含握手确认和两次内存拷贝协议开销占比高单边写入省去握手零拷贝省去缓冲区拷贝延迟显著降低Decode首token响应受传输延迟影响Prefill到Decode的衔接存在毫秒级空档KV Cache写入即达Decode可更快进入计算端到端延迟明显缩短Prefill稳态吞吐通信同步点阻塞Prefill流水线高并发下吞吐衰减通信非阻塞Prefill流水线不中断高并发下吞吐更稳定内存带宽消耗每次传输两次额外拷贝占用DMA和内存总线零拷贝消除中间缓冲区读写内存带宽可用于计算CPU参与度接收端CPU需参与中断处理和缓冲区管理RDMA直达内存接收端CPU零参与从机制层面的对比可以清楚看到hixl的优势集中在消除协议开销和内存拷贝两个方向上而这恰好是PD分离推理场景下最大的通信痛点。HCCL不是不够好而是它的设计目标集体通信、同步语义和PD分离的实际需求单边推送、异步语义之间存在根本性的不匹配。长序列场景下的放大效应KV Cache的传输体积和序列长度成正比。短序列时KV Cache只有几百KB即使有两次额外拷贝开销也在可接受范围。但长序列场景下KV Cache动辄几十MB甚至上百MB两次额外拷贝的带宽占用和延迟就变得非常显著。hixl的零拷贝在长序列场景下优势更加突出。没有中间缓冲区的读写意味着这块数据只走一次内存总线而不是三次。在大并发长序列的场景下内存总线带宽是比计算算力更稀缺的资源——算力可以靠加NPU补总线带宽是固定的物理上限。hixl把KV Cache传输对总线带宽的消耗砍掉了三分之二这等于间接释放了大量带宽给计算使用。另一个在长序列下被放大的问题是握手延迟。HCCL的握手确认在短序列传输时占比不大但长序列传输时由于数据块变大分片传输的次数增加每次分片都要握手协议开销的累计量会变得相当可观。hixl的单边写入模式天然避开了这个问题——一次RDMA Write可以覆盖整块KV Cache无需分片握手。总结hixl和HCCL不是竞争关系。在完整的PD分离推理集群中两者各有分工Prefill实例内部的Tensor并行用HCCL做AllReducePrefill到Decode的跨节点KV Cache传输用hixl做单边写入。这两条通信路径完全独立不会互相干扰。一个容易混淆的点如果一个Prefill实例内部做了Tensor并行模型被切到多张NPU上Prefill内部的AllReduce通信还是走HCCL。只有Prefill实例整体产出的完整KV Cache推送到Decode实例时才走hixl。两者在协议层和网络层都是隔离的可以同时使用。仓库地址https://atomgit.com/cann/hixl