L2.6 算子 Shape 泛化与性能分析

三维坐标　layer: L2（数据与算子）｜level: Senior｜pillar: 编程与编译 + 训推框架

上一篇我们把算子编译到了硬件上，本文要回答一个更工程化的问题：同一个算子在千变万化的输入 shape（变长序列、不同 batch）下，如何既不浪费算力、又能被精确度量与归因？ 你将建立两套核心武器——「Shape 泛化的工程权衡」与「Roofline + Profiler 的性能定位方法论」。

学习目标

前置知识：读过 L0 全层（尤其 L0.2 三大物理墙：算力墙 / 显存墙、Roofline 拐点的直觉）；读过 L1.1（知道 GPU Tensor Core「精度换吞吐」、warp 调度与 occupancy）；读过本 L2 层前序（L2.1~L2.5：算子如何被编译成 kernel、什么是 GEMM / Attention 算子）；会写基础 PyTorch、能跑通一段推理脚本即可，无需 CUDA 编程经验。
学完产出：① 能写出静态 Padding 的有效算力利用率公式（区分 FFN 线性项与 Attention 平方项），并据此量化「一条长序列拖垮整个 batch」浪费了多少 FLOPs；② 能讲清 Padding / Bucket / Padding-Free 三档泛化策略各自的收益与代价，并给出一棵随长度方差、QPS、团队工程能力变化的决策树；③ 能用 Roofline 模型把任意算子钉在斜坡区（Memory-Bound）或平顶区（Compute-Bound），手算 ridge point 并解释为什么 LLM decode 几乎全是 Memory-Bound；④ 能照「系统级 timeline → top-k kernel → Roofline 归因 → 单 kernel 深挖」这条定位漏斗，排除「GPU 利用率低就是 kernel 烂」之类的误判；⑤ 亲手用 torch.profiler 抓一段推理 trace，导出 chrome trace 并对 top-3 kernel 做 Memory/Compute-Bound 归因。
阅读姿势：盯住一条主线——「先量化，再决策；先定位，再优化」。无论是选 Shape 泛化策略，还是判 kernel 瓶颈，本文的全部方法论都在反对「凭直觉拍脑袋调优」：用一个利用率公式量化浪费、用 Roofline 一眼判定天花板、用 Profiler 漏斗逐层收敛——把性能工程从玄学变成可复算的工程。

背景与现状

在大模型推理与训练里，输入 shape 几乎从不固定：每个请求的序列长度（seq_len）不同、每个 batch 的样本数不同。而 GPU/NPU 的 kernel 往往针对「固定 shape」编译最优，于是诞生了一对永恒的矛盾——算力浪费（为了对齐 shape 而做无用计算）与 工程复杂度（为了不浪费而引入变长处理逻辑）。

围绕这对矛盾，业界的演进可概括为三个阶段：

静态 Padding 时代：把一个 batch 内所有序列补齐到最大长度（甚至补到模型支持的 max_seq_len），实现简单、kernel 规整，但浪费的算力随长度方差线性增长——一个 [8, 2048] 的 batch 里若真实有效 token 只占 30%，就有 70% 的 FLOPs 在算 padding。
动态 Bucket 时代：把相近长度的序列**分桶（bucketing）**后分别 padding，把浪费控制在桶内方差。代价是调度复杂、桶边界要调参、小桶利用率低。
Padding-Free / 变长时代：以 FlashAttention 的 varlen 接口、vLLM 的 continuous batching、FlashDecoding 为代表，用 cu_seqlens（累积序列长度）把多条变长序列拼成一条连续内存，彻底消灭 padding 算力浪费。代价是算子实现门槛极高、对 kernel 与内存布局有强约束。

业界信号：HuggingFace Transformers 在 4.x 引入 attn_implementation="flash_attention_2" 并支持 padding-free 训练；vLLM 把 continuous batching 做成默认调度策略。这说明「消灭 padding 浪费」已从论文走向生产标配——Senior 工程师的价值，正体现在能量化这笔浪费并选择正确的泛化策略。

而要「量化浪费」「定位瓶颈」，光靠直觉不行，必须有可观测工具与理论模型：

