1. 先用 roofline 思想判断瓶颈

GPU 的峰值 FLOPs 与 HBM 带宽都是有限的。一个算子若每搬运 1 byte 能做很多 FLOPs,可能接近计算受限;若每 byte 做的 FLOPs 少,就更可能被内存带宽限制。算术强度只是第一层近似,真实表现还受矩阵形状、并发、寄存器/共享内存占用、kernel launch、量化格式和通信等待影响。

算术强度 = 有效 FLOPs / 从主存搬运的字节 可达性能 ≤ min(硬件峰值 FLOPs, 算术强度 × HBM 带宽)

prefill 中的大 GEMM 往往有较高算术强度;decode 中的“小 batch × 大权重”更常受权重读取和 KV 读取限制。不能拿训练吞吐的 roofline 结论直接解释在线生成。

2. GEMM:Transformer 最核心的计算块

线性层本质是矩阵乘 C = AB。Q/K/V 投影、输出投影、MLP 的 up/down 投影、logits head 都是 GEMM 或其批量变种。NVIDIA Tensor Core 等专用矩阵单元在 BF16、FP8、FP4 等格式下可以提供很高吞吐,但必须有合适的 tile 形状和足够的工作并发。

实现层面要关心 M,N,K:其中一维对应 token batch,另一维是隐藏/中间维,最后一维是归约维。长 prefill 让 token 维变大,通常提高 GEMM 效率;单 token decode 的 token 维太小,往往需要 GEMV-like kernel、分组 GEMM、持久化 kernel 或把多个 sequence 拼在一起。

3. 量化不是只把权重“变小”

权重量化降低 HBM 流量和容量,占用常从 BF16 的约 2 bytes/参数下降到 FP8、INT8 或 4-bit 级别。但每种格式都带 scale、block/group metadata、反量化路径和精度误差。若 kernel 不支持直接在 Tensor Core 上高效消费该格式,反量化临时张量会抵消好处。

KV 量化、activation quantization 与 weight-only quantization 是三个不同问题。KV 量化直接影响长上下文 decode 的读取字节;权重量化主要影响模型常驻与读权重带宽;activation 量化还会影响数值稳定性与通信。部署报告必须逐项说明,而不是写“模型采用 INT4”。

4. Attention 的两类 kernel 路径

prefill attention 要让许多 query token 对同一段 KV 做乘加。FlashAttention 仍然计算精确 attention,只是通过 tile 和 online softmax 在 SRAM/寄存器内完成分块,避免把巨大的 score 矩阵反复写入 HBM。收益主要来自更少的 I/O,而非减少 attention 的数学项。

decode attention 通常是 Q 长度很短、KV 长度很长。所谓 paged attention 会按 KV block 遍历非连续物理页,避免因 sequence 长度不同造成预分配浪费;其 kernel 要处理 page table、不同长度、GQA head 映射和量化 KV。此处常见瓶颈是读 KV 的有效带宽,而不是 softmax 的 FLOPs。

FlashAttention 原论文中 HBM、SRAM 与 attention tile 的数据路径
原始/论文结构图:FlashAttention, Figure 1。左侧给出 HBM 与片上 SRAM 的层级,中间以 K/V 外层、Q 内层 tile 展示如何避免把完整 attention score matrix 写回 HBM;右侧是该论文条件下的实现对比。它解释的是 exact attention 的 I/O 重排,不保证所有 GPU、精度、序列长度或后续 kernel 都有相同加速比。

5. Kernel fusion:减少看不见的读写

一个逻辑层若分别执行 RMSNorm、QKV GEMM、RoPE、attention、输出投影、残差,每一步都会读写中间张量、启动 kernel、同步流。融合将相邻、不会改变依赖关系的操作放进同一个或少量 kernel,使中间值留在寄存器或共享内存。它通常改善 decode 的低 batch 延迟,也能减少 HBM 流量。

融合有边界:过大的 kernel 可能因寄存器压力降低 occupancy;不同模型 shape、精度和硬件需要不同实现;调试和数值验证也更复杂。评估时要固定模型、输入长度、batch、GPU、精度和编译版本,不能把 fusion 的收益泛化成一切 workload 的常数。

6. RMSNorm、RoPE 与 elementwise 链条为何仍重要

RMSNorm、SiLU、残差加法、RoPE 和采样的单次 FLOPs 很少,但它们是每层、每 token 都要做的 elementwise 操作。对 decode 而言,大 GEMM 的并行度下降后,这些小算子会暴露出 launch 及内存访问开销。RoPE 还受位置、缩放策略和长上下文配置影响;它不能随意在 prefix cache 与不同配置之间混用。

常见实现会融合 norm 与线性输入准备,把 RoPE 放入 Q/K 写入路径,并将残差或激活与相邻操作合并。单项收益看似不大,但大量短 kernel 合并后,decode 路径会少掉可观的启动和读写开销。

