排行榜     FlashInfer 追踪     GitHub

你有没有想过一个能够自我改进的 AI 系统？

— 不是科幻小说中的机器霸主，但一个能够利用 AI 本身来扩展其能力的 AI 系统听起来还是挺酷的。

— 这正是我们正在构建的：FlashInfer-Bench——一个基准测试和基础设施，它为 AI 加速实际世界中的 AI 部署开辟了道路。

AI 代理已经变得非常强大，能够编写复杂的代码，甚至构建复杂的系统。如此强大的能力自然让我们思考：AI 代理能否优化它们所运行的生产系统本身？在这些 AI 系统的核心中，最密集的部分是 GPU 内核——执行 AI 模型核心操作的低级程序。最近我们看到了令人惊叹的进展，表明 LLM 可以生成合理的 GPU 内核（https://scalingintelligence.stanford.edu/blogs/kernelbench/）。

这促使我们提出下一个自然的问题：我们如何系统地让 AI 代理改进它们所依赖的 AI 系统本身？我们知道实现这个终极梦想可能仍有障碍，但现在是时候准备好构建一条清晰的未来之路了。我们构建了 FlashInfer-Bench，这是一个由实际 AI 系统驱动的 GPU 工作负载基准测试，更重要的是，它提供了一个基础设施和工作流程，可以将 AI 生成的内核零日部署到生产环境中。虽然今天的 AI 可能尚未达到专家工程师的水平，但我们希望做好基础设施准备，并为加速这一终极目标开辟道路。相同的工作流程也将使机器学习系统工程受益，帮助我们构建和发展下一代内核库。

FlashInfer-Bench cover illustration

概述：系统地利用 AI 改进 AI 系统

当我们接近“AI 改进 AI 系统”的目标时，我们自然发现我们需要的不仅仅是基准测试。实际的生产环境非常复杂：它们涉及大量具有不同 API 设计和输入签名的复杂内核，而每个内核的性能也会根据其接收到的输入而变化。我们需要向 AI 代理和相关工程师清晰地传达这些需求；更重要的是，我们还需要能够快速将代理的解决方案重新应用于生产环境。从根本上说，我们需要考虑以下三个要素来系统地解决我们的问题

  1. 澄清现状。我们需要一种标准化的方式来描述工作负载，包括它们的输入和输出，以及相关的统计数据和特征（例如，形状、不规则性等）。
  2. 设定正确的目标。我们需要捕获实际、生产级工作负载的基准测试——换句话说，是 LLM 在实践中部署时发生的场景。
  3. 建立零日生产路径。AI 生成和人工编写的代码可以立即无缝部署到真实的 LLM 引擎中，干预越少越好。
Overview diagram of the virtuous cycle

FlashInfer-Bench 专门设计用于解决这三个挑战。如上图架构图所示，其核心组件的工作方式如下

  • 澄清现状 — FlashInfer Trace。为了形式化地描述 GPU 工作负载，我们开发了 FlashInfer Trace：一个系统化的、不断发展的开放模式，用于捕获 LLM 推理场景中的内核定义、实现和评估。这种格式通过四个关键组件解决了关键需求：定义内核签名和计算的内核定义；描述内核实际输入的负载；内核的解决方案；以及解决方案的评估结果。

  • 设定正确的目标 — FlashInfer-Bench 数据集。我们收集并整理了生产中使用的最重要内核及其真实工作负载，以创建数据集。这使我们能够在真实的流量条件下衡量内核性能，并进一步使改进生产系统成为可能。

  • 建立零日生产路径 — FlashInfer 一级集成。我们与 FlashInfer 建立了深入集成——一个广泛用于主要 LLM 推理引擎的开放 LLM 内核库。我们的解决方案可以动态地用 FlashInfer Trace 和 FlashInfer-Bench 数据集评估出的最佳性能内核替换 FlashInfer 内核。这使得在 LLM 引擎中激活最佳内核并以最小的努力测试端到端性能成为可能。

在本文的其余部分，我们将详细介绍这三个主要元素。

FlashInfer Trace：LLM 内核的标准化和演进模式

Comic for FlashInfer Trace

正如在概述部分所讨论的，我们需要一种标准化的方式，让人类、基础设施工作流和 AI 代理能够在整个生产生命周期中就上下文进行沟通。我们特别需要支持生产生命周期中的不同阶段，包括将问题交给代理/工程师、记录答案、评估和检查输出。

我们定义了 FlashInfer Trace，这是一组 JSON 模式，用于正式描述内核的定义、基准测试工作负载、解决方案和评估结果。下图显示了 FlashInfer Trace 的组件

FlashInfer Trace components

