完整文档索引见 llms.txt。在任意 URL 后追加 `.md` 即可查看该页面的 Markdown 版本。

选择合适的 kernel 优化工具

如今，大多数团队并不会从零开始直接编写原始 CUDA 来优化 kernel。它们通常会借助一整套位于模型代码与 GPU 之间的工具栈来完成这项工作。

实际问题很简单：如果你需要为 LLM 推理获得更快的 kernel，应该先从这个工具栈的哪一层入手？

理解工具栈的一个简单心智模型

你可以把 kernel 工具链看作一把“控制权梯子”。

最上层是 cuBLAS、cuDNN 这类高层 vendor library。它们通过稳定 API 暴露优化好的 kernel。再往下一层是 TVM、XLA 这类 AI compiler，它们试图从更高层的 graph 或 intermediate representation 中生成优化 kernel。再下面是 Triton 这类 kernel DSL，在这里你仍然需要自己写 kernel，但编程模型比原始 CUDA 友好得多。最底层则是手写自定义 kernel，在这里你可以直接控制内存访问模式、tiling 策略和调度方式。

Mojo 和 MAX 这类较新的全栈方案则试图模糊这些边界，把语言、编译器、runtime 和 kernel 工具整合到同一个栈中。

当你沿着这把梯子向下走时，通常会发生三件事：

你会获得对 tiling、memory movement 和执行细节的更多控制
你会需要更多 GPU 专业知识和更多工程时间
你能“免费”得到的可移植性会更少

高层 CUDA 库

最常见的起点是 NVIDIA 自家的库栈，其中包括 cuBLAS 和 cuDNN 等库。

这些库把高度优化的 kernel 封装在稳定 API 背后。如果你的 workload 符合常见 operator 模式，那么它们往往是不亲自做 kernel 工作却获得强性能的最快方式。

这一层覆盖的操作包括：

cuBLAS 和 cuBLASLt 中的 GEMM 与 tensor contraction
cuDNN 中的神经网络 primitive
越来越多的 fused attention 以及其他常见深度学习模式

这就是为什么 PyTorch、TensorRT-LLM 以及许多推理栈在底层都依赖它们。它们继承了多年针对特定架构（例如 Ampere、Hopper）的调优成果，而大多数团队都无法证明自己投入同等成本是合理的。

这种调优非常重要。一个在某代 NVIDIA GPU 上表现良好的 kernel，并不会自动在下一代上达到峰值性能。新 GPU 具有不同的 memory hierarchy、tensor core 行为、调度细节和指令集。NVIDIA 的库团队会针对这些变化重新调优，从而让用户无需亲自处理。

其好处显而易见。如果你的 workload 由标准操作构成，这些库通常开箱即用就已经足够好。

不过，它们也有几个限制：

覆盖范围局限于已知模式。它们不会自动支持每一种新的 attention 变体、fusion 策略或模型专用 operator。当模型演化时，库支持往往会滞后。
跨 operator 的优化能力有限。跨越不同库边界时很难做定制。如果你的关键路径横跨多个 operator，而最优解是一个新的 fused kernel，那么单纯调用库可能就不够了。
可移植性受限。CUDA 库绑定在 NVIDIA 生态之中。它们很强大，但这份性能收益无法直接迁移到 AMD GPU 或其他加速器上。

高层库能帮你快速起步，但它们只覆盖 NVIDIA 选择优化并通过这些库暴露出来的 operator。一旦你的 workload 超出这份目录，通常就需要转向 compiler、DSL 或自定义 kernel 路径。

AI compiler

AI compiler 是位于高层模型代码（例如用 PyTorch 或 TensorFlow 编写）与实际执行硬件之间的一层软件。它们的职责是在无需你手工重写任何内容的前提下，自动变换并优化计算，使其运行得更快。

一个简单例子是 matmul + bias + activation。朴素实现可能会把它们作为独立 kernel 运行，每一步都把中间结果写回 global memory。compiler 往往可以把这串操作融合起来，把更多中间数据保留在片上，减少额外 HBM traffic，从而提升效率。

为什么需要 AI compiler？

现代 AI 框架暴露出数以千计的 operator。想要针对每个硬件目标都手写并重新调优每一个 operator，这在现实中并不可行。原因包括：

