AI 算子与编译器

• 掌握各算子的功能与计算特征
• 理解各算子如何适配 GPU 并行
• 理解 KV Cache 与注意力实现的关联

LLM 推理的核心算子包括：GEMM（矩阵乘），用于线性层与注意力投影，通过 tiling + 共享内存 + Tensor Core 实现高吞吐；Attention（FlashAttention 系列），用分块+在线 softmax 减少 HBM 读写，把注意力计算放到片上；RMSNorm，对激活做逐行归一化，是 elementwise 规约算子，可用并行规约实现；RoPE（旋转位置编码），把位置信息编码进 Q/K 向量，通过旋转（可分解为少数三角运算）实现，仍是逐元素/块级运算；KV Cache 用于缓存历史 K、V，避免重复计算，推理时按需读取并与当前 Q 计算注意力。这些算子相互配合：Q=GEMM(输入,Wq) 后做 RoPE，再与 KV Cache 做 FlashAttention，最后经 GEMM 输出。GPU 实现上，GEMM 与 Attention 是计算/访存密集的关键，用 Tensor Core 与显存优化；RMSNorm 与 RoPE 相对简单，用规约与向量化实现。

理解这些算子要抓住两点：每个算子"算的是什么"和"访存/计算特征如何、如何并行"。GEMM 与 Attention 是性能瓶颈，RMSNorm/RoPE 是轻量但需融合的算子，KV Cache 是内存管理核心。

• 理解 FlashAttention 各版本的演进优化
• 掌握 tiling、recomputation、warp specialization 的机制
• 理解如何逼近 IO 复杂度下限

FlashAttention-1 用 tiling 与在线 softmax 把 HBM 读写降到 O(N²)。FlashAttention-2 进一步优化：减少非矩阵乘操作（非线性），把并行化从序列维度扩展到 head 维度，提高并行度与 Tensor Core 利用率；同时优化了任务分配，让每个线程块处理一个 head 的更多分块，减少子块间通信。FlashAttention-3 引入 warp specialization（warp 分工），让不同 warp 分别负责数据加载（TMA/异步拷贝）与矩阵乘计算，实现计算与访存流水线重叠；配合 recomputation（重计算，不保存中间 S/P 矩阵，反向时重新计算以减少显存）与更精细的调度，把注意力运行的 IO 与计算开销逼近理论下限。这些优化共同把注意力推向"瓶颈由 HBM 带宽/算力决定而非算法冗余"的极限。

各版本优化的主线是"更小的 IO 开销 + 更充分的算力利用 + 更强的并行流水"。理解 tiling 解决 IO、recomputation 省显存、warp specialization 做流水重叠，就能把握 FlashAttention 系列的全貌。

• 理解算子库与编译器的分工
• 掌握各自适用场景
• 判断手写 kernel 的必要性

算子库（cuBLAS/cuDNN/oneDNN）针对常见算子（矩阵乘、卷积、归一化等）提供预编译、高度优化的实现，覆盖广、性能好，是首选；但它们是"算子级"的，无法跨算子融合，且对非常规或自定义算子支持有限。编译器（XLA/TVM/Triton）在更高层次做图优化、算子融合、布局变换与代码生成，能把多个算子融合成更高效的 kernel，适合动态/自定义计算图。两者边界在于：算子库优化的是"单个算子内的性能"，编译器优化的是"算子间的协调与代码生成"。手写 kernel 的必要场景：当算子库/编译器无法覆盖（如新的自定义算子、融合特殊逻辑）、需要极致性能（如达到理论下限）、或需要利用特定硬件特性（TMA、WGMMA）时，才需要手写 CUDA/Triton kernel。多数情况下先用算子库与编译器，性能瓶颈处再手写。

决策顺序通常是"算子库优先，编译器融合，最后手写"。算子库省心、编译器省力、手写最灵活但成本高。手写 kernel 是收益/成本权衡的结果，而非默认选择。

• 理解 PagedAttention 的块表机制
• 理解它如何减少显存碎片
• 理解它如何支持连续批处理

