算子

算子#

KernelGenBench 中算子来源的参考列表。

ATen 算子（110 个）#

选择标准#

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

  • 均匀采样 60 个长尾算子

示例#

类别

算子

算术

add, div, floor_divide, cos, sin

归约

sum, mean, argmax, amin

线性

matmul、linear、bmm

归一化

_softmax, softmax

激活

hardsigmoid, prelu, heaviside

形状操作

expand, repeat, view

vLLM 算子（50 个）#

类别#

类别

描述

注意力

PagedAttention 变体

KV 缓存

缓存管理内核

量化

FP8、AWQ 内核

归一化

RMS 归一化、融合归一化

示例#

算子

描述

paged_attention_v1

内存高效注意力

fused_add_rms_norm

融合归一化

rotary_embedding

位置编码

cuBLAS 算子（50 个）#

GEMM 系列（14 种变体）#

精度

标准

StridedBatched

Batched

Float32

cublasSgemm_v2

cublasSgemmStridedBatched

cublasSgemmBatched_64

Float64

—

cublasDgemmStridedBatched

cublasDgemmBatched

Complex64

cublasCgemm_v2

cublasCgemmStridedBatched

—

Complex128

—

cublasZgemmStridedBatched

cublasZgemmBatched

Float16

—

cublasHgemmStridedBatched

cublasHgemmBatched

其他 BLAS 例程#

  • GEMV（矩阵-向量乘法）

  • SYRK（对称秩-k 更新）

  • TRSM（三角求解）

  • 以及跨多种精度的其他例程

算子命名约定#

来源

前缀

示例

ATen

aten::

aten::add.Tensor

vLLM

vllm13::

vllm13::rms_norm

cuBLAS

cublas::

cublas::cublasSgemm_v2