hixl:PD 分离背后的单边通信库,到底是个什么东西

有个架构层面的东西我一直没太搞明白,就是昇腾 NPU 的通信库。hccl、hcomm、hixl,一堆名字分不清干什么用的。

前阵子因为业务需要研究了 PD 分离(Prefill-Decode 分离),才知道 hixl 在这里面的角色。搞清楚之后发现,这个库的设计思路挺有意思的,今天把它写出来。


通信库全家桶,先搞清楚谁是谁

在说 hixl 之前,先把昇腾 NPU 生态里几个通信相关的库理一遍,不然容易混。

hccl:集合通信库,AllReduce、Broadcast、AllGather 这类操作全靠它。多卡训练必须用 hccl 做梯度同步,地位相当于 NVIDIA 的 NCCL。这个大家应该都熟。

hcomm:昇腾的通信抽象层,对上封装 hccl,对下对接硬件,相当于昇腾 NPU 上的通信中间件。如果你做多卡推理或者并行计算,基本都会打交道。

hixl:今天的主角,单边通信库。注意这个"单边"——它和 hccl/hcomm 的最大区别是,不需要对方卡配合操作,数据发送方自己就能完成传输,接收方不需要主动参与。这听起来很简单,但这个特性在 PD 分离场景里非常关键。

简单类比一下:hccl/hcomm 像打电话,必须双方同时在线才能说上话;hixl 像发快递,寄件人自己就能把东西送出去,收件人什么时候取都行。


PD 分离为什么需要 hixl

PD 分离是 LLM 推理优化里的一种重要架构,把 Prefill(输入处理)和 Decode(逐 token 生成)两个阶段拆分到不同硬件上跑。

为什么拆?因为这两个阶段的计算特性完全不同:Prefill 是 compute-bound,要处理整个输入序列,计算量大;Decode 是 memory-bound,每次只生成一个 token,访存为主。如果硬塞在同一张卡上,互相抢资源,两边都跑不快。

拆开之后有个问题:Prefill 阶段输出的 KV Cache 要传 Decode 阶段用。这个 KV Cache 的数据量可以非常大——8192 token 上下文,30+ 层,heads=32,head_dim=128,光 K 或 V 矩阵就有 4096×4096×128×2Bytes ≈ 4GB,K+V 一起接近 8GB。

这时候问题来了:Prefill 卡和 Decode 卡怎么传这个数据?

如果你用 hccl/hcomm 的集合通信,必须两边同时调用同步接口。但 Prefill 和 Decode 的执行进度天然不同步——Prefill 算完可能 Decode 还在跑上一个 token,反过来也可能。这种强耦合的通信方式在 PD 分离架构里会引入严重的等待开销。

hixl 的单边通信特性在这里就派上用场了:Prefill 卡直接把 KV Cache 写进 Decode 卡的显存,不需要 Decode 卡配合,自己就能完成传输。Decode 卡有空的时候直接读就行,不用双方同时等待。

这就是 hixl 解决的核心问题:让不同执行阶段之间的大块数据传输变成异步的、单边驱动的


零拷贝:hixl 性能高的原因

hixl 另一个重要特性是支持零拷贝数据传输。

传统的数据传输:数据从发送端显存 → 复制到通信 buffer → 通过 PCIe/CCIe 发送到接收端 → 复制到接收端显存。两次显内复制,加上一次跨硬件互联传输,开销不小。

hixl 的零拷贝:利用昇腾 NPU 的共享内存机制(shmem),数据直接在两个卡之间传递,不需要中转 buffer。具体来说,hixl 会在发送端和接收端的共享内存区域建立直接的数据通道,发送端写入本地共享内存,接收端直接从自己的共享内存读。

# hixl 零拷贝数据传输示例(简化版)
import hixl

# 初始化 hixl 上下文,指定本端和远端设备
ctx = hixl.Context(device_id=0, peer_device_id=1)