7. Logits 与采样:在小 batch 上也会挡住首 token

末层 hidden state 需要投影到 vocabulary logits。词表较大时,这本身也是一次矩阵乘,之后还可能执行 temperature、top-k、top-p、repetition penalty、grammar/JSON mask、随机采样或 argmax。低并发、严格低延迟的服务通常把这些操作留在 GPU 上,避免往返 CPU;约束解码、工具 schema 和安全过滤则会增加额外路径。

采样的正确性和吞吐同等重要。top-p 需要处理排序或选择,随机数生成要可复现,grammar mask 要与 tokenizer 状态同步。任何“每 token 延迟”图表若没有说明采样、约束解码是否开启,就不是完整端到端指标。

8. MoE kernel 与专家的 grouped GEMM

router 给 token 分配专家后,同一专家获得的 token 数各不相同。若为每个专家单独发一个 GEMM,会产生大量小 kernel;grouped GEMM 把多个形状相近的专家矩阵乘在一次调用中组织起来,提高 launch 效率。token packing 与排序让同一专家的 token 连续,combine 阶段再恢复原 sequence 位置。

但是 router 的负载不均衡仍会造成 straggler:一个专家接到过多 token,其他 GPU 已完成却在等待。capacity factor、动态专家复制、token 重路由和 batch 内平衡都是性能与模型语义之间的取舍。不能只看平均 token/expert。

9. All-to-all:MoE 的算子图会跨越网络

当专家位于不同 GPU,dispatch 将 token activation 按目标 GPU 交换,expert GEMM 完成后 combine 再交换回来。collective 的有效吞吐不只取决于链路标称带宽:小消息、不同 rank 的数据量、拓扑、拥塞控制和同步等待都影响尾部。EP group 跨一个 NVSwitch 域与跨 RoCE fabric 的行为可以完全不同。

MoE 端到端延迟 = router/pack + dispatch all-to-all + expert GEMM + combine all-to-all + 等待最慢 rank。只优化 expert GEMM 不能保证 P99 改善。

07 通信与拓扑将把这些 collective 置于 NVLink、NVSwitch、RDMA 以及数据中心网络的实际约束中。

DeepSpeed-MoE 论文中分层 all-to-all 的 token 重排和节点内外交换
原始/论文结构图:DeepSpeed-MoE, Figure 8。图把 token 先作 local transform,再执行节点内 all-to-all、并发节点间 all-to-all,最后恢复局部布局;它说明为什么 MoE dispatch 不是单一步骤。该图是 2022 的分层 all-to-all 基线,不能把其通信库、GPU 数量或带宽结论直接替代 2026 DeepEP/特定 fabric 的实现。

10. 流、事件与计算通信重叠

GPU runtime 可使用多个 stream:在某一块 token 的专家结果返回时,计算另一块的局部算子;或在一个 layer 的部分通信未完成时准备下一步可独立的工作。重叠要有真实的数据依赖空隙,否则只会产生竞争。还要预留足够 HBM 作为双缓冲/通信 buffer,并避免过多 stream 导致调度开销。

报告“通信被隐藏”时,应同时给出未重叠基线、不同 batch 下的 profiling,以及 P50/P99;在低 batch 或长尾专家场景,理论可重叠的比例可能消失。

11. 如何做一次有意义的 profiling

端到端时间线应标出 tokenize、排队、prefill、每轮 decode、KV transfer、采样和网络等待。Nsight Systems 等工具可检查 CPU、GPU、NCCL/RDMA 流上的空洞,kernel profile 则用于核对实际矩阵 shape、memory throughput、Tensor Core 利用率与 occupancy。MoE 还要按 expert rank 分桶,否则均值会掩盖热点。

实验矩阵至少覆盖短/长 prompt、短/长输出、小/大并发、prefix hit/miss、不同 KV 精度、节点内/跨节点 EP。先确定瓶颈类型,才选择 FlashAttention、量化、fusion、专家复制或网络调参;不要把所有优化同时打开后无法解释因果。

12. 算子选择怎样上升为系统配置

若 prefill 的 GEMM/attention 是主瓶颈,增加 prefill token batch、选择更适合的算子或将其独立到 P worker 往往有效。若 decode 被 KV 读取压住,应先计算 bytes/token,考虑 GQA/MLA、KV 格式、分页与 cache locality。若 MoE all-to-all 是瓶颈,要检查 EP group、拓扑、专家热度和通信重叠。若 GPU 时间线有大片空洞,则可能是 scheduler 或 CPU 前处理问题。

下一页 04 KV cache 解释 why decode 经常是存储问题;后面的 05 Runtime 解释如何把这些不同算子形状混合进同一服务。

13. 来源与边界