forked from metax-maca/op_optimization
Compare commits
31 Commits
beckylu-de
...
master
| Author | SHA1 | Date |
|---|---|---|
 Beckylu
|
64b7883e5d | |
yyyymmm
|
7b9de194ab | |
|
Beckylu
|
57f56aabb8 | |
|
Beckylu
|
2d38aa3d1b | |
|
xiao
|
8e6231b6ab | |
|
Beckylu
|
6d623763d6 | |
|
Ke Xiao (i26293)
|
1dbb422167 | |
|
Ke Xiao (i26293)
|
790db292d7 | |
|
何沐川
|
12b2af4460 | |
|
Beckylu
|
2e60758a4f | |
|
xiao
|
d4ac7f5a2f | |
|
Ke Xiao (i26293)
|
398ec687c3 | |
|
yyyymmm
|
6b2e4a709c | |
|
yyyymmm
|
cc3d597feb | |
|
yyyymmm
|
a452e70d15 | |
|
Beckylu
|
59a4fb8f8c | |
|
xiaoke-shu
|
f302a83b62 | |
|
Beckylu
|
dc5841a10c | |
|
yyyymmm
|
167b72b2e1 | |
|
yyyymmm
|
567c8c42a0 | |
|
yyyymmm
|
8070486586 | |
|
yyyymmm
|
7c13dd377f | |
|
yyyymmm
|
b85e2637b6 | |
|
yyyymmm
|
cc05079507 | |
|
yyyymmm
|
0385d36047 | |
|
yyyymmm
|
ea98b9a069 | |
|
yyyymmm
|
4e045a3e7d | |
|
yyyymmm
|
6382da7296 | |
|
Beckylu
|
53880f68c3 | |
|
xiao
|
392362ba88 | |
|
Beckylu
|
9f27203d0d |
10
README.md
10
README.md
|
|
@ -18,12 +18,14 @@
|
|||
|
||||
初赛攻坚 Fused Moe Gemm、决赛冲刺 MLA/NSA，全程基于全国产技术栈开发优化，打造自主可控的低成本大模型推理算子底座。
|
||||
|
||||
|
||||
**赛题一相关资料**
|
||||
|
||||
- [赛题一方案：基于国产软件栈的推理前沿算子优化比赛方案](基于国产软件栈大模型推理前沿算子优化/基于国产软件栈的推理前沿算子优化比赛方案.md)
|
||||
- [TileLang MACA 构建指南（模力方舟）](基于国产软件栈大模型推理前沿算子优化/tilelang_maca_build_guide_模力方舟.md)
|
||||
- [TileLang 算子 sample 跑通测试指南](基于国产软件栈大模型推理前沿算子优化/race_tests_run_guide基于tilelang算子sample跑通测试.md)
|
||||
- [赛题一国产软件栈样例说明](基于国产软件栈大模型推理前沿算子优化/race-1-domestic-software-stack-samples.md)
|
||||
- [赛题一国产软件栈样例说明](基于国产软件栈大模型推理前沿算子优化/基于国产软件栈的大模型推理前沿算子优化（TileLang）Sample%20算子获取说明.md)
|
||||
- [TileLang 训练营学习资料](https://www.gitlink.org.cn/ccf-ai-infra/Intro-ops/about)
|
||||
|
||||
### 赛题二：基于 AI Agent 开发范式的国产 GPU 大模型推理算子库优化
|
||||
|
||||
|
|
@ -47,6 +49,8 @@
|
|||
- [模力方舟 Agent 部署准备教程](基于AI%20Agent开发范式的国产GPU大模型推理算子库优化/模力方舟Agent部署准备教程.md)
|
||||
- [赛题二说明及资料参考](基于AI%20Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md)
|
||||
|
||||
###**两个赛题统一使用模力方舟上的镜像PyTorch-Agent / 2.8.0 / Python 3.12 / maca 3.7.2.1**
|
||||
|
||||
## 参赛对象
|
||||
|
||||
面向全国全日制专科、本科、硕博在校生（非在职），40 周岁以下青年科技人才均可参与；
|
||||
|
|
@ -99,6 +103,8 @@
|
|||
|
||||
沐曦开发者社区活动页面，完成新人礼任务，提前熟悉 C500 在线算力使用。https://developer.metax-tech.com/activities/6
|
||||
|
||||
学生可通过登录启悟社区后跳转沐曦开发者社区领取100元算力代金券。(https://developer.metax-tech.com/activities/11)
|
||||
|
||||
报名赛事后还可额外获得300元算力代金券：https://developer.metax-tech.com/activities/17
|
||||
|
||||
使用指南：[模力方舟快速使用 SOP](模力方舟快速使用SOP.md)
|
||||
|
|
@ -109,7 +115,7 @@
|
|||
|
||||
加入沐曦股份“揭榜挂帅”赛题，深耕国产 GPU 算子优化，用技术降低每 Token 推理成本，共建自主可控 AI 算力新生态！
|
||||
|
||||
👇 报名通道：登录挑战杯官网 [www.tiaozhanbei.net](https://link.wtturl.cn/?target=https%3A%2F%2Fwww.tiaozhanbei.net&scene=im&aid=497858&lang=zh) 揭榜挂帅入口报名
|
||||
👇 报名通道：登录挑战杯官网[2026年度中国青年科技创新“揭榜挂帅”擂台赛（学生赛道）](https://2026.tiaozhanbei.net/)揭榜挂帅入口报名
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,20 +1,26 @@
|
|||
# FlashAttention Baseline 入门：从环境验证到 KV-Cache Benchmark
|
||||
# Flashattention 迁移 Baseline 实战：从性能基线到 XPU-OJ 评测
|
||||
|
||||
## 一、教程定位
|
||||
|
||||
本教程是参赛训练课程的 FlashAttention Baseline 入门模块，主要帮助用户快速跑通 FlashAttention 的 KV-Cache 推理性能基准测试最小可运行流程。
|
||||
本教程是参赛训练课程的 **FlashAttention Baseline 入门与评测提交衔接** 模块，主要帮助用户跑通 FlashAttention paged KV‑cache 推理核函数 `flash_attn_with_kvcache` 的基准测试流程，理解 baseline 的输入输出、性能指标和评测含义，并基于 XPU‑OJ 题包完成一个最小正确版 `run_kernel` 的实现与提交。
|
||||
|
||||
完成本教程后，用户应能够：
|
||||
需要特别说明：本教程中的 Baseline 主要用于帮助参赛者理解目标算子的调用方式、输入输出结构和性能基线，Baseline 不是最终提交物。最终评测以 XPU‑OJ 题包为准，参赛者需要根据题包中的接口约定，实现自己的 `run_kernel`，并在输出结果对齐 baseline 参考结果的前提下提升性能。
|
||||
|
||||
* 完成环境验证和依赖检查
|
||||
完成本教程后，学员应能够：
|
||||
|
||||
* 完成环境验证与依赖检查
|
||||
|
||||
* 理解并配置基准测试参数
|
||||
* 理解并配置 KV‑Cache Benchmark 的核心参数
|
||||
|
||||
* 明确Baseline的解读，理解为什么KV-Cache是性能瓶颈
|
||||
* 明确 Baseline 的含义，理解 KV‑Cache 为何成为推理性能瓶颈
|
||||
|
||||
* 运行 KV-Cache Benchmark 测试
|
||||
* 运行 `flash_attn_with_kvcache` 的 Benchmark 测试并获取性能数据
|
||||
|
||||
* 输出一份 baseline 性能结果记录表，为后续算子优化提供对比基准
|
||||
* 输出一份 Baseline 性能结果记录表，为后续算子优化提供对比基准
|
||||
|
||||
* 理解 XPU‑OJ 评测的 `run_kernel` 接口规范与精度要求
|
||||
|
||||
* 基于 OJ 题包接口实现一个最小正确版 `run_kernel`，通过正确性校验
|
||||
|
||||
|
||||
> Baseline 解读：为什么本教程基于 KV Cache 做性能基线？
|
||||
|
|
@ -125,17 +131,20 @@
|
|||
|
||||
完成本模块后，你将能够：
|
||||
|
||||
1. 理解 FlashAttention `flash_attn_with_kvcache` 核函数的基本作用与应用场景；
|
||||
1. **理解 FlashAttention Paged KV‑Cache 算子的作用**
|
||||
明白 `flash_attn_with_kvcache` 在 LLM 推理 decode 阶段如何高效利用分页 KV Cache，减少显存碎片并提升吞吐。
|
||||
|
||||
2. 利用预装专属镜像，完成沐曦 GPU 硬件环境的快速验证；
|
||||
2. **完成环境准备与 Benchmark 运行**
|
||||
安装所需依赖，运行 `benchmark_kvcache.py`，生成包含执行时间与显存带宽的 CSV 性能记录。
|
||||
|
||||
3. 深入理解 Baseline 的概念，掌握性能测试与正确性测试的联系与区别；
|
||||
3. **读懂 Baseline 并定位性能瓶颈**
|
||||
分析不同 `batch_size`、`seq_len_kv` 下的带宽曲线，理解显存带宽对 decode 阶段的影响。
|
||||
|
||||
4. 跑通 KV-Cache Benchmark 基准测试脚本；
|
||||
4. **理清 Baseline 与 OJ 题包的关系**
|
||||
明确基准脚本用于性能参照，OJ 题包定义最终提交接口、数据范围和精度校验标准。
|
||||
|
||||
5. 完成多种 `batch_size × seq_len_kv` 组合的性能测试；
|
||||
|
||||
6. 输出带宽性能结果 CSV 文件并进行结果分析。
|
||||
5. **实现并提交一个最小正确版** `**run_kernel**`
|
||||
根据题包中的接口约定编写 CUDA 算子，通过 OJ 正确性校验并记录首次提交耗时。
|
||||
|
||||
|
||||
---
|
||||
|
|
@ -163,37 +172,41 @@
|
|||
---
|
||||
|
||||
## 四、前置准备
|
||||
|
||||
开始实战前，请确认你已经完成以下准备：
|
||||
|
||||
### 环境准备
|
||||
|
||||
* **设置领取与兑换算力券**
|
||||
|
||||
* 前往沐曦开发者社区注册账号并完成邮箱验证，申请并获取 MACA 算力代金券兑换码。
|
||||
|
||||
|
||||
1. 前往沐曦开发者社区注册账号并完成邮箱验证，申请并获取 MACA 算力代金券兑换码。
|
||||
|
||||
链接：[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6)
|
||||
链接：https://developer.metax-tech.com/activities/6
|
||||
|
||||
* 登录 模力方舟平台 (Gitee AI)，在“费用中心 -> 算力券”页面输入兑换码完成充值。
|
||||
|
||||
|
||||
链接：[https://ai.gitee.com/](https://ai.gitee.com/)
|
||||
2. 登录 模力方舟平台 (Gitee AI)，在“费用中心 -> 算力券”页面输入兑换码完成充值。 链接：https://ai.gitee.com/
|
||||
|
||||
|
||||
* **创建并启动实例**
|
||||
|
||||
* 进入 算力市场，筛选“沐曦”芯片厂商，选择合适的 GPU 规格（推荐 曦云 C500 节点）。
|
||||
|
||||
* **关键配置：** 在预装镜像处，务必选择专属开发镜像（`PyTorch Agent/2.8.0/Python 3.12/maca 3.7.2.1`）。
|
||||
|
||||
* 创建完成后，进入算力容器，点击“工具-lab”即可打开 JupyterLab 终端开始项目创作。
|
||||
|
||||
|
||||
**说明：**由于本次使用的是预装的专属镜像，环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后，无需再进行繁琐的依赖库版本验证即可直接进入测试环节。
|
||||
1. 进入 算力市场，筛选“沐曦”芯片厂商，选择合适的 GPU 规格（推荐 曦云 C500 节点）。
|
||||
|
||||
2. 关键配置： 在预装镜像处，务必选择专属开发镜像（PyTorch Agent/2.8.0/Python 3.12/maca 3.7.2.1）。
|
||||
|
||||
3. 创建完成后，进入算力容器，点击“工具-lab”即可打开 JupyterLab 终端开始项目创作。
|
||||
|
||||
|
||||
\*\*说明：\*\*由于本次使用的是预装的专属镜像，环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后，无需再进行繁琐的依赖库版本验证即可直接进入测试环节。
|
||||
|
||||
### 代码准备
|
||||
|
||||
* 已获取基准测试脚本 `benchmark_kvcache.py`
|
||||
* 获取目标源码（包含 `benchmark_kvcache.py` 及 OJ 题包）
|
||||
|
||||
* 已进入项目目录 `/data/flashattn_baseline`
|
||||
|
||||
* 准备 Benchmark 脚本与 OJ 测试脚本
|
||||
|
||||
|
||||
---
|
||||
|
||||
|
|
@ -201,69 +214,63 @@
|
|||
|
||||
### 关键术语
|
||||
|
||||
* **KV-Cache**：缓存历史 Token 的 Key/Value 向量，避免 Transformer 推理时重复计算。
|
||||
* **KV-Cache：**缓存历史 Token 的 Key/Value 向量，避免 Transformer 推理时重复计算。
|
||||
|
||||
* **Paged KV-Cache**：将 KV-Cache 分页管理，减少显存碎片，提高利用率。
|
||||
* **Paged KV-Cache：**将 KV-Cache 分页管理，减少显存碎片，提高利用率。
|
||||
|
||||
* **Batch Size**：一次处理的样本数，越大并行度越高，但显存占用越大。
|
||||
* **Batch Size：**一次处理的样本数，越大并行度越高，但显存占用越大。
|
||||
|
||||
* **seq\_len\_kv**：KV-Cache 中已缓存的历史 Token 数量。
|
||||
* **seq\_len\_kv：**KV-Cache 中已缓存的历史 Token 数量。
|
||||
|
||||
* **headdim**：注意力头维度，常见为 64/128/256。
|
||||
* **headdim：**注意力头维度，常见为 64/128/256。
|
||||
|
||||
|
||||
### 核心知识
|
||||
|
||||
* **正确性测试 ：**验证算子输出结果的数学精度是否与标准实现一致，这是**绝对底线**。
|
||||
* **正确性测试 ：**验证算子输出结果的数学精度是否与标准实现一致，这是绝对底线。
|
||||
|
||||
* **性能测试 ：**在正确的前提下，测算速度与吞吐。通过建立 **Baseline（基线），**才能量化后续每次代码修改带来的真实收益（加速比）。
|
||||
* **性能测试 ：**在正确的前提下，测算速度与吞吐。通过建立Baseline（基线），才能量化后续每次代码修改带来的真实收益（加速比）。
|
||||
|
||||
* **Benchmark （基准测试）：**在固定条件下反复运行同一任务，获取可重复的性能指标，用于建立基线、量化优化效果和定位瓶颈。
|
||||
|
||||
### 关键指标
|
||||
* **XPU‑OJ**：比赛官方在线评测平台，最终评测会调用参赛者提交代码中的 `run_kernel`。
|
||||
|
||||
* **Kernel 执行时间**：GPU 核函数运行耗时（ms），使用 GPU 端同步计时获得。
|
||||
|
||||
### 关键指标
|
||||
|
||||
* **Kernel 执行时间：**GPU 核函数运行耗时（ms），使用 GPU 端同步计时获得。
|
||||
|
||||
* **有效带宽：**`数据传输量 (GB) ÷ Kernel 时间 (s)`，越接近理论峰值说明显存带宽利用越充分。
|
||||
* **有效带宽：**数据传输量 (GB) ÷ Kernel 时间 (s)，越接近理论峰值说明显存带宽利用越充分。
|
||||
|
||||
|
||||
### 其他要点
|
||||
|
||||
* **Warmup**：预热若干次（不记录），使 GPU 进入稳定状态。
|
||||
* **Warmup：**预热若干次（不记录），使 GPU 进入稳定状态。
|
||||
|
||||
* **Repeat**：正式运行多次，取平均值或中位数以消除波动。
|
||||
* **Repeat：**正式运行多次，取平均值或中位数以消除波动。
|
||||
|
||||
* **同步**：调用 `torch.cuda.synchronize()` 确保精确计时。
|
||||
* **同步：**调用 torch.cuda.synchronize() 确保精确计时。
|
||||
|
||||
* **数据类型**：本教程使用 `bfloat16`，在精度和性能取得平衡。
|
||||
* **数据类型：**本教程使用 bfloat16，在精度和性能取得平衡。
|
||||
|
||||
* **显存占用估算**：`KV-Cache ≈ batch × seq_len_kv × num_heads_k × headdim × 2(K+V) × 字节数`。
|
||||
* **显存占用估算：**KV-Cache ≈ batch × seq\_len\_kv × num\_heads\_k × headdim × 2(K+V) × 字节数。
|
||||
|
||||
* **OOM 应对**：减小 batch/seq\_len\_kv、使用更小 dtype 或释放中间变量。
|
||||
* **OOM 应对：**减小 batch/seq\_len\_kv、使用更小 dtype 或释放中间变量。
|
||||
|
||||
* **Tensor Core**：现代 GPU（含沐曦 C500）的矩阵乘法专用单元，要求维度对齐为 8 或 16 的倍数。
|
||||
* **Tensor Core：**现代 GPU（含沐曦 C500）的矩阵乘法专用单元，要求维度对齐为 8 或 16 的倍数。
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 六、项目实践：FlashAttention KV-Cache Benchmark
|
||||
## 六、项目实践：FlashAttention KV‑Cache Benchmark & OJ 评测
|
||||
|
||||
### 项目目标
|
||||
|
||||
对 FlashAttention 的 paged KV-cache 推理核函数（`flash_attn_with_kvcache`）进行自动化性能基准测试，覆盖多种 `batch_size × seq_len_kv` 组合，输出执行时间和有效显存带宽。
|
||||
|
||||
基准测试脚本 `benchmark_kvcache.py` 的最小闭环包括：
|
||||
|
||||
1. 准备 paged KV-cache 张量和 block table。
|
||||
|
||||
2. 通过 `torch.profiler` 对 kernel 进行计时。
|
||||
|
||||
3. 计算有效显存带宽（GB/s）。
|
||||
|
||||
4. 将结果写入带时间戳的 CSV 文件。
|
||||
* 对 FlashAttention 的 paged KV-cache 推理核函数（`flash_attn_with_kvcache`）进行自动化性能基准测试，覆盖多种 `batch_size × seq_len_kv` 组合，输出执行时间和有效显存带宽。
|
||||
|
||||
|
||||
---
|
||||
* 理解 XPU‑OJ 题包的 `run_kernel` 接口，实现一个最小正确版 CUDA 算子，通过所有 OJ 正确性测试用例。
|
||||
|
||||
|
||||
### 步骤 0：进入创建的实例环境
|
||||
|
||||
|
|
@ -271,7 +278,7 @@
|
|||
|
||||
选择工具-lab进入实例环境
|
||||
|
||||
|
||||
|
||||
|
||||
### 步骤 1：检查运行环境
|
||||
|
||||
|
|
@ -279,7 +286,7 @@
|
|||
|
||||
在JupyterLab Terminal中检查运行环境的配置。
|
||||
|
||||
|
||||
|
||||
|
||||
**操作：** 检查 GPU 状态、Python 版本和依赖版本。
|
||||
|
||||
|
|
@ -307,22 +314,22 @@ python -c "import einops; print('einops OK')"
|
|||
* `mx-smi` 显示沐曦 GPU 信息
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
* Python 版本 = 3.8
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
* `torch.cuda.is_available()` 返回 `True`
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
* 所有依赖版本符合要求
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
**常见问题：**
|
||||
|
||||
|
|
@ -336,44 +343,46 @@ python -c "import einops; print('einops OK')"
|
|||
|
||||
### 步骤 2：进入项目目录
|
||||
|
||||
**目标：** 进入本模块所需的项目目录。
|
||||
目标：进入本模块所需的源码目录。
|
||||
|
||||
1. 克隆代码仓库
|
||||
1. 克隆代码仓库
|
||||
|
||||
```Bash
|
||||
git clone https://gitlink.org.cn/metax-maca/op_optimization.git
|
||||
```
|
||||
|
||||
2. 准备flashattn\_baseline
|
||||
|
||||
在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashattn_baseline` 文件夹。可以将 `flashattn_baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
|
||||
**下一步操作：** 切换到 Flashattn\_Baseline 项目目录。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
cd flashattn_baseline/Flashattn_Baselinels -la
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```Plain
|
||||
total xx
|
||||
drwxr-xr-x 2 root root 4096 Jun 1 09:00 __MACOSX
|
||||
|
||||
- rw-r--r-- 1 root root 5232 Jun 1 09:00 benchmark_kvcache.py
|
||||
|
||||
|
||||
- rw-r--r-- 1 root root 1440 Jun 1 09:00 benchmark_kvcache_20260526_150953.csv
|
||||
|
||||
...
|
||||
|
||||
```
|
||||
|
||||
3. 将Flashattn\_Baseline文件加入JupyterLab。
|
||||
|
||||
|
||||
```Bash
|
||||
git clone https://gitlink.org.cn/metax-maca/op_optimization.git
|
||||
```
|
||||
|
||||
|
||||
2. 准备flashattn_baseline
|
||||
|
||||
在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashattn_baseline` 文件夹。可以将 `flashattn_baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
|
||||
|
||||
|
||||
**操作：** 切换到基准测试脚本所在目录。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```Bash
|
||||
cd data/flashattn_baseline
|
||||
ls -la
|
||||
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```Plain
|
||||
total xx
|
||||
drwxr-xr-x 2 root root 4096 Jun 1 09:00 __MACOSX
|
||||
|
||||
- rw-r--r-- 1 root root 5232 Jun 1 09:00 benchmark_kvcache.py
|
||||
|
||||
|
||||
- rw-r--r-- 1 root root 1440 Jun 1 09:00 benchmark_kvcache_20260526_150953.csv
|
||||
|
||||
...
|
||||
|
||||
```
|
||||
---
|
||||
|
||||
### 步骤 3：配置基准测试参数
|
||||
|
|
@ -550,40 +559,434 @@ repeat = 200
|
|||
ms = run_with_profiler(run_fn, warmup=warmup, reps=repeat, print_result=True, target_kernels=["flash"])
|
||||
|
||||
```
|
||||
---
|
||||
|
||||
## 七、Agent 使用说明
|
||||
### 步骤 7:从 Baseline 到 XPU-OJ 提交
|
||||
|
||||
在本模块中，Agent 可用于以下场景：
|
||||
Baseline benchmark 用于理解目标算子的调用方式、输入输出 shape 和性能基线;XPU-OJ 题包用于定义最终评测接口、数据范围、参考输出和精度要求。
|
||||
|
||||
### Prompt 模板
|
||||
跑完 baseline 后,选手需要完成以下转换:
|
||||
|
||||
**环境验证：**
|
||||
1. 从 benchmark 脚本中理解目标 API,本任务对应的是 `flash_attn.flash_attn_interface` 中的 `flash_attn_with_kvcache`(paged KV cache 布局);
|
||||
|
||||
2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口;
|
||||
|
||||
3. 对照题包中的输入 shape、数据范围和精度要求;
|
||||
|
||||
4. 编写自己的 `run_kernel(...)`;
|
||||
|
||||
5. 提交 OJ,先通过正确性;
|
||||
|
||||
6. 正确性通过后,再对比 baseline / OJ 耗时继续优化。
|
||||
|
||||
|
||||
```Plain
|
||||
请帮我验证当前环境是否满足 FlashAttention KV-Cache Benchmark 的运行要求，包括：
|
||||
1. 沐曦 GPU 是否可见
|
||||
2. PyTorch 版本和 CUDA 支持
|
||||
3. flash-attn 和 einops 是否已安装
|
||||
### 题目说明(FlashAttention KV Cache Decode)
|
||||
|
||||
> 注意:每个子题的接口参数、数据范围和精度要求可能不同,正式要求以对应 XPU-OJ 题包为准。本节以 **FlashAttention KV Cache Decode** 为例,演示从 baseline benchmark 到 XPU-OJ 提交的完整流程。
|
||||
|
||||
* **对应 baseline 脚本**:`flashattn_baseline/baseline/benchmark_kvcache.py`
|
||||
|
||||
* **对应 FlashAttention API**:`flash_attn_with_kvcache`(paged KV cache 版本)
|
||||
|
||||
* **算子说明**:实现 paged KV cache 下的 decode 注意力,每个 batch 只有 1 个 query token,KV cache 按 page 存储,长度由 `seqlen_k` 决定。
|
||||
|
||||
* **对应 OJ 题包**:XPU-OJ 上 `FlashAttention KV Cache Decode` 题(题号 20005)。
|
||||
|
||||
|
||||
### 步骤 8:理解 XPU-OJ 评测接口与精度要求
|
||||
|
||||
**目标**:明确 Baseline 与最终评测提交之间的关系,理解选手需要实现什么。
|
||||
|
||||
完成 baseline benchmark 后,需要注意:baseline 脚本主要用于建立性能基线,**不是最终提交物**。最终评测以 XPU-OJ 题包为准,评测程序会调用选手提交代码中的 `run_kernel`,并将输出结果与 baseline 参考结果进行比较。
|
||||
|
||||
下面以 **FlashAttention KV Cache Decode** 题为例,题包目录中通常包含以下文件:
|
||||
|
||||
* `zh_CN/00_题目描述.md`:说明需要实现的算子功能;
|
||||
|
||||
* `zh_CN/01_接口约定.cuda.md`:说明必须实现的 `run_kernel` 函数签名;
|
||||
|
||||
* `zh_CN/05_数据范围与提示.md`:说明测试范围和精度要求;
|
||||
|
||||
* `testcase_config.py`:定义测试数据生成、baseline 参考实现和正确性校验方式。
|
||||
|
||||
|
||||
#### 1. 必须实现的接口
|
||||
|
||||
选手需要在提交的 CUDA 源码中提供如下 C 符号,函数名、参数类型、顺序必须完全一致,并使用 `extern "C"` 防止 name mangling:
|
||||
|
||||
```cpp
|
||||
#include <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
extern "C" void run_kernel(
|
||||
const __nv_bfloat16* q,
|
||||
const __nv_bfloat16* k_cache_paged,
|
||||
const __nv_bfloat16* v_cache_paged,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* cache_seqlens,
|
||||
const int32_t* block_table,
|
||||
int64_t batch_size,
|
||||
int64_t seqlen_k,
|
||||
int64_t seqlen_q,
|
||||
int64_t num_heads,
|
||||
int64_t num_heads_k,
|
||||
int64_t headdim,
|
||||
int64_t page_block_size,
|
||||
int64_t num_blocks,
|
||||
int64_t causal
|
||||
);
|
||||
|
||||
```
|
||||
|
||||
**参数配置建议：**
|
||||
#### 2. 参数说明
|
||||
|
||||
```Plain
|
||||
我需要测试 headdim=128 和 headdim=256 的性能差异，请帮我推荐合适的 batch_sizes 和 seq_lens_kv 扫描范围。
|
||||
| 参数 | 说明 |
|
||||
| --- | --- |
|
||||
| `q` | decode query tensor,shape `(batch_size, seqlen_q, num_heads, headdim)`,连续 `bf16` |
|
||||
| `k_cache_paged` | paged key cache,shape `(num_blocks, page_block_size, num_heads_k, headdim)`,连续 `bf16` |
|
||||
| `v_cache_paged` | paged value cache,shape `(num_blocks, page_block_size, num_heads_k, headdim)`,连续 `bf16` |
|
||||
| `output` | 输出缓冲区,shape `(batch_size, seqlen_q, num_heads, headdim)`,连续 `bf16` |
|
||||
| `cache_seqlens` | 每个 batch 的 KV 长度,shape `(batch_size)`,连续 `int32` |
|
||||
| `block_table` | 每个 batch 的 page 映射表,shape `(batch_size, num_blocks / batch_size)`,连续 `int32` |
|
||||
| `seqlen_q` | query 长度,评测中固定为 `1` |
|
||||
| `page_block_size` | page size,评测中固定为 `16` |
|
||||
| `causal` | 是否启用 causal mask,评测中固定为 `0` |
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确,**不建议在** `**run_kernel**` **内部做** `**cudaDeviceSynchronize()**` **或显式同步**。
|
||||
|
||||
#### 3. KV cache 布局
|
||||
|
||||
KV cache layout 固定为 `flash_attn_with_kvcache` 的 paged cache 布局:`(num_blocks, page_block_size, num_heads_k, headdim)`。
|
||||
|
||||
第 `t` 个 KV token 位于 `block_table[batch_idx, t / page_block_size]` 指向的物理 page 中,page 内偏移为 `t % page_block_size`。
|
||||
|
||||
例如 `batch_size = 1`、`seqlen_k = 512`、`page_block_size = 16` 时,每个序列需要访问 `32` 个有效 page。
|
||||
|
||||
#### 4. 评测数据范围
|
||||
|
||||
`testcase_config.py` 中定义了本题的测试配置(摘自题包):
|
||||
|
||||
```python
|
||||
HEAD_DIMS = [128]
|
||||
BATCH_SIZES = [1, 4, 16]
|
||||
SEQ_LENS_KV = [1024, 4096, 8192, 16384]
|
||||
SEQ_LEN_Q = 1
|
||||
NUM_HEADS = 8
|
||||
NUM_HEADS_K = 8
|
||||
PAGE_BLOCK_SIZE = 16
|
||||
CAUSAL = 0
|
||||
|
||||
```
|
||||
|
||||
**结果分析：**
|
||||
`num_blocks = max(1024, ceil(seqlen_k / page_block_size) * batch_size * 3)`。
|
||||
|
||||
```Plain
|
||||
请帮我分析这份 benchmark 结果 CSV 文件，找出峰值带宽配置和 OOM 边界。
|
||||
#### 5. 精度要求
|
||||
|
||||
当前 `FlashAttention KV Cache Decode` 题的校验方式为:
|
||||
|
||||
```python
|
||||
torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
||||
|
||||
```
|
||||
---
|
||||
|
||||
## 八、常见问题
|
||||
也就是说,选手实现的输出需要在上述容差范围内与 baseline 输出一致。不同算子的容差可能不同,**正式精度要求以对应 OJ 题包说明为准**。
|
||||
|
||||
### 步骤 9：登录 XPU-OJ 并进入题目页面
|
||||
|
||||
使用组委会统一发放的账号登录 XPU-OJ，并进入对应赛题页面。
|
||||
|
||||
1.打开 XPU-OJ 平台：https://xpuoj.com/
|
||||
|
||||
2.使用组委会统一发放的账号和初始密码登录【后续发布】；
|
||||
|
||||
3.登录后进入比赛 / 题目列表页面；
|
||||
|
||||
4.找到对应题目，例如 `20005 FlashAttention KV Cache Decode`；
|
||||
|
||||
5.点击进入题目详情页，查看题目描述、接口约定、数据范围和提交入口。
|
||||
|
||||
### 步骤 10：在Agent的帮助下提交 OJ 冒烟代码
|
||||
|
||||
**目标：**完成一次最小提交，确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。
|
||||
|
||||
**操作：**
|
||||
|
||||
1. 在语言下拉框中选择本题支持的提交语言，例如 `MXMACA C++`、`TileLang` 、 `Triton`；
|
||||
|
||||
2. 将实现了题目要求接口的代码复制到提交框中；
|
||||
|
||||
如果你还没有 run\_kernel，应该从哪里开始？
|
||||
|
||||
OJ 最终评测不会直接运行 baseline 脚本，而是调用你提交代码中的 `run_kernel(...)`。 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路。
|
||||
|
||||
3. 借助 Agent 从题包生成 run\_kernel 初版
|
||||
|
||||
|
||||
### 10.1 在镜像终端中安装并启动 OpenCode
|
||||
|
||||
1. 首先返回镜像JupyterLab Terminal中, 打开容器内的 **Terminal**(终端)
|
||||
|
||||
2. 执行下面的命令安装 OpenCode:
|
||||
|
||||
```bash
|
||||
curl -fsSL https://opencode.ai/install | bash
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
3. 安装完成后,`cd` 进入赛题文件夹(根据实际目录调整):
|
||||
|
||||
```bash
|
||||
cd xpuoj_problem/
|
||||
|
||||
```
|
||||
|
||||
4. 在该目录下直接输入 `opencode` 并回车,即可进入 OpenCode 的 Agent 界面:
|
||||
|
||||
```bash
|
||||
opencode
|
||||
|
||||
```
|
||||
|
||||
|
||||
|
||||
参考prompt：
|
||||
|
||||
```python
|
||||
本题是 FlashAttention paged KV cache decode 的 CUDA C++ 前向算子。
|
||||
输入输出规范如下：
|
||||
- q: [batch, 1, num_heads, head_dim] (float)
|
||||
- k_cache_paged: [num_blocks, num_heads, page_size, head_dim] (float)
|
||||
- v_cache_paged: [num_blocks, num_heads, page_size, head_dim] (float)
|
||||
- cache_seqlen: [batch] (int)
|
||||
- block_table: [batch, max_num_blocks] (int)
|
||||
- output: [batch, 1, num_heads, head_dim]
|
||||
需调用 run_kernel(q, k_cache_paged, v_cache_paged, cache_seqlen, block_table, output)
|
||||
你需要根据 cache_seqlen 和 block_table 从 paged cache 中取出对应的 K, V，计算 attention 结果，存入 output。
|
||||
参考：out = flash_attn_with_kvcache(q, k_cache_paged, v_cache_paged, cache_seqlen=..., block_table=...)
|
||||
要求你的实现与这个 API 的计算结果一致（误差允许1e-5）。
|
||||
请生成一个 run_kernel 函数，优先保证正确性。
|
||||
```
|
||||
|
||||
|
||||
|
||||
输出：
|
||||
|
||||
|
||||
|
||||
为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常:
|
||||
|
||||
```python
|
||||
#include <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define PAGE_SIZE 16
|
||||
#define HEAD_DIM 128
|
||||
|
||||
__global__ void paged_attention_kernel(
|
||||
const __nv_bfloat16* __restrict__ q,
|
||||
const __nv_bfloat16* __restrict__ k_cache_paged,
|
||||
const __nv_bfloat16* __restrict__ v_cache_paged,
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
const int32_t* __restrict__ cache_seqlens,
|
||||
const int32_t* __restrict__ block_table,
|
||||
int64_t batch_size,
|
||||
int64_t seqlen_q,
|
||||
int64_t num_heads,
|
||||
int64_t num_heads_k,
|
||||
int64_t headdim,
|
||||
int64_t page_block_size,
|
||||
int64_t blocks_per_batch)
|
||||
{
|
||||
int batch_idx = blockIdx.x / num_heads;
|
||||
int head_idx = blockIdx.x % num_heads;
|
||||
if (batch_idx >= batch_size || head_idx >= num_heads) return;
|
||||
|
||||
int seqlen = cache_seqlens[batch_idx];
|
||||
if (seqlen <= 0) {
|
||||
// 无有效 KV，输出 0
|
||||
int64_t out_base = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim;
|
||||
for (int i = threadIdx.x; i < headdim; i += blockDim.x)
|
||||
output[out_base + i] = __float2bfloat16(0.0f);
|
||||
return;
|
||||
}
|
||||
|
||||
int tid = threadIdx.x;
|
||||
|
||||
// 加载对应 head 的 query 元素（每个线程负责一个维度）
|
||||
int64_t q_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim;
|
||||
float q_val = __bfloat162float(q[q_offset + tid]);
|
||||
|
||||
// 全局 softmax 状态（每个线程维护自己维度的累加器）
|
||||
float max_val = -1e38f;
|
||||
float sum_exp = 0.0f;
|
||||
float out_acc = 0.0f;
|
||||
float scale = rsqrtf(static_cast<float>(headdim));
|
||||
|
||||
// 共享内存布局：
|
||||
// K_tile[PAGE_SIZE][HEAD_DIM] (bf16)
|
||||
// V_tile[PAGE_SIZE][HEAD_DIM] (bf16)
|
||||
// partial_scores[PAGE_SIZE][HEAD_DIM] (float, 用于归约点积)
|
||||
__shared__ __nv_bfloat16 K_tile[PAGE_SIZE][HEAD_DIM];
|
||||
__shared__ __nv_bfloat16 V_tile[PAGE_SIZE][HEAD_DIM];
|
||||
__shared__ float partial_scores[PAGE_SIZE][HEAD_DIM];
|
||||
|
||||
int total_pages = (seqlen + PAGE_SIZE - 1) / PAGE_SIZE;
|
||||
|
||||
for (int page = 0; page < total_pages; ++page) {
|
||||
int physical_block = block_table[batch_idx * blocks_per_batch + page];
|
||||
int tokens_this_page = min(seqlen - page * PAGE_SIZE, PAGE_SIZE);
|
||||
|
||||
// 1. 将当前 page 的 K 和 V 从全局显存加载到共享内存
|
||||
// 每个线程负责加载所有 token 的同一个 head 维度
|
||||
const int64_t kv_stride = num_heads_k * headdim; // 每个 (block, offset) 的 stride
|
||||
for (int j = 0; j < tokens_this_page; ++j) {
|
||||
int64_t offset = (physical_block * PAGE_SIZE + j) * kv_stride + head_idx * headdim + tid;
|
||||
K_tile[j][tid] = k_cache_paged[offset];
|
||||
V_tile[j][tid] = v_cache_paged[offset];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 2. 计算该 page 内每个 token 与 Q 的部分点积，存入 partial_scores
|
||||
for (int j = 0; j < tokens_this_page; ++j) {
|
||||
float k_val = __bfloat162float(K_tile[j][tid]);
|
||||
partial_scores[j][tid] = q_val * k_val;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 3. 对 partial_scores 做 tree reduction，得到每个 token 的完整点积
|
||||
#pragma unroll
|
||||
for (int stride = HEAD_DIM / 2; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < tokens_this_page; ++j) {
|
||||
partial_scores[j][tid] += partial_scores[j][tid + stride];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// 4. 在线 safe softmax 更新 + V 累加
|
||||
// 4.1 找出本 page 内点积的最大值，结合全局 max 得到 new_max
|
||||
float local_max = -1e38f;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < tokens_this_page; ++j) {
|
||||
local_max = fmaxf(local_max, partial_scores[j][0]);
|
||||
}
|
||||
float new_max = fmaxf(max_val, local_max);
|
||||
|
||||
// 4.2 用旧的 max 对全局状态进行重缩放
|
||||
float rescale = expf(max_val - new_max);
|
||||
sum_exp *= rescale;
|
||||
out_acc *= rescale;
|
||||
max_val = new_max;
|
||||
|
||||
// 4.3 累加 V，并更新 sum_exp
|
||||
#pragma unroll
|
||||
for (int j = 0; j < tokens_this_page; ++j) {
|
||||
float score = partial_scores[j][0] * scale;
|
||||
float weight = expf(score - new_max);
|
||||
sum_exp += weight;
|
||||
|
||||
float v_val = __bfloat162float(V_tile[j][tid]);
|
||||
out_acc += weight * v_val;
|
||||
}
|
||||
|
||||
__syncthreads(); // 准备下一个 page 的共享内存加载
|
||||
}
|
||||
|
||||
// 5. 最终归一化并写回
|
||||
if (sum_exp > 0.0f) {
|
||||
out_acc /= sum_exp;
|
||||
} else {
|
||||
out_acc = 0.0f;
|
||||
}
|
||||
|
||||
int64_t out_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim + tid;
|
||||
output[out_offset] = __float2bfloat16(out_acc);
|
||||
}
|
||||
|
||||
extern "C" void run_kernel(
|
||||
const __nv_bfloat16* q,
|
||||
const __nv_bfloat16* k_cache_paged,
|
||||
const __nv_bfloat16* v_cache_paged,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* cache_seqlens,
|
||||
const int32_t* block_table,
|
||||
int64_t batch_size,
|
||||
int64_t seqlen_k,
|
||||
int64_t seqlen_q,
|
||||
int64_t num_heads,
|
||||
int64_t num_heads_k,
|
||||
int64_t headdim,
|
||||
int64_t page_block_size,
|
||||
int64_t num_blocks,
|
||||
int64_t causal)
|
||||
{
|
||||
int64_t blocks_per_batch = num_blocks / batch_size;
|
||||
dim3 grid(batch_size * num_heads);
|
||||
dim3 block(HEAD_DIM);
|
||||
|
||||
paged_attention_kernel<<<grid, block>>>(
|
||||
q, k_cache_paged, v_cache_paged, output,
|
||||
cache_seqlens, block_table,
|
||||
batch_size, seqlen_q, num_heads, num_heads_k, headdim,
|
||||
page_block_size, blocks_per_batch
|
||||
);
|
||||
}
|
||||
```
|
||||
|
||||
以上代码仅用于说明接口结构，不代表最优实现，也不作为评分参考
|
||||
|
||||
5. 点击提交，等待评测结果返回；
|
||||
|
||||
评测时间与题目测试点数量、队列状态和平台负载有关，通常需要等待数十秒到数分钟。以平台实际返回为准。
|
||||
|
||||
|
||||
|
||||
OJ 对每次提交大致会走这个流程：
|
||||
|
||||
```plaintext
|
||||
XPU-OJ 的一次评测大致流程如下：
|
||||
1. 选手提交代码；
|
||||
2. 平台按所选语言编译或加载提交代码；
|
||||
3. 评测程序构造测试输入；
|
||||
4. 调用选手代码中的 `run_kernel(...)`；
|
||||
5. 调用 `testcase_config.py` 中的 `baseline()` / 参考实现生成 `output_ref`；
|
||||
6. 将 `run_kernel(...)` 的输出与 `output_ref` 做正确性校验；
|
||||
7. 正确性通过后，统计运行耗时或性能指标；
|
||||
8. 根据题目评分规则换算该题得分；
|
||||
9. 更新该题历史最好成绩；
|
||||
10. 汇总各题最好成绩，得到排行榜总分。
|
||||
```
|
||||
|
||||
6. 查看结果
|
||||
|
||||
提交详情会显示状态、得分、时间、内存、编译信息以及各测试点结果。
|
||||
|
||||
|
||||
|
||||
## 从 Baseline 到参赛作品的路径回顾：
|
||||
|
||||
1. 跑通 baseline benchmark，记录原始性能；
|
||||
|
||||
2. 提交冒烟代码，确认 OJ 链路正常；
|
||||
|
||||
3. 使用 Agent 阅读题包，理解输入输出、数据范围和精度要求；
|
||||
|
||||
4. 在 `run_kernel(...)` 中实现最小正确版算子；
|
||||
|
||||
5. 提交 OJ，先通过正确性；
|
||||
|
||||
6. 正确性通过后，再让 Agent 辅助分析性能瓶颈；
|
||||
|
||||
7. 围绕访存、softmax、线程划分、head\_dim 特化、K/V 复用等方向迭代优化；
|
||||
|
||||
8. 保存每轮 Agent Prompt、代码改动、OJ 结果和性能变化，形成可复现的 Agent/Skill 优化流程。
|
||||
|
||||
|
||||
## 七、常见问题
|
||||
|
||||
### Q1: 运行时提示 `ModuleNotFoundError: No module named 'flash_attn '`
|
||||
|
||||
|
|
@ -601,9 +1004,23 @@ ms = run_with_profiler(run_fn, warmup=warmup, reps=repeat, print_result=True, ta
|
|||
|
||||
**原因：** 可能是 warmup 不足或 GPU 未达到稳态。 **解决：** 增加 `warmup` 次数，如 `warmup = 20`。
|
||||
|
||||
---
|
||||
### Q5: 评测状态显示 `**Compile Error**`或提示 `**Undefined reference to run_kernel**`
|
||||
|
||||
## 九、下一步学习建议
|
||||
**原因：**C++函数名被修饰（Name Mangling）或参数类型/顺序与接口约定不符。 **解决：**在 run\_kernel 前添加 extern "C"，并严格逐字核对参数的类型和修饰符。
|
||||
|
||||
#### Q6: 评测状态显示 `**Wrong Answer**`，提示 torch.allclose校验失败
|
||||
|
||||
**原因：**bfloat16精度截断溢出、线程同步缺失或无效Token（Padding区域）处理错误。 **解决：**累加和 Softmax 强制转为 float32 计算；检查 \_\_syncthreads() 逻辑；增加 seqlen 的边界判空。
|
||||
|
||||
#### Q7: 评测状态显示 `**Runtime Error**` (非法内存访问/段错误)
|
||||
|
||||
**原因：**Paged KV 地址映射索引错误、尾部 Page 越界读取，或线程块维度超限。 **解决：**仔细核对物理块寻址公式；增加当前 Page 有效 Token 数量的越界判断；检查单 Block 线程数配置。
|
||||
|
||||
#### Q8: 评测状态显示`**Time Limit Exceeded**` (评测超时)
|
||||
|
||||
**原因：**发散分支内的 \_\_syncthreads() 导致内核死锁、误加主机端同步指令或并行度划分错误导致串行。 **解决：**确保同步指令在所有线程必经路径上；移除主机端多余的 cudaDeviceSynchronize()；优化 <<<grid, block>>> 参数以提升并行度。
|
||||
|
||||
## 八、下一步学习建议
|
||||
|
||||
完成本模块后，建议继续学习以下内容：
|
||||
|
||||
|
|
@ -615,5 +1032,6 @@ ms = run_with_profiler(run_fn, warmup=warmup, reps=repeat, print_result=True, ta
|
|||
|
||||
4. **性能对比分析** — 将 baseline 结果与优化后结果进行对比
|
||||
|
||||
|
||||
---
|
||||
5. **OJ 题包深度解析** — 学习阅读题包中的 `testcase_config.py`，掌握本地构造边界用例与独立 Debug 的能力
|
||||
|
||||
6. **评测打榜与极限优化** — 在通过 OJ 正确性校验的基础上，挑战排行榜（Leaderboard），不断逼近硬件理论带宽极限
|
||||
|
|
@ -1,28 +1,30 @@
|
|||
# FlashInfer 迁移 Baseline 实战
|
||||
# FlashInfer 关键算子迁移与优化
|
||||
|
||||
## 一、教程定位
|
||||
|
||||
本教程是参赛训练课程的 **FlashInfer Baseline 入门**模块，主要帮助用户快速跑通 FlashInfer 的最小可运行流程。完成本教程后，用户应能够完成源码编译、API 调用、正确性测试和 Benchmark 测试，并记录一份 baseline 性能结果，为后续算子优化提供对比基准。
|
||||
本教程是参赛训练课程的 **FlashInfer Baseline 入门** 模块，主要帮助用户快速跑通 FlashInfer 的最小可运行流程。完成本教程后，用户应能够完成源码编译、API 调用、正确性测试和 Benchmark 测试，并记录一份 baseline 性能结果，为后续算子优化提供对比基准。
|
||||
|
||||
## 二、学习目标
|
||||
## 二、完成本模块你将能够
|
||||
|
||||
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`。
|
||||
|
||||
|
||||
## 三、适用对象
|
||||
|
||||
适合人群
|
||||
**适合人群**
|
||||
|
||||
* 参赛选手：需要完成 Baseline 入门模块，为后续算子优化做准备
|
||||
* 软件开发者和Vibe Coding开发：希望从事AI相关行业开发，以及用智能体方式来做开发工作
|
||||
* LLM 推理开发者：希望了解 FlashInfer Attention Kernel 的性能表现
|
||||
* 算子优化工程师：希望基于MXMACA软件栈在沐曦国产 GPU 上做算子迁移和优化
|
||||
|
||||
前置基础
|
||||
**前置基础**
|
||||
|
||||
* Python 基础：能够运行和修改 Python 脚本
|
||||
* PyTorch 基础：了解MXMACA软化栈的使用
|
||||
|
|
@ -34,65 +36,45 @@
|
|||
|
||||
开始实战前，请确认你已经完成以下准备：
|
||||
|
||||
### GPU准备
|
||||
### 获得 GPU
|
||||
|
||||
* 步骤1：获取算力券[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6)
|
||||
1. [点击获取算力券](https://developer.metax-tech.com/activities/6)，首次登录需要使用邮箱或者手机号进行注册
|
||||
|
||||
2. 登录成功后验证邮箱
|
||||
|
||||
* 首次登录需要先进行注册（使用邮箱或者手机号进行注册）
|
||||
|
||||
* 登录成功后进行第二步-邮箱验证，填入自己的邮箱。
|
||||
|
||||
* 第三步，提交申请。
|
||||
3. 提交申请获得兑换码
|
||||
|
||||
4. 兑换算力和登陆平台：
|
||||
|
||||
* 获得兑换码
|
||||
|
||||
- [访问模力方舟官网](https://ai.gitee.com/)，进入费用中心 - 算力券 ，点击右上角 “兑换”。
|
||||
|
||||
* 步骤2：兑换算力和登陆平台
|
||||
- 进入算力容器，选择沐曦，租用算力，建议优先选 16G 显存 / 32G 显存，如下图
|
||||
|
||||
* 访问模力方舟官网：[https://ai.gitee.com/](https://ai.gitee.com/)
|
||||
|
||||
* 进入费用中心 - 算力券 ，点击右上角 “兑换”。
|
||||
|
||||
|
||||
* 步骤3：租用算力
|
||||
|
||||
|
||||
- 进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图：
|
||||
* 步骤4：创建实例
|
||||
- 创建实例。基础镜像：`maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64`
|
||||
|
||||
- 基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64
|
||||
* 步骤5：选择工具-lab进入实例环境
|
||||
- 选择工具-lab进入实例环境
|
||||
|
||||
- 步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用 `mx-smi` 命令查看
|
||||
|
||||
|
||||
* 步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用`mx-smi`命令查看
|
||||
### Python 环境
|
||||
|
||||
|
||||
### 环境依赖准备
|
||||
|
||||
```bash
|
||||
``` bash
|
||||
pip install flashinfer torch pandas numpy
|
||||
```
|
||||
|
||||
### opencode Agent安装
|
||||
### OpenCode 安装
|
||||
|
||||
```shell
|
||||
``` bash
|
||||
curl -fsSL https://opencode.ai/install | bash
|
||||
```
|
||||
|
||||
说明：
|
||||
|
||||
自动下载最新版本 自动配置环境 适用于 macOS / Linux / WSL 安装完成后，还会提示你怎么用，OpenCode 包含免费模式，使用方式：
|
||||
|
||||
```shell
|
||||
cd # 进入项目目录
|
||||
opencode
|
||||
```
|
||||
|
||||
### 代码准备
|
||||
|
||||
* 已进入指定项目目录；
|
||||
|
||||
* 准备测试脚本和 Benchmark 脚本。
|
||||
|
||||
测试脚本和 Benchmark 脚本:
|
||||
|
||||
| 命令 | 说明 |
|
||||
| --- | --- |
|
||||
|
|
@ -103,65 +85,105 @@ opencode
|
|||
|
||||
## 五、知识预备
|
||||
|
||||
#### LLM推理阶段重要概念：
|
||||
### LLM 推理阶段重要概念：
|
||||
|
||||
* **Prefill阶段**：prefill 阶段是指处理输入 prompt 的阶段
|
||||
- **Prefill 阶段**：Prefill 阶段是指处理输入 prompt 的阶段
|
||||
|
||||
* 输入：用户一次性给出的完整 prompt，长度为 seq\_len
|
||||
- 输入：用户一次性给出的完整 prompt，长度为 seq\_len
|
||||
|
||||
* 计算：对 prompt 中的每个 token 并行计算注意力，生成第一个输出 token 及 KV cache
|
||||
- 计算：对 prompt 中的每个 token 并行计算注意力，生成第一个输出 token 及 KV cache
|
||||
|
||||
* 特点：这是**计算密集型（compute-bound）**阶段，因为需要做完整的 seq\_len x seq\_len 注意力矩阵乘法
|
||||
- 特点：这是 **计算密集型（compute-bound）** 阶段，因为需要做完整的 `seq\_len * seq\_len` 注意力矩阵乘法
|
||||
|
||||
* **decode阶段**：
|
||||
- **Decode 阶段**：
|
||||
|
||||
* 每次只生成 1 个 token，利用 prefill 阶段填充好的 KV cache 做自回归生成
|
||||
|
||||
* **显存带宽密集型（memory-bound**），瓶颈在从显存读取 KV cache 而非计算
|
||||
- 每次只生成 1 个 token，利用 prefill 阶段填充好的 KV cache 做自回归生成
|
||||
|
||||
- **显存带宽密集型（memory-bound）**，瓶颈在从显存读取 KV cache 而非计算
|
||||
|
||||
prefill = 并行处理用户输入，decode = 逐个生成回答 token
|
||||
**Prefill = 并行处理用户输入，Decode = 逐个生成回答 token**
|
||||
|
||||
## 六、项目实践--FlashInfer-Baseline
|
||||
## 六、项目实践 -- FlashInfer Baseline
|
||||
|
||||
### Step 1：检查运行环境
|
||||
**目标：** 在赛事镜像中完成 FlashInfer Ragged Prefill 算子的 Baseline Benchmark，理解从 Benchmark 到 XPU-OJ 评测提交的完整流程，为后续算子优化建立性能基线。
|
||||
|
||||
**目标：** 进入terminal确认当前环境满足本模块运行要求。
|
||||
### 在赛事镜像中运行 FlashInfer Baseline Benchmark
|
||||
|
||||
**操作：** 检查 GPU、Python、编译工具和依赖版本。
|
||||
#### 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')"
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
- `mx-smi` 显示沐曦 GPU 信息
|
||||
|
||||
|
||||
|
||||
- Python 环境正常
|
||||
|
||||
|
||||
|
||||
- `torch.cuda.is_available()` 返回 `True`
|
||||
|
||||
|
||||
|
||||
- 所有依赖版本符合要求
|
||||
|
||||
|
||||
|
||||
|
||||
**常见问题：**
|
||||
|
||||
| 问题 | 解决方法 |
|
||||
| --- | --- |
|
||||
| `mx-smi: command not found` | 使用 `python -c "import torch; print(torch.cuda.get_device_name(0))"` |
|
||||
| `No GPUs are available` | 检查 MXMACA 驱动是否正确安装 |
|
||||
| `ModuleNotFoundError: No module named 'flashinfer'` | 执行 `pip install flashinfer` 安装 |
|
||||
| `mx-smi: command not found` | 确认已配置沐曦 GPU 驱动环境 |
|
||||
| `No GPUs are available` | 检查 MXMACA 驱动是否正确安装/环境变量是否正确配置 |
|
||||
| `ModuleNotFoundError: No module named 'xxx'` | `pip install xxx` |
|
||||
|
||||
### Step 2：进入项目目录
|
||||
#### 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)。
|
||||
|
||||
1. 克隆代码仓库
|
||||
|
||||
```Bash
|
||||
```bash
|
||||
git clone https://gitlink.org.cn/metax-maca/op_optimization.git
|
||||
```
|
||||
|
||||
2. 准备flashinfer_baseline
|
||||
|
||||
在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashinfer_baseline` 文件夹。可以将 `flashinfer_baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
从仓库根目录开始，在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashinfer_baseline` 文件夹。可以将 `flashinfer_baseline` 整个目录复制到工作目录 `data/` 下。
|
||||
```bash
|
||||
mkdir data
|
||||
cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline" data/
|
||||
```
|
||||
|
||||
**下一步操作：** 切换到 FlashInfer Baseline 项目目录。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
cd /data/flashinfer_baseline
|
||||
ls -la
|
||||
```
|
||||
3. 切换到 FlashInfer_Baseline 项目目录
|
||||
```bash
|
||||
cd data/flashinfer_baseline/FlashInfer_Baseline
|
||||
ls -al
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
|
|
@ -173,10 +195,9 @@ bench_batch_prefill_ragged.py
|
|||
bench_batch_mla.py
|
||||
README.md
|
||||
...
|
||||
|
||||
```
|
||||
|
||||
### Step 3：验证项目脚本
|
||||
#### Step 3：验证项目脚本
|
||||
|
||||
**目标：** 确认所有基准测试脚本可正常执行。
|
||||
|
||||
|
|
@ -190,7 +211,6 @@ python -c "import os; scripts = ['bench_common.py', 'bench_batch_decode.py', 'be
|
|||
|
||||
# 测试脚本导入
|
||||
python -c "from bench_common import setup_workspace, get_csv_path; print('脚本导入正常')"
|
||||
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
|
@ -202,32 +222,28 @@ python -c "from bench_common import setup_workspace, get_csv_path; print('脚本
|
|||
✓ bench_batch_prefill_ragged.py
|
||||
✓ bench_batch_mla.py
|
||||
脚本导入正常
|
||||
|
||||
```
|
||||
|
||||
### Step 4：运行单算子 Benchmark
|
||||
#### Step 4：运行单算子 Benchmark 并查看测试结果
|
||||
|
||||
**目标：** 执行基准测试，获取 Baseline 性能数据。
|
||||
**目标：** 执行基准测试，获取 Baseline 性能数据，查看并分析 Benchmark 输出结果。
|
||||
|
||||
**操作：** 运行 BatchDecode 基准测试脚本。
|
||||
**操作：** 运行 Ragged Prefill 基准测试脚本，读取生成的 CSV 结果文件。
|
||||
|
||||
**命令示例：**
|
||||
**运行Benchmark命令示例：**
|
||||
|
||||
```bash
|
||||
python bench_batch_decode.py
|
||||
|
||||
python bench_batch_prefill_ragged.py
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```plaintext
|
||||
[BatchDecodeWithPagedKVCacheWrapper] Starting benchmark, total cases: 144
|
||||
[1/144] bs=1, kv_len=512, hd=64: 0.020ms, 51.53 GB/s, 0.82 TFLOPs
|
||||
[2/144] bs=1, kv_len=512, hd=128: 0.023ms, 45.28 GB/s, 0.72 TFLOPs
|
||||
预期结果：（并非真实数据）
|
||||
``` 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 BatchDecodeWithPagedKVCacheWrapper_20260525_xxxxxx.csv
|
||||
|
||||
Results saved to BatchPrefillWithRaggedKVCacheWrapper_20260626_xxxxxx.csv
|
||||
```
|
||||
|
||||
**常见问题：**
|
||||
|
|
@ -237,46 +253,430 @@ Results saved to BatchDecodeWithPagedKVCacheWrapper_20260525_xxxxxx.csv
|
|||
| `out of memory` | 减小 batch\_size 或 seq\_len 参数 |
|
||||
| 运行时间过长 | 脚本会自动调整重复次数，耐心等待 |
|
||||
|
||||
### Step 5：查看测试结果
|
||||
***
|
||||
|
||||
**目标：** 查看并分析 Benchmark 输出结果。
|
||||
|
||||
**操作：** 读取生成的 CSV 结果文件。
|
||||
|
||||
**命令示例：**
|
||||
**查看结果命令示例：**
|
||||
|
||||
```bash
|
||||
# 查看最新生成的 CSV 文件
|
||||
ls -lt *.csv | head -1
|
||||
# 列出所有 CSV 结果文件（按修改时间排序，最新的在最上面）
|
||||
ls -lt *.csv 2>/dev/null || echo "未找到 CSV 文件，请先运行 benchmark"
|
||||
|
||||
# 使用 Python 查看结果
|
||||
python -c "
|
||||
import pandas as pd
|
||||
import glob
|
||||
csv_files = glob.glob('BatchDecodeWithPagedKVCacheWrapper_*.csv')
|
||||
if csv_files:
|
||||
df = pd.read_csv(max(csv_files))
|
||||
print(df[['batch_size', 'seq_len_kv', 'head_dim', 'time_ms', 'bandwidth_GB_s', 'tflops']].head(10))
|
||||
# 使用 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
|
||||
batch_size seq_len_kv head_dim time_ms bandwidth_GB_s tflops
|
||||
0 1 512 64 0.020429 51.528822 0.821253
|
||||
1 1 512 128 0.023340 45.278052 0.718833
|
||||
2 1 512 256 0.031949 66.153846 1.050256
|
||||
...
|
||||
|
||||
读取文件: 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
|
||||
...
|
||||
```
|
||||
|
||||
### XPU-OJ 在线评测教程
|
||||
|
||||
#### Step 5: 从 Baseline 到 XPU-OJ 提交
|
||||
|
||||
## 七、常见问题
|
||||
##### 5.1：Baseline 与 XPU-OJ 的关系
|
||||
|
||||
### 环境相关问题
|
||||
赛事镜像中的 Baseline Benchmark 和 XPU-OJ 在线评测任务不同。Baseline Benchmark 主要用于理解算子调用方式和建立性能基线；XPU-OJ 在线评测用于统一检查选手提交代码的正确性和性能。
|
||||
|
||||
| 维度 | Baseline 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 性能记录 | 排行榜得分 |
|
||||
|
||||
跑完 baseline 后，选手需要完成以下转换：
|
||||
1. 从 benchmark 脚本中理解目标 API，例如 BatchPrefillWithRaggedKVCacheWrapper；
|
||||
2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口；
|
||||
3. 对照题包中的输入 shape、数据范围和精度要求；
|
||||
4. 编写自己的 `run_kernel(...)`；
|
||||
5. 提交 OJ，先通过正确性；
|
||||
6. 正确性通过后，再对比 baseline / OJ 耗时继续优化。
|
||||
|
||||
***
|
||||
##### 5.2：选择目标算子
|
||||
|
||||
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** 为例演示从 baseline benchmark 到 XPU-OJ 提交的完整流程。
|
||||
|
||||
#### Step 6：理解 XPU-OJ 评测接口与精度要求
|
||||
**目标：** 明确 Baseline 与最终评测提交之间的关系，理解选手需要实现的内容。
|
||||
|
||||
> 完成 baseline benchmark 后，需要注意 baseline 脚本主要用于建立性能基线，并不需要最终提交。
|
||||
> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 baseline 参考结果进行比较。
|
||||
|
||||
下面以 FlashInfer Ragged Prefill 题为例，其中：
|
||||
- `zh_CN/00_题目描述.md`：说明需要实现的算子功能；
|
||||
- `zh_CN/01_接口约定.md`：说明必须实现的 `run_kernel` 函数签名；
|
||||
- `zh_CN/02_数据范围.md`：说明测试范围和精度要求；
|
||||
- `testcase_config.py`：定义测试数据生成、baseline 参考实现和正确性校验方式。
|
||||
|
||||
FlashInfer Ragged Prefill 的校验方式为：
|
||||
```python
|
||||
torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
|
||||
```
|
||||
|
||||
选手实现的输出需要在上述容差范围内与 baseline 输出一致。不同算子的容差可能不同，正式精度要求以对应 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 最终评测不会直接运行 baseline 脚本，而是调用你提交代码中的 `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. 查看结果
|
||||
|
||||
**50 分 / 10 分 / 与 baseline 加速比对比 / 榜单**：以 OJ 平台实际结果为准
|
||||
|
||||
## 七、Agent 使用样例
|
||||
|
||||
**目标：** 在本模块中，Agent 可以帮助你完成环境检查、运行 Benchmark、分析结果、理解 OJ 接口、生成和调试 `run_kernel` 等任务。以下是参考 prompt。
|
||||
|
||||
### 环境检查
|
||||
|
||||
```plaintext
|
||||
请帮我检查当前环境是否满足 FlashInfer 运行要求，包括：
|
||||
1. 沐曦 GPU 是否可见（mx-smi）
|
||||
2. Python 版本和 PyTorch CUDA 支持
|
||||
3. flashinfer、pandas、numpy 依赖是否已安装
|
||||
```
|
||||
|
||||
### 运行 Benchmark
|
||||
|
||||
``` plaintext
|
||||
请帮我运行 bench_batch_prefill_ragged.py 脚本，执行 Ragged Prefill 的基准测试。
|
||||
```
|
||||
|
||||
### 分析结果
|
||||
|
||||
``` plaintext
|
||||
请帮我读取最新的 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置，并与沐曦 GPU 理论峰值带宽做对比。
|
||||
```
|
||||
|
||||
### 理解 OJ 题包接口
|
||||
|
||||
``` plaintext
|
||||
请帮我阅读 FlashInfer Ragged Prefill 题包（problem_20001）中的以下文件：
|
||||
- zh_CN/00_题目描述.md
|
||||
- zh_CN/01_接口约定.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 时需要注意的边界条件
|
||||
```
|
||||
|
||||
### 生成 `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)
|
||||
```
|
||||
|
||||
### 调试 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 更新逻辑是否正确
|
||||
```
|
||||
|
||||
### 问题排查
|
||||
|
||||
``` plaintext
|
||||
运行 bench_batch_prefill_ragged.py 时报错 out of memory，请帮我分析原因并给出解决方案。
|
||||
代码理解
|
||||
请帮我解释 bench_batch_prefill_ragged.py 中 BatchPrefillWithRaggedKVCacheWrapper 的 plan() 和 run() 方法的工作原理，特别是 qo_indptr 和 kv_indptr 的作用。
|
||||
整理优化日志
|
||||
请帮我整理本次优化的记录，包括：
|
||||
1. 原始 baseline 性能数据（从 CSV 中提取关键配置的 time_ms 和 tflops）
|
||||
2. 优化后的性能数据（从 OJ 评测结果中提取）
|
||||
3. 加速比 = baseline_time / optimized_time
|
||||
4. 以表格形式输出：配置参数 | baseline 耗时 | 优化后耗时 | 加速比
|
||||
```
|
||||
|
||||
## 八、常见问题
|
||||
|
||||
### 流程问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| 登录后看不到题目 | 未使用赛用账号登录 | 七月份组委会统一发放 XPU-OJ 账号，请确认你使用的是组委会统一发放的账号，而不是自行注册账号；如仍无法看到题目，请联系助教或赛事运营确认账号权限。 |
|
||||
|
||||
### 环境问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
|
|
@ -284,7 +684,7 @@ if csv_files:
|
|||
| `ModuleNotFoundError: No module named 'flashinfer'` | flashinfer 未安装 | 执行 `pip install flashinfer` |
|
||||
| `out of memory` | GPU 显存不足 | 减小 `batch_size` 或 `seq_len` 参数 |
|
||||
|
||||
### 运行相关问题
|
||||
### 运行问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
|
|
@ -292,21 +692,32 @@ if csv_files:
|
|||
| `KeyError: 'BatchPrefillWithPagedKVCacheKernel'` | profiler 未捕获目标 kernel | 检查 `target_kernels` 配置是否正确 |
|
||||
| CSV 文件为空 | 测试未正常完成 | 检查 GPU 显存是否充足，重新运行 |
|
||||
|
||||
### 代码相关问题
|
||||
### 代码问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| `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` 配置 |
|
||||
|
||||
### 性能相关问题
|
||||
### 性能问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| TFLOPs 数值异常低 | 工作负载过小，kernel 启动开销占比大 | 增大 `batch_size` 或 `seq_len` |
|
||||
| 带宽数值异常低 | 数据未正确加载到 GPU | 检查 Tensor 是否在 CUDA 设备上 |
|
||||
|
||||
## 八、下一步学习建议
|
||||
### 评测问题
|
||||
|
||||
提交 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` 的 `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` |
|
||||
|
||||
|
||||
## 九、下一步学习建议
|
||||
|
||||
### 1. 保存你的 Baseline 结果
|
||||
|
||||
|
|
@ -316,7 +727,6 @@ if csv_files:
|
|||
# 建议创建 results 目录保存
|
||||
mkdir -p results
|
||||
mv *.csv results/
|
||||
|
||||
```
|
||||
|
||||
### 2. 深入理解 FlashInfer 核心概念
|
||||
|
|
@ -346,11 +756,18 @@ mv *.csv results/
|
|||
* FlashInfer 文档：https://flashinfer.ai
|
||||
|
||||
|
||||
### 5. 记录优化过程
|
||||
### 5. 记录优化流程
|
||||
|
||||
建议维护一份优化日志，记录每次优化的改动和性能变化：
|
||||
|
||||
| 优化项 | 改动内容 | Baseline | 优化后 | 提升比例 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% |
|
||||
| 完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。 | | | | |
|
||||
|
||||
完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。
|
||||
|
||||
> 使用 Agent 整理优化日志，可形成可复现的 Agent/Skill 优化流程
|
||||
|
||||
### 6. 使用多语言完成算子优化加速
|
||||
|
||||
可以使用 Triton 或 TileLang 语言实现 `run_kernel` 接口完成算子优化，对比不同的语言对于性能加速的影响。
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,160,0.0580,45.25
|
||||
2,512,8,160,0.0611,85.84
|
||||
4,512,8,160,0.0656,159.96
|
||||
8,512,8,160,0.0699,300.40
|
||||
16,512,8,160,0.1321,317.92
|
||||
32,512,8,160,0.2002,419.52
|
||||
64,512,8,160,0.3383,496.43
|
||||
128,512,8,160,0.6669,503.61
|
||||
1,1024,8,160,0.1129,46.45
|
||||
2,1024,8,160,0.1190,88.18
|
||||
4,1024,8,160,0.1224,171.49
|
||||
8,1024,8,160,0.1287,326.07
|
||||
16,1024,8,160,0.2479,338.49
|
||||
32,1024,8,160,0.3767,445.54
|
||||
64,1024,8,160,0.6419,523.01
|
||||
128,1024,8,160,1.2804,524.37
|
||||
1,2048,8,160,0.2270,46.20
|
||||
2,2048,8,160,0.2299,91.22
|
||||
4,2048,8,160,0.2349,178.63
|
||||
8,2048,8,160,0.2447,342.96
|
||||
16,2048,8,160,0.4773,351.60
|
||||
32,2048,8,160,0.7279,461.07
|
||||
64,2048,8,160,1.2559,534.49
|
||||
128,2048,8,160,2.5613,524.15
|
||||
1,4096,8,160,0.4460,47.02
|
||||
2,4096,8,160,0.4513,92.94
|
||||
4,4096,8,160,0.4593,182.64
|
||||
8,4096,8,160,0.4813,348.64
|
||||
16,4096,8,160,0.9363,358.43
|
||||
32,4096,8,160,1.4552,461.21
|
||||
64,4096,8,160,2.5615,524.05
|
||||
128,4096,8,160,5.1420,522.11
|
||||
1,8192,8,160,0.8847,47.41
|
||||
2,8192,8,160,0.8944,93.80
|
||||
4,8192,8,160,0.9094,184.51
|
||||
8,8192,8,160,0.9625,348.64
|
||||
16,8192,8,160,1.8550,361.80
|
||||
32,8192,8,160,2.9567,453.97
|
||||
64,8192,8,160,5.1398,522.30
|
||||
128,8192,8,160,10.2972,521.41
|
||||
1,16384,8,160,1.7608,47.64
|
||||
2,16384,8,160,1.7786,94.33
|
||||
4,16384,8,160,1.8143,184.95
|
||||
8,16384,8,160,1.9317,347.42
|
||||
16,16384,8,160,3.7301,359.83
|
||||
32,16384,8,160,5.9216,453.33
|
||||
64,16384,8,160,10.2668,522.94
|
||||
128,16384,8,160,20.6062,521.09
|
||||
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,192,0.0458,68.82
|
||||
2,512,8,192,0.0515,122.32
|
||||
4,512,8,192,0.0574,219.28
|
||||
8,512,8,192,0.0607,414.80
|
||||
16,512,8,192,0.1147,439.27
|
||||
32,512,8,192,0.1763,571.40
|
||||
64,512,8,192,0.2978,676.79
|
||||
128,512,8,192,0.5874,686.11
|
||||
1,1024,8,192,0.0946,66.55
|
||||
2,1024,8,192,0.1033,121.85
|
||||
4,1024,8,192,0.1073,234.66
|
||||
8,1024,8,192,0.1131,445.23
|
||||
16,1024,8,192,0.2165,465.24
|
||||
32,1024,8,192,0.3347,601.80
|
||||
64,1024,8,192,0.5701,706.63
|
||||
128,1024,8,192,1.1302,712.88
|
||||
1,2048,8,192,0.1943,64.79
|
||||
2,2048,8,192,0.1992,126.38
|
||||
4,2048,8,192,0.2059,244.52
|
||||
8,2048,8,192,0.2174,463.13
|
||||
16,2048,8,192,0.4202,479.24
|
||||
32,2048,8,192,0.6503,619.36
|
||||
64,2048,8,192,1.1158,721.93
|
||||
128,2048,8,192,2.2250,724.05
|
||||
1,4096,8,192,0.3834,65.65
|
||||
2,4096,8,192,0.3904,128.95
|
||||
4,4096,8,192,0.4043,249.04
|
||||
8,4096,8,192,0.4267,471.92
|
||||
16,4096,8,192,0.8271,486.90
|
||||
32,4096,8,192,1.2840,627.28
|
||||
64,4096,8,192,2.2148,727.29
|
||||
128,4096,8,192,4.3819,735.21
|
||||
1,8192,8,192,0.7566,66.52
|
||||
2,8192,8,192,0.7712,130.54
|
||||
4,8192,8,192,0.7974,252.49
|
||||
8,8192,8,192,0.8433,477.47
|
||||
16,8192,8,192,1.6433,490.09
|
||||
32,8192,8,192,2.5573,629.84
|
||||
64,8192,8,192,4.3785,735.73
|
||||
128,8192,8,192,8.7303,737.99
|
||||
1,16384,8,192,1.5068,66.81
|
||||
2,16384,8,192,1.5350,131.16
|
||||
4,16384,8,192,1.5868,253.76
|
||||
8,16384,8,192,1.6778,479.99
|
||||
16,16384,8,192,3.2750,491.81
|
||||
32,16384,8,192,5.0659,635.88
|
||||
64,16384,8,192,8.7435,736.85
|
||||
128,16384,8,192,17.5040,736.13
|
||||
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,224,0.1254,29.29
|
||||
2,512,8,224,0.1412,52.05
|
||||
4,512,8,224,0.1497,98.13
|
||||
8,512,8,224,0.1533,191.70
|
||||
16,512,8,224,0.1913,307.17
|
||||
32,512,8,224,0.3292,357.08
|
||||
64,512,8,224,0.5187,453.28
|
||||
128,512,8,224,0.9522,493.84
|
||||
1,1024,8,224,0.2727,26.93
|
||||
2,1024,8,224,0.2836,51.78
|
||||
4,1024,8,224,0.2890,101.63
|
||||
8,1024,8,224,0.2959,198.55
|
||||
16,1024,8,224,0.3696,317.93
|
||||
32,1024,8,224,0.6408,366.75
|
||||
64,1024,8,224,1.0081,466.21
|
||||
128,1024,8,224,1.8548,506.78
|
||||
1,2048,8,224,0.5515,26.63
|
||||
2,2048,8,224,0.5575,52.67
|
||||
4,2048,8,224,0.5666,103.65
|
||||
8,2048,8,224,0.5803,202.42
|
||||
16,2048,8,224,0.7250,324.05
|
||||
32,2048,8,224,1.2593,373.14
|
||||
64,2048,8,224,1.9890,472.48
|
||||
128,2048,8,224,3.6905,509.28
|
||||
1,4096,8,224,1.0939,26.84
|
||||
2,4096,8,224,1.1044,53.18
|
||||
4,4096,8,224,1.1219,104.69
|
||||
8,4096,8,224,1.1500,204.26
|
||||
16,4096,8,224,1.4390,326.48
|
||||
32,4096,8,224,2.4992,375.97
|
||||
64,4096,8,224,4.0082,468.86
|
||||
128,4096,8,224,7.3372,512.26
|
||||
1,8192,8,224,2.1775,26.97
|
||||
2,8192,8,224,2.1989,53.41
|
||||
4,8192,8,224,2.2338,105.15
|
||||
8,8192,8,224,2.3268,201.90
|
||||
16,8192,8,224,2.8806,326.18
|
||||
32,8192,8,224,5.0187,374.43
|
||||
64,8192,8,224,8.0323,467.90
|
||||
128,8192,8,224,14.6300,513.78
|
||||
1,16384,8,224,4.3360,27.09
|
||||
2,16384,8,224,4.3820,53.60
|
||||
4,16384,8,224,4.5006,104.38
|
||||
8,16384,8,224,4.6987,199.96
|
||||
16,16384,8,224,5.7361,327.59
|
||||
32,16384,8,224,10.1291,371.03
|
||||
64,16384,8,224,16.0745,467.60
|
||||
128,16384,8,224,OOM,OOM
|
||||
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,256,0.0877,47.89
|
||||
2,512,8,256,0.0921,91.17
|
||||
4,512,8,256,0.0940,178.74
|
||||
8,512,8,256,0.0964,348.52
|
||||
16,512,8,256,0.1450,463.27
|
||||
32,512,8,256,0.2250,597.21
|
||||
64,512,8,256,0.3609,744.43
|
||||
128,512,8,256,0.6932,775.25
|
||||
1,1024,8,256,0.1747,48.04
|
||||
2,1024,8,256,0.1762,95.27
|
||||
4,1024,8,256,0.1784,188.22
|
||||
8,1024,8,256,0.1817,369.53
|
||||
16,1024,8,256,0.2796,480.25
|
||||
32,1024,8,256,0.4339,619.00
|
||||
64,1024,8,256,0.6960,771.73
|
||||
128,1024,8,256,1.3439,799.36
|
||||
1,2048,8,256,0.3410,49.21
|
||||
2,2048,8,256,0.3439,97.60
|
||||
4,2048,8,256,0.3469,193.52
|
||||
8,2048,8,256,0.3533,379.94
|
||||
16,2048,8,256,0.5461,491.67
|
||||
32,2048,8,256,0.8493,632.28
|
||||
64,2048,8,256,1.3667,785.82
|
||||
128,2048,8,256,2.6465,811.64
|
||||
1,4096,8,256,0.6742,49.77
|
||||
2,4096,8,256,0.6777,99.03
|
||||
4,4096,8,256,0.6836,196.36
|
||||
8,4096,8,256,0.6950,386.31
|
||||
16,4096,8,256,1.0803,497.02
|
||||
32,4096,8,256,1.6794,639.44
|
||||
64,4096,8,256,2.7101,792.50
|
||||
128,4096,8,256,5.2543,817.52
|
||||
1,8192,8,256,1.3375,50.18
|
||||
2,8192,8,256,1.3448,99.81
|
||||
4,8192,8,256,1.3564,197.91
|
||||
8,8192,8,256,1.3799,389.08
|
||||
16,8192,8,256,2.1465,500.25
|
||||
32,8192,8,256,3.3342,644.12
|
||||
64,8192,8,256,5.3983,795.67
|
||||
128,8192,8,256,10.4691,820.55
|
||||
1,16384,8,256,2.6697,50.28
|
||||
2,16384,8,256,2.6817,100.10
|
||||
4,16384,8,256,2.7049,198.49
|
||||
8,16384,8,256,2.7533,390.00
|
||||
16,16384,8,256,4.2789,501.89
|
||||
32,16384,8,256,6.6476,646.11
|
||||
64,16384,8,256,10.7723,797.43
|
||||
128,16384,8,256,OOM,OOM
|
||||
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,32,0.0257,20.45
|
||||
2,512,8,32,0.0256,41.02
|
||||
4,512,8,32,0.0258,81.28
|
||||
8,512,8,32,0.0265,158.45
|
||||
16,512,8,32,0.0396,212.30
|
||||
32,512,8,32,0.0516,325.43
|
||||
64,512,8,32,0.0721,465.83
|
||||
128,512,8,32,0.1270,529.03
|
||||
1,1024,8,32,0.0461,22.75
|
||||
2,1024,8,32,0.0465,45.15
|
||||
4,1024,8,32,0.0477,88.04
|
||||
8,1024,8,32,0.0548,153.23
|
||||
16,1024,8,32,0.0734,228.71
|
||||
32,1024,8,32,0.0958,350.42
|
||||
64,1024,8,32,0.1334,503.15
|
||||
128,1024,8,32,0.2381,564.04
|
||||
1,2048,8,32,0.0872,24.06
|
||||
2,2048,8,32,0.0904,46.42
|
||||
4,2048,8,32,0.1028,81.59
|
||||
8,2048,8,32,0.1067,157.25
|
||||
16,2048,8,32,0.1428,235.10
|
||||
32,2048,8,32,0.1818,369.13
|
||||
64,2048,8,32,0.2554,525.57
|
||||
128,2048,8,32,0.4622,580.86
|
||||
1,4096,8,32,0.1730,24.25
|
||||
2,4096,8,32,0.1955,42.91
|
||||
4,4096,8,32,0.2020,83.05
|
||||
8,4096,8,32,0.2140,156.83
|
||||
16,4096,8,32,0.2777,241.65
|
||||
32,4096,8,32,0.3542,378.99
|
||||
64,4096,8,32,0.4990,538.05
|
||||
128,4096,8,32,0.9099,590.13
|
||||
1,8192,8,32,0.3820,21.96
|
||||
2,8192,8,32,0.3913,42.88
|
||||
4,8192,8,32,0.4127,81.31
|
||||
8,8192,8,32,0.4224,158.88
|
||||
16,8192,8,32,0.5490,244.51
|
||||
32,8192,8,32,0.6960,385.70
|
||||
64,8192,8,32,0.9870,543.98
|
||||
128,8192,8,32,1.8100,593.25
|
||||
1,16384,8,32,0.7655,21.92
|
||||
2,16384,8,32,0.8067,41.59
|
||||
4,16384,8,32,0.8228,81.56
|
||||
8,16384,8,32,0.8397,159.85
|
||||
16,16384,8,32,1.0910,246.04
|
||||
32,16384,8,32,1.3824,388.37
|
||||
64,16384,8,32,1.9663,546.08
|
||||
128,16384,8,32,3.6107,594.78
|
||||
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,512,0.3588,23.40
|
||||
2,512,8,512,0.3651,46.00
|
||||
4,512,8,512,0.3736,89.89
|
||||
8,512,8,512,0.3856,174.22
|
||||
16,512,8,512,0.7472,179.80
|
||||
32,512,8,512,1.1447,234.72
|
||||
64,512,8,512,1.9549,274.89
|
||||
128,512,8,512,3.8962,275.85
|
||||
1,1024,8,512,0.7261,23.12
|
||||
2,1024,8,512,0.7354,45.65
|
||||
4,1024,8,512,0.7496,89.57
|
||||
8,1024,8,512,0.7746,173.35
|
||||
16,1024,8,512,1.5049,178.46
|
||||
32,1024,8,512,2.3111,232.42
|
||||
64,1024,8,512,3.9538,271.70
|
||||
128,1024,8,512,7.8811,272.62
|
||||
1,2048,8,512,1.4636,22.93
|
||||
2,2048,8,512,1.4826,45.27
|
||||
4,2048,8,512,1.5109,88.86
|
||||
8,2048,8,512,1.5549,172.68
|
||||
16,2048,8,512,3.0237,177.60
|
||||
32,2048,8,512,4.6439,231.27
|
||||
64,2048,8,512,7.9560,269.99
|
||||
128,2048,8,512,15.8741,270.63
|
||||
1,4096,8,512,2.9312,22.90
|
||||
2,4096,8,512,2.9675,45.24
|
||||
4,4096,8,512,3.0243,88.77
|
||||
8,4096,8,512,3.1127,172.50
|
||||
16,4096,8,512,6.0753,176.76
|
||||
32,4096,8,512,9.3182,230.49
|
||||
64,4096,8,512,15.9642,269.07
|
||||
128,4096,8,512,31.8313,269.89
|
||||
1,8192,8,512,5.8843,22.81
|
||||
2,8192,8,512,5.9344,45.24
|
||||
4,8192,8,512,6.0465,88.80
|
||||
8,8192,8,512,6.2334,172.27
|
||||
16,8192,8,512,12.1594,176.62
|
||||
32,8192,8,512,18.6826,229.90
|
||||
64,8192,8,512,32.0055,268.41
|
||||
128,8192,8,512,OOM,OOM
|
||||
1,16384,8,512,11.8153,22.72
|
||||
2,16384,8,512,11.9237,45.03
|
||||
4,16384,8,512,12.1671,88.25
|
||||
8,16384,8,512,12.4948,171.88
|
||||
16,16384,8,512,24.3414,176.45
|
||||
32,16384,8,512,37.3907,229.74
|
||||
64,16384,8,512,OOM,OOM
|
||||
128,16384,8,512,OOM,OOM
|
||||
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,64,0.0404,25.99
|
||||
2,512,8,64,0.0399,52.60
|
||||
4,512,8,64,0.0413,101.69
|
||||
8,512,8,64,0.0482,174.25
|
||||
16,512,8,64,0.0540,310.86
|
||||
32,512,8,64,0.0629,533.75
|
||||
64,512,8,64,0.0833,806.14
|
||||
128,512,8,64,0.1104,1216.59
|
||||
1,1024,8,64,0.0747,28.08
|
||||
2,1024,8,64,0.0766,54.77
|
||||
4,1024,8,64,0.0891,94.17
|
||||
8,1024,8,64,0.0918,182.94
|
||||
16,1024,8,64,0.1044,321.41
|
||||
32,1024,8,64,0.1179,569.43
|
||||
64,1024,8,64,0.1566,857.28
|
||||
128,1024,8,64,0.2078,1292.17
|
||||
1,2048,8,64,0.1455,28.84
|
||||
2,2048,8,64,0.1684,49.82
|
||||
4,2048,8,64,0.1730,97.01
|
||||
8,2048,8,64,0.1850,181.39
|
||||
16,2048,8,64,0.2009,334.18
|
||||
32,2048,8,64,0.2268,592.01
|
||||
64,2048,8,64,0.3002,894.44
|
||||
128,2048,8,64,0.4027,1333.64
|
||||
1,4096,8,64,0.3265,25.69
|
||||
2,4096,8,64,0.3322,50.51
|
||||
4,4096,8,64,0.3522,95.27
|
||||
8,4096,8,64,0.3632,184.79
|
||||
16,4096,8,64,0.3942,340.56
|
||||
32,4096,8,64,0.4456,602.47
|
||||
64,4096,8,64,0.5927,905.94
|
||||
128,4096,8,64,0.7938,1352.87
|
||||
1,8192,8,64,0.6508,25.78
|
||||
2,8192,8,64,0.6879,48.78
|
||||
4,8192,8,64,0.7008,95.77
|
||||
8,8192,8,64,0.7199,186.44
|
||||
16,8192,8,64,0.7786,344.79
|
||||
32,8192,8,64,0.8798,610.25
|
||||
64,8192,8,64,1.1745,914.30
|
||||
128,8192,8,64,1.5728,1365.50
|
||||
1,16384,8,64,1.3524,24.81
|
||||
2,16384,8,64,1.3698,48.99
|
||||
4,16384,8,64,1.3923,96.40
|
||||
8,16384,8,64,1.4267,188.16
|
||||
16,16384,8,64,1.5451,347.47
|
||||
32,16384,8,64,1.7622,609.32
|
||||
64,16384,8,64,2.3392,918.09
|
||||
128,16384,8,64,3.1332,1370.84
|
||||
|
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,96,0.0407,38.67
|
||||
2,512,8,96,0.0398,79.02
|
||||
4,512,8,96,0.0431,146.08
|
||||
8,512,8,96,0.0495,254.61
|
||||
16,512,8,96,0.0698,360.64
|
||||
32,512,8,96,0.1117,450.87
|
||||
64,512,8,96,0.1780,566.16
|
||||
128,512,8,96,0.3329,605.28
|
||||
1,1024,8,96,0.0732,43.01
|
||||
2,1024,8,96,0.0794,79.29
|
||||
4,1024,8,96,0.0871,144.54
|
||||
8,1024,8,96,0.0934,269.52
|
||||
16,1024,8,96,0.1297,388.14
|
||||
32,1024,8,96,0.2114,476.36
|
||||
64,1024,8,96,0.3379,596.08
|
||||
128,1024,8,96,0.6327,636.68
|
||||
1,2048,8,96,0.1505,41.80
|
||||
2,2048,8,96,0.1619,77.76
|
||||
4,2048,8,96,0.1713,146.94
|
||||
8,2048,8,96,0.1780,282.84
|
||||
16,2048,8,96,0.2492,404.09
|
||||
32,2048,8,96,0.4088,492.55
|
||||
64,2048,8,96,0.6575,612.55
|
||||
128,2048,8,96,1.2457,646.61
|
||||
1,4096,8,96,0.3099,40.61
|
||||
2,4096,8,96,0.3259,77.23
|
||||
4,4096,8,96,0.3346,150.42
|
||||
8,4096,8,96,0.3467,290.41
|
||||
16,4096,8,96,0.4888,411.94
|
||||
32,4096,8,96,0.8055,499.94
|
||||
64,4096,8,96,1.3209,609.72
|
||||
128,4096,8,96,2.4810,649.25
|
||||
1,8192,8,96,0.6343,39.68
|
||||
2,8192,8,96,0.6437,78.20
|
||||
4,8192,8,96,0.6601,152.50
|
||||
8,8192,8,96,0.6826,294.97
|
||||
16,8192,8,96,0.9688,415.64
|
||||
32,8192,8,96,1.6057,501.55
|
||||
64,8192,8,96,2.6527,607.19
|
||||
128,8192,8,96,4.9464,651.27
|
||||
1,16384,8,96,1.2581,40.01
|
||||
2,16384,8,96,1.2812,78.57
|
||||
4,16384,8,96,1.3112,153.55
|
||||
8,16384,8,96,1.3653,294.92
|
||||
16,16384,8,96,1.9351,416.16
|
||||
32,16384,8,96,3.2277,499.01
|
||||
64,16384,8,96,5.3192,605.60
|
||||
128,16384,8,96,9.8747,652.44
|
||||
|
|
|
@ -0,0 +1,330 @@
|
|||
# guide
|
||||
|
||||
# FlashAttention KV Cache Decode - 参赛指南
|
||||
|
||||
## 一、登录 XPU-OJ 平台
|
||||
|
||||
打开浏览器，访问：\*\*https://xpuoj.com/\*\*
|
||||
|
||||
1. 等待组委会统一发放 XPU-OJ 账号;
|
||||
|
||||
2. 使用分配的用户名和初始密码登录平台
|
||||
|
||||
|
||||
|
||||
|
||||
点击 **"登录"** 进入平台。
|
||||
|
||||
## 二、进入比赛
|
||||
|
||||
### 2.1 点击导航栏"比赛"
|
||||
|
||||
登录成功后,会跳转到 XPUOJ 平台首页。
|
||||
|
||||
找到页面顶部的导航栏,点击 **"比赛"** 标签,进入比赛列表页。
|
||||
|
||||
|
||||
|
||||
### 2.2 找到目标比赛
|
||||
|
||||
在比赛列表中,根据比赛状态(全部 / 未开始 / 进行中 / 已结束)找到目标比赛。
|
||||
|
||||
|
||||
|
||||
## 三、进入题目
|
||||
|
||||
### 3.1 在题目列表中找到目标题目
|
||||
|
||||
进入比赛后,页面会展示该比赛的**题目列表**。每道题都有编号(1 ~ N)和标题。
|
||||
|
||||
在列表中找到 **FlashAttention KV Cache Decode** 这道题(前缀为 `FlashAttention`),点击标题即可进入题目详情页。
|
||||
|
||||
|
||||
|
||||
### 3.2 查看题目要求
|
||||
|
||||
进入题目详情页后,可以看到以下几个区域:
|
||||
|
||||
* **左侧**:题目描述、接口约定、参数说明
|
||||
|
||||
* **右侧**:代码编辑器,用于编写并提交代码
|
||||
|
||||
|
||||
请仔细阅读左侧的 **题目描述** 和 **接口约定**,重点关注:
|
||||
|
||||
* 入口函数名(本题为 `run_kernel`)
|
||||
|
||||
* 必传的参数列表及其类型、顺序
|
||||
|
||||
* 编译/运行环境(语言选择,目标硬件 C500)
|
||||
|
||||
|
||||
### 3.3 编写并提交代码
|
||||
|
||||
在右侧的代码编辑器中,按照题目要求填入完整代码(可点击右上角"重置"恢复初始模板)。
|
||||
|
||||
#### 最小正确性代码(可先复制跑通)
|
||||
|
||||
为方便参赛者先跑通完整流程,这里提供一份 **最小正确性代码**,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常:
|
||||
|
||||
```cpp
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
|
||||
#define HEAD_DIM 128
|
||||
|
||||
|
||||
__global__ void paged_attention_kernel(
|
||||
|
||||
const __nv_bfloat16* q,
|
||||
|
||||
const __nv_bfloat16* k_cache_paged,
|
||||
|
||||
const __nv_bfloat16* v_cache_paged,
|
||||
|
||||
__nv_bfloat16* output,
|
||||
|
||||
const int32_t* cache_seqlens,
|
||||
|
||||
const int32_t* block_table,
|
||||
|
||||
int64_t batch_size,
|
||||
|
||||
int64_t seqlen_q,
|
||||
|
||||
int64_t num_heads,
|
||||
|
||||
int64_t num_heads_k,
|
||||
|
||||
int64_t headdim,
|
||||
|
||||
int64_t page_block_size,
|
||||
|
||||
int64_t blocks_per_batch)
|
||||
|
||||
{
|
||||
|
||||
int batch_idx = blockIdx.x / num_heads;
|
||||
|
||||
int head_idx = blockIdx.x % num_heads;
|
||||
|
||||
if (batch_idx >= batch_size || head_idx >= num_heads) return;
|
||||
|
||||
|
||||
int seqlen = cache_seqlens[batch_idx];
|
||||
|
||||
int tid = threadIdx.x;
|
||||
|
||||
|
||||
// 加载对应 head 的 query 元素
|
||||
|
||||
int64_t q_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim;
|
||||
|
||||
float q_val = __bfloat162float(q[q_offset + tid]);
|
||||
|
||||
|
||||
// Online safe softmax 状态
|
||||
|
||||
float max_val = -1e38f;
|
||||
|
||||
float sum_exp = 0.0f;
|
||||
|
||||
float out_acc = 0.0f;
|
||||
|
||||
float scale = 1.0f / sqrtf(static_cast<float>(headdim));
|
||||
|
||||
|
||||
// 静态共享内存,避免动态分配可能带来的兼容性问题
|
||||
|
||||
__shared__ float s_score[HEAD_DIM];
|
||||
|
||||
|
||||
for (int token = 0; token < seqlen; ++token) {
|
||||
|
||||
int page_idx = token / page_block_size;
|
||||
|
||||
int page_offset = token % page_block_size;
|
||||
|
||||
int physical_block = block_table[batch_idx * blocks_per_batch + page_idx];
|
||||
|
||||
|
||||
// 读取 key 元素
|
||||
|
||||
const __nv_bfloat16* k_ptr = k_cache_paged
|
||||
|
||||
+ (physical_block * page_block_size + page_offset) * (num_heads_k * headdim)
|
||||
|
||||
+ head_idx * headdim;
|
||||
|
||||
float k_val = __bfloat162float(k_ptr[tid]);
|
||||
|
||||
|
||||
// 点积 -> 共享内存归约
|
||||
|
||||
s_score[tid] = q_val * k_val;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
|
||||
for (int stride = HEAD_DIM >> 1; stride > 0; stride >>= 1) {
|
||||
|
||||
if (tid < stride) {
|
||||
|
||||
s_score[tid] += s_score[tid + stride];
|
||||
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
}
|
||||
|
||||
float score = s_score[0] * scale;
|
||||
|
||||
|
||||
// 更新 softmax 状态
|
||||
|
||||
float new_max = fmaxf(max_val, score);
|
||||
|
||||
float rescale = expf(max_val - new_max);
|
||||
|
||||
sum_exp = sum_exp * rescale + expf(score - new_max);
|
||||
|
||||
out_acc = out_acc * rescale;
|
||||
|
||||
max_val = new_max;
|
||||
|
||||
|
||||
// 读取 value 元素,并累加(用最新 max 的权重)
|
||||
|
||||
const __nv_bfloat16* v_ptr = v_cache_paged
|
||||
|
||||
+ (physical_block * page_block_size + page_offset) * (num_heads_k * headdim)
|
||||
|
||||
+ head_idx * headdim;
|
||||
|
||||
float v_val = __bfloat162float(v_ptr[tid]);
|
||||
|
||||
out_acc += expf(score - max_val) * v_val;
|
||||
|
||||
|
||||
__syncthreads(); // 确保下次迭代共享内存可安全复用
|
||||
|
||||
}
|
||||
|
||||
|
||||
if (seqlen > 0) {
|
||||
|
||||
out_acc /= sum_exp;
|
||||
|
||||
} else {
|
||||
|
||||
out_acc = 0.0f;
|
||||
|
||||
}
|
||||
|
||||
|
||||
int64_t out_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim + tid;
|
||||
|
||||
output[out_offset] = __float2bfloat16(out_acc);
|
||||
|
||||
}
|
||||
|
||||
|
||||
extern "C" void run_kernel(
|
||||
|
||||
const __nv_bfloat16* q,
|
||||
|
||||
const __nv_bfloat16* k_cache_paged,
|
||||
|
||||
const __nv_bfloat16* v_cache_paged,
|
||||
|
||||
__nv_bfloat16* output,
|
||||
|
||||
const int32_t* cache_seqlens,
|
||||
|
||||
const int32_t* block_table,
|
||||
|
||||
int64_t batch_size,
|
||||
|
||||
int64_t seqlen_k,
|
||||
|
||||
int64_t seqlen_q,
|
||||
|
||||
int64_t num_heads,
|
||||
|
||||
int64_t num_heads_k,
|
||||
|
||||
int64_t headdim,
|
||||
|
||||
int64_t page_block_size,
|
||||
|
||||
int64_t num_blocks,
|
||||
|
||||
int64_t causal)
|
||||
|
||||
{
|
||||
|
||||
int64_t blocks_per_batch = num_blocks / batch_size;
|
||||
|
||||
dim3 grid(batch_size * num_heads);
|
||||
|
||||
dim3 block(HEAD_DIM);
|
||||
|
||||
|
||||
paged_attention_kernel<<<grid, block>>>(
|
||||
|
||||
q, k_cache_paged, v_cache_paged, output,
|
||||
|
||||
cache_seqlens, block_table,
|
||||
|
||||
batch_size, seqlen_q, num_heads, num_heads_k, headdim,
|
||||
|
||||
page_block_size, blocks_per_batch
|
||||
|
||||
);
|
||||
|
||||
}
|
||||
```
|
||||
|
||||
完成后:
|
||||
|
||||
1. 在编辑器下方 **语言** 选项中,根据代码实际情况选择对应语言(可选 `Triton` / `CUDA Maca` / `MXMACA C++` / `TileLang` 等)
|
||||
|
||||
2. 选择 **目标硬件** 为 `C500`
|
||||
|
||||
3. 点击右上角 **"提交"** 按钮
|
||||
|
||||
|
||||
|
||||
|
||||
## 四、查看提交结果
|
||||
|
||||
点击 **"提交"** 后,系统会自动跳转到提交记录页,展示本次提交的详细信息。
|
||||
|
||||
页面顶部会显示一行汇总信息,包括:
|
||||
|
||||
* **状态** —— 评测结果(如 `Accepted` 表示通过)
|
||||
|
||||
* **分数** —— 本次提交获得的分数
|
||||
|
||||
* **题目** —— 对应的题目名称
|
||||
|
||||
* **用时** —— 程序运行耗时
|
||||
|
||||
* **内存** —— 占用内存大小
|
||||
|
||||
* **答案** —— 提交所用的语言/硬件
|
||||
|
||||
* **提交时间** —— 提交的时刻
|
||||
|
||||
|
||||
|
||||
|
||||
页面下方会展开 **编译信息** 和 **各测试点**(样例、测试点 #1 ~ #N)的结果,逐个显示:
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
{
|
||||
"id": 197,
|
||||
"displayId": 20005,
|
||||
"type": "Traditional",
|
||||
"isPublic": false,
|
||||
"locales": [
|
||||
"zh_CN"
|
||||
],
|
||||
"samples": [
|
||||
{
|
||||
"inputData": "1\n",
|
||||
"outputData": ""
|
||||
}
|
||||
],
|
||||
"problemTagIds": []
|
||||
}
|
||||
|
|
@ -0,0 +1,298 @@
|
|||
from __future__ import annotations
|
||||
|
||||
|
||||
HEAD_DIMS = [128]
|
||||
BATCH_SIZES = [1, 4, 16]
|
||||
SEQ_LENS_KV = [1024, 4096, 8192, 16384]
|
||||
SEQ_LEN_Q = 1
|
||||
NUM_HEADS = 8
|
||||
NUM_HEADS_K = 8
|
||||
PAGE_BLOCK_SIZE = 16
|
||||
CAUSAL = 0
|
||||
|
||||
|
||||
def _build_cases():
|
||||
cases = []
|
||||
for headdim in HEAD_DIMS:
|
||||
for seqlen_k in SEQ_LENS_KV:
|
||||
for batch_size in BATCH_SIZES:
|
||||
cases.append(
|
||||
(
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
SEQ_LEN_Q,
|
||||
NUM_HEADS,
|
||||
NUM_HEADS_K,
|
||||
headdim,
|
||||
PAGE_BLOCK_SIZE,
|
||||
CAUSAL,
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
TESTCASES = _build_cases()
|
||||
|
||||
|
||||
def getNumOfTestcases() -> int:
|
||||
return len(TESTCASES)
|
||||
|
||||
|
||||
try:
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
import math
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
KernelArg = Union[torch.Tensor, int, float]
|
||||
CURRENT_CASE = None
|
||||
|
||||
def _ensure_flashattn_importable():
|
||||
try:
|
||||
from flash_attn.flash_attn_interface import flash_attn_with_kvcache # noqa: F401
|
||||
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
here = Path(__file__).resolve()
|
||||
for parent in here.parents:
|
||||
candidate = parent / "flashattn"
|
||||
if (candidate / "flash_attn").is_dir():
|
||||
sys.path.insert(0, str(candidate))
|
||||
return
|
||||
|
||||
def _get_testcase_index() -> int:
|
||||
try:
|
||||
raw = input().strip()
|
||||
except EOFError:
|
||||
return 0
|
||||
if raw == "":
|
||||
return 0
|
||||
try:
|
||||
testcase_id = int(raw.split()[0])
|
||||
except ValueError:
|
||||
return 0
|
||||
if 1 <= testcase_id <= len(TESTCASES):
|
||||
return testcase_id - 1
|
||||
if 0 <= testcase_id < len(TESTCASES):
|
||||
return testcase_id
|
||||
return 0
|
||||
|
||||
def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int:
|
||||
workload = batch_size * seq_len * head_dim
|
||||
if workload < 1e5:
|
||||
return base_reps
|
||||
if workload < 1e6:
|
||||
return base_reps // 2
|
||||
if workload < 1e7:
|
||||
return base_reps // 4
|
||||
if workload < 1e8:
|
||||
return base_reps // 8
|
||||
if workload < 1e9:
|
||||
return base_reps // 16
|
||||
return base_reps // 32
|
||||
|
||||
def _get_num_blocks(batch_size: int, seqlen_k: int, page_block_size: int) -> int:
|
||||
num_blocks = math.ceil(seqlen_k / page_block_size) * batch_size * 3
|
||||
return max(1024, num_blocks)
|
||||
|
||||
def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]:
|
||||
testcase_id = _get_testcase_index()
|
||||
global CURRENT_CASE
|
||||
(
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
causal,
|
||||
) = TESTCASES[testcase_id]
|
||||
num_blocks = _get_num_blocks(batch_size, seqlen_k, page_block_size)
|
||||
blocks_per_batch = num_blocks // batch_size
|
||||
CURRENT_CASE = (
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
causal,
|
||||
20260720 + testcase_id,
|
||||
)
|
||||
warmup = 3
|
||||
iters = max(1, _compute_reps(batch_size, seqlen_k, headdim))
|
||||
return [
|
||||
(batch_size, seqlen_q, num_heads, headdim),
|
||||
(num_blocks, page_block_size, num_heads_k, headdim),
|
||||
(num_blocks, page_block_size, num_heads_k, headdim),
|
||||
(batch_size, seqlen_q, num_heads, headdim),
|
||||
(batch_size,),
|
||||
(batch_size, blocks_per_batch),
|
||||
(), (), (), (), (), (), (), (), (),
|
||||
], (warmup, iters)
|
||||
|
||||
def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]:
|
||||
del testcase_sizes
|
||||
(
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
causal,
|
||||
seed,
|
||||
) = CURRENT_CASE
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
blocks_per_batch = num_blocks // batch_size
|
||||
q = torch.randn(
|
||||
batch_size,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
headdim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
k_cache_paged = torch.randn(
|
||||
num_blocks,
|
||||
page_block_size,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
v_cache_paged = torch.randn(
|
||||
num_blocks,
|
||||
page_block_size,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
output = torch.empty(
|
||||
batch_size,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
headdim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
cache_seqlens = torch.full((batch_size,), seqlen_k, dtype=torch.int32, device=device)
|
||||
block_table = torch.randperm(num_blocks, dtype=torch.int32, device=device, generator=gen).reshape(
|
||||
batch_size,
|
||||
blocks_per_batch,
|
||||
)
|
||||
return [
|
||||
q,
|
||||
k_cache_paged,
|
||||
v_cache_paged,
|
||||
output,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
causal,
|
||||
]
|
||||
|
||||
def baseline(
|
||||
q,
|
||||
k_cache_paged,
|
||||
v_cache_paged,
|
||||
output,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
causal,
|
||||
):
|
||||
_ensure_flashattn_importable()
|
||||
from flash_attn.flash_attn_interface import flash_attn_with_kvcache
|
||||
|
||||
out = flash_attn_with_kvcache(
|
||||
q,
|
||||
k_cache_paged,
|
||||
v_cache_paged,
|
||||
None,
|
||||
None,
|
||||
cache_seqlens=cache_seqlens,
|
||||
cache_batch_idx=None,
|
||||
block_table=block_table,
|
||||
causal=bool(causal),
|
||||
window_size=(-1, -1),
|
||||
rotary_interleaved=False,
|
||||
alibi_slopes=None,
|
||||
num_splits=1,
|
||||
)
|
||||
output.copy_(out)
|
||||
return [
|
||||
q,
|
||||
k_cache_paged,
|
||||
v_cache_paged,
|
||||
output,
|
||||
cache_seqlens,
|
||||
block_table,
|
||||
batch_size,
|
||||
seqlen_k,
|
||||
seqlen_q,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
headdim,
|
||||
page_block_size,
|
||||
num_blocks,
|
||||
causal,
|
||||
]
|
||||
|
||||
def check(
|
||||
testcase_sizes,
|
||||
original_input_tensors,
|
||||
target_kernel_input_tensors,
|
||||
baseline_input_tensors,
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
) -> bool:
|
||||
del testcase_sizes, original_input_tensors
|
||||
output_t = target_kernel_input_tensors[3]
|
||||
output_ref = baseline_input_tensors[3]
|
||||
if output_t.shape != output_ref.shape:
|
||||
print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr)
|
||||
return False
|
||||
if output_t.dtype != output_ref.dtype:
|
||||
print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr)
|
||||
return False
|
||||
if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol):
|
||||
diff = (output_t.float() - output_ref.float()).abs()
|
||||
print(
|
||||
f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, "
|
||||
f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
---
|
||||
sectionTitle: "题目描述"
|
||||
type: "Text"
|
||||
---
|
||||
你需要实现 FlashAttention paged KV cache decode 的 CUDA C++ 前向算子。
|
||||
|
||||
本题输入采用 `flash_attn_with_kvcache` 在 `flashattn/benchmarks/benchmark_kvcache.py` 中使用的 paged KV cache 配置。每个 batch 只有 1 个 query token，KV cache 长度为 `seqlen_k`，K/V cache 按 page 存储。
|
||||
|
||||
评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `cache_seqlens` 和 `block_table` 读取 paged KV cache，并将结果写入 `output`。
|
||||
|
||||
baseline 使用 benchmark 中的 FlashAttention Python API：
|
||||
|
||||
```python
|
||||
out = flash_attn_with_kvcache(
|
||||
q, k_cache_paged, v_cache_paged, None, None,
|
||||
cache_seqlens=cache_seqlens,
|
||||
cache_batch_idx=None,
|
||||
block_table=block_table,
|
||||
causal=False,
|
||||
window_size=(-1, -1),
|
||||
rotary_interleaved=False,
|
||||
alibi_slopes=None,
|
||||
num_splits=1,
|
||||
)
|
||||
output.copy_(out)
|
||||
```
|
||||
|
||||
如何提交代码详见[评测指南](/d/2)。
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
---
|
||||
sectionTitle: "接口约定"
|
||||
type: "codeSample"
|
||||
lang: "cuda"
|
||||
---
|
||||
你必须在提交的 CUDA 源码中提供如下 **C 符号**，函数名、参数类型、顺序必须完全一致，并使用 `extern "C"` 防止 name mangling：
|
||||
|
||||
```cpp
|
||||
#include <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
extern "C" void run_kernel(
|
||||
const __nv_bfloat16* q,
|
||||
const __nv_bfloat16* k_cache_paged,
|
||||
const __nv_bfloat16* v_cache_paged,
|
||||
__nv_bfloat16* output,
|
||||
const int32_t* cache_seqlens,
|
||||
const int32_t* block_table,
|
||||
int64_t batch_size,
|
||||
int64_t seqlen_k,
|
||||
int64_t seqlen_q,
|
||||
int64_t num_heads,
|
||||
int64_t num_heads_k,
|
||||
int64_t headdim,
|
||||
int64_t page_block_size,
|
||||
int64_t num_blocks,
|
||||
int64_t causal
|
||||
);
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
* `q`：decode query tensor，shape `(batch_size, seqlen_q, num_heads, headdim)`，连续 `bf16`
|
||||
* `k_cache_paged`：paged key cache，shape `(num_blocks, page_block_size, num_heads_k, headdim)`，连续 `bf16`
|
||||
* `v_cache_paged`：paged value cache，shape `(num_blocks, page_block_size, num_heads_k, headdim)`，连续 `bf16`
|
||||
* `output`：输出缓冲区，shape `(batch_size, seqlen_q, num_heads, headdim)`，连续 `bf16`
|
||||
* `cache_seqlens`：每个 batch 的 KV 长度，shape `(batch_size)`，连续 `int32`
|
||||
* `block_table`：每个 batch 的 page 映射表，shape `(batch_size, num_blocks / batch_size)`，连续 `int32`
|
||||
* `seqlen_q`：query 长度，评测中固定为 `1`
|
||||
* `page_block_size`：page size，评测中固定为 `16`
|
||||
* `causal`：是否启用 causal mask，评测中固定为 `0`
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 `run_kernel` 内部做 `cudaDeviceSynchronize()` 或显式同步。
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
---
|
||||
sectionTitle: "接口约定"
|
||||
type: "codeSample"
|
||||
lang: "tilelang"
|
||||
---
|
||||
你必须在提交的 Python 代码中提供 `run_kernel` 函数，函数名、参数顺序、类型必须完全一致：
|
||||
|
||||
```python
|
||||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang import jit
|
||||
|
||||
real_kernel = None
|
||||
|
||||
@jit
|
||||
def build_kernel(*args):
|
||||
@T.prim_func
|
||||
def kernel(*args):
|
||||
...
|
||||
return kernel
|
||||
|
||||
def run_kernel(
|
||||
q, # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
|
||||
k_cache_paged, # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
|
||||
v_cache_paged, # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
|
||||
output, # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
|
||||
cache_seqlens, # Tensor[int32], shape (batch_size)
|
||||
block_table, # Tensor[int32], shape (batch_size, num_blocks / batch_size)
|
||||
batch_size, # int64
|
||||
seqlen_k, # int64
|
||||
seqlen_q, # int64
|
||||
num_heads, # int64
|
||||
num_heads_k, # int64
|
||||
headdim, # int64
|
||||
page_block_size, # int64
|
||||
num_blocks, # int64
|
||||
causal, # int64
|
||||
):
|
||||
global real_kernel
|
||||
if real_kernel is None:
|
||||
real_kernel = build_kernel(...)
|
||||
real_kernel(q, k_cache_paged, v_cache_paged, output,
|
||||
cache_seqlens, block_table,
|
||||
batch_size, seqlen_k, seqlen_q, num_heads,
|
||||
num_heads_k, headdim, page_block_size, num_blocks, causal)
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
* `q`：decode query tensor，连续 `bfloat16`
|
||||
* `k_cache_paged/v_cache_paged`：paged KV cache，连续 `bfloat16`
|
||||
* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果
|
||||
* `cache_seqlens/block_table`：paged KV metadata，连续 `int32`
|
||||
* `page_block_size`：评测中固定为 `16`
|
||||
* `causal`：评测中固定为 `0`
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 TileLang kernel。
|
||||
|
|
@ -0,0 +1,45 @@
|
|||
---
|
||||
sectionTitle: "接口约定"
|
||||
type: "codeSample"
|
||||
lang: "triton"
|
||||
---
|
||||
你必须在提交的 Python 代码中提供 `run_kernel` 函数，函数名、参数顺序、类型必须完全一致：
|
||||
|
||||
```python
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
@triton.jit
|
||||
def your_kernel(...):
|
||||
...
|
||||
|
||||
def run_kernel(
|
||||
q, # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
|
||||
k_cache_paged, # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
|
||||
v_cache_paged, # Tensor[bf16], shape (num_blocks, page_block_size, num_heads_k, headdim)
|
||||
output, # Tensor[bf16], shape (batch_size, seqlen_q, num_heads, headdim)
|
||||
cache_seqlens, # Tensor[int32], shape (batch_size)
|
||||
block_table, # Tensor[int32], shape (batch_size, num_blocks / batch_size)
|
||||
batch_size, # int64
|
||||
seqlen_k, # int64
|
||||
seqlen_q, # int64
|
||||
num_heads, # int64
|
||||
num_heads_k, # int64
|
||||
headdim, # int64
|
||||
page_block_size, # int64
|
||||
num_blocks, # int64
|
||||
causal, # int64
|
||||
):
|
||||
...
|
||||
```
|
||||
|
||||
### 参数说明
|
||||
|
||||
* `q`：decode query tensor，连续 `bfloat16`
|
||||
* `k_cache_paged/v_cache_paged`：paged KV cache，连续 `bfloat16`
|
||||
* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果
|
||||
* `cache_seqlens/block_table`：paged KV metadata，连续 `int32`
|
||||
* `page_block_size`：评测中固定为 `16`
|
||||
* `causal`：评测中固定为 `0`
|
||||
|
||||
`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 Triton kernel。
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
---
|
||||
sectionTitle: "输入格式"
|
||||
type: "Text"
|
||||
---
|
||||
本题输入由评测程序在 GPU 上构造，并按接口约定中的顺序传入 `run_kernel`。
|
||||
|
||||
`q/k_cache_paged/v_cache_paged/output` 均为连续 `torch.bfloat16` CUDA tensor，`cache_seqlens/block_table` 均为连续 `torch.int32` CUDA tensor。
|
||||
|
||||
KV cache layout 固定为 `flash_attn_with_kvcache` 的 paged cache 布局：`(num_blocks, page_block_size, num_heads_k, headdim)`。
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
---
|
||||
sectionTitle: "输出格式"
|
||||
type: "Text"
|
||||
---
|
||||
输出写入 `output`，shape 为 `(batch_size, 1, num_heads, headdim)`，类型为 `bfloat16`。
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
---
|
||||
sectionTitle: "样例"
|
||||
type: "Text"
|
||||
---
|
||||
若 `batch_size = 1`、`seqlen_k = 512`、`page_block_size = 16`，则每个序列需要访问 `32` 个有效 page：
|
||||
|
||||
```text
|
||||
cache_seqlens = [512]
|
||||
block_table.shape = (1, num_blocks)
|
||||
```
|
||||
|
||||
第 `t` 个 KV token 位于 `block_table[0, t / 16]` 指向的物理 page 中，page 内偏移为 `t % 16`。
|
||||
|
|
@ -0,0 +1 @@
|
|||
FlashAttention KV Cache Decode
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,16 @@
|
|||
{
|
||||
"id": 193,
|
||||
"displayId": 20001,
|
||||
"type": "Traditional",
|
||||
"isPublic": false,
|
||||
"locales": [
|
||||
"zh_CN"
|
||||
],
|
||||
"samples": [
|
||||
{
|
||||
"inputData": "1\n",
|
||||
"outputData": ""
|
||||
}
|
||||
],
|
||||
"problemTagIds": []
|
||||
}
|
||||
|
|
@ -0,0 +1,307 @@
|
|||
from __future__ import annotations
|
||||
|
||||
|
||||
HEAD_DIM_CONFIGS = [(128, 128)]
|
||||
BATCH_SIZES = [1, 4, 16]
|
||||
SEQ_LENS = [1024, 4096, 8192, 16384]
|
||||
NUM_QO_HEADS = 32
|
||||
NUM_KV_HEADS = 4
|
||||
CAUSAL = 1
|
||||
|
||||
|
||||
def _build_cases():
|
||||
cases = []
|
||||
for head_dim_qk, head_dim_vo in HEAD_DIM_CONFIGS:
|
||||
for batch_size in BATCH_SIZES:
|
||||
for seq_len in SEQ_LENS:
|
||||
cases.append(
|
||||
(
|
||||
batch_size,
|
||||
seq_len,
|
||||
NUM_QO_HEADS,
|
||||
NUM_KV_HEADS,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
CAUSAL,
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
TESTCASES = _build_cases()
|
||||
|
||||
|
||||
def getNumOfTestcases() -> int:
|
||||
return len(TESTCASES)
|
||||
|
||||
|
||||
try:
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple, Union
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
KernelArg = Union[torch.Tensor, int, float]
|
||||
CURRENT_CASE = None
|
||||
|
||||
def _ensure_flashinfer_importable():
|
||||
try:
|
||||
import flashinfer # noqa: F401
|
||||
|
||||
return
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
here = Path(__file__).resolve()
|
||||
for parent in here.parents:
|
||||
candidate = parent / "McFlashInfer"
|
||||
if (candidate / "flashinfer").is_dir():
|
||||
sys.path.insert(0, str(candidate))
|
||||
return
|
||||
|
||||
def _get_testcase_index() -> int:
|
||||
try:
|
||||
raw = input().strip()
|
||||
except EOFError:
|
||||
return 0
|
||||
if raw == "":
|
||||
return 0
|
||||
try:
|
||||
testcase_id = int(raw.split()[0])
|
||||
except ValueError:
|
||||
return 0
|
||||
if 1 <= testcase_id <= len(TESTCASES):
|
||||
return testcase_id - 1
|
||||
if 0 <= testcase_id < len(TESTCASES):
|
||||
return testcase_id
|
||||
return 0
|
||||
|
||||
def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int:
|
||||
workload = batch_size * seq_len * head_dim
|
||||
if workload < 1e5:
|
||||
return base_reps
|
||||
if workload < 1e6:
|
||||
return base_reps // 2
|
||||
if workload < 1e7:
|
||||
return base_reps // 4
|
||||
if workload < 1e8:
|
||||
return base_reps // 8
|
||||
if workload < 1e9:
|
||||
return base_reps // 16
|
||||
return base_reps // 32
|
||||
|
||||
def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]:
|
||||
testcase_id = _get_testcase_index()
|
||||
global CURRENT_CASE
|
||||
batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal = TESTCASES[testcase_id]
|
||||
CURRENT_CASE = (
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
20260610 + testcase_id,
|
||||
)
|
||||
qo_len = batch_size * seq_len
|
||||
kv_len = batch_size * seq_len
|
||||
warmup = 3
|
||||
iters = max(1, _compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo))
|
||||
return [
|
||||
(qo_len, num_qo_heads, head_dim_qk),
|
||||
(kv_len, num_kv_heads, head_dim_qk),
|
||||
(kv_len, num_kv_heads, head_dim_vo),
|
||||
(qo_len, num_qo_heads, head_dim_vo),
|
||||
(batch_size + 1,),
|
||||
(batch_size + 1,),
|
||||
(), (), (), (), (), (), (),
|
||||
], (warmup, iters)
|
||||
|
||||
def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]:
|
||||
del testcase_sizes
|
||||
batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal, seed = CURRENT_CASE
|
||||
gen = torch.Generator(device=device)
|
||||
gen.manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
qo_len = batch_size * seq_len
|
||||
q = torch.rand(
|
||||
qo_len,
|
||||
num_qo_heads,
|
||||
head_dim_qk,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
kv_len = batch_size * seq_len
|
||||
k = torch.rand(
|
||||
kv_len,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
v = torch.rand(
|
||||
kv_len,
|
||||
num_kv_heads,
|
||||
head_dim_vo,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
generator=gen,
|
||||
).contiguous()
|
||||
output = torch.empty(
|
||||
qo_len,
|
||||
num_qo_heads,
|
||||
head_dim_vo,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len
|
||||
kv_indptr = qo_indptr.clone()
|
||||
return [
|
||||
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,
|
||||
]
|
||||
|
||||
def baseline(
|
||||
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,
|
||||
):
|
||||
_ensure_flashinfer_importable()
|
||||
import flashinfer
|
||||
|
||||
workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device)
|
||||
wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
|
||||
workspace_buffer,
|
||||
kv_layout="NHD",
|
||||
backend="auto",
|
||||
)
|
||||
wrapper.plan(
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
int(num_qo_heads),
|
||||
int(num_kv_heads),
|
||||
int(head_dim_qk),
|
||||
int(head_dim_vo),
|
||||
causal=bool(causal),
|
||||
q_data_type=torch.bfloat16,
|
||||
kv_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, k, v, out=output)
|
||||
return [
|
||||
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,
|
||||
]
|
||||
|
||||
def check(
|
||||
testcase_sizes,
|
||||
original_input_tensors,
|
||||
target_kernel_input_tensors,
|
||||
baseline_input_tensors,
|
||||
rtol=1e-2,
|
||||
atol=1e-2,
|
||||
) -> bool:
|
||||
del testcase_sizes, original_input_tensors
|
||||
output_t = target_kernel_input_tensors[3]
|
||||
output_ref = baseline_input_tensors[3]
|
||||
if output_t.shape != output_ref.shape:
|
||||
print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr)
|
||||
return False
|
||||
if output_t.dtype != output_ref.dtype:
|
||||
print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr)
|
||||
return False
|
||||
if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol):
|
||||
diff = (output_t.float() - output_ref.float()).abs()
|
||||
print(
|
||||
f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, "
|
||||
f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
INPUT_CLASS = [
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"OUTPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
"INPUT",
|
||||
]
|
||||
|
||||
|
||||
def getWorkload(testcase_sizes) -> dict:
|
||||
raw_sizes = testcase_sizes[0] if isinstance(testcase_sizes, tuple) and len(testcase_sizes) == 2 else testcase_sizes
|
||||
q_shape, k_shape, v_shape, output_shape, qo_indptr_shape, kv_indptr_shape = raw_sizes[:6]
|
||||
qo_len, num_qo_heads, head_dim_qk = q_shape
|
||||
kv_len, num_kv_heads, k_dim = k_shape
|
||||
v_len, v_heads, head_dim_vo = v_shape
|
||||
assert k_dim == head_dim_qk
|
||||
assert v_len == kv_len
|
||||
assert v_heads == num_kv_heads
|
||||
assert qo_len == kv_len
|
||||
assert output_shape == (qo_len, num_qo_heads, head_dim_vo)
|
||||
assert qo_indptr_shape == kv_indptr_shape
|
||||
batch_size = qo_indptr_shape[0] - 1
|
||||
seq_len = kv_len // batch_size
|
||||
flops = batch_size * seq_len * seq_len * num_qo_heads * (head_dim_qk + head_dim_vo)
|
||||
memory_bytes = (
|
||||
qo_len * num_qo_heads * head_dim_qk * 2
|
||||
+ kv_len * num_kv_heads * head_dim_qk * 2
|
||||
+ kv_len * num_kv_heads * head_dim_vo * 2
|
||||
+ qo_len * num_qo_heads * head_dim_vo * 2
|
||||
+ (batch_size + 1) * 4 * 2
|
||||
)
|
||||
return {
|
||||
"flops": flops,
|
||||
"memory_bytes": memory_bytes,
|
||||
"dtype": "bf16",
|
||||
}
|
||||
|
||||
|
||||
DESIGNED_VRAM_SIZE = 48
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
---
|
||||
sectionTitle: "题目描述"
|
||||
type: "Text"
|
||||
---
|
||||
你需要实现 FlashInfer ragged KV cache prefill 的CUDA C++前向算子。
|
||||
|
||||
本题输入采用 FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` 的 ragged `NHD` 布局。每个 batch 中有 `seq_len` 个 query token，KV cache 中也有 `seq_len` 个 token：
|
||||
|
||||
其中 query heads 采用 GQA 布局：`num_qo_heads` 个 query/output heads 共享 `num_kv_heads` 个 KV heads，`G = num_qo_heads / num_kv_heads`。
|
||||
|
||||
评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `qo_indptr` 和 `kv_indptr` 读取 ragged Q/K/V，并将结果写入 `output`。
|
||||
|
||||
baseline 使用 FlashInfer ragged prefill 的 Python API：
|
||||
|
||||
```python
|
||||
wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto")
|
||||
wrapper.plan(qo_indptr, kv_indptr, num_qo_heads, num_kv_heads,
|
||||
head_dim_qk, head_dim_vo, causal=True,
|
||||
q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16)
|
||||
wrapper.run(q, k, v, out=output)
|
||||
```
|
||||
|
||||
如何提交代码详见[评测指南](/d/2)。
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue