多源算子#

KernelGenBench 评估来自三个不同来源的 210 个算子的内核生成能力,每个来源代表不同的复杂度和实际应用场景。

概述#

来源

算子数量

描述

ATen

110

PyTorch 原生算子

vLLM

50

大语言模型推理内核

cuBLAS

50

闭源库重实现

ATen 算子#

PyTorch ATen 算子是深度学习框架中使用的核心计算构建块。

选择标准#

  • 从 2,907 个开源模型训练轨迹中选择使用频率最高的前 50 个算子

  • 均匀采样 60 个长尾算子

  • 总计:从 900+ ATen API 中选择 110 个算子

示例#

softmaxmatmulembeddingcumsumadd.Tensor

提示词构建#

  • 动态提取 FunctionSchema

  • 官方 PyTorch 文档字符串

  • 所有重载变体作为独立问题

基准#

PyTorch 原生 C++ 实现

vLLM 算子#

来自 vLLM(v0.13.0)的生产级 大语言模型 推理内核。

覆盖范围#

  • 注意力机制(PagedAttention v1)

  • KV 缓存管理

  • 混合精度量化(FP8/AWQ)

挑战#

复杂的内存布局管理和算法逻辑使功能正确性极具挑战。

目标#

验证生成实际推理加速内核的能力。

cuBLAS 算子#

针对 cuBLAS(v12.4)的闭源库重实现。

选择策略#

  • 通过性能分析轨迹选择调用最频繁的前 10 个例程

  • 扩展到不同精度(S/D/C/Z/H)和批处理模式

  • 对多样化 BLAS 例程进行策略性采样

API 变体#

单个 GEMM 系列产生 14 个独立问题:

精度

标准

StridedBatched

Batched

64 位索引

Float32

cublasSgemm

Float64

Complex64

Complex128

Float16

挑战#

达到数十年专家手工优化的性能极为困难。

基准#

通过 ctypes.cdll 直接加载 libcublas.so,绕过高级封装。