如何选择合适的内核优化工具

如今,大多数团队并不是通过从头编写原生 CUDA 来优化内核,而是借助介于模型代码与 GPU 之间的一整套工具栈。

实际问题很简单:如果你需要为 LLM 推理更快的内核,应该先使用工具栈中的哪一层?

工具栈的简单心智模型

你可以把内核工具视为一把控制力从高到低的梯子。

最顶层是 cuBLAS、cuDNN 等高级厂商库,它们在稳定的 API 背后暴露优化内核。其下是 TVM、XLA 等 AI 编译器,它们试图从更高级的图或中间表示生成优化内核。再往下是 Triton 等内核 DSL,你仍然自己编写内核,但编程模型比原生 CUDA 更友好。最底层是手写的自定义内核,你直接控制内存访问模式、分块策略和调度。

Mojo 和 MAX 等较新的全栈方案试图将语言、编译器、运行时和内核工具融合到一个栈中,从而模糊这些边界。

沿着这把梯子向下,通常会发生三件事:

你对分块、内存搬运和执行细节的控制力更强
你需要更多 GPU 专业知识和更多工程时间
你能免费获得的可移植性更少

高级 CUDA 库

最常见的起点是 NVIDIA 自己的库栈,包括 cuBLAS 和 cuDNN 等库。

这些库在稳定的 API 背后封装了高度优化的内核。如果你的工作负载符合常见的算子模式,它们通常是在不自己做内核工作的情况下获得强劲性能的最快途径。

这一层覆盖以下运算:

cuBLAS 和 cuBLASLt 中的 GEMM 与张量收缩
cuDNN 中的神经网络原语
越来越多地涵盖融合注意力和其他常见的深度学习模式

这就是为什么 PyTorch、TensorRT-LLM 等框架和许多推理栈在底层依赖它们。它们继承了多年针对具体架构(如 Ampere、Hopper)的调优工作,这是大多数团队自己做都不划算的。

这些调优很重要。在某一代 NVIDIA 产品上正常工作的内核,并不会自动在下一代上达到峰值性能。新 GPU 有不同的内存层级、Tensor Core 行为、调度细节和指令集。NVIDIA 的库团队会针对这些变化重新调优,让用户无需亲自处理。

好处显而易见。如果你的工作负载由标准运算构成,这些库开箱即用通常就足够了。

然而,它们有几个局限:

覆盖范围仅限于已知模式:它们不会自动支持每一种新的注意力变体、融合策略或模型专属算子。当模型演进时,库的支持往往滞后。
跨算子优化受限:它们很难跨库边界进行定制。如果你的关键路径跨越多个算子,而最佳方案是新的融合内核,那么库调用可能不再够用。
可移植性受限:CUDA 库与 NVIDIA 生态绑定。它们很强大,但其性能优势无法转移到 AMD GPU 或其他加速器上。

高级库能帮你快速上手,但它们只覆盖 NVIDIA 选择通过库来优化和暴露的算子。一旦你的工作负载超出这个目录,往往就需要编译器、DSL 或自定义内核的路径。

AI 编译器

AI 编译器是位于高层模型代码(用 PyTorch 或 TensorFlow 等编写)与实际运行它的硬件之间的一层软件。它们的职责是自动转换和优化计算,使其运行得更快,而无需你手动重写任何东西。

一个简单的例子是 matmul + bias + activation。朴素实现可能会把它们作为独立内核运行,每次都把中间结果写回全局内存。编译器通常可以融合这一序列,把更多中间数据留在片上,减少多余的 HBM 流量,从而提高效率。

为什么需要 AI 编译器?

现代 AI 框架暴露了数以千计的算子。要针对每一种硬件目标手写并重新调优每一个算子是不现实的。原因在于:

运算可能包含许多小步骤和大量数据搬运(在 GPU 上很慢)。
不同硬件(NVIDIA GPU、移动端 ARM 芯片、Google TPU)需要不同的优化。
手工调优一切既耗时又难以扩展。

AI 编译器自动完成内核融合、内存布局调整、循环优化和硬件专属调优等优化。它们通常使用搜索算法甚至机器学习来寻找调度计算的最佳方式。

AI 编译器如何工作?

大多数 AI 编译器遵循相似的流水线:

导入模型:来自 PyTorch、TensorFlow 和 ONNX 等框架。
图级优化:简化计算图(例如移除不必要的部分、融合运算)。
算子级优化:分解为张量运算,并为目标硬件优化每个(或每组)运算。
代码生成:生成底层代码(例如使用 LLVM、CUDA 或自定义内核),并编译成可部署的模块。
运行时执行:高效地运行优化后的模型。

它们为你提供了从高层代码到优化执行、且无需直接编写内核的路径。

下面是几个知名的 AI 编译器。

Apache TVM

Apache TVM 是该领域最知名的开源编译器栈之一。你用 PyTorch 或 TensorFlow 等框架编写模型,TVM 将其转换为自己的图 IR(中间表示,Intermediate Representation),然后为目标硬件生成优化内核。

TVM 的关键创新之一是自动调优(auto-tuning)。TVM 会尝试数千种不同的分块和调度方式,逐一进行基准测试,并选出最快的一种。这就像用暴力搜索来寻找最佳算法,而不是手工编写。

然而,TVM 也面临几个重要挑战。

难以达到峰值性能:现代 GPU 依赖专用单元(如 Tensor Core)和复杂的内存行为。自动生成的内核往往不及高度调优的厂商库或手写实现。
生态碎片化:由于 TVM 是真实的实现(而不仅仅是标准),硬件厂商经常为它分叉并定制。这导致了分化和长期维护挑战。
难以应对快速演进的 AI 工作负载:TVM 最适合相对稳定的算子模式。现代 GenAI 工作负载往往需要新算法和紧耦合的内核设计(如 FlashAttention),编译器无法立即捕捉。
自动调优成本高:探索大型搜索空间会导致编译和调优时间漫长,拖慢迭代。

XLA 与 OpenXLA

XLA 来自 Google 生态,与 TensorFlow、JAX 等框架以及 TPU 执行紧密集成。它擅长该栈内的图级重写、布局优化和后端专属降级(lowering)。

对于已经在使用 JAX 或 TensorFlow 的用户来说,这种紧密集成是一个重大优势。当编译器能看到完整的计算图时,XLA 可以激进地融合和重排运算,以获得强劲性能。

注意,人们通常在两个密切相关的情境下谈论 XLA:

内部 XLA(面向 TPU):Google 内部用于 TPU 工作负载的编译器栈
OpenXLA / 开源 XLA:面向 CPU、GPU、TPU 及其他后端的公共编译器生态和相关项目

它们共享 StableHLO 等重要组件,但公共栈和内部栈并不完全相同,后端的成熟度也因框架和目标而异。

与 TVM 一样,XLA 也面临几个实际局限:

硬件抽象限制了控制力:XLA 的设计初衷是抽象硬件细节,但现代 GenAI 工作负载往往需要对内存、数据类型、自定义内核和执行模式进行细粒度控制。
依赖外部内核:在 GPU 上,峰值性能往往来自调用 CUDA 库,而不是完全由编译器生成的内核。

更多关于 AI 编译器的内容,请参阅 Chris Lattner 的博客文章:

AI 编译器试图通过隐藏底层细节让事情变得简单。然而,GenAI 工作负载并不让这变得容易。要获得顶级性能,你仍然需要:

对内存和执行的细粒度控制
编写自定义内核的能力

这带来了一个权衡:更多的抽象让事情更简单,但也限制了性能。

CUDA 给了你完全的控制,但很难用。由于如今大多数 AI 开发都在 Python 中进行,一个自然的想法是:

保留 Python 体验
同时仍然生成快速的 GPU 内核

这就是 DSL 的用武之地。

Triton 与 Python 内核 DSL

DSL(领域特定语言,Domain-Specific Language)是一种为特定目的而设计、而非通用用途的编程语言。例如:

SQL 是用于查询数据库的 DSL
正则表达式是用于模式匹配的 DSL

内核 DSL 是专门为编写 GPU 内核而设计的语言。

Triton 是最突出的例子,它介于原生 CUDA 与完全由编译器驱动的系统之间。

使用 Triton,你仍然要编写内核。区别在于,你是在更高级的、面向分块(tile)的编程模型中编写,而不是在 CUDA C++ 的线程级代码中编写。

这使得 Triton 比原生 CUDA 更容易上手。你仍然需要考虑内存访问、分块和并行工作划分,但直接管理底层机制的时间更少。