定义提供了内核的完整规范，包括内核元数据、输入和输出规范、输入和输出张量的轴，以及用 Python 编写的参考实现，定义了内核的计算行为。内核定义提供了 AI 生成内核所需的所有信息。这是一个通用矩阵乘法 (GEMM) 内核定义的示例

GEMM kernel definition

在内核定义中，我们旨在捕获尽可能多的内核细节，尤其是常量轴值。这是因为具有不同输入形状的内核通常需要不同的最优实现。例如，尺寸为 4096 的 GEMM 可能与尺寸为 128 的 GEMM 具有非常不同的实现。AI 系统中的某些轴是动态的，例如序列长度。这些维度将保持可变，以便工作负载可以提供其具体值。

我们还设计了 op_type，它在高级别指定输入、输出和计算，例如 GEMM（如 gemm）和分页分组查询注意力（如 gqa_paged）。通过定义 op_type，我们可以将具有不同轴但共享相似计算的多个内核定义组合在一起。我们为常见的 LLM 运算符定义了 op_type，我们的目标是与社区一起发展这个标准。

工作负载捕获了在实际流量中观察到的内核输入，包括张量输入以及标量或标志输入。如果张量的特定值不影响评估，则可以将其存储为随机以节省空间；否则，可以转储原始张量以重现最真实的生产场景。

GEMM workload

解决方案给出了内核的具体实现。它必须执行与内核定义中参考相同的计算。我们支持多种内核编程语言，包括 CUDA、Triton、Python 和 TVM（即将推出！），等等。我们还可以指定内核的作者、目标硬件、兼容的软件版本以及其他相关信息。

GEMM kernel solution

评估是在特定定义和工作负载上对指定解决方案的基准测试结果。评估所使用的硬件和软件，以及内核的正确性和运行时性能都经过了全面测量。当内核编译失败或出现运行时错误时，我们会相应地分配不同的评估状态。

GEMM kernel evaluation

通过 FlashInfer Trace 的标准化设计，我们可以轻松地在整个生产生命周期中交换信息

  • 工作负载追踪器可以从 LLM 推理引擎获取流量并生成工作负载追踪；
  • AI 代理可以获取定义以进一步生成内核解决方案；
  • 基准测试工具可以获取工作负载和解决方案并填写评估字段；
  • 排行榜将获取追踪集合并可视化结果。

我们按照这一理念构建了 flashinfer-bench Python 包，其中包含对 FlashInfer trace 和跨各个阶段的工具的支持，包括模式验证、基准测试、工作负载追踪以及与 LLM 引擎和 FlashInfer 的集成。FlashInfer Trace 的设计独立于特定的工具，我们鼓励社区开发更多功能的工具链。

真实的 CUDA 工程师也可以从 FlashInfer Trace 中受益。它标准化了内核接口和计算，减少了沟通开销，并实现了来自不同来源的内核之间的公平比较。它提供了测试工具来衡量内核的性能，并提供了追踪工具来轻松捕获实际工作负载。总之，这些大大简化了内核开发的负担。

我们提供了 FlashInfer Trace 和 Op Types 的标准化且清晰的文档。

数据集和基准测试 — 来自真实世界的衡量标准

Dataset comics

我们根据一个原则来整理数据集：实际世界相关性。这意味着关注 LLM 引擎中最重要的内核，并在真实的生产流量下记录它们的工作负载。我们选择了最流行的模型，包括 Llama 3、DeepSeek V3 和 Qwen 3，并记录了它们的主要内核，包括注意力、GEMM、MoE、归一化、采样等。我们还努力确保 LLM 引擎配置的真实性。例如，对于 DeepSeek-V3，我们使用原始的 FP8 量化，并启用张量并行 = 8 和专家并行 = 8。我们从真实数据集（包括 ShareGPT 等）提供输入。这些确保我们衡量的工作负载尽可能接近真实世界的场景。

Dataset workflow

通过 FlashInfer 实现零日集成

FIB apply comics

良性循环的最后一步是立即将 AI 生成的内核投入生产。我们通过与 FlashInfer（一个被主要 LLM 引擎广泛采用的内核库）和 SGLang 的一级集成来实现这一点。与 vLLM 的集成也在快速推进中，预计很快完成。

我们可以通过最少的努力，用我们评估中表现最佳的内核动态替换 FlashInfer API 中的内核。只需在 LLM 引擎中导入 flashinfer_bench 并启用环境变量 FIB_ENABLE_APPLY，内核就可以自动替换为本地数据库中最佳的内核。

FIB apply workflow

在幕后，动态替换由 flashinfer_bench.apply() 装饰器支持

