From 496c8cc691ba61ccb1be8176319a29a4872f58d7 Mon Sep 17 00:00:00 2001 From: "Yuliang Feng (i26389)" Date: Tue, 23 Jun 2026 11:40:37 +0800 Subject: [PATCH] fiximage --- ...ashAttention关键算子迁移与优化.md | 1019 ++++++++--------- .../benchmark_kvcache_headdim128.csv | 0 .../benchmark_kvcache_headdim160.csv | 0 .../benchmark_kvcache_headdim192.csv | 0 .../benchmark_kvcache_headdim224.csv | 0 .../benchmark_kvcache_headdim256.csv | 0 .../benchmark_kvcache_headdim32.csv | 0 .../benchmark_kvcache_headdim512.csv | 0 .../benchmark_kvcache_headdim64.csv | 0 .../benchmark_kvcache_headdim96.csv | 0 .../benchmark_kvcache.py | 0 .../{ => starter}/guide.md | 0 .../starter/smoke_test_code.txt | 119 ++ 13 files changed, 623 insertions(+), 515 deletions(-) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim128.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim160.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim192.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim224.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim256.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim32.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim512.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim64.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/baseline结果实例/benchmark_kvcache_headdim96.csv (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{Flashattn_Baseline => benchmark}/benchmark_kvcache.py (100%) rename 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/{ => starter}/guide.md (100%) create mode 100644 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test_code.txt diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md index c52af89..850452f 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashAttention关键算子迁移与优化.md @@ -1,296 +1,286 @@ -# Flashattention 迁移 Baseline 实战：从性能基线到 XPU-OJ 评测 +# Flashattention 迁移 Benchmark 实战:从性能基线到 XPU-OJ 评测 -## 一、教程定位 +## 1. 教程定位 -本教程是参赛训练课程的 **FlashAttention Baseline 入门与评测提交衔接** 模块，主要帮助用户跑通 FlashAttention paged KV‑cache 推理核函数 `flash_attn_with_kvcache` 的基准测试流程，理解 baseline 的输入输出、性能指标和评测含义，并基于 XPU‑OJ 题包完成一个最小正确版 `run_kernel` 的实现与提交。 +本教程是参赛训练课程的 **FlashAttention Benchmark 入门与评测提交衔接** 模块,主要帮助用户跑通 FlashAttention paged KV-cache 推理核函数 `flash_attn_with_kvcache` 的基准测试流程,理解 benchmark 脚本 的输入输出、性能指标和评测含义,并基于 XPU-OJ 题包完成一个最小正确版 `run_kernel` 的实现与提交。 -需要特别说明：本教程中的 Baseline 主要用于帮助参赛者理解目标算子的调用方式、输入输出结构和性能基线，Baseline 不是最终提交物。最终评测以 XPU‑OJ 题包为准，参赛者需要根据题包中的接口约定，实现自己的 `run_kernel`，并在输出结果对齐 baseline 参考结果的前提下提升性能。 +需要特别说明:本教程中的 benchmark 脚本 主要用于帮助参赛者理解目标算子的调用方式、输入输出结构和性能基线,benchmark 脚本 不是最终提交物。最终评测以 XPU-OJ 题包为准,参赛者需要根据题包中的接口约定,实现自己的 `run_kernel`,并在输出结果对齐 OJ 参考结果 的前提下提升性能。 -完成本教程后，学员应能够： +完成本教程后,学员应能够: * 完成环境验证与依赖检查 - -* 理解并配置 KV‑Cache Benchmark 的核心参数 - -* 明确 Baseline 的含义，理解 KV‑Cache 为何成为推理性能瓶颈 - -* 运行 `flash_attn_with_kvcache` 的 Benchmark 测试并获取性能数据 - -* 输出一份 Baseline 性能结果记录表，为后续算子优化提供对比基准 - -* 理解 XPU‑OJ 评测的 `run_kernel` 接口规范与精度要求 - -* 基于 OJ 题包接口实现一个最小正确版 `run_kernel`，通过正确性校验 - -> Baseline 解读：为什么本教程基于 KV Cache 做性能基线？ +* 理解并配置 KV-Cache Benchmark 的核心参数 -#### 什么是 Baseline（性能基线） +* 明确 benchmark 脚本 的含义,理解 KV-Cache 为何成为推理性能瓶颈 -> Baseline是在引入任何优化代码前，系统处于初始可用状态时的参考数据。它是衡量后续所有优化收益的“锚点”。记录内容通常包括：执行时间、吞吐量、带宽、显存占用。 +* 运行 `flash_attn_with_kvcache` 的 Benchmark 测试并获取性能数据 -#### 为什么需要 Baseline +* 输出一份 性能基线结果 记录表,为后续算子优化提供对比基准 -> Baseline 是未做任何优化前的参考性能指标，它可以回答： +* 理解 XPU-OJ 评测的 `run_kernel` 接口规范与精度要求 -* 当前性能处于什么水平？ - -* 是否存在明显性能瓶颈？ - -* 后续优化是否有效？ - -* 性能提升有多大？ - +* 基于 OJ 题包接口实现一个最小正确版 `run_kernel`,通过正确性校验 -> 没有 Baseline，则无法量化优化效果。 -> **注意**：Baseline ≠ Benchmark。 +### 1.0 性能基线解读:为什么本教程基于 KV Cache 做基准测试 -* **Benchmark** 是测量性能的手段（脚本、工具）。 - -* **Baseline** 是测量得到的具体结果数值。 - +### 1.1 什么是性能基线(baseline) -> 本教程运行的 `benchmark_kvcache.py` 是一个 Benchmark 脚本，它输出的 CSV 文件就是Baseline。 +> 性能基线 是在引入任何优化代码前,系统处于初始可用状态时的参考数据。它是衡量后续所有优化收益的"锚点"。记录内容通常包括:执行时间、吞吐量、带宽、显存占用。 -#### 为什么FlashAttention里面基线是针对KV Cache做benchmark性能验证？ +### 1.2 为什么需要性能基线 -> 因为KV Cache是最核心的性能瓶颈，尤其是在大模型推理的解码阶段。 +> 性能基线 是未做任何优化前的参考性能指标,它可以回答: -##### 瓶颈从计算转移到显存访问 +* 当前性能处于什么水平? -> Transformer 推理分为两个截然不同的阶段： +* 是否存在明显性能瓶颈? -| > 阶段 | > 特点 | > 计算模式 | > 主要限制 | +* 后续优化是否有效? + +* 性能提升有多大? + + +> 没有 性能基线,则无法量化优化效果。 + +> **注意**:性能基线 ≠ benchmark 脚本。 + +* **benchmark 脚本** 是测量性能的手段(脚本、工具)。 + +* **性能基线** 是测量得到的具体结果数值。 + + +> 本教程运行的 `benchmark_kvcache.py` 是一个 benchmark 脚本,它输出的 CSV 文件就是 性能基线结果。 + +### 1.3 为什么 FlashAttention 的 benchmark 脚本针对 KV Cache 做性能验证 + +> 因为KV Cache是最核心的性能瓶颈,尤其是在大模型推理的解码阶段。 + +#### 1.3.1 瓶颈从计算转移到显存访问 + +> Transformer 推理分为两个截然不同的阶段: + +| 阶段 | 特点 | 计算模式 | 主要限制 | | --- | --- | --- | --- | -| > **Prefill（预填充）** | > 一次性处理全部输入 token | > 矩阵运算密集，Tensor Core 利用率高 | > **计算受限（Compute Bound）** | -| > **Decode（解码）** | > 逐个生成新 token，每步只算一个 token | > 每次都要**读取全部历史 KV Cache** | > **显存访问受限（Memory Bound）** | +| **Prefill(预填充)** | 一次性处理全部输入 token | 矩阵运算密集,Tensor Core 利用率高 | **计算受限(Compute Bound)** | +| **Decode(解码)** | 逐个生成新 token,每步只算一个 token | 每次都要**读取全部历史 KV Cache** | **显存访问受限(Memory Bound)** | -* Prefill 阶段计算量大但形态规整，通常能较好利用 GPU 算力，不是主要瓶颈。 - -* Decode 阶段占推理过程的大部分时间（尤其是长上下文交互），每个 token 的生成都需要搬运整个 KV Cache。随着序列增长，显存访问开销占比越来越高。 - +* Prefill 阶段计算量大但形态规整,通常能较好利用 GPU 算力,不是主要瓶颈。 -##### KV‑Cache 的访存密集型特征 +* Decode 阶段占推理过程的大部分时间(尤其是长上下文交互),每个 token 的生成都需要搬运整个 KV Cache。随着序列增长,显存访问开销占比越来越高。 -> 在 Decode 阶段，flash\_attn\_with\_kvcache 内核的工作是： -* 读取当前 token 的 Q 向量（很小） - -* **反复读取整个历史的 KV Cache**（很大，线性增长） - -* 执行 FlashAttention 计算，然后将新的 K/V 追加写入缓存。 - +#### 1.3.2 KV-Cache 的访存密集型特征 -> 95% 以上的时间花在读取 KV Cache 上。因此，Decode 阶段的性能完全由显存带宽决定，而不是 GPU 算力。 +> 在 Decode 阶段,flash\_attn\_with\_kvcache 内核的工作是: -##### KV Cache 显存容量直接限制并发能力 +* 读取当前 token 的 Q 向量(很小) -> 大模型推理服务需要同时处理多个请求（batch\_size）。每个请求都有自己的 KV‑Cache，显存总占用量与 `batch_size × 序列长度` 成正比。因此，KV Cache 的显存开销直接决定了系统可以同时服务多少用户。 +* **反复读取整个历史的 KV Cache**(很大,线性增长) -> 通过 Benchmark 对不同 `batch_size × seq_len_kv` 组合进行压力测试，可以： +* 执行 FlashAttention 计算,然后将新的 K/V 追加写入缓存。 + + +> 95% 以上的时间花在读取 KV Cache 上。因此,Decode 阶段的性能完全由显存带宽决定,而不是 GPU 算力。 + +#### 1.3.3 KV Cache 显存容量直接限制并发能力 + +> 大模型推理服务需要同时处理多个请求(batch\_size)。每个请求都有自己的 KV-Cache,显存总占用量与 `batch_size × 序列长度` 成正比。因此,KV Cache 的显存开销直接决定了系统可以同时服务多少用户。 + +> 通过 Benchmark 对不同 `batch_size × seq_len_kv` 组合进行压力测试,可以: + +* 找出 **OOM 边界**:哪些参数组合会导致显存溢出,无法运行 -* 找出 **OOM 边界**：哪些参数组合会导致显存溢出，无法运行 - * 量化每个请求的平均显存开销 - -* 为后续 **分页 KV‑Cache（PagedAttention）** 等优化提供基线对比。 - -##### 为什么选择 flash\_attn\_with\_kvcache作为测试对象 +* 为后续 **分页 KV-Cache(PagedAttention)** 等优化提供基线对比。 -> 它是 FlashAttention 专门为推理阶段设计的核心算子。它融合了高效的 KV Cache 读取与更新逻辑、FlashAttention 的分块、重计算技术以节省显存 -> 推理场景下精确的带宽优化。选择它作为 Benchmark 对象，可以直接回答： +#### 1.3.4 为什么选择 flash_attn_with_kvcache 作为测试对象 -* GPU带宽利用率是多少？是否接近理论峰值？ - -* 哪些参数组合（batch\_size, seq\_len\_kv, headdim）达到峰值性能？ - -* OOM边界在哪里？ - -* 后续优化（如分页缓存、算子融合）是否有效？ - +> 它是 FlashAttention 专门为推理阶段设计的核心算子。它融合了高效的 KV Cache 读取与更新逻辑、FlashAttention 的分块、重计算技术以节省显存 -#### 为什么要记录 Baseline 结果 +> 推理场景下精确的带宽优化。选择它作为 Benchmark 对象,可以直接回答: -> Baseline（基线）是优化前的参考性能数据，记录 baseline 的意义在于： +* GPU带宽利用率是多少?是否接近理论峰值? -* **量化优化收益**：优化后对比 baseline，计算加速比（speedup = baseline\_time / optimized\_time） - -* **防止性能回退**：代码变更后重跑 benchmark，确认没有引入性能退化 - -* **建立测试矩阵**：记录不同参数组合下的 baseline，全面了解性能特征 - +* 哪些参数组合(batch\_size, seq\_len\_kv, headdim)达到峰值性能? -> 在本教程中，baseline 结果以 CSV 文件保存，包含每种 `batch_size × seq_len_kv` 配置的执行时间和带宽，为后续算子优化提供对比基准。 +* OOM边界在哪里? + +* 后续优化(如分页缓存、算子融合)是否有效? + + +### 1.4 为什么要记录性能基线结果 + +> 性能基线(baseline)是优化前的参考性能数据,记录 性能基线 的意义在于: + +* **量化优化收益**:优化后对比 性能基线,计算加速比(speedup = baseline\_time / optimized\_time) + +* **防止性能回退**:代码变更后重跑 benchmark,确认没有引入性能退化 + +* **建立测试矩阵**:记录不同参数组合下的 性能基线,全面了解性能特征 + + +> 在本教程中,性能基线结果 以 CSV 文件保存,包含每种 `batch_size × seq_len_kv` 配置的执行时间和带宽,为后续算子优化提供对比基准。 --- -## 二、学习目标 +## 2. 学习目标 -完成本模块后，你将能够： +完成本模块后,你将能够: + +1. **理解 FlashAttention Paged KV-Cache 算子的作用** + 明白 `flash_attn_with_kvcache` 在 LLM 推理 decode 阶段如何高效利用分页 KV Cache,减少显存碎片并提升吞吐。 + +2. **完成环境准备与 Benchmark 运行** + 安装所需依赖,运行 `benchmark_kvcache.py`,生成包含执行时间与显存带宽的 CSV 性能记录。 + +3. **读懂 benchmark 脚本 并定位性能瓶颈** + 分析不同 `batch_size`、`seq_len_kv` 下的带宽曲线,理解显存带宽对 decode 阶段的影响。 + +4. **理清 benchmark 脚本 与 OJ 题包的关系** + 明确 benchmark 脚本 用于建立性能基线,OJ 题包定义最终提交接口、数据范围和精度校验标准。 + +5. **实现并提交一个最小正确版 `run_kernel`** + 根据题包中的接口约定编写 CUDA 算子,通过 OJ 正确性校验并记录首次提交耗时。 -1. **理解 FlashAttention Paged KV‑Cache 算子的作用** - 明白 `flash_attn_with_kvcache` 在 LLM 推理 decode 阶段如何高效利用分页 KV Cache，减少显存碎片并提升吞吐。 - -2. **完成环境准备与 Benchmark 运行** - 安装所需依赖，运行 `benchmark_kvcache.py`，生成包含执行时间与显存带宽的 CSV 性能记录。 - -3. **读懂 Baseline 并定位性能瓶颈** - 分析不同 `batch_size`、`seq_len_kv` 下的带宽曲线，理解显存带宽对 decode 阶段的影响。 - -4. **理清 Baseline 与 OJ 题包的关系** - 明确基准脚本用于性能参照，OJ 题包定义最终提交接口、数据范围和精度校验标准。 - -5. **实现并提交一个最小正确版** `**run_kernel**` - 根据题包中的接口约定编写 CUDA 算子，通过 OJ 正确性校验并记录首次提交耗时。 - --- -## 三、适用对象 +## 3. 适用对象 -本模块适合以下人员： +本模块适合以下人员: -* 参与 AI 基础设施竞赛的参赛者 - -* 对 GPU 算子性能优化感兴趣的开发者 - -* 需要了解 FlashAttention KV-Cache 推理性能的研究人员 - +* 参与 AI 基础设施竞赛的参赛者 +* 对 GPU 算子性能优化感兴趣的开发者 +* 需要了解 FlashAttention KV-Cache 推理性能的研究人员 -**基础知识要求：** +**基础知识要求:** -* 了解 Python 编程基础 - -* 了解 PyTorch 基本用法 - -* 了解 GPU 推理的基本概念 - +* 了解 Python 编程基础 +* 了解 PyTorch 基本用法 +* 了解 GPU 推理的基本概念 --- -## 四、前置准备 +## 4. 前置准备 -开始实战前，请确认你已经完成以下准备： +开始实战前,请确认你已经完成以下准备: -### 环境准备 +### 4.1 环境准备 * **设置领取与兑换算力券** - -1. 前往沐曦开发者社区注册账号并完成邮箱验证，申请并获取 MACA 算力代金券兑换码。 - - 链接：https://developer.metax-tech.com/activities/6 - -2. 登录 模力方舟平台 (Gitee AI)，在“费用中心 -> 算力券”页面输入兑换码完成充值。 链接：https://ai.gitee.com/ - + +1. 前往沐曦开发者社区注册账号并完成邮箱验证,申请并获取 MACA 算力代金券兑换码。 + + 链接:https://developer.metax-tech.com/activities/6 + +2. 登录 模力方舟平台 (Gitee AI),在"费用中心 -> 算力券"页面输入兑换码完成充值。 链接:https://ai.gitee.com/ + * **创建并启动实例** - -1. 进入 算力市场，筛选“沐曦”芯片厂商，选择合适的 GPU 规格（推荐 曦云 C500 节点）。 - -2. 关键配置： 在预装镜像处，务必选择专属开发镜像（PyTorch Agent/2.8.0/Python 3.12/maca 3.7.2.1）。 - -3. 创建完成后，进入算力容器，点击“工具-lab”即可打开 JupyterLab 终端开始项目创作。 - -\*\*说明：\*\*由于本次使用的是预装的专属镜像，环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后，无需再进行繁琐的依赖库版本验证即可直接进入测试环节。 +1. 进入 算力市场,筛选"沐曦"芯片厂商,选择合适的 GPU 规格(推荐 曦云 C500 节点)。 -### 代码准备 +2. 关键配置: 在预装镜像处,务必选择专属开发镜像(PyTorch Agent/2.8.0/Python 3.12/maca 3.7.2.1)。 + +3. 创建完成后,进入算力容器,点击"工具-lab"即可打开 JupyterLab 终端开始项目创作。 + + +\*\*说明:\*\*由于本次使用的是预装的专属镜像,环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后,无需再进行繁琐的依赖库版本验证即可直接进入测试环节。 + +### 4.2 代码准备 + +* 获取目标源码(包含 `benchmark_kvcache.py` 及 OJ 题包) + +* 已进入项目目录 `/data/flashattn_baseline` + +* 准备 Benchmark 脚本与 OJ 测试脚本 -* 获取目标源码（包含 `benchmark_kvcache.py` 及 OJ 题包） - -* 已进入项目目录 `/data/flashattn_baseline` - -* 准备 Benchmark 脚本与 OJ 测试脚本 - --- -## 五、知识速览 +## 5. 知识速览 -### 关键术语 +### 5.1 关键术语 -* **KV-Cache：**缓存历史 Token 的 Key/Value 向量，避免 Transformer 推理时重复计算。 - -* **Paged KV-Cache：**将 KV-Cache 分页管理，减少显存碎片，提高利用率。 - -* **Batch Size：**一次处理的样本数，越大并行度越高，但显存占用越大。 - -* **seq\_len\_kv：**KV-Cache 中已缓存的历史 Token 数量。 - -* **headdim：**注意力头维度，常见为 64/128/256。 - +* **KV-Cache:**缓存历史 Token 的 Key/Value 向量,避免 Transformer 推理时重复计算。 -### 核心知识 +* **Paged KV-Cache:**将 KV-Cache 分页管理,减少显存碎片,提高利用率。 -* **正确性测试 ：**验证算子输出结果的数学精度是否与标准实现一致，这是绝对底线。 - -* **性能测试 ：**在正确的前提下，测算速度与吞吐。通过建立Baseline（基线），才能量化后续每次代码修改带来的真实收益（加速比）。 - -* **Benchmark （基准测试）：**在固定条件下反复运行同一任务，获取可重复的性能指标，用于建立基线、量化优化效果和定位瓶颈。 - -* **XPU‑OJ**：比赛官方在线评测平台，最终评测会调用参赛者提交代码中的 `run_kernel`。 - +* **Batch Size:**一次处理的样本数,越大并行度越高,但显存占用越大。 -### 关键指标 +* **seq\_len\_kv:**KV-Cache 中已缓存的历史 Token 数量。 -* **Kernel 执行时间：**GPU 核函数运行耗时（ms），使用 GPU 端同步计时获得。 - -* **有效带宽：**数据传输量 (GB) ÷ Kernel 时间 (s)，越接近理论峰值说明显存带宽利用越充分。 - +* **headdim:**注意力头维度,常见为 64/128/256。 -### 其他要点 -* **Warmup：**预热若干次（不记录），使 GPU 进入稳定状态。 - -* **Repeat：**正式运行多次，取平均值或中位数以消除波动。 - -* **同步：**调用 torch.cuda.synchronize() 确保精确计时。 - -* **数据类型：**本教程使用 bfloat16，在精度和性能取得平衡。 - -* **显存占用估算：**KV-Cache ≈ batch × seq\_len\_kv × num\_heads\_k × headdim × 2(K+V) × 字节数。 - -* **OOM 应对：**减小 batch/seq\_len\_kv、使用更小 dtype 或释放中间变量。 - -* **Tensor Core：**现代 GPU（含沐曦 C500）的矩阵乘法专用单元，要求维度对齐为 8 或 16 的倍数。 - +### 5.2 核心知识 + +* **正确性测试 :**验证算子输出结果的数学精度是否与标准实现一致,这是绝对底线。 + +* **性能测试 :**在正确的前提下,测算速度与吞吐。通过建立性能基线(baseline),才能量化后续每次代码修改带来的真实收益(加速比)。 + +* **Benchmark (基准测试):**在固定条件下反复运行同一任务,获取可重复的性能指标,用于建立基线、量化优化效果和定位瓶颈。 + +* **XPU-OJ**:比赛官方在线评测平台,最终评测会调用参赛者提交代码中的 `run_kernel`。 + + +### 5.3 关键指标 + +* **Kernel 执行时间:**GPU 核函数运行耗时(ms),使用 GPU 端同步计时获得。 + +* **有效带宽:**数据传输量 (GB) ÷ Kernel 时间 (s),越接近理论峰值说明显存带宽利用越充分。 + + +### 5.4 其他要点 + +* **Warmup:**预热若干次(不记录),使 GPU 进入稳定状态。 + +* **Repeat:**正式运行多次,取平均值或中位数以消除波动。 + +* **同步:**调用 torch.cuda.synchronize() 确保精确计时。 + +* **数据类型:**本教程使用 bfloat16,在精度和性能取得平衡。 + +* **显存占用估算:**KV-Cache ≈ batch × seq\_len\_kv × num\_heads\_k × headdim × 2(K+V) × 字节数。 + +* **OOM 应对:**减小 batch/seq\_len\_kv、使用更小 dtype 或释放中间变量。 + +* **Tensor Core:**现代 GPU(含沐曦 C500)的矩阵乘法专用单元,要求维度对齐为 8 或 16 的倍数。 + --- -## 六、项目实践：FlashAttention KV‑Cache Benchmark & OJ 评测 +## 6. 项目实践:FlashAttention KV-Cache Benchmark & OJ 评测 -### 项目目标 +* 对 FlashAttention 的 paged KV-cache 推理核函数(`flash_attn_with_kvcache`)进行自动化性能基准测试,覆盖多种 `batch_size × seq_len_kv` 组合,输出执行时间和有效显存带宽。 -* 对 FlashAttention 的 paged KV-cache 推理核函数（`flash_attn_with_kvcache`）进行自动化性能基准测试，覆盖多种 `batch_size × seq_len_kv` 组合，输出执行时间和有效显存带宽。 - +* 理解 XPU-OJ 题包的 `run_kernel` 接口,实现一个最小正确版 CUDA 算子,通过所有 OJ 正确性测试用例。 -* 理解 XPU‑OJ 题包的 `run_kernel` 接口，实现一个最小正确版 CUDA 算子，通过所有 OJ 正确性测试用例。 - +### 6.1 步骤 0:进入创建的实例环境 -### 步骤 0：进入创建的实例环境 - -模力方舟链接：https://ai.gitee.com/fwlhecko/dashboard/compute/instances +模力方舟链接:https://ai.gitee.com/fwlhecko/dashboard/compute/instances 选择工具-lab进入实例环境 -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/b5d7d783-3106-4a4d-97f5-c7ef4d7fa537.png) +![b5d7d783 3106 4a4d 97f5 c7ef4d7fa537](https://origin.picgo.net/2026/06/18/b5d7d783-3106-4a4d-97f5-c7ef4d7fa5379e806361950ffeed.png) -### 步骤 1：检查运行环境 +### 6.2 步骤 1:检查运行环境 -**目标：** 确认当前环境满足本模块运行要求。 +**目标:** 确认当前环境满足本模块运行要求。 -在JupyterLab Terminal中检查运行环境的配置。 +在JupyterLab Terminal中检查运行环境的配置。 -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/339ac3f8-e31d-43c6-992b-07e20e35ef95.png) +![339ac3f8 e31d 43c6 992b 07e20e35ef95](https://origin.picgo.net/2026/06/18/339ac3f8-e31d-43c6-992b-07e20e35ef95004ae954fadfe267.png) -**操作：** 检查 GPU 状态、Python 版本和依赖版本。 +**操作:** 检查 GPU 状态、Python 版本和依赖版本。 -**命令示例：** +**命令示例:** ```Bash # 检查沐曦 GPU 状态 @@ -309,114 +299,114 @@ python -c "import einops; print('einops OK')" ``` -**预期结果：** +**预期结果:** -* `mx-smi` 显示沐曦 GPU 信息 - +* `mx-smi` 显示沐曦 GPU 信息 -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/016119b5-960b-478d-9119-af27b9f3c727.png) -* Python 版本 = 3.8 - +![016119b5 960b 478d 9119 af27b9f3c727](https://origin.picgo.net/2026/06/18/016119b5-960b-478d-9119-af27b9f3c727ac0af039c38a9f5e.png) -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/568886e4-85a7-41c6-a94c-89a09913e0f7.png) +* Python 版本 = 3.8 -* `torch.cuda.is_available()` 返回 `True` - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/d1c97030-d40f-416f-89b3-b81caff88966.png) +![568886e4 85a7 41c6 a94c 89a09913e0f7](https://origin.picgo.net/2026/06/18/568886e4-85a7-41c6-a94c-89a09913e0f712d9e076005617b3.png) + +* `torch.cuda.is_available()` 返回 `True` + + +![d1c97030 d40f 416f 89b3 b81caff88966](https://origin.picgo.net/2026/06/18/d1c97030-d40f-416f-89b3-b81caff88966e271eefaf7ce35bc.png) * 所有依赖版本符合要求 - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/352b41ee-8a87-405e-991b-1c3714bdeff0.png) -**常见问题：** +![352b41ee 8a87 405e 991b 1c3714bdeff0](https://origin.picgo.net/2026/06/18/352b41ee-8a87-405e-991b-1c3714bdeff0e5fcd7b9b87d7b59.png) + +**常见问题:** | 问题 | 解决方法 | | --- | --- | -| `torch.cuda.is_available()` 返回 `False` | 检查 MXMACA 环境变量是否正确配置 | -| `ModuleNotFoundError: No module named 'flash_attn '` | 确认 flash-attn 已安装且版本= 2.6 | -| `mx-smi` 命令不存在 | 确认已配置沐曦 GPU 驱动环境 | +| `torch.cuda.is_available()` 返回 `False` | 检查 MXMACA 环境变量是否正确配置 | +| `ModuleNotFoundError: No module named 'flash\_attn '` | 确认 flash-attn 已安装且版本= 2.6 | +| `mx-smi` 命令不存在 | 确认已配置沐曦 GPU 驱动环境 | --- -### 步骤 2：进入项目目录 +### 6.3 步骤 2:进入项目目录 -目标：进入本模块所需的源码目录。 +目标:进入本模块所需的源码目录。 1. 克隆代码仓库 - + ```Bash git clone https://gitlink.org.cn/metax-maca/op_optimization.git ``` - + 2. 准备flashattn\_baseline - - 在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下，找到 `flashattn_baseline` 文件夹。可以将 `flashattn_baseline` 整个目录复制到工作目录 `data/` 下。 - - **下一步操作：** 切换到 Flashattn\_Baseline 项目目录。 - - **命令示例：** - + + 在仓库目录 `op_optimization/基于AI Agent开发范式的国产GPU大模型推理算子库优化` 下,找到 `flashattn_baseline` 文件夹。可以将 `flashattn_baseline` 整个目录复制到工作目录 `data/` 下。 + + **下一步操作:** 切换到 Flashattn\_Baseline 项目目录。 + + **命令示例:** + ```bash cd flashattn_baseline/Flashattn_Baselinels -la ``` - - **预期结果：** - + + **预期结果:** + ```Plain total xx drwxr-xr-x 2 root root 4096 Jun 1 09:00 __MACOSX - - - rw-r--r-- 1 root root 5232 Jun 1 09:00 benchmark_kvcache.py - - - - rw-r--r-- 1 root root 1440 Jun 1 09:00 benchmark_kvcache_20260526_150953.csv - - ... - - ``` - -3. 将Flashattn\_Baseline文件加入JupyterLab。 - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/3eb4d4fe-e948-405e-8d11-6301b4d8f16e.png) + - rw-r--r-- 1 root root 5232 Jun 1 09:00 benchmark_kvcache.py + + + - rw-r--r-- 1 root root 1440 Jun 1 09:00 benchmark_kvcache_20260526_150953.csv + + ... + + ``` + +3. 将Flashattn\_Baseline文件加入JupyterLab。 + + +![3eb4d4fe e948 405e 8d11 6301b4d8f16e](https://origin.picgo.net/2026/06/18/3eb4d4fe-e948-405e-8d11-6301b4d8f16e342580914be51bbf.png) --- -### 步骤 3：配置基准测试参数 +### 6.4 步骤 3:配置基准测试参数 -**目标：** 根据测试需求配置基准测试参数。 +**目标:** 根据测试需求配置基准测试参数。 -**操作：** 在 `benchmark_kvcache.py` 的 `main()` 函数中配置基准测试参数。 +**操作:** 在 `benchmark_kvcache.py` 的 `main()` 函数中配置基准测试参数。 -**参数说明：** +**参数说明:** | 参数 | 默认值 | 说明 | | --- | --- | --- | -| `headdims` | `[256]` | head dimension，可改为 `[128]` 等 | -| `page_block_size` | `16` | paged KV-cache 的 block 大小 | -| `batch_sizes` | `[1, 2, 4, 8, 16, 32, 64, 128]` | 批大小扫描范围 | -| `seq_lens_kv` | `[512, 1024, 2048, 4096, 8192, 16384]` | KV 序列长度扫描范围 | -| `num_heads` | `8` | query head 数量 | -| `num_heads_k` | `8` | KV head 数量 | -| `seqlen_q` | `1` | query 序列长度（单 token 推理） | +| `headdims` | `[256]` | head dimension,可改为 `[128]` 等 | +| `page_block_size` | `16` | paged KV-cache 的 block 大小 | +| `batch_sizes` | `[1, 2, 4, 8, 16, 32, 64, 128]` | 批大小扫描范围 | +| `seq_lens_kv` | `[512, 1024, 2048, 4096, 8192, 16384]` | KV 序列长度扫描范围 | +| `num_heads` | `8` | query head 数量 | +| `num_heads_k` | `8` | KV head 数量 | +| `seqlen_q` | `1` | query 序列长度(单 token 推理) | | `dtype` | `torch.bfloat16` | 数据类型 | -| `causal` | `False` | 是否启用 causal mask | +| `causal` | `False` | 是否启用 causal mask | | `warmup` | `10` | 预热迭代次数 | -| `repeat` | `100` | 正式 profiling 迭代次数 | +| `repeat` | `100` | 正式 profiling 迭代次数 | -**配置示例：** +**配置示例:** -如需测试 `headdim=128`，修改对应列表： +如需测试 `headdim=128`,修改对应列表: ```Python headdims = [128] ``` -如需启用 causal mask： +如需启用 causal mask: ```Python causal = True @@ -424,13 +414,13 @@ causal = True ``` --- -### 步骤 4：运行基准测试 +### 6.5 步骤 4:运行基准测试 -**目标：** 运行基准测试脚本，收集性能数据。 +**目标:** 运行基准测试脚本,收集性能数据。 -**操作：** 执行基准测试脚本并观察输出。 +**操作:** 执行基准测试脚本并观察输出。 -**命令示例：** +**命令示例:** ```Bash cd flashattn_baseline @@ -438,9 +428,9 @@ python benchmark_kvcache.py ``` -**预期结果：** +**预期结果:** -脚本运行时会在终端实时打印结果表格： +脚本运行时会在终端实时打印结果表格: ```Plain batch_size seq_len_kv heads headdim time_ms bandwidth_GB_s @@ -451,35 +441,35 @@ batch_size seq_len_kv heads headdim time_ms bandwidth_GB_s ``` -同时生成带时间戳的 CSV 文件，命名格式为 `benchmark_kvcache_YYYYMMDD_HHMMSS.csv`。 +同时生成带时间戳的 CSV 文件,命名格式为 `benchmark_kvcache_YYYYMMDD_HHMMSS.csv`。 -**常见问题：** +**常见问题:** | 问题 | 解决方法 | | --- | --- | -| 出现 `OOM` 标记 | 该配置超出 GPU 显存容量，可减小 batch\_size 或 seq\_len\_kv | -| 脚本运行缓慢 | 减少 `repeat` 次数或缩小扫描范围 | +| 出现 `OOM` 标记 | 该配置超出 GPU 显存容量,可减小 batch\_size 或 seq\_len\_kv | +| 脚本运行缓慢 | 减少 `repeat` 次数或缩小扫描范围 | --- -### 步骤 5：查看与分析结果 +### 6.6 步骤 5:查看与分析结果 -**目标：** 理解输出结果格式，分析性能数据。 +**目标:** 理解输出结果格式,分析性能数据。 -**CSV 输出格式：** +**CSV 输出格式:** | 列名 | 说明 | | --- | --- | | `batch_size` | 批大小 | -| `seq_len_kv` | KV 序列长度 | -| `heads` | head 数量 | -| `headdim` | head 维度 | -| `time_ms` | kernel 执行时间（毫秒） | -| `bandwidth_GB_s` | 有效显存带宽（GB/s） | +| `seq_len_kv` | KV 序列长度 | +| `heads` | head 数量 | +| `headdim` | head 维度 | +| `time_ms` | kernel 执行时间(毫秒) | +| `bandwidth_GB_s` | 有效显存带宽(GB/s) | -如果某个配置因显存不足而失败，对应的 `time_ms` 和 `bandwidth_GB_s` 列会标记为 `OOM`。 +如果某个配置因显存不足而失败,对应的 `time_ms` 和 `bandwidth_GB_s` 列会标记为 `OOM`。 -**带宽计算公式：** +**带宽计算公式:** ```Plain total_bytes = q_bytes + kv_bytes @@ -489,22 +479,22 @@ bandwidth = (total_bytes / 1e9) / (time_ms / 1e3) [GB/s] ``` -其中 `bytes_per_elem` 在 `bfloat16` 下为 2，`float32` 下为 4。 +其中 `bytes_per_elem` 在 `bfloat16` 下为 2,`float32` 下为 4。 -**性能观察：** +**性能观察:** -* **小 batch 时带宽较低**：batch\_size=1 时 kernel 无法充分利用 GPU 并行度，带宽通常 < 100 GB/s - -* **大 batch + 长序列时带宽较高**：batch\_size=128 时可接近 GPU 显存带宽上限 - -* **headdim 增大时 OOM 风险增加**：headdim=256 的显存占用是 headdim=128 的两倍，大 batch + 长序列更容易 OOM - +* **小 batch 时带宽较低**:batch\_size=1 时 kernel 无法充分利用 GPU 并行度,带宽通常 < 100 GB/s -**测试结果示例：** +* **大 batch + 长序列时带宽较高**:batch\_size=128 时可接近 GPU 显存带宽上限 -#### headdim=128（2026-05-26） +* **headdim 增大时 OOM 风险增加**:headdim=256 的显存占用是 headdim=128 的两倍,大 batch + 长序列更容易 OOM -所有 48 个配置均成功运行，峰值带宽约 1251 GB/s。 + +**测试结果示例:** + +#### 6.6.1 headdim=128(2026-05-26) + +所有 48 个配置均成功运行,峰值带宽约 1251 GB/s。 | batch\_size | seq\_len\_kv | time\_ms | bandwidth\_GB\_s | | --- | --- | --- | --- | @@ -513,9 +503,9 @@ bandwidth = (total_bytes / 1e9) / (time_ms / 1e3) [GB/s] | 1 | 16384 | 0.8356 | 80.32 | | 128 | 16384 | 6.8668 | 1250.98 | -#### headdim=256（2026-05-27） +#### 6.6.2 headdim=256(2026-05-27) -48 个配置中有 3 个因显存不足（OOM）而失败，峰值带宽约 807 GB/s。 +48 个配置中有 3 个因显存不足(OOM)而失败,峰值带宽约 807 GB/s。 | batch\_size | seq\_len\_kv | time\_ms | bandwidth\_GB\_s | | --- | --- | --- | --- | @@ -527,25 +517,25 @@ bandwidth = (total_bytes / 1e9) / (time_ms / 1e3) [GB/s] --- -### 步骤 6：自定义扩展（可选） +### 6.7 步骤 6:自定义扩展(可选) -如需进一步测试，可参考以下扩展方法： +如需进一步测试,可参考以下扩展方法: -**修改 headdim：** +**修改 headdim:** ```Python headdims = [128, 256] ``` -**启用 causal mask：** +**启用 causal mask:** ```Python causal = True ``` -**调整 profiling 精度：** +**调整 profiling 精度:** ```Python warmup = 20 @@ -553,65 +543,64 @@ repeat = 200 ``` -**启用详细 profiler 输出：** +**启用详细 profiler 输出:** ```Python ms = run_with_profiler(run_fn, warmup=warmup, reps=repeat, print_result=True, target_kernels=["flash"]) ``` -### 步骤 7:从 Baseline 到 XPU-OJ 提交 +### 6.8 步骤 7:从 benchmark 脚本 到 XPU-OJ 提交 -Baseline benchmark 用于理解目标算子的调用方式、输入输出 shape 和性能基线;XPU-OJ 题包用于定义最终评测接口、数据范围、参考输出和精度要求。 +benchmark 脚本 用于理解目标算子的调用方式、输入输出 shape 和性能基线;XPU-OJ 题包用于定义最终评测接口、数据范围、参考输出和精度要求。 -跑完 baseline 后,选手需要完成以下转换: +跑完 benchmark 后,选手需要完成以下转换: -1. 从 benchmark 脚本中理解目标 API,本任务对应的是 `flash_attn.flash_attn_interface` 中的 `flash_attn_with_kvcache`(paged KV cache 布局); - -2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口; - -3. 对照题包中的输入 shape、数据范围和精度要求; - -4. 编写自己的 `run_kernel(...)`; - -5. 提交 OJ,先通过正确性; - -6. 正确性通过后,再对比 baseline / OJ 耗时继续优化。 - +1. 从 benchmark 脚本中理解目标 API,本任务对应的是 `flash_attn.flash_attn_interface` 中的 `flash_attn_with_kvcache`(paged KV cache 布局); -### 题目说明(FlashAttention KV Cache Decode) +2. 在对应 OJ 题包中查看 `run_kernel(...)` 接口; -> 注意:每个子题的接口参数、数据范围和精度要求可能不同,正式要求以对应 XPU-OJ 题包为准。本节以 **FlashAttention KV Cache Decode** 为例,演示从 baseline benchmark 到 XPU-OJ 提交的完整流程。 +3. 对照题包中的输入 shape、数据范围和精度要求; -* **对应 baseline 脚本**:`flashattn_baseline/baseline/benchmark_kvcache.py` - -* **对应 FlashAttention API**:`flash_attn_with_kvcache`(paged KV cache 版本) - -* **算子说明**:实现 paged KV cache 下的 decode 注意力,每个 batch 只有 1 个 query token,KV cache 按 page 存储,长度由 `seqlen_k` 决定。 - -* **对应 OJ 题包**:XPU-OJ 上 `FlashAttention KV Cache Decode` 题(题号 20005)。 - +4. 编写自己的 `run_kernel(...)`; -### 步骤 8:理解 XPU-OJ 评测接口与精度要求 +5. 提交 OJ,先通过正确性; -**目标**:明确 Baseline 与最终评测提交之间的关系,理解选手需要实现什么。 +6. 正确性通过后,再对比 性能基线 / OJ 耗时继续优化。 -完成 baseline benchmark 后,需要注意:baseline 脚本主要用于建立性能基线,**不是最终提交物**。最终评测以 XPU-OJ 题包为准,评测程序会调用选手提交代码中的 `run_kernel`,并将输出结果与 baseline 参考结果进行比较。 -下面以 **FlashAttention KV Cache Decode** 题为例,题包目录中通常包含以下文件: +#### 6.8.1 题目说明(FlashAttention KV Cache Decode) + +> 注意:每个子题的接口参数、数据范围和精度要求可能不同,正式要求以对应 XPU-OJ 题包为准。本节以 **FlashAttention KV Cache Decode** 为例,演示从 baseline benchmark 到 XPU-OJ 提交的完整流程。 + +* **对应 benchmark 脚本**:`flashattn_baseline/baseline/benchmark_kvcache.py` + +* **对应 FlashAttention API**:`flash_attn_with_kvcache`(paged KV cache 版本) + +* **算子说明**:实现 paged KV cache 下的 decode 注意力,每个 batch 只有 1 个 query token,KV cache 按 page 存储,长度由 `seqlen_k` 决定。 + +* **对应 OJ 题包**:XPU-OJ 上 `FlashAttention KV Cache Decode` 题(题号 20005)。 + +### 6.9 步骤 8:理解 XPU-OJ 评测接口与精度要求 + +**目标**:明确 benchmark 脚本 与最终评测提交之间的关系,理解选手需要实现什么。 + +完成 benchmark 后,需要注意:benchmark 脚本 主要用于建立性能基线,**不是最终提交物**。最终评测以 XPU-OJ 题包为准,评测程序会调用选手提交代码中的 `run_kernel`,并将输出结果与 OJ 参考结果进行比较。 + +下面以 **FlashAttention KV Cache Decode** 题为例,题包目录中通常包含以下文件: * `zh_CN/00_题目描述.md`:说明需要实现的算子功能; - -* `zh_CN/01_接口约定.cuda.md`:说明必须实现的 `run_kernel` 函数签名; - + +* `zh_CN/01_接口约定.cuda.md`:说明必须实现的 `run_kernel` 函数签名; + * `zh_CN/05_数据范围与提示.md`:说明测试范围和精度要求; - -* `testcase_config.py`:定义测试数据生成、baseline 参考实现和正确性校验方式。 - -#### 1. 必须实现的接口 +* `testcase_config.py`:定义测试数据生成、**OJ 参考实现**(reference implementation)和正确性校验方式。 -选手需要在提交的 CUDA 源码中提供如下 C 符号,函数名、参数类型、顺序必须完全一致,并使用 `extern "C"` 防止 name mangling: + +#### 6.9.1 必须实现的接口 + +选手需要在提交的 CUDA 源码中提供如下 C 符号,函数名、参数类型、顺序必须完全一致,并使用 `extern "C"` 防止 name mangling: ```cpp #include @@ -637,33 +626,33 @@ extern "C" void run_kernel( ``` -#### 2. 参数说明 +#### 6.9.2 参数说明 | 参数 | 说明 | | --- | --- | -| `q` | decode query tensor,shape `(batch_size, seqlen_q, num_heads, headdim)`,连续 `bf16` | -| `k_cache_paged` | paged key cache,shape `(num_blocks, page_block_size, num_heads_k, headdim)`,连续 `bf16` | -| `v_cache_paged` | paged value cache,shape `(num_blocks, page_block_size, num_heads_k, headdim)`,连续 `bf16` | -| `output` | 输出缓冲区,shape `(batch_size, seqlen_q, num_heads, headdim)`,连续 `bf16` | -| `cache_seqlens` | 每个 batch 的 KV 长度,shape `(batch_size)`,连续 `int32` | -| `block_table` | 每个 batch 的 page 映射表,shape `(batch_size, num_blocks / batch_size)`,连续 `int32` | -| `seqlen_q` | query 长度,评测中固定为 `1` | -| `page_block_size` | page size,评测中固定为 `16` | -| `causal` | 是否启用 causal mask,评测中固定为 `0` | +| `q` | decode query tensor,shape `(batch_size, seqlen_q, num_heads, headdim)`,连续 `bf16` | +| `k_cache_paged` | paged key cache,shape `(num_blocks, page_block_size, num_heads_k, headdim)`,连续 `bf16` | +| `v_cache_paged` | paged value cache,shape `(num_blocks, page_block_size, num_heads_k, headdim)`,连续 `bf16` | +| `output` | 输出缓冲区,shape `(batch_size, seqlen_q, num_heads, headdim)`,连续 `bf16` | +| `cache_seqlens` | 每个 batch 的 KV 长度,shape `(batch_size)`,连续 `int32` | +| `block_table` | 每个 batch 的 page 映射表,shape `(batch_size, num_blocks / batch_size)`,连续 `int32` | +| `seqlen_q` | query 长度,评测中固定为 `1` | +| `page_block_size` | page size,评测中固定为 `16` | +| `causal` | 是否启用 causal mask,评测中固定为 `0` | -`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确,**不建议在** `**run_kernel**` **内部做** `**cudaDeviceSynchronize()**` **或显式同步**。 +`run_kernel` 内部需要自行计算合适的 launch 配置并启动 CUDA kernel。为保证计时准确,**不建议在** `**run_kernel**` **内部做** `**cudaDeviceSynchronize()**` **或显式同步**。 -#### 3. KV cache 布局 +#### 6.9.3 KV cache 布局 -KV cache layout 固定为 `flash_attn_with_kvcache` 的 paged cache 布局:`(num_blocks, page_block_size, num_heads_k, headdim)`。 +KV cache layout 固定为 `flash_attn_with_kvcache` 的 paged cache 布局:`(num_blocks, page_block_size, num_heads_k, headdim)`。 -第 `t` 个 KV token 位于 `block_table[batch_idx, t / page_block_size]` 指向的物理 page 中,page 内偏移为 `t % page_block_size`。 +第 `t` 个 KV token 位于 `block_table[batch_idx, t / page_block_size]` 指向的物理 page 中,page 内偏移为 `t % page_block_size`。 -例如 `batch_size = 1`、`seqlen_k = 512`、`page_block_size = 16` 时,每个序列需要访问 `32` 个有效 page。 +例如 `batch_size = 1`、`seqlen_k = 512`、`page_block_size = 16` 时,每个序列需要访问 `32` 个有效 page。 -#### 4. 评测数据范围 +#### 6.9.4 评测数据范围 -`testcase_config.py` 中定义了本题的测试配置(摘自题包): +`testcase_config.py` 中定义了本题的测试配置(摘自题包): ```python HEAD_DIMS = [128] @@ -677,113 +666,113 @@ CAUSAL = 0 ``` -`num_blocks = max(1024, ceil(seqlen_k / page_block_size) * batch_size * 3)`。 +`num_blocks = max(1024, ceil(seqlen_k / page_block_size) * batch_size * 3)`。 -#### 5. 精度要求 +#### 6.9.5 精度要求 -当前 `FlashAttention KV Cache Decode` 题的校验方式为: +当前 `FlashAttention KV Cache Decode` 题的校验方式为: ```python torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) ``` -也就是说,选手实现的输出需要在上述容差范围内与 baseline 输出一致。不同算子的容差可能不同,**正式精度要求以对应 OJ 题包说明为准**。 +也就是说,选手实现的输出需要在上述容差范围内与 OJ 参考输出 一致。不同算子的容差可能不同,**正式精度要求以对应 OJ 题包说明为准**。 -### 步骤 9：登录 XPU-OJ 并进入题目页面 +### 6.10 步骤 9:登录 XPU-OJ 并进入题目页面 -使用组委会统一发放的账号登录 XPU-OJ，并进入对应赛题页面。 +使用组委会统一发放的账号登录 XPU-OJ,并进入对应赛题页面。 -1.打开 XPU-OJ 平台：https://xpuoj.com/ +1.打开 XPU-OJ 平台:https://xpuoj.com/ -2.使用组委会统一发放的账号和初始密码登录【后续发布】； +2.使用组委会统一发放的账号和初始密码登录【后续发布】; -3.登录后进入比赛 / 题目列表页面； +3.登录后进入比赛 / 题目列表页面; -4.找到对应题目，例如 `20005 FlashAttention KV Cache Decode`； +4.找到对应题目,例如 `20005 FlashAttention KV Cache Decode`; -5.点击进入题目详情页，查看题目描述、接口约定、数据范围和提交入口。 +5.点击进入题目详情页,查看题目描述、接口约定、数据范围和提交入口。 -### 步骤 10：在Agent的帮助下提交 OJ 冒烟代码 +### 6.11 步骤 10:在Agent的帮助下提交 OJ 冒烟代码 -**目标：**完成一次最小提交，确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。 +**目标:**完成一次最小提交,确认 OJ 提交链路、语言环境和 `run_kernel(...)` 接口可用。 -**操作：** +**操作:** -1. 在语言下拉框中选择本题支持的提交语言，例如 `MXMACA C++`、`TileLang` 、 `Triton`； - -2. 将实现了题目要求接口的代码复制到提交框中； - - 如果你还没有 run\_kernel，应该从哪里开始？ - - OJ 最终评测不会直接运行 baseline 脚本，而是调用你提交代码中的 `run_kernel(...)`。   本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。  如果你还没有自己的 `run_kernel`，可以先让 Agent 阅读题包，并生成一个最小正确版实现思路。 - -3. 借助 Agent 从题包生成 run\_kernel 初版 - +1. 在语言下拉框中选择本题支持的提交语言,例如 `MXMACA C++`、`TileLang` 、 `Triton`; -### 10.1 在镜像终端中安装并启动 OpenCode +2. 将实现了题目要求接口的代码复制到提交框中; + + 如果你还没有 run\_kernel,应该从哪里开始? + + OJ 最终评测不会直接运行 benchmark 脚本,而是调用你提交代码中的 `run_kernel(...)`。 本赛题鼓励参赛者使用 AI Agent 辅助完成代码阅读、接口理解、初版实现、错误定位和性能优化。 如果你还没有自己的 `run_kernel`,可以先让 Agent 阅读题包,并生成一个最小正确版实现思路。 + +3. 借助 Agent 从题包生成 run\_kernel 初版 + + +#### 6.11.1 在镜像终端中安装并启动 OpenCode + +1. 首先返回镜像JupyterLab Terminal中, 打开容器内的 **Terminal**(终端) + +2. 执行下面的命令安装 OpenCode: -1. 首先返回镜像JupyterLab Terminal中, 打开容器内的 **Terminal**(终端) - -2. 执行下面的命令安装 OpenCode: - ```bash curl -fsSL https://opencode.ai/install | bash - + ``` - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/Lk3lbmbEQX2A6Om9/img/b4d06ad0-183b-47ce-a673-27a6a69ef348.png) - -3. 安装完成后,`cd` 进入赛题文件夹(根据实际目录调整): - + + ![b4d06ad0 183b 47ce a673 27a6a69ef348](https://origin.picgo.net/2026/06/18/b4d06ad0-183b-47ce-a673-27a6a69ef34860481172500d1bb3.png) + +3. 安装完成后,`cd` 进入赛题文件夹(根据实际目录调整): + ```bash cd xpuoj_problem/ - + ``` - -4. 在该目录下直接输入 `opencode` 并回车,即可进入 OpenCode 的 Agent 界面: - + +4. 在该目录下直接输入 `opencode` 并回车,即可进入 OpenCode 的 Agent 界面: + ```bash opencode - + ``` - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/Lk3lbmbEQX2A6Om9/img/60afdbc5-c334-4c78-8df9-80433091ddaf.png) - - 参考prompt： - + + ![60afdbc5 c334 4c78 8df9 80433091ddaf](https://origin.picgo.net/2026/06/18/60afdbc5-c334-4c78-8df9-80433091ddafd05eb034d28df07f.png) + + 参考prompt: + ```python 本题是 FlashAttention paged KV cache decode 的 CUDA C++ 前向算子。 - 输入输出规范如下： + 输入输出规范如下: - q: [batch, 1, num_heads, head_dim] (float) - - k_cache_paged: [num_blocks, num_heads, page_size, head_dim] (float) + - k_cache_paged: [num_blocks, num_heads, page_size, head_dim] (float) - v_cache_paged: [num_blocks, num_heads, page_size, head_dim] (float) - - cache_seqlen: [batch] (int) - - block_table: [batch, max_num_blocks] (int) - - output: [batch, 1, num_heads, head_dim] + - cache_seqlen: [batch] (int) + - block_table: [batch, max_num_blocks] (int) + - output: [batch, 1, num_heads, head_dim] 需调用 run_kernel(q, k_cache_paged, v_cache_paged, cache_seqlen, block_table, output) - 你需要根据 cache_seqlen 和 block_table 从 paged cache 中取出对应的 K, V，计算 attention 结果，存入 output。 - 参考：out = flash_attn_with_kvcache(q, k_cache_paged, v_cache_paged, cache_seqlen=..., block_table=...) - 要求你的实现与这个 API 的计算结果一致（误差允许1e-5）。 - 请生成一个 run_kernel 函数，优先保证正确性。 + 你需要根据 cache_seqlen 和 block_table 从 paged cache 中取出对应的 K, V,计算 attention 结果,存入 output。 + 参考:out = flash_attn_with_kvcache(q, k_cache_paged, v_cache_paged, cache_seqlen=..., block_table=...) + 要求你的实现与这个 API 的计算结果一致(误差允许1e-5)。 + 请生成一个 run_kernel 函数,优先保证正确性。 ``` - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/Lk3lbmbEQX2A6Om9/img/2efeb3c0-77b0-407b-9bff-815670bcbf00.png) - - 输出： - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/Lk3lbmbEQX2A6Om9/img/095af1c8-aa7e-44e3-ad0e-620206aaba1e.png) - + + ![2efeb3c0 77b0 407b 9bff 815670bcbf00](https://origin.picgo.net/2026/06/18/2efeb3c0-77b0-407b-9bff-815670bcbf0063a5b5b673b82532.png) + + 输出: + + ![095af1c8 aa7e 44e3 ad0e 620206aaba1e](https://origin.picgo.net/2026/06/18/095af1c8-aa7e-44e3-ad0e-620206aaba1e23a43b2eb7b0ead3.png) + 为方便参赛者先跑通完整流程,这里直接提供一份完整的冒烟代码,可直接复制粘贴到右侧编辑器,用于验证提交链路是否正常: - + ```python #include #include #include - + #define PAGE_SIZE 16 #define HEAD_DIM 128 - + __global__ void paged_attention_kernel( const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ k_cache_paged, @@ -802,42 +791,42 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) int batch_idx = blockIdx.x / num_heads; int head_idx = blockIdx.x % num_heads; if (batch_idx >= batch_size || head_idx >= num_heads) return; - + int seqlen = cache_seqlens[batch_idx]; if (seqlen <= 0) { - // 无有效 KV，输出 0 + // 无有效 KV,输出 0 int64_t out_base = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim; for (int i = threadIdx.x; i < headdim; i += blockDim.x) output[out_base + i] = __float2bfloat16(0.0f); return; } - + int tid = threadIdx.x; - - // 加载对应 head 的 query 元素（每个线程负责一个维度） + + // 加载对应 head 的 query 元素(每个线程负责一个维度) int64_t q_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim; float q_val = __bfloat162float(q[q_offset + tid]); - - // 全局 softmax 状态（每个线程维护自己维度的累加器） + + // 全局 softmax 状态(每个线程维护自己维度的累加器) float max_val = -1e38f; float sum_exp = 0.0f; float out_acc = 0.0f; float scale = rsqrtf(static_cast(headdim)); - - // 共享内存布局： + + // 共享内存布局: // K_tile[PAGE_SIZE][HEAD_DIM] (bf16) // V_tile[PAGE_SIZE][HEAD_DIM] (bf16) // partial_scores[PAGE_SIZE][HEAD_DIM] (float, 用于归约点积) __shared__ __nv_bfloat16 K_tile[PAGE_SIZE][HEAD_DIM]; __shared__ __nv_bfloat16 V_tile[PAGE_SIZE][HEAD_DIM]; __shared__ float partial_scores[PAGE_SIZE][HEAD_DIM]; - + int total_pages = (seqlen + PAGE_SIZE - 1) / PAGE_SIZE; - + for (int page = 0; page < total_pages; ++page) { int physical_block = block_table[batch_idx * blocks_per_batch + page]; int tokens_this_page = min(seqlen - page * PAGE_SIZE, PAGE_SIZE); - + // 1. 将当前 page 的 K 和 V 从全局显存加载到共享内存 // 每个线程负责加载所有 token 的同一个 head 维度 const int64_t kv_stride = num_heads_k * headdim; // 每个 (block, offset) 的 stride @@ -847,15 +836,15 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) V_tile[j][tid] = v_cache_paged[offset]; } __syncthreads(); - - // 2. 计算该 page 内每个 token 与 Q 的部分点积，存入 partial_scores + + // 2. 计算该 page 内每个 token 与 Q 的部分点积,存入 partial_scores for (int j = 0; j < tokens_this_page; ++j) { float k_val = __bfloat162float(K_tile[j][tid]); partial_scores[j][tid] = q_val * k_val; } __syncthreads(); - - // 3. 对 partial_scores 做 tree reduction，得到每个 token 的完整点积 + + // 3. 对 partial_scores 做 tree reduction,得到每个 token 的完整点积 #pragma unroll for (int stride = HEAD_DIM / 2; stride > 0; stride >>= 1) { if (tid < stride) { @@ -866,47 +855,47 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) } __syncthreads(); } - + // 4. 在线 safe softmax 更新 + V 累加 - // 4.1 找出本 page 内点积的最大值，结合全局 max 得到 new_max + // 4.1 找出本 page 内点积的最大值,结合全局 max 得到 new_max float local_max = -1e38f; #pragma unroll for (int j = 0; j < tokens_this_page; ++j) { local_max = fmaxf(local_max, partial_scores[j][0]); } float new_max = fmaxf(max_val, local_max); - + // 4.2 用旧的 max 对全局状态进行重缩放 float rescale = expf(max_val - new_max); sum_exp *= rescale; out_acc *= rescale; max_val = new_max; - - // 4.3 累加 V，并更新 sum_exp + + // 4.3 累加 V,并更新 sum_exp #pragma unroll for (int j = 0; j < tokens_this_page; ++j) { float score = partial_scores[j][0] * scale; float weight = expf(score - new_max); sum_exp += weight; - + float v_val = __bfloat162float(V_tile[j][tid]); out_acc += weight * v_val; } - + __syncthreads(); // 准备下一个 page 的共享内存加载 } - + // 5. 最终归一化并写回 if (sum_exp > 0.0f) { out_acc /= sum_exp; } else { out_acc = 0.0f; } - + int64_t out_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim + tid; output[out_offset] = __float2bfloat16(out_acc); } - + extern "C" void run_kernel( const __nv_bfloat16* q, const __nv_bfloat16* k_cache_paged, @@ -927,7 +916,7 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) int64_t blocks_per_batch = num_blocks / batch_size; dim3 grid(batch_size * num_heads); dim3 block(HEAD_DIM); - + paged_attention_kernel<<>>( q, k_cache_paged, v_cache_paged, output, cache_seqlens, block_table, @@ -936,102 +925,102 @@ torch.allclose(output_t.float(), output_ref.float(), rtol=1e-2, atol=1e-2) ); } ``` - - 以上代码仅用于说明接口结构，不代表最优实现，也不作为评分参考 - -5. 点击提交，等待评测结果返回； - - 评测时间与题目测试点数量、队列状态和平台负载有关，通常需要等待数十秒到数分钟。以平台实际返回为准。 - - ![0e9d0dccd68a1ddf1419973bc0c4e4bb.png](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/Yvenve5yZMWEwloy/img/95a36602-dcd7-48ea-9010-0dc4e7bb45d6.png) - - OJ 对每次提交大致会走这个流程： - + + 以上代码仅用于说明接口结构,不代表最优实现,也不作为评分参考 + +5. 点击提交,等待评测结果返回; + + 评测时间与题目测试点数量、队列状态和平台负载有关,通常需要等待数十秒到数分钟。以平台实际返回为准。 + + ![95a36602 dcd7 48ea 9010 0dc4e7bb45d6](https://origin.picgo.net/2026/06/18/95a36602-dcd7-48ea-9010-0dc4e7bb45d647d8de7b01ae032b.png) + + OJ 对每次提交大致会走这个流程: + ```plaintext - XPU-OJ 的一次评测大致流程如下： - 1. 选手提交代码； - 2. 平台按所选语言编译或加载提交代码； - 3. 评测程序构造测试输入； - 4. 调用选手代码中的 `run_kernel(...)`； - 5. 调用 `testcase_config.py` 中的 `baseline()` / 参考实现生成 `output_ref`； - 6. 将 `run_kernel(...)` 的输出与 `output_ref` 做正确性校验； - 7. 正确性通过后，统计运行耗时或性能指标； - 8. 根据题目评分规则换算该题得分； - 9. 更新该题历史最好成绩； - 10. 汇总各题最好成绩，得到排行榜总分。 + XPU-OJ 的一次评测大致流程如下: + 1. 选手提交代码; + 2. 平台按所选语言编译或加载提交代码; + 3. 评测程序构造测试输入; + 4. 调用选手代码中的 `run_kernel(...)`; + 5. OJ 后台参考实现生成 `output_ref`; + 6. 将 `run_kernel(...)` 的输出与 `output_ref` 做正确性校验; + 7. 正确性通过后,统计运行耗时或性能指标; + 8. 根据题目评分规则换算该题得分; + 9. 更新该题历史最好成绩; + 10. 汇总各题最好成绩,得到排行榜总分。 ``` - + 6. 查看结果 - + 提交详情会显示状态、得分、时间、内存、编译信息以及各测试点结果。 - + ![cd119ff85c024846b91d69b9ca9a1e00.png](https://origin.picgo.net/2026/06/17/cd119ff85c024846b91d69b9ca9a1e006456033e499779c6.png) -## 从 Baseline 到参赛作品的路径回顾： +## 7. 从性能基线到参赛作品的路径回顾 -1. 跑通 baseline benchmark，记录原始性能； - -2. 提交冒烟代码，确认 OJ 链路正常； - -3. 使用 Agent 阅读题包，理解输入输出、数据范围和精度要求； - -4. 在 `run_kernel(...)` 中实现最小正确版算子； - -5. 提交 OJ，先通过正确性； - -6. 正确性通过后，再让 Agent 辅助分析性能瓶颈； - -7. 围绕访存、softmax、线程划分、head\_dim 特化、K/V 复用等方向迭代优化； - -8. 保存每轮 Agent Prompt、代码改动、OJ 结果和性能变化，形成可复现的 Agent/Skill 优化流程。 - +1. 跑通 benchmark,建立性能基线; -## 七、常见问题 +2. 提交冒烟代码,确认 OJ 链路正常; -### Q1: 运行时提示 `ModuleNotFoundError: No module named 'flash_attn '` +3. 使用 Agent 阅读题包,理解输入输出、数据范围和精度要求; -**原因：** flash-attn 未安装或版本不兼容。 **解决：** 确认已安装 flash-attn = 2.6，可使用 `pip show flash-attn` 检查。 +4. 在 `run_kernel(...)` 中实现最小正确版算子; -### Q2: 所有配置都显示 OOM +5. 提交 OJ,先通过正确性; -**原因：** GPU 显存不足或 batch\\_size 设置过大。 **解决：** 减小 `batch_sizes` 和 `seq_lens_kv` 的范围重新测试。 +6. 正确性通过后,再让 Agent 辅助分析性能瓶颈; -### Q3: `mx-smi` 命令无输出或报错 +7. 围绕访存、softmax、线程划分、head\_dim 特化、K/V 复用等方向迭代优化; -**原因：** 沐曦 GPU 驱动未正确安装或环境变量未配置。 **解决：** 确认已正确配置 MXMACA 环境，检查 `/usr/local/maca` 目录是否存在。 +8. 保存每轮 Agent Prompt、代码改动、OJ 结果和性能变化,形成可复现的 Agent/Skill 优化流程。 -### Q4: 带宽数值异常低 -**原因：** 可能是 warmup 不足或 GPU 未达到稳态。 **解决：** 增加 `warmup` 次数，如 `warmup = 20`。 +## 8. 常见问题 -### Q5: 评测状态显示 `**Compile Error**`或提示 `**Undefined reference to run_kernel**` +### 8.1 Q1: 运行时提示 `ModuleNotFoundError: No module named 'flash_attn'` -**原因：**C++函数名被修饰（Name Mangling）或参数类型/顺序与接口约定不符。 **解决：**在 run\_kernel 前添加 extern "C"，并严格逐字核对参数的类型和修饰符。 +**原因:** flash-attn 未安装或版本不兼容。 **解决:** 确认已安装 flash-attn = 2.6,可使用 `pip show flash-attn` 检查。 -#### Q6: 评测状态显示 `**Wrong Answer**`，提示 torch.allclose校验失败 +### 8.2 Q2: 所有配置都显示 OOM -**原因：**bfloat16精度截断溢出、线程同步缺失或无效Token（Padding区域）处理错误。 **解决：**累加和 Softmax 强制转为 float32 计算；检查 \_\_syncthreads() 逻辑；增加 seqlen 的边界判空。 +**原因:** GPU 显存不足或 batch\\_size 设置过大。 **解决:** 减小 `batch_sizes` 和 `seq_lens_kv` 的范围重新测试。 -#### Q7: 评测状态显示 `**Runtime Error**`  (非法内存访问/段错误) +### 8.3 Q3: `mx-smi` 命令无输出或报错 -**原因：**Paged KV 地址映射索引错误、尾部 Page 越界读取，或线程块维度超限。 **解决：**仔细核对物理块寻址公式；增加当前 Page 有效 Token 数量的越界判断；检查单 Block 线程数配置。 +**原因:** 沐曦 GPU 驱动未正确安装或环境变量未配置。 **解决:** 确认已正确配置 MXMACA 环境,检查 `/usr/local/maca` 目录是否存在。 -#### Q8: 评测状态显示`**Time Limit Exceeded**` (评测超时) +### 8.4 Q4: 带宽数值异常低 -**原因：**发散分支内的 \_\_syncthreads() 导致内核死锁、误加主机端同步指令或并行度划分错误导致串行。 **解决：**确保同步指令在所有线程必经路径上；移除主机端多余的 cudaDeviceSynchronize()；优化 <<>> 参数以提升并行度。 +**原因:** 可能是 warmup 不足或 GPU 未达到稳态。 **解决:** 增加 `warmup` 次数,如 `warmup = 20`。 -## 八、下一步学习建议 +### 8.5 Q5: 评测状态显示 `Compile Error` 或提示 `Undefined reference to run_kernel` -完成本模块后，建议继续学习以下内容： +**原因:**C++函数名被修饰(Name Mangling)或参数类型/顺序与接口约定不符。 **解决:**在 run\_kernel 前添加 extern "C",并严格逐字核对参数的类型和修饰符。 -1. **算子优化基础** — 了解如何分析 kernel 性能瓶颈 - -2. **FlashAttention 源码解析** — 深入理解 `flash_attn_with_kvcache` 的实现原理 - -3. **自定义 kernel 开发** — 学习如何编写和优化沐曦 GPU 上的算子 - -4. **性能对比分析** — 将 baseline 结果与优化后结果进行对比 - -5. **OJ 题包深度解析** — 学习阅读题包中的 `testcase_config.py`，掌握本地构造边界用例与独立 Debug 的能力 - -6. **评测打榜与极限优化** — 在通过 OJ 正确性校验的基础上，挑战排行榜（Leaderboard），不断逼近硬件理论带宽极限 \ No newline at end of file +### 8.6 Q6: 评测状态显示 `Wrong Answer`，提示 torch.allclose 校验失败 + +**原因:**bfloat16精度截断溢出、线程同步缺失或无效Token(Padding区域)处理错误。 **解决:**累加和 Softmax 强制转为 float32 计算;检查 \_\_syncthreads() 逻辑;增加 seqlen 的边界判空。 + +### 8.7 Q7: 评测状态显示 `Runtime Error`（非法内存访问/段错误） + +**原因:**Paged KV 地址映射索引错误、尾部 Page 越界读取,或线程块维度超限。 **解决:**仔细核对物理块寻址公式;增加当前 Page 有效 Token 数量的越界判断;检查单 Block 线程数配置。 + +### 8.8 Q8: 评测状态显示 `Time Limit Exceeded`（评测超时） + +**原因:**发散分支内的 \_\_syncthreads() 导致内核死锁、误加主机端同步指令或并行度划分错误导致串行。 **解决:**确保同步指令在所有线程必经路径上;移除主机端多余的 cudaDeviceSynchronize();优化 <<>> 参数以提升并行度。 + +## 9. 下一步学习建议 + +完成本模块后,建议继续学习以下内容: + +1. **算子优化基础** - 了解如何分析 kernel 性能瓶颈 + +2. **FlashAttention 源码解析** - 深入理解 `flash_attn_with_kvcache` 的实现原理 + +3. **自定义 kernel 开发** - 学习编写和优化沐曦 GPU 上的算子 + +4. **性能对比分析** - 将 性能基线结果 与优化后结果进行对比 + +5. **OJ 题包深度解析** - 学习阅读题包中的 `testcase_config.py`,掌握本地构造边界用例与独立 Debug 的能力 + +6. **评测打榜与极限优化** - 在通过 OJ 正确性校验的基础上,挑战排行榜(Leaderboard),不断逼近硬件理论带宽极限 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim128.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim128.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim128.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim128.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim160.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim160.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim160.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim160.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim192.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim192.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim192.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim192.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim224.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim224.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim224.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim224.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim256.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim256.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim256.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim256.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim32.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim32.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim32.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim32.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim512.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim512.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim512.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim512.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim64.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim64.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim64.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim64.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim96.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim96.csv similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/baseline结果实例/benchmark_kvcache_headdim96.csv rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/baseline结果实例/benchmark_kvcache_headdim96.csv diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/benchmark_kvcache.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark_kvcache.py similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/Flashattn_Baseline/benchmark_kvcache.py rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/benchmark/benchmark_kvcache.py diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/guide.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/guide.md similarity index 100% rename from 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/guide.md rename to 基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/guide.md diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test_code.txt b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test_code.txt new file mode 100644 index 0000000..a8e1cbc --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/operator_task_package/flashattn_task_package/starter/smoke_test_code.txt @@ -0,0 +1,119 @@ +#include +#include +#include + +#define HEAD_DIM 128 + +__global__ void paged_attention_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* k_cache_paged, + const __nv_bfloat16* v_cache_paged, + __nv_bfloat16* output, + const int32_t* cache_seqlens, + const int32_t* block_table, + int64_t batch_size, + int64_t seqlen_q, + int64_t num_heads, + int64_t num_heads_k, + int64_t headdim, + int64_t page_block_size, + int64_t blocks_per_batch) +{ + int batch_idx = blockIdx.x / num_heads; + int head_idx = blockIdx.x % num_heads; + if (batch_idx >= batch_size || head_idx >= num_heads) return; + + int seqlen = cache_seqlens[batch_idx]; + int tid = threadIdx.x; + + // 加载对应 head 的 query 元素 + int64_t q_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim; + float q_val = __bfloat162float(q[q_offset + tid]); + + // Online safe softmax 状态 + float max_val = -1e38f; + float sum_exp = 0.0f; + float out_acc = 0.0f; + float scale = 1.0f / sqrtf(static_cast(headdim)); + + // 静态共享内存，避免动态分配可能带来的兼容性问题 + __shared__ float s_score[HEAD_DIM]; + + for (int token = 0; token < seqlen; ++token) { + int page_idx = token / page_block_size; + int page_offset = token % page_block_size; + int physical_block = block_table[batch_idx * blocks_per_batch + page_idx]; + + // 读取 key 元素 + const __nv_bfloat16* k_ptr = k_cache_paged + + (physical_block * page_block_size + page_offset) * (num_heads_k * headdim) + + head_idx * headdim; + float k_val = __bfloat162float(k_ptr[tid]); + + // 点积 -> 共享内存归约 + s_score[tid] = q_val * k_val; + __syncthreads(); + + for (int stride = HEAD_DIM >> 1; stride > 0; stride >>= 1) { + if (tid < stride) { + s_score[tid] += s_score[tid + stride]; + } + __syncthreads(); + } + float score = s_score[0] * scale; + + // 更新 softmax 状态 + float new_max = fmaxf(max_val, score); + float rescale = expf(max_val - new_max); + sum_exp = sum_exp * rescale + expf(score - new_max); + out_acc = out_acc * rescale; + max_val = new_max; + + // 读取 value 元素，并累加（用最新 max 的权重） + const __nv_bfloat16* v_ptr = v_cache_paged + + (physical_block * page_block_size + page_offset) * (num_heads_k * headdim) + + head_idx * headdim; + float v_val = __bfloat162float(v_ptr[tid]); + out_acc += expf(score - max_val) * v_val; + + __syncthreads(); // 确保下次迭代共享内存可安全复用 + } + + if (seqlen > 0) { + out_acc /= sum_exp; + } else { + out_acc = 0.0f; + } + + int64_t out_offset = ((batch_idx * seqlen_q + 0) * num_heads + head_idx) * headdim + tid; + output[out_offset] = __float2bfloat16(out_acc); +} + +extern "C" void run_kernel( + const __nv_bfloat16* q, + const __nv_bfloat16* k_cache_paged, + const __nv_bfloat16* v_cache_paged, + __nv_bfloat16* output, + const int32_t* cache_seqlens, + const int32_t* block_table, + int64_t batch_size, + int64_t seqlen_k, + int64_t seqlen_q, + int64_t num_heads, + int64_t num_heads_k, + int64_t headdim, + int64_t page_block_size, + int64_t num_blocks, + int64_t causal) +{ + int64_t blocks_per_batch = num_blocks / batch_size; + dim3 grid(batch_size * num_heads); + dim3 block(HEAD_DIM); + + paged_attention_kernel<<>>( + q, k_cache_paged, v_cache_paged, output, + cache_seqlens, block_table, + batch_size, seqlen_q, num_heads, num_heads_k, headdim, + page_block_size, blocks_per_batch + ); +} \ No newline at end of file -- 2.34.1