工具 / 模型	解决什么问题	粒度
Nsight Systems	系统级 timeline：CPU/GPU 时间线、kernel 启动间隙、H2D/D2H 拷贝	系统级
Nsight Compute	单 kernel 深挖：SM 占用率、内存吞吐、warp stall 原因	kernel 级
Ascend Insight	昇腾 NPU 的算子级 profiling（对标 Nsight）	算子 / kernel 级
PyTorch Profiler	框架内一站式：aten 算子 → 底层 kernel 的映射 + chrome trace	算子 → kernel
Roofline 模型	理论判定 kernel 是 Memory-Bound 还是 Compute-Bound	理论

原理与架构

2.1 静态 Padding 的算力浪费——可量化

设一个 batch 有 B 条序列，真实长度为 L_i，padding 后统一到 L_max = max(L_i)。对 Transformer，Attention 的计算量约正比于序列长度的平方，FFN 正比于序列长度。有效算力占比可近似为：

利用率 ≈ Σ L_i / (B · L_max)     # FFN 主导（线性）
利用率 ≈ Σ L_i² / (B · L_max²)   # Attention 主导（平方，浪费更狠）

当长度方差大时，L_max 被极端长序列拉高，分母膨胀，利用率骤降——这正是「一条长序列拖垮整个 batch」的根因。Bucket 的本质就是把方差大的整体拆成方差小的子集，Padding-Free 则直接令分母 B·L_max → Σ L_i，利用率回归 100%。

工程权衡铁律：Padding-Free 不是「免费午餐」。它要求算子原生支持 cu_seqlens、要求 KV Cache 按 token 而非按样本组织（PagedAttention）、要求 loss/mask 逻辑同步改写。在「序列长度方差小」的场景，静态 padding 的简单性反而更划算——先 profile 量化浪费，再决定是否上变长。

2.2 Roofline 模型——判定 Memory-Bound vs Compute-Bound

光知道某个 kernel「慢」还不够，要知道它「为什么慢、优化天花板在哪」。Roofline 模型 用两个硬件常数和一个算子属性，把任意 kernel 钉在一张图上：

峰值算力 π（FLOPS）：硬件每秒最多做多少次浮点运算。
峰值带宽 β（Bytes/s）：硬件每秒最多从显存搬多少字节。
算术强度 I（Arithmetic Intensity）= FLOPs / Bytes：算子每搬运 1 字节数据，能做多少次浮点运算。这是算子自身的属性，与硬件无关。

可达性能上界为：

可达 FLOPS = min( π ,  β × I )

二者相等处的算术强度叫 ridge point（脊点 / 拐点）：I_ridge = π / β。

一个具体算例：以 A100（FP16 Tensor Core 峰值 ≈ 312 TFLOPS，HBM2e 带宽 ≈ 2.0 TB/s）为例，I_ridge = 312e12 / 2.0e12 ≈ 156 FLOPs/Byte。

大矩阵乘 [4096,4096]×[4096,4096]：FLOPs ≈ 2·4096³ ≈ 1.37e11，读写字节 ≈ 3·4096²·2 ≈ 1.0e8，I ≈ 1370 FLOPs/Byte ≫ 156 → Compute-Bound，应吃满 Tensor Core。
LayerNorm / Softmax / 残差加：每个元素只做几次运算就要读写一遍，I 通常 < 5 → Memory-Bound，瓶颈在带宽，优化靠融合（kernel fusion）减少显存往返，而非堆算力。

这里的关键在于：LLM 推理的 decode 阶段（逐 token 生成）几乎全是 Memory-Bound——batch=1 时 GEMM 退化成 GEMV，算术强度极低，瓶颈是把权重从 HBM 搬出来。这就是为什么推理优化的核心是显存带宽与权重量化（INT8/INT4 减少搬运字节），而非提升峰值算力。Roofline 一眼看穿这件事。

2.3 工具分工：从系统到 kernel 的定位漏斗

这条漏斗的纪律：永远从系统级 timeline 入手（先看有没有 GPU 空跑、CPU launch bound、拷贝瓶颈），再钻到具体 kernel，最后用 Roofline 判定方向、用 Nsight Compute 验证细节。跳过系统级直接调单 kernel，常常是「优化了一个本来就不在关键路径上的算子」。

动手实践：极简代码实操

实验目标：用 torch.profiler 抓一段推理的 timeline，导出 chrome trace 可视化，按设备耗时排序定位 top-3 kernel，并用 Roofline 直觉对每个 kernel 做 Memory/Compute-Bound 归因。产出物：一张 trace.json（可拖进 chrome://tracing / edge://tracing 或 Perfetto 查看）+ top-3 kernel 的归因表。

