Skip to content

Commit 8c76754

Browse files
baominghellybaominghelly
committed
Init benchmark script
1 parent 1400daf commit 8c76754

3 files changed

Lines changed: 3345 additions & 0 deletions

File tree

‎scripts/BENCHMARK.md‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
# InfiniOps 性能评估脚本
2+
3+
独立于 pytest 的系统化性能评估工具，测量 InfiniOps 算子与 PyTorch 参考实现的耗时对比。
4+
5+
## 快速开始
6+
7+
```bash
8+
# 列出可用算子和设备
9+
python scripts/benchmark.py --list
10+
11+
# 快速验证单个算子
12+
python scripts/benchmark.py --ops add --device cpu --mode quick --no-json
13+
14+
# GPU 上跑全部 ntops 算子（推荐）
15+
python scripts/benchmark.py --category ntops --device cuda --mode standard --output ntops_results.json
16+
```
17+
18+
## 命令行参数
19+
20+
| 参数 | 说明 | 默认值 |
21+
|------|------|--------|
22+
| `--ops` | 指定算子名（可多个） | 全部可用算子 |
23+
| `--category` | `native` / `torch` / `ntops` / `all` | `all` |
24+
| `--device` | 设备（可多个，如 `cuda cpu`） | 全部可用设备 |
25+
| `--dtype` | 数据类型：`float16` `bfloat16` `float32` | fp16 + bf16 |
26+
| `--mode` | `quick` / `standard` / `thorough` | `standard` |
27+
| `--warmup` | Warmup 轮数 | `3` |
28+
| `--min-time` | 最小测量时间（秒） | `0.1` |
29+
| `--output` | JSON 输出路径 | 无（不输出 JSON） |
30+
| `--json-only` | 仅输出 JSON，不打印表格 | 关闭 |
31+
| `--no-json` | 仅打印表格，不输出 JSON | 关闭 |
32+
| `--list` | 列出可用算子并退出 | 关闭 |
33+
34+
## 三种模式
35+
36+
| 模式 | Shape 数量 | Dtype | 用途 |
37+
|------|-----------|-------|------|
38+
| `quick` | 每算子 1 个 | fp16, bf16 | 快速验证流程 |
39+
| `standard` | 2-3 个（LLaMA decode/prefill） | fp16, bf16 | 日常性能评估（推荐） |
40+
| `thorough` | 4-7 个（全量 LLM shape） | fp16, bf16, fp32 | 完整性能报告 |
41+
42+
Shape 基于 LLaMA-7B / 70B 模型的真实推理维度（hidden=4096/8192, FFN=11008/28672 等）。默认 shape 足够大（64M 元素），确保 kernel 计算时间远大于 pybind11 调度开销。
43+
44+
## 算子类别
45+
46+
### Native（手写原生算子，14 个）
47+
48+
有专用 C++/CUDA 后端实现的算子：
49+
50+
`add` `mul` `cast` `cat` `gemm` `matmul` `linear` `rms_norm` `causal_softmax` `swiglu` `flash_attention` `rotary_embedding` `add_rms_norm` `reshape_and_cache`
51+
52+
### ntops（ATen Fallback 算子，50 个）
53+
54+
通过 ATen `_out` API（slot=8）分发的 LLM 推理常用算子，与 PyTorch 原生算子性能对比：
55+
56+
**逐元素一元**：`abs` `neg` `exp` `rsqrt` `sigmoid` `silu` `gelu` `tanh` `sin` `cos` `bitwise_not` `softmax` `sqrt` `reciprocal` `log_softmax` `ceil` `floor` `log` `sign` `round` `hardtanh`
57+
58+
**逐元素二元**：`add` `sub` `div` `pow` `eq` `ne` `lt` `le` `gt` `ge` `bitwise_and` `bitwise_or` `maximum` `minimum`
59+
60+
**归约**：`sum` `mean` `cumsum` `amax` `amin` `argmax`
61+
62+
**矩阵运算**：`mm` `bmm` `addmm`
63+
64+
**排序/索引**：`topk` `sort` `gather` `index_select`
65+
66+
**归一化**：`rms_norm`
67+
68+
**池化**：`avg_pool2d`
69+
70+
### Torch（全部 PyTorch Fallback 算子，500+ 个）
71+
72+
基于 `torch_ops_metadata.json` 的完整 ATen fallback 算子集。
73+
74+
## 性能对比方法
75+
76+
### 测量方式
77+
78+
使用 `torch.utils.benchmark.Timer.blocked_autorange()` 进行多次测量取中位数。InfiniOps 和 PyTorch ref 均使用 `out=` 预分配输出 API，避免 tensor 分配开销影响对比。
79+
80+
### 调度路径差异
81+
82+
- **InfiniOps**：Python → pybind11（~15-20us）→ C++ dispatch → ATen kernel
83+
- **PyTorch**：Python → 原生 C binding（~5us）→ ATen kernel
84+
85+
两者底层调用相同的 ATen CUDA kernel，差异仅在前端调度开销。当 kernel 计算时间远大于 ~20us 时，speedup 趋近 1.00x。
86+
87+
### 已知特殊情况
88+
89+
| 算子 | 现象 | 原因 |
90+
|------|------|------|
91+
| `silu` `hardtanh` | speedup ~1.9x | PyTorch Python API 不支持 `out=`，ref 多了 tensor 分配+copy 开销 |
92+
| `add` | speedup <1.0x | 走 native slot 0 自定义后端，op cache + descriptor 开销比 ATen slot 8 重 |
93+
| `rms_norm` | speedup <1.0x | 走 native slot 0 自定义后端，原生 kernel 在大 tensor 下性能不如 PyTorch |
94+
95+
## 使用示例
96+
97+
```bash
98+
# ntops 算子性能对比（推荐用法）
99+
python scripts/benchmark.py --category ntops --device cuda --mode standard --output ntops_results.json
100+
101+
# 单个算子快速验证
102+
python scripts/benchmark.py --ops sigmoid mm --device cuda --dtype float16 --mode quick --no-json
103+
104+
# 仅原生算子
105+
python scripts/benchmark.py --category native --device cuda --mode standard --output native.json
106+
107+
# 全量测试（耗时较长）
108+
python scripts/benchmark.py --mode thorough --device cuda --output full.json
109+
110+
# 调整测量参数
111+
python scripts/benchmark.py --category ntops --warmup 5 --min-time 0.3 --output results.json
112+
```
113+
114+
## 输出格式
115+
116+
### 控制台表格
117+
118+
```
119+
==========================================================================================================
120+
InfiniOps Performance Benchmark | Device: cuda | Mode: quick
121+
==========================================================================================================
122+
Category Operator Shape dtype InfiniOps(us) PyTorch(us) Speedup
123+
----------------------------------------------------------------------------------------------------------
124+
ntops abs [8192,8192] float16 158.75 158.24 1.00x
125+
ntops silu [8192,8192] float16 174.29 329.52 1.89x
126+
ntops mm [128,4096]x[4096,11008] float16 104.60 104.41 1.00x
127+
ntops rms_norm [8192,8192] float16 1.05 2.33 0.45x
128+
----------------------------------------------------------------------------------------------------------
129+
Summary: 50 benchmarks (0 native + 0 torch ops) | Avg speedup: 1.01x
130+
==========================================================================================================
131+
```
132+
133+
### JSON 文件
134+
135+
```json
136+
{
137+
"metadata": {
138+
"timestamp": "2026-05-26T...",
139+
"device": ["cuda"],
140+
"mode": "standard",
141+
"torch_version": "2.6.0+cu124"
142+
},
143+
"results": [
144+
{
145+
"category": "ntops",
146+
"operator": "abs",
147+
"device": "cuda",
148+
"dtype": "torch.float16",
149+
"shape_description": "[8192,8192]",
150+
"infiniops_median_us": 158.75,
151+
"reference_median_us": 158.24,
152+
"speedup": 1.00,
153+
"status": "ok"
154+
}
155+
],
156+
"summary": {
157+
"total": 50,
158+
"ok": 50,
159+
"avg_speedup": 1.01
160+
}
161+
}
162+
```
163+
164+
## 错误处理
165+
166+
- **算子不可用**：自动跳过，标记 `SKIP`
167+
- **dtype 不支持**：捕获 RuntimeError，跳过
168+
- **GPU OOM**：捕获异常，跳过该配置
169+
- **Metadata 不存在**：仅 benchmark 原生算子
170+
171+
## 文件说明
172+
173+
| 文件 | 说明 |
174+
|------|------|
175+
| `scripts/benchmark.py` | 主脚本（CLI、测量引擎、算子 setup、输出格式化） |
176+
| `scripts/benchmark_configs.py` | 原生算子 LLM shape 配置、FLOPS 计算、标量默认值 |
177+
178+
## 依赖
179+
180+
- `torch`（含 `torch.utils.benchmark`）
181+
- `infiniops`
182+
- `tests/utils.py`（复用 `randn_strided`、`empty_strided`、`get_stream` 等工具函数）
183+
184+
## 前置条件
185+
186+
ntops 算子需要 `feat/torch-operator-bases` 分支并启用 `WITH_TORCH` 编译：
187+
188+
```bash
189+
git checkout feat/torch-operator-bases
190+
CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) pip install . --no-build-isolation \
191+
-C cmake.define.WITH_CPU=ON \
192+
-C cmake.define.WITH_NVIDIA=ON
193+
```

0 commit comments

Comments
 (0)