请求合并 #47
File diff suppressed because it is too large
Load Diff
|
|
@ -1,887 +0,0 @@
|
|||
# FlashInfer 关键算子迁移与优化
|
||||
|
||||
## 1. 教程定位
|
||||
|
||||
本教程是赛题二FlashInfer任务的“benchmark 性能基线与 XPU-OJ 提交衔接”模块，主要帮助学员跑通目标算子的 benchmark 脚本，理解原库 API、输入输出结构、性能指标和性能基线结果，并进一步读懂 XPU-OJ 题目包中的接口约定、测试数据、参考输出和精度要求。
|
||||
|
||||
需要特别说明：
|
||||
|
||||
- 本教程不提供可直接提交的 OJ 参考实现源代码或标准答案代码。
|
||||
- 本教程仅提供冒烟级 starter 示例代码，用于验证环境、语言、提交链路和 `run_kernel(...)` 接口。
|
||||
- benchmark 脚本用于建立性能基线，不是最终提交物。
|
||||
- XPU-OJ 题包中的 `baseline()` 属于 OJ 后台参考实现，用于生成 `output_ref`，不是选手提交代码。
|
||||
- 选手最终需要自行实现 `run_kernel(...)`，并在正确性通过后继续优化性能。
|
||||
|
||||
完成本教程后，学员应能够跑通 benchmark 脚本，记录性能基线结果，读懂 XPU-OJ 题包，理解 OJ 的测试输入与参考实现，并完成一次冒烟级 OJ 提交。
|
||||
|
||||
## 2. 完成本模块你将能够
|
||||
|
||||
- 理解 FlashInfer Attention Kernel 的基本作用与适用场景；
|
||||
|
||||
- 完成 FlashInfer 环境、工具链的准备与源码编译；
|
||||
|
||||
- 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例；
|
||||
|
||||
- 完成各算子在不同参数配置下的 Benchmark 测试；
|
||||
|
||||
- 输出各算子的性能基线结果记录表，为后续算子优化提供对比基准。
|
||||
|
||||
- 理解 XPU-OJ 评测 `run_kernel` 接口与精度要求。
|
||||
|
||||
- 理解 benchmark 与 XPU-OJ 评测题包之间的关系，能够根据题包接口实现一个最小正确版 `run_kernel`。
|
||||
|
||||
|
||||
## 3. 适用对象
|
||||
|
||||
### 3.1 **适合人群**
|
||||
|
||||
* 参赛选手：需要完成 Benchmark 入门模块，为后续算子优化做准备
|
||||
* 软件开发者和Vibe Coding开发：希望从事AI相关行业开发，以及用智能体方式来做开发工作
|
||||
* LLM 推理开发者：希望了解 FlashInfer Attention Kernel 的性能表现
|
||||
* 算子优化工程师：希望基于MXMACA软件栈在沐曦国产 GPU 上做算子迁移和优化
|
||||
|
||||
### **3.2 前置基础**
|
||||
|
||||
* Python 基础：能够运行和修改 Python 脚本
|
||||
* PyTorch 基础：了解MXMACA软化栈的使用
|
||||
* Linux 命令行：能够使用终端执行命令
|
||||
* 了解 Attention 机制：理解 Q/K/V、KV Cache 等基本概念
|
||||
|
||||
|
||||
## 4. 前置准备
|
||||
|
||||
开始实战前，请确认你已经完成以下准备：
|
||||
|
||||
### 4.1 获得 GPU
|
||||
|
||||
1. [点击获取算力券](https://developer.metax-tech.com/activities/6)，首次登录需要使用邮箱或者手机号进行注册
|
||||
|
||||
2. 登录成功后验证邮箱
|
||||
|
||||
3. 提交申请获得兑换码
|
||||
|
||||
4. 兑换算力和登陆平台：
|
||||
|
||||
- [访问模力方舟官网](https://ai.gitee.com/)，进入费用中心 - 算力券 ，点击右上角 “兑换”。
|
||||
|
||||
- 进入算力容器，选择沐曦，租用算力，建议优先选 16G 显存 / 32G 显存，如下图
|
||||
|
||||
|
||||
|
||||
- 创建实例。基础镜像：`PyTorch-Agent/2.8.0/Python 3.12/maca 3.7.2.1`
|
||||
|
||||
|
||||
|
||||
- 选择工具-lab进入实例环境
|
||||
|
||||
- 步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用 `mx-smi` 命令查看
|
||||
|
||||
|
||||
### 4.2 Python 环境
|
||||
|
||||
``` bash
|
||||
pip install flashinfer torch pandas numpy
|
||||
```
|
||||
|
||||
### 4.3 OpenCode 安装
|
||||
|
||||
``` bash
|
||||
curl -fsSL https://opencode.ai/install | bash
|
||||
opencode
|
||||
```
|
||||
|
||||
### 4.4 代码准备
|
||||
|
||||
测试脚本和 Benchmark 脚本:
|
||||
|
||||
| 命令 | 说明 |
|
||||
| --- | --- |
|
||||
| `python bench_batch_decode.py` | 运行 Batch Decode 基准测试 |
|
||||
| `python bench_batch_prefill_paged.py` | 运行 Batch Prefill (Paged KV Cache) 基准测试 |
|
||||
| `python bench_batch_prefill_ragged.py` | 运行 Batch Prefill (Ragged KV Cache) 基准测试 |
|
||||
| `python bench_batch_mla.py` | 运行 MLA (Multi-head Latent Attention) 基准测试 |
|
||||
|
||||
## 5. 知识预备
|
||||
|
||||
### 5.1 LLM 推理阶段重要概念
|
||||
|
||||
- **Prefill 阶段**：Prefill 阶段是指处理输入 prompt 的阶段
|
||||
|
||||
- 输入：用户一次性给出的完整 prompt，长度为 seq\_len
|
||||
|
||||
- 计算：对 prompt 中的每个 token 并行计算注意力，生成第一个输出 token 及 KV cache
|
||||
|
||||
- 特点：这是 **计算密集型（compute-bound）** 阶段，因为需要做完整的 `seq\_len * seq\_len` 注意力矩阵乘法
|
||||
|
||||
- **Decode 阶段**：
|
||||
|
||||
- 每次只生成 1 个 token，利用 prefill 阶段填充好的 KV cache 做自回归生成
|
||||
|
||||
- **显存带宽密集型（memory-bound）**，瓶颈在从显存读取 KV cache 而非计算
|
||||
|
||||
**Prefill = 并行处理用户输入，Decode = 逐个生成回答 token**
|
||||
|
||||
## 6. 项目实践 -- FlashInfer Benchmark
|
||||
|
||||
**目标：** 在赛事镜像中完成 FlashInfer Ragged Prefill 算子的 Benchmark，理解从 Benchmark 到 XPU-OJ 评测提交的完整流程，为后续算子优化建立性能基线。
|
||||
|
||||
### 6.1 在赛事镜像中运行 FlashInfer Benchmark
|
||||
|
||||
#### Step 1：检查运行环境
|
||||
|
||||
**目标：** 确认当前环境满足本模块运行要求。
|
||||
|
||||
**操作：** 进入 Terminal 检查 GPU、Python、编译工具和依赖版本。
|
||||
|
||||
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```Bash
|
||||
# 检查沐曦 GPU 状态
|
||||
mx-smi
|
||||
|
||||
# 检查 Python 版本
|
||||
python --version
|
||||
|
||||
# 检查 PyTorch 是否能识别 GPU
|
||||
python -c "import torch; print(f'GPU available: {torch.cuda.is_available()}'); print(f'GPU count: {torch.cuda.device_count()}')"
|
||||
|
||||
# 检查依赖版本
|
||||
python -c "import torch; print(f'PyTorch {torch.__version__}')"
|
||||
python -c "import einops; print('einops OK')"
|
||||
# 安装必要依赖
|
||||
pip install pandas
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
- `mx-smi` 显示沐曦 GPU 信息
|
||||
|
||||
|
||||
|
||||
- Python 环境正常
|
||||
|
||||
|
||||
|
||||
- `torch.cuda.is_available()` 返回 `True`
|
||||
|
||||
|
||||
|
||||
- 所有依赖版本符合要求
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
**常见问题：**
|
||||
|
||||
| 问题 | 解决方法 |
|
||||
| --- | --- |
|
||||
| `mx-smi: command not found` | 确认已配置沐曦 GPU 驱动环境 |
|
||||
| `No GPUs are available` | 检查 MXMACA 驱动是否正确安装/环境变量是否正确配置 |
|
||||
| `ModuleNotFoundError: No module named 'xxx'` | `pip install xxx` |
|
||||
|
||||
#### Step 2：进入项目目录
|
||||
|
||||
**目标：** 进入本模块所需的源码目录 [flashinfer_baseline](https://www.gitlink.org.cn/metax-maca/op_optimization/tree/master/基于AI Agent开发范式的国产GPU大模型推理算子库优化%2Foperator_task_package%2Fflashinfer_task_package%2FFlashInfer_Baseline)。
|
||||
|
||||
1. 克隆代码仓库
|
||||
|
||||
```bash
|
||||
git clone https://gitlink.org.cn/metax-maca/op_optimization.git
|
||||
```
|
||||
|
||||
2. 准备 FlashInfer_Baseline
|
||||
|
||||
从仓库根目录开始，在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `operator_task_package/flashinfer_task_package` 文件夹。可以将 `FlashInfer_Baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
|
||||
```bash
|
||||
mkdir data
|
||||
cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package" data/
|
||||
```
|
||||
|
||||
3. 切换到项目目录 `FlashInfer_Baseline`
|
||||
```bash
|
||||
cd data/operator_task_package/flashinfer_task_package/FlashInfer_Baseline
|
||||
ls
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```plaintext
|
||||
bench_common.py
|
||||
bench_batch_decode.py
|
||||
bench_batch_prefill_paged.py
|
||||
bench_batch_prefill_ragged.py
|
||||
bench_batch_mla.py
|
||||
README.md
|
||||
...
|
||||
```
|
||||
|
||||
#### Step 3：验证项目脚本
|
||||
|
||||
**目标：** 确认所有基准测试脚本可正常执行。
|
||||
|
||||
**操作：** 检查脚本文件是否存在且可读。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
# 检查脚本文件
|
||||
python -c "import os; scripts = ['bench_common.py', 'bench_batch_decode.py', 'bench_batch_prefill_paged.py', 'bench_batch_prefill_ragged.py', 'bench_batch_mla.py']; [print(f'✓ {s}') if os.path.exists(s) else print(f'✗ {s} missing') for s in scripts]"
|
||||
|
||||
# 测试脚本导入
|
||||
python -c "from bench_common import setup_workspace, get_csv_path; print('脚本导入正常')"
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```plaintext
|
||||
✓ bench_common.py
|
||||
✓ bench_batch_decode.py
|
||||
✓ bench_batch_prefill_paged.py
|
||||
✓ bench_batch_prefill_ragged.py
|
||||
✓ bench_batch_mla.py
|
||||
脚本导入正常
|
||||
```
|
||||
|
||||
#### Step 4：运行单算子 Benchmark 并查看测试结果
|
||||
|
||||
**目标：** 执行基准测试，获取性能基线数据，查看并分析 Benchmark 输出结果。
|
||||
|
||||
**操作：** 运行 Ragged Prefill 基准测试脚本，读取生成的 CSV 结果文件。
|
||||
|
||||
**运行Benchmark命令示例：**
|
||||
|
||||
```bash
|
||||
python bench_batch_prefill_ragged.py
|
||||
```
|
||||
预期结果：（并非真实数据）
|
||||
``` plaintext
|
||||
[BatchPrefillWithRaggedKVCacheWrapper] Starting benchmark, total cases: 48
|
||||
[1/48] bs=1, sl=1024, hd=[128,128]: 0.032ms, 66.67 GB/s, 272.00 TFLOPs
|
||||
[2/48] bs=1, sl=4096, hd=[128,128]: 0.042ms, 197.83 GB/s, 3238.06 TFLOPs
|
||||
[3/48] bs=1, sl=8192, hd=[128,128]: 0.064ms, 261.50 GB/s, 8559.25 TFLOPs
|
||||
...
|
||||
|
||||
Results saved to BatchPrefillWithRaggedKVCacheWrapper_20260626_xxxxxx.csv
|
||||
```
|
||||
|
||||
**常见问题：**
|
||||
|
||||
| 问题 | 解决方法 |
|
||||
| --- | --- |
|
||||
| `out of memory` | 减小 batch\_size 或 seq\_len 参数 |
|
||||
| 运行时间过长 | 脚本会自动调整重复次数，耐心等待 |
|
||||
|
||||
***
|
||||
|
||||
**查看结果命令示例：**
|
||||
|
||||
```bash
|
||||
# 列出所有 CSV 结果文件（按修改时间排序，最新的在最上面）
|
||||
ls -lt *.csv 2>/dev/null || echo "未找到 CSV 文件，请先运行 benchmark"
|
||||
|
||||
# 使用 Python 查看最新结果（自动适配所有 benchmark 类型的列名）
|
||||
python3 -c "
|
||||
import pandas as pd, glob, os
|
||||
|
||||
# 找到所有 CSV 文件，按修改时间取最新的
|
||||
csv_files = sorted(glob.glob('*.csv'), key=os.path.getmtime, reverse=True)
|
||||
if not csv_files:
|
||||
print('未找到 CSV 文件，请先运行 benchmark 脚本')
|
||||
else:
|
||||
latest = csv_files[0]
|
||||
print(f'读取文件: {latest}')
|
||||
df = pd.read_csv(latest)
|
||||
# 动态选择列名：优先显示通用列 + 时间/性能列
|
||||
perf_cols = ['time_ms', 'bandwidth_GB_s', 'tflops']
|
||||
avail_cols = [c for c in df.columns if c in perf_cols or c not in ['api']]
|
||||
# 只保留有意义的分析列（排除 api、seq_len_q 等辅助列）
|
||||
display_cols = [c for c in avail_cols if c not in ('seq_len_q',)]
|
||||
print(df[display_cols].head(10).to_string(index=False))
|
||||
"
|
||||
```
|
||||
|
||||
**预期结果（以 Ragged Prefill 为例）：**
|
||||
|
||||
```plaintext
|
||||
读取文件: BatchPrefillWithRaggedKVCacheWrapper_20260626_145454.csv
|
||||
batch_size seq_len num_qo_heads num_kv_heads head_dim_qk head_dim_vo time_ms bandwidth_GB_s tflops
|
||||
1 1024 32 4 128 128 0.031580 66.666667 272.004150
|
||||
1 4096 32 4 128 128 0.042445 197.828709 3238.063402
|
||||
1 8192 32 4 128 128 0.064123 261.499213 8559.251770
|
||||
4 1024 32 4 128 128 0.050221 167.754590 2737.757998
|
||||
4 4096 32 4 128 128 0.101234 332.907816 21734.876630
|
||||
16 1024 32 4 128 128 0.149876 224.887654 14683.437981
|
||||
...
|
||||
```
|
||||
|
||||
### 6.2 XPU-OJ 在线评测教程
|
||||
|
||||
#### Step 5：从 Benchmark 到 XPU-OJ 提交
|
||||
|
||||
**Benchmark 与 XPU-OJ 的关系**
|
||||
|
||||
赛事镜像中的 Benchmark 和 XPU-OJ 在线评测任务不同。Benchmark 主要用于理解算子调用方式和建立性能基线；XPU-OJ 在线评测用于统一检查选手提交代码的正确性和性能。
|
||||
|
||||
| 维度 | Benchmark 脚本 | XPU-OJ 提交 |
|
||||
|------|---------------|------------|
|
||||
| **目的** | 理解算子接口、建立性能基线 | 统一环境下的正确性+性能评测 |
|
||||
| **接口形式** | Python API（`wrapper.plan()` + `wrapper.run()`） | C 接口（`extern "C" void run_kernel(...)`） |
|
||||
| **数据范围** | 多种 head_dim/batch_size/seq_len 组合 | 固定参数范围（以题包为准） |
|
||||
| **验证** | 无自动正确性校验 | 强制通过 `torch.allclose(rtol=1e-2, atol=1e-2)` |
|
||||
| **输出** | CSV 性能记录 | 排行榜得分 |
|
||||
|
||||
跑完 benchmark、建立性能基线后，选手需要完成以下转换：
|
||||
1. 从 benchmark 脚本中理解目标 API，例如 BatchPrefillWithRaggedKVCacheWrapper；
|
||||
2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口；
|
||||
3. 对照题包中的输入 shape、数据范围和精度要求；
|
||||
4. 编写自己的 `run_kernel(...)`；
|
||||
5. 提交 OJ，先通过正确性；
|
||||
6. 正确性通过后，再对比 benchmark 耗时 / OJ 耗时继续优化。
|
||||
|
||||
**选择目标算子**
|
||||
|
||||
FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benchmark 脚本、OJ 题包、`run_kernel(...)` 接口和数据范围可能不同。
|
||||
|
||||
| OJ 题号 | 算子类型 | 核心特点 | Benchmark 脚本 | FlashInfer API |
|
||||
|---------|---------------|----------------|----------|----------|
|
||||
| **1** | Ragged Prefill | GQA布局，Q/K/V平坦存储，causal=1 | `bench_batch_prefill_ragged.py` | `BatchPrefillWithRaggedKVCacheWrapper` |
|
||||
| **2** | Paged Prefill | KV Cache分页存储，需解析page table | `bench_batch_prefill_paged.py` | `BatchPrefillWithPagedKVCacheWrapper` |
|
||||
| **3** | MLA Paged Attention | DeepSeek MLA特有，双路Q(nope+pe)/双路Cache(ckv+kpe) |`bench_batch_mla.py` | `BatchMLAPagedAttentionWrapper` |
|
||||
| **4** | Paged Decode | 每次只1个query token，memory-bound |`bench_batch_decode.py` | `BatchDecodeWithPagedKVCacheWrapper` |
|
||||
|
||||
**每个子题的接口参数、数据范围和精度要求以对应 XPU-OJ 题包为准。** 下文以题目 **1 Ragged Prefill** 为例演示从 benchmark 到 XPU-OJ 提交的完整流程。
|
||||
|
||||
#### Step 6：理解 XPU-OJ 评测接口与精度要求
|
||||
**目标：** 明确 Benchmark 与最终评测提交之间的关系，理解选手需要实现的内容。
|
||||
|
||||
> 完成 benchmark 后，需要注意 benchmark 脚本主要用于建立性能基线，并不需要最终提交。
|
||||
> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 OJ 参考实现结果进行比较。
|
||||
|
||||
下面以 FlashInfer Ragged Prefill 题为例，其中：
|
||||
- `zh_CN/00_题目描述.md`：说明需要实现的算子功能；
|
||||
- `zh_CN/01_接口约定cuda.md`：说明必须实现cuda的 `run_kernel` 函数签名；
|
||||
- `zh_CN/01_接口约定triton.md`：说明必须实现triton的 `run_kernel` 函数签名；
|
||||
- `zh_CN/01_接口约定tilelang.md`：说明必须实现tilelang的 `run_kernel` 函数签名；
|
||||
- `zh_CN/02_数据范围.md`：说明测试范围和精度要求；
|
||||
- `testcase_config.py`：定义测试数据生成、OJ 参考实现和正确性校验方式。
|
||||
|
||||
FlashInfer Ragged Prefill 的校验方式为：
|
||||
```python
|
||||
torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
||||
```
|
||||
|
||||
选手实现的输出需要在上述容差范围内与 OJ 参考实现输出一致。不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。
|
||||
|
||||
#### Step 7：登录 XPU-OJ 并进入题目页面
|
||||
|
||||
使用组委会统一发放的账号登录 XPU-OJ，并进入对应赛题页面。
|
||||
|
||||
1. 打开 XPU-OJ 平台：https://xpuoj.com/
|
||||
2. 使用组委会统一发放的账号和初始密码登录 **【后续发布】**；
|
||||
|
||||
|
||||
3. 登录后进入比赛 / 题目列表页面；
|
||||
|
||||
4. 找到对应题目，例如 1 FlashInfer Ragged Prefill；
|
||||
|
||||
5. 点击进入题目详情页，查看题目描述、接口约定、数据范围和提交入口。
|
||||
|
||||
|
||||
#### Step 8：提交 OJ 冒烟代码
|
||||
**目标**：完成一次最小提交，确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。
|
||||
|
||||
1. 在语言下拉框中选择本题支持的提交语言，例如 `MXMACA C++`、`TileLang` 或后续开放的 `Triton`；
|
||||
2. 将实现了题目要求接口的代码复制到提交框中；
|
||||
> 如果你还没有 `run_kernel`，应该从哪里开始？
|
||||
>
|
||||
> - OJ 最终评测不会直接运行 benchmark 脚本，而是调用你提交代码中的 `run_kernel(...)`。
|
||||
> - 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。
|
||||
> - 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路。
|
||||
3. 借助 Agent 从题包生成 `run_kernel` 初版
|
||||
在下方参考 prompt 的引导下，Agent 会：
|
||||
1. 读取对应 OJ 题包中的接口约定文档（`01_接口约定.md`），提取 `run_kernel` 函数签名；
|
||||
2. 读取数据范围文档（`02_数据范围.md`），了解输入张量 shape 和精度要求；
|
||||
3. 生成一个能编译通过的最小 `run_kernel` 实现，优先保证接口正确性，不追求性能。
|
||||
|
||||
生成的代码可在本地编译验证后，直接提交到 OJ 上冒烟测试，确认提交链路可用。
|
||||
|
||||
**参考 prompt**
|
||||
|
||||
```plaintext
|
||||
请帮我为 FlashInfer Ragged Prefill 题（OJ 题号 20001）生成一个最小可运行的 run_kernel 实现，要求：
|
||||
1. 阅读题包中的 01_接口约定.md，理解 run_kernel 的函数签名（参数类型、顺序、const 修饰）；
|
||||
2. 阅读 02_数据范围.md，了解 head_dim_qk、head_dim_vo 的可能取值；
|
||||
3. 阅读 00_题目描述.md，理解需要实现的注意力计算逻辑；
|
||||
4. 生成一个只使用简单双重循环的 naive 实现（不加 tiling、不加 shared memory），确保：
|
||||
- 函数签名为 extern "C" void run_kernel(...)
|
||||
- 包含必要的头文件（cuda_bf16.h、cuda_runtime.h、stdint.h、math.h）
|
||||
- 支持 GQA（Group Query Attention）的头的映射
|
||||
- 支持 causal mask
|
||||
- 使用 bfloat16 数据类型
|
||||
- scale = 1/sqrt(head_dim_qk)
|
||||
```
|
||||
|
||||
|
||||
|
||||
**OJ 冒烟代码：**用于最小链路验证。
|
||||
|
||||
```cpp
|
||||
#include <stdint.h>
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <math.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
__device__ __forceinline__ float warp_sum(float x)
|
||||
{
|
||||
for (int offset = 16; offset > 0; offset >>= 1)
|
||||
{
|
||||
x += __shfl_down_sync(0xffffffffu, x, offset);
|
||||
}
|
||||
return __shfl_sync(0xffffffffu, x, 0);
|
||||
}
|
||||
|
||||
__global__ void ragged_prefill_smoke_kernel(
|
||||
const __nv_bfloat16 *__restrict__ q,
|
||||
const __nv_bfloat16 *__restrict__ k,
|
||||
const __nv_bfloat16 *__restrict__ v,
|
||||
__nv_bfloat16 *__restrict__ output,
|
||||
const int32_t *__restrict__ qo_indptr,
|
||||
const int32_t *__restrict__ kv_indptr,
|
||||
int64_t batch_size,
|
||||
int64_t seq_len,
|
||||
int64_t num_qo_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t head_dim_qk,
|
||||
int64_t head_dim_vo,
|
||||
int64_t causal)
|
||||
{
|
||||
const int lane = threadIdx.x & 31;
|
||||
const int warp_id = threadIdx.x >> 5;
|
||||
const int warps_per_block = blockDim.x >> 5;
|
||||
|
||||
int64_t work = static_cast<int64_t>(blockIdx.x) * warps_per_block + warp_id;
|
||||
const int64_t total = batch_size * seq_len * num_qo_heads;
|
||||
if (work >= total)
|
||||
return;
|
||||
|
||||
const int64_t qo_head = work % num_qo_heads;
|
||||
work /= num_qo_heads;
|
||||
const int64_t q_pos = work % seq_len;
|
||||
const int64_t batch = work / seq_len;
|
||||
|
||||
const int64_t qo_begin = qo_indptr[batch];
|
||||
const int64_t qo_len = qo_indptr[batch + 1] - qo_begin;
|
||||
if (q_pos >= qo_len)
|
||||
return;
|
||||
|
||||
const int64_t kv_begin = kv_indptr[batch];
|
||||
const int64_t kv_len = kv_indptr[batch + 1] - kv_begin;
|
||||
int64_t visible = kv_len;
|
||||
if (causal)
|
||||
{
|
||||
visible = kv_len - qo_len + q_pos + 1;
|
||||
if (visible < 0)
|
||||
visible = 0;
|
||||
if (visible > kv_len)
|
||||
visible = kv_len;
|
||||
}
|
||||
|
||||
const int64_t group = num_qo_heads / num_kv_heads;
|
||||
const int64_t kv_head = qo_head / group;
|
||||
const int64_t q_row = qo_begin + q_pos;
|
||||
const float scale = rsqrtf(static_cast<float>(head_dim_qk));
|
||||
|
||||
const __nv_bfloat16 *q_ptr = q + (q_row * num_qo_heads + qo_head) * head_dim_qk;
|
||||
float qv[4];
|
||||
float acc[4];
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
const int d = lane + i * 32;
|
||||
qv[i] = (d < head_dim_qk) ? __bfloat162float(q_ptr[d]) : 0.0f;
|
||||
acc[i] = 0.0f;
|
||||
}
|
||||
|
||||
float m = -1.0e20f;
|
||||
float l = 0.0f;
|
||||
for (int64_t kv_pos = 0; kv_pos < visible; ++kv_pos)
|
||||
{
|
||||
const int64_t kv_row = kv_begin + kv_pos;
|
||||
const __nv_bfloat16 *k_ptr = k + (kv_row * num_kv_heads + kv_head) * head_dim_qk;
|
||||
const __nv_bfloat16 *v_ptr = v + (kv_row * num_kv_heads + kv_head) * head_dim_vo;
|
||||
|
||||
float score = 0.0f;
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
const int d = lane + i * 32;
|
||||
if (d < head_dim_qk)
|
||||
{
|
||||
score += qv[i] * __bfloat162float(k_ptr[d]);
|
||||
}
|
||||
}
|
||||
score = warp_sum(score) * scale;
|
||||
|
||||
const float m_new = fmaxf(m, score);
|
||||
const float alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f;
|
||||
const float beta = __expf(score - m_new);
|
||||
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
const int d = lane + i * 32;
|
||||
if (d < head_dim_vo)
|
||||
{
|
||||
acc[i] = acc[i] * alpha + beta * __bfloat162float(v_ptr[d]);
|
||||
}
|
||||
}
|
||||
l = l * alpha + beta;
|
||||
m = m_new;
|
||||
}
|
||||
|
||||
__nv_bfloat16 *out_ptr = output + (q_row * num_qo_heads + qo_head) * head_dim_vo;
|
||||
const float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f;
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
const int d = lane + i * 32;
|
||||
if (d < head_dim_vo)
|
||||
{
|
||||
out_ptr[d] = __float2bfloat16(acc[i] * inv_l);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
extern "C" void run_kernel(
|
||||
const __nv_bfloat16 *q,
|
||||
const __nv_bfloat16 *k,
|
||||
const __nv_bfloat16 *v,
|
||||
__nv_bfloat16 *output,
|
||||
const int32_t *qo_indptr,
|
||||
const int32_t *kv_indptr,
|
||||
int64_t batch_size,
|
||||
int64_t seq_len,
|
||||
int64_t num_qo_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t head_dim_qk,
|
||||
int64_t head_dim_vo,
|
||||
int64_t causal)
|
||||
{
|
||||
constexpr int kThreads = 128;
|
||||
constexpr int kWarpsPerBlock = kThreads / 32;
|
||||
const int64_t total = batch_size * seq_len * num_qo_heads;
|
||||
const int blocks = static_cast<int>((total + kWarpsPerBlock - 1) / kWarpsPerBlock);
|
||||
ragged_prefill_smoke_kernel<<<blocks, kThreads>>>(
|
||||
q, k, v, output, qo_indptr, kv_indptr, batch_size, seq_len,
|
||||
num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal);
|
||||
}
|
||||
```
|
||||
以上代码仅用于说明接口结构，不代表最优实现，也不作为评分参考。
|
||||
|
||||
4. 点击提交，等待评测结果返回
|
||||
|
||||
|
||||
|
||||
评测时间与题目测试点数量、队列状态和平台负载有关，通常需要等待数十秒到数分钟。以平台实际返回为准。
|
||||
|
||||
**OJ 评测流程**：
|
||||
|
||||
1. 选手提交代码；
|
||||
2. 平台按所选语言编译或加载提交代码；
|
||||
3. 评测程序构造测试输入；
|
||||
4. 调用选手代码中的 `run_kernel(...)`；
|
||||
5. 将 `run_kernel(...)` 的输出与 `output_ref` 做正确性校验；
|
||||
6. 正确性通过后，统计运行耗时或性能指标；
|
||||
7. 更新该题历史最好成绩；
|
||||
8. 汇总各题最好成绩，得到排行榜总分。
|
||||
|
||||
5. 查看结果
|
||||
|
||||
OJ 平台的每个测试用例都会会返回两行结果信息（以 Problem 20001 Ragged Prefill 冒烟测试第 1 个测试用例为例）：
|
||||
|
||||
```
|
||||
OJCHAL v1 7P82PDM8iqVVD4hu6nNkKg==
|
||||
OJRESULT v1 cd1dc1cc5abf7d59552d9a82db40c83041527d232fcb7760f8d2048e76045220 eyJ0aW1lX21zIjo4LjEzNiwic3BlZWR1cCI6MC4wMTYyMjQsInRrX3RpbWVfbXMiOjguMTM2LCJ0Yl90aW1lX21zIjowLjEzMiwidGhfdGltZV9tcyI6MC4wMzU3OTEsInNjb3JlX3JhdGlvIjowLjAxMTczOCwicGFzcyI6dHJ1ZX0=
|
||||
```
|
||||
|
||||
- **OJCHAL**：挑战相关元信息（参赛者无需关注）；
|
||||
- **OJRESULT**：核心结果。分为两部分：
|
||||
- 前缀（`OJRESULT v1` 后的第一段 base64）：评测平台签名哈希，用于防篡改校验（参赛者无需关注）；
|
||||
- 后缀（第二段 base64）：JSON 格式的性能数据。
|
||||
|
||||
**下面演示如何提取性能数据：**
|
||||
|
||||
- 反编码 OJRESULT 获取性能数据
|
||||
|
||||
- 将上面的后缀 base64 复制下来，使用以下命令解码：
|
||||
|
||||
``` bash
|
||||
python3 -c "
|
||||
import base64, json
|
||||
|
||||
# 替换为你实际 OJRESULT 第二段 base64 字符串
|
||||
data = 'eyJ0aW1lX21zIjo4LjEzNiwic3BlZWR1cCI6MC4wMTYyMjQsInRrX3RpbWVfbXMiOjguMTM2LCJ0Yl90aW1lX21zIjowLjEzMiwidGhfdGltZV9tcyI6MC4wMzU3OTEsInNjb3JlX3JhdGlvIjowLjAxMTczOCwicGFzcyI6dHJ1ZX0='
|
||||
|
||||
# 解码并格式化输出
|
||||
decoded = base64.b64decode(data.encode('utf-8')).decode('utf-8')
|
||||
result = json.loads(decoded)
|
||||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||||
"
|
||||
```
|
||||
|
||||
- **解码输出如下**
|
||||
|
||||
``` json
|
||||
{
|
||||
"time_ms": 8.136,
|
||||
"speedup": 0.016224,
|
||||
"tk_time_ms": 8.136,
|
||||
"tb_time_ms": 0.132,
|
||||
"th_time_ms": 0.035791,
|
||||
"score_ratio": 0.011738,
|
||||
"pass": true
|
||||
}
|
||||
```
|
||||
|
||||
- **字段含义**
|
||||
|
||||
| 字段 | 含义 | 说明 |
|
||||
|------|------|------|
|
||||
| `pass` | 正确性校验 | `true` 表示输出通过了 `torch.allclose(rtol=1e-2, atol=1e-2)` 校验，可以参与排名 |
|
||||
| `tb_time_ms` | OJ 参考实现耗时 (ms) | OJ 内置参考实现在该测试用例上的 GPU 执行时间 |
|
||||
| `tk_time_ms` | 提交耗时 (ms) | 你提交的 `run_kernel` 在该测试用例上的 GPU 执行时间 |
|
||||
| `speedup` | 加速比 | `speedup = tb_time_ms / tk_time_ms`，**>1 表示优于 OJ 参考实现，<1 表示慢于 OJ 参考实现** |
|
||||
| `time_ms` | 总耗时 (ms) | 通常等于 `tk_time_ms`，即 kernel 执行总时间 |
|
||||
| `th_time_ms` | Host 辅助耗时 (ms) | host 端数据准备、kernel launch、结果拷贝等开销 |
|
||||
| `score_ratio` | 归一化得分 | 综合正确性和加速比的加权评分（0~1），用于排行榜排名 |
|
||||
|
||||
- **加速比解读与冒烟测试预期**
|
||||
|
||||
以冒烟测试为例，上述结果中：
|
||||
|
||||
- `tb_time_ms = 0.132ms`（OJ 参考实现）
|
||||
- `tk_time_ms = 8.136ms`（冒烟 kernel）
|
||||
- `speedup = 0.016224`，即约慢 62 倍
|
||||
|
||||
冒烟测试加速比 <1 是完全正常的。后续优化目标是逐步提升加速比，直至超越 OJ 参考实现。
|
||||
|
||||
加速比计算公式：
|
||||
|
||||
```
|
||||
加速比 = tb_time_ms / tk_time_ms
|
||||
= OJ 参考实现耗时 / 你的 kernel 耗时
|
||||
```
|
||||
|
||||
例：若 `tb_time_ms = 0.132`，优化后 `tk_time_ms = 0.080`，则 `speedup = 0.132 / 0.080 = 1.65`，即比 OJ 参考实现快 1.65 倍。
|
||||
|
||||
- **查看所有测试用例结果**
|
||||
|
||||
OJ 平台每次提交会评测所有测试用例，每个测试用例返回一组独立的 `OJRESULT`。建议将每组结果反编码后整理成表格，方便跟踪优化进展：
|
||||
|
||||
| 测试用例 | batch_size | seq_len | tb_time_ms | tk_time_ms | speedup | pass |
|
||||
|----------|------------|---------|------------|------------|---------|------|
|
||||
| 1 | 1 | 1024 | — | — | — | — |
|
||||
| 2 | 4 | 1024 | 0.132 | 8.136 | 0.016 | true |
|
||||
| 3 | 16 | 1024 | — | — | — | — |
|
||||
| ... | ... | ... | ... | ... | ... | ... |
|
||||
|
||||
* **与榜单的关系**
|
||||
* OJ 平台根据各测试用例的 `score_ratio` 加权计算总分；
|
||||
* 榜单排名以加权总分为依据，`pass=false` 的提交不参与排名。
|
||||
|
||||
## 7. Agent 使用样例
|
||||
|
||||
**目标：** 在本模块中，Agent 可以帮助你完成环境检查、运行 Benchmark、分析结果、理解 OJ 接口、生成和调试 `run_kernel` 等任务。以下是参考 prompt。
|
||||
|
||||
### 7.1 环境检查
|
||||
|
||||
```plaintext
|
||||
请帮我检查当前环境是否满足 FlashInfer 运行要求，包括：
|
||||
1. 沐曦 GPU 是否可见（mx-smi）
|
||||
2. Python 版本和 PyTorch CUDA 支持
|
||||
3. flashinfer、pandas、numpy 依赖是否已安装
|
||||
```
|
||||
|
||||
### 7.2 运行 Benchmark
|
||||
|
||||
``` plaintext
|
||||
请帮我运行 bench_batch_prefill_ragged.py 脚本，执行 Ragged Prefill 的基准测试。
|
||||
```
|
||||
|
||||
### 7.3 分析结果
|
||||
|
||||
``` plaintext
|
||||
请帮我读取最新的性能基线 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置，并与沐曦 GPU 理论峰值带宽做对比。
|
||||
```
|
||||
|
||||
### 7.4 理解 OJ 题包接口
|
||||
|
||||
``` plaintext
|
||||
请帮我阅读 FlashInfer Ragged Prefill 题包（problem_20001）中的以下文件：
|
||||
- zh_CN/00_题目描述.md
|
||||
- zh_CN/01_接口约定cuda.md
|
||||
- zh_CN/01_接口约定triton.md
|
||||
- zh_CN/01_接口约定tilelang.md
|
||||
- zh_CN/02_数据范围.md
|
||||
然后帮我总结：
|
||||
1. run_kernel 的函数签名和每个参数的含义
|
||||
2. 输入张量的形状约定（q/k/v 的 layout、indptr 的作用）
|
||||
3. head_dim_qk 和 head_dim_vo 的可能取值
|
||||
4. 精度要求（rtol/atol）
|
||||
5. causal=1 时需要注意的边界条件
|
||||
```
|
||||
|
||||
### 7.5 生成 `run_kernel` 初版
|
||||
|
||||
``` plaintext
|
||||
请帮我为 FlashInfer Ragged Prefill 题（OJ 题号 20001）生成一个最小可运行的 run_kernel 实现，要求：
|
||||
1. 阅读题包中的 01_接口约定.md，理解 run_kernel 的函数签名（参数类型、顺序、const 修饰）；
|
||||
2. 阅读 02_数据范围.md，了解 head_dim_qk、head_dim_vo 的可能取值；
|
||||
3. 生成一个只使用简单双重循环的 naive 实现（不加 tiling、不加 shared memory），确保：
|
||||
- 函数签名为 extern "C" void run_kernel(...)
|
||||
- 包含必要的头文件（cuda_bf16.h、cuda_runtime.h、stdint.h、math.h）
|
||||
- 支持 GQA（Group Query Attention）的头映射：hkv = hq * num_kv_heads / num_qo_heads
|
||||
- 支持 causal mask
|
||||
- 使用 bfloat16 数据类型
|
||||
- scale = 1.0f / sqrtf(head_dim_qk)
|
||||
```
|
||||
|
||||
### 7.6 调试 OJ 提交错误
|
||||
``` plaintext
|
||||
我的 run_kernel 提交到 OJ 后显示 Wrong Answer，请帮我对比以下信息：
|
||||
1. 这是我的 submit 代码：[粘贴你的 run_kernel 实现]
|
||||
2. 题包的接口约定在这里：[粘贴或引用 01_接口约定.md]
|
||||
3. 题包的测试配置在这里：[粘贴或引用 testcase_config.py]
|
||||
请帮我逐项检查：
|
||||
- 函数签名是否完全匹配
|
||||
- GQA 头映射公式是否正确
|
||||
- causal mask 边界条件是否正确
|
||||
- float4 向量化加载的偏移是否正确
|
||||
- online softmax 的 m/l 更新逻辑是否正确
|
||||
```
|
||||
|
||||
### 7.7 问题排查
|
||||
|
||||
``` plaintext
|
||||
运行 bench_batch_prefill_ragged.py 时报错 out of memory，请帮我分析原因并给出解决方案。
|
||||
代码理解
|
||||
请帮我解释 bench_batch_prefill_ragged.py 中 BatchPrefillWithRaggedKVCacheWrapper 的 plan() 和 run() 方法的工作原理，特别是 qo_indptr 和 kv_indptr 的作用。
|
||||
整理优化日志
|
||||
请帮我整理本次优化的记录，包括：
|
||||
1. 原始 benchmark 结果性能数据（从 CSV 中提取关键配置的 time_ms 和 tflops）
|
||||
2. 优化后的性能数据（从 OJ 评测结果中提取）
|
||||
3. 加速比 = benchmark_time / optimized_time
|
||||
4. 以表格形式输出：配置参数 | benchmark 耗时 | 优化后耗时 | 加速比
|
||||
```
|
||||
|
||||
## 8. 常见问题
|
||||
|
||||
### 8.1 流程问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| 登录后看不到题目 | 未使用赛用账号登录 | 七月份组委会统一发放 XPU-OJ 账号，请确认你使用的是组委会统一发放的账号，而不是自行注册账号；如仍无法看到题目，请联系助教或赛事运营确认账号权限。 |
|
||||
|
||||
### 8.2 环境问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| `No GPUs are available` | MXMACA 驱动未安装或 GPU 不可见 | 检查驱动安装，运行 `python -c "import torch; print(torch.cuda.device_count())"` 验证 |
|
||||
| `ModuleNotFoundError: No module named 'flashinfer'` | flashinfer 未安装 | 执行 `pip install flashinfer` |
|
||||
| `out of memory` | GPU 显存不足 | 减小 `batch_size` 或 `seq_len` 参数 |
|
||||
|
||||
### 8.3 运行问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| Benchmark 运行时间过长 | 参数组合过多， workload 较大 | 耐心等待，脚本会自动调整重复次数 |
|
||||
| `KeyError: 'BatchPrefillWithPagedKVCacheKernel'` | profiler 未捕获目标 kernel | 检查 `target_kernels` 配置是否正确 |
|
||||
| CSV 文件为空 | 测试未正常完成 | 检查 GPU 显存是否充足，重新运行 |
|
||||
|
||||
### 8.4 代码问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| `ImportError: cannot import name 'xxx' from 'bench_common'` | 函数名拼写错误 | 检查 `bench_common.py` 中的函数名 |
|
||||
| `RuntimeError: error: device-side assert triggered` | 输入参数超出范围 | 检查 `num_qo_heads`、`num_kv_heads`、`head_dim` 配置 |
|
||||
|
||||
### 8.5 性能问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| TFLOPs 数值异常低 | 工作负载过小，kernel 启动开销占比大 | 增大 `batch_size` 或 `seq_len` |
|
||||
| 带宽数值异常低 | 数据未正确加载到 GPU | 检查 Tensor 是否在 CUDA 设备上 |
|
||||
|
||||
### 8.6 评测问题
|
||||
|
||||
提交 XPU-OJ 后可能遇到的异常评测结果及排查方向：
|
||||
|
||||
| 问题 | 可能原因 | 解决办法 |
|
||||
|------|------|---------|
|
||||
| **Compilation Error** 编译错误 | 1. `run_kernel` 签名与 OJ 接口约定不一致（参数类型、顺序、数量不匹配）；2. 缺少 `extern "C"` 声明导致 C++ name mangling；3. 缺少必要头文件（`cuda_bf16.h`、`cuda_runtime.h`、`math.h`）；3. 使用了 OJ 环境不支持的语法或 API | 1. 逐行对照对应题包的「01_接口约定.md」，确认参数类型（`int64_t` vs `int`、`const` 修饰）、顺序完全一致；2. 在 `run_kernel` 前加 `extern "C"`；3. 确认文件顶部 include 了 `<cuda_bf16.h>`、`<cuda_runtime.h>`、`<stdint.h>`、`<math.h>`；4. 去掉 `printf`、`assert` 等调试代码后重新提交 |
|
||||
| **Time Limit Exceeded** 运行超时 | 1. `run_kernel` 内部调用了 `cudaDeviceSynchronize()` 导致额外等待；2. kernel 中存在死循环（for 循环边界条件错误）；3. `__syncthreads()` 放在条件分支内导致线程死锁；4. grid 配置过大，启动的 block 数量远超合理范围 | 1. 删除 `run_kernel` 函数体内的 `cudaDeviceSynchronize()` 调用——评测器会在外部自行同步；2. 检查 kernel 中所有 for 循环的终止条件，确保 `kv_start <= block_max_q` 等边界正确；3. 将所有 `__syncthreads()` 移到 if/else 分支之外；4. 检查 grid 计算：`(seq_len + Br - 1) / Br`，确认 `Br` 取值合理 |
|
||||
| **Wrong Answer** 答案错误 | 1. 注意力计算公式错误（score、scale、softmax 实现有偏差）；2. GQA 头映射错误：`hkv = hq / (num_qo_heads / num_kv_heads)` 计算不对；3. Causal mask 未正确实现（`causal=1` 时 query 看到了不该看的未来 token）；4. Online softmax 的 m/l 更新逻辑有误；5. float4 向量化加载的偏移计算错误，导致 K/V 数据错位；6. 输出写入偏移错误，或对无效位置写了垃圾值 | 1. 本地用题包中的 PyTorch 参考实现对拍：运行 `testcase_config.py` 的 OJ 参考实现 `baseline()` 与你 kernel 输出做 `torch.allclose(rtol=1e-2, atol=1e-2)` 比对；2. GQA 公式：`int hkv = hq * num_kv_heads / num_qo_heads`（整数除法）；3. Causal 逻辑：`kv_end = min(kv_start + Bc, q_idx + 1)`，注意 +1 的处理；4. 对照论文 FlashAttention 的 Algorithm 1 逐行验证 online softmax；5. float4 加载偏移公式：`(cur_kv_start + i) * num_kv_heads * D_QK + hkv * D_QK + d_idx * 8`，确认 `num_kv_heads` 而非 `num_qo_heads` |
|
||||
|
||||
|
||||
## 9. 下一步学习建议
|
||||
|
||||
### 9.1 保存你的性能基线结果
|
||||
|
||||
将本次运行生成的 CSV 文件妥善保存，后续优化时需要以此作为对比基准。
|
||||
|
||||
```bash
|
||||
# 建议创建 results 目录保存
|
||||
mkdir -p results
|
||||
mv *.csv results/
|
||||
```
|
||||
|
||||
### 9.2 深入理解 FlashInfer 核心概念
|
||||
|
||||
* 阅读 FlashInfer 官方文档，理解 Paged KV Cache、Ragged KV Cache 的设计理念
|
||||
|
||||
* 学习 MLA (Multi-head Latent Attention) 的原理，了解 DeepSeek 的注意力优化方案
|
||||
|
||||
* 理解 `plan()` 和 `run()` 两阶段设计的作用
|
||||
|
||||
|
||||
### 9.3 进入算子优化模块
|
||||
|
||||
参考后续优化模块，学习以下优化技术：
|
||||
|
||||
* **Kernel Tuning**：调整 Block Size、Thread Count 等参数
|
||||
|
||||
* **Memory Optimization**：减少显存占用、优化数据搬运
|
||||
|
||||
* **Compute Optimization**：提升计算效率
|
||||
|
||||
|
||||
### 9.4 参考资源
|
||||
|
||||
* FlashInfer 官方仓库：https://github.com/flashinfer-ai/flashinfer
|
||||
|
||||
* FlashInfer 文档：https://flashinfer.ai
|
||||
|
||||
|
||||
### 9.5 记录优化流程
|
||||
|
||||
建议维护一份优化日志，记录每次优化的改动和性能变化：
|
||||
|
||||
| 优化项 | 改动内容 | 性能基线 | 优化后 | 提升比例 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% |
|
||||
|
||||
完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。
|
||||
|
||||
> 使用 Agent 整理优化日志，可形成可复现的 Agent/Skill 优化流程
|
||||
|
||||
### 9.6 使用多语言完成算子优化加速
|
||||
|
||||
可以使用 Triton 或 TileLang 语言实现 `run_kernel` 接口完成算子优化，对比不同的语言对于性能加速的影响。
|
||||
Loading…
Reference in New Issue