3.1 环境准备

# 推荐 Python 3.11；用 uv 或 venv 隔离
python3 -m venv .venv && source .venv/bin/activate
# 有 NVIDIA GPU：装 CUDA 版（自动选）
pip install torch torchvision
# 无 GPU（Mac/CPU）：装 CPU 版，实验同样跑得通
# pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

3.2 代码：抓一段推理 timeline 并定位 top-3 kernel

import torch
import torch.nn as nn
from torch.profiler import profile, ProfilerActivity, schedule

dev = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[device] {dev}")

# 一个小 Transformer Encoder 作为推理负载（也可换成 torchvision.models.resnet18）
model = nn.TransformerEncoder(
    nn.TransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048,
                               batch_first=True),
    num_layers=4,
).to(dev).eval()

x = torch.randn(8, 256, 512, device=dev)  # [batch, seq_len, d_model]

# 预热：触发 cuBLAS handle / kernel JIT，避免首次开销污染测量
with torch.inference_mode():
    for _ in range(5):
        _ = model(x)
    if dev == "cuda":
        torch.cuda.synchronize()

# 有 GPU 抓 CUDA 活动，无 GPU 抓 CPU 活动（aten 算子）
acts = [ProfilerActivity.CPU] + ([ProfilerActivity.CUDA] if dev == "cuda" else [])

with torch.inference_mode():
    with profile(
        activities=acts,
        record_shapes=True,     # 记录每个算子的输入 shape，便于归因
        with_stack=False,
        profile_memory=True,    # 记录显存分配，辅助判 Memory-Bound
    ) as prof:
        for _ in range(10):
            _ = model(x)
        if dev == "cuda":
            torch.cuda.synchronize()

# ① 按设备耗时排序，打印 top-k 算子表
sort_key = "cuda_time_total" if dev == "cuda" else "cpu_time_total"
print(prof.key_averages(group_by_input_shape=True).table(
    sort_by=sort_key, row_limit=10))

# ② 导出 chrome trace —— 拖进 chrome://tracing 或 https://ui.perfetto.dev 看 timeline
prof.export_chrome_trace("trace.json")
print("[ok] 已导出 trace.json，可用 Perfetto / chrome://tracing 打开看 kernel timeline")

3.3 运行与观察

python profile_infer.py

打开 trace.json（拖进 chrome://tracing 或 Perfetto UI），结合终端打印的算子表，按设备耗时取 top-3 kernel 并归因：

Top kernel（示例，实际名因 GPU 架构而异）	算子来源	Roofline 归因	优化方向
`ampere_gemm_` / `cutlass_*`	FFN / QKV 投影的大 GEMM	Compute-Bound（I 高，吃 Tensor Core）	低精度（FP16/BF16）、提升占用率
`softmax` / `layer_norm`	Attention softmax、LayerNorm	Memory-Bound（I 低，访存型）	算子融合（FlashAttention 已融合）、减少往返
`elementwise_kernel` / `add` / 拷贝	残差加、激活、H2D/D2H	Memory-Bound	融合进相邻算子、避免冗余拷贝

NVIDIA GPU 路径：表格里 Self CUDA % 最高的几行即关键 kernel；在 timeline 上能看到 GEMM 是连续大块（compute 占满），而 softmax/norm 是密集小块（典型访存型）。需要深挖单 kernel 的 SM 占用率与访存效率时，用 Nsight Compute（ncu）；想看系统级 GPU 空隙与 launch 开销，用 Nsight Systems（nsys）。
Mac / CPU 替代方案：无 GPU 时 activities 只含 ProfilerActivity.CPU，排序键改为 cpu_time_total，你会看到 aten::addmm、aten::native_layer_norm、aten::_softmax 等 aten 算子及其耗时占比。方法论完全一致——定位 top-3、判断访存型还是计算型、找融合机会。trace.json 同样可在 Perfetto 中查看 CPU 算子 timeline。
昇腾 NPU 路径：把上述「定位 top-3 + 归因」流程平移到 Ascend Insight，它提供等价的算子级 timeline 与 kernel 耗时分解。

踩坑预警 (Gotchas)

