1709 字
9 分钟
XDNA 空间数据流 LLM Decode

AMD XDNA NPU 将一层 Transformer 组织成静态数据流:权重流式进入,激活尽量留在片上,DMA 搬运与计算重叠。一次提交运行完整层,可以减少 batch=1 decode 的调度和内存访问开销。

背景:为什么 Decode 慢#

LLM 推理的 decode 阶段每次只生成一个 token,每一步都要用大量权重处理一个较小的隐藏向量:

  • 输入/输出 hidden 向量约为 KB 级
  • 单层权重可能达到百 MB 级(例如 Qwen3-8B Q4 量化后一层约 115 MB)
  • 每个权重只服务一个 token,算术强度低,典型瓶颈是内存带宽而不是峰值算力

GPU 在 batch=1 decode 场景下会遇到三个结构性问题:

  1. 中间结果反复写回 HBM:一层 Transformer 被拆成 RMSNorm、QKV projection、Attention、O projection、MLP 等多个 kernel,中间激活反复写回/读出。
  2. Kernel launch 延迟叠加:小 batch 下单个 kernel 计算时间短,固定提交开销占比上升。
  3. 并行资源利用不足:batch=1 只有一条当前 token 向量,SIMT 执行资源难以充分利用。

核心矛盾是:GPU 擅长高并行度、高算术强度任务;decode 更需要高带宽利用率和少中间搬运。

XDNA 的核心模型:空间数据流架构#

XDNA NPU 采用空间数据流架构(Spatial Dataflow Architecture):

  • 底层 Shim 负责主存接口和数据进出
  • 中间 Memtile 负责片上缓存、DMA 分发与汇聚
  • 上层 Compute Tile 运行编译好的小程序

它更像固定流水线工厂:每个 tile 是一个工位,数据沿编译时确定的物理连线流动,中间尽量不回主存。

典型 Tile 角色#

角色职责复用方式
main tiles量化矩阵向量乘,负责 Q/K/V/O/Up/Gate/Down 投影同一组 tile 按时间轮转跑多个投影
vector tile持有完整 hidden 向量,做 RMSNorm 与 residual add整层内多阶段复用
post tileQ/K head-wise norm 与 RoPEQKV 阶段使用
attention tiles分块 score、online softmax、加权求和Attention 阶段使用
swiglu tileSiLU(gate) × up 激活MLP 阶段使用
memtile bridge/hub权重分发、compact 汇聚、KV 中转由 DMA/BD 配置驱动

关键硬件约束#

XDNA 编程除了编写 kernel,还需要根据硬件约束显式设计数据通路:

约束编程影响
Compute Tile 只有约 64 KB 本地内存不能加载完整权重,必须流式处理
Tile 之间内存不可见不能共享内存,必须显式通过 stream/DMA 搬运
DMA 与 Compute 独立可以 ping-pong 双缓冲,让搬运与计算重叠
同步靠硬件 lock/计数器数据就绪后自动触发,无需软件调度每一步
Stream/路由编译时确定运行时不能临时改拓扑
BD bank、packet ID、DMA channel 有物理规则配错可能静默死锁或数据送错

全层融合:一次提交跑完整层#

XDNA 更适合把一层 Transformer 组织为一次静态数据流提交:

Host hidden → Shim → vector RMSNorm
→ main16 Q/K/V projection
→ post Q/K norm + RoPE + KV writeback
→ attention tiles online softmax
→ main16 O projection → vector residual add
→ vector post RMSNorm
→ main16 Up/Gate projection → swiglu
→ main16 Down projection → vector residual add
→ Shim → Host output

理想路径中,中间激活(Q/K/V/O/attention/up/gate/down 等)尽量在片上流转,不为每个 operator 单独回主存,也不为每个 operator 单独承担 host 调度开销。

五个编程思维转换#

1. 权重必须“流过”,不能“加载”#

