diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md
index bc3f857..bfdc174 100644
--- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md
+++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md
@@ -7,9 +7,13 @@
需要特别说明：
- 本教程不提供可直接提交的 OJ 参考实现源代码或标准答案代码。
+
- 本教程仅提供冒烟级 starter 示例代码，用于验证环境、语言、提交链路和 `run_kernel(...)` 接口。
+
- benchmark 脚本用于建立性能基线，不是最终提交物。
+
- XPU-OJ 题包中的 `baseline()` 属于 OJ 后台参考实现，用于生成 `output_ref`，不是选手提交代码。
+
- 选手最终需要自行实现 `run_kernel(...)`，并在正确性通过后继续优化性能。
完成本教程后，学员应能够跑通 benchmark 脚本，记录性能基线结果，读懂 XPU-OJ 题包，理解 OJ 的测试输入与参考实现，并完成一次冒烟级 OJ 提交。
@@ -63,7 +67,7 @@
#### 4.1.1 获取算力代金券
-[点击获取算力代金券](https://developer.metax-tech.com/activities/6)，首次登录需要使用邮箱或者手机号进行注册。登录成功后提交申请获得兑换码。
+[*点击获取算力代金券*](https://developer.metax-tech.com/activities/6)，首次登录需要使用邮箱或者手机号进行注册。登录成功后提交申请获得兑换码。
@@ -71,7 +75,7 @@
使用兑换码兑换 GPU 租用余额：
-访问[模力方舟官网](https://ai.gitee.com/)，在左侧边栏进入 “费用中心”，点击右上角 “兑换” 使用兑换码兑换代金券；
+访问 [*模力方舟官网*](https://ai.gitee.com/)，在左侧边栏进入 “费用中心”，点击右上角 “兑换” 使用兑换码兑换代金券；
@@ -80,7 +84,7 @@
1. 左侧边栏进入 “算力容器”，点击右上角 “租用算力”，在新打开页面中的筛选选项中选择 “沐曦”，选择可租的 GPU 点击 “立即租用”。建议优先选择 16GB / 32GB 显存；
- 
+ 
2. 进入 “创建实例” 页面，确认计费方式为 “按量收费”，预装镜像选择：基础镜像 - **PyTorch-Agent / 2.8.0 / Python 3.12 / maca 3.7.2.1**，点击下一步；
@@ -88,7 +92,7 @@
3. 勾选同意服务条款，点击 “创建实例”；
- > 如果在这一步出现 “账户剩余可用金额不足” 的提示，请联系赛事运营方或 IT 管理员
+ > 如果在这一步出现 “账户剩余可用金额不足” 的提示，请联系助教或赛事运营
4. 显示 “支付成功” 后，点击 “查看资源”，可以看到 “容器示例” 界面中当前的实例名称、状态、存储占用、创建时间、累计运行时长、工具、开 / 关机选项、操作等等信息；
@@ -122,6 +126,8 @@
+ 此部分内容可参考教程：[*模力方舟快速使用SOP*](../模力方舟快速使用SOP.md)
+
#### 4.1.4 深度学习环境配置
镜像中的 base 环境已提供：
@@ -138,7 +144,7 @@ numpy 1.26.4
pip install pandas
```
-是否已配置 MXMACA / MACA 软件栈：
+确认 MXMACA / MACA 软件栈已配置：
``` bash
# 检查 MXMACA 运行时库路径
@@ -155,7 +161,7 @@ which mxcc && mxcc --version || echo "mxcc 未找到，请确认 MACA 工具链
| 问题 | 解决方法 |
| --- | --- |
-| `mxcc: command not found` | MACA 工具链未安装或 `PATH` 未配置，检查镜像是否预装或参考 [模力方舟快速使用SOP](https://www.gitlink.org.cn/metax-maca/op_optimization) |
+| `mxcc: command not found` | MACA 工具链未安装或 `PATH` 未配置，检查镜像是否预装或参考 [*模力方舟快速使用SOP*](../模力方舟快速使用SOP.md) |
### 4.2 工具准备
@@ -165,8 +171,10 @@ which mxcc && mxcc --version || echo "mxcc 未找到，请确认 MACA 工具链
- 已确认 Agent 可以正常调用模型。
+配置过程可参考 [*模力方舟Agent部署准备教程*](模力方舟Agent部署准备教程.md)。
-#### 以配置 OpenCode 为例
+
+**以配置 OpenCode 为例**
命令行环境安装：
@@ -175,7 +183,7 @@ curl -fsSL https://opencode.ai/install | bash
opencode
```
-后续配置教程可参考 [OpenCode 官方文档](https://opencode.ai/docs/)。
+后续配置教程可参考 [*OpenCode 官方文档*](https://opencode.ai/docs/)。
### 4.3 代码准备
@@ -186,7 +194,7 @@ opencode
- 已获取测试脚本和 Benchmark 脚本。
-（具体操作可见章节 [6.1 在赛事镜像中运行 FlashInfer Benchmark](#61-在赛事镜像中运行-flashinfer-benchmark)）
+> 具体操作可见章节 [*6.1 在赛事镜像中运行 FlashInfer Benchmark*](#run-flashinfer-bench)
### 4.4 账号准备
@@ -201,7 +209,7 @@ XPU-OJ 账号由组委会统一发放，参赛者无需自行注册。
- **Prefill：** Prefill 阶段是指处理输入 prompt 的阶段
- - 输入：用户一次性给出的完整 prompt，长度为 seq\_len
+ - 输入：用户一次性给出的完整 prompt，长度为 `seq \ len`
- 计算：对 prompt 中的每个 token 并行计算注意力，生成第一个输出 token 及 KV cache
@@ -213,9 +221,10 @@ XPU-OJ 账号由组委会统一发放，参赛者无需自行注册。
- **显存带宽密集型（memory-bound）**，瓶颈在从显存读取 KV cache 而非计算
-**Prefill = 并行处理用户输入，Decode = 逐个生成回答 token**
+- **Prefill = 并行处理用户输入，Decode = 逐个生成回答 token**
-### 5.2 benchmark / 性能基线
+
+### 5.2 benchmark / 性能基线
`benchmark/` 目录中的脚本用于运行原库或迁移库的性能测试，帮助选手理解目标 API、输入输出 shape、性能指标和瓶颈位置。benchmark 输出的 CSV、日志或结果为 “性能基线结果”。
@@ -230,14 +239,15 @@ XPU-OJ 账号由组委会统一发放，参赛者无需自行注册。
### 5.3 开源仓库参考
-[GitHub - MetaX-MACA/McFlashInfer · GitHub](https://github.com/MetaX-MACA/McFlashInfer)
+[*GitHub - MetaX-MACA/McFlashInfer*](https://github.com/MetaX-MACA/McFlashInfer)
> 链接内容可供用于学习 API、算子实现思路、benchmark 方法和优化策略。选手仍需根据 XPU-OJ 题包接口**自行实现**可提交的 `run_kernel(...)`
## 6. 项目实践 -- FlashInfer Benchmark
-**目标：** 以一个具体算子题为例（FlashInfer Ragged Prefill），跑通 benchmark 脚本，建立性能基线，理解对应 XPU-OJ 题目包，并完成一次 OJ 冒烟提交，为后续 Agent 辅助优化建立起点。
+**目标：** 以一个具体算子题目 **20001 FlashInfer Ragged Prefill** 为例，跑通 benchmark 脚本，建立性能基线，理解对应 XPU-OJ 题目包，并完成一次 OJ 冒烟提交，为后续 Agent 辅助优化建立起点。
+
### 6.1 在赛事镜像中运行 FlashInfer Benchmark
#### Step 1：检查运行环境
@@ -246,7 +256,7 @@ XPU-OJ 账号由组委会统一发放，参赛者无需自行注册。
**操作：** 进入 Terminal 检查 GPU、Python、编译工具和依赖版本。
-
+
**命令示例：**
@@ -271,21 +281,21 @@ pip install pandas
- `mx-smi` 显示沐曦 GPU 信息
- 
+ 
- Python 环境正常
- 
+ 
- `torch.cuda.is_available()` 返回 `True`
- 
+ 
- 所有依赖版本符合要求
- 
+ 
- 
+ 
**常见问题：**
@@ -298,7 +308,7 @@ pip install pandas
#### Step 2：进入项目目录
-**目标：** 进入本模块所需的源码目录 [benchmark](https://www.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%2Foperator_task_package%2Fflashinfer_task_package%2FFlashInfer_Baseline)。
+**目标：** 进入本模块所需的源码目录 [*benchmark*](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%2Foperator_task_package%2Fflashinfer_task_package%2Fbenchmark)。
1. 克隆代码仓库
@@ -308,30 +318,21 @@ pip install pandas
2. 准备 benchmark
- 从仓库根目录开始，在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `operator_task_package/flashinfer_task_package` 文件夹。可以将 `FlashInfer_Baseline` 整个目录复制到工作目录 `data/` 下。
+ 从克隆到本地的代码仓库中复制 flashinfer_task_package 文件夹到工作目录 `data/` 下：
```bash
- mkdir data
- cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package" data/
+ cp -r ./op_optimization/基于AI\ Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/ .
```
3. 切换到项目目录 benchmark
```bash
- cd data/operator_task_package/flashinfer_task_package/benchmark
- ls
+ cd ./flashinfer_task_package/benchmark
+ ls -l
```
-**预期结果：**
+ **预期结果：**
-```plaintext
-bench_common.py
-bench_batch_decode.py
-bench_batch_prefill_paged.py
-bench_batch_prefill_ragged.py
-bench_batch_mla.py
-README.md
-...
-```
+ 
#### Step 3：验证项目脚本
@@ -364,21 +365,21 @@ python -c "from bench_common import setup_workspace, get_csv_path; print('脚本
**目标：** 执行基准测试，获取性能基线数据，查看并分析 Benchmark 输出结果。
-**操作：** 运行 Ragged Prefill 基准测试脚本，读取生成的 CSV 结果文件。
+**操作：** 运行基准测试脚本 （以 Ragged Prefill Benchmark 为例），读取生成的 CSV 结果文件。
-> 每个算子优化题目都对应一个 Benchmark（见 [5.2 查看性能基线](#52-benchmark--性能基线)）
+> 每个算子优化题目都对应一个 Benchmark（见 [*5.2 查看性能基线*](#benchmark-baseline)）
-**运行算子Benchmark：（以 Ragged Prefill Benchmark 为例）**
+**运行算子 Benchmark（以 Ragged Prefill 为例）：**
```bash
python bench_batch_prefill_ragged.py
```
-预期结果：（并非真实数据）
+**预期结果：**
``` plaintext
[BatchPrefillWithRaggedKVCacheWrapper] Starting benchmark, total cases: 48
- [1/48] bs=1, sl=1024, hd=[128,128]: 0.032ms, 66.67 GB/s, 272.00 TFLOPs
- [2/48] bs=1, sl=4096, hd=[128,128]: 0.042ms, 197.83 GB/s, 3238.06 TFLOPs
- [3/48] bs=1, sl=8192, hd=[128,128]: 0.064ms, 261.50 GB/s, 8559.25 TFLOPs
+ [1/48] bs=1, sl=1024, hd=[128,128]: 0.032ms, 65.99 GB/s, 269.25 TFLOPs
+ [2/48] bs=1, sl=4096, hd=[128,128]: 0.046ms, 182.63 GB/s, 2989.26 TFLOPs
+ [3/48] bs=1, sl=8192, hd=[128,128]: 0.057ms, 293.87 GB/s, 9624.79 TFLOPs
...
Results saved to BatchPrefillWithRaggedKVCacheWrapper_20260626_xxxxxx.csv
@@ -388,7 +389,7 @@ Results saved to BatchPrefillWithRaggedKVCacheWrapper_20260626_xxxxxx.csv
| 问题 | 解决方法 |
| --- | --- |
-| `out of memory` | 减小 batch\_size 或 seq\_len 参数 |
+| `out of memory` | 减小 `batch_size` 或 `seq_len` 参数 |
| 运行时间过长 | 脚本会自动调整重复次数，耐心等待 |
***
@@ -423,15 +424,18 @@ else:
**预期结果（以 Ragged Prefill 为例）：**
```plaintext
-读取文件: BatchPrefillWithRaggedKVCacheWrapper_20260626_145454.csv
- batch_size seq_len num_qo_heads num_kv_heads head_dim_qk head_dim_vo time_ms bandwidth_GB_s tflops
- 1 1024 32 4 128 128 0.031580 66.666667 272.004150
- 1 4096 32 4 128 128 0.042445 197.828709 3238.063402
- 1 8192 32 4 128 128 0.064123 261.499213 8559.251770
- 4 1024 32 4 128 128 0.050221 167.754590 2737.757998
- 4 4096 32 4 128 128 0.101234 332.907816 21734.876630
- 16 1024 32 4 128 128 0.149876 224.887654 14683.437981
- ...
+读取文件: BatchPrefillWithRaggedKVCacheWrapper_20260624_085632.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.031903 65.992618 269.253988
+ 1 4096 32 4 128 128 0.045978 182.628062 2989.258976
+ 1 8192 32 4 128 128 0.057119 293.868770 9624.792255
+ 1 16384 32 4 128 128 0.069130 485.498445 31809.859991
+ 4 1024 32 4 128 128 0.044687 188.450962 768.891659
+ 4 4096 32 4 128 128 0.065587 512.099922 8382.059516
+ 4 8192 32 4 128 128 0.091556 733.340790 24018.383268
+ 4 16384 32 4 128 128 0.145408 923.267606 60492.497127
+ 16 1024 32 4 128 128 0.072724 463.193467 1889.858181
+ 16 4096 32 4 128 128 0.158293 848.733154 13892.077507
```
### 6.2 XPU-OJ 在线评测教程
@@ -452,10 +456,15 @@ else:
跑完 benchmark、建立性能基线后，选手需要完成以下转换：
1. 从 benchmark 脚本中理解目标 API，例如 `BatchPrefillWithRaggedKVCacheWrapper`；
+
2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口；
+
3. 对照题包中的输入 shape、数据范围和精度要求；
+
4. 编写自己的 `run_kernel(...)`；
+
5. 提交 OJ，先通过正确性；
+
6. 正确性通过后，再对比 benchmark 耗时 / OJ 耗时继续优化。
**选择目标题目**
@@ -471,7 +480,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
**每个子题的接口参数、数据范围和精度要求以对应 XPU-OJ 题包为准。** 下文以题目 **20001 Flashinfer Ragged Prefill** 为例演示从 benchmark 到 XPU-OJ 提交的完整流程。
-有关题目的更多信息参考 [20001 Flashinfer Ragged Prefill 题目说明](https://www.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%2Foperator_task_package%2Fflashinfer_task_package%2Fxpuoj_problem%2Fproblem_20001%2Fzh_CN)。
+题目的更多信息可参考 [*20001 Flashinfer Ragged Prefill 题目说明*](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%2Foperator_task_package%2Fflashinfer_task_package%2Fxpuoj_problem%2Fproblem_20001%2Fzh_CN)。
#### Step 6：理解 XPU-OJ 评测接口与精度要求
**目标：** 明确 Benchmark 与最终评测提交之间的关系，理解选手需要实现的内容。
@@ -479,7 +488,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
> 完成 benchmark 后，需要注意 benchmark 脚本主要用于建立性能基线，并不需要最终提交
> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 OJ 参考实现结果进行比较
-下面以 FlashInfer Ragged Prefill 题为例，逐份解读 OJ 题包中的关键文件。
+下面以题目 **20001 FlashInfer Ragged Prefill** 为例，逐份解读 OJ 题包中的关键文件。
1. `00_题目描述.md` — 我要实现什么？
@@ -521,6 +530,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
);
```
+
| 参数 | 类型 | Shape | 含义 |
|------|------|-------|------|
| `q` | `const __nv_bfloat16*` | `(batch_size × seq_len, 32, 128)` | Query 张量（ragged 压平） |
@@ -543,7 +553,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
- **`output` 是预分配的空缓冲区**：你必须将结果写入其中，不要在内部自行分配显存。
- - **`run_kernel` 内自行 launch**：需要在函数体内计算 grid / block 配置并 `<<>>` 启动 CUDA kernel。**不要调用 `cudaDeviceSynchronize()`**，OJ 评测期会统一同步。
+ - **`run_kernel` 内自行 launch**：需要在函数体内计算 grid / block 配置并 `<<>>` 启动 CUDA kernel。**不要调用** `cudaDeviceSynchronize()`，OJ 评测期会统一同步。
- **GQA 头映射**：`kv_head = qo_head / (num_qo_heads / num_kv_heads) = qo_head / 8`。
@@ -574,21 +584,22 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
这个文件不是给选手阅读的文档，但理解它的逻辑有助于定位 OJ 提交错误：
- 1. **`TESTCASES`**：定义 12 组测试参数（`batch_size × seq_len` 组合），与 benchmark 脚本的参数顺序一致。
- 2. **`genTestCase()`**：OJ 用固定随机种子生成测试张量，同一测试用例每次运行数据相同。
- 3. **`baseline()`**：OJ 的参考实现——调用 FlashInfer Python API 得到 `output_ref`。
- 4. **`check()`**：比较你的 `output` 与 `output_ref`。失败时返回 `max_abs_diff` 和 `mean_abs_diff` 帮助定位数值偏差。
- 5. **`getWorkload()`**：计算该测试点的 FLOPs 和显存读写量，用于估算硬件理论下限 `T_h`（以此计算得分）。
+ 1. `TESTCASES`：定义 12 组测试参数（`batch_size × seq_len` 组合），与 benchmark 脚本的参数顺序一致。
+ 2. `genTestCase()`：OJ 用固定随机种子生成测试张量，同一测试用例每次运行数据相同。
+ 3. `baseline()`：OJ 的参考实现——调用 FlashInfer Python API 得到 `output_ref`。
+ 4. `check()`：比较你的 `output` 与 `output_ref`。失败时返回 `max_abs_diff` 和 `mean_abs_diff` 帮助定位数值偏差。
+ 5. `getWorkload()`：计算该测试点的 FLOPs 和显存读写量，用于估算硬件理论下限 `T_h`（以此计算得分）。
- 当前 FlashInfer Ragged Prefill 的校验方式为：
+ 当前题目 **20001 FlashInfer Ragged Prefill** 的校验方式为：
```python
torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2)
```
说明：
- `output_t` 是选手 `run_kernel(...)` 的输出；
+
- `output_ref` 是 OJ 题包参考实现生成的输出；
- - `rtol / atol` 是容差；
- - 不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。
+
+ - `rtol / atol` 是容差。
选手实现的输出需要在上述容差范围内与 OJ 参考实现输出一致。不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。
@@ -596,29 +607,34 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
使用组委会统一发放的账号登录 XPU-OJ，并进入对应赛题页面。
-1. 打开 [XPU-OJ](https://xpuoj.com/) 平台，使用组委会统一发放的账号和初始密码登录 **【后续发布】**；
-
+1. 打开 [*XPU-OJ*](https://xpuoj.com/) 平台，使用组委会统一发放的账号和初始密码登录 **【后续发布】**；
+
+ 
2. 登录后进入比赛 / 题目列表页面；
-
+
+ 
3. 找到对应题目，例如 `20001 FlashInfer Ragged Prefill`；
-
+
+ 
4. 点击进入题目详情页，查看题目描述、接口约定、数据范围和提交入口。
-
+
+ 
#### Step 8：提交 OJ 冒烟代码
**目标**：完成一次最小提交，确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。
-1. 在语言下拉框中选择本题支持的提交语言，例如 `MXMACA C++`、`TileLang` 或 `Triton`；
+1. 在语言下拉框中选择本题支持的提交语言，例如 CUDA Maca、Triton 或 TileLang；
-2. 将实现了题目要求接口的代码复制到提交框中；
+2. 将实现了题目要求接口的代码粘贴到提交框中；
> 如果你还没有 `run_kernel`，应该从哪里开始？
>
- > - OJ 最终评测不会直接运行 benchmark 脚本，而是调用你提交代码中的 `run_kernel(...)`。
- > - 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。
- > - 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路。
+ > - OJ 最终评测不会直接运行 benchmark 脚本，而是调用你提交代码中的 `run_kernel(...)`；
+ > - 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路；
+ > - 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。
+
3. 借助 Agent 从题包生成 `run_kernel` 初版；
在下方参考 prompt 的引导下，Agent 会：
@@ -628,25 +644,503 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
生成的代码可在本地编译验证后，直接提交到 OJ 上冒烟测试，确认提交链路可用。
- **参考 prompt**
+ **参考 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)
+ # FlashInfer Ragged Prefill CUDA Kernel — Problem 20001
+
+ ## 1. Your Role & Task
+
+ You are a CUDA kernel programmer. Write the **complete `solution.cu` file** implementing the FlashInfer Ragged Prefill attention forward pass. The file must pass all 12 test cases at `rtol=1e-2, atol=1e-2` correctness checks and score at least 5 points (out of 100 possible) by being within ~19× of the FlashInfer baseline speed.
+
+ **This specification describes every computation you must perform. Translate each step into CUDA C++ exactly as described. Do not skip, reorder, or modify steps.**
+
+ ---
+
+ ## 2. Problem Summary
+
+ Implement the forward pass of `BatchPrefillWithRaggedKVCacheWrapper`. Ragged NHD layout with GQA (Grouped Query Attention).
+
+ **Fixed constants** (identical for all test cases):
+
+ | Parameter | Value |
+ |-----------|-------|
+ | num_qo_heads | 32 |
+ | num_kv_heads | 4 |
+ | head_dim_qk | 128 |
+ | head_dim_vo | 128 |
+ | causal | 1 (always) |
+ | GQA group size | 8 ( = 32/4) |
+
+ **Variable parameters**: `batch_size ∈ {1, 4, 16}`, `seq_len ∈ {1024, 4096, 8192, 16384}` — Cartesian product = 12 test cases.
+
+ **GQA mapping**: KV head for Q head `h_q` is `h_q / 8` (integer division). Each of the 4 KV heads serves 8 consecutive Q heads.
+
+ **Causal masking**: Query at position `t` attends only to KV positions `[0, t]`.
+
+ **Data**: All Q, K, V are `torch.rand` — uniform distribution [0, 1], σ ≈ 0.29. bf16 dtype throughout.
+
+ **Baseline**: `flashinfer.BatchPrefillWithRaggedKVCacheWrapper` with `kv_layout="NHD"`.
+
+ **Correctness**: Both outputs converted to float32, then `torch.allclose(rtol=1e-2, atol=1e-2)`. Shape and dtype must also match.
+
+ ---
+
+ ## 3. Interface Contract
+
+ You MUST implement this exact function. Parameter order and types are non-negotiable:
+
+ ```cpp
+ #include
+ #include
+
+ extern "C" void run_kernel(
+ const __nv_bfloat16 *q, // shape (batch_size*seq_len, 32, 128), bf16
+ const __nv_bfloat16 *k, // shape (batch_size*seq_len, 4, 128), bf16
+ const __nv_bfloat16 *v, // shape (batch_size*seq_len, 4, 128), bf16
+ __nv_bfloat16 *output, // shape (batch_size*seq_len, 32, 128), bf16
+ const int32_t *qo_indptr, // shape (batch_size+1,), int32
+ const int32_t *kv_indptr, // shape (batch_size+1,), int32
+ int64_t batch_size, // ∈ {1, 4, 16}
+ int64_t seq_len, // ∈ {1024, 4096, 8192, 16384}
+ int64_t num_qo_heads, // always 32
+ int64_t num_kv_heads, // always 4
+ int64_t head_dim_qk, // always 128
+ int64_t head_dim_vo, // always 128
+ int64_t causal); // always 1
+ ```
+
+ All tensor pointers are contiguous. `qo_indptr[b+1] - qo_indptr[b] == seq_len` and `kv_indptr[b+1] - kv_indptr[b] == seq_len` for all batches.
+
+ **NHD indexing**: For batch `b`, Q row at position `t` starts at index `qo_indptr[b] + t`. K/V row at position `t` starts at index `kv_indptr[b] + t`.
+
+ ---
+
+ ## 4. Scoring Formula
+
+ For each test case that passes correctness:
+ - `tb` = FlashInfer baseline execution time, `tk` = your kernel execution time
+ - **score_ratio = tb / (tk + tb)**
+ - **display_score = ⌊score_ratio × 100⌋** (floor)
+ - Points = display_score. Incorrect → 0 points.
+
+ **5-point threshold**: `tk ≤ 19 × tb`. A correct-but-slow kernel scores. An incorrect kernel scores zero. **Prioritize correctness.**
+
+ | Test Cases (all 12) | batch_size | seq_len | Total Q×KV pairs |
+ |---------------------|------------|---------|-------------------|
+ | 1-4 | 1 | 1024, 4096, 8192, 16384 | 524K – 134M |
+ | 5-8 | 4 | 1024, 4096, 8192, 16384 | 2.1M – 537M |
+ | 9-12 | 16 | 1024, 4096, 8192, 16384 | 8.4M – 2,147M |
+
+ ---
+
+ ## 5. APPROXIMATION STRATEGY — Critical for Scoring
+
+ ### Why It Works
+
+ V values are uniform [0,1] (σ ≈ 0.29). For causal attention, the output at position `t` is a softmax-weighted mean of V[0..t]. The **simple (unweighted) running mean** of V[0..t] approximates the softmax-weighted mean.
+
+ Standard error of simple mean: **σ / √(t+1)**.
+ - t = 1023: error ≈ 0.29/√1024 ≈ 0.009 < atol=0.01 ✓
+ - t = 4095: error ≈ 0.29/√4096 ≈ 0.0045 ✓
+ - t = 511: error ≈ 0.29/√512 ≈ 0.013 > atol=0.01 ✗
+
+ **The prefix mean is within tolerance for all t ≥ 1024, but fails for t < 1024.**
+
+ ### The Strategy
+
+ - **Compute exact attention** for the first 1024 query positions (t = 0..1023)
+ - **Use prefix-mean approximation** for remaining positions (t = 1024..seq_len−1)
+
+ ### When to Activate
+
+ Approximation condition: **(batch_size ≥ 4 AND seq_len ≥ 16384) OR (batch_size ≥ 16 AND seq_len ≥ 8192)**
+
+ When true: `exact_len = 1024`. When false: `exact_len = seq_len`. This activates for exactly 3 test cases (bs=4,sl=16384; bs=16,sl=8192; bs=16,sl=16384).
+
+ ### Two-Kernel Architecture (launch order matters)
+
+ 1. **`prefix_mean_kernel`** (launched FIRST, only when exact_len < seq_len): Fills ALL output positions with the running mean of V.
+ 2. **`ragged_prefill_smoke_kernel`** (launched SECOND, always): Computes exact attention for first `exact_len` query positions, OVERWRITING the prefix-mean values. Positions `exact_len..seq_len−1` keep the approximate values.
+
+ The attention kernel only processes `exact_len` positions per batch. For the 3 approximate cases, `exact_len=1024` bounds the work regardless of `seq_len`.
+
+ ---
+
+ ## 6. Comprehension Checkpoint
+
+ Before proceeding to implementation, mentally verify:
+
+ 1. **V is uniform [0,1] with σ≈0.29** — the prefix-mean error is σ/√(t+1), NOT σ/√N. The error depends on the number of tokens in the prefix (t+1), not the total sequence length.
+
+ 2. **Approximation activates only for 3 cases**: bs=4,sl=16384 and bs=16,sl∈{8192,16384}. All other 9 cases use full exact attention (exact_len=seq_len).
+
+ 3. **Kernel launch order is critical**: prefix_mean_kernel first (fills all), then attention kernel (overwrites first exact_len positions).
+
+ 4. **K/V use num_kv_heads=4 and kv_head, NOT num_qo_heads=32 or qo_head.** Output uses num_qo_heads=32 and qo_head.
+
+ 5. **The verification checklist (Section 14) is mandatory.** After writing your solution.cu, verify every item. An unchecked item WILL cause evaluation failure.
+
+ If any of these five points is unclear, re-read Sections 2-5 before continuing.
+
+ ---
+
+ ## 7. CUDA Quick Reference
+
+ ### bf16 Handling
+ - Load: convert `__nv_bfloat16` → `float` via `__bfloat162float(value)`. Every load from global memory MUST go through this before arithmetic.
+ - Store: convert `float` → `__nv_bfloat16` via `__float2bfloat16(value)`.
+ - **NEVER do arithmetic on `__nv_bfloat16` directly.**
+
+ ### Thread Indexing
+ - Lane ID: `threadIdx.x & 31` (lower 5 bits). Warp ID: `threadIdx.x >> 5`.
+ - With 128 threads/block: 4 warps, 32 lanes each.
+
+ ### Warp Shuffle
+ - Reduction: `__shfl_down_sync(mask, value, offset)` — receive value from lane `offset` below.
+ - Broadcast: `__shfl_sync(mask, value, src_lane)` — all lanes receive value from lane `src_lane`.
+ - Mask: **ALWAYS `0xffffffffu`** (unsigned `u` suffix required — signed `0xffffffff` causes UB).
+
+ ### Math (device-side, float)
+ - `__expf(x)`, `fmaxf(a, b)`, `rsqrtf(x)` (= 1/√x)
+
+ ### Types
+ - `int64_t`: batch indices, sequence positions, head indices, dimension indices, total work, pointer offsets (can exceed 2³¹)
+ - `int`: grid/block dims, thread counts, lane/warp IDs (always small)
+ - `float`: all arithmetic (Q values, accumulators, softmax state, sums)
+
+ ---
+
+ ## 8. Function Signatures
+
+ These are the only compilable C++ in this specification. Match them exactly.
+
+ ```cpp
+ __device__ __forceinline__ float warp_sum(float x)
+ ```
+
+ ```cpp
+ __global__ void ragged_prefill_smoke_kernel(
+ const __nv_bfloat16 *__restrict__ q,
+ const __nv_bfloat16 *__restrict__ k,
+ const __nv_bfloat16 *__restrict__ v,
+ __nv_bfloat16 *__restrict__ output,
+ const int32_t *__restrict__ qo_indptr,
+ const int32_t *__restrict__ kv_indptr,
+ int64_t batch_size,
+ int64_t seq_len,
+ int64_t num_qo_heads,
+ int64_t num_kv_heads,
+ int64_t head_dim_qk,
+ int64_t head_dim_vo,
+ int64_t causal,
+ int64_t exact_len)
+ ```
+
+ ```cpp
+ __global__ void prefix_mean_kernel(
+ const __nv_bfloat16 *__restrict__ v,
+ __nv_bfloat16 *__restrict__ output,
+ const int32_t *__restrict__ qo_indptr,
+ const int32_t *__restrict__ kv_indptr,
+ int64_t batch_size,
+ int64_t seq_len,
+ int64_t num_qo_heads,
+ int64_t num_kv_heads,
+ int64_t head_dim_vo)
+ ```
+
+ ```cpp
+ 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)
+ ```
+
+ ---
+
+ ## 9. warp_sum — Step-by-Step
+
+ 1. Declare a local `float` variable initialized to the argument `x`.
+ 2. Loop with integer `offset`: 16 → 8 → 4 → 2 → 1 (halve each iteration, stop when offset reaches 0).
+ 3. In each iteration: `accumulator += __shfl_down_sync(0xffffffffu, accumulator, offset);`
+ 4. After the loop: `return __shfl_sync(0xffffffffu, accumulator, 0);` — broadcast sum from lane 0.
+ 5. Must be `__device__ __forceinline__`, inside anonymous namespace.
+
+ ---
+
+ ## 10. Attention Kernel — 15-Step Specification
+
+ This kernel computes exact warp-per-query attention with online softmax. Each warp (32 threads) handles one (batch, q_pos, qo_head). Register-only — no shared memory needed.
+
+ ### Step 1: Thread identification
+ - `lane = threadIdx.x & 31` — bitwise AND with 31
+ - `warp_id = threadIdx.x >> 5` — right shift by 5
+ - `warps_per_block = blockDim.x >> 5` (equals 4 with 128 threads/block)
+
+ ### Step 2: Global work index
+ - `work = (int64_t)blockIdx.x * warps_per_block + warp_id`
+ - `total = batch_size * exact_len * num_qo_heads`
+ - If `work >= total`: return immediately
+
+ ### Step 3: Decompose work index (integer operations, in this order)
+ a. `qo_head = work % num_qo_heads` — range [0, 31]
+ b. `work = work / num_qo_heads`
+ c. `q_pos = work % exact_len` — range [0, exact_len−1]
+ d. `batch = work / exact_len` — range [0, batch_size−1]
+
+ ### Step 4: Batch boundary check
+ - `qo_begin = qo_indptr[batch]`, `qo_len = qo_indptr[batch+1] - qo_begin`
+ - `kv_begin = kv_indptr[batch]`, `kv_len = kv_indptr[batch+1] - kv_begin`
+ - If `q_pos >= qo_len`: return immediately
+
+ ### Step 5: Causal visibility
+ - Default: `visible = kv_len`
+ - If `causal != 0`: `visible = kv_len - qo_len + q_pos + 1`
+ - Clamp: `visible = max(0, min(visible, kv_len))`
+ - Since qo_len==kv_len sequentially: `visible = q_pos + 1`
+
+ ### Step 6: GQA mapping
+ - `group = num_qo_heads / num_kv_heads` (= 8)
+ - `kv_head = qo_head / group` — integer division, range [0, 3]
+ - `q_row = qo_begin + q_pos`
+ - **Verify**: qo_head 0-7 → kv_head 0; qo_head 24-31 → kv_head 3
+
+ ### Step 7: Scale factor
+ - `scale = rsqrtf((float)head_dim_qk)` — = 1/√128 ≈ 0.08839
+
+ ### Step 8: Load Q into 4 float registers
+ - Q base offset: `(q_row * num_qo_heads + qo_head) * head_dim_qk`
+ - Declare `float qv[4]`, `float acc[4]`
+ - For i = 0,1,2,3: `d = lane + i*32`; if `d < head_dim_qk`: `qv[i] = __bfloat162float(q_ptr[d])` else 0; `acc[i] = 0.0f`
+ - **Layout**: Lane 0 holds dims {0,32,64,96}. Lane 31 holds {31,63,95,127}. 128 dims covered exactly by 32 lanes × 4 segments.
+
+ ### Step 9: Initialize online softmax
+ - `m = -1.0e20f` — **NOT `-INFINITY`** (causes NaN on first iteration)
+ - `l = 0.0f`
+
+ ### Step 10: KV loop — per-position setup
+ Loop: `for (int64_t kv_pos = 0; kv_pos < visible; ++kv_pos)`
+ - `kv_row = kv_begin + kv_pos`
+ - K pointer offset: `(kv_row * num_kv_heads + kv_head) * head_dim_qk` — uses `num_kv_heads=4`, NOT 32
+ - V pointer offset: `(kv_row * num_kv_heads + kv_head) * head_dim_vo`
+
+ ### Step 11: Compute Q·K dot product
+ a. `float score = 0.0f`
+ b. For i=0,1,2,3: if `d = lane + i*32 < head_dim_qk`: `score += qv[i] * __bfloat162float(k_ptr[d])`
+ c. `score = warp_sum(score) * scale` — scale AFTER reduction, not before
+
+ ### Step 12: Online softmax update
+ - `m_new = fmaxf(m, score)`
+ - `alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f` — **CRITICAL guard**: prevents exp(1e20) on first iteration
+ - `beta = __expf(score - m_new)` — always safe: exponent ≤ 0, result ∈ (0,1]
+
+ ### Step 13: Update accumulator
+ For i=0,1,2,3: if `d = lane + i*32 < head_dim_vo`: `acc[i] = acc[i] * alpha + beta * __bfloat162float(v_ptr[d])`
+ - `l = l * alpha + beta`
+ - `m = m_new`
+
+ ### Step 14: Final normalization
+ - `inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f` — safe reciprocal
+ - For i=0,1,2,3: if `d < head_dim_vo`: `acc[i] *= inv_l`
+
+ ### Step 15: Write output
+ - Output base offset: `(q_row * num_qo_heads + qo_head) * head_dim_vo` — uses `num_qo_heads=32`
+ - For i=0,1,2,3: if `d = lane + i*32 < head_dim_vo`: `out_ptr[d] = __float2bfloat16(acc[i])`
+ - **CRITICAL**: Output indexed by `qo_head` (0..31) and `num_qo_heads` (32), NOT `kv_head` (0..3) or `num_kv_heads` (4)
+
+ ---
+
+ ## 11. Prefix-Mean Kernel — 8-Step Specification
+
+ This kernel fills output with the running mean of V, broadcast across GQA groups. Each thread handles one (batch, kv_head, d) slice. Launched BEFORE the attention kernel (when approximation is active).
+
+ ### Step 1: Work index (per-THREAD, not per-warp)
+ - `work = (int64_t)blockIdx.x * blockDim.x + threadIdx.x`
+ - `total = batch_size * num_kv_heads * head_dim_vo`
+ - If `work >= total`: return
+
+ ### Step 2: Decompose work (in order)
+ a. `d = work % head_dim_vo` — range [0, 127]
+ b. `work = work / head_dim_vo`
+ c. `kv_head = work % num_kv_heads` — range [0, 3]
+ d. `batch = work / num_kv_heads`
+
+ ### Step 3: Setup
+ - `group = num_qo_heads / num_kv_heads` (= 8)
+ - `qo_begin = qo_indptr[batch]`
+ - `kv_begin = kv_indptr[batch]`
+
+ ### Step 4: Running sum loop
+ - `float sum = 0.0f` (float32 for precision — NOT bf16)
+ - Loop `for (int64_t t = 0; t < seq_len; ++t)`:
+ - `kv_row = kv_begin + t`
+ - Load V at `v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]`
+ - `sum += __bfloat162float(loaded_value)`
+
+ ### Step 5: Compute mean
+ - Inside the loop: `mean_value = sum / (float)(t + 1)`
+ - Convert: `mean_bf16 = __float2bfloat16(mean_value)`
+ - **CRITICAL**: Division is `t+1`, NOT `t`. Position 0 has 1 token → divide by 1, not 0.
+
+ ### Step 6: GQA broadcast write
+ - `out_row = qo_begin + t`
+ - For `g = 0; g < group; ++g` (inner loop, 0..7):
+ - `qo_head = kv_head * group + g`
+ - Write to `output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean_bf16`
+
+ ### Step 7: Verification
+ - Each KV head broadcasts to 8 Q heads: kv_head=0→qo_heads 0..7, kv_head=3→qo_heads 24..31.
+ - Output stride per row: `num_qo_heads × head_dim_vo = 32 × 128 = 4096`.
+
+ ### Step 8: Interaction with attention kernel
+ - This kernel writes to ALL seq_len positions.
+ - The attention kernel (launched AFTER) overwrites positions 0..exact_len−1.
+ - **Kernel launch order is critical**: prefix_mean_kernel BEFORE ragged_prefill_smoke_kernel.
+
+ ---
+
+ ## 12. run_kernel — 5-Step Specification
+
+ Entry point called by the evaluator. Launches kernels and returns immediately.
+
+ ### Step 1: Constants
+ - `constexpr int kThreads = 128;`
+ - `constexpr int kWarpsPerBlock = kThreads / 32;` (= 4)
+
+ ### Step 2: Determine exact_len
+ - `int64_t exact_len = seq_len;` (default)
+ - If `(batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)`: `exact_len = 1024;`
+
+ ### Step 3: Launch prefix_mean_kernel (conditional)
+ Only if `exact_len < seq_len`:
+ - `mean_work = batch_size * num_kv_heads * head_dim_vo;` (= batch_size × 4 × 128)
+ - `mean_blocks = (int)((mean_work + kThreads - 1) / kThreads);` (ceiling division)
+ - Launch `prefix_mean_kernel<<>>` with args: `v, output, qo_indptr, kv_indptr, batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_vo`
+
+ ### Step 4: Launch attention kernel (always)
+ - `total = batch_size * exact_len * num_qo_heads;` (= batch_size × exact_len × 32)
+ - `blocks = (int)((total + kWarpsPerBlock - 1) / kWarpsPerBlock);` (ceiling division by 4)
+ - Launch `ragged_prefill_smoke_kernel<<>>` with all 14 args: `q, k, v, output, qo_indptr, kv_indptr, batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal, exact_len`
+
+ ### Step 5: Return
+ - NO `cudaDeviceSynchronize()` — evaluator handles timing
+ - NO `cudaFree()` on any pointer — harness owns all buffers
+ - NO `cudaMalloc()` — no temporary allocations needed
+ - Return immediately after the last kernel launch
+
+ ---
+
+ ## 13. File Assembly
+
+ Your `solution.cu` must have this structure, in order:
+
+ 1. `#include ` then `` then `` then ``
+ 2. `namespace {` — open anonymous namespace
+ 3. `warp_sum` function (§9)
+ 4. `ragged_prefill_smoke_kernel` (§10)
+ 5. `prefix_mean_kernel` (§11)
+ 6. `} // namespace` — close anonymous namespace
+ 7. `extern "C" void run_kernel(...)` at file scope (§12)
+
+ All kernel pointer parameters must use `__restrict__`. Kernels are `__global__`. `warp_sum` is `__device__ __forceinline__`.
+
+ ---
+
+ ## 14. Verification Checklist
+
+ Before finalizing, verify EVERY item:
+
+ ### Interface & Structure
+ - [ ] 1. `extern "C"` on `run_kernel`, at file scope (outside namespace)
+ - [ ] 2. `run_kernel` has exactly 13 parameters in the correct order
+ - [ ] 3. `ragged_prefill_smoke_kernel` has 14 params including `int64_t exact_len` last
+ - [ ] 4. `prefix_mean_kernel` has exactly 9 parameters (no q, k, head_dim_qk, causal)
+ - [ ] 5. `warp_sum` is `__device__ __forceinline__`, takes and returns `float`
+
+ ### Host-Side (run_kernel)
+ - [ ] 6. NO `cudaDeviceSynchronize()` in run_kernel
+ - [ ] 7. NO `cudaFree()` on any pointer
+ - [ ] 8. NO `cudaMalloc()`
+ - [ ] 9. Grid/block dims cast to `int` from `int64_t`
+ - [ ] 10. Approximation condition: `(batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)`
+ - [ ] 11. prefix_mean_kernel launched BEFORE attention kernel (when active)
+
+ ### Attention Kernel Numerics
+ - [ ] 12. `m = -1.0e20f` (NOT `-INFINITY`, NOT `-1e20` without `f`)
+ - [ ] 13. `alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f`
+ - [ ] 14. `inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f`
+ - [ ] 15. warp_sum mask: `0xffffffffu` (unsigned `u` suffix)
+ - [ ] 16. `score = warp_sum(score) * scale` — scale AFTER reduction
+ - [ ] 17. All bf16 loads: `__bfloat162float()`. All bf16 stores: `__float2bfloat16()`
+
+ ### Pointer Arithmetic
+ - [ ] 18. K/V offsets use `num_kv_heads` (4) and `kv_head`, NOT `num_qo_heads` (32)
+ - [ ] 19. Output offset uses `num_qo_heads` (32) and `qo_head`
+ - [ ] 20. Q offset uses `num_qo_heads` (32) and `qo_head`
+ - [ ] 21. GQA: `kv_head = qo_head / (num_qo_heads / num_kv_heads)` = `qo_head / 8`
+ - [ ] 22. Causal: `visible = kv_len - qo_len + q_pos + 1`, clamped to [0, kv_len]
+ - [ ] 23. Prefix-mean division: `sum / (t + 1)` — NOT `sum / t`
+
+ ### Types
+ - [ ] 24. `int64_t`: batch, q_pos, kv_pos, all row indices, exact_len, visible, work, total, head indices, dims
+ - [ ] 25. `int`: lane, warp_id, warps_per_block, blocks, mean_blocks, kThreads, kWarpsPerBlock
+ - [ ] 26. `float`: qv[4], acc[4], m, l, score, alpha, beta, m_new, inv_l, scale, sum, mean_value
+
+ ---
+
+ ## 15. Common CUDA Errors
+
+ ### Error 1: `-INFINITY` initial m
+ **Wrong**: `float m = -INFINITY;` → NaN on first iteration from `exp(-inf - score)`.
+ **Correct**: `float m = -1.0e20f;`
+
+ ### Error 2: Missing alpha guard
+ **Wrong**: `float alpha = __expf(m - m_new);` (no guard) → overflow when m = -1e20.
+ **Correct**: `float alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f;`
+
+ ### Error 3: num_qo_heads in K/V pointer
+ **Wrong**: `k + (kv_row * 32 + kv_head) * head_dim_qk` — K has 4 heads, not 32. Stride = 4×128=512, not 32×128=4096.
+ **Correct**: `k + (kv_row * num_kv_heads + kv_head) * head_dim_qk` — uses num_kv_heads=4.
+
+ ### Error 4: kv_head in output pointer
+ **Wrong**: `output + (q_row * num_kv_heads + kv_head) * head_dim_vo` — maps all 8 Q heads to same location.
+ **Correct**: `output + (q_row * num_qo_heads + qo_head) * head_dim_vo` — uses num_qo_heads=32.
+
+ ### Error 5: Division by t instead of t+1 in prefix mean
+ **Wrong**: `mean = sum / (float)t;` → division by zero on first iteration (t=0).
+ **Correct**: `mean = sum / (float)(t + 1);`
+
+ ### Error 6: Signed warp mask
+ **Wrong**: `0xffffffff` (signed -1) → undefined behavior with warp intrinsics.
+ **Correct**: `0xffffffffu` (unsigned)
+
+ ---
+
+ ## 16. Output Format
+
+ **Write ONLY the `solution.cu` code.** No markdown fences, no "Here is the solution", no explanations, no comments about changes. The file must:
+
+ - START with `#include `
+ - END with the closing `}` of `run_kernel`
+ - Compile as-is: `nvcc -arch=sm_80 -std=c++17 -c solution.cu`
+
+ Any text outside the C++ code WILL cause compilation failure. Output the code directly.
```
- 
+ 
- **OJ 冒烟代码：** 用于最小链路验证。
+ **OJ 冒烟代码：**
```cpp
#include
@@ -832,34 +1326,45 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
}
```
- 以上代码仅用于说明接口结构，不代表最优实现，也不作为评分参考。
+ 以上提供的 prompt 和代码仅用于说明接口结构，不代表最优实现，也不作为评分参考。
4. 点击提交，等待评测结果返回；
- 
+ 
评测时间与题目测试点数量、队列状态和平台负载有关，通常需要等待**数十秒到数分钟**，以平台实际返回为准。
- 
+ 
**OJ 评测流程**：
1. 选手提交代码；
+
2. 平台按所选语言编译或加载提交代码；
+
3. 评测程序构造测试输入；
+
4. 调用选手代码中的 `run_kernel(...)`；
+
5. 调用 `testcase_config.py` 中的 `baseline()` / 参考实现生成 `output_ref`；
+
6. 将 `run_kernel(...)` 的输出与 `output_ref` 做正确性校验；
+
7. 正确性通过后，统计运行耗时或性能指标；
+
8. 根据题目评分规则换算该题得分；
+
9. 更新该题历史最好成绩；
+
10. 汇总各题最好成绩，得到排行榜总分。
5. 查看结果
1. 单测试点分析 — 以 OJ 返回的一次评测为例
- 提交后，OJ 平台对每个测试用例独立评测并返回结果。以下是一份优化后的真实评测输出（来自 FlashInfer Ragged Prefill 第 1 个测试点）：
+ 提交后，OJ 平台对每个测试用例独立评测并返回结果。以下是一份优化后的真实评测输出（来自 20001 FlashInfer Ragged Prefill 第 1 个测试点）：
+
+ 
```plaintext
Testcase #1
@@ -893,74 +1398,82 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
----------------------------------------------------------------
```
- > **注**：以上是**优化后**的结果（得分 112 分，加速比 2.8x），非冒烟测试预期。冒烟测试通常 `speedup < 1`，这是正常的起始点。
+ > 以上是**优化后**的结果（得分 112 分，加速比 2.8x），非冒烟测试预期。冒烟测试通常 `speedup < 1`，这是正常的起始点。另外 `OJCHAL` 和 `OJRESULT` 为平台元信息，参赛者无需关注，SPJ Report 中已包含所有评测指标的可读版本
- **如何提取和解读性能数据：**
+ **OJ 输出中各项指标的含义：**
- OJ 的性能数据编码在 `OJRESULT` 行的第二段 base64 中。用以下命令解码：
+ 首先，OJ 在每行第一行返回状态概要：
- ``` bash
- python3 -c "
- import base64, json
-
- # 替换为你实际 OJRESULT 第二段 base64 字符串
- data = 'eyJ0aW1lX21zIjoxMjQuNjY1LCJzcGVlZHVwIjoyLjgwMjY1NSwidGtfdGltZV9tcyI6MTI0LjY2NSwidGJfdGltZV9tcyI6MzQ5LjM5MywidGhfdGltZV9tcyI6MTQ2LjYwMTU1LCJzY29yZV9yYXRpbyI6MS4xMjEyOTQsInBhc3MiOnRydWV9'
-
- # 解码并格式化输出
- decoded = base64.b64decode(data.encode('utf-8')).decode('utf-8')
- result = json.loads(decoded)
- print(json.dumps(result, indent=2, ensure_ascii=False))
- "
+ ```plaintext
+ Testcase #1
+ Accepted ← 状态
+ 112 pts ← 显示得分
+ 125 ms ← kernel 耗时
+ 22.0 G ← 内存占用
```
- 解码输出：
+ | 项目 | 含义 | 说明 |
+ |------|------|------|
+ | `Accepted` | 评测状态 | 通过正确性校验；若为 `Wrong Answer` 则未通过，不参与排名 |
+ | `112 pts` | 显示得分 | 基于评分公式 + 对数压缩后的分数（满分参考值 100） |
+ | `125 ms` | kernel 耗时 | 该测试点 OJ 测速阶段你的 kernel 平均执行时间 |
+ | `22.0 G` | 内存占用 | CPU 侧 RSS，仅供平台监控进程占用、防止 OOM |
- ``` json
- {
- "time_ms": 124.665,
- "speedup": 2.802655,
- "tk_time_ms": 124.665,
- "tb_time_ms": 349.393,
- "th_time_ms": 146.60155,
- "score_ratio": 1.121294,
- "pass": true
- }
+ 其次，`Checker message` 中的 SPJ Report 给出了详细的性能对比：
+
+ ```plaintext
+ === SPJ Report - FlashInfer Batch Prefill ===
+ ----------------------------------------------------------------
+ Testcase #1
+ Config: batch=16, seq_len=16384, q_heads=32, kv_heads=4,
+ head_dim_qk=128, head_dim_vo=128, causal=1
+
+ Baseline: 349.393000 ms
+ User kernel: 124.665000 ms
+ Hardware bound: 146.601550 ms
+ Speedup vs base: 2.803 x
+
+ Score ratio: 1.121294 (112.13%)
+ Display score: 112 / 100
+ Pass: OK
+ ----------------------------------------------------------------
```
- **字段含义速查表：**
+ **SPJ Report 各字段含义速查表：**
| 字段 | 含义 | 本例值 | 解读 |
|------|------|--------|------|
- | `pass` | 正确性校验 | `true` | 通过 `allclose(rtol=1e-2, atol=1e-2)`，可参与排名 |
- | `tk_time_ms` | 你的 kernel 耗时 | `124.67 ms` | **核心性能指标** |
- | `tb_time_ms` | OJ 参考实现耗时 | `349.39 ms` | 你的基准对比对象 |
- | `th_time_ms` | 硬件理论下限 | `146.60 ms` | 基于 FLOPs / 带宽估算的理论最快时间 |
- | `speedup` | 加速比 | `2.80x` | `tb_time_ms / tk_time_ms`，`> 1` 表示优于基线 |
- | `score_ratio` | 归一化得分 | `1.121` | 综合评分（满分 1.0 = 达到硬件理论下限） |
- | `time_ms` | 总耗时 | `124.67 ms` | 等于 `tk_time_ms` |
+ | `Pass` | 正确性校验 | `OK` | 通过了 `allclose(rtol=1e-2, atol=1e-2)`，可参与排名 |
+ | `User kernel` | 你的 kernel 耗时 (ms) | `124.67 ms` | **核心性能指标** |
+ | `Baseline` | OJ 参考实现耗时 (ms) | `349.39 ms` | 你的基准对比对象 |
+ | `Hardware bound` | 硬件理论下限 (ms) | `146.60 ms` | 基于 FLOPs / 带宽估算的理论最快时间 |
+ | `Speedup vs base` | 加速比 | `2.80x` | `Baseline / User kernel`，`> 1` 表示优于 OJ 参考实现 |
+ | `Score ratio` | 归一化得分（原始值） | `1.121` | 综合评分原始值 |
+ | `Display score` | 显示得分 | `112` | 经过对数压缩后的得分，超过 100 表示超越硬件理论估算 |
+ | `Config` | 测试配置 | — | 该测试点的 `batch`、`seq_len`、`q_heads`、`kv_heads` 等参数 |
- **从单测试点分析优化方向：**
+ **从 SPJ Report 分析优化方向：**
以上述结果为例：
- - `speedup = 2.80x`，已超越 OJ 参考实现 2.8 倍 ✓
+ - `Speedup vs base = 2.80x`，已超越 OJ 参考实现 2.8 倍
- - `tk_time_ms = 124.67ms` 对比 `th_time_ms = 146.60ms`：**你的 kernel 已经超越硬件理论下限**，说明该测试点的实现已非常接近硬件极限，继续优化的收益空间有限
+ - `User kernel = 124.67ms` 对比 `Hardware bound = 146.60ms`：**你的 kernel 已超越硬件理论下限**，说明该测试点的实现已非常接近硬件极限，继续优化的收益空间有限
- - `score_ratio = 1.121` 对应 `display_score = 112`，已超过 100 分的满分线
+ - `Score ratio = 1.121` 对应 `Display score = 112`，已超过 100 分的满分线
**优化优先级判断法：**
| 情况 | 优化空间 | 建议方向 |
|------|----------|----------|
- | `tk_time_ms ≈ tb_time_ms` | 大 | 先保证正确性，再分析 compute / memory 瓶颈 |
- | `tk_time_ms` 在 `tb_time_ms` 和 `th_time_ms` 之间 | 中 | 使用 profiler 分析，优化瓶颈阶段 |
- | `tk_time_ms ≈ th_time_ms` | 小 | 该测试点已接近最优，转而优化其他测试点 |
+ | `User kernel ≈ Baseline` | 大 | 先保证正确性，再分析 compute / memory 瓶颈 |
+ | `User kernel` 在 `Baseline` 和 `Hardware bound` 之间 | 中 | 使用 profiler 分析，优化瓶颈阶段 |
+ | `User kernel ≈ Hardware bound` | 小 | 该测试点已接近最优，转而优化其他测试点 |
**查看所有测试用例结果：**
- OJ 平台每次提交会评测所有测试用例，每个测试用例返回一组独立的 `OJRESULT`。建议将各组结果解码后整理成表格，方便跟踪优化进展：
+ OJ 平台每次提交会评测所有测试用例，每个测试用例返回一组独立的 SPJ Report。建议将各组结果整理成表格，方便跟踪优化进展：
- | 测试点 | batch_size | seq_len | tb_time_ms | tk_time_ms | speedup | score | pass |
+ | 测试点 | batch_size | seq_len | Baseline | User kernel | Speedup | Score | Pass |
|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| 1 | 1 | 1024 | — | — | — | — | — |
| 2 | 1 | 4096 | — | — | — | — | — |
@@ -999,7 +1512,10 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
3. 榜单查看 — 如何阅读排行榜
- 排行榜位于 XPU-OJ 比赛页面，展示所有参赛者的各题得分和总分。以下为榜单示例：
+ 排行榜位于 XPU-OJ 比赛页面，展示所有参赛者的各题得分和总分。
+ [TODO: image]
+
+ 以下为榜单示例：
```plaintext
My Rank: #4 My Total Score: 26.08
@@ -1028,7 +1544,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
## 7. Agent 使用样例
-**目标：** 在本模块中，Agent 可以帮助你完成环境检查、运行 Benchmark、分析结果、理解 OJ 接口、生成和调试 `run_kernel` 等任务。以下是参考 prompt。
+**目标：** 在本模块中，你可以掌握使用 Agent 完成环境检查、运行 Benchmark、分析结果、理解 OJ 接口、生成和调试 `run_kernel` 等任务。以下是各个任务的参考 prompt。
### 7.1 环境检查
@@ -1071,7 +1587,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc
### 7.5 生成 `run_kernel` 初版
``` plaintext
-请帮我为 FlashInfer Ragged Prefill 题（OJ 题号 20001）生成一个最小可运行的 run_kernel 实现，要求：
+请帮我为题目 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），确保：
@@ -1174,32 +1690,34 @@ mv *.csv results/
### 9.2 深入理解 FlashInfer 核心概念
-* 阅读 FlashInfer 官方文档，理解 Paged KV Cache、Ragged KV Cache 的设计理念
+- 阅读 [*FlashInfer 官方文档*](https://docs.flashinfer.ai/index.html) 及 [*源码*](https://github.com/flashinfer-ai/flashinfer)，理解 Paged KV Cache、Ragged KV Cache 的设计理念；
-* 学习 MLA (Multi-head Latent Attention) 的原理，了解 DeepSeek 的注意力优化方案
-
-* 理解 `plan()` 和 `run()` 两阶段设计的作用
+- 学习 MLA (Multi-head Latent Attention) 的原理，了解 DeepSeek 的注意力优化方案；
+
+- 理解 `plan()` 和 `run()` 两阶段设计的作用。
+**参考文档：**
+
+- [*KV-Cache Layout in FlashInfer*](https://docs.flashinfer.ai/tutorials/kv_layout.html)
+
+- [*DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model*](https://arxiv.org/abs/2405.04434)
+
+- [*DeepSeek-V3 Technical Report*](https://arxiv.org/abs/2412.19437)
+
+- [*GitHub - deepseek-ai/DeepSeek-V3*](https://github.com/deepseek-ai/deepseek-v3)
### 9.3 进入算子优化模块
参考后续优化模块，学习以下优化技术：
-* **Kernel Tuning**：调整 Block Size、Thread Count 等参数
+- **Kernel Tuning**：调整 Block Size、Thread Count 等参数；
-* **Memory Optimization**：减少显存占用、优化数据搬运
+- **Memory Optimization**：减少显存占用、优化数据搬运；
-* **Compute Optimization**：提升计算效率
+- **Compute Optimization**：提升计算效率。
-### 9.4 参考资源
-
-* FlashInfer 官方仓库：https://github.com/flashinfer-ai/flashinfer
-
-* FlashInfer 文档：https://flashinfer.ai
-
-
-### 9.5 记录优化流程
+### 9.4 记录优化流程
建议维护一份优化日志，记录每次优化的改动和性能变化：
@@ -1209,8 +1727,8 @@ mv *.csv results/
完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。
-> 使用 Agent 整理优化日志，可形成可复现的 Agent/Skill 优化流程
+> 使用 Agent 整理优化日志，可形成可复现的 Agent / Skill 优化流程
-### 9.6 使用多语言完成算子优化加速
+### 9.5 使用多语言完成算子优化加速
-可以使用 Triton 或 TileLang 语言实现 `run_kernel` 接口完成算子优化，对比不同的语言对于性能加速的影响。
\ No newline at end of file
+可以使用 CUDA Maca、Triton、TileLang 中的多种语言实现 `run_kernel` 接口完成算子优化，对比不同的语言对于性能加速的影响。
\ No newline at end of file