不预热 / 不 synchronize = 测了个寂寞：首次调用含 cuBLAS handle 初始化与 JIT；CUDA kernel 异步下发，不 torch.cuda.synchronize() 会把 kernel 耗时算成接近 0。务必先预热、计时段末尾同步。
cpu_time_total vs cuda_time_total 搞反：GPU 路径若按 cpu_time_total 排序，看到的全是 launch/调度开销，会误判瓶颈。设备与排序键必须对齐。
只跑 1 个 iteration 噪声大：profiler 有固定开销，单次迭代里它会占比虚高。循环 10 次取平均，或用 schedule(wait,warmup,active) 跳过预热段。
with_stack=True 拖慢且 trace 巨大：抓调用栈开销不小，定位 kernel 阶段先关掉，需要溯源到 Python 行号时再开。
chrome trace 打不开：文件过大时 chrome://tracing 会卡，改用 Perfetto UI（流式加载）更稳。
把 padding 浪费当成 kernel 慢：若 batch 内做了静态 padding，GEMM 的「绝对耗时高」可能是因为在算 padding，而非 kernel 本身低效——先核对 record_shapes 里的实际 shape 与有效 token 占比。

深入思考

下面三题每题先给题干，再用 <details> 折叠一份图文并茂的参考答案。建议先合上答案自己想 3 分钟，再展开对照。

思考题 1：泛化策略选型的决策树

你负责一个对话推理服务，线上请求序列长度从 16 到 4096 不等、方差极大。结合 2.1 节的利用率公式与三档泛化策略，给出你的决策树：在什么 QPS / 长度分布 / 团队工程能力下选「静态 Padding / Bucket / Padding-Free」哪一档？如何用一次 profile + 利用率公式量化「上 Padding-Free 能省多少算力」来支撑这个决策？

展开参考答案（含三档选型决策树图 + 算一遍）

结论：选型不是「越先进越好」，而是用 2.1 节的利用率公式先算出「当前浪费了多少 FLOPs」，再拿这笔浪费去对冲 Padding-Free 的工程复杂度——浪费大且长期跑、团队 hold 得住变长 kernel 才上 Padding-Free，否则 Bucket 往往是性价比最高的折中。

用具体数字算一遍（量化「上 Padding-Free 能省多少」，承接 2.1 节公式）：

抓一次真实流量的 profile，统计一个典型 batch 的真实长度，假设 B=8、长度为 [120, 2048, 96, 310, 64, 180, 220, 90]，则 L_max = 2048，Σ L_i = 3128。
FFN（线性项）利用率 ≈ Σ L_i / (B · L_max) = 3128 / (8 × 2048) ≈ 19.1%——即 FFN 有约 81% 的 FLOPs 在算 padding。
Attention（平方项）利用率 ≈ Σ L_i² / (B · L_max²)；Σ L_i² ≈ 4.41e6，B·L_max² = 8 × 2048² ≈ 3.36e7，利用率 ≈ 13.1%——Attention 浪费更狠（约 87%）。
决策支撑：Padding-Free 把分母 B·L_max 拉到 Σ L_i，利用率回归 ~~100%，意味着同样吞吐下算力需求可降到约 1/5（或同硬件吞吐提升约 4~~5 倍）。这笔账 × QPS × 全年时长，就是上 Padding-Free 的收益上界；只要它显著大于「改 KV Cache 组织 + 重写 mask/loss + 接 varlen 算子」的工程成本，就值得上。

一句话决策：方差小 → 静态 Padding 不折腾；方差大但低频 → Bucket 快速止血；方差大、高 QPS、团队能 hold → Padding-Free 一步到位。永远先 profile 量化浪费，再决定值不值得上变长。

思考题 2：Roofline 驱动的 decode 优化

profile 显示某 LLM decode 阶段（逐 token 生成）的 top-1 耗时是权重 GEMV，算术强度极低、明显 Memory-Bound。结合 2.2 节的 Roofline 模型，解释：为什么此时换更高峰值算力的卡几乎无济于事，而 权重 INT4 量化 + KV Cache 量化 才是正解？量化把 Roofline 图上的哪个量改变了？

展开参考答案（含 Roofline 量化前后移动图 + 算一遍）

结论：decode 的 GEMV 落在 Roofline 的斜坡区（受带宽 β 限制），此时性能上界是 β × I 而非峰值算力 π，所以堆 π（换算力更高的卡）撞不到天花板；量化的本质是减少每个权重要搬的字节数，等价于在 Roofline 图上同时抬高算术强度 I 并让有效带宽吃得更满——这才是斜坡区算子唯一有效的杠杆。