单个投影权重远大于 tile 本地内存。权重被切成小 chunk,通过 DMA 按顺序流入 tile;tile 在线反量化、乘加、累加,用完即覆盖。关键是不要把完整矩阵解压或装入 tile。

2. DMA 可独立循环运行#

DMA 由 Buffer Descriptor(BD)驱动。BD 描述地址、长度、启动 lock、完成 release、下一张 BD。多个 BD 可组成 ring,让数据搬运在硬件中自动循环执行,host 不需要逐 chunk 介入。

3. Attention 不物化完整 score 矩阵#

Attention 可按小 block 扫 KV cache。每个 block 产生包含局部权重、block max、block sum 的 carrier,下游用 online softmax 合并。这样只保存 running max、running sum、output accumulator 等小状态,而不是完整 score matrix。

4. 近似数学函数要可控#

Tile 处理器不适合直接搬 CPU/GPU 的通用数学库。RMSNorm 可用 Newton-Raphson 近似 rsqrt,softmax 的 exp 可用 e^x = 2^(x/ln2) + 多项式近似,SwiGLU 的 sigmoid 可用查表插值。目标是让 kernel 小、快、误差可控。

5. Packet 路由与 channel 所有权是物理约束#

多路数据可通过 packet ID 共享物理通路,但同一路由域内 packet ID 不能冲突。DMA channel、BD ring、lock pair 也必须有明确 owner;错误配置可能不会编译时报错,而是在运行时 timeout 或产生不易察觉的数值错误。

Record ABI:让并行 tile 输出可汇聚#

多个 main tile 并行计算不同输出行时,可定义稳定 record 格式:header 编码 phase/block/group/row,payload 存若干 bf16 输出。Memtile 再按 column compact → global compact 汇聚成完整 output block。

这种 ABI 有助于定位错误:

  • header 对、payload 错 → 多半是计算 kernel 问题
  • header 错 → 多半是路由或汇聚配置问题

Runtime 参数 patch#

NPU 编译成本高,不能每个 token 重编译。常见做法是编译最大容量拓扑,运行时只 patch 当前 token、block count、KV 写入位置等少量描述符/RTP。

关键陷阱:RTP 必须在 core 读取前写好。可用 runtime-start lock 让 core 启动后先阻塞,host 写完 RTP 后再 release,避免很多 token 都表现得像 token 0。

Debug 方法:三级验证#

XDNA 程序的错误通常表现为 timeout 或数值不对,缺少传统 stack trace。更实用的是分层验证:

  1. 结构检查:编译前检查 dataflow 图、record 常量、packet ID、channel owner、BD/lock 关系。
  2. 编译检查:确认硬件描述可生成二进制,但不要把“可编译”误认为“可运行正确”。
  3. 真机运行 + CPU reference:用 CPU oracle 模拟量化、bf16 舍入、RoPE、online attention,逐 lane 对比 cache writeback 与最终 output。

与 NPU-CPU 异构推理的关系#

这篇材料补充了 NPU-CPU 异构 LLM 推理 的 decode 侧机制:

  • 异构推理页面关注 CPU/NPU 如何按 Prefill/Decode 分工
  • XDNA 空间数据流关注 NPU 内部如何把 decode 变成低调度开销、低中间搬运的数据流管线

两者共同指向同一原则:端侧 LLM 推理不能只看峰值算力,必须围绕阶段特征、内存带宽、调度粒度和数据搬运路径设计。

核心洞察#

XDNA 将拓扑、路由、同步和数据复用尽量安排在编译期,以减少运行时工作。代价是灵活性降低,收益是更少的主存往返、更低的调度开销和更可预测的数据流。

参考#

XDNA 空间数据流 LLM Decode
https://blog.lpkt.cn/posts/concepts/xdna-spatial-dataflow-llm-decode/
作者
lollipopkit
发布于
2026-07-01
许可协议
CC BY-NC-SA 4.0