传统 KV Cache 按请求固定长度预分配连续显存，导致大量浪费与碎片。PagedAttention 借鉴操作系统分页思想，把 KV Cache 划分为固定大小的块（block），每个请求的 KV 通过一张块表（block table）记录其逻辑块到物理块的映射，物理块按需动态分配。这样：只有实际使用的块才占用显存，请求长度可变时不再整体预留，显著减少显存浪费与碎片；同时块表让不同请求的 KV 可以零散存放，便于在连续批处理（continuous batching）中高效调度——请求结束立即释放其块，新请求复用已释放块，实现请求级动态显存管理。稳定管理下，显存利用率大幅提升，吞吐得到显著提高。

PagedAttention 的核心是"按需分页 + 块表映射"，用显存碎片化换去利用率与调度灵活性。它把 KV 从"连续预留"变为"按需分配"，是连续批处理高吞吐的关键支撑。

• 理解 torch.compile 的编译流程
• 掌握 TorchInductor 与 Triton 的角色
• 理解图捕获、算子融合、自动调优的工程价值

torch.compile 是 PyTorch 2.0 引入的编译入口，把 eager 模式计算图捕获（graph capture）为中间表示，经 aot_autograd 生成反向图，TorchInductor 作为默认后端把图编译为目标代码；GPU 上 TorchInductor 生成 Triton kernel，CPU 上生成 C++（LLVM）代码。工程价值体现为：图捕获消除逐算子调度的 Python 开销，算子融合把多个小算子合并成单个 kernel，减少 kernel launch 与显存往返，自动调优（autotune）对 tile 大小、循环顺序等择优，提升实际性能。整体上 torch.compile 让研究者无需手写 kernel 即可获得接近手工优化的性能，大幅降低部署成本。

torch.compile 的价值是"编译自动化"：把 eager 的灵活性与编译的高性能结合。理解图捕获（减少调度开销）、融合（减少访问）、autotune（选择最优实现）是把握编译栈的关键。

• 掌握各框架的定位与特点
• 理解 KV cache 压缩、连续批处理、推测解码等共性技术
• 能对比各框架的差异

TensorRT-LLM 是 NVIDIA 的推理引擎，深度优化单卡/多卡性能，支持图编译、FP8/INT8 量化、KV cache 管理与 CUDA Graphs，适合需要极致延迟与吞吐、深度调优的场景。vLLM 以其 PagedAttention 块表管理 KV cache 和连续批处理著称，显存利用率高、吞吐高，开源生态活跃，是主流推理服务选型。SGLang 在高效推理基础上强调结构化输出与复杂控制流（如 Jump-Forward、约束解码），并优化了前缀缓存与 KV cache 复用。三者的共性技术：KV cache 压缩（KV cache 量化、块管理、前缀复用）减少显存与带宽；Continuous Batching（连续批处理）在请求粒度动态调度，提高吞吐；Speculative Decoding（推测解码）用草稿模型一次生成多个 token 再验证，减少解码步数。差异在于：TensorRT-LLM 偏性能深度优化、vLLM 偏吞吐与易用、SGLang 偏结构化推理与复杂场景。

对比框架要抓住"优化思路"而不只是功能列表。KV cache 压缩、连续批处理、推测解码是提升推理吞吐的三大抓手，各框架在不同方面强化，选型取决于延迟/吞吐/灵活性需求。

• 理解 FlashAttention 的 IO 优化核心
• 掌握分块与在线 softmax 的机制
• 理解其对通用算子优化的启发

FlashAttention 的 IO 优化思想是"把计算从慢速 HBM 搬到片上快速 SRAM"：通过分块（tiling）让数据块在片上处理，避免中间结果（S、P 矩阵）落盘回读；通过在线 softmax（分块更新 running max 与归一化系数）让 softmax 不需要完整矩阵即可正确计算，从而支持分块流式处理。这一思想对算子优化的启发是：优先考虑数据放置与访存路径（IO 复杂度），而非仅盯计算量；能用流式/在线算法避免保存大型中间量的尽量用；减少 HBM 往返往往比减少 FLOP 更有效，尤其对 memory-bound 算子。这些启发同样适用于归一化、Attention 变体、窗口注意力等：把"访存最小化"作为算子设计的第一原则。

FlashAttention 的普适启发是"算法设计要同时考虑 IO 与计算"。在线 softmax 是重难点：它让原本依赖全局统计量的运算可以分块增量计算，是各类流式算子设计的模板。

• 理解算子融合的原理
• 掌握融合减少 launch 与显存往返的机制
• 了解典型融合案例

