diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md index 826f9b4..b058ddf 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer关键算子迁移与优化.md @@ -1,30 +1,46 @@ # FlashInfer 关键算子迁移与优化 -## 一、教程定位 +## 1. 教程定位 -本教程是参赛训练课程的 **FlashInfer Baseline 入门** 模块，主要帮助用户快速跑通 FlashInfer 的最小可运行流程。完成本教程后，用户应能够完成源码编译、API 调用、正确性测试和 Benchmark 测试，并记录一份 baseline 性能结果，为后续算子优化提供对比基准。 +本教程是赛题二FlashInfer任务的“benchmark 性能基线与 XPU-OJ 提交衔接”模块，主要帮助学员跑通目标算子的 benchmark 脚本，理解原库 API、输入输出结构、性能指标和性能基线结果，并进一步读懂 XPU-OJ 题目包中的接口约定、测试数据、参考输出和精度要求。 -## 二、完成本模块你将能够 +需要特别说明： -1. 理解 FlashInfer Attention Kernel 的基本作用与适用场景； -2. 完成 FlashInfer 环境、工具链的准备与源码编译； -3. 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例； -4. 完成各算子在不同参数配置下的 Benchmark 测试； -5. 输出各算子的 Baseline 性能结果记录表，为后续算子优化提供对比基准。 -6. 理解 XPU-OJ 评测 `run_kernel` 接口与精度要求。 -7. 理解 Baseline 与 XPU-OJ 评测题包之间的关系，能够根据题包接口实现一个最小正确版 `run_kernel`。 +- 本教程不提供可直接提交的 OJ 参考实现源代码或标准答案代码。 +- 本教程仅提供冒烟级 starter 示例代码，用于验证环境、语言、提交链路和 `run_kernel(...)` 接口。 +- benchmark 脚本用于建立性能基线，不是最终提交物。 +- XPU-OJ 题包中的 `baseline()` 属于 OJ 后台参考实现，用于生成 `output_ref`，不是选手提交代码。 +- 选手最终需要自行实现 `run_kernel(...)`，并在正确性通过后继续优化性能。 + +完成本教程后，学员应能够跑通 benchmark 脚本，记录性能基线结果，读懂 XPU-OJ 题包，理解 OJ 的测试输入与参考实现，并完成一次冒烟级 OJ 提交。 + +## 2. 完成本模块你将能够 + +- 理解 FlashInfer Attention Kernel 的基本作用与适用场景； + +- 完成 FlashInfer 环境、工具链的准备与源码编译； + +- 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例； + +- 完成各算子在不同参数配置下的 Benchmark 测试； + +- 输出各算子的性能基线结果记录表，为后续算子优化提供对比基准。 + +- 理解 XPU-OJ 评测 `run_kernel` 接口与精度要求。 + +- 理解 benchmark 与 XPU-OJ 评测题包之间的关系，能够根据题包接口实现一个最小正确版 `run_kernel`。 -## 三、适用对象 +## 3. 适用对象 -**适合人群** +### 3.1 **适合人群** -* 参赛选手：需要完成 Baseline 入门模块，为后续算子优化做准备 +* 参赛选手：需要完成 Benchmark 入门模块，为后续算子优化做准备 * 软件开发者和Vibe Coding开发：希望从事AI相关行业开发，以及用智能体方式来做开发工作 * LLM 推理开发者：希望了解 FlashInfer Attention Kernel 的性能表现 * 算子优化工程师：希望基于MXMACA软件栈在沐曦国产 GPU 上做算子迁移和优化 -**前置基础** +### **3.2 前置基础** * Python 基础：能够运行和修改 Python 脚本 * PyTorch 基础：了解MXMACA软化栈的使用 @@ -32,11 +48,11 @@ * 了解 Attention 机制：理解 Q/K/V、KV Cache 等基本概念 -## 四、前置准备 +## 4. 前置准备 开始实战前，请确认你已经完成以下准备： -### 获得 GPU +### 4.1 获得 GPU 1. [点击获取算力券](https://developer.metax-tech.com/activities/6)，首次登录需要使用邮箱或者手机号进行注册 @@ -45,34 +61,36 @@ 3. 提交申请获得兑换码 4. 兑换算力和登陆平台： - + - [访问模力方舟官网](https://ai.gitee.com/)，进入费用中心 - 算力券 ，点击右上角 “兑换”。 - 进入算力容器，选择沐曦，租用算力，建议优先选 16G 显存 / 32G 显存，如下图 ![Weixin-Image_20260617180029](https://origin.picgo.net/2026/06/17/Weixin-Image_20260617180029_15_100c1e1fecb298bac54.png) - - - 创建实例。基础镜像：`maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64` + + - 创建实例。基础镜像：`PyTorch-Agent/2.8.0/Python 3.12/maca 3.7.2.1` + + ![pytorch agent](https://origin.picgo.net/2026/06/18/pytorch-agent70c0e9dce299e26c.png) - 选择工具-lab进入实例环境 - 步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用 `mx-smi` 命令查看 -### Python 环境 +### 4.2 Python 环境 ``` bash pip install flashinfer torch pandas numpy ``` -### OpenCode 安装 +### 4.3 OpenCode 安装 ``` bash curl -fsSL https://opencode.ai/install | bash opencode ``` -### 代码准备 +### 4.4 代码准备 测试脚本和 Benchmark 脚本: @@ -83,31 +101,31 @@ opencode | `python bench_batch_prefill_ragged.py` | 运行 Batch Prefill (Ragged KV Cache) 基准测试 | | `python bench_batch_mla.py` | 运行 MLA (Multi-head Latent Attention) 基准测试 | -## 五、知识预备 +## 5. 知识预备 -### LLM 推理阶段重要概念： +### 5.1 LLM 推理阶段重要概念 - **Prefill 阶段**：Prefill 阶段是指处理输入 prompt 的阶段 - + - 输入：用户一次性给出的完整 prompt，长度为 seq\_len - + - 计算：对 prompt 中的每个 token 并行计算注意力，生成第一个输出 token 及 KV cache - + - 特点：这是 **计算密集型（compute-bound）** 阶段，因为需要做完整的 `seq\_len * seq\_len` 注意力矩阵乘法 - **Decode 阶段**： - + - 每次只生成 1 个 token，利用 prefill 阶段填充好的 KV cache 做自回归生成 - + - **显存带宽密集型（memory-bound）**，瓶颈在从显存读取 KV cache 而非计算 **Prefill = 并行处理用户输入，Decode = 逐个生成回答 token** -## 六、项目实践 -- FlashInfer Baseline +## 6. 项目实践 -- FlashInfer Benchmark -**目标：** 在赛事镜像中完成 FlashInfer Ragged Prefill 算子的 Baseline Benchmark，理解从 Benchmark 到 XPU-OJ 评测提交的完整流程，为后续算子优化建立性能基线。 +**目标：** 在赛事镜像中完成 FlashInfer Ragged Prefill 算子的 Benchmark，理解从 Benchmark 到 XPU-OJ 评测提交的完整流程，为后续算子优化建立性能基线。 -### 在赛事镜像中运行 FlashInfer Baseline Benchmark +### 6.1 在赛事镜像中运行 FlashInfer Benchmark #### Step 1：检查运行环境 @@ -132,26 +150,30 @@ python -c "import torch; print(f'GPU available: {torch.cuda.is_available()}'); p # 检查依赖版本 python -c "import torch; print(f'PyTorch {torch.__version__}')" python -c "import einops; print('einops OK')" +# 安装必要依赖 +pip install pandas ``` **预期结果：** - `mx-smi` 显示沐曦 GPU 信息 - ![result mx smi](https://origin.picgo.net/2026/06/04/result-mx-smif86a3bed6681382e.png) + ![](https://origin.picgo.net/2026/06/18/-2026-06-18-164201---f51bf784081acbc6.png) - Python 环境正常 - ![result python](https://origin.picgo.net/2026/06/04/result-pythond15f856ddb84c649.png) + ![](https://origin.picgo.net/2026/06/18/-2026-06-18-164201----26f188cca92241bd2.png) - `torch.cuda.is_available()` 返回 `True` - ![result gpu available](https://origin.picgo.net/2026/06/04/result-gpu-available44e049addf5638fb.png) + ![](https://origin.picgo.net/2026/06/18/-2026-06-18-164201----360b71aedf7b90b21.png) - 所有依赖版本符合要求 - ![result dependency version](https://origin.picgo.net/2026/06/04/result-dependency-versionc17845898ee89434.png) - + ![](https://origin.picgo.net/2026/06/18/-2026-06-18-164201----4ffb68bf433f690a8.png) + + ![](https://origin.picgo.net/2026/06/18/-2026-06-18-164201----5fed48dda82423727.png) + **常见问题：** @@ -163,7 +185,7 @@ python -c "import einops; print('einops OK')" #### Step 2：进入项目目录 -**目标：** 进入本模块所需的源码目录[flashinfer_baseline](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%2Fbaselines%2Fflashinfer_baseline)。 +**目标：** 进入本模块所需的源码目录 [flashinfer_baseline](https://www.gitlink.org.cn/metax-maca/op_optimization/tree/master/基于AI Agent开发范式的国产GPU大模型推理算子库优化%2Foperator_task_package%2Fflashinfer_task_package%2FFlashInfer_Baseline)。 1. 克隆代码仓库 @@ -171,18 +193,19 @@ python -c "import einops; print('einops OK')" git clone https://gitlink.org.cn/metax-maca/op_optimization.git ``` -2. 准备flashinfer_baseline +2. 准备 FlashInfer_Baseline + + 从仓库根目录开始，在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `operator_task_package/flashinfer_task_package` 文件夹。可以将 `FlashInfer_Baseline` 整个目录复制到工作目录 `data/` 下。 - 从仓库根目录开始，在 `基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashinfer_baseline` 文件夹。可以将 `flashinfer_baseline` 整个目录复制到工作目录 `data/` 下。 ```bash mkdir data - cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline" data/ + cp -r "基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package" data/ ``` -3. 切换到 FlashInfer_Baseline 项目目录 +3. 切换到项目目录 `FlashInfer_Baseline` ```bash - cd data/flashinfer_baseline/FlashInfer_Baseline - ls -al + cd data/operator_task_package/flashinfer_task_package/FlashInfer_Baseline + ls ``` **预期结果：** @@ -226,7 +249,7 @@ python -c "from bench_common import setup_workspace, get_csv_path; print('脚本 #### Step 4：运行单算子 Benchmark 并查看测试结果 -**目标：** 执行基准测试，获取 Baseline 性能数据，查看并分析 Benchmark 输出结果。 +**目标：** 执行基准测试，获取性能基线数据，查看并分析 Benchmark 输出结果。 **操作：** 运行 Ragged Prefill 基准测试脚本，读取生成的 CSV 结果文件。 @@ -296,15 +319,15 @@ else: ... ``` -### XPU-OJ 在线评测教程 +### 6.2 XPU-OJ 在线评测教程 -#### Step 5: 从 Baseline 到 XPU-OJ 提交 +#### Step 5：从 Benchmark 到 XPU-OJ 提交 -##### 5.1：Baseline 与 XPU-OJ 的关系 +**Benchmark 与 XPU-OJ 的关系** -赛事镜像中的 Baseline Benchmark 和 XPU-OJ 在线评测任务不同。Baseline Benchmark 主要用于理解算子调用方式和建立性能基线；XPU-OJ 在线评测用于统一检查选手提交代码的正确性和性能。 +赛事镜像中的 Benchmark 和 XPU-OJ 在线评测任务不同。Benchmark 主要用于理解算子调用方式和建立性能基线；XPU-OJ 在线评测用于统一检查选手提交代码的正确性和性能。 -| 维度 | Baseline Benchmark | XPU-OJ 提交 | +| 维度 | Benchmark 脚本 | XPU-OJ 提交 | |------|---------------|------------| | **目的** | 理解算子接口、建立性能基线 | 统一环境下的正确性+性能评测 | | **接口形式** | Python API（`wrapper.plan()` + `wrapper.run()`） | C 接口（`extern "C" void run_kernel(...)`） | @@ -312,16 +335,15 @@ else: | **验证** | 无自动正确性校验 | 强制通过 `torch.allclose(rtol=1e-2, atol=1e-2)` | | **输出** | CSV 性能记录 | 排行榜得分 | -跑完 baseline 后，选手需要完成以下转换： +跑完 benchmark、建立性能基线后，选手需要完成以下转换： 1. 从 benchmark 脚本中理解目标 API，例如 BatchPrefillWithRaggedKVCacheWrapper； 2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口； 3. 对照题包中的输入 shape、数据范围和精度要求； 4. 编写自己的 `run_kernel(...)`； 5. 提交 OJ，先通过正确性； -6. 正确性通过后，再对比 baseline / OJ 耗时继续优化。 +6. 正确性通过后，再对比 benchmark 耗时 / OJ 耗时继续优化。 -*** -##### 5.2：选择目标算子 +**选择目标算子** FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benchmark 脚本、OJ 题包、`run_kernel(...)` 接口和数据范围可能不同。 @@ -332,26 +354,28 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc | **3** | MLA Paged Attention | DeepSeek MLA特有，双路Q(nope+pe)/双路Cache(ckv+kpe) |`bench_batch_mla.py` | `BatchMLAPagedAttentionWrapper` | | **4** | Paged Decode | 每次只1个query token，memory-bound |`bench_batch_decode.py` | `BatchDecodeWithPagedKVCacheWrapper` | -**每个子题的接口参数、数据范围和精度要求以对应 XPU-OJ 题包为准。** 下文以题目 **1 Ragged Prefill** 为例演示从 baseline benchmark 到 XPU-OJ 提交的完整流程。 +**每个子题的接口参数、数据范围和精度要求以对应 XPU-OJ 题包为准。** 下文以题目 **1 Ragged Prefill** 为例演示从 benchmark 到 XPU-OJ 提交的完整流程。 #### Step 6：理解 XPU-OJ 评测接口与精度要求 -**目标：** 明确 Baseline 与最终评测提交之间的关系，理解选手需要实现的内容。 +**目标：** 明确 Benchmark 与最终评测提交之间的关系，理解选手需要实现的内容。 -> 完成 baseline benchmark 后，需要注意 baseline 脚本主要用于建立性能基线，并不需要最终提交。 -> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 baseline 参考结果进行比较。 +> 完成 benchmark 后，需要注意 benchmark 脚本主要用于建立性能基线，并不需要最终提交。 +> 最终评测以 XPU-OJ 题包为准，评测程序会调用选手提交代码中的 `run_kernel`，并将输出结果与 OJ 参考实现结果进行比较。 下面以 FlashInfer Ragged Prefill 题为例，其中： - `zh_CN/00_题目描述.md`：说明需要实现的算子功能； -- `zh_CN/01_接口约定.md`：说明必须实现的 `run_kernel` 函数签名； +- `zh_CN/01_接口约定cuda.md`：说明必须实现cuda的 `run_kernel` 函数签名； +- `zh_CN/01_接口约定triton.md`：说明必须实现triton的 `run_kernel` 函数签名； +- `zh_CN/01_接口约定tilelang.md`：说明必须实现tilelang的 `run_kernel` 函数签名； - `zh_CN/02_数据范围.md`：说明测试范围和精度要求； -- `testcase_config.py`：定义测试数据生成、baseline 参考实现和正确性校验方式。 +- `testcase_config.py`：定义测试数据生成、OJ 参考实现和正确性校验方式。 FlashInfer Ragged Prefill 的校验方式为： ```python torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) ``` -选手实现的输出需要在上述容差范围内与 baseline 输出一致。不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。 +选手实现的输出需要在上述容差范围内与 OJ 参考实现输出一致。不同算子的容差可能不同，正式精度要求以对应 OJ 题包说明为准。 #### Step 7：登录 XPU-OJ 并进入题目页面 @@ -370,11 +394,12 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) #### Step 8：提交 OJ 冒烟代码 **目标**：完成一次最小提交，确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。 + 1. 在语言下拉框中选择本题支持的提交语言，例如 `MXMACA C++`、`TileLang` 或后续开放的 `Triton`； 2. 将实现了题目要求接口的代码复制到提交框中； > 如果你还没有 `run_kernel`，应该从哪里开始？ > - > - OJ 最终评测不会直接运行 baseline 脚本，而是调用你提交代码中的 `run_kernel(...)`。 + > - OJ 最终评测不会直接运行 benchmark 脚本，而是调用你提交代码中的 `run_kernel(...)`。 > - 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。 > - 如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路。 3. 借助 Agent 从题包生成 `run_kernel` 初版 @@ -403,32 +428,30 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) ![Weixin-Image_20260617170807](https://origin.picgo.net/2026/06/17/Weixin-Image_20260617170807_13_100b5fbaeba5856b3b2.png) - **OJ 冒烟代码** - - 用于最小链路验证。 + **OJ 冒烟代码：**用于最小链路验证。 ```cpp - #include - +#include + #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__ q, const __nv_bfloat16 *__restrict__ k, const __nv_bfloat16 *__restrict__ v, __nv_bfloat16 *__restrict__ output, @@ -445,24 +468,24 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) const int lane = threadIdx.x & 31; const int warp_id = threadIdx.x >> 5; const int warps_per_block = blockDim.x >> 5; - + int64_t work = static_cast(blockIdx.x) * warps_per_block + warp_id; - const int64_t total = batch_size * seq_len * num_qo_heads; + const int64_t total = batch_size * seq_len * num_qo_heads; if (work >= total) return; - + const int64_t qo_head = work % num_qo_heads; - work /= num_qo_heads; + work /= num_qo_heads; const int64_t q_pos = work % seq_len; const int64_t batch = work / seq_len; - + const int64_t qo_begin = qo_indptr[batch]; - const int64_t qo_len = qo_indptr[batch + 1] - qo_begin; + 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; + const int64_t kv_len = kv_indptr[batch + 1] - kv_begin; int64_t visible = kv_len; if (causal) { @@ -472,14 +495,14 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 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 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 qv[4]; float acc[4]; for (int i = 0; i < 4; ++i) { @@ -487,17 +510,17 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) qv[i] = (d < head_dim_qk) ? __bfloat162float(q_ptr[d]) : 0.0f; acc[i] = 0.0f; } - + float m = -1.0e20f; - float l = 0.0f; + 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) + for (int i = 0; i < 4; ++i) { const int d = lane + i * 32; if (d < head_dim_qk) @@ -506,13 +529,13 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) } } 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 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) { @@ -522,9 +545,9 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 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; + 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; @@ -534,11 +557,11 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) } } } - + } // namespace extern "C" void run_kernel( - const __nv_bfloat16 *q, + const __nv_bfloat16 *q, const __nv_bfloat16 *k, const __nv_bfloat16 *v, __nv_bfloat16 *output, @@ -562,16 +585,15 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) } ``` 以上代码仅用于说明接口结构，不代表最优实现，也不作为评分参考。 - -4. 点击提交，等待评测结果返回； + +4. 点击提交，等待评测结果返回 ![image 20260616155023443](https://origin.picgo.net/2026/06/16/image-202606161550234433e54d703d581f858.png) 评测时间与题目测试点数量、队列状态和平台负载有关，通常需要等待数十秒到数分钟。以平台实际返回为准。 - ![image iwEcAqNwbmcDAQTRBt8F0QOu](https://origin.picgo.net/2026/06/17/iwEcAqNwbmcDAQTRBt8F0QOuBrA3mrGf2xJ4fgoGfxH9GVsAB9MAAAABijZ-PQgACaJpbQoAC9IAAlPf771f660d13b07471.png) + ![image iwEcAqNwbmcDAQTRBt8F0QOu](https://origin.picgo.net/2026/06/17/iwEcAqNwbmcDAQTRBt8F0QOuBrA3mrGf2xJ4fgoGfxH9GVsAB9MAAAABijZ-PQgACaJpbQoAC9IAAlPf771f660d13b07471.png)**OJ 评测流程**： - **OJ 评测流程** 1. 选手提交代码； 2. 平台按所选语言编译或加载提交代码； 3. 评测程序构造测试输入； @@ -583,13 +605,103 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 5. 查看结果 - **50 分 / 10 分 / 与 baseline 加速比对比 / 榜单**：以 OJ 平台实际结果为准 + OJ 平台的每个测试用例都会会返回两行结果信息（以 Problem 20001 Ragged Prefill 冒烟测试第 1 个测试用例为例）： -## 七、Agent 使用样例 + ``` + OJCHAL v1 7P82PDM8iqVVD4hu6nNkKg== + OJRESULT v1 cd1dc1cc5abf7d59552d9a82db40c83041527d232fcb7760f8d2048e76045220 eyJ0aW1lX21zIjo4LjEzNiwic3BlZWR1cCI6MC4wMTYyMjQsInRrX3RpbWVfbXMiOjguMTM2LCJ0Yl90aW1lX21zIjowLjEzMiwidGhfdGltZV9tcyI6MC4wMzU3OTEsInNjb3JlX3JhdGlvIjowLjAxMTczOCwicGFzcyI6dHJ1ZX0= + ``` + + - **OJCHAL**：挑战相关元信息（参赛者无需关注）； + - **OJRESULT**：核心结果。分为两部分： + - 前缀（`OJRESULT v1` 后的第一段 base64）：评测平台签名哈希，用于防篡改校验（参赛者无需关注）； + - 后缀（第二段 base64）：JSON 格式的性能数据。 + + **下面演示如何提取性能数据：** + + - 反编码 OJRESULT 获取性能数据 + + - 将上面的后缀 base64 复制下来，使用以下命令解码： + + ``` bash + python3 -c " + import base64, json + + # 替换为你实际 OJRESULT 第二段 base64 字符串 + data = 'eyJ0aW1lX21zIjo4LjEzNiwic3BlZWR1cCI6MC4wMTYyMjQsInRrX3RpbWVfbXMiOjguMTM2LCJ0Yl90aW1lX21zIjowLjEzMiwidGhfdGltZV9tcyI6MC4wMzU3OTEsInNjb3JlX3JhdGlvIjowLjAxMTczOCwicGFzcyI6dHJ1ZX0=' + + # 解码并格式化输出 + decoded = base64.b64decode(data.encode('utf-8')).decode('utf-8') + result = json.loads(decoded) + print(json.dumps(result, indent=2, ensure_ascii=False)) + " + ``` + + - **解码输出如下** + + ``` json + { + "time_ms": 8.136, + "speedup": 0.016224, + "tk_time_ms": 8.136, + "tb_time_ms": 0.132, + "th_time_ms": 0.035791, + "score_ratio": 0.011738, + "pass": true + } + ``` + + - **字段含义** + + | 字段 | 含义 | 说明 | + |------|------|------| + | `pass` | 正确性校验 | `true` 表示输出通过了 `torch.allclose(rtol=1e-2, atol=1e-2)` 校验，可以参与排名 | + | `tb_time_ms` | OJ 参考实现耗时 (ms) | OJ 内置参考实现在该测试用例上的 GPU 执行时间 | + | `tk_time_ms` | 提交耗时 (ms) | 你提交的 `run_kernel` 在该测试用例上的 GPU 执行时间 | + | `speedup` | 加速比 | `speedup = tb_time_ms / tk_time_ms`，**>1 表示优于 OJ 参考实现，<1 表示慢于 OJ 参考实现** | + | `time_ms` | 总耗时 (ms) | 通常等于 `tk_time_ms`，即 kernel 执行总时间 | + | `th_time_ms` | Host 辅助耗时 (ms) | host 端数据准备、kernel launch、结果拷贝等开销 | + | `score_ratio` | 归一化得分 | 综合正确性和加速比的加权评分（0~1），用于排行榜排名 | + + - **加速比解读与冒烟测试预期** + + 以冒烟测试为例，上述结果中： + + - `tb_time_ms = 0.132ms`（OJ 参考实现） + - `tk_time_ms = 8.136ms`（冒烟 kernel） + - `speedup = 0.016224`，即约慢 62 倍 + + 冒烟测试加速比 <1 是完全正常的。后续优化目标是逐步提升加速比，直至超越 OJ 参考实现。 + + 加速比计算公式： + + ``` + 加速比 = tb_time_ms / tk_time_ms + = OJ 参考实现耗时 / 你的 kernel 耗时 + ``` + + 例：若 `tb_time_ms = 0.132`，优化后 `tk_time_ms = 0.080`，则 `speedup = 0.132 / 0.080 = 1.65`，即比 OJ 参考实现快 1.65 倍。 + + - **查看所有测试用例结果** + + OJ 平台每次提交会评测所有测试用例，每个测试用例返回一组独立的 `OJRESULT`。建议将每组结果反编码后整理成表格，方便跟踪优化进展： + + | 测试用例 | batch_size | seq_len | tb_time_ms | tk_time_ms | speedup | pass | + |----------|------------|---------|------------|------------|---------|------| + | 1 | 1 | 1024 | — | — | — | — | + | 2 | 4 | 1024 | 0.132 | 8.136 | 0.016 | true | + | 3 | 16 | 1024 | — | — | — | — | + | ... | ... | ... | ... | ... | ... | ... | + + * **与榜单的关系** + * OJ 平台根据各测试用例的 `score_ratio` 加权计算总分； + * 榜单排名以加权总分为依据，`pass=false` 的提交不参与排名。 + +## 7. Agent 使用样例 **目标：** 在本模块中，Agent 可以帮助你完成环境检查、运行 Benchmark、分析结果、理解 OJ 接口、生成和调试 `run_kernel` 等任务。以下是参考 prompt。 -### 环境检查 +### 7.1 环境检查 ```plaintext 请帮我检查当前环境是否满足 FlashInfer 运行要求，包括： @@ -598,24 +710,26 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 3. flashinfer、pandas、numpy 依赖是否已安装 ``` -### 运行 Benchmark +### 7.2 运行 Benchmark ``` plaintext 请帮我运行 bench_batch_prefill_ragged.py 脚本，执行 Ragged Prefill 的基准测试。 ``` -### 分析结果 +### 7.3 分析结果 ``` plaintext -请帮我读取最新的 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置，并与沐曦 GPU 理论峰值带宽做对比。 +请帮我读取最新的性能基线 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置，并与沐曦 GPU 理论峰值带宽做对比。 ``` -### 理解 OJ 题包接口 +### 7.4 理解 OJ 题包接口 ``` plaintext 请帮我阅读 FlashInfer Ragged Prefill 题包（problem_20001）中的以下文件： - zh_CN/00_题目描述.md -- zh_CN/01_接口约定.md +- zh_CN/01_接口约定cuda.md +- zh_CN/01_接口约定triton.md +- zh_CN/01_接口约定tilelang.md - zh_CN/02_数据范围.md 然后帮我总结： 1. run_kernel 的函数签名和每个参数的含义 @@ -625,7 +739,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 5. causal=1 时需要注意的边界条件 ``` -### 生成 `run_kernel` 初版 +### 7.5 生成 `run_kernel` 初版 ``` plaintext 请帮我为 FlashInfer Ragged Prefill 题（OJ 题号 20001）生成一个最小可运行的 run_kernel 实现，要求： @@ -640,7 +754,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) - scale = 1.0f / sqrtf(head_dim_qk) ``` -### 调试 OJ 提交错误 +### 7.6 调试 OJ 提交错误 ``` plaintext 我的 run_kernel 提交到 OJ 后显示 Wrong Answer，请帮我对比以下信息： 1. 这是我的 submit 代码：[粘贴你的 run_kernel 实现] @@ -654,7 +768,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) - online softmax 的 m/l 更新逻辑是否正确 ``` -### 问题排查 +### 7.7 问题排查 ``` plaintext 运行 bench_batch_prefill_ragged.py 时报错 out of memory，请帮我分析原因并给出解决方案。 @@ -662,21 +776,21 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) 请帮我解释 bench_batch_prefill_ragged.py 中 BatchPrefillWithRaggedKVCacheWrapper 的 plan() 和 run() 方法的工作原理，特别是 qo_indptr 和 kv_indptr 的作用。 整理优化日志 请帮我整理本次优化的记录，包括： -1. 原始 baseline 性能数据（从 CSV 中提取关键配置的 time_ms 和 tflops） +1. 原始 benchmark 结果性能数据（从 CSV 中提取关键配置的 time_ms 和 tflops） 2. 优化后的性能数据（从 OJ 评测结果中提取） -3. 加速比 = baseline_time / optimized_time -4. 以表格形式输出：配置参数 | baseline 耗时 | 优化后耗时 | 加速比 +3. 加速比 = benchmark_time / optimized_time +4. 以表格形式输出：配置参数 | benchmark 耗时 | 优化后耗时 | 加速比 ``` -## 八、常见问题 +## 8. 常见问题 -### 流程问题 +### 8.1 流程问题 | 问题 | 原因 | 解决办法 | | --- | --- | --- | | 登录后看不到题目 | 未使用赛用账号登录 | 七月份组委会统一发放 XPU-OJ 账号，请确认你使用的是组委会统一发放的账号，而不是自行注册账号；如仍无法看到题目，请联系助教或赛事运营确认账号权限。 | -### 环境问题 +### 8.2 环境问题 | 问题 | 原因 | 解决办法 | | --- | --- | --- | @@ -684,7 +798,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) | `ModuleNotFoundError: No module named 'flashinfer'` | flashinfer 未安装 | 执行 `pip install flashinfer` | | `out of memory` | GPU 显存不足 | 减小 `batch_size` 或 `seq_len` 参数 | -### 运行问题 +### 8.3 运行问题 | 问题 | 原因 | 解决办法 | | --- | --- | --- | @@ -692,21 +806,21 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) | `KeyError: 'BatchPrefillWithPagedKVCacheKernel'` | profiler 未捕获目标 kernel | 检查 `target_kernels` 配置是否正确 | | CSV 文件为空 | 测试未正常完成 | 检查 GPU 显存是否充足，重新运行 | -### 代码问题 +### 8.4 代码问题 | 问题 | 原因 | 解决办法 | | --- | --- | --- | | `ImportError: cannot import name 'xxx' from 'bench_common'` | 函数名拼写错误 | 检查 `bench_common.py` 中的函数名 | | `RuntimeError: error: device-side assert triggered` | 输入参数超出范围 | 检查 `num_qo_heads`、`num_kv_heads`、`head_dim` 配置 | -### 性能问题 +### 8.5 性能问题 | 问题 | 原因 | 解决办法 | | --- | --- | --- | | TFLOPs 数值异常低 | 工作负载过小，kernel 启动开销占比大 | 增大 `batch_size` 或 `seq_len` | | 带宽数值异常低 | 数据未正确加载到 GPU | 检查 Tensor 是否在 CUDA 设备上 | -### 评测问题 +### 8.6 评测问题 提交 XPU-OJ 后可能遇到的异常评测结果及排查方向： @@ -714,12 +828,12 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) |------|------|---------| | **Compilation Error** 编译错误 | 1. `run_kernel` 签名与 OJ 接口约定不一致（参数类型、顺序、数量不匹配）；2. 缺少 `extern "C"` 声明导致 C++ name mangling；3. 缺少必要头文件（`cuda_bf16.h`、`cuda_runtime.h`、`math.h`）；3. 使用了 OJ 环境不支持的语法或 API | 1. 逐行对照对应题包的「01_接口约定.md」，确认参数类型（`int64_t` vs `int`、`const` 修饰）、顺序完全一致；2. 在 `run_kernel` 前加 `extern "C"`；3. 确认文件顶部 include 了 ``、``、``、``；4. 去掉 `printf`、`assert` 等调试代码后重新提交 | | **Time Limit Exceeded** 运行超时 | 1. `run_kernel` 内部调用了 `cudaDeviceSynchronize()` 导致额外等待；2. kernel 中存在死循环（for 循环边界条件错误）；3. `__syncthreads()` 放在条件分支内导致线程死锁；4. grid 配置过大，启动的 block 数量远超合理范围 | 1. 删除 `run_kernel` 函数体内的 `cudaDeviceSynchronize()` 调用——评测器会在外部自行同步；2. 检查 kernel 中所有 for 循环的终止条件，确保 `kv_start <= block_max_q` 等边界正确；3. 将所有 `__syncthreads()` 移到 if/else 分支之外；4. 检查 grid 计算：`(seq_len + Br - 1) / Br`，确认 `Br` 取值合理 | -| **Wrong Answer** 答案错误 | 1. 注意力计算公式错误（score、scale、softmax 实现有偏差）；2. GQA 头映射错误：`hkv = hq / (num_qo_heads / num_kv_heads)` 计算不对；3. Causal mask 未正确实现（`causal=1` 时 query 看到了不该看的未来 token）；4. Online softmax 的 m/l 更新逻辑有误；5. float4 向量化加载的偏移计算错误，导致 K/V 数据错位；6. 输出写入偏移错误，或对无效位置写了垃圾值 | 1. 本地用题包中的 PyTorch 参考实现对拍：运行 `testcase_config.py` 的 `baseline()` 与你 kernel 输出做 `torch.allclose(rtol=1e-2, atol=1e-2)` 比对；2. GQA 公式：`int hkv = hq * num_kv_heads / num_qo_heads`（整数除法）；3. Causal 逻辑：`kv_end = min(kv_start + Bc, q_idx + 1)`，注意 +1 的处理；4. 对照论文 FlashAttention 的 Algorithm 1 逐行验证 online softmax；5. float4 加载偏移公式：`(cur_kv_start + i) * num_kv_heads * D_QK + hkv * D_QK + d_idx * 8`，确认 `num_kv_heads` 而非 `num_qo_heads` | +| **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` | -## 九、下一步学习建议 +## 9. 下一步学习建议 -### 1. 保存你的 Baseline 结果 +### 9.1 保存你的性能基线结果 将本次运行生成的 CSV 文件妥善保存，后续优化时需要以此作为对比基准。 @@ -729,7 +843,7 @@ mkdir -p results mv *.csv results/ ``` -### 2. 深入理解 FlashInfer 核心概念 +### 9.2 深入理解 FlashInfer 核心概念 * 阅读 FlashInfer 官方文档，理解 Paged KV Cache、Ragged KV Cache 的设计理念 @@ -738,7 +852,7 @@ mv *.csv results/ * 理解 `plan()` 和 `run()` 两阶段设计的作用 -### 3. 进入算子优化模块 +### 9.3 进入算子优化模块 参考后续优化模块，学习以下优化技术： @@ -749,18 +863,18 @@ mv *.csv results/ * **Compute Optimization**：提升计算效率 -### 4. 参考资源 +### 9.4 参考资源 * FlashInfer 官方仓库：https://github.com/flashinfer-ai/flashinfer * FlashInfer 文档：https://flashinfer.ai -### 5. 记录优化流程 +### 9.5 记录优化流程 建议维护一份优化日志，记录每次优化的改动和性能变化： -| 优化项 | 改动内容 | Baseline | 优化后 | 提升比例 | +| 优化项 | 改动内容 | 性能基线 | 优化后 | 提升比例 | | --- | --- | --- | --- | --- | | 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% | @@ -768,6 +882,6 @@ mv *.csv results/ > 使用 Agent 整理优化日志，可形成可复现的 Agent/Skill 优化流程 -### 6. 使用多语言完成算子优化加速 +### 9.6 使用多语言完成算子优化加速 可以使用 Triton 或 TileLang 语言实现 `run_kernel` 接口完成算子优化，对比不同的语言对于性能加速的影响。 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/starter/示例冒烟代码.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/starter/示例冒烟代码.md new file mode 100644 index 0000000..59ee384 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashinfer_task_package/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 +); +``` +