用具体数字算一遍（以一个 7B 模型 decode 一个 token 为例，数量级估算）：

decode batch=1 时，每生成 1 个 token 要把全部权重读一遍。7B 模型 FP16 权重 ≈ 7e9 × 2 Byte = 14 GB。
在 HBM 带宽 ≈ 2.0 TB/s 的卡上，光搬权重 ≈ 14e9 / 2.0e12 ≈ 7 ms/token——这是带宽给的硬下限，与算力无关。
换算力翻倍的卡（π↑、β 不变）：搬权重还是 7 ms，token 速率几乎不动——钱花在了用不上的 π 上。
权重量化到 INT4（每权重 0.5 Byte）：权重 ≈ 3.5 GB，搬运 ≈ 3.5e9 / 2.0e12 ≈ 1.75 ms/token，理论约 4 倍 提速；KV Cache 量化进一步削减 decode 时随上下文增长的 KV 搬运量。

对应 Roofline 图：量化没有改 π（峰值算力），改的是算子的算术强度 I（FLOPs 不变、Bytes 变小 → I 变大），把工作点沿斜坡向右推、并减少了实际要搬的字节绝对量。斜坡区的算子，杠杆永远在「少搬字节」，不在「堆算力」。

思考题 3：跨栈定位陷阱

一位同学反馈「GPU 利用率只有 40%，肯定是 kernel 写得烂」。结合 2.3 节的「系统级 → kernel 级」定位漏斗，列出在断言「kernel 慢」之前必须排除的系统级原因（CPU launch bound、数据加载、H2D 拷贝、kernel 启动间隙、CPU-GPU 同步点），并说明各自在 Nsight Systems / PyTorch Profiler timeline 上的可观测特征。

展开参考答案（含系统级排查漏斗图 + 对比表）

结论：「GPU 利用率低」绝大多数时候不是单个 kernel 写得烂，而是 GPU 在「等」——等 CPU 下发、等数据加载、等 H2D 拷贝、等同步点；必须先沿 2.3 节的漏斗从系统级 timeline 排除这些「GPU 空跑」成因，确认关键路径确实卡在某个 kernel 上，才轮得到去优化 kernel 本身。

五类系统级原因的可观测特征对比表：

系统级原因	timeline 上的特征	典型修法
CPU launch bound	GPU 行大片空白，CPU 行被 Python/aten 调度排满；kernel 又小又稀	用 CUDA Graph / 算子融合减少 launch 次数；增大 batch
数据加载阻塞	每个 step 开头 GPU 固定空等一段；CPU 上 DataLoader/解码占满	`num_workers`↑、`pin_memory`、预取（prefetch）
H2D / D2H 拷贝	`Memcpy HtoD/DtoH` 行很长，compute 被推到拷贝之后串行	数据常驻 GPU、异步拷贝 + stream 重叠、减少回传
kernel 启动间隙	一串小 kernel 之间有等宽的固定空隙，空隙累计可观	融合小算子、CUDA Graph 把序列固化
CPU-GPU 同步点	周期性出现 `cudaStreamSynchronize` / GPU 突然停摆，常对应 `.item()`/`.cpu()`/`print`	移除训练热路径里的同步调用、延后取标量

为什么必须按这个顺序：跳过系统级直接调单 kernel，最常见的翻车是「优化了一个根本不在关键路径上的算子」——比如真正的瓶颈是 DataLoader 喂不上数，你却花两周把一个只占 3% 时间的 softmax kernel 调快了 20%，端到端几乎没动。40% 利用率往往不是 kernel 算得慢，而是 GPU 在等；只有当系统级原因全部排除、关键路径确实压在某个 kernel 上时，才轮到 Nsight Compute 去看它的 SM 占用率与访存效率（呼应 2.3 节漏斗的第 ④ 步）。

学习目标​

背景与现状​

原理与架构​

2.1 静态 Padding 的算力浪费——可量化​

2.2 Roofline 模型——判定 Memory-Bound vs Compute-Bound​

2.3 工具分工：从系统到 kernel 的定位漏斗​

动手实践：极简代码实操​

3.1 环境准备​

3.2 代码：抓一段推理 timeline 并定位 top-3 kernel​

3.3 运行与观察​

踩坑预警 (Gotchas)​

深入思考​

思考题 1：泛化策略选型的决策树​

思考题 2：Roofline 驱动的 decode 优化​

思考题 3：跨栈定位陷阱​

延伸阅读​