从高层看,Triton 内核看起来是这样的:

@triton.jit
def add_kernel(x_ptr, y_ptr, output_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    x = tl.load(x_ptr + offsets)
    y = tl.load(y_ptr + offsets)
    tl.store(output_ptr + offsets, x + y)

需要关注的重点不是语法,而是编程模型:

你处理数据块或分块(tl.arange、tl.load)
你显式地加载和存储数据
你不直接管理 threadIdx、blockIdx 或大多数 warp 级细节

这使 Triton 非常适合现代 AI 内核工作。它常用于自定义注意力内核、融合的逐元素运算,以及超出标准库目录的研究驱动内核。

Triton 也说明了内核 DSL 在实践中为何有用。它降低了门槛,但并不假装问题已经变简单。你仍然需要 GPU 直觉,仍然需要理解内存流量、占用率以及分块形状如何与工作负载交互。

它也有局限:

调试仍然可能很困难
性能调优不是自动的
超出主要后端的可移植性取决于后端成熟度和编译器支持

这就是为什么许多人把 Triton 看作中间层。它帮助那些会写内核的人,但并不会消除内核工作的需求。

自定义内核

有时候正确的做法很简单:去掉抽象,为确切的瓶颈编写自定义内核。

FlashAttention 之类的突破正是从这里来的。

这是 LLM 推理中的一个重要模式。新的模型架构往往带来新的注意力变体、张量布局或融合机会。库可能滞后,编译器优化得也不好。到那时,自定义内核就成为前沿路径。

然而,它们难以大规模推广。这些工作很难跨模型、跨硬件代际或跨团队转移。主要原因包括:

开发成本高:写好一个内核不仅仅是编码,还需要理解:
- GPU 架构(线程、warp、内存层级)
- 如何平衡计算与内存带宽
- 如何对工作负载分块并避免瓶颈
  
  这种专业水平很稀缺,需要时间积累。
难以维护:每一次模型或硬件变更都可能使内核中固化的假设失效:预期的内存布局、针对特定 GPU 调优的分块大小,或下一代架构上不存在的内建函数。
硬件锁定:当今 LLM 生态中的许多自定义内核都是用 CUDA 编写的,这意味着:
- 它们只能在 NVIDIA GPU 上运行
- 移植到 AMD 或其他加速器需要重写
- 团队被绑定在单一硬件生态上

这就是为什么 AI 社区一直在寻找更好的方案:既要有手写内核的高性能,又要没有让它们难以规模化推广的专业壁垒和硬件依赖。

Mojo 与 MAX:一次全栈尝试

到目前为止,工具栈的每一层都迫使我们做出某些权衡:

厂商库帮你快速上手,但把你锁定在 NVIDIA 硬件上。
AI 编译器提高了灵活性和可移植性,但牺牲了峰值性能所需的细粒度控制。
Triton 提供了更友好的编程模型,但你仍然要管理分块并调试底层细节。
自定义内核给了你完全的控制,但代价是时间、专业知识和可移植性上的高成本。

Modular 通过 Mojo 和 MAX 的尝试采取了不同的方法。它们不把内核、编译器和运行时视为分离的层,而是将它们组合成一个纵向集成的单一系统。

Mojo 是一种具有 Python 风格语法、构建于 MLIR(多层次中间表示,Multi-Level Intermediate Representation)之上、并将 GPU 编程作为一等公民支持的编程语言。其目标是让底层优化变得易用,而不强迫开发者在每个性能关键路径上都退回到 CUDA C++。

简而言之,你可以:

用熟悉的风格编写高层代码
在需要时逐步引入底层控制
在整个工具栈中只使用一种语言

最大的区别之一是硬件可移植性。传统的自定义内核用 CUDA 编写,只能在 NVIDIA 上运行。对于其他硬件,你必须重写或维护单独的栈。

使用 Mojo,同一份内核代码可以在 NVIDIA、AMD 等更多平台上运行。MLIR 编译步骤负责翻译为每家厂商的特定指令。这样,你就可以在每种目标上获得感知硬件的性能和优化,而无需复制内核代码。例如,编译器可以面向以下内容进行编译,而不需要为每家厂商单独实现。

Warp 同步
Tensor Core 指令
内存层级细节

如下面的代码所示,一个 warp 同步原语可以根据目标平台(NVIDIA 与 AMD)被编译成不同的实现:

# Compile-time warp synchronization per hardware

@always_inline("nodebug")
fn syncwarp(mask: Int = -1):
    """Synchronizes threads within a warp using a barrier."""

    @parameter
    if is_nvidia_gpu():
        __mlir_op.`nvvm.bar.warp.sync`(
            __mlir_op.`index.casts`[_type = __mlir_type.i32](
                mask._mlir_value
            )
        )
    elif is_amd_gpu():
        # In AMD GPU this is a nop (everything executed in lock-step).
        return

MAX 是其底层的执行与编译层,负责完整的推理路径:

模型定义(使用类似 PyTorch 的 API)
图编译(通过 MLIR)
在硬件上生成内核并执行

MAX 不需要把 PyTorch、vLLM、CUDA 和自定义内核拼接在一起,而是提供了一个从模型图到 GPU 内核的单一栈。这使得优化可以跨层进行,而不仅仅是在单层内。

请注意,Mojo 和 MAX 仍在成熟过程中,它们周围的生态只是 CUDA 十五年积累的一小部分。然而,它们与前人有着本质上的不同。看待 Mojo 和 MAX 的正确方式,不是把它们当作已定的赢家,而是看作一次认真弥合旧工具遗留空白的尝试:过多的碎片化、过多的重写工作,以及峰值性能与特定后端之间过强的耦合。

对于需要在混合 GPU 集群中游走、应对不断上升的硬件成本,或承受手写 CUDA 内核维护负担的团队来说,值得去了解这个栈能做什么。

你应该如何在这些工具中做选择

大多数团队不应从工具栈的底层开始。

如果你的工作负载符合常见的算子模式,就从你的框架或推理引擎所带的厂商库开始。
如果图级优化能消除内存往返或融合重复模式,AI 编译器可能会给你带来优势。
如果你需要自定义内核,但想要比 CUDA C++ 更好的编写模型,Triton 通常是下一步。
如果你既想要内核级控制,又想要更广的硬件可移植性,Mojo 和 MAX 提供了另一条路径。它们让你用更高级的语言编写自定义内核,同时依靠编译器栈在支持的地方面向不同后端。
如果某个关键瓶颈在特定硬件目标上仍然需要最大的控制,那么手写 CUDA 内核可能值得付出成本。

正确的选择取决于几个实际问题,答案会把你引向工具栈中的不同层。

该算子是标准的还是新颖的?
你更在意单一后端上的峰值性能,还是跨多个后端的可移植性?
你的团队是具备编译器专长、内核专长,还是两者都没有?
你是在构建通用的产品基础设施,还是在构建前沿模型内核?

常见问题

大多数推理团队应该编写自定义内核吗?

通常不需要。

大多数团队应该从框架、厂商库或现有推理引擎已提供的内核开始。当某个关键工作负载无法被这些工具很好地满足,并且性能提升足够大到能证明工程成本合理时,自定义内核才有意义。

厂商库对现代 LLM 推理来说足够吗?

对于许多标准运算来说,它们是必不可少的,但对于每一个前沿工作负载来说,它们并不够。

一旦某个模型引入了新的注意力模式、新的融合机会或不同寻常的算子组合,你可能就需要编译器帮助或自定义内核路径。

Triton 与 TVM 或 XLA 等 AI 编译器有何不同?

Triton 是一种内核编写模型。开发者仍然要编写内核逻辑,并在相当低的层面控制数据搬运。

TVM 和 XLA 工作在更高的图或编译器层面。它们试图在众多运算之间自动化更多的降级和优化过程。

工具栈的简单心智模型​

高级 CUDA 库​

AI 编译器​

为什么需要 AI 编译器?​

AI 编译器如何工作?​

Apache TVM​

XLA 与 OpenXLA​

Triton 与 Python 内核 DSL​

自定义内核​

Mojo 与 MAX:一次全栈尝试​

你应该如何在这些工具中做选择​

常见问题​

大多数推理团队应该编写自定义内核吗?​

厂商库对现代 LLM 推理来说足够吗?​

Triton 与 TVM 或 XLA 等 AI 编译器有何不同?​

扩展阅读​