操作本身可能包含许多小步骤，并伴随大量数据移动（这在 GPU 上会很慢）
不同硬件（NVIDIA GPU、移动端 ARM 芯片、Google TPU）需要不同的优化
手工调优所有东西既耗时，也无法规模化

AI compiler 会自动完成 kernel fusion、memory layout 调整、循环优化以及硬件特定调优等工作。它们通常会使用搜索算法，甚至机器学习，来寻找最优的计算调度方式。

AI compiler 如何工作？

大多数 AI compiler 都遵循类似的流水线：

导入模型。来源可以是 PyTorch、TensorFlow、ONNX 等框架。
Graph-level optimization。简化计算图（例如删除不必要部分、融合操作）。
Operator-level optimization。将计算分解为 tensor operation，并针对目标硬件优化每一个（或一组）操作。
代码生成。生成低层代码（例如基于 LLVM、CUDA 或自定义 kernel），并将其编译为可部署模块。
Runtime 执行。高效运行优化后的模型。

它们提供了一条从高层代码走向优化执行的路径，而无需直接编写 kernel。

下面是几个知名的 AI compiler。

Apache TVM

Apache TVM 是这一领域最知名的开源 compiler stack 之一。你先在 PyTorch 或 TensorFlow 这类框架中编写模型，TVM 再将其转换成自己的 graph IR（Intermediate Representation），然后为目标硬件生成优化 kernel。

TVM 的一个关键创新是 auto-tuning。TVM 会尝试成千上万种不同的 tiling 与调度方式，对每一种进行基准测试，然后选出最快的那个。这有点像通过暴力搜索找到最佳算法，而不是手工把它写出来。

不过，TVM 也面临几个重要挑战。

很难达到峰值性能。现代 GPU 依赖专用单元（例如 Tensor Core）和复杂的内存行为。自动生成的 kernel 往往不如高度调优的 vendor library 或手写实现。
生态碎片化。由于 TVM 是一个真实实现（而不仅仅是一个标准），硬件厂商经常会为自身需求 fork 并定制它。这会导致分叉和长期维护挑战。
难以跟上快速演进的 AI workload。TVM 最擅长的是相对稳定的 operator 模式。现代 GenAI workload 往往需要新算法以及紧密耦合的 kernel 设计（例如 FlashAttention），而 compiler 无法立即捕捉这些需求。
auto-tuning 成本高。探索巨大的搜索空间会导致编译和调优时间变长，从而拖慢迭代速度。

XLA 与 OpenXLA

XLA 来自 Google 生态，并与 TensorFlow、JAX 以及 TPU 执行紧密集成。它在该栈内尤其擅长 graph-level rewrite、layout 优化以及 backend-specific lowering。

对于已经在使用 JAX 或 TensorFlow 的用户来说，这种紧密集成是一个显著优势。当 compiler 能看到完整计算图时， XLA 就可以激进地进行 fusion 和重排操作，从而获得很强的性能。

需要注意的是，人们谈论 XLA 时，通常会指两个彼此密切相关的语境：

内部 XLA（以 TPU 为中心）：Google 内部用于 TPU workload 的 compiler stack
OpenXLA / 开源 XLA：面向 CPU、GPU、TPU 及其他 backend 的公共 compiler 生态和相关项目

它们共享 StableHLO 等重要组成部分，但公开栈与内部栈并不完全相同，而且不同 framework 与 target 上的 backend 成熟度也不一样。

与 TVM 类似，XLA 也面临一些现实限制：

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

想进一步了解 AI compiler，可以阅读 Chris Lattner 的博客文章：

AI compiler 试图通过隐藏底层细节来降低难度。但 GenAI workload 并不会让这件事变得简单。为了拿到顶级性能，你仍然需要：

对内存与执行有细粒度控制
具备编写自定义 kernel 的能力

这就形成了一个权衡：抽象越多，事情越容易做，但性能上限也越受限。

CUDA 给你完全控制权，但使用起来很难。由于如今大多数 AI 开发都发生在 Python 中，一个很自然的想法就是：

保留 Python 的开发体验
仍然生成快速的 GPU kernel

这正是 DSL 发挥作用的地方。

Triton 与 Python kernel DSL

DSL（Domain-Specific Language）是一种为特定目的而设计的编程语言，而不是通用语言。例如：

SQL 是用于数据库查询的 DSL
Regex 是用于模式匹配的 DSL

kernel DSL 则是专门为编写 GPU kernel 设计的语言。

