题目重点：拆分框架前端、图捕获或 eager 执行、IR、优化与 lowering、kernel 和 GPU runtime，纠正‘所有模型必须先转 ONNX’的常见误解。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- 以一个张量模型的推理为训练例；实际流程随 PyTorch eager、torch.compile、JAX jit、ONNX Runtime 等框架和后端改变。
- 区分模型权重与输入数据的设备放置、程序编译和 GPU kernel launch。
解题路线
- 先从 Python/框架层定义操作与参数：构造模块或纯函数，确定输入形状、dtype、设备与权重。框架可能逐 op eager 调度，也可能捕获函数为图。
- 编译路径将图表示为中间表示，做常量折叠、布局选择、融合等合法优化，再降低为目标设备 kernel 或调用现有高性能库；图捕获失败的片段可能回退 eager。并非每个算子都生成一份新的 GPU 汇编。
- runtime 分配和复用显存，传输或引用权重/输入，调度 kernel，处理同步与输出回传。计时需明确是否包含主机到设备传输和编译。
- ONNX 是显式导出或交换模型的一种选择；PyTorch torch.compile 与 JAX jit 各有自己的图/IR 路径。说‘一定经过 ONNX’会误导。
手动走一遍
训练例 `y = relu(x @ W + b)`：输入和权重先放在目标设备；框架或编译器识别矩阵乘、偏置与激活；后端可选择库调用和融合 elementwise kernel；runtime launch 后输出仍在设备上，只有显式读取到 CPU 时才发生相应同步/传输。具体融合取决于后端与形状，不能从这个表达式保证一定产生单个 kernel。
规模与取舍
矩阵乘若 `x` 为 B×D、`W` 为 D×H，算术量 O(BDH)，输入/权重/输出至少涉及 O(BD+DH+BH) 个元素；编译时间与图规模、后端和缓存相关，不能由这一式给固定渐近界。总延迟还包括传输、launch 与内存访问。
易错点与追问
容易失分的边界
- 把 ONNX 写成所有框架的强制中间表示。
- 忽略 eager 与编译路径不同，或以为所有计算都必须生成新 kernel。
- 基准测试中混淆首次编译、GPU 异步 launch、同步和数据传输。
继续深挖
- 怎样区别 graph capture 与 kernel compilation？
- 输入 shape 改变可能在哪个阶段触发重编译？
- 如何用 profiler 证明瓶颈在 kernel、launch 还是数据传输？
技术参考
- JAX Ahead-of-Time lowering：技术参考：JAX lowering 阶段。
- PyTorch Compiler 官方文档：技术参考：TorchDynamo/Inductor 编译路径。
- PyTorch ONNX Exporter 官方文档：技术参考：ONNX 是显式导出流程。