算子融合把多个相邻算子合并为单个 kernel，从而减少 kernel launch 次数（降低启动开销）与中间结果的显存读写（避免中间张量写回 HBM 再读回）。典型融合案例：GEMM + Bias + Activation + LayerNorm 融合（把线性层、激活、归一化合成一个 kernel，减少中间张量往返）；FlashAttention 融合了 QK^T、softmax、PV 三部分；融合归一化（fused RMSNorm）把规约与缩放合并；逐元素运算链（如 y = sigmoid(x) * w + b）融合为单个 kernel。融合的收益来自减少访存与启动开销，但需注意：过度融合会增大 kernel 复杂度、降低算子复用性，且寄存器/共享内存不足时可能受限。编译器（如 torch.compile、Triton）常自动做此类融合。

融合的核心收益是"减少中间张量落盘"与"减少 kernel launch"。判断哪些算子可融合：数据流连续、可算子内直接传递、无副作用时。典型案例如 GEMM+bias+activation 是融合的经典模板。

• 理解算子融合与访存优化的内在联系
• 掌握 FlashAttention 作为融合算子的体现
• 理解访存优化对算子设计的影响

AI 算子融合与访存优化高度相关：融合的本质是把原本需要多次 HBM 往返的中间结果留在片上，从而减少访存量。FlashAttention 是融合的典范，它把 QK^T、softmax、PV 三个算子融合为单个 kernel，中间结果（S、P 矩阵）不落回 HBM，直接在片上计算与消费，把 IO 复杂度从 O(N²) 降到 O(N)。访存优化是融合的动机与手段：融合的收益主要来自减少访存而非减少 FLOP，因此访存密集的算子链（如 GEMM 后的激活、归一化、注意力）最值得融合。反过来，访存优化也常通过融合实现——把需要复用数据的算子合并，使数据在共享内存/寄存器中流转。理解这一关系，就能在设计算子时优先考虑"能否融合减少访存"。

融合与访存优化是"一体两面"：融合提供了减少访存的路径，访存优化是融合的收益来源。FlashAttention 明确展示了如何用融合把中间结果留在片上，是 AI 算子优化的范式。

• 理解 tiling 的数据复用原理
• 掌握寄存器分块（register blocking）的机制
• 理解访存优化对卷积/矩阵乘的作用

卷积与矩阵乘都是访存密集型算子，tiling（分块）与寄存器分块（register blocking）是核心访存优化。tiling 把输出矩阵/张量划分为小块，让每个块在计算时只加载所需输入块到共享内存，多个线程复用同一块数据，减少对全局内存（HBM）的重复访问；块大小与共享内存容量匹配，块越大复用率越高但受容量限制。寄存器分块让每个线程在寄存器中累积多个输出元素（如 4x4 或 8x8），从而减少从共享内存/全局内存重复读取输入元素的次数，把最频繁的累加留在寄存器中。两者结合：tiling 负责块级数据复用（减少 HBM 访问），寄存器分块负责线程内数据复用（减少共享内存访问），配合向量化加载与合并访问，能显著提升访存效率与带宽利用率。矩阵乘中这一套组合（如 cutlass 的 tile + warp-level register blocking）是高性能 GEMM 的基础。

访存优化围绕"数据复用"展开：tiling 在块级复用，寄存器分块在线程内复用。理解的层次是"哪一级数据被复用多少次、减少哪种内存的访问"，从而判断优化是否有效。

• 理解量化算子的反量化机制
• 掌握 dequant 与 GEMM 的融合
• 理解量化误差对注意力的影响

量化推理中，权重与激活通常以 FP8/INT8 存储，GEMM 在低精度下计算，但累加器保持在更高精度（如 FP32/INT32）。反量化（dequant）把低精度结果按 scale/zero-point 还原为高精度。性能优化的关键是 dequant 与 GEMM 融合：把 scale 融合进 GEMM 的累加或输出阶段，而非单独做一次 dequant kernel，从而减少内存往返与额外 kernel launch。同时量化在 GEMM 前用 scale 缩放，避免逐元素反量化。量化误差对注意力输出的影响体现在：注意力分数 QK^T 与 softmax 对精度敏感，低精度表示的 Q/K 会引入误差，可能导致 softmax 峰值偏移、注意力分布不准，进而影响输出生成质量；尤其在长序列与多轮时误差可能累积。因此注意力投影常使用更高精度或更精细的量化策略（如 per-head/per-token scale）来降低误差。

