多源算子#
KernelGenBench 评估来自三个不同来源的 210 个算子的内核生成能力,每个来源代表不同的复杂度和实际应用场景。
概述#
ATen 算子#
PyTorch ATen 算子是深度学习框架中使用的核心计算构建块。
选择标准#
从 2,907 个开源模型训练轨迹中选择使用频率最高的前 50 个算子
均匀采样 60 个长尾算子
总计:从 900+ ATen API 中选择 110 个算子
示例#
softmax、matmul、embedding、cumsum、add.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,绕过高级封装。