LLM 推理的内核优化
当人们谈论加速 LLM 推理时,通常会想到批处理(batching)、缓存(caching)或路由(routing)。这很合理,因为它们都是可见的、系统层面的手段。
在这一切之下,有一个更安静的层在做实际工作:GPU 内核(kernel)。
什么是内核优化?
当 LLM 生成一个 token 时,GPU 会执行一系列高度并行的运算,例如矩阵乘法、注意力(attention)、归一化和激活函数。这些运算通过 GPU 内核实现:即同时在大量 GPU 线程上运行的函数。
有些运算映射到单个内核,有些则由多个内核组合而成。在许多情况下,多个运算还可以融合成一个内核,这是最具影响力的优化手段之一。
可以把内核看作 LLM 推理的内层循环。它们承担每个 token 生成步骤背后的核心数值计算,也是实际消耗大部分计算时间的地方。
内核优化就是通过以下方式让这些函数运行得更快、更高效:
- 降低每次运算的延迟
- 提高硬件利用率(计算单元、内存带宽等)
- 最小化内存搬运和数据传输
从宏观上看,根据你对 GPU 控制程度的诉求,内核优化有两种常见的路径。
-
手写内核(hand-written kernels):这是传统路径。你直接编写并调优 GPU 内核,通常使用 CUDA 或像 Triton 这样更高级的内核 DSL。例如:
- 自定义 CUDA 内核
- 优化的注意力实现,如 FlashAttention
这种方式提供了细粒度的控制。你可以从头设计和调优内存访问、线程布局和计算调度。然而,它也很难,因为需要对 GPU 架构有深入的了解。
备注对于 Triton,你仍然在编写内核,但基于分块(tile)的编程模型隐藏了许多底层细节:无需手动管理 warp,也无需为常见模式显式调度共享内存。这使得 Triton 比编写原生 CUDA 内核更容易,同时又比完全由编译器驱动的方法提供了更多控制。
-
编译器驱动优化(compiler-driven optimization):你依赖 TVM 或 XLA 等编译系统,从更高级的计算描述中生成优化内核。这些系统可以:
- 将多个运算融合到一个内核中
- 重排计算顺序以提高效率
- 自动生成优化内核
在这种模式下,你描述要计算什么,编译器决定如何在目标硬件后端上高效执行。这大大降低了入门门槛。然而,它对底层执行的直接控制较少,在支持新的模型架构或专门的计算模式时,可能落后于手写内核。
内核优化不是什么
人们很容易把它与推理优化章节中介绍的其他优化技术混淆。内核优化不是:
- 连续批处理或动态批处理
- 前缀缓存
- KV 缓存路由或卸载
- 跨 GPU 负载均衡
- 多区域或多云部署
这些技术作用于内核之上、位于系统层和运行时层。内核优化则深入一层,进入真正执行计算的 GPU 代码。
一种简单的区分方式:
- 内核优化:单个运算在 GPU 上运行的速度
- 系统优化:请求在整个服务基础设施中的流转效率
两者都很重要,只是解决的问题根本不同。
为什么大多数团队从不触碰这一层
大多数 AI 工程师从不编写 GPU 内核,这是有充分理由的。vLLM、SGLang 和 TensorRT-LLM 等推理框架抽象掉了所有底层执行细节。
这意味着你可以简单地:
- 加载模型
- 调用 API
- 专注于产品功能
这种抽象是有意为之。现代推理框架的目标之一就是让 GPU 的复杂性变得不可见。此外,如果某个框架已经提供了可接受的性能,就没有明显的理由再深入下去。
为什么内核优化仍然重要
即使你自己从不编写内核,内核层也决定了其上方所有层的性能上限。
新架构需要新内核
当一个新的模型架构出现时,它往往会带来新的模式:
- 新的注意力机制
- 新的张量形状
- 新的计算顺序
支持它不仅仅是用 PyTorch 实现。你往往需要新的内核或新的运算融合方式。
vLLM 或 SGLang 等框架依赖预先构建、预先优化的内核库。它们不会在第一天就自动支持每一个新想法。在有人为新的计算模式编写并集成高效内核之前,框架要么退回到更慢的通用实现,要么根本不支持该架构。
这就是为什么在以下两者之间常常存在差距:
- 模型发布
- 该模型在生产环境中高效运行
FlashAttention 本身就是一个很好的例子。它并不是某个推理框架发明的功能,而是一项独立的内核优化,在证明其价值之后才被各框架采纳。
性能上限是真实存在的
在某种意义上,高层框架是构建在内核集合之上的编排层。vLLM、SGLang 和 TensorRT-LLM 最终都要分发内核来执行实际计算。这意味着推理速度不可能超过底层内核所允许的上限。
系统级优化可以提供帮助,但无法完全弥补低效内核造成的损失。如果瓶颈出在内核本身,任何运行时层面的改动都无法弥合这一差距。
硬件可移植性
如果你编写自定义 CUDA 内核,它们只能在 NVIDIA GPU 上运行。这会造成锁定(lock-in):
- 迁移到 AMD GPU 或其他加速器变得困难
- 你可能需要从头重写内核
这就是可移植层和基于编译器的解决方案变得越来越重要的原因。它们试图解决一个艰难的权衡:
- 在特定后端上获得峰值性能
- 在多个硬件后端之间获得更广的可移植性
大多数团队两者都想要。今天,很少有解决方案能完全满足。为一种 GPU 架构精心调优的内核,往往需要大量返工才能在另一种架构上获得良好性能。
常见问题
什么是 GPU 内核?
GPU 内核是一种小型程序,它在大量 GPU 线程上并行运行。
在 LLM 推理中,内核负责实际数学计算,例如:
- 矩阵乘法
- 注意力计算
- 归一化和激活函数
你可以把内核看作在 GPU 上执行的工作单元。当你的模型生成一个 token 时,它在底层会触发许多内核。
什么是内核融合?
内核融合(kernel fusion)是指将多个运算组合到单个 GPU 内核中。其关键好处是减少内存搬运和内核启动开销。
以做三明治为例。没有融合时,链条中的每个步骤都是独立的任务:
- 任务 1:去冰箱拿面包,把它拿到台面上,烤好,再放回去
- 任务 2:去冰箱拿烤好的面包,把它拿回台面,加上火腿,再放回去
- 任务 3:再去冰箱,拿出火腿三明治,把它拿回台面,加上奶酪,再放回去
用 GPU 的术语来说,冰箱代表全局内存(HBM),台面代表快速的片上内存(寄存器、共享内存、缓存)。关键问题在于,HBM 与片上内存之间的数据搬运,其代价远高于计算本身。
没有融合时:
- 每个运算通常会把中间结果写回全局内存
- 下一个运算再重新读取它们
- 这意味着反复的内存往返和高带宽压力
有了融合:
- 多个运算被组合到单个内核中
- 中间结果通常可以留在快速的片上内存中
- 你避免了不必要的全局内存读写