diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md index a7db9db..730ee7f 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md @@ -2,21 +2,21 @@ ## 1. 教程定位 -本教程是赛题二 FlashInfer 任务的 “benchmark 性能基线与 XPU-OJ 提交衔接” 模块，主要帮助学员跑通目标算子的 benchmark 脚本，理解原库 API、输入输出结构、性能指标和性能基线结果，并进一步读懂 XPU-OJ 题目包中的接口约定、测试数据、参考输出和精度要求。 +本教程是赛题二 FlashInfer 任务的 “benchmark 性能基线与 XPU-OJ 提交衔接” 模块，主要帮助学员跑通目标算子的 benchmark 脚本，理解原库 API、输入输出结构、性能指标和性能基线结果，并进一步读懂 XPU-OJ 题目中的接口约定、测试数据、参考输出和精度要求。 需要特别说明： -- 本教程不提供可直接提交的 OJ 参考实现源代码或标准答案代码。 +- 本教程不提供可直接提交的 OJ 参考实现源代码或标准答案代码； -- 本教程仅提供冒烟级 starter 示例代码，用于验证环境、语言、提交链路和 `run_kernel(...)` 接口。 +- 本教程仅提供冒烟级 starter 示例代码，用于验证环境、语言、提交链路和 `run_kernel(...)` 接口； -- benchmark 脚本用于建立性能基线，不是最终提交物。 +- benchmark 脚本用于建立性能基线，不是最终提交物； -- XPU-OJ 题包中的 `baseline()` 属于 OJ 后台参考实现，用于生成 `output_ref`，不是选手提交代码。 +- XPU-OJ 上题目说明中的 `baseline()` 属于 OJ 后台参考实现，用于生成 `output_ref`，不是选手提交代码； - 选手最终需要自行实现 `run_kernel(...)`，并在正确性通过后继续优化性能。 -完成本教程后，学员应能够跑通 benchmark 脚本，记录性能基线结果，读懂 XPU-OJ 题包，理解 OJ 的测试输入与参考实现，并完成一次冒烟级 OJ 提交。 +完成本教程后，学员应能够跑通 benchmark 脚本，记录性能基线结果，读懂 XPU-OJ 上的题目，理解 OJ 的测试输入与参考实现，并完成一次冒烟级 OJ 提交。 ## 2. 学习目标 @@ -28,11 +28,11 @@ 3. 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 benchmark 脚本，记录性能基线结果，为后续算子优化提供对比基准； -4. 读懂对应 XPU-OJ 题包中的题目描述、接口约定、数据范围和精度要求； +4. 读懂对应 XPU-OJ 题目中的题目描述、接口约定、数据范围和精度要求； 5. 完成一次冒烟级 `run_kernel(...)` 提交，确认 OJ 链路、语言环境和接口调用正常； -6. 使用 AI Agent 辅助阅读题包、生成初版实现、定位错误并规划性能优化方向。 +6. 使用 AI Agent 辅助阅读题目、生成初版实现、定位错误并规划性能优化方向。 ## 3. 适用对象 @@ -202,6 +202,21 @@ XPU-OJ 账号由组委会统一发放，参赛者无需自行注册。 > 如果登录后看不到题目，请联系助教或赛事运营确认账号是否已加入对应比赛 / 用户组 +### 4.5 评分规则概要 + +提交作品的最终评价采用 **100 分制**，由客观评测和专家评审共同组成： + +| 类别 | 评审维度 | 权重 | 说明 | +|------|---------|:---:|------| +| 客观评测 | **性能提升效果** | **60%** | 基于 XPU-OJ 各任务榜单排名计分。各任务独立计分，**最终仅取得分最高的一个任务**计入"性能提升效果"，不叠加计分。进入前 30 名且 OJ 分数 > 50 的提交按排名计分：第 1 名得 60 分，每降 1 名扣 2 分，第 30 名得 2 分。未进前 30 或 OJ 分数 ≤ 50 的提交不得分。未通过正确性或稳定性测试的作品记 0 分 | +| 客观评测 | **Agent/Skill 可复现性** | **20%** | Agent 应真实参与源码理解、代码生成、性能分析、自动调优、Benchmark 和多轮迭代等优化过程。分四档：1. 功能可复现得 5 分；2. 性能复现达提交标称 60% 以上得 10 分；3. 80% 以上得 15 分；4. 90% 以上得 20 分 | +| 客观+主观 | **文档说明与演示报告** | **20%** | 根据技术报告、README、运行说明、性能测试报告、Agent/Skill 说明文档、演示视频及答辩材料的完整性、规范性、技术表达质量和工程可复现性综合评分 | + + +> - OJ 榜单分数与比赛最终得分属于不同分值体系。OJ 内部以各算子 baseline 为 50 分基准、硬件理论上限约 100 分；比赛性能部分满分 60 分，按排名映射； +> - FlashInfer 方向包含多个子任务，任选一种或多种提交，各子任务独立计分，取通过正确性和稳定性测试的最高成绩参与排名； +> - 组委会对提交内容严格审查，严禁抄袭，一经发现取消成绩。 + ## 5. 知识预备 ### 5.1 LLM 推理阶段重要概念 @@ -239,11 +254,11 @@ XPU-OJ 账号由组委会统一发放，参赛者无需自行注册。 [*GitHub - MetaX-MACA/McFlashInfer*](https://github.com/MetaX-MACA/McFlashInfer) -> 链接内容可供用于学习 API、算子实现思路、benchmark 方法和优化策略。选手仍需根据 XPU-OJ 题包接口**自行实现**可提交的 `run_kernel(...)` +> 链接内容可供用于学习 API、算子实现思路、benchmark 方法和优化策略。选手仍需根据 XPU-OJ 题目接口**自行实现**可提交的 `run_kernel(...)` ## 6. 项目实践 -- FlashInfer Benchmark -**目标：** 以一个具体算子题目 **20001 FlashInfer Ragged Prefill** 为例，跑通 benchmark 脚本，建立性能基线，理解对应 XPU-OJ 题目包，并完成一次 OJ 冒烟提交，为后续 Agent 辅助优化建立起点。 +**目标：** 以一个具体算子题目 **20001 FlashInfer Ragged Prefill** 为例，理解对应题目，跑通 benchmark 脚本，建立性能基线，并完成一次 OJ 冒烟提交，为后续 Agent 辅助优化建立起点。 ### 6.1 在赛事镜像中运行 FlashInfer Benchmark @@ -441,32 +456,30 @@ else: **Benchmark 与 XPU-OJ 的关系** -赛事镜像中的 Benchmark 和 XPU-OJ 在线评测任务不同。Benchmark 脚本用于理解目标算子的调用方式、输入输出 shape 和性能基线；XPU-OJ 题包用于定义最终评测接口、数据范围、参考输出和精度要求。 +赛事镜像中的 Benchmark 和 XPU-OJ 在线评测任务不同。Benchmark 脚本用于理解目标算子的调用方式、输入输出 shape 和性能基线；XPU-OJ 题目说明用于定义最终评测接口、数据范围、参考输出和精度要求。 | 维度 | Benchmark 脚本 | XPU-OJ 提交 | |------|---------------|------------| | **目的** | 理解算子接口、建立性能基线 | 统一环境下的正确性+性能评测 | | **接口形式** | Python API（`wrapper.plan()` + `wrapper.run()`） | C 接口（`extern "C" void run_kernel(...)`） | -| **数据范围** | 多种 head_dim / batch_size / seq_len 组合 | 固定参数范围（以题包为准） | +| **数据范围** | 多种 head_dim / batch_size / seq_len 组合 | 固定参数范围（以 XPU-OJ 题目说明为准） | | **验证** | 无自动正确性校验 | 强制通过 `torch.allclose(rtol=1e-2, atol=1e-2)` | -| **输出** | CSV 性能记录 | 排行榜得分 | +| **输出** | CSV 性能记录 | 排行榜得分（XPU-OJ 内部得分） | 跑完 benchmark、建立性能基线后，选手需要完成以下转换： 1. 从 benchmark 脚本中理解目标 API，例如 `BatchPrefillWithRaggedKVCacheWrapper`； -2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口； +2. 在 XPU-OJ 上查看对应题目说明的 `run_kernel(...)` 接口； -3. 对照题包中的输入 shape、数据范围和精度要求； +3. 对照题目说明的输入 shape、数据范围和精度要求，编写自己的 `run_kernel(...)`； -4. 编写自己的 `run_kernel(...)`； +4. 提交 OJ，先保证正确性； -5. 提交 OJ，先通过正确性； - -6. 正确性通过后，再对比 benchmark 耗时 / OJ 耗时继续优化。 +5. 正确性通过后，再对比 benchmark 耗时 / OJ 耗时继续优化。 **选择目标题目** -FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benchmark 脚本、OJ 题包、`run_kernel(...)` 接口和数据范围可能不同。 +FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benchmark 脚本、题目说明、`run_kernel(...)` 接口和数据范围。 | OJ 题号 | 算子类型 | 核心特点 | Benchmark 脚本 | FlashInfer API | |---------|---------------|----------------|----------|----------| @@ -475,21 +488,20 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc | **20003** | MLA Paged Attention | DeepSeek MLA特有，双路Q(nope+pe)/双路Cache(ckv+kpe) |`bench_batch_mla.py` | `BatchMLAPagedAttentionWrapper` | | **20004** | Paged Decode | 每次只1个query token，memory-bound |`bench_batch_decode.py` | `BatchDecodeWithPagedKVCacheWrapper` | -**每个子题的接口参数、数据范围和精度要求以对应 XPU-OJ 题包为准。** 下文以题目 **20001 Flashinfer Ragged Prefill** 为例演示从 benchmark 到 XPU-OJ 提交的完整流程。 +**每个子题的接口参数、数据范围和精度要求以对应题目说明为准。** 下文以题目 **20001 Flashinfer Ragged Prefill** 为例演示从 benchmark 到 XPU-OJ 提交的完整流程。 -题目的更多信息可参考 [*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 与最终评测提交之间的关系，理解选手需要实现的内容。 > 完成 benchmark 后，需要注意 benchmark 脚本主要用于建立性能基线，并不需要最终提交 -> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 OJ 参考实现结果进行比较 +> 最终评测以 XPU-OJ 为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 OJ 后台参考实现结果进行比较 -下面以题目 **20001 FlashInfer Ragged Prefill** 为例，逐份解读 OJ 题包中的关键文件。 +下面以题目 **20001 FlashInfer Ragged Prefill** 为例，从本地题目文档 [*Agent 推理算子库优化 - FlashInfer Ragged Prefill*](./operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20001/Agent%20推理算子库优化%20-%20FlashInfer%20Ragged%20Prefill.md) 中逐节解读关键信息。 -1. `00_题目描述.md` — 我要实现什么？ +1. `## 1. 题目描述` — 我要实现什么？ - 该文件明确了三个核心信息： + 该章节明确了三个核心信息： 1. **算子功能**：实现 FlashInfer Ragged KV Cache Prefill 的前向 CUDA C++ 算子。给定扁平存储的 Q / K / V，计算带 causal mask 的 scaled dot-product attention，结果写入 `output` 张量。 @@ -505,7 +517,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc 你的 `run_kernel` 需要实现等效的注意力计算，最终输出与上述 FlashInfer API 在 `rtol = 1e-2, atol = 1e-2` 容差内一致。 -2. `01_接口约定.md` — 我提交的函数签名是什么？ +2. `## 2. 接口约定` — 我提交的函数签名是什么？ 以 CUDA C++ 为例，`run_kernel` 的精确 C 符号如下（函数名、参数类型、顺序、`const` 修饰均不可修改）： @@ -555,52 +567,22 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc - **GQA 头映射**：`kv_head = qo_head / (num_qo_heads / num_kv_heads) = qo_head / 8`。 - - **Triton / TileLang 接口**：函数名同样是 `run_kernel`，参数顺序相同，数据类型映射遵循对应语言的沙箱规则（详见题包中的 `01_接口约定triton.md` 和 `01_接口约定tilelang.md`）。 + - **Triton / TileLang 接口**：函数名同样是 `run_kernel`，参数顺序相同，数据类型映射遵循对应语言的沙箱规则（详见题目文档中 `### 2.2 Triton` 和 `### 2.3 TileLang` 章节）。 -3. `02_数据范围.md` — 输入规模有多大？ +3. `## 6. 数据范围与提示` — 输入规模有多大？ - 该文件定义了测试用例的参数组合和精度容差： + 该章节定义了测试用例的参数组合和精度容差： - **固定参数**：`num_qo_heads = 32`、`num_kv_heads = 4`、`head_dim_qk = 128`、`head_dim_vo = 128`、`causal = 1`、数据类型 `bfloat16` - - **可变参数**（共 **12 个测试用例**）： - | batch_size | seq_len | 估算 Q 张量大小 | 估算 KV 张量大小 | - |:---:|:---:|:---:|:---:| - | 1 | 1024 | 1×1024×32×128×2B = 8 MB | 2×1×1024×4×128×2B ≈ 2 MB | - | 1 | 4096 | 32 MB | 8 MB | - | 1 | 8192 | 64 MB | 16 MB | - | 1 | 16384 | 128 MB | 32 MB | - | 4 | 1024~16384 | 32 MB ~ 512 MB | 8 MB ~ 128 MB | - | 16 | 1024~16384 | 128 MB ~ 2 GB | 32 MB ~ 512 MB | + - **可变参数**：batch_size 覆盖 {1, 2, 4, 15, 16, 27, 33}，seq_len 上界覆盖 {1, 65, 123, 873, 987, 1024, 1280, 2048, 4096, 16384}，共 **15 个测试用例**，覆盖等长长序列、变长 ragged、`q_len < kv_len`、短段和非 2 的幂长度。 - - **精度要求**：`torch.allclose(output.float(), output_ref.float(), rtol = 1e-2, atol = 1e-2)` - - **显存上限**：OJ 评测环境设计 `VRAM_SIZE = 48 GB`（见 `testcase_config.py`） + - **精度要求**：`torch.allclose(output.float(), output_ref.float(), rtol = 1.6e-2, atol = 1.6e-2)`，且允许不超过 1% 的元素超差（匹配率需 ≥ 0.99）。 + + - **显存上限**：OJ 评测环境设计 `VRAM_SIZE = 48 GB`（OJ 后台配置） > **优化提示**：`batch_size = 16, seq_len = 16384` 是最大 workload（Q ≈ 2 GB），需要特别关注显存使用和 compute 效率 -4. `testcase_config.py` — OJ 后台如何评测我？ - - 这个文件不是给选手阅读的文档，但理解它的逻辑有助于定位 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`（以此计算得分）。 - - 当前题目 **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 参考实现输出一致。不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。 - #### Step 7：登录 XPU-OJ 并进入题目页面 使用组委会统一发放的账号登录 XPU-OJ，并进入对应赛题页面。 @@ -626,10 +608,13 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc 1. 在语言下拉框中选择本题支持的提交语言，例如 CUDA Maca、Triton 或 TileLang； -2. 借助 Agent 从题包生成 `run_kernel` 初版； +2. 借助 Agent 阅读 XPU-OJ 的题目信息生成 `run_kernel` 初版； 在下方参考 prompt 的引导下，Agent 会： - 1. 读取对应 OJ 题包中的接口约定文档（`01_接口约定.md`），提取 `run_kernel` 函数签名； - 2. 读取数据范围文档（`02_数据范围.md`），了解输入张量 shape 和精度要求； + + 1. 读取对应题目文档中的接口约定章节（`## 2. 接口约定`），提取 `run_kernel` 函数签名； + + 2. 读取数据范围章节（`## 6. 数据范围与提示`），了解输入张量 shape 和精度要求； + 3. 生成一个能编译通过的最小 `run_kernel` 实现，优先保证接口正确性，不追求性能。 生成的代码可在本地编译验证后，直接提交到 OJ 上冒烟测试，确认提交链路可用。 @@ -641,7 +626,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc - OJ 最终评测不会直接运行 benchmark 脚本，而是调用你提交代码中的 `run_kernel(...)`； - - 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路； + - 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题目文档，并生成一个最小正确版实现思路； - 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。 @@ -691,7 +676,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc 4. 调用选手代码中的 `run_kernel(...)`； - 5. 调用 `testcase_config.py` 中的 `baseline()` / 参考实现生成 `output_ref`； + 5. 调用 OJ 后台参考实现生成 `output_ref`； 6. 将 `run_kernel(...)` 的输出与 `output_ref` 做正确性校验； @@ -825,9 +810,15 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc | ... | ... | ... | ... | ... | ... | ... | ... | | 12 | 16 | 16384 | 349.39 | 124.67 | 2.80 | 112.13 | ✓ | - 2. 总得分查看 — 评分公式详解 + 2. XPU-OJ 平台评分机制（仅供参考） - OJ 平台对每次提交的所有测试点评分后，按以下规则计算总得分。 + 以下为 XPU-OJ 平台的内部评分算法，用于生成榜单中每个题目的单题分数。请注意： + + - **OJ 榜单分数 ≠ 比赛最终得分。** 比赛最终得分为 100 分制（详见 [*4.5 评分规则概要*](#45%20评分规则概要)），其中性能部分取 OJ 排名映射换算，满分 60 分； + + - OJ 内部评分以各算子 baseline 为 50 分基准，硬件理论上限约 100 分。超过 100 分说明超越理论估算； + + - 未通过正确性或稳定性测试的作品不参与排名，客观评测得分记为 0 分。 **单测试点评分公式**（参考 Sol-ExecBench）： @@ -857,36 +848,20 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc 3. 榜单查看 — 如何阅读排行榜 - 排行榜位于 XPU-OJ 比赛页面，展示所有参赛者的各题得分和总分。 - [TODO: image] - - 以下为榜单示例： + 排行榜位于 XPU-OJ 比赛页面，展示各参赛者在每个题目的 OJ 内部得分。请注意：**OJ 榜单中各题的分数与比赛最终得分不是同一体系。** 比赛最终成绩中，性能提升效果仅取你得分最高的一个任务，按该任务的 OJ 排名换算（详见 [*4.5 评分规则概要*](#45%20评分规则概要)）。 - ```plaintext - My Rank: #4 My Total Score: 26.08 - - # User | Ragged Prefill | Paged Prefill | MLA | Paged Decode | Flashattn | Fused MoE | Total - -----|---------------|----------------|---------------|------------|--------------|-----------|-----------|-------- - 1 | muxitest006 | 0 | 0 | 84.75 | 0 | 0 | 0 | 84.75 - 2 | ceerrep | 20.42 | 0 | 0 | 42.33 | 0 | 0 | 62.75 - 3 | muxitest001 | 1.17 | 4.83 | 0.04 | 8.87 | 40.42 | 0 | 54.16 - 4 | muxitest007 | 20.42 | 4.50 | 4.13 | 1.53 | 0 | 0 | 26.08 - 5 | muxitest005 | 5.17 | 0 | 0 | 0 | 0 | 0 | 5.17 - ... - ``` + ![ranklist](https://origin.picgo.net/2026/07/07/-2026-07-07-173642c2d1ab65ebb06d55.png) **榜单解读**： - 6 列对应 6 个 OJ 子题：Ragged Prefill (20001)、Paged Prefill (20002)、MLA Attention (20003)、Paged Decode (20004)、FlashAttention KV Cache Decode (20005)、Fused MoE i8 tn (20006) - - **Total = 6 个子题得分的总和**，排名按 Total 降序 + - 各题分数为 OJ 内部评分（OJ baseline = 50 分），排名按 Total（各题 OJ 得分之和）降序。**比赛最终只取单一最高分任务**，而非多任务累加 - 每列下方的括号数值（如 `(48)`）表示该题提交次数 - `pass = false` 的提交不参与排名（对应 0 分） - **刷榜策略建议**：优先确保**每一题至少有一次 `pass = true` 的提交**（0 分 ≠ 未参与），然后逐个突破有性能瓶颈的题目，逐步提升各题得分和总分。 - ## 7. Agent 使用样例 **目标：** 在本模块中，你可以掌握使用 Agent 完成环境检查、运行 Benchmark、分析结果、理解 OJ 接口、生成和调试 `run_kernel` 等任务。以下是各个任务的参考 prompt。 @@ -912,15 +887,13 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc 请帮我读取最新的性能基线 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置，并与沐曦 GPU 理论峰值带宽做对比。 ``` -### 7.4 理解 OJ 题包接口 +### 7.4 理解 OJ 题目接口 ``` plaintext -请帮我阅读 FlashInfer Ragged Prefill 题包（problem_20001）中的以下文件： -- zh_CN/00_题目描述.md -- zh_CN/01_接口约定cuda.md -- zh_CN/01_接口约定triton.md -- zh_CN/01_接口约定tilelang.md -- zh_CN/02_数据范围.md +请帮我阅读 FlashInfer Ragged Prefill 题目文档（`problem_20001/Agent 推理算子库优化 - FlashInfer Ragged Prefill.md`）： +- `## 1. 题目描述` +- `## 2. 接口约定`（含 CUDA / Triton / TileLang 三语言接口） +- `## 6. 数据范围与提示` 然后帮我总结： 1. run_kernel 的函数签名和每个参数的含义 2. 输入张量的形状约定（q/k/v 的 layout、indptr 的作用） @@ -933,8 +906,8 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc ``` plaintext 请帮我为题目 FlashInfer Ragged Prefill（OJ 题号 20001）生成一个最小可运行的 run_kernel 实现，要求： -1. 阅读题包中的 01_接口约定.md，理解 run_kernel 的函数签名（参数类型、顺序、const 修饰）； -2. 阅读 02_数据范围.md，了解 head_dim_qk、head_dim_vo 的可能取值； +1. 阅读题目文档中的 `## 2. 接口约定` 章节，理解 run_kernel 的函数签名（参数类型、顺序、const 修饰）； +2. 阅读 `## 6. 数据范围与提示` 章节，了解 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） @@ -948,8 +921,7 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc ``` plaintext 我的 run_kernel 提交到 OJ 后显示 Wrong Answer，请帮我对比以下信息： 1. 这是我的 submit 代码：[粘贴你的 run_kernel 实现] -2. 题包的接口约定在这里：[粘贴或引用 01_接口约定.md] -3. 题包的测试配置在这里：[粘贴或引用 testcase_config.py] +2. 题目文档的接口约定在这里：[粘贴或引用 `## 2. 接口约定` 章节] 请帮我逐项检查： - 函数签名是否完全匹配 - GQA 头映射公式是否正确 @@ -1016,9 +988,9 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc | 问题 | 可能原因 | 解决办法 | |------|------|---------| -| **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` 等调试代码后重新提交 | +| **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. 逐行对照对应题目文档的「## 2. 接口约定」章节，确认参数类型（`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` 的 OJ 参考实现 `baseline()` 与你 kernel 输出做 `torch.allclose(rtol=1e-2, atol=1e-2)` 比对；2. GQA 公式：`int hkv = hq * num_kv_heads / num_qo_heads`（整数除法）；3. Causal 逻辑：`kv_end = min(kv_start + Bc, q_idx + 1)`，注意 +1 的处理；4. 对照论文 FlashAttention 的 Algorithm 1 逐行验证 online softmax；5. float4 加载偏移公式：`(cur_kv_start + i) * num_kv_heads * D_QK + hkv * D_QK + d_idx * 8`，确认 `num_kv_heads` 而非 `num_qo_heads` | +| **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. 参考题目文档中的 `## 8. PyTorch 参考实现` 进行本地对拍：用 PyTorch 参考实现与你 kernel 输出做 `torch.allclose(rtol=1e-2, atol=1e-2)` 比对；2. GQA 公式：`int hkv = hq * num_kv_heads / num_qo_heads`（整数除法）；3. Causal 逻辑：`kv_end = min(kv_start + Bc, q_idx + 1)`，注意 +1 的处理；4. 对照论文 FlashAttention 的 Algorithm 1 逐行验证 online softmax；5. float4 加载偏移公式：`(cur_kv_start + i) * num_kv_heads * D_QK + hkv * D_QK + d_idx * 8`，确认 `num_kv_heads` 而非 `num_qo_heads` | ## 9. 下一步学习建议 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20001/Agent 推理算子库优化 - FlashInfer Ragged Prefill.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20001/Agent 推理算子库优化 - FlashInfer Ragged Prefill.md new file mode 100644 index 0000000..a914ce6 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20001/Agent 推理算子库优化 - FlashInfer Ragged Prefill.md @@ -0,0 +1,360 @@ +# Agent 推理算子库优化 - FlashInfer Ragged Prefill + +当前题目说明来源为 [*XPU-OJ 20001*](https://xpuoj.com/contest/2/problem/1)，内容以 XPU-OJ 线上为准。 + +## 1. 题目描述 +你需要实现 FlashInfer ragged KV cache prefill 的CUDA C++前向算子。 + +本题输入采用 FlashInfer BatchPrefillWithRaggedKVCacheWrapper 的 ragged NHD 布局。每个 batch 段的 query/KV 长度由 qo_indptr 和 kv_indptr 给出；seq_len 只是所有段长度的上界，真实总长度分别是 qo_indptr[batch_size] 和 kv_indptr[batch_size]。 + +其中 query heads 采用 GQA 布局：num_qo_heads 个 query/output heads 共享 num_kv_heads 个 KV heads，G = num_qo_heads / num_kv_heads。 + +评测程序会调用你提交代码中的 run_kernel 函数。你需要根据 qo_indptr 和 kv_indptr 读取 ragged Q/K/V，并将结果写入 output。 + +baseline 使用 FlashInfer ragged prefill 的 Python API： + +``` python +wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto") +wrapper.plan(qo_indptr, kv_indptr, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, causal=True, + q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) +wrapper.run(q, k, v, out=output) + +``` + +如何提交代码详见 [*评测指南*](https://xpuoj.com/d/2)。 + +## 2. 接口约定 + +### 2.1 CUDA + +你必须在提交的 CUDA 源码中提供如下 C 符号，函数名、参数类型、顺序必须完全一致，并使用 extern "C" 防止 name mangling： + +``` cpp +#include +#include + +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 +); + +``` + +**参数说明** + +- q：query tensor，shape (total_q, num_qo_heads, head_dim_qk)，连续 bf16，其中 total_q = qo_indptr[batch_size] +- k：key tensor，shape (total_kv, num_kv_heads, head_dim_qk)，连续 bf16，其中 total_kv = kv_indptr[batch_size] +- v：value tensor，shape (total_kv, num_kv_heads, head_dim_vo)，连续 bf16 +- output：输出缓冲区，shape (total_q, num_qo_heads, head_dim_vo)，连续 bf16 +- qo_indptr：query/output ragged indptr，shape (batch_size + 1)，连续 int32 +- kv_indptr：KV ragged indptr，shape (batch_size + 1)，连续 int32 +- seq_len：所有 query/KV 段长度的上界，可用于 launch grid；真实段长必须由 indptr 读取 +- causal：是否启用 causal mask，评测中固定为 1 + +部分测试点是等长段，但也包含 q_len != kv_len 和不同 batch 段长度不相等的 ragged 测试点。实现不能假设 qo_indptr[b + 1] - qo_indptr[b] == seq_len 或 kv_indptr[b + 1] - kv_indptr[b] == seq_len。 + +run_kernel 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 run_kernel 内部做 cudaDeviceSynchronize() 或显式同步。 + +### 2.2 Triton + +你必须在提交的 Python 代码中提供 run_kernel 函数，函数名、参数顺序、类型必须完全一致： + +``` python +import triton +import triton.language as tl + +@triton.jit +def your_kernel(...): + ... + +def run_kernel( + q, # Tensor[bf16], shape (total_q, num_qo_heads, head_dim_qk) + k, # Tensor[bf16], shape (total_kv, num_kv_heads, head_dim_qk) + v, # Tensor[bf16], shape (total_kv, num_kv_heads, head_dim_vo) + output, # Tensor[bf16], shape (total_q, 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, max segment length bound + 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；真实段长和 total_q/total_kv 以 indptr 为准 +- causal：是否启用 causal mask，评测中固定为 1 + +run_kernel 内部需要自行计算合适的 grid/block，并 launch 你实现的 Triton kernel。 + +### 2.3 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 (total_q, num_qo_heads, head_dim_qk) + k, # Tensor[bf16], shape (total_kv, num_kv_heads, head_dim_qk) + v, # Tensor[bf16], shape (total_kv, num_kv_heads, head_dim_vo) + output, # Tensor[bf16], shape (total_q, 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, max segment length bound + 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；真实段长和 total_q/total_kv 以 indptr 为准 +- causal：是否启用 causal mask，评测中固定为 1 + +run_kernel 内部需要自行计算合适的 grid/block，并 launch 你实现的 TileLang kernel。 + +## 3. 输入格式 + +本题输入由评测程序在 GPU 上构造，并按接口约定中的顺序传入 run_kernel。 + +所有 q/k/v/output 均为连续 torch.bfloat16 CUDA tensor，qo_indptr/kv_indptr 为连续 torch.int32 CUDA tensor。 + +张量布局固定为 FlashInfer ragged prefill 的 NHD 布局。 + +## 4. 输出格式 + +输出写入 output，shape 为 (total_q, num_qo_heads, head_dim_vo)，类型为 bfloat16，其中 total_q = qo_indptr[batch_size]。 + +## 5. 样例 + +若 batch_size = 1、seq_len = 4、num_qo_heads = 1、num_kv_heads = 1，则： + +``` +qo_indptr = [0, 4] +kv_indptr = [0, 4] +``` + +第 t 个 query 会访问同一 batch 内的 KV token 前缀；启用 causal mask 时，只能看到位置不超过 t 的 token。例如 t = 2 时： + +``` +attention = softmax(q[2, 0, :] @ k[0:3, 0, :].T / sqrt(head_dim_qk)) +output[2, 0, :] = attention @ v[0:3, 0, :] +``` + +若某个 varlen case 中 q_len=2、kv_len=4，则 causal mask 采用 FlashInfer/sol-execbench 的 bottom-right 对齐：第 t 个 query 可见的 KV 上界为 t + 1 + (kv_len - q_len)。例如 t=0 时可见 k[0:3]，t=1 时可见 k[0:4]。 + +## 6. 数据范围与提示 + +- 数据类型：q/k/v/output 均为 bfloat16 +- KV layout：NHD +- num_qo_heads = 32 +- num_kv_heads = 4 +- causal = 1 +- head_dim_qk, head_dim_vo 取值为 (128, 128) +- batch_size 取值随测试点变化，覆盖 1, 2, 4, 15, 16, 27, 33 +- seq_len 参数表示所有 query/KV 段长度的上界，各测试点的段长上界覆盖 1, 65, 123, 873, 987, 1024, 1280, 2048, 4096, 16384（变长测试点内部还包含 512、640 等更短的真实段长） +- total_q = qo_indptr[batch_size] +- total_kv = kv_indptr[batch_size] + +注意： + +- G = num_qo_heads / num_kv_heads，同一个 KV head 服务连续的 G 个 query heads。 +- 对 query head h_q，对应的 KV head 为 h_q / G。 +- 真实段长必须从 qo_indptr 和 kv_indptr 读取，不能假设每段长度相同。 +- 启用 causal mask 后，采用 bottom-right 对齐。若当前段 q_len != kv_len，第 t 个 query 可访问的位置满足 kv_pos < t + 1 + (kv_len - q_len)。 +- 输出校验容差为 rtol=1.6e-2, atol=1.6e-2，且允许不超过 1% 的元素超差（匹配率需 ≥ 0.99）。 +- 被容忍的超差元素其绝对误差仍不得超过 8 × (atol + rtol · |ref|)，避免个别段被整段算错而蒙混通过。 +- 单 token 边界（用例 14）和非 2 的幂尾段（用例 15）为小规模确定性用例，要求逐元素通过（匹配率需 = 1.0）。 +- q/k/v 使用标准正态分布生成，避免均匀正输入导致长序列 softmax 退化成近似 prefix mean。 + +## 7. 测试用例尺寸 + +测试点顺序与 testcase_config.py 的 TESTCASES 一致。共 15 个测试点，全部 head_dim_qk = head_dim_vo = 128，覆盖等长长序列、变长 ragged、q_len < kv_len、短段和非 2 的幂长度。 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
测试用例ID类型batchtotal_qtotal_kvmax_qmax_kvheadshead_dim
1混合 ragged 长序列331629498732/4128/128
2等长序列11024
34096
416384
5440961024
6163844096
7161024
8327682048
9变长 q_len < kv_len4204840965121024
10混合变长 q_len < kv_len153635846401280
11双段变长 q_len < kv_len2102430725122048
12混合 ragged 中长序列2712251873
13混合 ragged 短序列15969123
14单 token 边界1
15非 2 的幂尾段29865
+ +说明：变长测试点的真实段长由 qo_indptr 和 kv_indptr 给出；参赛实现应始终以 indptr 为准，而不是从 seq_len、total_q 或 total_kv 反推出每段长度。 + +## 8. PyTorch 参考实现 + +``` python +def baseline(q, k, v, output, qo_indptr, kv_indptr, + batch_size, seq_len, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, causal): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + backend="auto", + ) + wrapper.plan( + qo_indptr, + kv_indptr, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal=bool(causal), + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + wrapper.run(q, k, v, out=output) + +``` diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20002/Agent 推理算子库优化 - FlashInfer Paged Prefill.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20002/Agent 推理算子库优化 - FlashInfer Paged Prefill.md new file mode 100644 index 0000000..9c8cc78 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20002/Agent 推理算子库优化 - FlashInfer Paged Prefill.md @@ -0,0 +1,189 @@ +# Agent 推理算子库优化 - FlashInfer Paged Prefill + +当前题目说明来源为 [*XPU-OJ 20002*](https://xpuoj.com/contest/2/problem/2)，内容以 XPU-OJ 线上为准。 + +## 1. 题目描述 +你需要实现 FlashInfer paged KV cache prefill 的CUDA C++前向算子。 + +本题输入采用 FlashInfer BatchPrefillWithPagedKVCacheWrapper 的 paged NHD 布局。每个 batch 中有 seq_len 个 query token，KV cache 也有 seq_len 个 token，并按 page 存储。 + +评测程序会调用你提交代码中的 run_kernel 函数。你需要根据 qo_indptr、kv_indptr、kv_indices 和 last_page_len 读取 paged KV cache，并将结果写入 output。 + +baseline 使用 FlashInfer paged prefill 的 Python API： + +``` python +wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto") +wrapper.plan(qo_indptr, kv_indptr, kv_indices, last_page_len, + num_qo_heads, num_kv_heads, head_dim, page_block_size, + causal=bool(causal), + q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) +wrapper.run(q, kv_data, out=output) + +``` + +如何提交代码详见 [*评测指南*](https://xpuoj.com/d/2)。 + +## 2. 接口约定 + +### 2.1 CUDA + +你必须在提交的 CUDA 源码中提供如下 C 符号，函数名、参数类型、顺序必须完全一致，并使用 extern "C" 防止 name mangling： + +``` cpp +#include +#include + +extern "C" void run_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* kv_data, + __nv_bfloat16* output, + const int32_t* qo_indptr, + const int32_t* kv_indptr, + const int32_t* kv_indices, + const int32_t* last_page_len, + int64_t batch_size, + int64_t seq_len, + int64_t num_qo_heads, + int64_t num_kv_heads, + int64_t head_dim, + int64_t page_block_size, + int64_t causal +); + +``` + +**参数说明** + +- q：query tensor，shape (batch_size * seq_len, num_qo_heads, head_dim)，连续 bf16 +- kv_data：paged KV cache，shape (num_blocks, 2, page_block_size, num_kv_heads, head_dim)，连续 bf16，其中 kv_data[:, 0] 为 key，kv_data[:, 1] 为 value +- output：输出缓冲区，shape (batch_size * seq_len, num_qo_heads, head_dim)，连续 bf16 +- qo_indptr：query/output indptr，shape (batch_size + 1)，连续 int32 +- kv_indptr：paged KV indptr，shape (batch_size + 1)，连续 int32 +- kv_indices：page index，shape (num_blocks)，连续 int32 +- last_page_len：每个 batch 最后一个 page 的有效 token 数，shape (batch_size)，连续 int32 +- page_block_size：page size，评测中固定为 16 +- causal：是否启用 causal mask，本题按 benchmark case 固定为 0 + +run_kernel 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 run_kernel 内部做 cudaDeviceSynchronize() 或显式同步。 + +### 2.2 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。 + +### 2.3 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。 + +## 3. 输入格式 + +本题输入由评测程序在 GPU 上构造，并按接口约定中的顺序传入 run_kernel。 + +q/kv_data/output 均为连续 torch.bfloat16 CUDA tensor，qo_indptr/kv_indptr/kv_indices/last_page_len 均为连续 torch.int32 CUDA tensor。 + +KV layout 固定为 FlashInfer paged prefill 的 NHD 布局，page size 固定为 16。 + +## 4. 输出格式 + +输出写入 output，shape 为 (batch_size * seq_len, num_qo_heads, head_dim)，类型为 bfloat16。 + +## 5. 样例 + +若 batch_size = 1、seq_len = 32、page_block_size = 16，则： + +``` +qo_indptr = [0, 32] +kv_indptr = [0, 2] +kv_indices = [0, 1] +last_page_len = [16] +``` + +第 0 个 batch 的 KV token 存放在 page 0 和 page 1 中，每个 page 有 16 个 token。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20003/Agent 推理算子库优化 - FlashInfer MLA Paged Attention.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20003/Agent 推理算子库优化 - FlashInfer MLA Paged Attention.md new file mode 100644 index 0000000..5aa9d8e --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20003/Agent 推理算子库优化 - FlashInfer MLA Paged Attention.md @@ -0,0 +1,194 @@ +# Agent 推理算子库优化 - FlashInfer MLA Paged Attention + +当前题目说明来源为 [*XPU-OJ 20003*](https://xpuoj.com/contest/2/problem/3)，内容以 XPU-OJ 线上为准。 + +## 1. 题目描述 +你需要实现 FlashInfer MLA paged attention 的 CUDA C++ 前向算子。 + +本题使用 BatchMLAPagedAttentionWrapper 的 DeepSeek MLA decode 配置：q_nope 表示不含 RoPE 的 query，q_pe 表示 RoPE 部分的 query，ckv 与 kpe 分别表示 compressed KV cache 与 RoPE KV cache。MLA page size 固定为 1。 + +评测程序会调用你提交代码中的 run_kernel 函数。你需要根据 q_indptr、kv_indptr、kv_indices 和 kv_lens 读取 cache，并将结果写入 output。 + +baseline 使用 FlashInfer MLA 的 Python API： + +``` python +wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace, backend="auto") +wrapper.plan(q_indptr, kv_indptr, kv_indices, kv_lens, + num_heads, head_dim_ckv, head_dim_kpe, + page_size, False, sm_scale, + q_nope.dtype, ckv.dtype) +wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) + +``` + +如何提交代码详见 [*评测指南*](https://xpuoj.com/d/2)。 + +## 2. 接口约定 + +### 2.1 CUDA + +你必须在提交的 CUDA 源码中提供如下 C 符号，函数名、参数类型、顺序必须完全一致，并使用 extern "C" 防止 name mangling： + +``` cpp +#include +#include + +extern "C" void run_kernel( + const __nv_bfloat16* q_nope, + const __nv_bfloat16* q_pe, + const __nv_bfloat16* ckv, + const __nv_bfloat16* kpe, + __nv_bfloat16* output, + const int32_t* q_indptr, + const int32_t* kv_indptr, + const int32_t* kv_indices, + const int32_t* kv_lens, + int64_t batch_size, + int64_t seq_len, + int64_t num_heads, + int64_t head_dim_ckv, + int64_t head_dim_kpe, + int64_t page_size, + int64_t causal +); + +``` + +**参数说明** + +- q_nope：query 的 compressed/nope 部分，shape (batch_size, num_heads, head_dim_ckv)，连续 bf16 +- q_pe：query 的 RoPE 部分，shape (batch_size, num_heads, head_dim_kpe)，连续 bf16 +- ckv：compressed KV cache，shape (batch_size * seq_len, 1, head_dim_ckv)，连续 bf16 +- kpe：RoPE KV cache，shape (batch_size * seq_len, 1, head_dim_kpe)，连续 bf16 +- output：输出缓冲区，shape (batch_size, num_heads, head_dim_ckv)，连续 bf16 +- q_indptr：decode query indptr，shape (batch_size + 1)，内容为 [0, 1, ..., batch_size] +- kv_indptr：KV indptr，shape (batch_size + 1)，每段长度为 seq_len +- kv_indices：page index，shape (batch_size * seq_len)，连续 int32 +- kv_lens：每个 batch 的 KV 长度，shape (batch_size)，连续 int32 +- page_size：评测中固定为 1 +- causal：评测中固定为 0 + +run_kernel 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 run_kernel 内部做 cudaDeviceSynchronize() 或显式同步。 + +### 2.2 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。 + +### 2.3 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。 + +## 3. 输入格式 + +本题输入由评测程序在 GPU 上构造，并按接口约定中的顺序传入 run_kernel。 + +q_nope/q_pe/ckv/kpe/output 均为连续 torch.bfloat16 CUDA tensor，q_indptr/kv_indptr/kv_indices/kv_lens 均为连续 torch.int32 CUDA tensor。 + +## 4. 输出格式 + +输出写入 output，shape 为 (batch_size, num_heads, head_dim_ckv)，类型为 bfloat16。 + +## 5. 样例 + +若 batch_size = 2、seq_len = 4，则： + +``` +q_indptr = [0, 1, 2] +kv_indptr = [0, 4, 8] +kv_indices = [0, 1, 2, 3, 4, 5, 6, 7] +kv_lens = [4, 4] +``` + +每个 batch 只有 1 个 decode query，会访问对应 batch 的全部 KV cache。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20004/Agent 推理算子库优化 - FlashInfer Paged Decode.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20004/Agent 推理算子库优化 - FlashInfer Paged Decode.md new file mode 100644 index 0000000..4459eb0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20004/Agent 推理算子库优化 - FlashInfer Paged Decode.md @@ -0,0 +1,178 @@ +# Agent 推理算子库优化 - FlashInfer Paged Decode + +当前题目说明来源为 [*XPU-OJ 20004*](https://xpuoj.com/contest/2/problem/4)，内容以 XPU-OJ 线上为准。 + +## 1. 题目描述 +你需要实现 FlashInfer paged KV cache decode 的 CUDA C++ 前向算子。 + +本题输入采用 FlashInfer BatchDecodeWithPagedKVCacheWrapper 的 paged NHD 布局。每个 batch 只有 1 个 query token，KV cache 长度为 seq_len_kv。 + +评测程序会调用你提交代码中的 run_kernel 函数。你需要根据 kv_indptr、kv_indices 和 last_page_len 读取 paged KV cache，并将结果写入 output。 + +baseline 使用 FlashInfer paged decode 的 Python API： + +``` python +wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace, kv_layout="NHD", use_tensor_cores=True) +wrapper.plan(kv_indptr, kv_indices, last_page_len, + num_qo_heads, num_kv_heads, head_dim, page_block_size, + data_type=torch.bfloat16, q_data_type=torch.bfloat16) +wrapper.run(q, kv_data, out=output) + +``` + +如何提交代码详见 [*评测指南*](https://xpuoj.com/d/2)。 + +## 2. 接口约定 + +### 2.1 CUDA + +你必须在提交的 CUDA 源码中提供如下 C 符号，函数名、参数类型、顺序必须完全一致，并使用 extern "C" 防止 name mangling： + +``` cpp +#include +#include + +extern "C" void run_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* kv_data, + __nv_bfloat16* output, + const int32_t* kv_indptr, + const int32_t* kv_indices, + const int32_t* last_page_len, + int64_t batch_size, + int64_t seq_len_kv, + int64_t num_qo_heads, + int64_t num_kv_heads, + int64_t head_dim, + int64_t page_block_size +); + +``` + +**参数说明** + +- q：decode query tensor，shape (batch_size, num_qo_heads, head_dim)，连续 bf16 +- kv_data：paged KV cache，shape (num_blocks, 2, page_block_size, num_kv_heads, head_dim)，连续 bf16，其中 kv_data[:, 0] 为 key，kv_data[:, 1] 为 value +- output：输出缓冲区，shape (batch_size, num_qo_heads, head_dim)，连续 bf16 +- kv_indptr：paged KV indptr，shape (batch_size + 1)，连续 int32 +- kv_indices：page index，shape (num_blocks)，连续 int32 +- last_page_len：每个 batch 最后一个 page 的有效 token 数，shape (batch_size)，连续 int32 +- page_block_size：page size，评测中固定为 16 + +run_kernel 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确，不建议在 run_kernel 内部做 cudaDeviceSynchronize() 或显式同步。 + +### 2.2 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。 + +### 2.3 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。 + +## 3. 输入格式 + +本题输入由评测程序在 GPU 上构造，并按接口约定中的顺序传入 run_kernel。 + +q/kv_data/output 均为连续 torch.bfloat16 CUDA tensor，kv_indptr/kv_indices/last_page_len 均为连续 torch.int32 CUDA tensor。 + +KV layout 固定为 FlashInfer paged decode 的 NHD 布局，page size 固定为 16。 + +## 4. 输出格式 + +输出写入 output，shape 为 (batch_size, num_qo_heads, head_dim)，类型为 bfloat16。 + +## 5. 样例 + +若 batch_size = 1、seq_len_kv = 32、page_block_size = 16，则： + +``` +kv_indptr = [0, 2] +kv_indices = [0, 1] +last_page_len = [16] +``` + +decode query 会访问该 batch 的全部 32 个 KV token。