# 申请零拷贝缓冲区,大小按实际 KV Cache 尺寸定
# 这里假设 KV Cache 大小约 512MB
buf = ctx.allocate(size=512 * 1024 * 1024, flags=hixl.MEM_ZERO_COPY)

# Prefill 完成后,把 KV Cache 数据写入 hixl buffer
src_addr = kv_cache_ptr  # KV Cache 在本卡显存中的地址
bytes_to_send = kv_cache_size  # 实际数据大小
ctx.write(src_addr, buf, bytes_to_send)  # 单边写,本卡独立完成

# 远端 Decode 卡读取数据(可以在另一个时间点独立执行)
# ctx = hixl.Context(device_id=1, peer_device_id=0)
# dst_addr = decode_kv_cache_buffer
# ctx.read(buf, dst_addr, bytes_to_send)  # Decode 卡的单边读

这里的关键是 ctx.write() 这一步——这是单边操作,发送端自己调用,不需要接收端同步参与。只要传输通道建立好,数据什么时候发、什么时候收,两端各自独立调度,互不阻塞。


hixl 和 shmem 是什么关系

有人会问:昇腾 NPU 上还有个 shmem 仓库,专门做共享内存,hixl 和它是什么关系?

简单说:hixl 依赖 shmem,但比 shmem 更高层

shmem 是昇腾 NPU 的共享内存基础设施,提供了跨设备共享内存的底层能力——申请共享内存区域、建立映射、管理生命周期这些事情。hixl 在这个基础上封装了具体的数据传输协议和零拷贝逻辑。

你可以理解成:shmem 是修路的,hixl 是在路上跑的车。路是基础设施,但光有路没用,得有车才能运货。

昇腾异构计算架构里,它们的关系是这样的:

AscendCL
  └─ hixl(单边通信,高层接口)
      └─ shmem(共享内存,底层支撑)
          └─ 昇腾硬件层
hccl / hcomm(集合通信,依赖 hixl 做数据交换)

另外,昇腾还有一套 ascend-boost-comm(算子公共平台),也涉及到跨设备数据传输。它的定位和 hixl 不同:hixl 做的是原始数据搬运,ascend-boost-comm 做的是算子级别的融合和数据复用——是把"车"整合成"物流系统"的中间件。


几个踩坑经验

坑一:跨 NUMA 节点的 hixl 性能会降。 如果你的 Prefill 卡和 Decode 卡不在同一个 NUMA 节点上,hixl 的零拷贝路径可能走不通,会退化成普通拷贝。部署的时候最好确认两张卡在物理上是邻近的,否则 hixl 的性能优势会大打折扣。

坑二:buffer 生命周期要管好。 hixl 的零拷贝 buffer 申请和释放必须成对出现,不然显存泄漏。昇腾 NPU 的显存比普通 GPU 紧张得多,一个泄漏的 buffer 可能直接导致后续任务 OOM。建议用 context manager 封装:

import hixl

with hixl.Context(device_id=0, peer_device_id=1) as ctx:
    buf = ctx.allocate(size=512 * 1024 * 1024, flags=hixl.MEM_ZERO_COPY)
    # ... 传输逻辑
# buffer 自动释放,不用手动 free

坑三:hixl 的传输是非原子性的。 hixl 的单边 write 是非原子操作,如果接收端在 write 过程中读取 buffer,可能拿到不完整的数据。PD 分离架构里,Prefill 卡写完 KV Cache 之后,需要通过信号量或者独立的消息通知 Decode 卡"可以读了",不要让 Decode 卡自己轮询 buffer 的状态。


结尾

hixl 在昇腾 CANN 生态里是个偏底层的库,大多数人日常不一定直接打交道。但如果你在优化 LLM 推理性能、研究 PD 分离架构、理解昇腾 NPU 的通信机制,它是绕不开的一环。

和 hcomm(集合通信)、ascend-boost-comm(算子平台)一起,构成了昇腾 NPU 完整的通信体系。每个库定位不同,用错地方性能差别很大。

源码在 https://atomgit.com/cann/hixl

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