算子与编译器:Profiling、调试与数值稳定
很多 kernel 优化最后不是败在“不会写更快代码”,而是败在三件事上:没有正确测量,优化方向从一开始就错了;没有可靠验证,性能上去了但结果悄悄错了;没有处理数值稳定性,速度和精度之间出现不可接受的偏差。
成熟算子工程不是“会写 kernel”,而是建立从热点识别、性能定位、正确性验证到数值验收的完整闭环。
Profiling 必须先于优化。否则你可能把时间花在看起来重要、实际不在热路径的 kernel 上。正确流程是先定位瓶颈,再改实现,再验证数值和端到端收益。
家里水压低,不应该先随便换水龙头。要先看是总阀、水管堵塞、热水器还是某个弯头漏水。Profiling 就是这套排查流程,避免优化错位置。
Profiling 先于优化
AI 系统里的直觉经常误导人:
- 以为 GEMM 是瓶颈,实际热点是 layout transform;
- 以为 decode 慢是模型太大,实际是 kernel launch 太碎;
- 以为量化没提速是算法问题,实际是 dequant kernel 带宽受限;
- 以为 occupancy 太低,实际卡在 shared memory bank conflict。
任何 serious optimization 都应从 profile 开始,而不是从改代码开始。优化前要明确:问题发生在系统调度、单 kernel、内存访问、微架构利用率,还是数值回退路径。
三层 Profiling 视角
| 层级 | 关注点 | 常用工具 |
|---|---|---|
| 系统级 | 请求流、prefill/decode、通信、I/O、CPU/GPU 重叠 | Nsight Systems、serving trace、PyTorch profiler timeline |
| Kernel 级 | 单个 kernel 时间、launch 次数、shape 分布、调用路径 | PyTorch Profiler、框架 trace、自定义日志 |
| 微架构级 | 带宽、occupancy、warp stall、Tensor Core、cache 行为 | Nsight Compute、硬件计数器 |
只看系统级,知道哪段慢但不知道为什么;只看 kernel 级,知道哪个 kernel 慢但不知道是否被调度放大;只看微架构级,容易局部最优。有效工作流是在三层之间来回切换。
Microbenchmark 要可信
微基准很容易夸大收益。较稳妥的原则是:
- 充分 warmup,排除 JIT、memory pool、cache 和 graph capture 前置成本;
- 固定 shape、dtype、layout 和输入分布;
- 区分单 kernel 时间与端到端时间;
- 覆盖真实 shape 分布,而不是只测最漂亮的整齐形状;
- 报告平均值、p95/p99 和长尾 shape;
- 明确是否包含编译、autotune、cache miss 和 fallback。
真实系统中 shape 往往不固定。一个 kernel 在某个整齐 shape 下极快,到了 batch、sequence、head dim、量化 group 变化时可能明显退化。优化结论必须按真实 workload 分桶报告。
关键指标怎么读
| 指标类 | 看什么 | 常见解释 |
|---|---|---|
| 时间 | latency、throughput、launch 次数、TTFT、TPOT、step time | 判断端到端是否真的变快 |
| 资源 | SM 利用率、Tensor Core、DRAM throughput、shared memory、register、occupancy | 判断算力或带宽是否被用起来 |
| 结构 | warp stall、branch divergence、cache hit/miss、load/store efficiency | 判断为什么没跑满 |
| 数值 | 误差、溢出、归约差异、低精度 scale | 判断是否安全可替换 |
Roofline 是很好的先验判断器。若 kernel 算术强度低且 DRAM throughput 已接近上限,继续做算术优化收益很小;若理论上应是 Tensor Core 热点但 Tensor Core 利用率很低,说明 tile、layout 或 dtype 路径可能有问题。
正确性验证
任何优化都必须先通过正确性验证。基础做法包括:
- 与参考实现逐元素对比;
- 覆盖多种 shape、dtype、layout 和边界 mask;
- 测随机输入、极值输入、全零、全同值、非连续 stride;
- 测训练与推理路径;
- 测多步迭代误差是否放大;
- 固定 seed 和输入,保证可复现。
“看起来差不多”远远不够。一个归一化 kernel 某些边界块少算一个元素,单次误差可能很小,多层叠加后会明显偏航。低精度、softmax、attention、norm、optimizer kernel 都应有更严格的误差阈值和边界样本。
数值稳定性是一等公民
AI 算子不仅要算得快,还要在低精度、长序列和大动态范围下稳定。常见问题包括:
- softmax overflow / underflow;
- 方差计算误差;
- FP16/FP8 累加损失;
- quant scale 过大或过小;
- 原子加法和归约顺序导致非确定性;
- fused kernel 改变了高精度边界。
Softmax 通常写成:
减去最大值是为了避免指数爆炸。FlashAttention 这类分块实现还要维护在线最大值与归一化项,既减少 I/O,又保持数值精确。LayerNorm、RMSNorm、variance 计算则需要关注 Welford、FP32 accumulation 或其他稳定统计方法。
低精度验收
不同低精度类型风险不同:
| 类型 | 主要风险 | 验收重点 |
|---|---|---|
| FP16 | 动态范围窄,易 overflow/underflow | loss scale、FP32 accumulate、极值输入 |
| BF16 | 范围大但尾数少 | 长期误差、归约和 optimizer 路径 |
| FP8 | 强依赖 scaling 和校准 | amax、scale 粒度、溢出/下溢、敏感层 |
| INT8/INT4 | 依赖量化尺度和校准数据 | scale、zero point、outlier、累加精度 |
低精度 kernel 不能只看单层误差。要看长序列、多层堆叠、训练反向、端到端任务指标和异常 bucket。尤其是 FP8/INT 相关 kernel,scale 的读取、更新、广播和融合边界都可能成为真实 bug 来源。
调试闭环
一个实用闭环是:
- 用系统 trace 找到真实热点;
- 用 kernel profile 判断瓶颈类型;
- 用 microbenchmark 验证单点优化;
- 用参考实现做正确性和数值对齐;
- 用真实 workload 分桶验证端到端收益;
- 把 shape、dtype、误差阈值和性能基线写入回归测试;
- 记录硬件、驱动、编译器、框架和 kernel 版本。
算子优化的基本纪律是:没有 profiling,不谈优化;没有正确性,不谈性能;没有端到端验证,不谈上线收益。
- Title: 算子与编译器:Profiling、调试与数值稳定
- Author: Charles
- Created at : 2025-09-19 09:00:00
- Updated at : 2025-09-19 09:00:00
- Link: https://charles2530.github.io/2025/09/19/ai-files-operators-profiling-debugging-and-numerical-stability/
- License: This work is licensed under CC BY-NC-SA 4.0.