量化的核心是"用精度换吞吐"，关键是 dequant 与 GEMM 融合以省去独立反量化开销。注意力对量化误差敏感，需在精度与性能间权衡，常对注意力关键计算保留更高精度。

• 理解推测解码的原理
• 掌握草稿模型与验证模型的分工
• 理解加速比与接受率的关系

推测解码（speculative decoding）用一个小的草稿模型（draft model）快速自回归生成 K 个候选 token，再用大的目标模型（target model）并行验证这些候选：一次前向传播同时计算 K 个 token 的 logits，逐位置与草稿 token 比较，从第一个不匹配的位置截断，保留匹配的前缀并丢弃错误 token。由于目标模型自回归生成是逐位置的串行解码，而验证是并行的，因此能显著减少解码步数。加速比与接受率（每个位置草稿 token 被目标模型接受的概率）相关：设接受率 a，则一次验证平均能推进约 a 的几何相关个数，期望加速比近似为 1/(1-a) 量级（更精确地，期望推进长度 = 1/(1-a)）。接受率越高，加速比越大；当草稿模型与目标模型分布接近时接受率高，加速接近 K 的上限。实用上草稿模型需足够小以保证生成速度快，又要与目标模型分布接近以保证接受率高。

推测解码的本质是"用小模型试错、大模型并行验证"。加速比核心依赖接受率 a：期望推进 token 数约为 1/(1-a)。理解接受率与加速的数学关系，就能把握草稿模型"小且准"的取舍。

• 理解各编译器的 IR 设计
• 掌握 IR 对跨硬件可移植性的作用
• 理解不同 IR 的抽象层次

各编译器的 IR 设计决定了跨硬件可移植性。Triton 使用块级（block-level）IR，把 kernel 表示为对 tile 块的运算，编译器负责把块级操作映射到具体硬件（如 CUDA/ROCm），这种高抽象 IR 让同一 Triton 代码可移植到多种 GPU。Pallas（JAX）采用显式分配内存空间（SMEM/VMEM/HBM 等）的 IR，针对 TPU 等提供低层控制，可移植性较强但需平台特定知识。MLIR 是一个多级 IR 框架，提供从高层的 linalg/gpu 到底层 LLVM 的多级抽象，通过逐步 lowering 把计算图翻译到不同硬件后端，跨平台能力强。共同点是：IR 抽象层次越高越易移植但控制力弱，越低越难移植但优化空间大。IR 设计的目标是在"可移植性"与"性能可表达性"之间平衡，让编译器可以把高层计算降到多硬件。

IR 是编译器跨平台的核心：抽象层次决定可移植性与性能表达。Triton 用块级 IR、Pallas 用显式内存 IR、MLIR 用多级 IR，各有取舍。理解 IR 层次就能理解各编译器的可移植性来源。

• 理解 Triton 的块级编程模型
• 掌握 Triton 与 CUDA 的差异
• 理解 autotune 的机制

Triton 的编程模型以"块（tile/block）"为单位，程序员用 Python 编写对块级张量的运算，编译器自动处理线程/寄存器分配、共享内存与内存布局，无需手动管理 block/thread/warp 与共享内存；而 CUDA 需要显式管理线程层次、共享内存、内存拷贝与同步。Triton 因此更易写、更易移植，但牺牲了底层控制（如精确的寄存器分配、bank 冲突规避）。自动调优（autotune）通过 @triton.autotune 装饰器，为 kernel 提供多个候选配置（如 tile 大小、num_warps、num_stages），在运行时对每个配置做基准测试，选择最快的一个，并缓存结果避免重复测试。这使同一 kernel 能在不同硬件上自动选择最优参数，弥补了抽象层带来的性能不确定性。

Triton 用块级抽象换取易用性，用 autotune 弥补性能调优。理解差异在于"谁掌控线程/内存细节"：CUDA 手动、Triton 自动。autotune 通过基准测选择最优配置，是 Triton 性能的关键保障。

• 理解计算图优化的内容
• 掌握代码生成的流程
• 理解 TVM/MLIR 的分层

