RealisticTritonBench：现实世界人工智能框架中 Triton 内核生成的基准

在现代人工智能框架中，GPU 内核是整体系统性能的关键。 Triton 结合了可用性、可移植性和近乎手写的 CUDA 性能，被广泛用于实现 GPU 内核。最近的进展显示了 大语言模型 (LLM) 自动生成 Triton 内核的潜力，减少了专家内核开发人员所需的手动工作。多个基准测试评估 LLM 生成的 Triton 内核。然而，它们面临三个关键限制：（1）它们将任务限制为 PyTorch 到 Triton 的转换，无法反映现实世界 Triton 任务的多样性和复杂性； （2）它们仅评估单个内核性能，而不是端到端性能，这是人工智能框架实际部署的核心标准； (3)它们依赖于针对各个内核手动编写的评估脚本，这些脚本可能包含模型可以利用的缺陷来绕过正确性检查并获得夸大的分数。为了解决这些限制，我们引入了 RealisticTritonBench，这是第一个从流行 AI 框架中的真实拉取请求导出 Triton 内核生成任务的基准，从而实现现实的、类似于生产的评估。 RealisticTritonBench 系统地从流行的开源 AI 框架中提取修改 Triton 内核的 PR，并将其转换为具有具体工程环境的生成任务。每个任务都以自然语言要求作为输入，并需要相应的 Triton 内核实现，以及完整且可重现的评估环境。与之前专注于独立内核性能的基准测试不同，RealisticTritonBench 将生成的内核集成到其原始框架中，并使用端到端测试对其进行评估，从而实现更忠实的评估。我们在 RealisticTritonBench 上评估了领先的 LLM，发现它们仍然难以完成现实世界的 Triton 内核生成任务。