到底什么是内核？

现在似乎每个人都在讨论推理性能与工程优化。

我的时间线上每隔几条帖子就能看到“内核“这个词，而推理领域的企业们也在四处融资。

因此，这里将介绍GPU内核的基础知识，以及为何编写自定义内核具有实用价值。

背景：“内核“的多义性

在操作系统领域，内核是连接软件程序与物理硬件（CPU、内存等）的核心程序，负责提供接口。

在GPU编程中，内核指的是可在成千上万个GPU线程中并行执行的函数。英伟达GPU通常预装cuBLAS库，其中包含矩阵乘法、矩阵向量乘法、点积等预编译内核函数。

这两种“内核“的定义截然不同。你可能更常听到后者的含义：当人们说“编写了某个内核“，或某个模型因“定制内核“而运行飞快时，指的都是编写自定义GPU内核函数。

内核的本质

内核是运行在GPU上的函数。GPU会让成千上万的线程同时执行相同的函数，每个线程被分配唯一的线程ID，从而能够识别自身及其操作范围。

这是一个向量加法的内核示例：

注意代码中没有循环结构。a[]、b[]、out[]存储在共享内存中，你可以将这个函数分配到GPU的数千个核心上并行运行。

这就是核心思想！每个线程会确定自己负责处理的共享内存片段，从共享内存读取数据，进行计算，然后将结果写回共享内存。

内核的调用机制

除非编写自定义内核，否则你不会直接调用它。

例如torch.matmul(a, b)的编译过程如下：

• PyTorch检测到两个数组位于GPU上，会选择对应的C++矩阵乘法函数
• 该函数调用英伟达预装的库（如用于矩阵乘法的cuBLAS、用于神经网络常见运算的cuDNN），或是PyTorch自己编写的内核
• 库函数向GPU驱动发出内核启动请求：“这是代码和数据指针，请用指定数量的线程运行”
• CPU将启动任务加入队列后继续执行下一行Python代码，GPU则按顺序处理队列

因此运行单次模型推理意味着CPU将大量内核启动任务加入队列，GPU则依次处理这些任务。如果使用性能分析工具记录GPU的工作状态，你会看到精确对应这个过程：由数百个小方块组成的时间线，每个方块代表一个内核。

多个小方块，每个代表一个内核

每个内核启动都有数微秒的开销，这些开销会不断累积！

编写自定义内核

初看时GPU内核似乎已是成熟领域；矩阵乘法内核应该早已优化到极致。但模型并非单一矩阵运算；它由数百个小型操作串联而成，优化空间存在于操作融合、针对不同输入形状的特化设计、减少内核启动次数，以及利用最新硬件特性等方面。

我们将重点探讨内核融合：

GPU由缓慢但容量大的高带宽内存（HBM）与快速的小容量SRAM构成。HBM与SRAM之间的带宽约为每秒数太字节，但计算核心的处理速度远高于此。因此推理过程中大多数阶段，计算核心都在等待数据传输的空闲状态中度过。

（Claude生成的GPU示意图，数据基于H100）

现在观察原生PyTorch模型的执行过程。假设连续进行三个操作：乘法运算、常数加法、激活函数计算。这需要三个独立的内核：每个内核都从慢速内存读取输入数据，执行少量计算，再将结果写回慢速内存。

自定义内核能将这三个步骤融合为单次操作：仅读取一次数据，在SRAM中完成乘法、加法和激活函数计算，最后统一写回结果。这就是“融合内核“的实质含义。

FlashAttention是经典案例：标准注意力机制会计算完整的注意力矩阵，写入HBM，再读取回来应用softmax并与V相乘。

而FlashAttention在SRAM中以分块方式计算注意力分数、softmax并与V相乘，完成分块计算后再累积结果。（注：分块处理实则相当困难，需要精巧的算法设计）

自定义内核虽酷却应最后尝试

需明确的是，虽然编写自定义内核是提升首token时间（TTFT）、单token延迟和吞吐量的炫酷方式，但多数性能提升实际上来自内核之上那些看似平凡的优化层：

• 批量处理：将多个请求组合成单次前向传播
• 量化：用8位或4位替代16位存储权重，减少每token的数据读取量
• 推测解码：用小模型预测若干token，再用大模型单次验证
• 编译器优化：让机器学习编译器在手工编码前自动融合简单操作
• 其他技巧：包括前缀缓存、预填充/解码分离、分页注意力等

推理系统速度缓慢往往是因为跳过了上述标准优化步骤。编写自定义内核工作量巨大，只有在所有上层优化都已完成后才值得考虑。

入门路径

针对英伟达GPU，Triton库允许你用Python编写数据块处理逻辑，自动生成GPU代码。它能解决大部分需求，而CUDA则提供更底层的完全控制权。

这个领域仍有大量待探索空间。多数工作超越了内核层面：需要优化整个推理栈的批处理、调度和内存管理。许多探索也围绕内核展开…Triton只是众多内核专用领域语言（DSL）之一，该领域发展迅猛，ThunderKittens、TileLang、CuTe、Mojo等项目都在控制粒度与易用性之间探索不同的平衡点。

Horace He撰写的《从第一性原理实现深度学习加速》是性能工程与推理优化的绝佳入门读物，即使你不特别关注内核编写，也强烈推荐阅读。