Triton 是其中最突出的例子。它位于原始 CUDA 与完全 compiler 驱动的系统之间。

在 Triton 中，你仍然是在写 kernel。不同之处在于，你使用的是一种更高层、面向 tile 的编程模型，而不是 CUDA C++ 那种 thread-level 代码风格。

这使 Triton 比原始 CUDA 更容易上手。你仍然需要思考内存访问、tiling 和并行工作划分，但你花在直接管理底层机制上的时间会更少。

从高层看，一个 Triton kernel 大致长这样：

@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)

真正重要的不是语法，而是它的编程模型：

你处理的是数据 block 或 tile（tl.arange、tl.load）
你会显式加载和存储数据
你不需要直接管理 threadIdx、blockIdx，以及大多数 warp 级细节

这使 Triton 很适合现代 AI kernel 工作。它常被用于自定义 attention kernel、fused pointwise operation，以及那些超出标准库目录的研究驱动型 kernel。

Triton 也很好地说明了 kernel DSL 在实践中为何有价值。它在降低门槛的同时，并不会假装问题已经变简单了。你仍然需要 GPU 直觉，仍然需要理解 memory traffic、occupancy，以及 tile shape 如何与 workload 相互作用。

它同样有边界：

调试仍然可能很困难
性能调优并不是自动完成的
超出主 backend 的可移植性，取决于 backend 成熟度和 compiler 支持

这就是为什么很多人把 Triton 视为中间层。它帮助那些本来就能写 kernel 的人更高效工作；但它并没有消除 kernel 工作本身的必要性。

自定义 kernel

有时候，正确的做法很直接：去掉抽象层，针对那个确切的瓶颈编写一个自定义 kernel。

像 FlashAttention 这样的突破，就是从这里诞生的。

这是 LLM 推理中的一个重要模式。新的模型架构经常会带来新的 attention 变体、tensor layout 或 fusion 机会。库可能跟不上，compiler 也未必能很好优化它们。到这个阶段，自定义 kernel 就成了前沿路径。

不过，这条路并不容易规模化。这些投入很难在不同模型、不同硬件代际或不同团队之间平滑迁移。主要原因包括：

开发成本高。写出一个优秀的 kernel 远不只是写代码。它要求你理解：
- GPU 架构（thread、warp、memory hierarchy）
- 如何平衡 compute 与 memory bandwidth
- 如何对 workload 做 tiling 并避免瓶颈
  
  这种水平的专业能力很稀缺，也需要时间积累。
难以维护。任何模型或硬件变化都可能使 kernel 中固化的假设失效：例如预期的 memory layout、针对某一代 GPU 调好的 tile size，或者下一代架构中根本不存在的 intrinsic。
硬件锁定。当今 LLM 生态中的许多自定义 kernel 都是用 CUDA 写的，这意味着：
- 它们只能运行在 NVIDIA GPU 上
- 迁移到 AMD 或其他加速器需要重写
- 团队会被绑定到单一硬件生态

这就是为什么 AI 社区一直在寻找更好的方案：既保留手写 kernel 的高性能，又摆脱让其难以规模化的专业门槛与硬件依赖。

Mojo 和 MAX：一种全栈式尝试

到目前为止，这个栈中的每一层都在迫使你接受某种权衡：

vendor library 能让你快速起步，但会把你锁进 NVIDIA 硬件生态
AI compiler 提升了灵活性和可移植性，但牺牲了峰值性能所需的细粒度控制
Triton 提供了更友好的编程模型，但你仍然需要自己管理 tiling 并调试底层细节
自定义 kernel 给你完整控制权，但代价是时间、专业门槛和可移植性

Modular 通过 Mojo 和 MAX 做出的尝试，采取了不同路径。它不再把 kernel、compiler 和 runtime 视为彼此独立的层，而是把它们组合成一个垂直整合的系统。

Mojo 是一种具备 Python 风格语法、构建于 MLIR（Multi-Level Intermediate Representation）之上的编程语言，并对 GPU 编程提供一等支持。它的目标是让低层优化变得可用，而不必强迫开发者在每一条性能关键路径上都切换到 CUDA C++。

简而言之，你可以：

用熟悉的风格编写高层代码
在需要时逐步引入低层控制
在整个栈中保持使用同一种语言

它最大的不同之一在于硬件可移植性。传统自定义 kernel 是用 CUDA 编写的，只能运行在 NVIDIA 上。对于其他硬件，你不得不重写代码，或者维护多套独立栈。

