From 374871a83800f54427344e4e996c85eb5353ee37 Mon Sep 17 00:00:00 2001 From: "Yuliang Feng (i26389)" Date: Wed, 8 Jul 2026 17:18:16 +0800 Subject: [PATCH] =?UTF-8?q?=E8=B5=9B=E9=A2=98=E5=85=A5=E5=8F=A3=E4=BF=AE?= =?UTF-8?q?=E6=94=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...ashAttention关键算子迁移与优化.md | 67 ++++++++++--------- 1 file changed, 37 insertions(+), 30 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md index bd03394..674fba7 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md @@ -1,8 +1,8 @@ -# Flashattention 迁移 Benchmark 实战:从性能基线到 XPU-OJ 评测 +# Agent推理算子库优化-FlashAttention KV Cache Decode Benchmark 实战：从性能基线到 XPU-OJ 评测 ## 1. 教程定位 -本教程是参赛训练课程的 **FlashAttention Benchmark 入门与评测提交衔接** 模块，主要帮助用户跑通 FlashAttention paged KV-cache 推理核函数 `flash_attn_with_kvcache` 的基准测试流程，理解 benchmark 脚本的输入输出、性能指标和评测含义，并基于 XPU-OJ 题包完成一个最小正确版 `run_kernel` 的实现与提交。 +本教程面向 **沐曦-揭榜挂帅-Agent推理算子库优化-FlashAttention任务**，是围绕题目 **Agent推理算子库优化-FlashAttention KV Cache Decode** 的 **FlashAttention Benchmark 入门与评测提交衔接** 模块，主要帮助用户跑通 FlashAttention paged KV-cache 推理核函数 `flash_attn_with_kvcache` 的基准测试流程，理解 benchmark 脚本的输入输出、性能指标和评测含义，并基于 XPU-OJ 题包完成一个最小正确版 `run_kernel` 的实现与提交。 需要特别说明：本教程中的 benchmark 脚本主要用于帮助参赛者理解目标算子的调用方式、输入输出结构和性能基线，benchmark 脚本不是最终提交物。最终评测以 XPU-OJ 题包为准，参赛者需要根据题包中的接口约定实现自己的 `run_kernel`，并在输出结果对齐 OJ 参考结果的前提下提升性能。 @@ -30,7 +30,7 @@ 本模块适合以下人员： -* 参与 AI 基础设施竞赛的参赛者 +* 参与 **沐曦-揭榜挂帅-Agent推理算子库优化-FlashAttention任务** 的参赛者 * 对 GPU 算子性能优化感兴趣的开发者 * 需要了解 FlashAttention KV-Cache 推理性能的研究人员 @@ -200,7 +200,7 @@ ls **预期结果：** ```text -benchmark结果实例 benchmark_kvcache.py +benchmark_kvcache.py ``` ### Step 4：配置基准测试参数 @@ -371,7 +371,7 @@ Benchmark 脚本用于理解目标算子的调用方式、输入输出 shape 和 跑完 benchmark、建立性能基线后，选手需要完成以下转换： 1. 从 benchmark 脚本中理解目标 API，本任务对应 `flash_attn.flash_attn_interface` 中的 `flash_attn_with_kvcache`，使用 paged KV cache 布局。 -2. 在 XPU-OJ 平台上查看 `FlashAttention KV Cache Decode` 题目的接口约定。 +2. 在 XPU-OJ 平台上查看 **Agent推理算子库优化-FlashAttention KV Cache Decode** 题目的接口约定。 3. 对照题包中的输入 shape、数据范围和精度要求。 4. 编写自己的 `run_kernel(...)`。 5. 提交 OJ，先通过正确性。 @@ -379,17 +379,17 @@ Benchmark 脚本用于理解目标算子的调用方式、输入输出 shape 和 ### Step 9：题目说明与提交入口 -注意：每个子题的接口参数、数据范围和精度要求可能不同，正式要求以对应 XPU-OJ 题目界面为准。本节以 **FlashAttention KV Cache Decode** 为例，演示从 benchmark 到 XPU-OJ 提交的完整流程。 +本教程只覆盖 **Agent推理算子库优化-FlashAttention KV Cache Decode** 这一题。正式接口、数据范围和精度要求以 XPU-OJ 题目界面为准。本节演示从 benchmark 到 XPU-OJ 提交的完整流程。 * **算子说明**：实现 paged KV cache 下的 decode 注意力，每个 batch 只有 1 个 query token，KV cache 按 page 存储，长度由 `seqlen_k` 决定。 -* **对应 OJ 题目**：XPU-OJ 上 `FlashAttention KV Cache Decode` 题（题号 20005）。 +* **对应 OJ 题目**：XPU-OJ 上 **Agent推理算子库优化-FlashAttention KV Cache Decode** 题。 -使用组委会统一发放的账号登录 XPU-OJ，并进入对应赛题页面。 +使用组委会统一发放的账号登录 XPU-OJ，并进入对应比赛页面和题目页面。 1. 打开 XPU-OJ 平台：https://xpuoj.com/ 2. 使用组委会统一发放的账号和初始密码登录。 -3. 登录后进入比赛 / 题目列表页面。 -4. 找到对应题目，例如 `20005 FlashAttention KV Cache Decode`。 +3. 登录后进入比赛 / 题目列表页面，找到比赛 **沐曦-揭榜挂帅-Agent推理算子库优化-FlashAttention任务**。 +4. 找到对应题目 **Agent推理算子库优化-FlashAttention KV Cache Decode**。 5. 点击进入题目详情页，查看题目描述、接口约定、数据范围和提交入口。 ![image](https://origin.picgo.net/2026/06/23/image123512a0b877cd31.png) @@ -468,7 +468,7 @@ PAGE_BLOCK_SIZE = 16 CAUSAL = 0 ``` -当前 `FlashAttention KV Cache Decode` 题的校验方式为： +当前 **Agent推理算子库优化-FlashAttention KV Cache Decode** 题的校验方式为： ```python torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) @@ -505,14 +505,14 @@ OJ 对每次提交大致会走以下流程： 7. 正确性通过后，统计运行耗时或性能指标 8. 根据题目评分规则换算该题得分 9. 更新该题历史最好成绩 -10. 汇总各题最好成绩，得到排行榜总分 +10. 在榜单中展示本题得分和排名 ``` ### Step 12：分析评测结果与评分机制 **目标：** 理解 OJ 评测结果的含义，分析性能表现。 -在 XPU-OJ 平台上查看提交结果。提交详情会显示状态、总得分、时间、内存、编译信息以及各测试点结果。 +在 XPU-OJ 平台上查看提交结果。提交详情会显示状态、本题得分、时间、内存、编译信息以及各测试点结果。 ![image](https://origin.picgo.net/2026/06/24/image158f29857ec8755e.png) @@ -555,15 +555,22 @@ OJ 对每次提交大致会走以下流程： OJ 平台对单测试点的评分遵循以下公式： -$$ +```math S(T_k) = \frac{100}{1 + \left(\frac{1}{0.5} - 1\right) \cdot \frac{T_k - T_h}{T_b - T_h}} -$$ +``` 其中： * $T_k$：你的 kernel 平均执行时间 * $T_b$：Baseline 参考实现平均执行时间，对应 50 分 -* $T_h$：硬件理论下限耗时，$T_h = \max\left(\frac{\text{FLOPs}}{\text{peak\_tflops}},\ \frac{\text{bytes}}{\text{peak\_bw}}\right)$，对应 100 分 +* $T_h$：硬件理论下限耗时，对应 100 分，计算方式为： + +```math +T_h = \max\left( + \frac{\mathrm{FLOPs}}{\mathrm{peak\_tflops}}, + \frac{\mathrm{bytes}}{\mathrm{peak\_bw}} +\right) +``` **关键分数节点：** @@ -576,11 +583,11 @@ $$ 当单测试点得分超过 150 分时，平台会按对数压缩规则显示： -$$ -S_{\text{display}} = 150 + 10 \cdot \log_{10}(S/150) -$$ +```math +S_{\mathrm{display}} = 150 + 10 \cdot \log_{10}\left(\frac{S}{150}\right) +``` -总得分为各测试点得分的算术平均，总耗时为各测试点 $T_k$ 的求和。 +本题得分为各测试点得分的算术平均，本题总耗时为各测试点 $T_k$ 的求和。 ### Step 13：榜单查看与初步优化方向 @@ -590,9 +597,9 @@ $$ ![image](https://origin.picgo.net/2026/06/24/image6f3c0164de8a994e.png) -* **总得分：** 各题目得分的总和，排名按总得分从高到低排序。 -* **个人排名：** 页面顶部会显示“我的排名”和“我的总分”，方便快速了解自己的位置。 -* **各题目得分：** 表格中每列对应一个题目的得分，帮助分析不同算子优化任务上的表现。 +* **本题得分：** 展示该题的历史最好成绩，排名按本题得分从高到低排序。 +* **个人排名：** 页面顶部会显示你的当前排名和当前得分，方便快速了解自己的位置。 +* **提交次数：** 分数下方括号中的数字表示该账号在本题下的提交次数。 初步优化方向包括： @@ -637,7 +644,7 @@ opencode | 任务阶段 | 参考 Prompt 模板 | 核心目的 | | --- | --- | --- | -| 题包解析 | 请阅读题号 20005 的题目界面和 FlashAttention 任务包材料，总结 CUDA Maca `run_kernel` 函数签名、Paged KV Cache 寻址公式、精度校验方式和测试数据范围。 | 提取接口契约，明确参数 shape 和边界条件 | +| 题包解析 | 请阅读 **Agent推理算子库优化-FlashAttention KV Cache Decode** 的题目界面和 FlashAttention 任务包材料，总结 CUDA Maca `run_kernel` 函数签名、Paged KV Cache 寻址公式、精度校验方式和测试数据范围。 | 提取接口契约，明确参数 shape 和边界条件 | | 生成冒烟代码 | 请生成一个最小可运行的 CUDA Maca `run_kernel` 实现，要求严格匹配 `extern "C"` 签名，支持 Paged KV Cache 的 `block_table` 寻址，支持 head 映射，优先保证正确性。 | 快速验证接口和环境 | | OJ 报错调试 | 我的代码提交后 `Wrong Answer`。这是我的代码和 SPJ Report。请检查 Paged KV 地址映射、bf16 到 float32 的计算转换、尾部 page 有效 token 判断是否正确。 | 结构化排查功能错误 | | 性能瓶颈分析 | 这是测试点的 SPJ Report。请分析 `User kernel` 与 `Hardware bound` 的差距，判断更接近 compute-bound 还是 memory-bound，并给出具体的 mctlass 或访存优化建议。 | 将 OJ 反馈转化为优化行动 | @@ -645,7 +652,7 @@ opencode **题包解析 Prompt：** ```text -请阅读 XPU-OJ 上题号 20005 FlashAttention KV Cache Decode 的题目说明，以及 flashattn_task_package 中与 benchmark / 提交相关的材料。 +请阅读 XPU-OJ 上 **Agent推理算子库优化-FlashAttention KV Cache Decode** 的题目说明，以及 flashattn_task_package 中与 benchmark / 提交相关的材料。 请输出以下内容： 1. CUDA Maca 版本 run_kernel 的完整函数签名； @@ -686,7 +693,7 @@ opencode **Wrong Answer 调试 Prompt：** ```text -我的 FlashAttention KV Cache Decode 代码提交后出现 Wrong Answer。 +我的 Agent推理算子库优化-FlashAttention KV Cache Decode 代码提交后出现 Wrong Answer。 这是我的代码： [粘贴代码] @@ -708,7 +715,7 @@ opencode **性能瓶颈分析 Prompt：** ```text -这是 FlashAttention KV Cache Decode 某个测试点的 SPJ Report： +这是 Agent推理算子库优化-FlashAttention KV Cache Decode 某个测试点的 SPJ Report： [粘贴报告] 请分析： @@ -770,11 +777,11 @@ opencode ### 10.4 结合 OJ Report 做定向优化 -不要只看总分。单测试点的 `Config`、`User kernel`、`Hardware bound` 和 `Speedup vs base` 更适合指导下一轮优化方向。长序列、大 batch、小 batch 的瓶颈可能完全不同。 +不要只看本题最终得分。单测试点的 `Config`、`User kernel`、`Hardware bound` 和 `Speedup vs base` 更适合指导下一轮优化方向。长序列、大 batch、小 batch 的瓶颈可能完全不同。 -### 10.5 扩展到其他赛题 +### 10.5 继续优化本题 -完成 FlashAttention KV Cache Decode 后，可以继续尝试 XPU-OJ 上的其他算子优化题目，例如 FlashInfer MLA Paged Attention 或 FlashInfer Paged Prefill +完成 Agent推理算子库优化-FlashAttention KV Cache Decode 的冒烟提交后，可以继续围绕不同 batch、KV 长度和 head dimension 配置做定向优化。 ## 附录:完整代码参考