From 794e15fbc616128e7b643a951f4e9a824d054ae4 Mon Sep 17 00:00:00 2001 From: MaseChen <93691652+MaseChen@users.noreply.github.com> Date: Tue, 7 Jul 2026 03:04:09 -0700 Subject: [PATCH 1/2] =?UTF-8?q?(flashinfer):=20=E5=8E=BB=E9=99=A4?= =?UTF-8?q?=E9=A2=98=E5=8C=85=E5=8F=8A=E5=85=B6=E6=8F=8F=E8=BF=B0=EF=BC=8C?= =?UTF-8?q?=E6=9B=B4=E6=96=B0=E6=96=87=E4=BB=B6=E7=BB=93=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...：从性能基线到XPU-OJ冒烟提交.md | 178 ++++----- ...ithPagedKVCacheWrapper_20260525_093143.csv | 145 +++++++ ...APagedAttentionWrapper_20260525_074200.csv | 33 ++ ...ithPagedKVCacheWrapper_20260526_150417.csv | 33 ++ ...thRaggedKVCacheWrapper_20260526_145454.csv | 49 +++ .../benchmark/bench_batch_decode.py | 101 +++++ .../benchmark/bench_batch_mla.py | 110 ++++++ .../benchmark/bench_batch_prefill_paged.py | 133 +++++++ .../benchmark/bench_batch_prefill_ragged.py | 136 +++++++ .../benchmark/bench_common.py | 91 +++++ .../starter/示例冒烟代码.md | 209 ++++++++++ .../xpuoj_problem/problem_20001/meta.json | 16 + .../problem_20001/testcase_config.py | 307 +++++++++++++++ .../problem_20001/zh_CN/00_题目描述.md | 23 ++ .../zh_CN/01_接口约定cuda.md | 41 ++ .../zh_CN/01_接口约定tilelang.md | 52 +++ .../zh_CN/01_接口约定triton.md | 41 ++ .../problem_20001/zh_CN/02_数据范围.md | 26 ++ .../problem_20001/zh_CN/03_参考实现.md | 28 ++ .../problem_20001/zh_CN/_title.txt | 1 + .../xpuoj_problem/problem_20002/meta.json | 16 + .../problem_20002/testcase_config.py | 273 +++++++++++++ .../problem_20002/zh_CN/00_题目描述.md | 22 ++ .../zh_CN/01_接口约定cuda.md | 42 ++ .../zh_CN/01_接口约定tilelang.md | 55 +++ .../zh_CN/01_接口约定triton.md | 44 +++ .../problem_20002/zh_CN/02_数据范围.md | 26 ++ .../problem_20002/zh_CN/03_参考实现.md | 31 ++ .../problem_20002/zh_CN/_title.txt | 1 + .../xpuoj_problem/problem_20003/meta.json | 16 + .../problem_20003/testcase_config.py | 245 ++++++++++++ .../problem_20003/zh_CN/00_题目描述.md | 22 ++ .../zh_CN/01_接口约定cuda.md | 46 +++ .../zh_CN/01_接口约定tilelang.md | 57 +++ .../zh_CN/01_接口约定triton.md | 45 +++ .../problem_20003/zh_CN/02_数据范围.md | 26 ++ .../problem_20003/zh_CN/03_参考实现.md | 30 ++ .../problem_20003/zh_CN/_title.txt | 1 + .../xpuoj_problem/problem_20004/meta.json | 16 + .../problem_20004/testcase_config.py | 248 ++++++++++++ .../problem_20004/zh_CN/00_题目描述.md | 22 ++ .../zh_CN/01_接口约定cuda.md | 38 ++ .../zh_CN/01_接口约定tilelang.md | 52 +++ .../zh_CN/01_接口约定triton.md | 41 ++ .../problem_20004/zh_CN/02_数据范围.md | 26 ++ .../problem_20004/zh_CN/03_参考实现.md | 29 ++ .../problem_20004/zh_CN/_title.txt | 1 + ...库优化 - FlashInfer Ragged Prefill.md | 360 ++++++++++++++++++ ...子库优化 - FlashInfer Paged Prefill.md | 189 +++++++++ ...优化 - FlashInfer MLA Paged Attention.md | 194 ++++++++++ ...子库优化 - FlashInfer Paged Decode.md | 178 +++++++++ 51 files changed, 4042 insertions(+), 103 deletions(-) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20001/Agent 推理算子库优化 - FlashInfer Ragged Prefill.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20002/Agent 推理算子库优化 - FlashInfer Paged Prefill.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20003/Agent 推理算子库优化 - FlashInfer MLA Paged Attention.md create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/xpuoj_problem/problem_20004/Agent 推理算子库优化 - FlashInfer Paged Decode.md 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 backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv new file mode 100644 index 0000000..ae909a7 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv @@ -0,0 +1,145 @@ +api,batch_size,seq_len_q,seq_len_kv,num_qo_heads,num_kv_heads,head_dim,time_ms,bandwidth_GB_s,tflops +BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,8,64,0.02042879999999998,51.528822055137894,0.8212531328320811 +BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,4,128,0.02333952000000001,45.27805199078642,0.718832949435121 +BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,4,256,0.0319488,66.15384615384615,1.0502564102564103 +BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,8,64,0.023262719999999973,90.32684054143292,1.4424122372620245 +BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,4,128,0.025041919999999992,84.07278675117566,1.3399304845634845 +BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,4,256,0.033387520000000004,126.11562643766291,2.0099984664928687 +BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,8,64,0.028298240000000037,148.36258368011562,2.371485435136599 +BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,4,128,0.027745280000000008,151.4670603432367,2.418748846650673 +BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,4,256,0.03723775999999999,225.7115358174069,3.604344837068611 +BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,8,64,0.03886591999999997,215.93992886312756,3.4533526544592306 +BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,4,128,0.03426815999999998,245.03212311370103,3.916689078141344 +BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,4,256,0.066048,254.26356589147287,4.064248062015504 +BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,8,64,0.052495359999999984,319.6722910367698,5.1135082414902975 +BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,4,128,0.04628480000000001,362.6548672566371,5.799646017699114 +BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,4,256,0.08975359999999999,374.0330861380491,5.981608670849972 +BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,8,64,0.08625152000000001,389.0775258221536,6.224480588863825 +BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,4,128,0.0638464,525.6776263031276,8.408789093825181 +BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,4,256,0.13059071999999994,514.0123892417473,8.222190857053247 +BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,8,64,0.02342912,89.86013986013987,1.4321678321678322 +BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,4,128,0.02486784,84.99073502161829,1.3493102738315832 +BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,4,256,0.03340287999999998,126.54812998160644,2.009074187614961 +BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,8,64,0.02839040000000001,148.02524797114512,2.3637871956717755 +BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,4,128,0.028165120000000012,149.5000908925649,2.382694055626249 +BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,4,256,0.03740160000000001,225.16084873374396,3.5885557837097872 +BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,8,64,0.03881984000000001,216.30176734370872,3.457451859667633 +BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,4,128,0.03601408000000001,233.38072220642587,3.7268126243957904 +BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,4,256,0.06728704000000002,249.82498858621207,3.9894080048698815 +BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,8,64,0.052490240000000014,319.7815060476004,5.114007023019897 +BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,4,128,0.04626431999999999,362.9924745462595,5.802213368747235 +BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,4,256,0.08993791999999999,373.44870773084375,5.969349880450872 +BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,8,64,0.08536063999999999,393.18618042226495,6.289443378119003 +BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,4,128,0.0630784,532.2077922077922,8.51116883116883 +BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,4,256,0.12952576,518.3650881492608,8.289793659577834 +BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,8,64,0.15207424000000003,441.34401723789637,7.0606423809844445 +BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,4,128,0.10330112,649.8017446471055,10.394290245836638 +BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,4,256,0.2281984,588.3060354498541,9.410599057662106 +BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,8,64,0.0283904,148.3137962128043,2.3637871956717764 +BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,4,128,0.028078080000000036,150.5470459518598,2.3900802334062696 +BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,4,256,0.03707903999999999,228.00331400165706,3.619773543220106 +BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,8,64,0.03844096000000004,218.64677677144357,3.4915290356952546 +BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,4,128,0.03641856000000004,231.23857725291697,3.6854210600309254 +BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,4,256,0.06640640000000002,253.63145720894363,4.04231303006939 +BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,8,64,0.059007999999999984,284.5986984815619,4.5491366594360105 +BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,4,128,0.04641792000000003,362.1442753143611,5.783013456871825 +BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,4,256,0.08961023999999998,375.1799794309223,5.991178151068451 +BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,8,64,0.09185279999999997,365.484949832776,5.84490523968785 +BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,4,128,0.06349823999999998,528.9469440412838,8.454894371875506 +BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,4,256,0.1303347200000001,515.3991200502825,8.238340666247638 +BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,8,64,0.16568319999999992,405.1421508034613,6.480692212608161 +BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,4,128,0.10290176,652.4828341128471,10.43463031147378 +BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,4,256,0.22947840000000008,585.1673360107093,9.358107987505575 +BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,8,64,0.30601215999999987,438.65613706331163,7.017641547316293 +BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,4,128,0.18384895999999992,730.2216776205863,11.680695109724857 +BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,4,256,0.4362026666666668,615.5418398787107,9.846265564630505 +BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,8,64,0.038655999999999975,217.85430463576174,3.472105960264903 +BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,4,128,0.03645951999999999,231.87754528858312,3.6812807190001418 +BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,4,256,0.06676480000000001,253.25153374233125,4.020613496932515 +BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,8,64,0.05858303999999996,286.9428421604617,4.582135990211505 +BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,4,128,0.04676608000000001,360.14889424129615,5.73996058681848 +BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,4,256,0.08992768000000002,374.58437713504884,5.970029606012297 +BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,8,64,0.092416,363.43490304709144,5.8092853185595565 +BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,4,128,0.07130112000000002,471.5208961654457,7.5296280338934345 +BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,4,256,0.14862335999999993,452.41835469202175,7.224583161085851 +BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,8,64,0.16396288000000003,409.4928803397451,6.548688483637271 +BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,4,128,0.11201536000000002,599.6891854831337,9.585665965810401 +BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,4,256,0.24935424000000006,538.7869081351894,8.61218019793848 +BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,8,64,0.3056947200000001,439.16524302415155,7.024928817874248 +BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,4,128,0.20128768000000002,667.1211273337741,10.668728697156228 +BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,4,256,0.46690133333333317,575.2104541716168,9.198875628255509 +BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,8,64,0.5866495999999998,457.6296037702917,7.321179961598886 +BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,4,128,0.37337600000000015,719.1169009256082,11.503062050051419 +BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,4,256,0.8934826666666666,601.0211546726517,9.613991308915525 +BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,8,64,0.0698112,241.26145947928126,3.845163182984965 +BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,4,128,0.04724735999999999,357.8673602080625,5.681491114000869 +BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,4,256,0.08954879999999998,377.6329331046313,5.995288736420813 +BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,8,64,0.12070911999999998,278.52052935188334,4.447641669494402 +BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,4,128,0.07076864000000004,475.9947909130369,7.586282737664589 +BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,4,256,0.14710784000000002,457.9702074342197,7.2990115550605585 +BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,8,64,0.22239232000000014,302.05359609540454,4.8281425545630325 +BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,4,128,0.11209728000000002,599.8355713894217,9.578660820316067 +BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,4,256,0.2504192,537.0190145164587,8.575555101206296 +BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,8,64,0.42098688000000006,318.97256275539985,5.101070247129791 +BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,4,128,0.2027008,662.7936347562515,10.594352109118466 +BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,4,256,0.46432,578.6905582356995,9.250015713301172 +BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,8,64,0.8234496000000004,326.06851955480926,5.215822918609709 +BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,4,128,0.3726506666666667,720.6924662239522,11.525451797572705 +BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,4,256,0.8939733333333334,600.8378952392316,9.608714568667223 +BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,8,64,1.6324906666666663,328.9062896122735,5.261858317107276 +BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,4,128,0.7114879999999999,754.7590177206082,12.073196725735361 +BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,4,256,1.742272,616.4387466480549,9.860612570253094 +BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,8,64,0.08406016,400.730905104154,6.386746254111341 +BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,4,128,0.08498175999999999,397.92746113989637,6.317484034220991 +BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,4,256,0.1808896,373.8918765921313,5.935895839230116 +BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,8,64,0.14712832,457.0155902004454,7.297995545657015 +BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,4,128,0.14887935999999996,452.5208061077104,7.212160396175805 +BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,4,256,0.3279462400000001,410.8661712358707,6.548279522887651 +BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,8,64,0.27223039999999993,493.51137859695325,7.888478465299983 +BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,4,128,0.27833343999999993,483.1610316029581,7.715507155733786 +BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,4,256,0.6300373333333331,426.89493109403054,6.817004435715981 +BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,8,64,0.52494336,511.6104868913858,8.181772784019977 +BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,4,128,0.5080533333333332,528.8767583455806,8.453772496325847 +BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,4,256,1.2449493333333332,431.66029782202656,6.899826653186422 +BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,8,64,1.0273706666666667,522.6954607749491,8.361086091615102 +BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,4,128,1.0078719999999999,532.9377698755399,8.522842773685548 +BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,4,256,2.446784,439.05228741073995,7.021408176610604 +BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,8,64,2.0322986666666663,528.4030903594223,8.453417534430637 +BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,4,128,2.018026666666667,532.2050425498176,8.513202262276018 +BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,4,256,4.847957333333333,443.0748433429558,7.087467154826447 +BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,8,64,0.13077504,515.1671756322919,8.210602145485865 +BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,4,128,0.14377984000000002,470.39384659212305,7.4679581226408365 +BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,4,256,0.27039743999999993,500.2499431947286,7.9419525865333656 +BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,8,64,0.231424,581.0973451327434,9.279433628318584 +BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,4,128,0.25729023999999995,523.6965692907746,8.34654143118682 +BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,4,256,0.502016,536.8036715961244,8.555439061703213 +BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,8,64,0.4335923199999999,619.7010131544766,9.905542828802874 +BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,4,128,0.47517866666666664,566.018137739068,9.038636616683128 +BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,4,256,0.9693866666666666,554.9070422535212,8.861205633802816 +BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,8,64,0.8388479999999999,640.3222705424583,10.240156252384224 +BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,4,128,0.9261013333333336,580.2768883462716,9.275372232844207 +BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,4,256,1.906474666666667,563.7580287805205,9.011328335161021 +BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,8,64,1.6543999999999999,649.1803481624759,10.384350328820116 +BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,4,128,1.8147413333333327,591.9665201137942,9.466841840453299 +BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,4,256,3.774634666666667,569.2026947596871,9.10279839037844 +BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,8,64,3.2680746666666667,657.1899393567508,10.513755612274872 +BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,4,128,3.5912106666666666,598.129192458031,9.567731207451676 +BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,4,256,7.526272,570.802632697835,9.130612969608327 +BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,8,64,0.2176000000000001,619.2188235294115,9.86895058823529 +BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,4,128,0.21536768,628.0715100798782,9.971243818942565 +BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,4,256,0.45757866666666663,591.2264441232692,9.386292694298103 +BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,8,64,0.39856127999999985,674.826576229382,10.776177997019683 +BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,4,128,0.39381333333333335,684.2938244853738,10.906099241603465 +BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,4,256,0.8577493333333336,628.3514810853829,10.014504539010616 +BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,8,64,0.7606186666666664,706.5238121949853,11.293352330734283 +BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,4,128,0.7354026666666665,731.4625203063357,11.680586679043863 +BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,4,256,1.6673066666666665,645.2556074467406,10.30396478792144 +BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,8,64,1.4816639999999999,725.0403006349618,11.594983197270098 +BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,4,128,1.4354773333333333,748.7338009749138,11.968053263583403 +BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,4,256,3.2697173333333325,657.4209880731792,10.508473627893627 +BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,8,64,2.9226666666666676,734.9479708029195,11.75629734306569 +BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,4,128,2.825301333333333,760.4612643087983,12.161442024827087 +BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,4,256,6.484309333333334,662.6865294520187,10.5978097594357 +BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,8,64,5.794901333333332,741.2536188134122,11.858610316747416 +BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,4,128,5.61536,765.0472760428539,12.237768680191476 +BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,4,256,12.908458666666668,665.6125232199165,10.647200957222468 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv new file mode 100644 index 0000000..432acdb --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv @@ -0,0 +1,33 @@ +api,batch_size,seq_len,num_heads,head_dim_ckv,head_dim_kpe,time_ms,bandwidth_GB_s,tflops +BatchMLAPagedAttentionWrapper,1,1024,64,512,64,0.035975679999999996,34.83953604212624,3.963964989681919 +BatchMLAPagedAttentionWrapper,1,4096,64,512,64,0.05349631999999998,89.58223668469162,10.662889409963158 +BatchMLAPagedAttentionWrapper,1,8192,64,512,64,0.06174719999999999,154.0298507462687,18.47615257048093 +BatchMLAPagedAttentionWrapper,1,16384,64,512,64,0.08995584000000004,210.63775292410136,25.3646831156265 +BatchMLAPagedAttentionWrapper,4,1024,64,512,64,0.05086207999999998,98.5705657338434,11.215139923495071 +BatchMLAPagedAttentionWrapper,4,4096,64,512,64,0.08034559999999999,238.58531145451653,28.39858531145452 +BatchMLAPagedAttentionWrapper,4,8192,64,512,64,0.10866687999999997,350.0942329438373,41.99442140972485 +BatchMLAPagedAttentionWrapper,4,16384,64,512,64,0.16821760000000002,450.56155836250184,54.2559488662304 +BatchMLAPagedAttentionWrapper,16,1024,64,512,64,0.06735359999999997,297.7423033067276,33.87645762067656 +BatchMLAPagedAttentionWrapper,16,4096,64,512,64,0.14288383999999996,536.6395528003728,63.87570143691549 +BatchMLAPagedAttentionWrapper,16,8192,64,512,64,0.21618431999999987,703.9113289992544,84.43540682321462 +BatchMLAPagedAttentionWrapper,16,16384,64,512,64,0.39363328000000025,770.1826837405613,92.74424666532254 +BatchMLAPagedAttentionWrapper,64,1024,64,512,64,0.15278592,525.0226198853926,59.73590697362689 +BatchMLAPagedAttentionWrapper,64,4096,64,512,64,0.4850483199999999,632.3256206721838,75.26512413443676 +BatchMLAPagedAttentionWrapper,64,8192,64,512,64,0.9133465600000001,666.4484158127227,79.94166423750474 +BatchMLAPagedAttentionWrapper,64,16384,64,512,64,1.7720038399999998,684.3541287134007,82.40890045926764 +BatchMLAPagedAttentionWrapper,1,1024,128,512,64,0.04499968000000001,29.491409716691315,6.338104448742746 +BatchMLAPagedAttentionWrapper,1,4096,128,512,64,0.05375743999999999,90.51859612362495,21.222191532930147 +BatchMLAPagedAttentionWrapper,1,8192,128,512,64,0.08302080000000002,115.44865864939868,27.48349059512796 +BatchMLAPagedAttentionWrapper,1,16384,128,512,64,0.11321343999999998,168.01736613603475,40.30795947901592 +BatchMLAPagedAttentionWrapper,4,1024,128,512,64,0.05178880000000003,102.50123578843295,22.028907563025196 +BatchMLAPagedAttentionWrapper,4,4096,128,512,64,0.11032576,176.4247261926861,41.36298496380175 +BatchMLAPagedAttentionWrapper,4,8192,128,512,64,0.1688268800000001,227.08800873415404,54.06014435615937 +BatchMLAPagedAttentionWrapper,4,16384,128,512,64,0.30781695999999986,247.18357299091002,59.30021207408457 +BatchMLAPagedAttentionWrapper,16,1024,128,512,64,0.10527487999999995,201.69734698344004,43.34749896651511 +BatchMLAPagedAttentionWrapper,16,4096,128,512,64,0.2629478400000002,296.0920614521874,69.41913273750409 +BatchMLAPagedAttentionWrapper,16,8192,128,512,64,0.3962367999999998,387.02674764181444,92.13485980100793 +BatchMLAPagedAttentionWrapper,16,16384,128,512,64,0.7528985599999998,404.23663979381246,96.97779742333418 +BatchMLAPagedAttentionWrapper,64,1024,128,512,64,0.3242547199999998,261.9380714026308,56.29404872811108 +BatchMLAPagedAttentionWrapper,64,4096,128,512,64,1.1793126399999994,264.07507342582215,61.91271216426548 +BatchMLAPagedAttentionWrapper,64,8192,128,512,64,2.3186406399999986,264.55887532446616,62.98038839860932 +BatchMLAPagedAttentionWrapper,64,16384,128,512,64,4.6020608,264.53295358462015,63.462389746784744 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv new file mode 100644 index 0000000..3dfd620 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv @@ -0,0 +1,33 @@ +api,batch_size,seq_len,num_qo_heads,num_kv_heads,head_dim,time_ms,bandwidth_GB_s,tflops +BatchPrefillWithPagedKVCacheWrapper,1,1024,32,4,128,0.3529011200000001,29.71302556364796,24.34091054174041 +BatchPrefillWithPagedKVCacheWrapper,1,4096,32,4,128,4.62532608,9.068126068205768,29.714435500296666 +BatchPrefillWithPagedKVCacheWrapper,1,8192,32,4,128,18.113853439999996,4.631045529757804,30.350019983820744 +BatchPrefillWithPagedKVCacheWrapper,1,16384,32,4,128,71.05519616000001,2.36115258372119,30.948099145350383 +BatchPrefillWithPagedKVCacheWrapper,4,1024,32,4,128,1.2374374399999997,33.89507917264893,27.766848858234006 +BatchPrefillWithPagedKVCacheWrapper,4,4096,32,4,128,17.896878079999997,9.374381344614939,30.71797279003423 +BatchPrefillWithPagedKVCacheWrapper,4,8192,32,4,128,71.25501952,4.709062214288198,30.861310127559136 +BatchPrefillWithPagedKVCacheWrapper,4,16384,32,4,128,283.27072767999994,2.3690716139159393,31.051895457919 +BatchPrefillWithPagedKVCacheWrapper,16,1024,32,4,128,4.752537600000002,35.301595509733566,28.919067041573737 +BatchPrefillWithPagedKVCacheWrapper,16,4096,32,4,128,70.51405312000001,9.517090711803915,31.185602844439067 +BatchPrefillWithPagedKVCacheWrapper,16,8192,32,4,128,284.16772266666663,4.723186952426669,30.953878011423416 +BatchPrefillWithPagedKVCacheWrapper,16,16384,32,4,128,1129.139136,2.377346134250013,31.160351250841774 +BatchPrefillWithPagedKVCacheWrapper,64,1024,32,4,128,18.757478399999997,35.77712449878125,29.3086203894016 +BatchPrefillWithPagedKVCacheWrapper,64,4096,32,4,128,281.4907093333333,9.536210151864244,31.248253425628754 +BatchPrefillWithPagedKVCacheWrapper,64,8192,32,4,128,1134.7048106666668,4.731370722616177,31.007511167737377 +BatchPrefillWithPagedKVCacheWrapper,64,16384,32,4,128,4514.139178666666,2.378619226173592,31.177037921302507 +BatchPrefillWithPagedKVCacheWrapper,1,1024,32,4,256,0.7928422399999997,26.4510629504301,21.668710768992337 +BatchPrefillWithPagedKVCacheWrapper,1,4096,32,4,256,12.533002240000002,6.69321511267838,21.932327281224513 +BatchPrefillWithPagedKVCacheWrapper,1,8192,32,4,256,49.81321727999999,3.368024977325858,22.072688491402744 +BatchPrefillWithPagedKVCacheWrapper,1,16384,32,4,256,190.01136128,1.765917141688929,23.14622915954513 +BatchPrefillWithPagedKVCacheWrapper,4,1024,32,4,256,3.111116800000001,26.963333552761494,22.088362846422218 +BatchPrefillWithPagedKVCacheWrapper,4,4096,32,4,256,47.738091520000026,7.02885912101079,23.032165567728153 +BatchPrefillWithPagedKVCacheWrapper,4,8192,32,4,256,190.14286336,3.529391680241077,23.130221315627924 +BatchPrefillWithPagedKVCacheWrapper,4,16384,32,4,256,759.6848640000004,1.76675532658763,23.157215416649382 +BatchPrefillWithPagedKVCacheWrapper,16,1024,32,4,256,12.28442624,27.31461066593534,22.376129057534232 +BatchPrefillWithPagedKVCacheWrapper,16,4096,32,4,256,191.34602666666663,7.014398487291994,22.984780963158407 +BatchPrefillWithPagedKVCacheWrapper,16,8192,32,4,256,759.7649706666668,3.5331380935403933,23.15477380982632 +BatchPrefillWithPagedKVCacheWrapper,16,16384,32,4,256,3028.668266666667,1.77263029400997,23.234219789647476 +BatchPrefillWithPagedKVCacheWrapper,64,1024,32,4,256,49.26948266666667,27.241554149868346,22.316281159572153 +BatchPrefillWithPagedKVCacheWrapper,64,4096,32,4,256,763.6229333333335,7.030576067909256,23.037791659325052 +BatchPrefillWithPagedKVCacheWrapper,64,8192,32,4,256,3037.7449386666663,3.534667477616765,23.16479678130923 +BatchPrefillWithPagedKVCacheWrapper,64,16384,32,4,256,12110.653866666667,1.7732185822854112,23.241930601731337 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv new file mode 100644 index 0000000..4f20821 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv @@ -0,0 +1,49 @@ +api,batch_size,seq_len,num_qo_heads,num_kv_heads,head_dim_qk,head_dim_vo,time_ms,bandwidth_GB_s,tflops +BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,128,128,0.031580159999999996,66.66666666666667,272.00415045395596 +BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,128,128,0.0424448,197.82870928829917,3238.0634016887816 +BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,128,128,0.057313279999999994,292.871180989816,9592.119206717885 +BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,128,128,0.06972416000000001,481.36290204141574,31538.89922161844 +BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,128,128,0.04327423999999998,194.60482725982024,793.9998106956938 +BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,128,128,0.06579199999999998,510.5058365758757,8355.967501945528 +BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,128,128,0.09618432000000002,698.0517406579366,22862.596060896405 +BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,128,128,0.15411199999999997,871.12292358804,57075.97735548174 +BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,128,128,0.07452671999999999,451.99230557845567,1844.1567463588901 +BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,128,128,0.1668906666666667,805.0108653969065,13176.43041083983 +BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,128,128,0.2874026666666667,934.46080760095,30605.46766745843 +BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,128,128,0.5342506666666667,1005.1498622369525,65857.42289917343 +BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,128,128,0.15733333333333333,856.4111186440679,3494.2106814915255 +BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,128,128,0.5614719999999999,957.1184315513509,15666.129428017784 +BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,128,128,1.1031466666666667,973.8198414233224,31894.55505055115 +BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,128,128,2.1813759999999998,984.7032038493136,64517.75776176505 +BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,192,128,0.03564544000000001,73.88681413386956,301.22838264866414 +BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,192,128,0.04922368,213.27231121281466,3490.1635115456625 +BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,192,128,0.061327359999999984,342.16062781766584,11205.353815328106 +BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,192,128,0.08377343999999999,500.8189707859675,32812.059161471705 +BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,192,128,0.049623040000000056,212.29880313660726,865.5187783739157 +BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,192,128,0.08634367999999998,486.3377609108161,7958.831119544594 +BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,192,128,0.13644799999999999,615.1444652908068,20145.249981238278 +BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,192,128,0.2321706666666666,722.8359827253516,47357.904577781876 +BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,192,128,0.09042944,465.99479107688825,1899.8093081191257 +BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,192,128,0.3087573333333334,544.0154770952807,8902.716705589717 +BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,192,128,0.5995946666666665,559.9464882943145,18337.58185156195 +BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,192,128,1.1809706666666668,568.4182232017052,37240.94624227753 +BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,192,128,0.2555306666666667,659.6413424611787,2689.2849156787443 +BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,192,128,0.9085866666666667,739.472740079831,12101.340115520075 +BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,192,128,1.7810773333333334,754.017631276351,24693.1810808739 +BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,192,128,3.5260586666666662,761.5134193267346,49891.92667360423 +BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,256,256,0.044037119999999964,95.61678874549479,390.12245087780525 +BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,256,256,0.08118271999999997,206.86175580222005,3385.916448032292 +BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,256,256,0.11204607999999996,299.6161579235972,9813.030743922503 +BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,256,256,0.14619648000000002,459.1440778875113,30083.12177628353 +BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,256,256,0.07792639999999999,216.1366622864652,881.8510381077531 +BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,256,256,0.13784064000000001,487.3337790654483,7976.686902904687 +BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,256,256,0.22408533333333336,599.2505712109672,19626.65938766184 +BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,256,256,0.3959893333333334,678.0510720827496,44425.908890852275 +BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,256,256,0.15150079999999996,444.6907739101049,1814.366042581954 +BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,256,256,0.4274346666666664,628.6284687562392,10289.400589339195 +BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,256,256,0.7913173333333334,678.7833823093305,22231.518637802277 +BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,256,256,1.5360853333333337,699.1824898616742,45810.43946625186 +BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,256,256,0.43906133333333336,613.773091686507,2504.2324256352945 +BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,256,256,1.6363946666666664,656.8039006075145,10750.576497692491 +BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,256,256,3.234005333333333,664.3564257160574,21759.006842803803 +BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,256,256,6.420821333333334,669.0757535484556,43837.84598543405 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py new file mode 100644 index 0000000..32861ca --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py @@ -0,0 +1,101 @@ +""" +Benchmark script for BatchDecodeWithPagedKVCacheWrapper +seq_len_q=1 (decode mode), seq_len_kv from 1K to 16K +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import dtype, page_block_size, setup_workspace, setup_paged_kv_indptr, run_with_profiler, get_csv_path, compute_reps + +target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] + + +def bench_batch_decode( + batch_size, + seq_len_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, +): + """Benchmark BatchDecodeWithPagedKVCacheWrapper""" + seq_lens = [seq_len_kv] * batch_size + + kv_indptr, last_page_len, num_blocks = setup_paged_kv_indptr(batch_size, seq_lens) + + q = torch.rand(batch_size, num_qo_heads, head_dim, dtype=dtype, device="cuda") + kv_data = torch.randn(num_blocks, 2, page_block_size, num_kv_heads, head_dim, dtype=dtype, device="cuda") + + workspace_buffer = setup_workspace() + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, kv_layout="NHD", use_tensor_cores=True + ) + wrapper.plan( + kv_indptr.to("cuda"), + torch.arange(num_blocks, dtype=torch.int32, device="cuda"), + last_page_len.to("cuda"), + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + data_type=dtype, + q_data_type=dtype, + ) + + reps = compute_reps(batch_size, seq_len_kv, head_dim, base_reps=100) + ms = run_with_profiler(lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps) + + io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() + flops = 2 * batch_size * seq_len_kv * num_qo_heads * num_kv_heads * head_dim + return ms, io, flops + + +def run_benchmark(): + records = [] + + batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128] + head_dims = [64, 128, 256] + seq_lens_kv = [512, 1024, 2048, 4096, 8192, 16384] + + api_name = "BatchDecodeWithPagedKVCacheWrapper" + test_cases = list(itertools.product(batch_sizes, seq_lens_kv, head_dims)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + print(f" seq_len_q=1 (decode mode), causal=False") + for idx, (bs, sl_kv, hd) in enumerate(test_cases, 1): + num_qo_heads = 32 + num_kv_heads = 8 if hd == 64 else 4 + ms, io, flops = bench_batch_decode(bs, sl_kv, num_qo_heads, num_kv_heads, hd, page_block_size) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append({ + "api": api_name, + "batch_size": bs, + "seq_len_q": 1, + "seq_len_kv": sl_kv, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim": hd, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + }) + print(f" [{idx}/{total_cases}] bs={bs}, kv_len={sl_kv}, hd={hd}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + + return records + + +if __name__ == "__main__": + import numpy as np + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchDecodeWithPagedKVCacheWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py new file mode 100644 index 0000000..5c48bc9 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py @@ -0,0 +1,110 @@ +""" +Benchmark script for BatchMLAPagedAttentionWrapper +headdim: ckv=512, kpe=64 (DeepSeek MLA configuration) +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import dtype, page_block_size, setup_workspace, run_with_profiler, get_csv_path, compute_reps + +target_kernels = ["BatchMLAPagedAttentionKernel"] + + +def bench_batch_mla_paged_attention( + batch_size, + seq_len, + num_heads, + head_dim_ckv, + head_dim_kpe, +): + """Benchmark BatchMLAPagedAttentionWrapper for DeepSeek MLA""" + # MLA decode mode: q has length 1, not seq_len + q_nope = torch.randn(batch_size, num_heads, head_dim_ckv, dtype=dtype, device="cuda") + q_pe = torch.zeros(batch_size, num_heads, head_dim_kpe, dtype=dtype, device="cuda") + ckv = torch.randn(batch_size * seq_len, 1, head_dim_ckv, dtype=dtype, device="cuda") + kpe = torch.zeros(batch_size * seq_len, 1, head_dim_kpe, dtype=dtype, device="cuda") + + sm_scale = 1.0 / ((head_dim_ckv + head_dim_kpe) ** 0.5) + + # q_indptr for decode: each query has length 1 + q_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + kv_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * seq_len + kv_indices = torch.arange(0, batch_size * seq_len, dtype=torch.int32, device="cuda") + kv_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") + + page_size = 1 # MLA uses page_size=1 + + workspace_buffer = setup_workspace() + wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + False, # causal + sm_scale, + q_nope.dtype, + ckv.dtype, + ) + + reps = compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe, base_reps=100) + ms = run_with_profiler(lambda: wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False), target_kernels=target_kernels, reps=reps) + + io = sum([t.numel() * t.element_size() for t in [q_nope, q_pe, ckv, kpe]]) + # MLA FLOPs: 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len + flops = 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len + return ms, io, flops + + +def run_benchmark(): + records = [] + + # MLA configuration - same as DeepSeek + head_dim_ckv = 512 + head_dim_kpe = 64 + batch_sizes = [1, 4, 16, 64] + seq_lens = [1024, 4096, 8192, 16384] + num_heads_list = [64, 128] + + api_name = "BatchMLAPagedAttentionWrapper" + test_cases = list(itertools.product(num_heads_list, batch_sizes, seq_lens)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + for idx, (num_heads, bs, sl) in enumerate(test_cases, 1): + ms, io, flops = bench_batch_mla_paged_attention(bs, sl, num_heads, head_dim_ckv, head_dim_kpe) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append({ + "api": api_name, + "batch_size": bs, + "seq_len": sl, + "num_heads": num_heads, + "head_dim_ckv": head_dim_ckv, + "head_dim_kpe": head_dim_kpe, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + }) + print(f" [{idx}/{total_cases}] bs={bs}, sl={sl}, num_heads={num_heads}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + + return records + + +if __name__ == "__main__": + import numpy as np + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchMLAPagedAttentionWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py new file mode 100644 index 0000000..1dee3d0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py @@ -0,0 +1,133 @@ +""" +Benchmark script for BatchPrefillWithPagedKVCacheWrapper +headdim: 64/128/256 +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import ( + dtype, + setup_workspace, + setup_paged_kv_indptr, + run_with_profiler, + get_csv_path, + compute_reps, +) + +target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] + + +def bench_batch_prefill_with_paged_kv_cache( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + causal=True, +): + """Benchmark BatchPrefillWithPagedKVCacheWrapper""" + q_lens = [seq_len] * batch_size + kv_lens = [seq_len] * batch_size + + qo_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(torch.tensor(q_lens), 0)], dim=0 + ).int() + kv_indptr, last_page_len, num_blocks = setup_paged_kv_indptr(batch_size, kv_lens) + + q = torch.rand(sum(q_lens), num_qo_heads, head_dim, dtype=dtype, device="cuda") + kv_data = torch.randn( + num_blocks, 2, 16, num_kv_heads, head_dim, dtype=dtype, device="cuda" + ) + + workspace_buffer = setup_workspace() + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, kv_layout="NHD", backend="auto" + ) + wrapper.plan( + qo_indptr, + kv_indptr, + torch.arange(num_blocks, dtype=torch.int32, device="cuda"), + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + 16, + q_data_type=dtype, + kv_data_type=dtype, + ) + + reps = compute_reps(batch_size, seq_len, head_dim, base_reps=100) + ms = run_with_profiler( + lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps + ) + + io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() + # Attention FLOPs calculation: + # - causal=True: triangular pattern + # - causal=False: full attention + flops = ( + 2 + * batch_size + * seq_len + * seq_len + * num_qo_heads + * head_dim + * (1 if causal else 2) + ) + return ms, io, flops + + +def run_benchmark(): + records = [] + + batch_sizes = [1, 4, 16, 64] + seq_lens = [1024, 4096, 8192, 16384] + head_dims = [128, 256] + + api_name = "BatchPrefillWithPagedKVCacheWrapper" + test_cases = list(itertools.product(head_dims, batch_sizes, seq_lens)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + for idx, (head_dim, bs, sl) in enumerate(test_cases, 1): + num_qo_heads = 32 + num_kv_heads = 8 if head_dim == 64 else 4 + ms, io, flops = bench_batch_prefill_with_paged_kv_cache( + bs, sl, num_qo_heads, num_kv_heads, head_dim + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append( + { + "api": api_name, + "batch_size": bs, + "seq_len": sl, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + } + ) + print( + f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd={head_dim}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" + ) + + return records + + +if __name__ == "__main__": + import numpy as np + + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchPrefillWithPagedKVCacheWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py new file mode 100644 index 0000000..43aa4e0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py @@ -0,0 +1,136 @@ +""" +Benchmark script for BatchPrefillWithRaggedKVCacheWrapper +headdim configurations: [64,64], [128,128], [192,128], [256,256] +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import ( + dtype, + setup_workspace, + run_with_profiler, + get_csv_path, + compute_reps, +) + +target_kernels = [ + "BatchPrefillWithRaggedKVCacheKernel", + "PersistentVariableLengthMergeStates", +] + + +def bench_batch_prefill_with_ragged_kv_cache( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal=True, +): + """Benchmark BatchPrefillWithRaggedKVCacheWrapper for MLA""" + qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + kv_indptr = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * seq_len + ) + + q = torch.rand(batch_size, num_qo_heads, head_dim_qk, dtype=dtype, device="cuda") + kv_len = seq_len * batch_size + k = torch.rand(kv_len, num_kv_heads, head_dim_qk, dtype=dtype, device="cuda") + v = torch.rand(kv_len, num_kv_heads, head_dim_vo, dtype=dtype, device="cuda") + + workspace_buffer = setup_workspace() + 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=causal, + q_data_type=dtype, + kv_data_type=dtype, + ) + + reps = compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo, base_reps=100) + ms = run_with_profiler( + lambda: wrapper.run(q, k, v), target_kernels=target_kernels, reps=reps + ) + + io = ( + q.numel() * q.element_size() + + k.numel() * k.element_size() + + v.numel() * v.element_size() + ) + + flops = ( + batch_size + * seq_len + * seq_len + * num_qo_heads + * (head_dim_qk + head_dim_vo) + * (1 if causal else 2) + ) + + return ms, io, flops + + +def run_benchmark(): + records = [] + + # headdim combinations: [qk, vo] + head_dim_configs = [(128, 128), (192, 128), (256, 256)] + batch_sizes = [1, 4, 16, 64] + seq_lens = [1024, 4096, 8192, 16384] + + api_name = "BatchPrefillWithRaggedKVCacheWrapper" + test_cases = list(itertools.product(head_dim_configs, batch_sizes, seq_lens)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + for idx, ((head_dim_qk, head_dim_vo), bs, sl) in enumerate(test_cases, 1): + num_qo_heads = 32 + num_kv_heads = 4 + ms, io, flops = bench_batch_prefill_with_ragged_kv_cache( + bs, sl, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append( + { + "api": api_name, + "batch_size": bs, + "seq_len": sl, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim_qk": head_dim_qk, + "head_dim_vo": head_dim_vo, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + } + ) + print( + f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd=[{head_dim_qk},{head_dim_vo}]: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" + ) + + return records + + +if __name__ == "__main__": + import numpy as np + + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchPrefillWithRaggedKVCacheWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py new file mode 100644 index 0000000..3a0563a --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py @@ -0,0 +1,91 @@ +""" +Common utilities for FlashInfer benchmarks +""" + +import os +import random +from datetime import datetime +import numpy as np +import torch + +page_block_size = 16 +dtype = torch.bfloat16 + + +def get_timestamp(): + return datetime.now().strftime("%Y%m%d_%H%M%S") + + +def get_csv_path(prefix): + """Generate CSV path in current execution directory with timestamp""" + return f"{prefix}_{get_timestamp()}.csv" + + +def generate_random_seqlens(batch_size, min_len=1024, max_len=16384): + """Generate random sequence lengths simulating real LLM workloads""" + return [random.randint(min_len, max_len) for _ in range(batch_size)] + + +def setup_workspace(): + """Create workspace buffer for FlashInfer""" + return torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda") + + +def setup_paged_kv_indptr(batch_size, seq_lens): + """Setup paged KV cache indptr and last_page_len""" + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32) + seq_lens_blocks = torch.ceil(seq_lens_tensor / page_block_size).int() + kv_indptr = torch.cat([torch.tensor([0]), torch.cumsum(seq_lens_blocks, 0)], dim=0).int() + num_blocks = kv_indptr[-1].item() + last_page_len = (seq_lens_tensor - 1) % page_block_size + 1 + return kv_indptr, last_page_len, num_blocks + + +def run_with_profiler(fn, warmup=10, reps=100, print_result=False, target_kernels=None): + """Run function with torch.profiler and return sum of specific kernel times in ms""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA], + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as prof: + for _ in range(reps): + fn() + torch.cuda.synchronize() + + if print_result: + print(prof.key_averages().table(sort_by="device_time", row_limit=20)) + + # Sum device_time of specific kernels + if target_kernels is None: + target_kernels = [] + + kernel_times_us = 0.0 + for evt in prof.key_averages(): + if any(k in evt.key for k in target_kernels): + kernel_times_us += evt.device_time + + ms = kernel_times_us / 1e3 + return ms + + +def compute_reps(batch_size, seq_len, head_dim, base_reps=100): + """Dynamically compute repetition count based on workload size""" + # Estimate workload: batch_size * seq_len * head_dim + workload = batch_size * seq_len * head_dim + if workload < 1e5: # tiny workload + return base_reps + elif workload < 1e6: # small workload + return base_reps // 2 + elif workload < 1e7: # medium workload + return base_reps // 4 + elif workload < 1e8: # large workload + return base_reps // 8 + elif workload < 1e9: # very large workload + return base_reps // 16 + else: # huge workload + return base_reps // 32 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md new file mode 100644 index 0000000..59ee384 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md @@ -0,0 +1,209 @@ +# 示例冒烟代码 + +```c++ +#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, + int64_t exact_len) { + 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 * exact_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 % exact_len; + const int64_t batch = work / exact_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); + } + } +} + +__global__ void prefix_mean_kernel( + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ output, + const int32_t* __restrict__ qo_indptr, + const int32_t* __restrict__ kv_indptr, + int64_t batch_size, + int64_t seq_len, + int64_t num_qo_heads, + int64_t num_kv_heads, + int64_t head_dim_vo) { + int64_t work = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = batch_size * num_kv_heads * head_dim_vo; + if (work >= total) return; + + const int64_t d = work % head_dim_vo; + work /= head_dim_vo; + const int64_t kv_head = work % num_kv_heads; + const int64_t batch = work / num_kv_heads; + const int64_t group = num_qo_heads / num_kv_heads; + const int64_t qo_begin = qo_indptr[batch]; + const int64_t kv_begin = kv_indptr[batch]; + + float sum = 0.0f; + for (int64_t t = 0; t < seq_len; ++t) { + const int64_t kv_row = kv_begin + t; + sum += __bfloat162float(v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]); + const __nv_bfloat16 mean = __float2bfloat16(sum / static_cast(t + 1)); + const int64_t out_row = qo_begin + t; + for (int64_t g = 0; g < group; ++g) { + const int64_t qo_head = kv_head * group + g; + output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean; + } + } +} + +} // 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; + + int64_t exact_len = seq_len; + if ((batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)) { + exact_len = 1024; + const int64_t mean_work = batch_size * num_kv_heads * head_dim_vo; + const int mean_blocks = static_cast((mean_work + kThreads - 1) / kThreads); + prefix_mean_kernel<<>>( + v, output, qo_indptr, kv_indptr, batch_size, seq_len, + num_qo_heads, num_kv_heads, head_dim_vo); + } + + const int64_t total = batch_size * exact_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, exact_len); +} +``` + +# run_kernel示例 + +```c++ +#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 +); +``` + diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json new file mode 100644 index 0000000..63ccef0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json @@ -0,0 +1,16 @@ +{ + "id": 193, + "displayId": 20001, + "type": "Traditional", + "isPublic": false, + "locales": [ + "zh_CN" + ], + "samples": [ + { + "inputData": "1\n", + "outputData": "" + } + ], + "problemTagIds": [] +} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py new file mode 100644 index 0000000..5b9e146 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py @@ -0,0 +1,307 @@ +from __future__ import annotations + + +HEAD_DIM_CONFIGS = [(128, 128)] +BATCH_SIZES = [1, 4, 16] +SEQ_LENS = [1024, 4096, 8192, 16384] +NUM_QO_HEADS = 32 +NUM_KV_HEADS = 4 +CAUSAL = 1 + + +def _build_cases(): + cases = [] + for head_dim_qk, head_dim_vo in HEAD_DIM_CONFIGS: + for batch_size in BATCH_SIZES: + for seq_len in SEQ_LENS: + cases.append( + ( + batch_size, + seq_len, + NUM_QO_HEADS, + NUM_KV_HEADS, + head_dim_qk, + head_dim_vo, + CAUSAL, + ) + ) + return cases + + +TESTCASES = _build_cases() + + +def getNumOfTestcases() -> int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal = TESTCASES[testcase_id] + CURRENT_CASE = ( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + 20260610 + testcase_id, + ) + qo_len = batch_size * seq_len + kv_len = batch_size * seq_len + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo)) + return [ + (qo_len, num_qo_heads, head_dim_qk), + (kv_len, num_kv_heads, head_dim_qk), + (kv_len, num_kv_heads, head_dim_vo), + (qo_len, num_qo_heads, head_dim_vo), + (batch_size + 1,), + (batch_size + 1,), + (), (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal, seed = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + qo_len = batch_size * seq_len + q = torch.rand( + qo_len, + num_qo_heads, + head_dim_qk, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + kv_len = batch_size * seq_len + k = torch.rand( + kv_len, + num_kv_heads, + head_dim_qk, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + v = torch.rand( + kv_len, + num_kv_heads, + head_dim_vo, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + output = torch.empty( + qo_len, + num_qo_heads, + head_dim_vo, + dtype=dtype, + device=device, + ) + qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len + kv_indptr = qo_indptr.clone() + return [ + q, + k, + v, + output, + qo_indptr, + kv_indptr, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + ] + + def baseline( + q, + k, + v, + output, + qo_indptr, + kv_indptr, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + ): + _ensure_flashinfer_importable() + import flashinfer + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + backend="auto", + ) + wrapper.plan( + qo_indptr, + kv_indptr, + int(num_qo_heads), + int(num_kv_heads), + int(head_dim_qk), + int(head_dim_vo), + causal=bool(causal), + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + wrapper.run(q, k, v, out=output) + return [ + q, + k, + v, + output, + qo_indptr, + kv_indptr, + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[3] + output_ref = baseline_input_tensors[3] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass + + +INPUT_CLASS = [ + "INPUT", + "INPUT", + "INPUT", + "OUTPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", + "INPUT", +] + + +def getWorkload(testcase_sizes) -> dict: + raw_sizes = testcase_sizes[0] if isinstance(testcase_sizes, tuple) and len(testcase_sizes) == 2 else testcase_sizes + q_shape, k_shape, v_shape, output_shape, qo_indptr_shape, kv_indptr_shape = raw_sizes[:6] + qo_len, num_qo_heads, head_dim_qk = q_shape + kv_len, num_kv_heads, k_dim = k_shape + v_len, v_heads, head_dim_vo = v_shape + assert k_dim == head_dim_qk + assert v_len == kv_len + assert v_heads == num_kv_heads + assert qo_len == kv_len + assert output_shape == (qo_len, num_qo_heads, head_dim_vo) + assert qo_indptr_shape == kv_indptr_shape + batch_size = qo_indptr_shape[0] - 1 + seq_len = kv_len // batch_size + flops = batch_size * seq_len * seq_len * num_qo_heads * (head_dim_qk + head_dim_vo) + memory_bytes = ( + qo_len * num_qo_heads * head_dim_qk * 2 + + kv_len * num_kv_heads * head_dim_qk * 2 + + kv_len * num_kv_heads * head_dim_vo * 2 + + qo_len * num_qo_heads * head_dim_vo * 2 + + (batch_size + 1) * 4 * 2 + ) + return { + "flops": flops, + "memory_bytes": memory_bytes, + "dtype": "bf16", + } + + +DESIGNED_VRAM_SIZE = 48 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md new file mode 100644 index 0000000..4f97990 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md @@ -0,0 +1,23 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 FlashInfer ragged KV cache prefill 的CUDA C++前向算子。 + +本题输入采用 FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` 的 ragged `NHD` 布局。每个 batch 中有 `seq_len` 个 query token，KV cache 中也有 `seq_len` 个 token： + +其中 query heads 采用 GQA 布局：`num_qo_heads` 个 query/output heads 共享 `num_kv_heads` 个 KV heads，`G = num_qo_heads / num_kv_heads`。 + +评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `qo_indptr` 和 `kv_indptr` 读取 ragged Q/K/V，并将结果写入 `output`。 + +baseline 使用 FlashInfer ragged prefill 的 Python API： + +```python +wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto") +wrapper.plan(qo_indptr, kv_indptr, num_qo_heads, num_kv_heads, + head_dim_qk, head_dim_vo, causal=True, + q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) +wrapper.run(q, k, v, out=output) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md new file mode 100644 index 0000000..06f1703 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md @@ -0,0 +1,41 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "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 `(batch_size * seq_len, num_qo_heads, head_dim_qk)`，连续 `bf16` +* `k`：key tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_qk)`，连续 `bf16` +* `v`：value tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_vo)`，连续 `bf16` +* `output`：输出缓冲区，shape `(batch_size * seq_len, 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` +* `causal`：是否启用 causal mask，评测中固定为 `1` + +本题测试中 `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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..35791da --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..d717add --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md new file mode 100644 index 0000000..1e57f2e --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：`q/kv_data/output` 均为 `bfloat16` +- KV layout：`NHD` +- `num_qo_heads = 32` +- `num_kv_heads = 4` +- `page_block_size = 16` +- `causal = 0` +- `head_dim` 固定为 `128` +- `batch_size` 取值为 `1, 4, 16` +- `seq_len` 取值为 `1024, 4096, 8192, 16384` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即： + +```python +for head_dim in [128]: + for batch_size in [1, 4, 16]: + for seq_len in [1024, 4096, 8192, 16384]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md new file mode 100644 index 0000000..b4417ae --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md @@ -0,0 +1,28 @@ +pytorch参考实现： + +```python +def baseline(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): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + 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) +``` + diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt new file mode 100644 index 0000000..1c5167b --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer Ragged Prefill diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json new file mode 100644 index 0000000..cd4684c --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json @@ -0,0 +1,16 @@ +{ + "id": 194, + "displayId": 20002, + "type": "Traditional", + "isPublic": false, + "locales": [ + "zh_CN" + ], + "samples": [ + { + "inputData": "1\n", + "outputData": "" + } + ], + "problemTagIds": [] +} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py new file mode 100644 index 0000000..4066882 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py @@ -0,0 +1,273 @@ +from __future__ import annotations + + +HEAD_DIMS = [128] +BATCH_SIZES = [1, 4, 16] +SEQ_LENS = [1024, 4096, 8192, 16384] +NUM_QO_HEADS = 32 +PAGE_BLOCK_SIZE = 16 +CAUSAL = 0 + + +def _build_cases(): + cases = [] + for head_dim in HEAD_DIMS: + for batch_size in BATCH_SIZES: + for seq_len in SEQ_LENS: + num_kv_heads = 8 if head_dim == 64 else 4 + cases.append( + ( + batch_size, + seq_len, + NUM_QO_HEADS, + num_kv_heads, + head_dim, + PAGE_BLOCK_SIZE, + CAUSAL, + ) + ) + return cases + + +TESTCASES = _build_cases() + + +def getNumOfTestcases() -> int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str): + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor") + kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device) + kv_indptr[0] = 0 + kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0) + num_blocks = int(kv_indptr[-1].item()) + last_page_len = (seq_lens - 1) % page_block_size + 1 + return kv_indptr, last_page_len, num_blocks + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim, page_block_size, causal = TESTCASES[testcase_id] + num_blocks = batch_size * ((seq_len + page_block_size - 1) // page_block_size) + qo_len = batch_size * seq_len + CURRENT_CASE = ( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal, + num_blocks, + 20260620 + testcase_id, + ) + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len, head_dim)) + return [ + (qo_len, num_qo_heads, head_dim), + (num_blocks, 2, page_block_size, num_kv_heads, head_dim), + (qo_len, num_qo_heads, head_dim), + (batch_size + 1,), + (batch_size + 1,), + (num_blocks,), + (batch_size,), + (), (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + ( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + causal, + num_blocks, + seed, + ) = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + qo_len = batch_size * seq_len + q = torch.rand(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous() + kv_data = torch.randn( + num_blocks, + 2, + page_block_size, + num_kv_heads, + head_dim, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + output = torch.empty(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device) + qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len + kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr( + batch_size, + seq_len, + page_block_size, + device, + ) + assert num_blocks_check == num_blocks + kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device) + return [ + 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, + ] + + def baseline( + 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, + ): + _ensure_flashinfer_importable() + import flashinfer + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + backend="auto", + ) + wrapper.plan( + qo_indptr, + kv_indptr, + kv_indices, + last_page_len, + int(num_qo_heads), + int(num_kv_heads), + int(head_dim), + int(page_block_size), + causal=bool(causal), + q_data_type=torch.bfloat16, + kv_data_type=torch.bfloat16, + ) + wrapper.run(q, kv_data, out=output) + return [ + 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, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[2] + output_ref = baseline_input_tensors[2] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md new file mode 100644 index 0000000..3e12eca --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md @@ -0,0 +1,22 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 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) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md new file mode 100644 index 0000000..48fbe06 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md @@ -0,0 +1,42 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..098c6da --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..d32ea1b --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md new file mode 100644 index 0000000..1e57f2e --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：`q/kv_data/output` 均为 `bfloat16` +- KV layout：`NHD` +- `num_qo_heads = 32` +- `num_kv_heads = 4` +- `page_block_size = 16` +- `causal = 0` +- `head_dim` 固定为 `128` +- `batch_size` 取值为 `1, 4, 16` +- `seq_len` 取值为 `1024, 4096, 8192, 16384` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即： + +```python +for head_dim in [128]: + for batch_size in [1, 4, 16]: + for seq_len in [1024, 4096, 8192, 16384]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md new file mode 100644 index 0000000..a55a635 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md @@ -0,0 +1,31 @@ +--- +sectionTitle: "输出格式" +type: "Text" +--- +PyTorch 参考实现 + +```python +def baseline(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): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, + 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) +``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt new file mode 100644 index 0000000..34e4801 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer Paged Prefill diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json new file mode 100644 index 0000000..f9593e3 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json @@ -0,0 +1,16 @@ +{ + "id": 195, + "displayId": 20003, + "type": "Traditional", + "isPublic": false, + "locales": [ + "zh_CN" + ], + "samples": [ + { + "inputData": "1\n", + "outputData": "" + } + ], + "problemTagIds": [] +} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py new file mode 100644 index 0000000..ee3ae71 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py @@ -0,0 +1,245 @@ +from __future__ import annotations + + +HEAD_DIM_CKV = 512 +HEAD_DIM_KPE = 64 +NUM_HEADS_LIST = [64, 128] +BATCH_SIZES = [1, 4, 16] +SEQ_LENS = [1024, 4096, 8192, 16384] +PAGE_SIZE = 1 +CAUSAL = 0 + + +def _build_cases(): + cases = [] + for num_heads in NUM_HEADS_LIST: + for batch_size in BATCH_SIZES: + for seq_len in SEQ_LENS: + cases.append( + ( + batch_size, + seq_len, + num_heads, + HEAD_DIM_CKV, + HEAD_DIM_KPE, + PAGE_SIZE, + CAUSAL, + ) + ) + return cases + + +TESTCASES = _build_cases() + + +def getNumOfTestcases() -> int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal = TESTCASES[testcase_id] + CURRENT_CASE = ( + batch_size, + seq_len, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + causal, + 20260630 + testcase_id, + ) + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe)) + kv_len = batch_size * seq_len + return [ + (batch_size, num_heads, head_dim_ckv), + (batch_size, num_heads, head_dim_kpe), + (kv_len, 1, head_dim_ckv), + (kv_len, 1, head_dim_kpe), + (batch_size, num_heads, head_dim_ckv), + (batch_size + 1,), + (batch_size + 1,), + (kv_len,), + (batch_size,), + (), (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal, seed = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + kv_len = batch_size * seq_len + q_nope = torch.randn(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous() + q_pe = torch.zeros(batch_size, num_heads, head_dim_kpe, dtype=dtype, device=device) + ckv = torch.randn(kv_len, 1, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous() + kpe = torch.zeros(kv_len, 1, head_dim_kpe, dtype=dtype, device=device) + output = torch.empty(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device) + q_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) + kv_indptr = q_indptr * seq_len + kv_indices = torch.arange(0, kv_len, dtype=torch.int32, device=device) + kv_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + return [ + 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, + ] + + def baseline( + 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, + ): + _ensure_flashinfer_importable() + import flashinfer + + sm_scale = 1.0 / ((int(head_dim_ckv) + int(head_dim_kpe)) ** 0.5) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device) + wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + int(num_heads), + int(head_dim_ckv), + int(head_dim_kpe), + int(page_size), + bool(causal), + sm_scale, + q_nope.dtype, + ckv.dtype, + ) + wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) + return [ + 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, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[4] + output_ref = baseline_input_tensors[4] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md new file mode 100644 index 0000000..c3bf938 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md @@ -0,0 +1,22 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 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) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md new file mode 100644 index 0000000..7cd07bb --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md @@ -0,0 +1,46 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..7196b09 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..43be38b --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md new file mode 100644 index 0000000..8919075 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：输入输出均为 `bfloat16` +- `head_dim_ckv = 512` +- `head_dim_kpe = 64` +- `page_size = 1` +- `causal = 0` +- `num_heads` 取值为 `64, 128` +- `batch_size` 取值为 `1, 4, 16` +- `seq_len` 取值为 `1024, 4096, 8192, 16384` +- `sm_scale = 1 / sqrt(head_dim_ckv + head_dim_kpe)` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_mla.py` 中的 cases 一致，即： + +```python +for num_heads in [64, 128]: + for batch_size in [1, 4, 16]: + for seq_len in [1024, 4096, 8192, 16384]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md new file mode 100644 index 0000000..23285e8 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md @@ -0,0 +1,30 @@ +--- +sectionTitle: "输出格式" +type: "Text" +--- +PyTorch 参考实现 + +```python +def baseline(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): + sm_scale = 1.0 / ((head_dim_ckv + head_dim_kpe) ** 0.5) + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device) + wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + bool(causal), + sm_scale, + q_nope.dtype, + ckv.dtype, + ) + wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) +``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt new file mode 100644 index 0000000..9f19bce --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer MLA Paged Attention diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json new file mode 100644 index 0000000..035e3ac --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json @@ -0,0 +1,16 @@ +{ + "id": 196, + "displayId": 20004, + "type": "Traditional", + "isPublic": false, + "locales": [ + "zh_CN" + ], + "samples": [ + { + "inputData": "1\n", + "outputData": "" + } + ], + "problemTagIds": [] +} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py new file mode 100644 index 0000000..3273aa3 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py @@ -0,0 +1,248 @@ +from __future__ import annotations + + +BATCH_SIZES = [1, 2, 4, 8, 16] +SEQ_LENS_KV = [512, 1024, 2048, 4096, 8192, 16384] +HEAD_DIMS = [128] +NUM_QO_HEADS = 32 +PAGE_BLOCK_SIZE = 16 + + +def _build_cases(): + cases = [] + for batch_size in BATCH_SIZES: + for seq_len_kv in SEQ_LENS_KV: + for head_dim in HEAD_DIMS: + num_kv_heads = 8 if head_dim == 64 else 4 + cases.append( + ( + batch_size, + seq_len_kv, + NUM_QO_HEADS, + num_kv_heads, + head_dim, + PAGE_BLOCK_SIZE, + ) + ) + return cases + + +TESTCASES = _build_cases() + + +def getNumOfTestcases() -> int: + return len(TESTCASES) + + +try: + from pathlib import Path + from typing import List, Tuple, Union + import sys + + import torch + + KernelArg = Union[torch.Tensor, int, float] + CURRENT_CASE = None + + def _ensure_flashinfer_importable(): + try: + import flashinfer # noqa: F401 + + return + except ImportError: + pass + + here = Path(__file__).resolve() + for parent in here.parents: + candidate = parent / "McFlashInfer" + if (candidate / "flashinfer").is_dir(): + sys.path.insert(0, str(candidate)) + return + + def _get_testcase_index() -> int: + try: + raw = input().strip() + except EOFError: + return 0 + if raw == "": + return 0 + try: + testcase_id = int(raw.split()[0]) + except ValueError: + return 0 + if 1 <= testcase_id <= len(TESTCASES): + return testcase_id - 1 + if 0 <= testcase_id < len(TESTCASES): + return testcase_id + return 0 + + def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: + workload = batch_size * seq_len * head_dim + if workload < 1e5: + return base_reps + if workload < 1e6: + return base_reps // 2 + if workload < 1e7: + return base_reps // 4 + if workload < 1e8: + return base_reps // 8 + if workload < 1e9: + return base_reps // 16 + return base_reps // 32 + + def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str): + seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) + seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor") + kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device) + kv_indptr[0] = 0 + kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0) + num_blocks = int(kv_indptr[-1].item()) + last_page_len = (seq_lens - 1) % page_block_size + 1 + return kv_indptr, last_page_len, num_blocks + + def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: + testcase_id = _get_testcase_index() + global CURRENT_CASE + batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size = TESTCASES[testcase_id] + num_blocks = batch_size * ((seq_len_kv + page_block_size - 1) // page_block_size) + CURRENT_CASE = ( + batch_size, + seq_len_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + num_blocks, + 20260710 + testcase_id, + ) + warmup = 3 + iters = max(1, _compute_reps(batch_size, seq_len_kv, head_dim)) + return [ + (batch_size, num_qo_heads, head_dim), + (num_blocks, 2, page_block_size, num_kv_heads, head_dim), + (batch_size, num_qo_heads, head_dim), + (batch_size + 1,), + (num_blocks,), + (batch_size,), + (), (), (), (), (), (), + ], (warmup, iters) + + def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: + del testcase_sizes + batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size, num_blocks, seed = CURRENT_CASE + gen = torch.Generator(device=device) + gen.manual_seed(seed) + dtype = torch.bfloat16 + q = torch.rand(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous() + kv_data = torch.randn( + num_blocks, + 2, + page_block_size, + num_kv_heads, + head_dim, + dtype=dtype, + device=device, + generator=gen, + ).contiguous() + output = torch.empty(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device) + kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr( + batch_size, + seq_len_kv, + page_block_size, + device, + ) + assert num_blocks_check == num_blocks + kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device) + return [ + 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, + ] + + def baseline( + 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, + ): + _ensure_flashinfer_importable() + import flashinfer + + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, + kv_layout="NHD", + use_tensor_cores=True, + ) + wrapper.plan( + kv_indptr, + kv_indices, + last_page_len, + int(num_qo_heads), + int(num_kv_heads), + int(head_dim), + int(page_block_size), + data_type=torch.bfloat16, + q_data_type=torch.bfloat16, + ) + wrapper.run(q, kv_data, out=output) + return [ + 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, + ] + + def check( + testcase_sizes, + original_input_tensors, + target_kernel_input_tensors, + baseline_input_tensors, + rtol=1e-2, + atol=1e-2, + ) -> bool: + del testcase_sizes, original_input_tensors + output_t = target_kernel_input_tensors[2] + output_ref = baseline_input_tensors[2] + if output_t.shape != output_ref.shape: + print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) + return False + if output_t.dtype != output_ref.dtype: + print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) + return False + if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): + diff = (output_t.float() - output_ref.float()).abs() + print( + f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " + f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", + file=sys.stderr, + ) + return False + return True +except Exception: + pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md new file mode 100644 index 0000000..74508b0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md @@ -0,0 +1,22 @@ +--- +sectionTitle: "题目描述" +type: "Text" +--- +你需要实现 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) +``` + +如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md new file mode 100644 index 0000000..c498b62 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md @@ -0,0 +1,38 @@ +--- +sectionTitle: "接口约定" +type: "codeSample" +lang: "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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md new file mode 100644 index 0000000..8d0fd87 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md new file mode 100644 index 0000000..d1751a9 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/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。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md new file mode 100644 index 0000000..9fe7bee --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md @@ -0,0 +1,26 @@ +--- +sectionTitle: "输入格式" +type: "Text" +--- +## 数据范围 + +- 数据类型：`q/kv_data/output` 均为 `bfloat16` +- KV layout：`NHD` +- `seq_len_q = 1` +- `num_qo_heads = 32` +- `page_block_size = 16` +- `head_dim` 固定为 `128` +- `num_kv_heads = 4` +- `batch_size` 取值为 `1, 2, 4, 8, 16` +- `seq_len_kv` 取值为 `512, 1024, 2048, 4096, 8192, 16384` + +测试点顺序与 `McFlashInfer/benchmarks/bench_batch_decode.py` 中的 cases 一致，即： + +```python +for batch_size in [1, 2, 4, 8, 16]: + for seq_len_kv in [512, 1024, 2048, 4096, 8192, 16384]: + for head_dim in [128]: + ... +``` + +输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md new file mode 100644 index 0000000..5d17935 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md @@ -0,0 +1,29 @@ +--- +sectionTitle: "输出格式" +type: "Text" +--- +PyTorch 参考实现 + +```python +def baseline(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): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, + 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) +``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt new file mode 100644 index 0000000..552b751 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt @@ -0,0 +1 @@ +FlashInfer Paged Decode 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。 From 483ca9ccbc83488e0ce3ce79e2aca047eeb68c41 Mon Sep 17 00:00:00 2001 From: MaseChen <93691652+MaseChen@users.noreply.github.com> Date: Tue, 7 Jul 2026 03:07:01 -0700 Subject: [PATCH 2/2] =?UTF-8?q?(flashinfer):=20=E5=8E=BB=E9=99=A4?= =?UTF-8?q?=E9=A2=98=E5=8C=85=E5=8F=8A=E5=85=B6=E6=8F=8F=E8=BF=B0=EF=BC=8C?= =?UTF-8?q?=E6=9B=B4=E6=96=B0=E6=96=87=E4=BB=B6=E7=BB=93=E6=9E=84-2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...ithPagedKVCacheWrapper_20260525_093143.csv | 145 --------- ...APagedAttentionWrapper_20260525_074200.csv | 33 -- ...ithPagedKVCacheWrapper_20260526_150417.csv | 33 -- ...thRaggedKVCacheWrapper_20260526_145454.csv | 49 --- .../benchmark/bench_batch_decode.py | 101 ------ .../benchmark/bench_batch_mla.py | 110 ------- .../benchmark/bench_batch_prefill_paged.py | 133 -------- .../benchmark/bench_batch_prefill_ragged.py | 136 -------- .../benchmark/bench_common.py | 91 ------ .../starter/示例冒烟代码.md | 209 ------------ .../xpuoj_problem/problem_20001/meta.json | 16 - .../problem_20001/testcase_config.py | 307 ------------------ .../problem_20001/zh_CN/00_题目描述.md | 23 -- .../zh_CN/01_接口约定cuda.md | 41 --- .../zh_CN/01_接口约定tilelang.md | 52 --- .../zh_CN/01_接口约定triton.md | 41 --- .../problem_20001/zh_CN/02_数据范围.md | 26 -- .../problem_20001/zh_CN/03_参考实现.md | 28 -- .../problem_20001/zh_CN/_title.txt | 1 - .../xpuoj_problem/problem_20002/meta.json | 16 - .../problem_20002/testcase_config.py | 273 ---------------- .../problem_20002/zh_CN/00_题目描述.md | 22 -- .../zh_CN/01_接口约定cuda.md | 42 --- .../zh_CN/01_接口约定tilelang.md | 55 ---- .../zh_CN/01_接口约定triton.md | 44 --- .../problem_20002/zh_CN/02_数据范围.md | 26 -- .../problem_20002/zh_CN/03_参考实现.md | 31 -- .../problem_20002/zh_CN/_title.txt | 1 - .../xpuoj_problem/problem_20003/meta.json | 16 - .../problem_20003/testcase_config.py | 245 -------------- .../problem_20003/zh_CN/00_题目描述.md | 22 -- .../zh_CN/01_接口约定cuda.md | 46 --- .../zh_CN/01_接口约定tilelang.md | 57 ---- .../zh_CN/01_接口约定triton.md | 45 --- .../problem_20003/zh_CN/02_数据范围.md | 26 -- .../problem_20003/zh_CN/03_参考实现.md | 30 -- .../problem_20003/zh_CN/_title.txt | 1 - .../xpuoj_problem/problem_20004/meta.json | 16 - .../problem_20004/testcase_config.py | 248 -------------- .../problem_20004/zh_CN/00_题目描述.md | 22 -- .../zh_CN/01_接口约定cuda.md | 38 --- .../zh_CN/01_接口约定tilelang.md | 52 --- .../zh_CN/01_接口约定triton.md | 41 --- .../problem_20004/zh_CN/02_数据范围.md | 26 -- .../problem_20004/zh_CN/03_参考实现.md | 29 -- .../problem_20004/zh_CN/_title.txt | 1 - 46 files changed, 3046 deletions(-) delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md delete mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv deleted file mode 100644 index ae909a7..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv +++ /dev/null @@ -1,145 +0,0 @@ -api,batch_size,seq_len_q,seq_len_kv,num_qo_heads,num_kv_heads,head_dim,time_ms,bandwidth_GB_s,tflops -BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,8,64,0.02042879999999998,51.528822055137894,0.8212531328320811 -BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,4,128,0.02333952000000001,45.27805199078642,0.718832949435121 -BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,4,256,0.0319488,66.15384615384615,1.0502564102564103 -BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,8,64,0.023262719999999973,90.32684054143292,1.4424122372620245 -BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,4,128,0.025041919999999992,84.07278675117566,1.3399304845634845 -BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,4,256,0.033387520000000004,126.11562643766291,2.0099984664928687 -BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,8,64,0.028298240000000037,148.36258368011562,2.371485435136599 -BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,4,128,0.027745280000000008,151.4670603432367,2.418748846650673 -BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,4,256,0.03723775999999999,225.7115358174069,3.604344837068611 -BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,8,64,0.03886591999999997,215.93992886312756,3.4533526544592306 -BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,4,128,0.03426815999999998,245.03212311370103,3.916689078141344 -BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,4,256,0.066048,254.26356589147287,4.064248062015504 -BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,8,64,0.052495359999999984,319.6722910367698,5.1135082414902975 -BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,4,128,0.04628480000000001,362.6548672566371,5.799646017699114 -BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,4,256,0.08975359999999999,374.0330861380491,5.981608670849972 -BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,8,64,0.08625152000000001,389.0775258221536,6.224480588863825 -BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,4,128,0.0638464,525.6776263031276,8.408789093825181 -BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,4,256,0.13059071999999994,514.0123892417473,8.222190857053247 -BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,8,64,0.02342912,89.86013986013987,1.4321678321678322 -BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,4,128,0.02486784,84.99073502161829,1.3493102738315832 -BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,4,256,0.03340287999999998,126.54812998160644,2.009074187614961 -BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,8,64,0.02839040000000001,148.02524797114512,2.3637871956717755 -BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,4,128,0.028165120000000012,149.5000908925649,2.382694055626249 -BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,4,256,0.03740160000000001,225.16084873374396,3.5885557837097872 -BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,8,64,0.03881984000000001,216.30176734370872,3.457451859667633 -BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,4,128,0.03601408000000001,233.38072220642587,3.7268126243957904 -BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,4,256,0.06728704000000002,249.82498858621207,3.9894080048698815 -BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,8,64,0.052490240000000014,319.7815060476004,5.114007023019897 -BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,4,128,0.04626431999999999,362.9924745462595,5.802213368747235 -BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,4,256,0.08993791999999999,373.44870773084375,5.969349880450872 -BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,8,64,0.08536063999999999,393.18618042226495,6.289443378119003 -BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,4,128,0.0630784,532.2077922077922,8.51116883116883 -BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,4,256,0.12952576,518.3650881492608,8.289793659577834 -BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,8,64,0.15207424000000003,441.34401723789637,7.0606423809844445 -BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,4,128,0.10330112,649.8017446471055,10.394290245836638 -BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,4,256,0.2281984,588.3060354498541,9.410599057662106 -BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,8,64,0.0283904,148.3137962128043,2.3637871956717764 -BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,4,128,0.028078080000000036,150.5470459518598,2.3900802334062696 -BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,4,256,0.03707903999999999,228.00331400165706,3.619773543220106 -BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,8,64,0.03844096000000004,218.64677677144357,3.4915290356952546 -BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,4,128,0.03641856000000004,231.23857725291697,3.6854210600309254 -BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,4,256,0.06640640000000002,253.63145720894363,4.04231303006939 -BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,8,64,0.059007999999999984,284.5986984815619,4.5491366594360105 -BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,4,128,0.04641792000000003,362.1442753143611,5.783013456871825 -BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,4,256,0.08961023999999998,375.1799794309223,5.991178151068451 -BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,8,64,0.09185279999999997,365.484949832776,5.84490523968785 -BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,4,128,0.06349823999999998,528.9469440412838,8.454894371875506 -BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,4,256,0.1303347200000001,515.3991200502825,8.238340666247638 -BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,8,64,0.16568319999999992,405.1421508034613,6.480692212608161 -BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,4,128,0.10290176,652.4828341128471,10.43463031147378 -BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,4,256,0.22947840000000008,585.1673360107093,9.358107987505575 -BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,8,64,0.30601215999999987,438.65613706331163,7.017641547316293 -BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,4,128,0.18384895999999992,730.2216776205863,11.680695109724857 -BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,4,256,0.4362026666666668,615.5418398787107,9.846265564630505 -BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,8,64,0.038655999999999975,217.85430463576174,3.472105960264903 -BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,4,128,0.03645951999999999,231.87754528858312,3.6812807190001418 -BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,4,256,0.06676480000000001,253.25153374233125,4.020613496932515 -BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,8,64,0.05858303999999996,286.9428421604617,4.582135990211505 -BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,4,128,0.04676608000000001,360.14889424129615,5.73996058681848 -BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,4,256,0.08992768000000002,374.58437713504884,5.970029606012297 -BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,8,64,0.092416,363.43490304709144,5.8092853185595565 -BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,4,128,0.07130112000000002,471.5208961654457,7.5296280338934345 -BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,4,256,0.14862335999999993,452.41835469202175,7.224583161085851 -BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,8,64,0.16396288000000003,409.4928803397451,6.548688483637271 -BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,4,128,0.11201536000000002,599.6891854831337,9.585665965810401 -BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,4,256,0.24935424000000006,538.7869081351894,8.61218019793848 -BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,8,64,0.3056947200000001,439.16524302415155,7.024928817874248 -BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,4,128,0.20128768000000002,667.1211273337741,10.668728697156228 -BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,4,256,0.46690133333333317,575.2104541716168,9.198875628255509 -BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,8,64,0.5866495999999998,457.6296037702917,7.321179961598886 -BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,4,128,0.37337600000000015,719.1169009256082,11.503062050051419 -BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,4,256,0.8934826666666666,601.0211546726517,9.613991308915525 -BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,8,64,0.0698112,241.26145947928126,3.845163182984965 -BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,4,128,0.04724735999999999,357.8673602080625,5.681491114000869 -BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,4,256,0.08954879999999998,377.6329331046313,5.995288736420813 -BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,8,64,0.12070911999999998,278.52052935188334,4.447641669494402 -BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,4,128,0.07076864000000004,475.9947909130369,7.586282737664589 -BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,4,256,0.14710784000000002,457.9702074342197,7.2990115550605585 -BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,8,64,0.22239232000000014,302.05359609540454,4.8281425545630325 -BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,4,128,0.11209728000000002,599.8355713894217,9.578660820316067 -BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,4,256,0.2504192,537.0190145164587,8.575555101206296 -BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,8,64,0.42098688000000006,318.97256275539985,5.101070247129791 -BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,4,128,0.2027008,662.7936347562515,10.594352109118466 -BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,4,256,0.46432,578.6905582356995,9.250015713301172 -BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,8,64,0.8234496000000004,326.06851955480926,5.215822918609709 -BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,4,128,0.3726506666666667,720.6924662239522,11.525451797572705 -BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,4,256,0.8939733333333334,600.8378952392316,9.608714568667223 -BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,8,64,1.6324906666666663,328.9062896122735,5.261858317107276 -BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,4,128,0.7114879999999999,754.7590177206082,12.073196725735361 -BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,4,256,1.742272,616.4387466480549,9.860612570253094 -BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,8,64,0.08406016,400.730905104154,6.386746254111341 -BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,4,128,0.08498175999999999,397.92746113989637,6.317484034220991 -BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,4,256,0.1808896,373.8918765921313,5.935895839230116 -BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,8,64,0.14712832,457.0155902004454,7.297995545657015 -BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,4,128,0.14887935999999996,452.5208061077104,7.212160396175805 -BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,4,256,0.3279462400000001,410.8661712358707,6.548279522887651 -BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,8,64,0.27223039999999993,493.51137859695325,7.888478465299983 -BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,4,128,0.27833343999999993,483.1610316029581,7.715507155733786 -BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,4,256,0.6300373333333331,426.89493109403054,6.817004435715981 -BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,8,64,0.52494336,511.6104868913858,8.181772784019977 -BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,4,128,0.5080533333333332,528.8767583455806,8.453772496325847 -BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,4,256,1.2449493333333332,431.66029782202656,6.899826653186422 -BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,8,64,1.0273706666666667,522.6954607749491,8.361086091615102 -BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,4,128,1.0078719999999999,532.9377698755399,8.522842773685548 -BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,4,256,2.446784,439.05228741073995,7.021408176610604 -BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,8,64,2.0322986666666663,528.4030903594223,8.453417534430637 -BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,4,128,2.018026666666667,532.2050425498176,8.513202262276018 -BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,4,256,4.847957333333333,443.0748433429558,7.087467154826447 -BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,8,64,0.13077504,515.1671756322919,8.210602145485865 -BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,4,128,0.14377984000000002,470.39384659212305,7.4679581226408365 -BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,4,256,0.27039743999999993,500.2499431947286,7.9419525865333656 -BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,8,64,0.231424,581.0973451327434,9.279433628318584 -BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,4,128,0.25729023999999995,523.6965692907746,8.34654143118682 -BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,4,256,0.502016,536.8036715961244,8.555439061703213 -BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,8,64,0.4335923199999999,619.7010131544766,9.905542828802874 -BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,4,128,0.47517866666666664,566.018137739068,9.038636616683128 -BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,4,256,0.9693866666666666,554.9070422535212,8.861205633802816 -BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,8,64,0.8388479999999999,640.3222705424583,10.240156252384224 -BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,4,128,0.9261013333333336,580.2768883462716,9.275372232844207 -BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,4,256,1.906474666666667,563.7580287805205,9.011328335161021 -BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,8,64,1.6543999999999999,649.1803481624759,10.384350328820116 -BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,4,128,1.8147413333333327,591.9665201137942,9.466841840453299 -BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,4,256,3.774634666666667,569.2026947596871,9.10279839037844 -BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,8,64,3.2680746666666667,657.1899393567508,10.513755612274872 -BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,4,128,3.5912106666666666,598.129192458031,9.567731207451676 -BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,4,256,7.526272,570.802632697835,9.130612969608327 -BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,8,64,0.2176000000000001,619.2188235294115,9.86895058823529 -BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,4,128,0.21536768,628.0715100798782,9.971243818942565 -BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,4,256,0.45757866666666663,591.2264441232692,9.386292694298103 -BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,8,64,0.39856127999999985,674.826576229382,10.776177997019683 -BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,4,128,0.39381333333333335,684.2938244853738,10.906099241603465 -BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,4,256,0.8577493333333336,628.3514810853829,10.014504539010616 -BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,8,64,0.7606186666666664,706.5238121949853,11.293352330734283 -BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,4,128,0.7354026666666665,731.4625203063357,11.680586679043863 -BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,4,256,1.6673066666666665,645.2556074467406,10.30396478792144 -BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,8,64,1.4816639999999999,725.0403006349618,11.594983197270098 -BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,4,128,1.4354773333333333,748.7338009749138,11.968053263583403 -BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,4,256,3.2697173333333325,657.4209880731792,10.508473627893627 -BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,8,64,2.9226666666666676,734.9479708029195,11.75629734306569 -BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,4,128,2.825301333333333,760.4612643087983,12.161442024827087 -BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,4,256,6.484309333333334,662.6865294520187,10.5978097594357 -BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,8,64,5.794901333333332,741.2536188134122,11.858610316747416 -BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,4,128,5.61536,765.0472760428539,12.237768680191476 -BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,4,256,12.908458666666668,665.6125232199165,10.647200957222468 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv deleted file mode 100644 index 432acdb..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchMLAPagedAttentionWrapper_20260525_074200.csv +++ /dev/null @@ -1,33 +0,0 @@ -api,batch_size,seq_len,num_heads,head_dim_ckv,head_dim_kpe,time_ms,bandwidth_GB_s,tflops -BatchMLAPagedAttentionWrapper,1,1024,64,512,64,0.035975679999999996,34.83953604212624,3.963964989681919 -BatchMLAPagedAttentionWrapper,1,4096,64,512,64,0.05349631999999998,89.58223668469162,10.662889409963158 -BatchMLAPagedAttentionWrapper,1,8192,64,512,64,0.06174719999999999,154.0298507462687,18.47615257048093 -BatchMLAPagedAttentionWrapper,1,16384,64,512,64,0.08995584000000004,210.63775292410136,25.3646831156265 -BatchMLAPagedAttentionWrapper,4,1024,64,512,64,0.05086207999999998,98.5705657338434,11.215139923495071 -BatchMLAPagedAttentionWrapper,4,4096,64,512,64,0.08034559999999999,238.58531145451653,28.39858531145452 -BatchMLAPagedAttentionWrapper,4,8192,64,512,64,0.10866687999999997,350.0942329438373,41.99442140972485 -BatchMLAPagedAttentionWrapper,4,16384,64,512,64,0.16821760000000002,450.56155836250184,54.2559488662304 -BatchMLAPagedAttentionWrapper,16,1024,64,512,64,0.06735359999999997,297.7423033067276,33.87645762067656 -BatchMLAPagedAttentionWrapper,16,4096,64,512,64,0.14288383999999996,536.6395528003728,63.87570143691549 -BatchMLAPagedAttentionWrapper,16,8192,64,512,64,0.21618431999999987,703.9113289992544,84.43540682321462 -BatchMLAPagedAttentionWrapper,16,16384,64,512,64,0.39363328000000025,770.1826837405613,92.74424666532254 -BatchMLAPagedAttentionWrapper,64,1024,64,512,64,0.15278592,525.0226198853926,59.73590697362689 -BatchMLAPagedAttentionWrapper,64,4096,64,512,64,0.4850483199999999,632.3256206721838,75.26512413443676 -BatchMLAPagedAttentionWrapper,64,8192,64,512,64,0.9133465600000001,666.4484158127227,79.94166423750474 -BatchMLAPagedAttentionWrapper,64,16384,64,512,64,1.7720038399999998,684.3541287134007,82.40890045926764 -BatchMLAPagedAttentionWrapper,1,1024,128,512,64,0.04499968000000001,29.491409716691315,6.338104448742746 -BatchMLAPagedAttentionWrapper,1,4096,128,512,64,0.05375743999999999,90.51859612362495,21.222191532930147 -BatchMLAPagedAttentionWrapper,1,8192,128,512,64,0.08302080000000002,115.44865864939868,27.48349059512796 -BatchMLAPagedAttentionWrapper,1,16384,128,512,64,0.11321343999999998,168.01736613603475,40.30795947901592 -BatchMLAPagedAttentionWrapper,4,1024,128,512,64,0.05178880000000003,102.50123578843295,22.028907563025196 -BatchMLAPagedAttentionWrapper,4,4096,128,512,64,0.11032576,176.4247261926861,41.36298496380175 -BatchMLAPagedAttentionWrapper,4,8192,128,512,64,0.1688268800000001,227.08800873415404,54.06014435615937 -BatchMLAPagedAttentionWrapper,4,16384,128,512,64,0.30781695999999986,247.18357299091002,59.30021207408457 -BatchMLAPagedAttentionWrapper,16,1024,128,512,64,0.10527487999999995,201.69734698344004,43.34749896651511 -BatchMLAPagedAttentionWrapper,16,4096,128,512,64,0.2629478400000002,296.0920614521874,69.41913273750409 -BatchMLAPagedAttentionWrapper,16,8192,128,512,64,0.3962367999999998,387.02674764181444,92.13485980100793 -BatchMLAPagedAttentionWrapper,16,16384,128,512,64,0.7528985599999998,404.23663979381246,96.97779742333418 -BatchMLAPagedAttentionWrapper,64,1024,128,512,64,0.3242547199999998,261.9380714026308,56.29404872811108 -BatchMLAPagedAttentionWrapper,64,4096,128,512,64,1.1793126399999994,264.07507342582215,61.91271216426548 -BatchMLAPagedAttentionWrapper,64,8192,128,512,64,2.3186406399999986,264.55887532446616,62.98038839860932 -BatchMLAPagedAttentionWrapper,64,16384,128,512,64,4.6020608,264.53295358462015,63.462389746784744 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv deleted file mode 100644 index 3dfd620..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv +++ /dev/null @@ -1,33 +0,0 @@ -api,batch_size,seq_len,num_qo_heads,num_kv_heads,head_dim,time_ms,bandwidth_GB_s,tflops -BatchPrefillWithPagedKVCacheWrapper,1,1024,32,4,128,0.3529011200000001,29.71302556364796,24.34091054174041 -BatchPrefillWithPagedKVCacheWrapper,1,4096,32,4,128,4.62532608,9.068126068205768,29.714435500296666 -BatchPrefillWithPagedKVCacheWrapper,1,8192,32,4,128,18.113853439999996,4.631045529757804,30.350019983820744 -BatchPrefillWithPagedKVCacheWrapper,1,16384,32,4,128,71.05519616000001,2.36115258372119,30.948099145350383 -BatchPrefillWithPagedKVCacheWrapper,4,1024,32,4,128,1.2374374399999997,33.89507917264893,27.766848858234006 -BatchPrefillWithPagedKVCacheWrapper,4,4096,32,4,128,17.896878079999997,9.374381344614939,30.71797279003423 -BatchPrefillWithPagedKVCacheWrapper,4,8192,32,4,128,71.25501952,4.709062214288198,30.861310127559136 -BatchPrefillWithPagedKVCacheWrapper,4,16384,32,4,128,283.27072767999994,2.3690716139159393,31.051895457919 -BatchPrefillWithPagedKVCacheWrapper,16,1024,32,4,128,4.752537600000002,35.301595509733566,28.919067041573737 -BatchPrefillWithPagedKVCacheWrapper,16,4096,32,4,128,70.51405312000001,9.517090711803915,31.185602844439067 -BatchPrefillWithPagedKVCacheWrapper,16,8192,32,4,128,284.16772266666663,4.723186952426669,30.953878011423416 -BatchPrefillWithPagedKVCacheWrapper,16,16384,32,4,128,1129.139136,2.377346134250013,31.160351250841774 -BatchPrefillWithPagedKVCacheWrapper,64,1024,32,4,128,18.757478399999997,35.77712449878125,29.3086203894016 -BatchPrefillWithPagedKVCacheWrapper,64,4096,32,4,128,281.4907093333333,9.536210151864244,31.248253425628754 -BatchPrefillWithPagedKVCacheWrapper,64,8192,32,4,128,1134.7048106666668,4.731370722616177,31.007511167737377 -BatchPrefillWithPagedKVCacheWrapper,64,16384,32,4,128,4514.139178666666,2.378619226173592,31.177037921302507 -BatchPrefillWithPagedKVCacheWrapper,1,1024,32,4,256,0.7928422399999997,26.4510629504301,21.668710768992337 -BatchPrefillWithPagedKVCacheWrapper,1,4096,32,4,256,12.533002240000002,6.69321511267838,21.932327281224513 -BatchPrefillWithPagedKVCacheWrapper,1,8192,32,4,256,49.81321727999999,3.368024977325858,22.072688491402744 -BatchPrefillWithPagedKVCacheWrapper,1,16384,32,4,256,190.01136128,1.765917141688929,23.14622915954513 -BatchPrefillWithPagedKVCacheWrapper,4,1024,32,4,256,3.111116800000001,26.963333552761494,22.088362846422218 -BatchPrefillWithPagedKVCacheWrapper,4,4096,32,4,256,47.738091520000026,7.02885912101079,23.032165567728153 -BatchPrefillWithPagedKVCacheWrapper,4,8192,32,4,256,190.14286336,3.529391680241077,23.130221315627924 -BatchPrefillWithPagedKVCacheWrapper,4,16384,32,4,256,759.6848640000004,1.76675532658763,23.157215416649382 -BatchPrefillWithPagedKVCacheWrapper,16,1024,32,4,256,12.28442624,27.31461066593534,22.376129057534232 -BatchPrefillWithPagedKVCacheWrapper,16,4096,32,4,256,191.34602666666663,7.014398487291994,22.984780963158407 -BatchPrefillWithPagedKVCacheWrapper,16,8192,32,4,256,759.7649706666668,3.5331380935403933,23.15477380982632 -BatchPrefillWithPagedKVCacheWrapper,16,16384,32,4,256,3028.668266666667,1.77263029400997,23.234219789647476 -BatchPrefillWithPagedKVCacheWrapper,64,1024,32,4,256,49.26948266666667,27.241554149868346,22.316281159572153 -BatchPrefillWithPagedKVCacheWrapper,64,4096,32,4,256,763.6229333333335,7.030576067909256,23.037791659325052 -BatchPrefillWithPagedKVCacheWrapper,64,8192,32,4,256,3037.7449386666663,3.534667477616765,23.16479678130923 -BatchPrefillWithPagedKVCacheWrapper,64,16384,32,4,256,12110.653866666667,1.7732185822854112,23.241930601731337 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv deleted file mode 100644 index 4f20821..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv +++ /dev/null @@ -1,49 +0,0 @@ -api,batch_size,seq_len,num_qo_heads,num_kv_heads,head_dim_qk,head_dim_vo,time_ms,bandwidth_GB_s,tflops -BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,128,128,0.031580159999999996,66.66666666666667,272.00415045395596 -BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,128,128,0.0424448,197.82870928829917,3238.0634016887816 -BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,128,128,0.057313279999999994,292.871180989816,9592.119206717885 -BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,128,128,0.06972416000000001,481.36290204141574,31538.89922161844 -BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,128,128,0.04327423999999998,194.60482725982024,793.9998106956938 -BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,128,128,0.06579199999999998,510.5058365758757,8355.967501945528 -BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,128,128,0.09618432000000002,698.0517406579366,22862.596060896405 -BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,128,128,0.15411199999999997,871.12292358804,57075.97735548174 -BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,128,128,0.07452671999999999,451.99230557845567,1844.1567463588901 -BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,128,128,0.1668906666666667,805.0108653969065,13176.43041083983 -BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,128,128,0.2874026666666667,934.46080760095,30605.46766745843 -BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,128,128,0.5342506666666667,1005.1498622369525,65857.42289917343 -BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,128,128,0.15733333333333333,856.4111186440679,3494.2106814915255 -BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,128,128,0.5614719999999999,957.1184315513509,15666.129428017784 -BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,128,128,1.1031466666666667,973.8198414233224,31894.55505055115 -BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,128,128,2.1813759999999998,984.7032038493136,64517.75776176505 -BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,192,128,0.03564544000000001,73.88681413386956,301.22838264866414 -BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,192,128,0.04922368,213.27231121281466,3490.1635115456625 -BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,192,128,0.061327359999999984,342.16062781766584,11205.353815328106 -BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,192,128,0.08377343999999999,500.8189707859675,32812.059161471705 -BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,192,128,0.049623040000000056,212.29880313660726,865.5187783739157 -BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,192,128,0.08634367999999998,486.3377609108161,7958.831119544594 -BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,192,128,0.13644799999999999,615.1444652908068,20145.249981238278 -BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,192,128,0.2321706666666666,722.8359827253516,47357.904577781876 -BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,192,128,0.09042944,465.99479107688825,1899.8093081191257 -BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,192,128,0.3087573333333334,544.0154770952807,8902.716705589717 -BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,192,128,0.5995946666666665,559.9464882943145,18337.58185156195 -BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,192,128,1.1809706666666668,568.4182232017052,37240.94624227753 -BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,192,128,0.2555306666666667,659.6413424611787,2689.2849156787443 -BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,192,128,0.9085866666666667,739.472740079831,12101.340115520075 -BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,192,128,1.7810773333333334,754.017631276351,24693.1810808739 -BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,192,128,3.5260586666666662,761.5134193267346,49891.92667360423 -BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,256,256,0.044037119999999964,95.61678874549479,390.12245087780525 -BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,256,256,0.08118271999999997,206.86175580222005,3385.916448032292 -BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,256,256,0.11204607999999996,299.6161579235972,9813.030743922503 -BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,256,256,0.14619648000000002,459.1440778875113,30083.12177628353 -BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,256,256,0.07792639999999999,216.1366622864652,881.8510381077531 -BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,256,256,0.13784064000000001,487.3337790654483,7976.686902904687 -BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,256,256,0.22408533333333336,599.2505712109672,19626.65938766184 -BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,256,256,0.3959893333333334,678.0510720827496,44425.908890852275 -BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,256,256,0.15150079999999996,444.6907739101049,1814.366042581954 -BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,256,256,0.4274346666666664,628.6284687562392,10289.400589339195 -BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,256,256,0.7913173333333334,678.7833823093305,22231.518637802277 -BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,256,256,1.5360853333333337,699.1824898616742,45810.43946625186 -BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,256,256,0.43906133333333336,613.773091686507,2504.2324256352945 -BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,256,256,1.6363946666666664,656.8039006075145,10750.576497692491 -BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,256,256,3.234005333333333,664.3564257160574,21759.006842803803 -BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,256,256,6.420821333333334,669.0757535484556,43837.84598543405 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py deleted file mode 100644 index 32861ca..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_decode.py +++ /dev/null @@ -1,101 +0,0 @@ -""" -Benchmark script for BatchDecodeWithPagedKVCacheWrapper -seq_len_q=1 (decode mode), seq_len_kv from 1K to 16K -""" - -import itertools -import pandas as pd -import torch - -import flashinfer -from bench_common import dtype, page_block_size, setup_workspace, setup_paged_kv_indptr, run_with_profiler, get_csv_path, compute_reps - -target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] - - -def bench_batch_decode( - batch_size, - seq_len_kv, - num_qo_heads, - num_kv_heads, - head_dim, - page_block_size, -): - """Benchmark BatchDecodeWithPagedKVCacheWrapper""" - seq_lens = [seq_len_kv] * batch_size - - kv_indptr, last_page_len, num_blocks = setup_paged_kv_indptr(batch_size, seq_lens) - - q = torch.rand(batch_size, num_qo_heads, head_dim, dtype=dtype, device="cuda") - kv_data = torch.randn(num_blocks, 2, page_block_size, num_kv_heads, head_dim, dtype=dtype, device="cuda") - - workspace_buffer = setup_workspace() - wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( - workspace_buffer, kv_layout="NHD", use_tensor_cores=True - ) - wrapper.plan( - kv_indptr.to("cuda"), - torch.arange(num_blocks, dtype=torch.int32, device="cuda"), - last_page_len.to("cuda"), - num_qo_heads, - num_kv_heads, - head_dim, - page_block_size, - data_type=dtype, - q_data_type=dtype, - ) - - reps = compute_reps(batch_size, seq_len_kv, head_dim, base_reps=100) - ms = run_with_profiler(lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps) - - io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() - flops = 2 * batch_size * seq_len_kv * num_qo_heads * num_kv_heads * head_dim - return ms, io, flops - - -def run_benchmark(): - records = [] - - batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128] - head_dims = [64, 128, 256] - seq_lens_kv = [512, 1024, 2048, 4096, 8192, 16384] - - api_name = "BatchDecodeWithPagedKVCacheWrapper" - test_cases = list(itertools.product(batch_sizes, seq_lens_kv, head_dims)) - total_cases = len(test_cases) - - print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") - print(f" seq_len_q=1 (decode mode), causal=False") - for idx, (bs, sl_kv, hd) in enumerate(test_cases, 1): - num_qo_heads = 32 - num_kv_heads = 8 if hd == 64 else 4 - ms, io, flops = bench_batch_decode(bs, sl_kv, num_qo_heads, num_kv_heads, hd, page_block_size) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append({ - "api": api_name, - "batch_size": bs, - "seq_len_q": 1, - "seq_len_kv": sl_kv, - "num_qo_heads": num_qo_heads, - "num_kv_heads": num_kv_heads, - "head_dim": hd, - "time_ms": ms, - "bandwidth_GB_s": bw, - "tflops": tflops, - }) - print(f" [{idx}/{total_cases}] bs={bs}, kv_len={sl_kv}, hd={hd}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") - - return records - - -if __name__ == "__main__": - import numpy as np - np.random.seed(42) - torch.random.manual_seed(42) - - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchDecodeWithPagedKVCacheWrapper") - df.to_csv(csv_path, index=False) - print(f"\nResults saved to {csv_path}") \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py deleted file mode 100644 index 5c48bc9..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_mla.py +++ /dev/null @@ -1,110 +0,0 @@ -""" -Benchmark script for BatchMLAPagedAttentionWrapper -headdim: ckv=512, kpe=64 (DeepSeek MLA configuration) -""" - -import itertools -import pandas as pd -import torch - -import flashinfer -from bench_common import dtype, page_block_size, setup_workspace, run_with_profiler, get_csv_path, compute_reps - -target_kernels = ["BatchMLAPagedAttentionKernel"] - - -def bench_batch_mla_paged_attention( - batch_size, - seq_len, - num_heads, - head_dim_ckv, - head_dim_kpe, -): - """Benchmark BatchMLAPagedAttentionWrapper for DeepSeek MLA""" - # MLA decode mode: q has length 1, not seq_len - q_nope = torch.randn(batch_size, num_heads, head_dim_ckv, dtype=dtype, device="cuda") - q_pe = torch.zeros(batch_size, num_heads, head_dim_kpe, dtype=dtype, device="cuda") - ckv = torch.randn(batch_size * seq_len, 1, head_dim_ckv, dtype=dtype, device="cuda") - kpe = torch.zeros(batch_size * seq_len, 1, head_dim_kpe, dtype=dtype, device="cuda") - - sm_scale = 1.0 / ((head_dim_ckv + head_dim_kpe) ** 0.5) - - # q_indptr for decode: each query has length 1 - q_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") - kv_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * seq_len - kv_indices = torch.arange(0, batch_size * seq_len, dtype=torch.int32, device="cuda") - kv_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") - - page_size = 1 # MLA uses page_size=1 - - workspace_buffer = setup_workspace() - wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") - wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - kv_lens, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - False, # causal - sm_scale, - q_nope.dtype, - ckv.dtype, - ) - - reps = compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe, base_reps=100) - ms = run_with_profiler(lambda: wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False), target_kernels=target_kernels, reps=reps) - - io = sum([t.numel() * t.element_size() for t in [q_nope, q_pe, ckv, kpe]]) - # MLA FLOPs: 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len - flops = 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len - return ms, io, flops - - -def run_benchmark(): - records = [] - - # MLA configuration - same as DeepSeek - head_dim_ckv = 512 - head_dim_kpe = 64 - batch_sizes = [1, 4, 16, 64] - seq_lens = [1024, 4096, 8192, 16384] - num_heads_list = [64, 128] - - api_name = "BatchMLAPagedAttentionWrapper" - test_cases = list(itertools.product(num_heads_list, batch_sizes, seq_lens)) - total_cases = len(test_cases) - - print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") - for idx, (num_heads, bs, sl) in enumerate(test_cases, 1): - ms, io, flops = bench_batch_mla_paged_attention(bs, sl, num_heads, head_dim_ckv, head_dim_kpe) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append({ - "api": api_name, - "batch_size": bs, - "seq_len": sl, - "num_heads": num_heads, - "head_dim_ckv": head_dim_ckv, - "head_dim_kpe": head_dim_kpe, - "time_ms": ms, - "bandwidth_GB_s": bw, - "tflops": tflops, - }) - print(f" [{idx}/{total_cases}] bs={bs}, sl={sl}, num_heads={num_heads}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") - - return records - - -if __name__ == "__main__": - import numpy as np - np.random.seed(42) - torch.random.manual_seed(42) - - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchMLAPagedAttentionWrapper") - df.to_csv(csv_path, index=False) - print(f"\nResults saved to {csv_path}") \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py deleted file mode 100644 index 1dee3d0..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_paged.py +++ /dev/null @@ -1,133 +0,0 @@ -""" -Benchmark script for BatchPrefillWithPagedKVCacheWrapper -headdim: 64/128/256 -""" - -import itertools -import pandas as pd -import torch - -import flashinfer -from bench_common import ( - dtype, - setup_workspace, - setup_paged_kv_indptr, - run_with_profiler, - get_csv_path, - compute_reps, -) - -target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] - - -def bench_batch_prefill_with_paged_kv_cache( - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim, - causal=True, -): - """Benchmark BatchPrefillWithPagedKVCacheWrapper""" - q_lens = [seq_len] * batch_size - kv_lens = [seq_len] * batch_size - - qo_indptr = torch.cat( - [torch.tensor([0]), torch.cumsum(torch.tensor(q_lens), 0)], dim=0 - ).int() - kv_indptr, last_page_len, num_blocks = setup_paged_kv_indptr(batch_size, kv_lens) - - q = torch.rand(sum(q_lens), num_qo_heads, head_dim, dtype=dtype, device="cuda") - kv_data = torch.randn( - num_blocks, 2, 16, num_kv_heads, head_dim, dtype=dtype, device="cuda" - ) - - workspace_buffer = setup_workspace() - wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( - workspace_buffer, kv_layout="NHD", backend="auto" - ) - wrapper.plan( - qo_indptr, - kv_indptr, - torch.arange(num_blocks, dtype=torch.int32, device="cuda"), - last_page_len, - num_qo_heads, - num_kv_heads, - head_dim, - 16, - q_data_type=dtype, - kv_data_type=dtype, - ) - - reps = compute_reps(batch_size, seq_len, head_dim, base_reps=100) - ms = run_with_profiler( - lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps - ) - - io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() - # Attention FLOPs calculation: - # - causal=True: triangular pattern - # - causal=False: full attention - flops = ( - 2 - * batch_size - * seq_len - * seq_len - * num_qo_heads - * head_dim - * (1 if causal else 2) - ) - return ms, io, flops - - -def run_benchmark(): - records = [] - - batch_sizes = [1, 4, 16, 64] - seq_lens = [1024, 4096, 8192, 16384] - head_dims = [128, 256] - - api_name = "BatchPrefillWithPagedKVCacheWrapper" - test_cases = list(itertools.product(head_dims, batch_sizes, seq_lens)) - total_cases = len(test_cases) - - print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") - for idx, (head_dim, bs, sl) in enumerate(test_cases, 1): - num_qo_heads = 32 - num_kv_heads = 8 if head_dim == 64 else 4 - ms, io, flops = bench_batch_prefill_with_paged_kv_cache( - bs, sl, num_qo_heads, num_kv_heads, head_dim - ) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append( - { - "api": api_name, - "batch_size": bs, - "seq_len": sl, - "num_qo_heads": num_qo_heads, - "num_kv_heads": num_kv_heads, - "head_dim": head_dim, - "time_ms": ms, - "bandwidth_GB_s": bw, - "tflops": tflops, - } - ) - print( - f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd={head_dim}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" - ) - - return records - - -if __name__ == "__main__": - import numpy as np - - np.random.seed(42) - torch.random.manual_seed(42) - - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchPrefillWithPagedKVCacheWrapper") - df.to_csv(csv_path, index=False) - print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py deleted file mode 100644 index 43aa4e0..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_batch_prefill_ragged.py +++ /dev/null @@ -1,136 +0,0 @@ -""" -Benchmark script for BatchPrefillWithRaggedKVCacheWrapper -headdim configurations: [64,64], [128,128], [192,128], [256,256] -""" - -import itertools -import pandas as pd -import torch - -import flashinfer -from bench_common import ( - dtype, - setup_workspace, - run_with_profiler, - get_csv_path, - compute_reps, -) - -target_kernels = [ - "BatchPrefillWithRaggedKVCacheKernel", - "PersistentVariableLengthMergeStates", -] - - -def bench_batch_prefill_with_ragged_kv_cache( - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim_qk, - head_dim_vo, - causal=True, -): - """Benchmark BatchPrefillWithRaggedKVCacheWrapper for MLA""" - qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") - kv_indptr = ( - torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * seq_len - ) - - q = torch.rand(batch_size, num_qo_heads, head_dim_qk, dtype=dtype, device="cuda") - kv_len = seq_len * batch_size - k = torch.rand(kv_len, num_kv_heads, head_dim_qk, dtype=dtype, device="cuda") - v = torch.rand(kv_len, num_kv_heads, head_dim_vo, dtype=dtype, device="cuda") - - workspace_buffer = setup_workspace() - 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=causal, - q_data_type=dtype, - kv_data_type=dtype, - ) - - reps = compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo, base_reps=100) - ms = run_with_profiler( - lambda: wrapper.run(q, k, v), target_kernels=target_kernels, reps=reps - ) - - io = ( - q.numel() * q.element_size() - + k.numel() * k.element_size() - + v.numel() * v.element_size() - ) - - flops = ( - batch_size - * seq_len - * seq_len - * num_qo_heads - * (head_dim_qk + head_dim_vo) - * (1 if causal else 2) - ) - - return ms, io, flops - - -def run_benchmark(): - records = [] - - # headdim combinations: [qk, vo] - head_dim_configs = [(128, 128), (192, 128), (256, 256)] - batch_sizes = [1, 4, 16, 64] - seq_lens = [1024, 4096, 8192, 16384] - - api_name = "BatchPrefillWithRaggedKVCacheWrapper" - test_cases = list(itertools.product(head_dim_configs, batch_sizes, seq_lens)) - total_cases = len(test_cases) - - print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") - for idx, ((head_dim_qk, head_dim_vo), bs, sl) in enumerate(test_cases, 1): - num_qo_heads = 32 - num_kv_heads = 4 - ms, io, flops = bench_batch_prefill_with_ragged_kv_cache( - bs, sl, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo - ) - bw = io / ms / 1e6 - tflops = flops / ms / 1e9 - records.append( - { - "api": api_name, - "batch_size": bs, - "seq_len": sl, - "num_qo_heads": num_qo_heads, - "num_kv_heads": num_kv_heads, - "head_dim_qk": head_dim_qk, - "head_dim_vo": head_dim_vo, - "time_ms": ms, - "bandwidth_GB_s": bw, - "tflops": tflops, - } - ) - print( - f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd=[{head_dim_qk},{head_dim_vo}]: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" - ) - - return records - - -if __name__ == "__main__": - import numpy as np - - np.random.seed(42) - torch.random.manual_seed(42) - - records = run_benchmark() - df = pd.DataFrame(records) - csv_path = get_csv_path("BatchPrefillWithRaggedKVCacheWrapper") - df.to_csv(csv_path, index=False) - print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py deleted file mode 100644 index 3a0563a..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/benchmark/bench_common.py +++ /dev/null @@ -1,91 +0,0 @@ -""" -Common utilities for FlashInfer benchmarks -""" - -import os -import random -from datetime import datetime -import numpy as np -import torch - -page_block_size = 16 -dtype = torch.bfloat16 - - -def get_timestamp(): - return datetime.now().strftime("%Y%m%d_%H%M%S") - - -def get_csv_path(prefix): - """Generate CSV path in current execution directory with timestamp""" - return f"{prefix}_{get_timestamp()}.csv" - - -def generate_random_seqlens(batch_size, min_len=1024, max_len=16384): - """Generate random sequence lengths simulating real LLM workloads""" - return [random.randint(min_len, max_len) for _ in range(batch_size)] - - -def setup_workspace(): - """Create workspace buffer for FlashInfer""" - return torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda") - - -def setup_paged_kv_indptr(batch_size, seq_lens): - """Setup paged KV cache indptr and last_page_len""" - seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32) - seq_lens_blocks = torch.ceil(seq_lens_tensor / page_block_size).int() - kv_indptr = torch.cat([torch.tensor([0]), torch.cumsum(seq_lens_blocks, 0)], dim=0).int() - num_blocks = kv_indptr[-1].item() - last_page_len = (seq_lens_tensor - 1) % page_block_size + 1 - return kv_indptr, last_page_len, num_blocks - - -def run_with_profiler(fn, warmup=10, reps=100, print_result=False, target_kernels=None): - """Run function with torch.profiler and return sum of specific kernel times in ms""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - - with torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CUDA], - record_shapes=False, - profile_memory=False, - with_stack=False, - ) as prof: - for _ in range(reps): - fn() - torch.cuda.synchronize() - - if print_result: - print(prof.key_averages().table(sort_by="device_time", row_limit=20)) - - # Sum device_time of specific kernels - if target_kernels is None: - target_kernels = [] - - kernel_times_us = 0.0 - for evt in prof.key_averages(): - if any(k in evt.key for k in target_kernels): - kernel_times_us += evt.device_time - - ms = kernel_times_us / 1e3 - return ms - - -def compute_reps(batch_size, seq_len, head_dim, base_reps=100): - """Dynamically compute repetition count based on workload size""" - # Estimate workload: batch_size * seq_len * head_dim - workload = batch_size * seq_len * head_dim - if workload < 1e5: # tiny workload - return base_reps - elif workload < 1e6: # small workload - return base_reps // 2 - elif workload < 1e7: # medium workload - return base_reps // 4 - elif workload < 1e8: # large workload - return base_reps // 8 - elif workload < 1e9: # very large workload - return base_reps // 16 - else: # huge workload - return base_reps // 32 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md deleted file mode 100644 index 59ee384..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/starter/示例冒烟代码.md +++ /dev/null @@ -1,209 +0,0 @@ -# 示例冒烟代码 - -```c++ -#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, - int64_t exact_len) { - 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 * exact_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 % exact_len; - const int64_t batch = work / exact_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); - } - } -} - -__global__ void prefix_mean_kernel( - const __nv_bfloat16* __restrict__ v, - __nv_bfloat16* __restrict__ output, - const int32_t* __restrict__ qo_indptr, - const int32_t* __restrict__ kv_indptr, - int64_t batch_size, - int64_t seq_len, - int64_t num_qo_heads, - int64_t num_kv_heads, - int64_t head_dim_vo) { - int64_t work = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = batch_size * num_kv_heads * head_dim_vo; - if (work >= total) return; - - const int64_t d = work % head_dim_vo; - work /= head_dim_vo; - const int64_t kv_head = work % num_kv_heads; - const int64_t batch = work / num_kv_heads; - const int64_t group = num_qo_heads / num_kv_heads; - const int64_t qo_begin = qo_indptr[batch]; - const int64_t kv_begin = kv_indptr[batch]; - - float sum = 0.0f; - for (int64_t t = 0; t < seq_len; ++t) { - const int64_t kv_row = kv_begin + t; - sum += __bfloat162float(v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]); - const __nv_bfloat16 mean = __float2bfloat16(sum / static_cast(t + 1)); - const int64_t out_row = qo_begin + t; - for (int64_t g = 0; g < group; ++g) { - const int64_t qo_head = kv_head * group + g; - output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean; - } - } -} - -} // 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; - - int64_t exact_len = seq_len; - if ((batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)) { - exact_len = 1024; - const int64_t mean_work = batch_size * num_kv_heads * head_dim_vo; - const int mean_blocks = static_cast((mean_work + kThreads - 1) / kThreads); - prefix_mean_kernel<<>>( - v, output, qo_indptr, kv_indptr, batch_size, seq_len, - num_qo_heads, num_kv_heads, head_dim_vo); - } - - const int64_t total = batch_size * exact_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, exact_len); -} -``` - -# run_kernel示例 - -```c++ -#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 -); -``` - diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json deleted file mode 100644 index 63ccef0..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/meta.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "id": 193, - "displayId": 20001, - "type": "Traditional", - "isPublic": false, - "locales": [ - "zh_CN" - ], - "samples": [ - { - "inputData": "1\n", - "outputData": "" - } - ], - "problemTagIds": [] -} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py deleted file mode 100644 index 5b9e146..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/testcase_config.py +++ /dev/null @@ -1,307 +0,0 @@ -from __future__ import annotations - - -HEAD_DIM_CONFIGS = [(128, 128)] -BATCH_SIZES = [1, 4, 16] -SEQ_LENS = [1024, 4096, 8192, 16384] -NUM_QO_HEADS = 32 -NUM_KV_HEADS = 4 -CAUSAL = 1 - - -def _build_cases(): - cases = [] - for head_dim_qk, head_dim_vo in HEAD_DIM_CONFIGS: - for batch_size in BATCH_SIZES: - for seq_len in SEQ_LENS: - cases.append( - ( - batch_size, - seq_len, - NUM_QO_HEADS, - NUM_KV_HEADS, - head_dim_qk, - head_dim_vo, - CAUSAL, - ) - ) - return cases - - -TESTCASES = _build_cases() - - -def getNumOfTestcases() -> int: - return len(TESTCASES) - - -try: - from pathlib import Path - from typing import List, Tuple, Union - import sys - - import torch - - KernelArg = Union[torch.Tensor, int, float] - CURRENT_CASE = None - - def _ensure_flashinfer_importable(): - try: - import flashinfer # noqa: F401 - - return - except ImportError: - pass - - here = Path(__file__).resolve() - for parent in here.parents: - candidate = parent / "McFlashInfer" - if (candidate / "flashinfer").is_dir(): - sys.path.insert(0, str(candidate)) - return - - def _get_testcase_index() -> int: - try: - raw = input().strip() - except EOFError: - return 0 - if raw == "": - return 0 - try: - testcase_id = int(raw.split()[0]) - except ValueError: - return 0 - if 1 <= testcase_id <= len(TESTCASES): - return testcase_id - 1 - if 0 <= testcase_id < len(TESTCASES): - return testcase_id - return 0 - - def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: - workload = batch_size * seq_len * head_dim - if workload < 1e5: - return base_reps - if workload < 1e6: - return base_reps // 2 - if workload < 1e7: - return base_reps // 4 - if workload < 1e8: - return base_reps // 8 - if workload < 1e9: - return base_reps // 16 - return base_reps // 32 - - def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: - testcase_id = _get_testcase_index() - global CURRENT_CASE - batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal = TESTCASES[testcase_id] - CURRENT_CASE = ( - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim_qk, - head_dim_vo, - causal, - 20260610 + testcase_id, - ) - qo_len = batch_size * seq_len - kv_len = batch_size * seq_len - warmup = 3 - iters = max(1, _compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo)) - return [ - (qo_len, num_qo_heads, head_dim_qk), - (kv_len, num_kv_heads, head_dim_qk), - (kv_len, num_kv_heads, head_dim_vo), - (qo_len, num_qo_heads, head_dim_vo), - (batch_size + 1,), - (batch_size + 1,), - (), (), (), (), (), (), (), - ], (warmup, iters) - - def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: - del testcase_sizes - batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal, seed = CURRENT_CASE - gen = torch.Generator(device=device) - gen.manual_seed(seed) - dtype = torch.bfloat16 - qo_len = batch_size * seq_len - q = torch.rand( - qo_len, - num_qo_heads, - head_dim_qk, - dtype=dtype, - device=device, - generator=gen, - ).contiguous() - kv_len = batch_size * seq_len - k = torch.rand( - kv_len, - num_kv_heads, - head_dim_qk, - dtype=dtype, - device=device, - generator=gen, - ).contiguous() - v = torch.rand( - kv_len, - num_kv_heads, - head_dim_vo, - dtype=dtype, - device=device, - generator=gen, - ).contiguous() - output = torch.empty( - qo_len, - num_qo_heads, - head_dim_vo, - dtype=dtype, - device=device, - ) - qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len - kv_indptr = qo_indptr.clone() - return [ - q, - k, - v, - output, - qo_indptr, - kv_indptr, - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim_qk, - head_dim_vo, - causal, - ] - - def baseline( - q, - k, - v, - output, - qo_indptr, - kv_indptr, - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim_qk, - head_dim_vo, - causal, - ): - _ensure_flashinfer_importable() - import flashinfer - - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) - wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( - workspace_buffer, - kv_layout="NHD", - backend="auto", - ) - wrapper.plan( - qo_indptr, - kv_indptr, - int(num_qo_heads), - int(num_kv_heads), - int(head_dim_qk), - int(head_dim_vo), - causal=bool(causal), - q_data_type=torch.bfloat16, - kv_data_type=torch.bfloat16, - ) - wrapper.run(q, k, v, out=output) - return [ - q, - k, - v, - output, - qo_indptr, - kv_indptr, - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim_qk, - head_dim_vo, - causal, - ] - - def check( - testcase_sizes, - original_input_tensors, - target_kernel_input_tensors, - baseline_input_tensors, - rtol=1e-2, - atol=1e-2, - ) -> bool: - del testcase_sizes, original_input_tensors - output_t = target_kernel_input_tensors[3] - output_ref = baseline_input_tensors[3] - if output_t.shape != output_ref.shape: - print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) - return False - if output_t.dtype != output_ref.dtype: - print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) - return False - if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): - diff = (output_t.float() - output_ref.float()).abs() - print( - f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " - f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", - file=sys.stderr, - ) - return False - return True -except Exception: - pass - - -INPUT_CLASS = [ - "INPUT", - "INPUT", - "INPUT", - "OUTPUT", - "INPUT", - "INPUT", - "INPUT", - "INPUT", - "INPUT", - "INPUT", - "INPUT", - "INPUT", - "INPUT", -] - - -def getWorkload(testcase_sizes) -> dict: - raw_sizes = testcase_sizes[0] if isinstance(testcase_sizes, tuple) and len(testcase_sizes) == 2 else testcase_sizes - q_shape, k_shape, v_shape, output_shape, qo_indptr_shape, kv_indptr_shape = raw_sizes[:6] - qo_len, num_qo_heads, head_dim_qk = q_shape - kv_len, num_kv_heads, k_dim = k_shape - v_len, v_heads, head_dim_vo = v_shape - assert k_dim == head_dim_qk - assert v_len == kv_len - assert v_heads == num_kv_heads - assert qo_len == kv_len - assert output_shape == (qo_len, num_qo_heads, head_dim_vo) - assert qo_indptr_shape == kv_indptr_shape - batch_size = qo_indptr_shape[0] - 1 - seq_len = kv_len // batch_size - flops = batch_size * seq_len * seq_len * num_qo_heads * (head_dim_qk + head_dim_vo) - memory_bytes = ( - qo_len * num_qo_heads * head_dim_qk * 2 - + kv_len * num_kv_heads * head_dim_qk * 2 - + kv_len * num_kv_heads * head_dim_vo * 2 - + qo_len * num_qo_heads * head_dim_vo * 2 - + (batch_size + 1) * 4 * 2 - ) - return { - "flops": flops, - "memory_bytes": memory_bytes, - "dtype": "bf16", - } - - -DESIGNED_VRAM_SIZE = 48 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md deleted file mode 100644 index 4f97990..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/00_题目描述.md +++ /dev/null @@ -1,23 +0,0 @@ ---- -sectionTitle: "题目描述" -type: "Text" ---- -你需要实现 FlashInfer ragged KV cache prefill 的CUDA C++前向算子。 - -本题输入采用 FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` 的 ragged `NHD` 布局。每个 batch 中有 `seq_len` 个 query token，KV cache 中也有 `seq_len` 个 token： - -其中 query heads 采用 GQA 布局：`num_qo_heads` 个 query/output heads 共享 `num_kv_heads` 个 KV heads，`G = num_qo_heads / num_kv_heads`。 - -评测程序会调用你提交代码中的 `run_kernel` 函数。你需要根据 `qo_indptr` 和 `kv_indptr` 读取 ragged Q/K/V，并将结果写入 `output`。 - -baseline 使用 FlashInfer ragged prefill 的 Python API： - -```python -wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(workspace, kv_layout="NHD", backend="auto") -wrapper.plan(qo_indptr, kv_indptr, num_qo_heads, num_kv_heads, - head_dim_qk, head_dim_vo, causal=True, - q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16) -wrapper.run(q, k, v, out=output) -``` - -如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md deleted file mode 100644 index 06f1703..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定cuda.md +++ /dev/null @@ -1,41 +0,0 @@ ---- -sectionTitle: "接口约定" -type: "codeSample" -lang: "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 `(batch_size * seq_len, num_qo_heads, head_dim_qk)`，连续 `bf16` -* `k`：key tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_qk)`，连续 `bf16` -* `v`：value tensor，shape `(batch_size * seq_len, num_kv_heads, head_dim_vo)`，连续 `bf16` -* `output`：输出缓冲区，shape `(batch_size * seq_len, 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` -* `causal`：是否启用 causal mask，评测中固定为 `1` - -本题测试中 `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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md deleted file mode 100644 index 35791da..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定tilelang.md +++ /dev/null @@ -1,52 +0,0 @@ ---- -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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md deleted file mode 100644 index d717add..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/01_接口约定triton.md +++ /dev/null @@ -1,41 +0,0 @@ ---- -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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md deleted file mode 100644 index 1e57f2e..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/02_数据范围.md +++ /dev/null @@ -1,26 +0,0 @@ ---- -sectionTitle: "输入格式" -type: "Text" ---- -## 数据范围 - -- 数据类型：`q/kv_data/output` 均为 `bfloat16` -- KV layout：`NHD` -- `num_qo_heads = 32` -- `num_kv_heads = 4` -- `page_block_size = 16` -- `causal = 0` -- `head_dim` 固定为 `128` -- `batch_size` 取值为 `1, 4, 16` -- `seq_len` 取值为 `1024, 4096, 8192, 16384` - -测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即： - -```python -for head_dim in [128]: - for batch_size in [1, 4, 16]: - for seq_len in [1024, 4096, 8192, 16384]: - ... -``` - -输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md deleted file mode 100644 index b4417ae..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/03_参考实现.md +++ /dev/null @@ -1,28 +0,0 @@ -pytorch参考实现： - -```python -def baseline(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): - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) - wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( - workspace_buffer, - 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) -``` - diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt deleted file mode 100644 index 1c5167b..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20001/zh_CN/_title.txt +++ /dev/null @@ -1 +0,0 @@ -FlashInfer Ragged Prefill diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json deleted file mode 100644 index cd4684c..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/meta.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "id": 194, - "displayId": 20002, - "type": "Traditional", - "isPublic": false, - "locales": [ - "zh_CN" - ], - "samples": [ - { - "inputData": "1\n", - "outputData": "" - } - ], - "problemTagIds": [] -} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py deleted file mode 100644 index 4066882..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/testcase_config.py +++ /dev/null @@ -1,273 +0,0 @@ -from __future__ import annotations - - -HEAD_DIMS = [128] -BATCH_SIZES = [1, 4, 16] -SEQ_LENS = [1024, 4096, 8192, 16384] -NUM_QO_HEADS = 32 -PAGE_BLOCK_SIZE = 16 -CAUSAL = 0 - - -def _build_cases(): - cases = [] - for head_dim in HEAD_DIMS: - for batch_size in BATCH_SIZES: - for seq_len in SEQ_LENS: - num_kv_heads = 8 if head_dim == 64 else 4 - cases.append( - ( - batch_size, - seq_len, - NUM_QO_HEADS, - num_kv_heads, - head_dim, - PAGE_BLOCK_SIZE, - CAUSAL, - ) - ) - return cases - - -TESTCASES = _build_cases() - - -def getNumOfTestcases() -> int: - return len(TESTCASES) - - -try: - from pathlib import Path - from typing import List, Tuple, Union - import sys - - import torch - - KernelArg = Union[torch.Tensor, int, float] - CURRENT_CASE = None - - def _ensure_flashinfer_importable(): - try: - import flashinfer # noqa: F401 - - return - except ImportError: - pass - - here = Path(__file__).resolve() - for parent in here.parents: - candidate = parent / "McFlashInfer" - if (candidate / "flashinfer").is_dir(): - sys.path.insert(0, str(candidate)) - return - - def _get_testcase_index() -> int: - try: - raw = input().strip() - except EOFError: - return 0 - if raw == "": - return 0 - try: - testcase_id = int(raw.split()[0]) - except ValueError: - return 0 - if 1 <= testcase_id <= len(TESTCASES): - return testcase_id - 1 - if 0 <= testcase_id < len(TESTCASES): - return testcase_id - return 0 - - def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: - workload = batch_size * seq_len * head_dim - if workload < 1e5: - return base_reps - if workload < 1e6: - return base_reps // 2 - if workload < 1e7: - return base_reps // 4 - if workload < 1e8: - return base_reps // 8 - if workload < 1e9: - return base_reps // 16 - return base_reps // 32 - - def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str): - seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) - seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor") - kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device) - kv_indptr[0] = 0 - kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0) - num_blocks = int(kv_indptr[-1].item()) - last_page_len = (seq_lens - 1) % page_block_size + 1 - return kv_indptr, last_page_len, num_blocks - - def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: - testcase_id = _get_testcase_index() - global CURRENT_CASE - batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim, page_block_size, causal = TESTCASES[testcase_id] - num_blocks = batch_size * ((seq_len + page_block_size - 1) // page_block_size) - qo_len = batch_size * seq_len - CURRENT_CASE = ( - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim, - page_block_size, - causal, - num_blocks, - 20260620 + testcase_id, - ) - warmup = 3 - iters = max(1, _compute_reps(batch_size, seq_len, head_dim)) - return [ - (qo_len, num_qo_heads, head_dim), - (num_blocks, 2, page_block_size, num_kv_heads, head_dim), - (qo_len, num_qo_heads, head_dim), - (batch_size + 1,), - (batch_size + 1,), - (num_blocks,), - (batch_size,), - (), (), (), (), (), (), (), - ], (warmup, iters) - - def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: - del testcase_sizes - ( - batch_size, - seq_len, - num_qo_heads, - num_kv_heads, - head_dim, - page_block_size, - causal, - num_blocks, - seed, - ) = CURRENT_CASE - gen = torch.Generator(device=device) - gen.manual_seed(seed) - dtype = torch.bfloat16 - qo_len = batch_size * seq_len - q = torch.rand(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous() - kv_data = torch.randn( - num_blocks, - 2, - page_block_size, - num_kv_heads, - head_dim, - dtype=dtype, - device=device, - generator=gen, - ).contiguous() - output = torch.empty(qo_len, num_qo_heads, head_dim, dtype=dtype, device=device) - qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) * seq_len - kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr( - batch_size, - seq_len, - page_block_size, - device, - ) - assert num_blocks_check == num_blocks - kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device) - return [ - 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, - ] - - def baseline( - 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, - ): - _ensure_flashinfer_importable() - import flashinfer - - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) - wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( - workspace_buffer, - kv_layout="NHD", - backend="auto", - ) - wrapper.plan( - qo_indptr, - kv_indptr, - kv_indices, - last_page_len, - int(num_qo_heads), - int(num_kv_heads), - int(head_dim), - int(page_block_size), - causal=bool(causal), - q_data_type=torch.bfloat16, - kv_data_type=torch.bfloat16, - ) - wrapper.run(q, kv_data, out=output) - return [ - 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, - ] - - def check( - testcase_sizes, - original_input_tensors, - target_kernel_input_tensors, - baseline_input_tensors, - rtol=1e-2, - atol=1e-2, - ) -> bool: - del testcase_sizes, original_input_tensors - output_t = target_kernel_input_tensors[2] - output_ref = baseline_input_tensors[2] - if output_t.shape != output_ref.shape: - print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) - return False - if output_t.dtype != output_ref.dtype: - print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) - return False - if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): - diff = (output_t.float() - output_ref.float()).abs() - print( - f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " - f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", - file=sys.stderr, - ) - return False - return True -except Exception: - pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md deleted file mode 100644 index 3e12eca..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/00_题目描述.md +++ /dev/null @@ -1,22 +0,0 @@ ---- -sectionTitle: "题目描述" -type: "Text" ---- -你需要实现 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) -``` - -如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md deleted file mode 100644 index 48fbe06..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定cuda.md +++ /dev/null @@ -1,42 +0,0 @@ ---- -sectionTitle: "接口约定" -type: "codeSample" -lang: "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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md deleted file mode 100644 index 098c6da..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定tilelang.md +++ /dev/null @@ -1,55 +0,0 @@ ---- -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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md deleted file mode 100644 index d32ea1b..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/01_接口约定triton.md +++ /dev/null @@ -1,44 +0,0 @@ ---- -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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md deleted file mode 100644 index 1e57f2e..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/02_数据范围.md +++ /dev/null @@ -1,26 +0,0 @@ ---- -sectionTitle: "输入格式" -type: "Text" ---- -## 数据范围 - -- 数据类型：`q/kv_data/output` 均为 `bfloat16` -- KV layout：`NHD` -- `num_qo_heads = 32` -- `num_kv_heads = 4` -- `page_block_size = 16` -- `causal = 0` -- `head_dim` 固定为 `128` -- `batch_size` 取值为 `1, 4, 16` -- `seq_len` 取值为 `1024, 4096, 8192, 16384` - -测试点顺序与 `McFlashInfer/benchmarks/bench_batch_prefill_paged.py` 中的 cases 一致，即： - -```python -for head_dim in [128]: - for batch_size in [1, 4, 16]: - for seq_len in [1024, 4096, 8192, 16384]: - ... -``` - -输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md deleted file mode 100644 index a55a635..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/03_参考实现.md +++ /dev/null @@ -1,31 +0,0 @@ ---- -sectionTitle: "输出格式" -type: "Text" ---- -PyTorch 参考实现 - -```python -def baseline(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): - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) - wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( - workspace_buffer, - 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) -``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt deleted file mode 100644 index 34e4801..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20002/zh_CN/_title.txt +++ /dev/null @@ -1 +0,0 @@ -FlashInfer Paged Prefill diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json deleted file mode 100644 index f9593e3..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/meta.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "id": 195, - "displayId": 20003, - "type": "Traditional", - "isPublic": false, - "locales": [ - "zh_CN" - ], - "samples": [ - { - "inputData": "1\n", - "outputData": "" - } - ], - "problemTagIds": [] -} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py deleted file mode 100644 index ee3ae71..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/testcase_config.py +++ /dev/null @@ -1,245 +0,0 @@ -from __future__ import annotations - - -HEAD_DIM_CKV = 512 -HEAD_DIM_KPE = 64 -NUM_HEADS_LIST = [64, 128] -BATCH_SIZES = [1, 4, 16] -SEQ_LENS = [1024, 4096, 8192, 16384] -PAGE_SIZE = 1 -CAUSAL = 0 - - -def _build_cases(): - cases = [] - for num_heads in NUM_HEADS_LIST: - for batch_size in BATCH_SIZES: - for seq_len in SEQ_LENS: - cases.append( - ( - batch_size, - seq_len, - num_heads, - HEAD_DIM_CKV, - HEAD_DIM_KPE, - PAGE_SIZE, - CAUSAL, - ) - ) - return cases - - -TESTCASES = _build_cases() - - -def getNumOfTestcases() -> int: - return len(TESTCASES) - - -try: - from pathlib import Path - from typing import List, Tuple, Union - import sys - - import torch - - KernelArg = Union[torch.Tensor, int, float] - CURRENT_CASE = None - - def _ensure_flashinfer_importable(): - try: - import flashinfer # noqa: F401 - - return - except ImportError: - pass - - here = Path(__file__).resolve() - for parent in here.parents: - candidate = parent / "McFlashInfer" - if (candidate / "flashinfer").is_dir(): - sys.path.insert(0, str(candidate)) - return - - def _get_testcase_index() -> int: - try: - raw = input().strip() - except EOFError: - return 0 - if raw == "": - return 0 - try: - testcase_id = int(raw.split()[0]) - except ValueError: - return 0 - if 1 <= testcase_id <= len(TESTCASES): - return testcase_id - 1 - if 0 <= testcase_id < len(TESTCASES): - return testcase_id - return 0 - - def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: - workload = batch_size * seq_len * head_dim - if workload < 1e5: - return base_reps - if workload < 1e6: - return base_reps // 2 - if workload < 1e7: - return base_reps // 4 - if workload < 1e8: - return base_reps // 8 - if workload < 1e9: - return base_reps // 16 - return base_reps // 32 - - def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: - testcase_id = _get_testcase_index() - global CURRENT_CASE - batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal = TESTCASES[testcase_id] - CURRENT_CASE = ( - batch_size, - seq_len, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - causal, - 20260630 + testcase_id, - ) - warmup = 3 - iters = max(1, _compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe)) - kv_len = batch_size * seq_len - return [ - (batch_size, num_heads, head_dim_ckv), - (batch_size, num_heads, head_dim_kpe), - (kv_len, 1, head_dim_ckv), - (kv_len, 1, head_dim_kpe), - (batch_size, num_heads, head_dim_ckv), - (batch_size + 1,), - (batch_size + 1,), - (kv_len,), - (batch_size,), - (), (), (), (), (), (), (), - ], (warmup, iters) - - def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: - del testcase_sizes - batch_size, seq_len, num_heads, head_dim_ckv, head_dim_kpe, page_size, causal, seed = CURRENT_CASE - gen = torch.Generator(device=device) - gen.manual_seed(seed) - dtype = torch.bfloat16 - kv_len = batch_size * seq_len - q_nope = torch.randn(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous() - q_pe = torch.zeros(batch_size, num_heads, head_dim_kpe, dtype=dtype, device=device) - ckv = torch.randn(kv_len, 1, head_dim_ckv, dtype=dtype, device=device, generator=gen).contiguous() - kpe = torch.zeros(kv_len, 1, head_dim_kpe, dtype=dtype, device=device) - output = torch.empty(batch_size, num_heads, head_dim_ckv, dtype=dtype, device=device) - q_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device=device) - kv_indptr = q_indptr * seq_len - kv_indices = torch.arange(0, kv_len, dtype=torch.int32, device=device) - kv_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) - return [ - 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, - ] - - def baseline( - 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, - ): - _ensure_flashinfer_importable() - import flashinfer - - sm_scale = 1.0 / ((int(head_dim_ckv) + int(head_dim_kpe)) ** 0.5) - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device) - wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") - wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - kv_lens, - int(num_heads), - int(head_dim_ckv), - int(head_dim_kpe), - int(page_size), - bool(causal), - sm_scale, - q_nope.dtype, - ckv.dtype, - ) - wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) - return [ - 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, - ] - - def check( - testcase_sizes, - original_input_tensors, - target_kernel_input_tensors, - baseline_input_tensors, - rtol=1e-2, - atol=1e-2, - ) -> bool: - del testcase_sizes, original_input_tensors - output_t = target_kernel_input_tensors[4] - output_ref = baseline_input_tensors[4] - if output_t.shape != output_ref.shape: - print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) - return False - if output_t.dtype != output_ref.dtype: - print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) - return False - if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): - diff = (output_t.float() - output_ref.float()).abs() - print( - f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " - f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", - file=sys.stderr, - ) - return False - return True -except Exception: - pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md deleted file mode 100644 index c3bf938..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/00_题目描述.md +++ /dev/null @@ -1,22 +0,0 @@ ---- -sectionTitle: "题目描述" -type: "Text" ---- -你需要实现 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) -``` - -如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md deleted file mode 100644 index 7cd07bb..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定cuda.md +++ /dev/null @@ -1,46 +0,0 @@ ---- -sectionTitle: "接口约定" -type: "codeSample" -lang: "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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md deleted file mode 100644 index 7196b09..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定tilelang.md +++ /dev/null @@ -1,57 +0,0 @@ ---- -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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md deleted file mode 100644 index 43be38b..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/01_接口约定triton.md +++ /dev/null @@ -1,45 +0,0 @@ ---- -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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md deleted file mode 100644 index 8919075..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/02_数据范围.md +++ /dev/null @@ -1,26 +0,0 @@ ---- -sectionTitle: "输入格式" -type: "Text" ---- -## 数据范围 - -- 数据类型：输入输出均为 `bfloat16` -- `head_dim_ckv = 512` -- `head_dim_kpe = 64` -- `page_size = 1` -- `causal = 0` -- `num_heads` 取值为 `64, 128` -- `batch_size` 取值为 `1, 4, 16` -- `seq_len` 取值为 `1024, 4096, 8192, 16384` -- `sm_scale = 1 / sqrt(head_dim_ckv + head_dim_kpe)` - -测试点顺序与 `McFlashInfer/benchmarks/bench_batch_mla.py` 中的 cases 一致，即： - -```python -for num_heads in [64, 128]: - for batch_size in [1, 4, 16]: - for seq_len in [1024, 4096, 8192, 16384]: - ... -``` - -输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md deleted file mode 100644 index 23285e8..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/03_参考实现.md +++ /dev/null @@ -1,30 +0,0 @@ ---- -sectionTitle: "输出格式" -type: "Text" ---- -PyTorch 参考实现 - -```python -def baseline(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): - sm_scale = 1.0 / ((head_dim_ckv + head_dim_kpe) ** 0.5) - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q_nope.device) - wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") - wrapper.plan( - q_indptr, - kv_indptr, - kv_indices, - kv_lens, - num_heads, - head_dim_ckv, - head_dim_kpe, - page_size, - bool(causal), - sm_scale, - q_nope.dtype, - ckv.dtype, - ) - wrapper.run(q_nope, q_pe, ckv, kpe, out=output, return_lse=False) -``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt deleted file mode 100644 index 9f19bce..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20003/zh_CN/_title.txt +++ /dev/null @@ -1 +0,0 @@ -FlashInfer MLA Paged Attention diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json deleted file mode 100644 index 035e3ac..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/meta.json +++ /dev/null @@ -1,16 +0,0 @@ -{ - "id": 196, - "displayId": 20004, - "type": "Traditional", - "isPublic": false, - "locales": [ - "zh_CN" - ], - "samples": [ - { - "inputData": "1\n", - "outputData": "" - } - ], - "problemTagIds": [] -} diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py deleted file mode 100644 index 3273aa3..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/testcase_config.py +++ /dev/null @@ -1,248 +0,0 @@ -from __future__ import annotations - - -BATCH_SIZES = [1, 2, 4, 8, 16] -SEQ_LENS_KV = [512, 1024, 2048, 4096, 8192, 16384] -HEAD_DIMS = [128] -NUM_QO_HEADS = 32 -PAGE_BLOCK_SIZE = 16 - - -def _build_cases(): - cases = [] - for batch_size in BATCH_SIZES: - for seq_len_kv in SEQ_LENS_KV: - for head_dim in HEAD_DIMS: - num_kv_heads = 8 if head_dim == 64 else 4 - cases.append( - ( - batch_size, - seq_len_kv, - NUM_QO_HEADS, - num_kv_heads, - head_dim, - PAGE_BLOCK_SIZE, - ) - ) - return cases - - -TESTCASES = _build_cases() - - -def getNumOfTestcases() -> int: - return len(TESTCASES) - - -try: - from pathlib import Path - from typing import List, Tuple, Union - import sys - - import torch - - KernelArg = Union[torch.Tensor, int, float] - CURRENT_CASE = None - - def _ensure_flashinfer_importable(): - try: - import flashinfer # noqa: F401 - - return - except ImportError: - pass - - here = Path(__file__).resolve() - for parent in here.parents: - candidate = parent / "McFlashInfer" - if (candidate / "flashinfer").is_dir(): - sys.path.insert(0, str(candidate)) - return - - def _get_testcase_index() -> int: - try: - raw = input().strip() - except EOFError: - return 0 - if raw == "": - return 0 - try: - testcase_id = int(raw.split()[0]) - except ValueError: - return 0 - if 1 <= testcase_id <= len(TESTCASES): - return testcase_id - 1 - if 0 <= testcase_id < len(TESTCASES): - return testcase_id - return 0 - - def _compute_reps(batch_size: int, seq_len: int, head_dim: int, base_reps: int = 100) -> int: - workload = batch_size * seq_len * head_dim - if workload < 1e5: - return base_reps - if workload < 1e6: - return base_reps // 2 - if workload < 1e7: - return base_reps // 4 - if workload < 1e8: - return base_reps // 8 - if workload < 1e9: - return base_reps // 16 - return base_reps // 32 - - def _setup_paged_kv_indptr(batch_size: int, seq_len: int, page_block_size: int, device: str): - seq_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device=device) - seq_lens_blocks = torch.div(seq_lens + page_block_size - 1, page_block_size, rounding_mode="floor") - kv_indptr = torch.empty((batch_size + 1,), dtype=torch.int32, device=device) - kv_indptr[0] = 0 - kv_indptr[1:] = torch.cumsum(seq_lens_blocks, dim=0) - num_blocks = int(kv_indptr[-1].item()) - last_page_len = (seq_lens - 1) % page_block_size + 1 - return kv_indptr, last_page_len, num_blocks - - def getTestCaseSize() -> Tuple[List[Tuple[int, ...]], Tuple[int, int]]: - testcase_id = _get_testcase_index() - global CURRENT_CASE - batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size = TESTCASES[testcase_id] - num_blocks = batch_size * ((seq_len_kv + page_block_size - 1) // page_block_size) - CURRENT_CASE = ( - batch_size, - seq_len_kv, - num_qo_heads, - num_kv_heads, - head_dim, - page_block_size, - num_blocks, - 20260710 + testcase_id, - ) - warmup = 3 - iters = max(1, _compute_reps(batch_size, seq_len_kv, head_dim)) - return [ - (batch_size, num_qo_heads, head_dim), - (num_blocks, 2, page_block_size, num_kv_heads, head_dim), - (batch_size, num_qo_heads, head_dim), - (batch_size + 1,), - (num_blocks,), - (batch_size,), - (), (), (), (), (), (), - ], (warmup, iters) - - def genTestCase(testcase_sizes, device: str = "cuda") -> List[KernelArg]: - del testcase_sizes - batch_size, seq_len_kv, num_qo_heads, num_kv_heads, head_dim, page_block_size, num_blocks, seed = CURRENT_CASE - gen = torch.Generator(device=device) - gen.manual_seed(seed) - dtype = torch.bfloat16 - q = torch.rand(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device, generator=gen).contiguous() - kv_data = torch.randn( - num_blocks, - 2, - page_block_size, - num_kv_heads, - head_dim, - dtype=dtype, - device=device, - generator=gen, - ).contiguous() - output = torch.empty(batch_size, num_qo_heads, head_dim, dtype=dtype, device=device) - kv_indptr, last_page_len, num_blocks_check = _setup_paged_kv_indptr( - batch_size, - seq_len_kv, - page_block_size, - device, - ) - assert num_blocks_check == num_blocks - kv_indices = torch.arange(num_blocks, dtype=torch.int32, device=device) - return [ - 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, - ] - - def baseline( - 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, - ): - _ensure_flashinfer_importable() - import flashinfer - - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) - wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( - workspace_buffer, - kv_layout="NHD", - use_tensor_cores=True, - ) - wrapper.plan( - kv_indptr, - kv_indices, - last_page_len, - int(num_qo_heads), - int(num_kv_heads), - int(head_dim), - int(page_block_size), - data_type=torch.bfloat16, - q_data_type=torch.bfloat16, - ) - wrapper.run(q, kv_data, out=output) - return [ - 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, - ] - - def check( - testcase_sizes, - original_input_tensors, - target_kernel_input_tensors, - baseline_input_tensors, - rtol=1e-2, - atol=1e-2, - ) -> bool: - del testcase_sizes, original_input_tensors - output_t = target_kernel_input_tensors[2] - output_ref = baseline_input_tensors[2] - if output_t.shape != output_ref.shape: - print(f"[FAIL] shape mismatch: target {output_t.shape}, ref {output_ref.shape}", file=sys.stderr) - return False - if output_t.dtype != output_ref.dtype: - print(f"[FAIL] dtype mismatch: target {output_t.dtype}, ref {output_ref.dtype}", file=sys.stderr) - return False - if not torch.allclose(output_t.float(), output_ref.float(), rtol=rtol, atol=atol): - diff = (output_t.float() - output_ref.float()).abs() - print( - f"[FAIL] allclose failed: max_abs_diff={float(diff.max().item()):.6f}, " - f"mean_abs_diff={float(diff.mean().item()):.6f} (rtol={rtol}, atol={atol})", - file=sys.stderr, - ) - return False - return True -except Exception: - pass diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md deleted file mode 100644 index 74508b0..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/00_题目描述.md +++ /dev/null @@ -1,22 +0,0 @@ ---- -sectionTitle: "题目描述" -type: "Text" ---- -你需要实现 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) -``` - -如何提交代码详见[评测指南](/d/2)。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md deleted file mode 100644 index c498b62..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定cuda.md +++ /dev/null @@ -1,38 +0,0 @@ ---- -sectionTitle: "接口约定" -type: "codeSample" -lang: "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()` 或显式同步。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md deleted file mode 100644 index 8d0fd87..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定tilelang.md +++ /dev/null @@ -1,52 +0,0 @@ ---- -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大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md deleted file mode 100644 index d1751a9..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/01_接口约定triton.md +++ /dev/null @@ -1,41 +0,0 @@ ---- -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。 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md deleted file mode 100644 index 9fe7bee..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/02_数据范围.md +++ /dev/null @@ -1,26 +0,0 @@ ---- -sectionTitle: "输入格式" -type: "Text" ---- -## 数据范围 - -- 数据类型：`q/kv_data/output` 均为 `bfloat16` -- KV layout：`NHD` -- `seq_len_q = 1` -- `num_qo_heads = 32` -- `page_block_size = 16` -- `head_dim` 固定为 `128` -- `num_kv_heads = 4` -- `batch_size` 取值为 `1, 2, 4, 8, 16` -- `seq_len_kv` 取值为 `512, 1024, 2048, 4096, 8192, 16384` - -测试点顺序与 `McFlashInfer/benchmarks/bench_batch_decode.py` 中的 cases 一致，即： - -```python -for batch_size in [1, 2, 4, 8, 16]: - for seq_len_kv in [512, 1024, 2048, 4096, 8192, 16384]: - for head_dim in [128]: - ... -``` - -输出参与 `torch.allclose` 校验，容差为 `rtol=1e-2, atol=1e-2`。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md deleted file mode 100644 index 5d17935..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/03_参考实现.md +++ /dev/null @@ -1,29 +0,0 @@ ---- -sectionTitle: "输出格式" -type: "Text" ---- -PyTorch 参考实现 - -```python -def baseline(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): - workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) - wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( - workspace_buffer, - 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) -``` \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt deleted file mode 100644 index 552b751..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package backup/xpuoj_problem/problem_20004/zh_CN/_title.txt +++ /dev/null @@ -1 +0,0 @@ -FlashInfer Paged Decode