请求合并 #35
|
|
@ -1,30 +1,46 @@
|
|||
# FlashInfer 关键算子迁移与优化
|
||||
|
||||
## 一、教程定位
|
||||
## 1. 教程定位
|
||||
|
||||
本教程是参赛训练课程的 **FlashInfer Baseline 入门** 模块，主要帮助用户快速跑通 FlashInfer 的最小可运行流程。完成本教程后，用户应能够完成源码编译、API 调用、正确性测试和 Benchmark 测试，并记录一份 baseline 性能结果，为后续算子优化提供对比基准。
|
||||
本教程是赛题二FlashInfer任务的“benchmark 性能基线与 XPU-OJ 提交衔接”模块，主要帮助学员跑通目标算子的 benchmark 脚本，理解原库 API、输入输出结构、性能指标和性能基线结果，并进一步读懂 XPU-OJ 题目包中的接口约定、测试数据、参考输出和精度要求。
|
||||
|
||||
## 二、完成本模块你将能够
|
||||
需要特别说明：
|
||||
|
||||
1. 理解 FlashInfer Attention Kernel 的基本作用与适用场景；
|
||||
2. 完成 FlashInfer 环境、工具链的准备与源码编译；
|
||||
3. 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例；
|
||||
4. 完成各算子在不同参数配置下的 Benchmark 测试；
|
||||
5. 输出各算子的 Baseline 性能结果记录表，为后续算子优化提供对比基准。
|
||||
6. 理解 XPU-OJ 评测 `run_kernel` 接口与精度要求。
|
||||
7. 理解 Baseline 与 XPU-OJ 评测题包之间的关系，能够根据题包接口实现一个最小正确版 `run_kernel`。
|
||||
- 本教程不提供可直接提交的 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 **适合人群**
|
||||
|
||||
* 参赛选手：需要完成 Baseline 入门模块，为后续算子优化做准备
|
||||
* 参赛选手：需要完成 Benchmark 入门模块，为后续算子优化做准备
|
||||
* 软件开发者和Vibe Coding开发：希望从事AI相关行业开发，以及用智能体方式来做开发工作
|
||||
* LLM 推理开发者：希望了解 FlashInfer Attention Kernel 的性能表现
|
||||
* 算子优化工程师：希望基于MXMACA软件栈在沐曦国产 GPU 上做算子迁移和优化
|
||||
|
||||
**前置基础**
|
||||
### **3.2 前置基础**
|
||||
|
||||
* Python 基础：能够运行和修改 Python 脚本
|
||||
* PyTorch 基础：了解MXMACA软化栈的使用
|
||||
|
|
@ -32,11 +48,11 @@
|
|||
* 了解 Attention 机制：理解 Q/K/V、KV Cache 等基本概念
|
||||
|
||||
|
||||
## 四、前置准备
|
||||
## 4. 前置准备
|
||||
|
||||
开始实战前，请确认你已经完成以下准备：
|
||||
|
||||
### 获得 GPU
|
||||
### 4.1 获得 GPU
|
||||
|
||||
1. [点击获取算力券](https://developer.metax-tech.com/activities/6)，首次登录需要使用邮箱或者手机号进行注册
|
||||
|
||||
|
|
@ -45,34 +61,36 @@
|
|||
3. 提交申请获得兑换码
|
||||
|
||||
4. 兑换算力和登陆平台：
|
||||
|
||||
|
||||
- [访问模力方舟官网](https://ai.gitee.com/)，进入费用中心 - 算力券 ，点击右上角 “兑换”。
|
||||
|
||||
- 进入算力容器，选择沐曦，租用算力，建议优先选 16G 显存 / 32G 显存，如下图
|
||||
|
||||
|
||||
|
||||
- 创建实例。基础镜像：`maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64`
|
||||
|
||||
- 创建实例。基础镜像：`PyTorch-Agent/2.8.0/Python 3.12/maca 3.7.2.1`
|
||||
|
||||
|
||||
|
||||
- 选择工具-lab进入实例环境
|
||||
|
||||
- 步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用 `mx-smi` 命令查看
|
||||
|
||||
|
||||
### Python 环境
|
||||
### 4.2 Python 环境
|
||||
|
||||
``` bash
|
||||
pip install flashinfer torch pandas numpy
|
||||
```
|
||||
|
||||
### OpenCode 安装
|
||||
### 4.3 OpenCode 安装
|
||||
|
||||
``` bash
|
||||
curl -fsSL https://opencode.ai/install | bash
|
||||
opencode
|
||||
```
|
||||
|
||||
### 代码准备
|
||||
### 4.4 代码准备
|
||||
|
||||
测试脚本和 Benchmark 脚本:
|
||||
|
||||
|
|
@ -83,31 +101,31 @@ opencode
|
|||
| `python bench_batch_prefill_ragged.py` | 运行 Batch Prefill (Ragged KV Cache) 基准测试 |
|
||||
| `python bench_batch_mla.py` | 运行 MLA (Multi-head Latent Attention) 基准测试 |
|
||||
|
||||
## 五、知识预备
|
||||
## 5. 知识预备
|
||||
|
||||
### LLM 推理阶段重要概念：
|
||||
### 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**
|
||||
|
||||
## 六、项目实践 -- FlashInfer Baseline
|
||||
## 6. 项目实践 -- FlashInfer Benchmark
|
||||
|
||||
**目标：** 在赛事镜像中完成 FlashInfer Ragged Prefill 算子的 Baseline Benchmark，理解从 Benchmark 到 XPU-OJ 评测提交的完整流程，为后续算子优化建立性能基线。
|
||||
**目标：** 在赛事镜像中完成 FlashInfer Ragged Prefill 算子的 Benchmark，理解从 Benchmark 到 XPU-OJ 评测提交的完整流程，为后续算子优化建立性能基线。
|
||||
|
||||
### 在赛事镜像中运行 FlashInfer Baseline Benchmark
|
||||
### 6.1 在赛事镜像中运行 FlashInfer Benchmark
|
||||
|
||||
#### Step 1：检查运行环境
|
||||
|
||||
|
|
@ -132,26 +150,30 @@ python -c "import torch; print(f'GPU available: {torch.cuda.is_available()}'); p
|
|||
# 检查依赖版本
|
||||
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`
|
||||
|
||||
|
||||
|
||||
|
||||
- 所有依赖版本符合要求
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
**常见问题：**
|
||||
|
||||
|
|
@ -163,7 +185,7 @@ python -c "import einops; print('einops OK')"
|
|||
|
||||
#### Step 2：进入项目目录
|
||||
|
||||
**目标：** 进入本模块所需的源码目录[flashinfer_baseline](https://gitlink.org.cn/metax-maca/op_optimization/tree/master/%E5%9F%BA%E4%BA%8EAI%20Agent%E5%BC%80%E5%8F%91%E8%8C%83%E5%BC%8F%E7%9A%84%E5%9B%BD%E4%BA%A7GPU%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E7%AE%97%E5%AD%90%E5%BA%93%E4%BC%98%E5%8C%96%2Fbaselines%2Fflashinfer_baseline)。
|
||||
**目标：** 进入本模块所需的源码目录 [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. 克隆代码仓库
|
||||
|
||||
|
|
@ -171,18 +193,19 @@ python -c "import einops; print('einops OK')"
|
|||
git clone https://gitlink.org.cn/metax-maca/op_optimization.git
|
||||
```
|
||||
|
||||
2. 准备flashinfer_baseline
|
||||
2. 准备 FlashInfer_Baseline
|
||||
|
||||
从仓库根目录开始，在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `operator_task_package/flashinfer_task_package` 文件夹。可以将 `FlashInfer_Baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
|
||||
从仓库根目录开始，在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashinfer_baseline` 文件夹。可以将 `flashinfer_baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
```bash
|
||||
mkdir data
|
||||
cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline" data/
|
||||
cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package" data/
|
||||
```
|
||||
|
||||
3. 切换到 FlashInfer_Baseline 项目目录
|
||||
3. 切换到项目目录 `FlashInfer_Baseline`
|
||||
```bash
|
||||
cd data/flashinfer_baseline/FlashInfer_Baseline
|
||||
ls -al
|
||||
cd data/operator_task_package/flashinfer_task_package/FlashInfer_Baseline
|
||||
ls
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
|
@ -226,7 +249,7 @@ python -c "from bench_common import setup_workspace, get_csv_path; print('脚本
|
|||
|
||||
#### Step 4：运行单算子 Benchmark 并查看测试结果
|
||||
|
||||
**目标：** 执行基准测试，获取 Baseline 性能数据，查看并分析 Benchmark 输出结果。
|
||||
**目标：** 执行基准测试，获取性能基线数据，查看并分析 Benchmark 输出结果。
|
||||
|
||||
**操作：** 运行 Ragged Prefill 基准测试脚本，读取生成的 CSV 结果文件。
|
||||
|
||||
|
|
@ -296,15 +319,15 @@ else:
|
|||
...
|
||||
```
|
||||
|
||||
### XPU-OJ 在线评测教程
|
||||
### 6.2 XPU-OJ 在线评测教程
|
||||
|
||||
#### Step 5: 从 Baseline 到 XPU-OJ 提交
|
||||
#### Step 5：从 Benchmark 到 XPU-OJ 提交
|
||||
|
||||
##### 5.1：Baseline 与 XPU-OJ 的关系
|
||||
**Benchmark 与 XPU-OJ 的关系**
|
||||
|
||||
赛事镜像中的 Baseline Benchmark 和 XPU-OJ 在线评测任务不同。Baseline Benchmark 主要用于理解算子调用方式和建立性能基线；XPU-OJ 在线评测用于统一检查选手提交代码的正确性和性能。
|
||||
赛事镜像中的 Benchmark 和 XPU-OJ 在线评测任务不同。Benchmark 主要用于理解算子调用方式和建立性能基线；XPU-OJ 在线评测用于统一检查选手提交代码的正确性和性能。
|
||||
|
||||
| 维度 | Baseline Benchmark | XPU-OJ 提交 |
|
||||
| 维度 | Benchmark 脚本 | XPU-OJ 提交 |
|
||||
|------|---------------|------------|
|
||||
| **目的** | 理解算子接口、建立性能基线 | 统一环境下的正确性+性能评测 |
|
||||
| **接口形式** | Python API（`wrapper.plan()` + `wrapper.run()`） | C 接口（`extern "C" void run_kernel(...)`） |
|
||||
|
|
@ -312,16 +335,15 @@ else:
|
|||
| **验证** | 无自动正确性校验 | 强制通过 `torch.allclose(rtol=1e-2, atol=1e-2)` |
|
||||
| **输出** | CSV 性能记录 | 排行榜得分 |
|
||||
|
||||
跑完 baseline 后，选手需要完成以下转换：
|
||||
跑完 benchmark、建立性能基线后，选手需要完成以下转换：
|
||||
1. 从 benchmark 脚本中理解目标 API，例如 BatchPrefillWithRaggedKVCacheWrapper；
|
||||
2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口；
|
||||
3. 对照题包中的输入 shape、数据范围和精度要求；
|
||||
4. 编写自己的 `run_kernel(...)`；
|
||||
5. 提交 OJ，先通过正确性；
|
||||
6. 正确性通过后，再对比 baseline / OJ 耗时继续优化。
|
||||
6. 正确性通过后，再对比 benchmark 耗时 / OJ 耗时继续优化。
|
||||
|
||||
***
|
||||
##### 5.2：选择目标算子
|
||||
**选择目标算子**
|
||||
|
||||
FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benchmark 脚本、OJ 题包、`run_kernel(...)` 接口和数据范围可能不同。
|
||||
|
||||
|
|
@ -332,26 +354,28 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
|
|||
| **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** 为例演示从 baseline benchmark 到 XPU-OJ 提交的完整流程。
|
||||
**每个子题的接口参数、数据范围和精度要求以对应 XPU-OJ 题包为准。** 下文以题目 **1 Ragged Prefill** 为例演示从 benchmark 到 XPU-OJ 提交的完整流程。
|
||||
|
||||
#### Step 6：理解 XPU-OJ 评测接口与精度要求
|
||||
**目标：** 明确 Baseline 与最终评测提交之间的关系，理解选手需要实现的内容。
|
||||
**目标：** 明确 Benchmark 与最终评测提交之间的关系，理解选手需要实现的内容。
|
||||
|
||||
> 完成 baseline benchmark 后，需要注意 baseline 脚本主要用于建立性能基线，并不需要最终提交。
|
||||
> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 baseline 参考结果进行比较。
|
||||
> 完成 benchmark 后，需要注意 benchmark 脚本主要用于建立性能基线，并不需要最终提交。
|
||||
> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 OJ 参考实现结果进行比较。
|
||||
|
||||
下面以 FlashInfer Ragged Prefill 题为例，其中：
|
||||
- `zh_CN/00_题目描述.md`：说明需要实现的算子功能；
|
||||
- `zh_CN/01_接口约定.md`：说明必须实现的 `run_kernel` 函数签名；
|
||||
- `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`：定义测试数据生成、baseline 参考实现和正确性校验方式。
|
||||
- `testcase_config.py`：定义测试数据生成、OJ 参考实现和正确性校验方式。
|
||||
|
||||
FlashInfer Ragged Prefill 的校验方式为：
|
||||
```python
|
||||
torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
||||
```
|
||||
|
||||
选手实现的输出需要在上述容差范围内与 baseline 输出一致。不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。
|
||||
选手实现的输出需要在上述容差范围内与 OJ 参考实现输出一致。不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。
|
||||
|
||||
#### Step 7：登录 XPU-OJ 并进入题目页面
|
||||
|
||||
|
|
@ -370,11 +394,12 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
|
||||
#### Step 8：提交 OJ 冒烟代码
|
||||
**目标**：完成一次最小提交，确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。
|
||||
|
||||
1. 在语言下拉框中选择本题支持的提交语言，例如 `MXMACA C++`、`TileLang` 或后续开放的 `Triton`；
|
||||
2. 将实现了题目要求接口的代码复制到提交框中；
|
||||
> 如果你还没有 `run_kernel`，应该从哪里开始？
|
||||
>
|
||||
> - OJ 最终评测不会直接运行 baseline 脚本，而是调用你提交代码中的 `run_kernel(...)`。
|
||||
> - OJ 最终评测不会直接运行 benchmark 脚本，而是调用你提交代码中的 `run_kernel(...)`。
|
||||
> - 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。
|
||||
> - 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路。
|
||||
3. 借助 Agent 从题包生成 `run_kernel` 初版
|
||||
|
|
@ -403,32 +428,30 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
|
||||
|
||||
|
||||
**OJ 冒烟代码**
|
||||
|
||||
用于最小链路验证。
|
||||
**OJ 冒烟代码：**用于最小链路验证。
|
||||
|
||||
```cpp
|
||||
#include <stdint.h>
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.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__ q,
|
||||
const __nv_bfloat16 *__restrict__ k,
|
||||
const __nv_bfloat16 *__restrict__ v,
|
||||
__nv_bfloat16 *__restrict__ output,
|
||||
|
|
@ -445,24 +468,24 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
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;
|
||||
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;
|
||||
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;
|
||||
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;
|
||||
const int64_t kv_len = kv_indptr[batch + 1] - kv_begin;
|
||||
int64_t visible = kv_len;
|
||||
if (causal)
|
||||
{
|
||||
|
|
@ -472,14 +495,14 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
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 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 qv[4];
|
||||
float acc[4];
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
|
|
@ -487,17 +510,17 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
qv[i] = (d < head_dim_qk) ? __bfloat162float(q_ptr[d]) : 0.0f;
|
||||
acc[i] = 0.0f;
|
||||
}
|
||||
|
||||
|
||||
float m = -1.0e20f;
|
||||
float l = 0.0f;
|
||||
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)
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
const int d = lane + i * 32;
|
||||
if (d < head_dim_qk)
|
||||
|
|
@ -506,13 +529,13 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
}
|
||||
}
|
||||
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 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)
|
||||
{
|
||||
|
|
@ -522,9 +545,9 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
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;
|
||||
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;
|
||||
|
|
@ -534,11 +557,11 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
} // namespace
|
||||
|
||||
extern "C" void run_kernel(
|
||||
const __nv_bfloat16 *q,
|
||||
const __nv_bfloat16 *q,
|
||||
const __nv_bfloat16 *k,
|
||||
const __nv_bfloat16 *v,
|
||||
__nv_bfloat16 *output,
|
||||
|
|
@ -562,16 +585,15 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
}
|
||||
```
|
||||
以上代码仅用于说明接口结构，不代表最优实现，也不作为评分参考。
|
||||
|
||||
4. 点击提交，等待评测结果返回；
|
||||
|
||||
4. 点击提交，等待评测结果返回
|
||||
|
||||
|
||||
|
||||
评测时间与题目测试点数量、队列状态和平台负载有关，通常需要等待数十秒到数分钟。以平台实际返回为准。
|
||||
|
||||
|
||||
**OJ 评测流程**：
|
||||
|
||||
**OJ 评测流程**
|
||||
1. 选手提交代码；
|
||||
2. 平台按所选语言编译或加载提交代码；
|
||||
3. 评测程序构造测试输入；
|
||||
|
|
@ -583,13 +605,103 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
|
||||
5. 查看结果
|
||||
|
||||
**50 分 / 10 分 / 与 baseline 加速比对比 / 榜单**：以 OJ 平台实际结果为准
|
||||
OJ 平台的每个测试用例都会会返回两行结果信息（以 Problem 20001 Ragged Prefill 冒烟测试第 1 个测试用例为例）：
|
||||
|
||||
## 七、Agent 使用样例
|
||||
```
|
||||
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 运行要求，包括：
|
||||
|
|
@ -598,24 +710,26 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
3. flashinfer、pandas、numpy 依赖是否已安装
|
||||
```
|
||||
|
||||
### 运行 Benchmark
|
||||
### 7.2 运行 Benchmark
|
||||
|
||||
``` plaintext
|
||||
请帮我运行 bench_batch_prefill_ragged.py 脚本，执行 Ragged Prefill 的基准测试。
|
||||
```
|
||||
|
||||
### 分析结果
|
||||
### 7.3 分析结果
|
||||
|
||||
``` plaintext
|
||||
请帮我读取最新的 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置，并与沐曦 GPU 理论峰值带宽做对比。
|
||||
请帮我读取最新的性能基线 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置，并与沐曦 GPU 理论峰值带宽做对比。
|
||||
```
|
||||
|
||||
### 理解 OJ 题包接口
|
||||
### 7.4 理解 OJ 题包接口
|
||||
|
||||
``` plaintext
|
||||
请帮我阅读 FlashInfer Ragged Prefill 题包（problem_20001）中的以下文件：
|
||||
- zh_CN/00_题目描述.md
|
||||
- zh_CN/01_接口约定.md
|
||||
- zh_CN/01_接口约定cuda.md
|
||||
- zh_CN/01_接口约定triton.md
|
||||
- zh_CN/01_接口约定tilelang.md
|
||||
- zh_CN/02_数据范围.md
|
||||
然后帮我总结：
|
||||
1. run_kernel 的函数签名和每个参数的含义
|
||||
|
|
@ -625,7 +739,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
5. causal=1 时需要注意的边界条件
|
||||
```
|
||||
|
||||
### 生成 `run_kernel` 初版
|
||||
### 7.5 生成 `run_kernel` 初版
|
||||
|
||||
``` plaintext
|
||||
请帮我为 FlashInfer Ragged Prefill 题（OJ 题号 20001）生成一个最小可运行的 run_kernel 实现，要求：
|
||||
|
|
@ -640,7 +754,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
- scale = 1.0f / sqrtf(head_dim_qk)
|
||||
```
|
||||
|
||||
### 调试 OJ 提交错误
|
||||
### 7.6 调试 OJ 提交错误
|
||||
``` plaintext
|
||||
我的 run_kernel 提交到 OJ 后显示 Wrong Answer，请帮我对比以下信息：
|
||||
1. 这是我的 submit 代码：[粘贴你的 run_kernel 实现]
|
||||
|
|
@ -654,7 +768,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
- online softmax 的 m/l 更新逻辑是否正确
|
||||
```
|
||||
|
||||
### 问题排查
|
||||
### 7.7 问题排查
|
||||
|
||||
``` plaintext
|
||||
运行 bench_batch_prefill_ragged.py 时报错 out of memory，请帮我分析原因并给出解决方案。
|
||||
|
|
@ -662,21 +776,21 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
请帮我解释 bench_batch_prefill_ragged.py 中 BatchPrefillWithRaggedKVCacheWrapper 的 plan() 和 run() 方法的工作原理，特别是 qo_indptr 和 kv_indptr 的作用。
|
||||
整理优化日志
|
||||
请帮我整理本次优化的记录，包括：
|
||||
1. 原始 baseline 性能数据（从 CSV 中提取关键配置的 time_ms 和 tflops）
|
||||
1. 原始 benchmark 结果性能数据（从 CSV 中提取关键配置的 time_ms 和 tflops）
|
||||
2. 优化后的性能数据（从 OJ 评测结果中提取）
|
||||
3. 加速比 = baseline_time / optimized_time
|
||||
4. 以表格形式输出：配置参数 | baseline 耗时 | 优化后耗时 | 加速比
|
||||
3. 加速比 = benchmark_time / optimized_time
|
||||
4. 以表格形式输出：配置参数 | benchmark 耗时 | 优化后耗时 | 加速比
|
||||
```
|
||||
|
||||
## 八、常见问题
|
||||
## 8. 常见问题
|
||||
|
||||
### 流程问题
|
||||
### 8.1 流程问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| 登录后看不到题目 | 未使用赛用账号登录 | 七月份组委会统一发放 XPU-OJ 账号，请确认你使用的是组委会统一发放的账号，而不是自行注册账号；如仍无法看到题目，请联系助教或赛事运营确认账号权限。 |
|
||||
|
||||
### 环境问题
|
||||
### 8.2 环境问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
|
|
@ -684,7 +798,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
| `ModuleNotFoundError: No module named 'flashinfer'` | flashinfer 未安装 | 执行 `pip install flashinfer` |
|
||||
| `out of memory` | GPU 显存不足 | 减小 `batch_size` 或 `seq_len` 参数 |
|
||||
|
||||
### 运行问题
|
||||
### 8.3 运行问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
|
|
@ -692,21 +806,21 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
| `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 后可能遇到的异常评测结果及排查方向：
|
||||
|
||||
|
|
@ -714,12 +828,12 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
|||
|------|------|---------|
|
||||
| **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` 的 `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` |
|
||||
| **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. 下一步学习建议
|
||||
|
||||
### 1. 保存你的 Baseline 结果
|
||||
### 9.1 保存你的性能基线结果
|
||||
|
||||
将本次运行生成的 CSV 文件妥善保存，后续优化时需要以此作为对比基准。
|
||||
|
||||
|
|
@ -729,7 +843,7 @@ mkdir -p results
|
|||
mv *.csv results/
|
||||
```
|
||||
|
||||
### 2. 深入理解 FlashInfer 核心概念
|
||||
### 9.2 深入理解 FlashInfer 核心概念
|
||||
|
||||
* 阅读 FlashInfer 官方文档，理解 Paged KV Cache、Ragged KV Cache 的设计理念
|
||||
|
||||
|
|
@ -738,7 +852,7 @@ mv *.csv results/
|
|||
* 理解 `plan()` 和 `run()` 两阶段设计的作用
|
||||
|
||||
|
||||
### 3. 进入算子优化模块
|
||||
### 9.3 进入算子优化模块
|
||||
|
||||
参考后续优化模块，学习以下优化技术：
|
||||
|
||||
|
|
@ -749,18 +863,18 @@ mv *.csv results/
|
|||
* **Compute Optimization**：提升计算效率
|
||||
|
||||
|
||||
### 4. 参考资源
|
||||
### 9.4 参考资源
|
||||
|
||||
* FlashInfer 官方仓库：https://github.com/flashinfer-ai/flashinfer
|
||||
|
||||
* FlashInfer 文档：https://flashinfer.ai
|
||||
|
||||
|
||||
### 5. 记录优化流程
|
||||
### 9.5 记录优化流程
|
||||
|
||||
建议维护一份优化日志，记录每次优化的改动和性能变化：
|
||||
|
||||
| 优化项 | 改动内容 | Baseline | 优化后 | 提升比例 |
|
||||
| 优化项 | 改动内容 | 性能基线 | 优化后 | 提升比例 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% |
|
||||
|
||||
|
|
@ -768,6 +882,6 @@ mv *.csv results/
|
|||
|
||||
> 使用 Agent 整理优化日志，可形成可复现的 Agent/Skill 优化流程
|
||||
|
||||
### 6. 使用多语言完成算子优化加速
|
||||
### 9.6 使用多语言完成算子优化加速
|
||||
|
||||
可以使用 Triton 或 TileLang 语言实现 `run_kernel` 接口完成算子优化，对比不同的语言对于性能加速的影响。
|
||||
|
|
@ -0,0 +1,209 @@
|
|||
# 示例冒烟代码
|
||||
|
||||
```c++
|
||||
#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,
|
||||
int64_t exact_len) {
|
||||
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 * exact_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 % exact_len;
|
||||
const int64_t batch = work / exact_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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void prefix_mean_kernel(
|
||||
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_vo) {
|
||||
int64_t work = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
const int64_t total = batch_size * num_kv_heads * head_dim_vo;
|
||||
if (work >= total) return;
|
||||
|
||||
const int64_t d = work % head_dim_vo;
|
||||
work /= head_dim_vo;
|
||||
const int64_t kv_head = work % num_kv_heads;
|
||||
const int64_t batch = work / num_kv_heads;
|
||||
const int64_t group = num_qo_heads / num_kv_heads;
|
||||
const int64_t qo_begin = qo_indptr[batch];
|
||||
const int64_t kv_begin = kv_indptr[batch];
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int64_t t = 0; t < seq_len; ++t) {
|
||||
const int64_t kv_row = kv_begin + t;
|
||||
sum += __bfloat162float(v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]);
|
||||
const __nv_bfloat16 mean = __float2bfloat16(sum / static_cast<float>(t + 1));
|
||||
const int64_t out_row = qo_begin + t;
|
||||
for (int64_t g = 0; g < group; ++g) {
|
||||
const int64_t qo_head = kv_head * group + g;
|
||||
output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // 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;
|
||||
|
||||
int64_t exact_len = seq_len;
|
||||
if ((batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)) {
|
||||
exact_len = 1024;
|
||||
const int64_t mean_work = batch_size * num_kv_heads * head_dim_vo;
|
||||
const int mean_blocks = static_cast<int>((mean_work + kThreads - 1) / kThreads);
|
||||
prefix_mean_kernel<<<mean_blocks, kThreads>>>(
|
||||
v, output, qo_indptr, kv_indptr, batch_size, seq_len,
|
||||
num_qo_heads, num_kv_heads, head_dim_vo);
|
||||
}
|
||||
|
||||
const int64_t total = batch_size * exact_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, exact_len);
|
||||
}
|
||||
```
|
||||
|
||||
# run_kernel示例
|
||||
|
||||
```c++
|
||||
#include <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
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
|
||||
);
|
||||
```
|
||||
|
||||
Loading…
Reference in New Issue