@apply(lambda A, B: f"gemm_n_{B.shape[0]}_k_{B.shape[1]}")
def gemm_bf16(
    A: torch.Tensor,
    B: torch.Tensor,
) -> torch.Tensor:
    # Fallback Implementation
    return torch.matmul(A, B)

它接受一个 lambda 函数，将输入映射到特定的内核定义名称，然后替换原始计算函数。替换后的函数根据输入从数据库中筛选出性能最佳的内核解决方案，执行它，并返回其结果。如果没有找到匹配的内核解决方案，则运行原始函数作为回退。

通过使用 apply()，我们可以替换 FlashInfer 中的大量运算符。如果你想替换自己的实现或其他库中的运算符，也可以使用 apply() 轻松实现。

FlashInfer-Bench 排行榜 — LLM 内核优化的竞技场

FlashInfer-Bench leaderboard screenshot

FlashInfer-Bench 为 AI 代理提供了改进 AI 系统的标准化基准测试。因此，我们还评估了不同模型的能力，并构建了 FlashInfer-Bench 排行榜。它直接利用从基准测试套件生成的 FlashInfer Trace，并提供了结果的可视化显示。

我们采用最初由 KernelBench 提出的 \(\text{fast}_p\) 指标来比较不同内核的性能。我们使用 FlashInfer 内核作为比较基线。\(\text{fast}_p\) 表示给定内核运行速度比 FlashInfer 快 p 倍的工作负载比例。如果 p > 1，则意味着我们发现了一个优于最先进内核库的内核！在实践中，大多数内核的 p < 1，表明当前模型仍有很大的改进空间。对于不同的 p 值，\(\text{fast}_p\) 形成一条曲线。曲线下的面积越大，内核的性能越好。

我们还发现了少量加速比大于 1 的 AI 生成内核。我们深知确保内核实现正确性的重要性，因此我们正在手动验证每个内核的正确性。我们将在审核后很快发布加速比大于 1 的内核。

我们还为每个内核提供了一个单独的排行榜，以便更容易地为每个内核选择最佳实现。排行榜还允许按工作负载进一步过滤，以选择特定工作负载下表现最佳的内核。

FlashInfer-Bench kernel page

此外，模型架构概述显示了模型使用的所有内核，并指示哪些内核已被 FlashInfer-Bench 覆盖

FlashInfer-Bench model architecture overview

总结

FlashInfer-Bench 提供了一个系统化的解决方案来构建改进 AI 系统的 AI。有了它，我们已经为 AI 构建了一个自我优化的循环——一个良性循环。AI 改进自身的革命将继续下去——让我们看看它会走向何方！

展望未来，我们计划不断扩大模型和内核的覆盖范围。我们将与 GPU 工程师和 LLM 引擎开发人员密切合作，以适应该领域快速发展的需求。这一旅程是一项社区努力，我们欢迎您的讨论和贡献。

欲了解更多信息，请访问以下链接

致谢

FlashInfer-Bench 是与 CMU Catalyst、NVIDIA 和 Bosch 合作发起的一项研究工作。我们正在作为 FlashInfer 社区的一部分建立一个开放社区，并欢迎 ML 系统社区的贡献。

我们感谢整个 FlashInfer-Bench 团队为该项目所做的贡献

  • Shanli Xing* (UW, CMU)：核心组件和网页开发
  • Yiyan Zhai* (CMU)：FlashInfer-Trace 数据集、工作负载追踪系统
  • Alexander Jiang* (CMU)：基准测试系统、代理设计
  • Yixin Dong* (CMU)：核心思想、整体架构设计
  • Yong Wu (NVIDIA)：RMSNorm、采样、融合 MoE
  • Zihao Ye (NVIDIA)：FlashInfer 支持
  • Charlie Ruan (UC Berkeley)：工作负载追踪系统
  • Yingyi Huang (CMU)：融合 MoE
  • Yineng Zhang (独立研究员)：SGLang 集成
  • Liangsheng Yin (独立研究员)：SGLang 集成
  • Aksara Bayyapu (CMU)：LLM 代理
  • Luis Ceze (UW, NVIDIA)：项目指导和建议
  • Tianqi Chen (CMU, NVIDIA)：项目指导和建议

我们还要感谢 Mark Saroufim、Zhuoming Chen、Weihua Du、Bohan Hou、Hongyi Jin、Ruihang Lai、Xinyu Yang、Yilong Zhao、Haizhong Zheng、FlashInfer 社区、GPUMODE 社区、HuggingFace 社区、SGLang 社区、TensorRT-LLM 社区、vLLM 社区、Databricks 和 xAI 提出的富有洞察力的反馈。