从寄信到直投:hixl单边通信库如何拆掉PD分离场景中的数据搬运墙——昇腾CANN计算基础层的跨步通信原语深度拆解
前言
大模型推理正在经历一场架构层面的变革。当模型参数量突破千亿,单卡已经无法装下完整的推理流程,PD分离架构应运而生——Prefill阶段和Decode阶段分别部署在不同的昇腾NPU上。这种拆分让两个阶段可以独立扩缩容,但也带来了一个物理层面的问题:Prefill产出的KV Cache必须搬到Decode所在的卡上,数据搬运的延迟直接吃掉了分离架构的收益。
hixl就是为解决这个数据搬运问题而生的通信库。它位于昇腾CANN生态的第五层——计算基础层,提供put/get语义的单边通信能力,支持零拷贝和跨步内存访问。与上层的hccl集合通信库不同,hixl更像是通信领域的"快递专线":不需要收件人签收,寄件人直接把东西放进对方信箱,收件人随时来取。
通信的两种姿势:双边与单边
理解hixl之前,需要先搞清楚"单边通信"到底是什么意思。这里用一个日常场景做类比。
想象两个人之间传递一份文件。双边通信就像打电话:你必须拨号、对方必须接听,两个人同时在线才能完成一次传递。任何一方不在,通信就阻塞。hccl的集合通信操作(AllReduce、Broadcast等)就属于这种模式——所有参与方必须同步到场,数据才能流动。
单边通信则像往对方信箱里投信:你把信塞进去就可以走了,不需要对方此刻站在信箱旁。对方什么时候来取信、取哪封信,都由对方自己决定。hixl的put操作就是这个"投信"动作——发起方直接把数据写入远端内存,不需要远端进程参与。get操作则反过来——发起方直接从远端内存读取数据,远端进程同样无感。
这两种模式的差异在PD分离场景下被放大。Prefill卡产出一批KV Cache后,如果用双边通信搬运,Decode卡必须同步等待,两边的计算节奏被强制耦合。换成单边通信,Prefill卡算完直接put到Decode卡的内存,Decode卡按自己的节奏消费数据,两边的计算流水线解耦了。这就是hixl在PD分离架构中的核心价值——用通信模式的切换,换取计算流水线的独立。
零拷贝:省掉那趟多余的搬运
单边通信解决了"谁参与"的问题,零拷贝解决的是"搬几次"的问题。
继续用类比。假设你要把一摞文件从A办公室搬到B办公室。常规做法是:先从A的柜子里拿出来放到推车上,推到B办公室,再从推车放到B的柜子里。文件被搬了两次——从柜子到推车、从推车到柜子。这就是传统通信中的拷贝开销:应用缓冲区到通信库内部缓冲区是一次拷贝,通信库内部缓冲区到对端应用缓冲区又是一次拷贝。
零拷贝的做法是:直接把A柜子里的文件标记为"通信区",通信硬件从这块区域直接读取数据发到对端,对端硬件直接写入B柜子中预先标记好的区域。文件始终在柜子里,没有被搬到推车上过。
hixl的零拷贝实现依赖昇腾NPU的DMA能力。应用在初始化时注册好内存区域,后续的put/get操作直接在这些注册内存上工作,数据从源端内存经HCCS链路直接到达目的端内存,中间不经过任何软件缓冲区。对于KV Cache这种动辄数百兆的大块数据,省掉一次拷贝就意味着省掉数百兆的内存带宽和对应的搬运时间。
HXLResult ret=hxl_put(src_npu,dst_npu,src_addr,dst_addr,size);这行代码的语义是"把src_npu上src_addr处的size字节数据,直接写入dst_npu的dst_addr处"。没有中间缓冲区参与,没有对端进程的配合。src_addr和dst_addr都必须是预先注册过的通信内存区域,hixl在内部把这些地址映射到硬件可识别的DMA描述符,由昇腾NPU的通信引擎执行实际的数据搬运。整个过程中,CPU侧只负责下发这条指令,不参与数据搬运本身。
跨步内存访问:不连续数据的连续传输
KV Cache在内存中的布局往往不是连续的。以Transformer的多头注意力为例,每个头的KV数据在内存中是按头维度连续存储的,但头与头之间可能隔着其他数据。如果要把多个头的KV Cache搬过去,最直观的做法是逐头搬运——每头发一次put,每次搬运一个头的连续数据。但这样通信次数太多,每次通信都有启动开销,小包密集传输的效率很低。
跨步内存访问(Strided Access)解决的就是这个问题。它的思路是:告诉硬件"起始地址在哪、每个元素的长度是多少、相邻元素之间的间距是多少",硬件就能自动跳着步子把不连续的数据拼成连续的流发出去。类比来看,就像你告诉快递员"这栋楼3层到10层,每隔一层的门口都有一个包裹,你按这个规律走一遍全收走",而不是你自己跑八趟每趟拿一个。
hixl通过HXLStridedDescriptor来描述这种跨步内存布局:
HXLStridedDescriptor desc;desc.count=num_heads;desc.size=head_dim*sizeof(float);desc.stride=head_stride;ret=hxl_put_strided(src_npu,dst_npu,src_base,dst_base,&desc);count是要搬运的元素个数(这里是注意力头数),size是每个元素的字节长度(一个头的KV数据长度),stride是相邻元素起始地址之间的间距(包含头间填充)。硬件根据这三个参数,自动计算每次访问的地址:第i个元素的地址是base + i * stride。这样一次hxl_put_strided调用就完成了所有头的搬运,通信次数从num_heads次降为一次,启动开销被摊薄到所有元素上。
跨步访问还有一个隐含的约束:64字节对齐。昇腾NPU的DMA引擎以64字节为最小传输单元,所以src_base、dst_base以及stride都必须是64字节的整数倍。这不是hixl的任意限制,而是硬件通信引擎的物理要求。如果应用层的内存布局不满足对齐条件,需要在注册通信内存时做对齐处理。
hixl与hccl:底层原语与上层编排
在CANN的通信体系中,hixl和hccl的关系类似于汇编指令和高级语言的关系。
hccl提供的是集合通信原语:AllReduce、AllGather、ReduceScatter、Broadcast等。这些操作关注的是"一组NPU之间如何协同完成某个全局计算",比如AllReduce是把所有卡上的梯度求和再分发回去。hccl内部会根据NPU数量、拓扑结构和数据量,自动编排通信步骤——谁先发、谁后发、数据走哪条链路。这种编排依赖底层的点对点通信能力,而hixl就是提供这种点对点能力的底层库。
从调用链来看,hccl的集合通信操作最终会分解为大量的点对点send/recv或put/get操作。hixl为这些底层操作提供了更高效的实现路径:零拷贝减少了中间拷贝,单边语义减少了对端参与开销,跨步访问减少了通信次数。hccl可以选择性地使用hixl作为其底层传输引擎,也可以使用其他传输实现——这取决于具体的硬件平台和通信拓扑。
对于应用开发者来说,直接使用hixl的场景是:你需要精细控制点对点通信的行为,比如PD分离中KV Cache的定向搬运。直接使用hccl的场景是:你需要集合通信的语义,比如训练中的梯度同步。两者不是替代关系,而是不同抽象层次上的互补。
hixl与shmem:跨步通信与共享内存
CANN生态中另一个容易与hixl混淆的通信库是shmem。两者都支持单边通信,但设计侧重点不同。
shmem的模型是共享内存:所有参与通信的NPU把一部分内存映射到一个共享地址空间,每个NPU可以像访问本地内存一样访问远端数据。这种模型适合需要频繁读写远端内存的场景,比如图计算中的顶点数据更新——每个NPU都可能随时读写任意远端顶点的数据,共享内存提供了最自然的编程接口。
hixl的模型是消息传递:通信以显式的put/get操作发起,每次操作搬运一块(或跨步的多块)数据。这种模型适合数据搬运模式可预测的场景——你知道数据从哪来、到哪去、搬多少。PD分离中的KV Cache搬运就是典型的可预测模式:Prefill产出的KV Cache地址和大小在计算前就能确定,Decode侧的消费地址也可以预先分配好。
跨步访问是hixl区别于shmem的关键能力。shmem的共享内存访问是按地址直接读写,如果数据布局不连续,应用需要自己逐个元素访问。hixl的跨步描述符让硬件自动完成跳步访问,对于KV Cache这种固定步长的不连续布局,通信效率显著更高。
选择hixl还是shmem,取决于你的通信模式:如果数据搬运路径明确且步长规则,用hixl;如果需要随机访问远端内存,用shmem。
PD分离场景中的hixl工作流
把前面的概念串起来,看hixl在PD分离推理中的完整工作流。
Prefill阶段:用户请求到达Prefill卡,模型执行前向计算,产出当前请求的KV Cache。这些KV Cache按照多头注意力的布局存储,每个头的数据在头维度内连续,头与头之间有固定间距。Prefill卡在产出KV Cache后,通过hixl的跨步put操作,将KV Cache直接写入Decode卡的预留内存区域。
HXLStridedDescriptor kv_desc;kv_desc.count=num_layers*num_heads;kv_desc.size=head_kv_bytes;kv_desc.stride=head_stride_bytes;for(intl=0;l<num_layers;l++){void*src=kv_base_prefill[l];void*dst=kv_base_decode[l];hxl_put_strided(prefill_npu,decode_npu,src,dst,&kv_desc);}这里按层循环,每层执行一次跨步put。count设为num_heads(每层的头数),stride是头间步长。一次调用就把该层所有头的KV Cache从Prefill卡搬到Decode卡,不需要逐头搬运。dst地址是Decode卡上预先分配好的KV Cache存储区,Prefill卡直接写入,Decode侧无需参与。Prefill卡在put返回后就可以继续处理下一个请求的Prefill计算,不会因为等待Decode侧而阻塞。
Decode阶段:Decode卡从自己的内存中读取已经就位的KV Cache,执行token生成计算。每生成一个token,只需要访问本地内存中的KV Cache,不需要跨卡通信。只有当新的Prefill请求的KV Cache被put过来时,Decode卡才需要把新的KV Cache纳入注意力计算的范围。
这个工作流的关键特征是:Prefill和Decode的计算完全解耦。Prefill只管产出和投递,Decode只管消费和生成。hixl的单边通信让这种解耦成为可能——双边通信要求两边同步,那Prefill就必须等Decode准备好才能发,两边的流水线又被耦合回去了。
64字节对齐的工程实践
前面提到hixl要求所有通信地址和步长64字节对齐,这在工程实践中需要特别处理。
KV Cache的内存分配通常由框架的内存池管理。框架在分配内存时,需要确保KV Cache的起始地址和头间步长都满足64字节对齐。起始地址的对齐比较容易——在内存池初始化时按64字节边界分配即可。步长的对齐需要根据模型结构计算:如果head_dim是128(float16下每头256字节),步长天然满足对齐;如果head_dim是96(float16下每头192字节),步长是192字节,不满足64字节对齐,需要在头间填充16字节使步长对齐到256字节。
这种填充会带来少量内存浪费,但相比跨步通信的性能收益,浪费是值得的。而且填充后的内存布局对Decode侧的注意力计算也有好处——对齐的访存地址让NPU的向量单元可以更高效地加载数据。
使用前后的效率对比
| 维度 | 使用传统RMA方式 | 使用hixl跨步通信原语 | 差异来源 |
|---|---|---|---|
| 数据传输粒度 | 每次传输只能操作连续内存块,非连续数据需要先打包再复制 | 支持自定义跨度和跳步,直接操作非连续地址空间 | hixl的跨步通信原语在硬件层面处理地址偏移计算 |
| 中间缓冲区 | 非连续数据传输前需要申请中间缓冲区完成数据整理 | 不需要中间缓冲区,源地址到目标地址直写 | Put语义的硬件单边访问消除了缓冲区拷贝 |
| PD分离KV交换场景 | 需要多轮RMA操作完成注意力头数据的组合传输 | 一次跨步Put即可传输完整层级的KV Cache数据 | stride参数的灵活配置适配了多层注意力头的数据布局 |
| 编程复杂度 | 开发者需要手动编排数据打包和缓冲区管理逻辑 | 只需配置src、dst、count、size、stride五个参数 | 简洁的API参数设计将地址计算和偏移管理的底层细节封装在内 |
| 多NPU同步 | 需要在每次数据传输后跨NPU同步,时间开销较大 | 跨步通信完成后即可直接访问,hixl保证数据可见性 | 单边通信语义本身就保证了数据一致性,不需要额外同步操作 |
下表对比了PD分离场景中,使用hixl前后在关键指标上的差异:
| 指标 | 使用常规拷贝通信 | 使用hixl单边零拷贝通信 |
|---|---|---|
| KV Cache搬运路径 | 应用缓冲区→通信库缓冲区→对端应用缓冲区,两次拷贝 | 源端内存→对端内存,DMA直搬,零拷贝 |
| 通信参与方 | Prefill卡和Decode卡必须同步,双边参与 | Prefill卡单边发起,Decode卡无感 |
| 不连续数据搬运方式 | 逐头发起通信,通信次数等于头数 | 跨步描述符一次发起,通信次数为一 |
| 通信启动开销 | 每次通信均有启动开销,总开销与头数成正比 | 启动开销仅一次,被所有头摊薄 |
| Prefill与Decode耦合度 | 强耦合,Prefill必须等Decode就绪 | 解耦,Prefill算完即投递,不等待 |
| 内存带宽占用 | 拷贝过程消耗额外内存带宽 | 无额外拷贝,内存带宽全部留给计算 |
| 对齐要求 | 无特殊对齐要求 | 地址和步长需64字节对齐 |
从表中可以看出,hixl在搬运路径、参与方、通信次数、耦合度、内存带宽五个维度上都有明显优势。代价是64字节对齐的约束,这在框架层面的内存分配器中很容易满足,工程代价很小。
hixl在CANN分层架构中的位置
CANN的软件栈分为五层:最上层是推理框架适配层(如MindSpore Lite),往下是计算图编译层(GE),再往下是算子库层(Ascend C),再往下是计算基础层,最底层是硬件驱动层。hixl位于计算基础层,与hccl、shmem同层。
这一层的库有两个共同特征:直接面向NPU硬件能力编程,不经过上层的图编译或算子调度;被上层库或框架间接调用,应用开发者通常不直接使用。hixl也不例外——大多数开发者通过hccl间接使用hixl的传输能力,只有在需要精细控制点对点通信时(如PD分离场景)才会直接调用hixl的API。
计算基础层的库之间也有分工:hccl管集合通信的编排,hixl管点对点单边通信的执行,shmem管共享内存的映射与访问。三者共同构成了CANN通信能力的完整基础,上层的分布式训练和分布式推理都构建在这套基础之上。
从概念到实践的映射
回顾全文,hixl的核心概念可以用一条线索串起来:单边通信解耦了通信参与方,零拷贝消除了中间搬运,跨步访问合并了不连续数据的通信次数。这三个能力叠加在一起,让PD分离场景中的KV Cache搬运从"同步、多次、有中间拷贝"变为"异步、一次、直搬"。
对于初次接触hixl的开发者,建议从三个维度建立认知:通信语义维度(put/get与send/recv的差异)、内存布局维度(连续访问与跨步访问的选择)、系统架构维度(hixl在CANN分层中的位置和与hccl/shmem的关系)。这三个维度交叉起来,就能判断自己的场景是否适合使用hixl,以及应该使用hixl的哪些能力。
仓库地址:https://atomgit.com/cann/hixl