而在 Mojo 中，同一份 kernel 代码可以运行在 NVIDIA、AMD 等多种硬件上。MLIR 编译步骤会负责将其转换为各家厂商的特定指令。这样一来，你无需复制 kernel 代码，就能在各个目标上获得具备硬件感知的性能与优化。例如，compiler 可以针对以下能力生成目标代码，而不需要你为每个厂商分别实现：

Warp synchronization
Tensor core instructions
Memory hierarchy details

如下方代码所示，一个 warp synchronization primitive 可以根据目标硬件（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）
graph 编译（通过 MLIR）
在硬件上生成并执行 kernel

MAX 不再把 PyTorch、vLLM、CUDA 和自定义 kernel 拼接在一起，而是提供一条从 model graph 到 GPU kernel 的统一栈。这使得优化可以跨层进行，而不是只局限在某一层内部。

需要注意的是，Mojo 和 MAX 仍在成熟过程中，其生态规模也远小于 CUDA 过去十五年积累的成果。不过，它们与此前的方案在性质上有所不同。理解 Mojo 和 MAX 的正确方式，不是把它们视为已经尘埃落定的赢家，而是把它们看作一种严肃尝试，试图弥补老工具留下的缺口：过多的碎片化、过多的重写工作，以及峰值性能与单一 backend 之间过强的耦合。

对于那些需要管理混合 GPU 集群、应对不断上升的硬件成本，或者承担手写 CUDA kernel 维护负担的团队来说，值得认真理解这套栈到底能做什么。

应该如何在这些工具之间做选择

大多数团队都不应该从栈底开始。

如果你的 workload 符合常见 operator 模式，就先通过 framework 或 inference engine 使用 vendor library。
如果 graph-level optimization 能消除内存往返、或者融合重复模式，那么 AI compiler 可能会给你带来杠杆效应。
如果你需要自定义 kernel，但又想要比 CUDA C++ 更好的编写模型，Triton 往往是下一步。
如果你既想要 kernel-level 控制，又想要更广泛的硬件可移植性，那么 Mojo 和 MAX 提供了另一条路径。它们让你用更高层的语言编写自定义 kernel，同时依赖 compiler stack 去面向不同 backend 生成目标代码。
如果某个关键瓶颈在特定硬件目标上仍然需要最大控制权，那么手写 CUDA kernel 可能值得付出相应成本。

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

这个 operator 是标准的，还是全新的？
你更在意单一 backend 上的峰值性能，还是跨多个 backend 的可移植性？
你的团队具备 compiler 专长、kernel 专长，还是两者都没有？
你是在构建通用产品基础设施，还是前沿模型 kernel？

常见问题

大多数推理团队应该自己写自定义 kernel 吗？

通常不应该。

大多数团队应该先使用 framework、vendor library 或现有 inference engine 已经提供的 kernel。只有当某个关键 workload 不能被这些工具很好支持，并且性能收益足以证明工程成本合理时，自定义 kernel 才有意义。

vendor library 对现代 LLM 推理来说够用吗？

对于许多标准操作来说，它们至关重要；但对于每一种前沿 workload，它们都不够。

一旦模型引入新的 attention 模式、新的 fusion 机会，或者不寻常的 operator 组合，你可能就需要 compiler 的帮助，或者走自定义 kernel 路径。

Triton 与 TVM、XLA 这类 AI compiler 有什么不同？

Triton 是一种 kernel 编写模型。开发者仍然需要编写 kernel 逻辑，并在相当低的层次上控制数据移动。

TVM 和 XLA 工作在更高的 graph 或 compiler 层。它们试图在跨多个操作的范围内，自动完成更多 lowering 与优化过程。

理解工具栈的一个简单心智模型​

高层 CUDA 库​

AI compiler​

为什么需要 AI compiler？​

AI compiler 如何工作？​

Apache TVM​

XLA 与 OpenXLA​

Triton 与 Python kernel DSL​

自定义 kernel​

Mojo 和 MAX：一种全栈式尝试​

应该如何在这些工具之间做选择​

常见问题​

大多数推理团队应该自己写自定义 kernel 吗？​

vendor library 对现代 LLM 推理来说够用吗？​

Triton 与 TVM、XLA 这类 AI compiler 有什么不同？​

延伸资源​