面向 LLM 推理的 kernel 优化

人们在讨论如何加速 LLM 推理时，通常会首先想到 batching、caching 或 routing。这很合理，因为这些都是可见的系统层杠杆。

但在这些之下，还有一个更不显眼却真正完成计算的层次：GPU kernel。

什么是 kernel 优化？

当 LLM 生成一个 token 时，GPU 会执行一系列高度并行的操作，例如矩阵乘法、attention、normalization 和 activation function。 这些操作通过 GPU kernel 来实现：也就是在大量 GPU thread 上同时运行的函数。

有些操作会映射到单个 kernel，而另一些则由多个 kernel 组成。在很多情况下，多个操作还可以融合进单个 kernel， 而这正是影响最大的优化之一。

可以把 kernel 看作是 LLM 推理的内层循环。它们承担每一步 token 生成背后的核心数值计算，也是绝大多数 compute time 真正消耗的地方。

kernel 优化，就是通过以下方式让这些函数运行得更快、更高效的实践：

• 降低每个操作的 latency
• 提高硬件利用率（计算单元、内存带宽等）
• 尽量减少内存移动和数据传输

从高层看，kernel 优化通常有两种常见路径，区别在于你希望对 GPU 拥有多少控制权。

• 手写 kernel。这是传统路径。你直接编写并调优 GPU kernel，通常会使用 CUDA，或像 Triton 这样的更高层 kernel DSL。示例包括：

  • 自定义 CUDA kernel

  • 像 FlashAttention 这样的优化 attention 实现

    这样可以获得细粒度控制。你可以从零开始设计并调优内存访问、thread 布局以及计算调度。 但这也很难，因为它要求你对 GPU 架构有深入理解。

    对 Triton 来说，你仍然是在写 kernel，但基于 tile 的编程模型隐藏了许多底层细节： 不需要手工管理 warp，也不需要为常见模式显式安排 shared memory 调度。 这让 Triton 比直接编写原始 CUDA kernel 更容易上手。同时，它又比完全由编译器驱动的方法提供了更多控制权。

• 编译器驱动的优化。你依赖 TVM 或 XLA 之类的编译器系统，从更高层的计算描述中生成优化 kernel。 这些系统可以：

  • 将多个操作融合成一个 kernel

  • 重新安排计算顺序以提高效率

  • 自动生成优化 kernel

    在这种模式下，你描述“要计算什么”，而编译器决定“如何在目标硬件 backend 上高效执行”。 这显著降低了门槛。但它对底层执行的直接控制更少，而且在支持新模型架构或特定计算模式时，可能落后于手写 kernel。

kernel 优化不是什么

它很容易与 Inference Optimization 章节中的其他优化技术混淆。kernel 优化并不是：

• continuous batching 或 dynamic batching
• prefix caching
• KV cache routing 或 offloading
• GPU 间的负载均衡
• 多地域或多云部署

这些都工作在 kernel 之上的系统层或 runtime 层。kernel 优化则更深入一层，进入真正执行计算的 GPU 代码。

一个简单的区分方式是：

• kernel 优化：GPU 上单个操作运行得有多快
• 系统优化：请求在整个服务基础设施中流动得有多高效

两者都重要，只是解决的是根本不同的问题。

为什么大多数团队从不触碰这一层

大多数 AI 工程师从不编写 GPU kernel，而且这完全有道理。vLLM、SGLang 和 TensorRT-LLM 这类推理框架已经把所有底层执行细节都抽象掉了。

这意味着你只需要：

• 加载模型
• 调用 API
• 专注于产品功能

这种抽象是有意设计的。现代推理框架的目标之一，就是让 GPU 的复杂性对用户不可见。 另外，如果一个框架已经能提供可接受的性能，也确实没有明显理由再往更底层走。

为什么 kernel 优化仍然重要

即使你从不亲自写 kernel，kernel 这一层仍然决定了其上所有层的性能上限。

新架构需要新 kernel

当新的模型架构出现时，它往往会带来新的模式：

• 新的 attention 机制
• 新的 tensor shape
• 新的计算顺序

支持这些新模式，并不只是用 PyTorch 把它写出来那么简单。你通常还需要新的 kernel，或者新的操作融合方式。

像 vLLM 或 SGLang 这样的框架依赖的是预构建、预优化的 kernel 库。它们不可能在第一天就自动支持每一个新想法。 在有人为新的计算模式编写并集成高效 kernel 之前，框架要么退回到较慢的通用实现，要么根本无法支持该架构。

这就是为什么下面两件事之间通常会存在时间差：

• 一个模型被发布
• 这个模型能高效地在生产环境中运行

FlashAttention 本身就是一个很好的例子。它并不是某个推理框架发明出来的功能，而是一项独立的 kernel 优化，在证明其价值之后才被各个框架采纳。

性能上限是真实存在的

从某种意义上说，高层框架只是构建在一组 kernel 之上的编排层。vLLM、SGLang 和 TensorRT-LLM 最终都要分发 kernel 来完成实际计算。 这意味着推理速度不可能超过底层 kernel 所允许的上限。

系统层优化当然有帮助，但它们无法完全弥补低效 kernel 带来的损失。如果瓶颈就在 kernel 本身，那么任何 runtime 层面的改动都无法填平这个差距。

硬件可移植性

如果你编写自定义 CUDA kernel，它们就只能运行在 NVIDIA GPU 上。这会带来锁定效应：

• 迁移到 AMD GPU 或其他加速器会变得困难
• 你可能需要从头重写 kernel

这正是为什么可移植层与基于编译器的方案越来越重要。它们试图解决一个棘手的权衡：

• 特定 backend 上的峰值性能
• 跨硬件 backend 的更广泛可移植性

大多数团队都想同时得到这两者。但今天几乎没有方案能完全做到。一个为某种 GPU 架构精心调优的 kernel， 往往在另一种架构上表现不佳，除非进行大量重构。

常见问题

什么是 GPU kernel？

GPU kernel 是一种会在大量 GPU thread 上并行运行的小程序。

在 LLM 推理中，kernel 负责真正的数学计算，例如：

• 矩阵乘法
• attention 计算
• normalization 和 activation function

你可以把 kernel 看作 GPU 上执行的一个工作单元。当模型生成一个 token 时，底层会触发很多个 kernel。

什么是 kernel fusion？

kernel fusion 指的是把多个操作合并进单个 GPU kernel。它的核心收益是减少内存移动和 kernel launch 开销。

可以把它类比为做三明治。没有 fusion 时，一连串操作中的每一步都是一个单独任务：

1. 任务 1：走到冰箱前，拿出面包，放到台面上，烤一下，然后再放回去
2. 任务 2：再走到冰箱前，拿出烤好的面包，放回台面上，加火腿，然后再放回去
3. 任务 3：又一次走到冰箱前，拿出火腿三明治，放回台面上，加奶酪，然后再放回去

在 GPU 的语境中，冰箱代表 global memory（HBM），台面代表高速片上内存（register、shared memory、cache）。 关键问题在于，数据在 HBM 与片上内存之间移动的成本，远高于计算本身。

没有 fusion 时：

• 每个操作通常都要把中间结果写回 global memory
• 下一个操作又会把它们重新读出来
• 这意味着反复的内存往返和很高的带宽压力

有了 fusion 之后：

• 多个操作被合并进单个 kernel
• 中间结果通常可以保留在高速片上内存中
• 从而避免不必要的 global memory 读写

延伸资源

• CUDA C++ Programming Guide
• 究竟什么是“CUDA”？