From 398ec687c33bf7d67312b44a578169a0438ff8e2 Mon Sep 17 00:00:00 2001 From: "Ke Xiao (i26293)" Date: Thu, 18 Jun 2026 10:06:54 +0800 Subject: [PATCH 1/2] =?UTF-8?q?=E6=9B=B4=E6=96=B0FlashInfer=E6=95=99?= =?UTF-8?q?=E7=A8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../FlashInfer关键算子迁移与优化.md | 852 +++++++++++------- 1 file changed, 521 insertions(+), 331 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md index 36de497..826f9b4 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md @@ -1,30 +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。 +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软化栈的使用 @@ -36,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/)，进入费用中心 - 算力券 ，点击右上角 “兑换”。 + + - 进入算力容器，选择沐曦，租用算力，建议优先选 16G 显存 / 32G 显存，如下图 + + ![Weixin-Image_20260617180029](https://origin.picgo.net/2026/06/17/Weixin-Image_20260617180029_15_100c1e1fecb298bac54.png) - * 获得兑换码 - + - 创建实例。基础镜像：`maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64` -* 步骤2：兑换算力和登陆平台 + - 选择工具-lab进入实例环境 - * 访问模力方舟官网：[https://ai.gitee.com/](https://ai.gitee.com/) - - * 进入费用中心 - 算力券 ，点击右上角 “兑换”。 - - -* 步骤3：租用算力 - - - 进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图： -* 步骤4：创建实例 - - - 基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64 -* 步骤5：选择工具-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 脚本: | 命令 | 说明 | | --- | --- | @@ -105,44 +85,81 @@ 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 -## 本地项目运行 +**目标：** 在赛事镜像中完成 FlashInfer Ragged Prefill 算子的 Baseline Benchmark，理解从 Benchmark 到 XPU-OJ 评测提交的完整流程，为后续算子优化建立性能基线。 + +### 在赛事镜像中运行 FlashInfer Baseline Benchmark #### Step 1：检查运行环境 -**目标：** 进入terminal确认当前环境满足本模块运行要求。 +**目标：** 确认当前环境满足本模块运行要求。 -**操作：** 检查 GPU、Python、编译工具和依赖版本。 +**操作：** 进入 Terminal 检查 GPU、Python、编译工具和依赖版本。 ![giteeai 实例 12](https://origin.picgo.net/2026/06/04/giteeai--12c1772b12867f6be0.png) +**命令示例：** + +```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 信息 + + ![result mx smi](https://origin.picgo.net/2026/06/04/result-mx-smif86a3bed6681382e.png) + +- Python 环境正常 + + ![result python](https://origin.picgo.net/2026/06/04/result-pythond15f856ddb84c649.png) + +- `torch.cuda.is_available()` 返回 `True` + + ![result gpu available](https://origin.picgo.net/2026/06/04/result-gpu-available44e049addf5638fb.png) + +- 所有依赖版本符合要求 + + ![result dependency version](https://origin.picgo.net/2026/06/04/result-dependency-versionc17845898ee89434.png) + + **常见问题：** | 问题 | 解决方法 | | --- | --- | -| `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：进入项目目录 @@ -150,22 +167,23 @@ prefill = 并行处理用户输入，decode = 逐个生成回答 token 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 baselines\flashinfer_baseline\FlashInfer_Baseline -ls -la -``` +3. 切换到 FlashInfer_Baseline 项目目录 + ```bash + cd data/flashinfer_baseline/FlashInfer_Baseline + ls -al + ``` **预期结果：** @@ -177,7 +195,6 @@ bench_batch_prefill_ragged.py bench_batch_mla.py README.md ... - ``` #### Step 3：验证项目脚本 @@ -194,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('脚本导入正常')" - ``` **预期结果：** @@ -206,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 ``` **常见问题：** @@ -241,265 +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 评测与提交** +### XPU-OJ 在线评测教程 -#### Step 1：访问XPUOJ平台(https://xpuoj.com/)进行登录 +#### Step 5: 从 Baseline 到 XPU-OJ 提交 -- 完成 baseline benchmark 后，要到统一评测工具XPU-OJ完成提交，进行评测。 +##### 5.1：Baseline 与 XPU-OJ 的关系 -- 等待组委会统一发放 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. 使用组委会统一发放的账号和初始密码登录 **【后续发布】**； ![image 20260616152804368](https://origin.picgo.net/2026/06/16/image-20260616152804368aa0d44b72c4f9572.png) -#### Step 2：进入比赛页面 - -- 登录后进入对应赛题页面 -- 选择对应的赛题 - +3. 登录后进入比赛 / 题目列表页面； ![image 20260616152953345](https://origin.picgo.net/2026/06/16/image-20260616152953345986ce39fda69da55.png) - -#### Step 3：选择对应的算子题目进行评测 - +4. 找到对应题目，例如 1 FlashInfer Ragged Prefill； ![image 20260616153142383](https://origin.picgo.net/2026/06/16/image-202606161531423833b0f0428edf2e35e.png) +5. 点击进入题目详情页，查看题目描述、接口约定、数据范围和提交入口。 +![Screenshot_17-6-2026_165526_xpuoj](https://origin.picgo.net/2026/06/17/Screenshot_17-6-2026_165526_xpuoj.com9a82c272c8ea5699.jpeg) -#### Step 4：阅读题目要求，接口约定，编写符合`run_kernel`接口规范的CUDA算子 - - 以下代码仅用于说明接口结构，不代表最优实现，也不作为评分参考： - -```python -FlashInfer Ragged Prefill示例代码 -#include -#include -#include -#include - -#ifndef max -#define max(a,b) ((a) > (b) ? (a) : (b)) -#endif - -#ifndef min -#define min(a,b) ((a) < (b) ? (a) : (b)) -#endif - -template -__global__ void ragged_prefill_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 seq_len, - int64_t num_qo_heads, - int64_t num_kv_heads, - float scale, - int64_t causal -) { - int b = blockIdx.x; - int hq = blockIdx.y; - int hkv = hq / (num_qo_heads / num_kv_heads); - - int ty = threadIdx.y; // 0~31, Block 内的 query 索引 (对应 Br) - int tx = threadIdx.x; // 0~31, Warp 内的线程 (lane) - - // 使用 int64_t 防止长序列+大Batch组合下的整数溢出 - int64_t q_idx = (int64_t)blockIdx.z * Br + ty; - bool is_valid = q_idx < seq_len; - - int64_t cur_qo_start = qo_indptr[b]; - int64_t cur_kv_start = kv_indptr[b]; - - float r_Q[D_QK / 32]; - if (is_valid) { - // 显式使用 int64_t 计算全局偏移量，防止超过 2GB 显存边界 - int64_t q_offset_base = (cur_qo_start + q_idx) * num_qo_heads * D_QK + hq * D_QK; - #pragma unroll - for (int d = 0; d < D_QK / 32; ++d) { - r_Q[d] = __bfloat162float(q[q_offset_base + tx + d * 32]); - } - } - - float r_O[D_VO / 32]; - #pragma unroll - for (int d = 0; d < D_VO / 32; ++d) r_O[d] = 0.0f; - - float m = -1e20f; - float l = 0.0f; - - extern __shared__ __nv_bfloat16 smem[]; - __nv_bfloat16* shared_K = smem; - __nv_bfloat16* shared_V = smem + Bc * D_QK; - - int tid = ty * 32 + tx; - int total_threads = Br * 32; - - int64_t block_max_q = min((int64_t)seq_len - 1, (int64_t)(blockIdx.z * Br + Br - 1)); - - int total_k_ldg = (Bc * D_QK) / 8; - int total_v_ldg = (Bc * D_VO) / 8; - - for (int64_t kv_start = 0; kv_start <= block_max_q; kv_start += Bc) { - - // 协同加载 K - for (int i = tid; i < total_k_ldg; i += total_threads) { - int token_idx = i / (D_QK / 8); - int d_idx = i % (D_QK / 8); - int64_t global_kv_idx = kv_start + token_idx; - - if (global_kv_idx < seq_len) { - // 同样强制为 int64_t - int64_t k_offset = (cur_kv_start + global_kv_idx) * num_kv_heads * D_QK + hkv * D_QK + d_idx * 8; - ((float4*)shared_K)[i] = *(reinterpret_cast(&k[k_offset])); - } else { - ((float4*)shared_K)[i] = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - } - } - - // 协同加载 V - for (int i = tid; i < total_v_ldg; i += total_threads) { - int token_idx = i / (D_VO / 8); - int d_idx = i % (D_VO / 8); - int64_t global_kv_idx = kv_start + token_idx; - - if (global_kv_idx < seq_len) { - int64_t v_offset = (cur_kv_start + global_kv_idx) * num_kv_heads * D_VO + hkv * D_VO + d_idx * 8; - ((float4*)shared_V)[i] = *(reinterpret_cast(&v[v_offset])); - } else { - ((float4*)shared_V)[i] = make_float4(0.0f, 0.0f, 0.0f, 0.0f); - } - } - __syncthreads(); - - if (is_valid) { - int64_t kv_end = causal ? min(kv_start + Bc, q_idx + 1) : min(kv_start + Bc, seq_len); - - for (int j = 0; j < kv_end - kv_start; ++j) { - float score = 0.0f; - #pragma unroll - for (int d = 0; d < D_QK / 32; ++d) { - score += r_Q[d] * __bfloat162float(shared_K[j * D_QK + tx + d * 32]); - } - - #pragma unroll - for (int offset = 16; offset > 0; offset /= 2) { - score += __shfl_down_sync(0xffffffff, score, offset); - } - score = __shfl_sync(0xffffffff, score, 0); - score *= scale; - - float m_prev = m; - m = max(m, score); - float P = __expf(score - m); - l = l * __expf(m_prev - m) + P; - - #pragma unroll - for (int d = 0; d < D_VO / 32; ++d) { - r_O[d] = r_O[d] * __expf(m_prev - m) + P * __bfloat162float(shared_V[j * D_VO + tx + d * 32]); - } - } - } - __syncthreads(); - } - - if (is_valid) { - int64_t out_offset_base = (cur_qo_start + q_idx) * num_qo_heads * D_VO + hq * D_VO; - #pragma unroll - for (int d = 0; d < D_VO / 32; ++d) { - float out_val = r_O[d] / l; - output[out_offset_base + tx + d * 32] = __float2bfloat16(out_val); - } - } -} - -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 -) { - float scale = 1.0f / sqrtf((float)head_dim_qk); - const int Br = 32; +#### 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` 实现，优先保证接口正确性，不追求性能。 - dim3 grid(batch_size, num_qo_heads, (seq_len + Br - 1) / Br); - dim3 block(32, Br); + 生成的代码可在本地编译验证后，直接提交到 OJ 上冒烟测试，确认提交链路可用。 - // 根据维度配置绝对安全的共享内存策略 - if (head_dim_qk == 128 && head_dim_vo == 128) { - const int Bc = 64; - int smem = Bc * 128 * 2 + Bc * 128 * 2; // 32 KB - cudaFuncSetAttribute(ragged_prefill_kernel<128, 128, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - ragged_prefill_kernel<128, 128, Br, Bc><<>>( - q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal); - } else if (head_dim_qk == 192 && head_dim_vo == 128) { - const int Bc = 64; - int smem = Bc * 192 * 2 + Bc * 128 * 2; // 40 KB - cudaFuncSetAttribute(ragged_prefill_kernel<192, 128, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - ragged_prefill_kernel<192, 128, Br, Bc><<>>( - q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal); - } else if (head_dim_qk == 256 && head_dim_vo == 256) { - // 关键降级：256维度下，将 Bc 从 64 降为 32。 - // 将所需的共享内存从 64KB 压回到绝对安全的 32KB，防止在部分环境触发 Launch Failure (OOM/Crash)。 - const int Bc = 32; - int smem = Bc * 256 * 2 + Bc * 256 * 2; // 32 KB - cudaFuncSetAttribute(ragged_prefill_kernel<256, 256, Br, Bc>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - ragged_prefill_kernel<256, 256, Br, Bc><<>>( - q, k, v, output, qo_indptr, kv_indptr, seq_len, num_qo_heads, num_kv_heads, scale, causal); + **参考 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) + ``` + + ![Weixin-Image_20260617170807](https://origin.picgo.net/2026/06/17/Weixin-Image_20260617170807_13_100b5fbaeba5856b3b2.png) + + **OJ 冒烟代码** + + 用于最小链路验证。 + + ```cpp + #include + + #include + #include + + #include + + 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(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(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((total + kWarpsPerBlock - 1) / kWarpsPerBlock); + ragged_prefill_smoke_kernel<<>>( + 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. 点击提交，等待评测结果返回； + + ![image 20260616155023443](https://origin.picgo.net/2026/06/16/image-202606161550234433e54d703d581f858.png) + + 评测时间与题目测试点数量、队列状态和平台负载有关，通常需要等待数十秒到数分钟。以平台实际返回为准。 + + ![image iwEcAqNwbmcDAQTRBt8F0QOu](https://origin.picgo.net/2026/06/17/iwEcAqNwbmcDAQTRBt8F0QOuBrA3mrGf2xJ4fgoGfxH9GVsAB9MAAAABijZ-PQgACaJpbQoAC9IAAlPf771f660d13b07471.png) + + **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 依赖是否已安装 ``` -#### Step 5：提交代码开始评测 +### 运行 Benchmark -![image 20260616155023443](https://origin.picgo.net/2026/06/16/image-202606161550234433e54d703d581f858.png) +``` 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 账号，请确认你使用的是组委会统一发放的账号，而不是自行注册账号；如仍无法看到题目，请联系助教或赛事运营确认账号权限。 | + +### 环境问题 | 问题 | 原因 | 解决办法 | | --- | --- | --- | @@ -507,7 +684,7 @@ extern "C" void run_kernel( | `ModuleNotFoundError: No module named 'flashinfer'` | flashinfer 未安装 | 执行 `pip install flashinfer` | | `out of memory` | GPU 显存不足 | 减小 `batch_size` 或 `seq_len` 参数 | -### 运行相关问题 +### 运行问题 | 问题 | 原因 | 解决办法 | | --- | --- | --- | @@ -515,21 +692,32 @@ extern "C" void run_kernel( | `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 了 ``、``、``、``；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 结果 @@ -539,7 +727,6 @@ extern "C" void run_kernel( # 建议创建 results 目录保存 mkdir -p results mv *.csv results/ - ``` ### 2. 深入理解 FlashInfer 核心概念 @@ -569,15 +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`接口完成算子优化，对比不同的语言对于性能加速的影响 \ No newline at end of file +可以使用 Triton 或 TileLang 语言实现 `run_kernel` 接口完成算子优化，对比不同的语言对于性能加速的影响。 \ No newline at end of file From 1dbb422167778322f28ae11c5c10955f0fe49fc5 Mon Sep 17 00:00:00 2001 From: "Ke Xiao (i26293)" Date: Thu, 18 Jun 2026 10:59:48 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E9=A2=98=E7=9B=AE?= =?UTF-8?q?=E5=8C=85=E5=92=8CFlashinfer=E6=95=99=E7=A8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...接口约定.md => 01_接口约定cuda.md} | 0 .../zh_CN/01_接口约定tilelang.md | 52 +++++++++++++++++ .../zh_CN/01_接口约定triton.md | 41 +++++++++++++ ...接口约定.md => 01_接口约定cuda.md} | 0 .../zh_CN/01_接口约定tilelang.md | 55 ++++++++++++++++++ .../zh_CN/01_接口约定triton.md | 44 ++++++++++++++ ...接口约定.md => 01_接口约定cuda.md} | 0 .../zh_CN/01_接口约定tilelang.md | 57 +++++++++++++++++++ .../zh_CN/01_接口约定triton.md | 45 +++++++++++++++ ...接口约定.md => 01_接口约定cuda.md} | 0 .../zh_CN/01_接口约定tilelang.md | 52 +++++++++++++++++ .../zh_CN/01_接口约定triton.md | 41 +++++++++++++ 12 files changed, 387 insertions(+) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/{01_接口约定.md => 01_接口约定cuda.md} (100%) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/{01_接口约定.md => 01_接口约定cuda.md} (100%) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/{01_接口约定.md => 01_接口约定cuda.md} (100%) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/{01_接口约定.md => 01_接口约定cuda.md} (100%) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..35791da --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md @@ -0,0 +1,52 @@ +--- +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 * seq_len, num_qo_heads, head_dim_qk) + k, # Tensor[bf16], shape (batch_size * seq_len, num_kv_heads, head_dim_qk) + v, # Tensor[bf16], shape (batch_size * seq_len, num_kv_heads, head_dim_vo) + output, # Tensor[bf16], shape (batch_size * seq_len, num_qo_heads, head_dim_vo) + qo_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + batch_size, # int64 + seq_len, # int64 + num_qo_heads, # int64 + num_kv_heads, # int64 + head_dim_qk, # int64 + head_dim_vo, # int64 + causal, # int64 +): + global real_kernel + if real_kernel is None: + real_kernel = build_kernel(...) + real_kernel(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) +``` + +### 参数说明 + +* `q/k/v`：FlashInfer ragged prefill 输入 tensor，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `qo_indptr/kv_indptr`：ragged indptr，连续 `int32` +* `causal`：是否启用 causal mask，评测中固定为 `1` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 TileLang kernel。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..d717add --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md @@ -0,0 +1,41 @@ +--- +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 * seq_len, num_qo_heads, head_dim_qk) + k, # Tensor[bf16], shape (batch_size * seq_len, num_kv_heads, head_dim_qk) + v, # Tensor[bf16], shape (batch_size * seq_len, num_kv_heads, head_dim_vo) + output, # Tensor[bf16], shape (batch_size * seq_len, num_qo_heads, head_dim_vo) + qo_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + batch_size, # int64 + seq_len, # int64 + num_qo_heads, # int64 + num_kv_heads, # int64 + head_dim_qk, # int64 + head_dim_vo, # int64 + causal, # int64 +): + ... +``` + +### 参数说明 + +* `q/k/v`：FlashInfer ragged prefill 输入 tensor，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `qo_indptr/kv_indptr`：ragged indptr，连续 `int32` +* `causal`：是否启用 causal mask，评测中固定为 `1` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 Triton kernel。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..098c6da --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md @@ -0,0 +1,55 @@ +--- +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 * seq_len, num_qo_heads, head_dim) + kv_data, # Tensor[bf16], shape (num_blocks, 2, page_block_size, num_kv_heads, head_dim) + output, # Tensor[bf16], shape (batch_size * seq_len, num_qo_heads, head_dim) + qo_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indices, # Tensor[int32], shape (num_blocks) + last_page_len, # Tensor[int32], shape (batch_size) + batch_size, # int64 + seq_len, # int64 + num_qo_heads, # int64 + num_kv_heads, # int64 + head_dim, # int64 + page_block_size, # int64 + causal, # int64 +): + global real_kernel + if real_kernel is None: + real_kernel = build_kernel(...) + real_kernel(q, kv_data, output, qo_indptr, kv_indptr, kv_indices, last_page_len, + batch_size, seq_len, num_qo_heads, num_kv_heads, + head_dim, page_block_size, causal) +``` + +### 参数说明 + +* `q`：query tensor，连续 `bfloat16` +* `kv_data`：paged KV cache，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `qo_indptr/kv_indptr/kv_indices/last_page_len`：paged KV metadata，连续 `int32` +* `page_block_size`：评测中固定为 `16` +* `causal`：评测中固定为 `0` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 TileLang kernel。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..d32ea1b --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md @@ -0,0 +1,44 @@ +--- +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 * seq_len, num_qo_heads, head_dim) + kv_data, # Tensor[bf16], shape (num_blocks, 2, page_block_size, num_kv_heads, head_dim) + output, # Tensor[bf16], shape (batch_size * seq_len, num_qo_heads, head_dim) + qo_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indices, # Tensor[int32], shape (num_blocks) + last_page_len, # Tensor[int32], shape (batch_size) + batch_size, # int64 + seq_len, # int64 + num_qo_heads, # int64 + num_kv_heads, # int64 + head_dim, # int64 + page_block_size, # int64 + causal, # int64 +): + ... +``` + +### 参数说明 + +* `q`：query tensor，连续 `bfloat16` +* `kv_data`：paged KV cache，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `qo_indptr/kv_indptr/kv_indices/last_page_len`：paged KV metadata，连续 `int32` +* `page_block_size`：评测中固定为 `16` +* `causal`：评测中固定为 `0` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 Triton kernel。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..7196b09 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md @@ -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_nope, # Tensor[bf16], shape (batch_size, num_heads, head_dim_ckv) + q_pe, # Tensor[bf16], shape (batch_size, num_heads, head_dim_kpe) + ckv, # Tensor[bf16], shape (batch_size * seq_len, 1, head_dim_ckv) + kpe, # Tensor[bf16], shape (batch_size * seq_len, 1, head_dim_kpe) + output, # Tensor[bf16], shape (batch_size, num_heads, head_dim_ckv) + q_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indices, # Tensor[int32], shape (batch_size * seq_len) + kv_lens, # Tensor[int32], shape (batch_size) + batch_size, # int64 + seq_len, # int64 + num_heads, # int64 + head_dim_ckv, # int64 + head_dim_kpe, # int64 + page_size, # int64 + causal, # int64 +): + global real_kernel + if real_kernel is None: + real_kernel = build_kernel(...) + real_kernel(q_nope, q_pe, ckv, kpe, output, + q_indptr, kv_indptr, kv_indices, kv_lens, + batch_size, seq_len, num_heads, + head_dim_ckv, head_dim_kpe, page_size, causal) +``` + +### 参数说明 + +* `q_nope/q_pe/ckv/kpe`：MLA attention 输入 tensor，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `q_indptr/kv_indptr/kv_indices/kv_lens`：paged attention metadata，连续 `int32` +* `page_size`：评测中固定为 `1` +* `causal`：评测中固定为 `0` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 TileLang kernel。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..43be38b --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md @@ -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_nope, # Tensor[bf16], shape (batch_size, num_heads, head_dim_ckv) + q_pe, # Tensor[bf16], shape (batch_size, num_heads, head_dim_kpe) + ckv, # Tensor[bf16], shape (batch_size * seq_len, 1, head_dim_ckv) + kpe, # Tensor[bf16], shape (batch_size * seq_len, 1, head_dim_kpe) + output, # Tensor[bf16], shape (batch_size, num_heads, head_dim_ckv) + q_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indices, # Tensor[int32], shape (batch_size * seq_len) + kv_lens, # Tensor[int32], shape (batch_size) + batch_size, # int64 + seq_len, # int64 + num_heads, # int64 + head_dim_ckv, # int64 + head_dim_kpe, # int64 + page_size, # int64 + causal, # int64 +): + ... +``` + +### 参数说明 + +* `q_nope/q_pe/ckv/kpe`：MLA attention 输入 tensor，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `q_indptr/kv_indptr/kv_indices/kv_lens`：paged attention metadata，连续 `int32` +* `page_size`：评测中固定为 `1` +* `causal`：评测中固定为 `0` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 Triton kernel。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..8d0fd87 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md @@ -0,0 +1,52 @@ +--- +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, num_qo_heads, head_dim) + kv_data, # Tensor[bf16], shape (num_blocks, 2, page_block_size, num_kv_heads, head_dim) + output, # Tensor[bf16], shape (batch_size, num_qo_heads, head_dim) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indices, # Tensor[int32], shape (num_blocks) + last_page_len, # Tensor[int32], shape (batch_size) + batch_size, # int64 + seq_len_kv, # int64 + num_qo_heads, # int64 + num_kv_heads, # int64 + head_dim, # int64 + page_block_size, # int64 +): + global real_kernel + if real_kernel is None: + real_kernel = build_kernel(...) + real_kernel(q, kv_data, output, kv_indptr, kv_indices, last_page_len, + batch_size, seq_len_kv, num_qo_heads, + num_kv_heads, head_dim, page_block_size) +``` + +### 参数说明 + +* `q`：decode query tensor，连续 `bfloat16` +* `kv_data`：paged KV cache，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `kv_indptr/kv_indices/last_page_len`：paged KV metadata，连续 `int32` +* `page_block_size`：评测中固定为 `16` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 TileLang kernel。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..d1751a9 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md @@ -0,0 +1,41 @@ +--- +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, num_qo_heads, head_dim) + kv_data, # Tensor[bf16], shape (num_blocks, 2, page_block_size, num_kv_heads, head_dim) + output, # Tensor[bf16], shape (batch_size, num_qo_heads, head_dim) + kv_indptr, # Tensor[int32], shape (batch_size + 1) + kv_indices, # Tensor[int32], shape (num_blocks) + last_page_len, # Tensor[int32], shape (batch_size) + batch_size, # int64 + seq_len_kv, # int64 + num_qo_heads, # int64 + num_kv_heads, # int64 + head_dim, # int64 + page_block_size, # int64 +): + ... +``` + +### 参数说明 + +* `q`：decode query tensor，连续 `bfloat16` +* `kv_data`：paged KV cache，连续 `bfloat16` +* `output`：输出缓冲区，连续 `bfloat16`，需要写入结果 +* `kv_indptr/kv_indices/last_page_len`：paged KV metadata，连续 `int32` +* `page_block_size`：评测中固定为 `16` + +`run_kernel` 内部需要自行计算合适的 grid/block，并 launch 你实现的 Triton kernel。