AI 编译器（TVM/MLIR）把高层计算图优化并生成可执行代码。计算图优化包括：算子融合（把相邻算子合并，减少访存与 launch）、常量折叠与精简、布局变换（把布局转换为对目标硬件佳的布局）、内存规划（复用内存缓冲）、代数化简等，这些在图级别（graph-level）进行。代码生成则把优化后的图逐出为设备代码：TVM 用 TensorIR/Relax 表示循环与调度，进行 tiling、向量化、并行化等调度变换，再生成 CUDA/ROCm/LLVM 代码；MLIR 通过多级 lowering（从 linalg/gpu 到底层 LLVM）逐级转换，最终生成目标代码。图优化管"算法级"，代码生成管"实现级"，两者配合让高层计算高效映射到硬件。TVM 还引入自动调度（meta schedule）与 auto-tuning 来自动搜索最优实现。

编译器分"图优化"与"代码生成"两大部分：前者在计算图级别做融合/布局/内存规划，后者把算子降为具体硬件代码。理解这种分层，才能知道性能问题出在算法层还是实现层。

• 理解算子融合的规则与收益
• 掌握布局转换的机制
• 理解两者对性能的影响

算子融合是图优化中降低访存与启动开销的关键手段。TVM/MLIR 按规则把满足条件的相邻算子合并为单个算子：例如 elementwise 算子与其前一算子融合（如 GEMM+elementwise 融合）、可逐元素对位计算的算子链融合。融合后中间结果不再写回全局内存，减少 HBM 往返与 kernel launch。布局转换（layout transformation）把数据从一种内存布局（如 NCHW，通道优先）转换为对目标硬件更高效的布局（如 NHWC 或 NHWC 的向量化友好布局），或转换为算子实现所需的布局。编译器会评估布局转换的代价（额外的数据搬移）与收益（算子效率提升），选择有利的布局。两者常配合：先融合减少中间量，再通过布局转换让算子访问更高效。TVM 的 AlterOpLayout 与 MLIR 的 layout 转换 pass 即体现此机制。

算子融合减少访存，布局转换提高访问效率。理解两者要关注"代价-收益"：融合省去中间写回，布局变换可能引入搬移但提升算子效率。编译器在图上做全局决策。

• 理解分块（tiling）的机制
• 掌握在线 softmax 的作用
• 理解 IO 复杂度降低

FlashAttention 的分块与 IO 优化原理：把 Q、K、V 按块（block）切分，每次只把一个小块加载到片上 SRAM（共享内存），在片上完成该块的 QK^T、softmax、PV 计算，避免把整个 S、P 矩阵写回 HBM。核心难点是 softmax 需要全局最大值与归一化，FlashAttention 用在线 softmax 解决：分块处理时维护 running max 与 running sum，用增量公式在每块到达时更新累计结果，从而在不需完整矩阵的前提下得到正确的 softmax 输出。最终 HBM 读写量从标准 Attention 的 O(N²)（S、P 落盘）降到 O(N)（只读写 Q、K、V）。分块大小受 SRAM 容量约束，选择与容量匹配的块大小能最大化复用、最小化 HBM 往返。这就是 FlashAttention 能把 Attention 的 IO 复杂度压到理论下限的原理。

分块解决"数据放哪"，在线 softmax 解决"如何流式算 softmax"。两者结合让中间结果不出片、HBM 只读写 Q、K、V，IO 复杂度由 O(N²) 降到 O(N)。

• 理解调度的内容
• 掌握内存规划的作用
• 理解 kernel 生成流程

AI 编译器在调度阶段决定算子如何实现，包括循环变换（tiling、vectorization、parallelize、unroll）与内存规划。内存规划（memory planning）决定张量在内存中的布局与复用：为中间张量分配缓冲区，并尽量复用已释放的缓冲区，减少显存占用；同时决定数据放在寄存器、共享内存还是全局内存，以及推理时的内存池管理。kernel 生成（codegen）把调度后的算子翻译为实际设备代码：TVM 生成 CUDA/OpenCL/LLVM 代码，Triton 生成目标 kernel，MLIR 逐级 lowering 最终生成 LLVM IR。调度决定"怎么算"，内存规划决定"数据放哪"，kernel 生成决定"最终代码"。三者配合决定算子的实际性能与显存效率，是编译器把高层计算落到高效硬件执行的关键。

调度是编译器性能的核心：循环变换决定计算效率，内存规划决定显存与复用，kernel 生成决定落地代码。理解调度就能理解为何同一算子在不同配置下性能差异巨大。