From 081a6c0a32cbde46b1f469c5c23d3814abda0d03 Mon Sep 17 00:00:00 2001 From: yyyymmm Date: Wed, 3 Jun 2026 18:35:33 +0800 Subject: [PATCH] =?UTF-8?q?Delete=20FlashAttention=20Baseline=20=E5=85=A5?= =?UTF-8?q?=E9=97=A8=EF=BC=9A=E4=BB=8E=E7=8E=AF=E5=A2=83=E9=AA=8C=E8=AF=81?= =?UTF-8?q?=E5=88=B0=20KV-Cache=20Benchmark=20=E7=BB=93=E6=9E=9C=E8=AE=B0?= =?UTF-8?q?=E5=BD=95.md?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...证到 KV-Cache Benchmark 结果记录.md | 829 ------------------ 1 file changed, 829 deletions(-) delete mode 100644 FlashAttention Baseline 入门：从环境验证到 KV-Cache Benchmark 结果记录.md diff --git a/FlashAttention Baseline 入门：从环境验证到 KV-Cache Benchmark 结果记录.md b/FlashAttention Baseline 入门：从环境验证到 KV-Cache Benchmark 结果记录.md deleted file mode 100644 index db1653c..0000000 --- a/FlashAttention Baseline 入门：从环境验证到 KV-Cache Benchmark 结果记录.md +++ /dev/null @@ -1,829 +0,0 @@ -# FlashAttention Baseline 入门：从环境验证到 KV-Cache Benchmark 结果记录（2） - -## 一、教程定位 - -本教程是参赛训练课程的 FlashAttention Baseline 入门模块，主要帮助用户快速跑通 FlashAttention 的 KV-Cache 推理性能基准测试最小可运行流程。 - -完成本教程后，用户应能够： - -* 完成环境验证和依赖检查 - -* 理解并配置基准测试参数 - -* 明确Baseline的解读，理解为什么KV-Cache是性能瓶颈 - -* 运行 KV-Cache Benchmark 测试 - -* 输出一份 baseline 性能结果记录表，为后续算子优化提供对比基准 - - -> Baseline 解读：为什么本教程基于 KV Cache 做性能基线？ - -#### 什么是 Baseline（性能基线） - -> Baseline是在引入任何优化代码前，系统处于初始可用状态时的参考数据。它是衡量后续所有优化收益的“锚点”。记录内容通常包括：执行时间、吞吐量、带宽、显存占用。 - -#### 为什么需要 Baseline - -> Baseline 是未做任何优化前的参考性能指标，它可以回答： - -* 当前性能处于什么水平？ - -* 是否存在明显性能瓶颈？ - -* 后续优化是否有效？ - -* 性能提升有多大？ - - -> 没有 Baseline，则无法量化优化效果。 - -> **注意**：Baseline ≠ Benchmark。 - -* **Benchmark** 是测量性能的手段（脚本、工具）。 - -* **Baseline** 是测量得到的具体结果数值。 - - -> 本教程运行的 `benchmark_kvcache.py` 是一个 Benchmark 脚本，它输出的 CSV 文件就是Baseline。 - -#### 为什么FlashAttention里面基线是针对KV Cache做benchmark性能验证？ - -> 因为KV Cache是最核心的性能瓶颈，尤其是在大模型推理的解码阶段。 - -##### 瓶颈从计算转移到显存访问 - -> Transformer 推理分为两个截然不同的阶段： - -| > 阶段 | > 特点 | > 计算模式 | > 主要限制 | -| --- | --- | --- | --- | -| > **Prefill（预填充）** | > 一次性处理全部输入 token | > 矩阵运算密集，Tensor Core 利用率高 | > **计算受限（Compute Bound）** | -| > **Decode（解码）** | > 逐个生成新 token，每步只算一个 token | > 每次都要**读取全部历史 KV Cache** | > **显存访问受限（Memory Bound）** | - -* Prefill 阶段计算量大但形态规整，通常能较好利用 GPU 算力，不是主要瓶颈。 - -* Decode 阶段占推理过程的大部分时间（尤其是长上下文交互），每个 token 的生成都需要搬运整个 KV Cache。随着序列增长，显存访问开销占比越来越高。 - - -##### KV‑Cache 的访存密集型特征 - -> 在 Decode 阶段，flash\_attn\_with\_kvcache 内核的工作是： - -* 读取当前 token 的 Q 向量（很小） - -* **反复读取整个历史的 KV Cache**（很大，线性增长） - -* 执行 FlashAttention 计算，然后将新的 K/V 追加写入缓存。 - - -> 95% 以上的时间花在读取 KV Cache 上。因此，Decode 阶段的性能完全由显存带宽决定，而不是 GPU 算力。 - -##### KV Cache 显存容量直接限制并发能力 - -> 大模型推理服务需要同时处理多个请求（batch\_size）。每个请求都有自己的 KV‑Cache，显存总占用量与 `batch_size × 序列长度` 成正比。因此，KV Cache 的显存开销直接决定了系统可以同时服务多少用户。 - -> 通过 Benchmark 对不同 `batch_size × seq_len_kv` 组合进行压力测试，可以： - -* 找出 **OOM 边界**：哪些参数组合会导致显存溢出，无法运行 - -* 量化每个请求的平均显存开销 - -* 为后续 **分页 KV‑Cache（PagedAttention）** 等优化提供基线对比。 - - -##### 为什么选择 flash\_attn\_with\_kvcache作为测试对象 - -> 它是 FlashAttention 专门为推理阶段设计的核心算子。它融合了高效的 KV Cache 读取与更新逻辑、FlashAttention 的分块、重计算技术以节省显存 - -> 推理场景下精确的带宽优化。选择它作为 Benchmark 对象，可以直接回答： - -* GPU带宽利用率是多少？是否接近理论峰值？ - -* 哪些参数组合（batch\_size, seq\_len\_kv, headdim）达到峰值性能？ - -* OOM边界在哪里？ - -* 后续优化（如分页缓存、算子融合）是否有效？ - - -#### 为什么要记录 Baseline 结果 - -> Baseline（基线）是优化前的参考性能数据，记录 baseline 的意义在于： - -* **量化优化收益**：优化后对比 baseline，计算加速比（speedup = baseline\_time / optimized\_time） - -* **防止性能回退**：代码变更后重跑 benchmark，确认没有引入性能退化 - -* **建立测试矩阵**：记录不同参数组合下的 baseline，全面了解性能特征 - - -> 在本教程中，baseline 结果以 CSV 文件保存，包含每种 `batch_size × seq_len_kv` 配置的执行时间和带宽，为后续算子优化提供对比基准。 - ---- - -## 二、学习目标 - -完成本模块后，你将能够： - -1. 理解 FlashAttention `flash_attn_with_kvcache` 核函数的基本作用与应用场景； - -2. 利用预装专属镜像，完成沐曦 GPU 硬件环境的快速验证； - -3. 深入理解 Baseline 的概念，掌握性能测试与正确性测试的联系与区别； - -4. 跑通 KV-Cache Benchmark 基准测试脚本； - -5. 完成多种 `batch_size × seq_len_kv` 组合的性能测试； - -6. 输出带宽性能结果 CSV 文件并进行结果分析。 - - ---- - -## 三、适用对象 - -本模块适合以下人员： - -* 参与 AI 基础设施竞赛的参赛者 - -* 对 GPU 算子性能优化感兴趣的开发者 - -* 需要了解 FlashAttention KV-Cache 推理性能的研究人员 - - -**基础知识要求：** - -* 了解 Python 编程基础 - -* 了解 PyTorch 基本用法 - -* 了解 GPU 推理的基本概念 - - ---- - -## 四、前置准备 - -开始实战前，请确认你已经完成以下准备： - -### 环境准备 - -#### 开发环境设置 - -1. **在沐曦开发者社区领取算力券** - - * 领取链接：[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6) - - * 登录平台 - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/2df2d4c4-1806-4c91-bc86-9a734f5e42cf.png) - - * 首次登录需要先进行注册（使用邮箱或者手机号进行注册） - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/3d2045a0-50de-42c0-9fef-26d1bad5f4e9.png) - - * 登录成功后进行第二步-邮箱验证，填入自己的邮箱。 - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/f0d8b34d-64d0-47bf-a7e6-579ca53eacee.png) - - * 第三步，提交申请。 - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/85740ff6-4b0b-40ae-9c85-3278c7a17e31.png) - - * 获得兑换码 - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/dd8b0766-1721-488a-b415-39f37d9671bb.png) - -2. **在模力方舟平台兑换算力券** - - * 平台链接：[https://ai.gitee.com/](https://ai.gitee.com/) - - * 1.登录模力方舟平台 - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/c404f350-9dc1-43e5-98bc-bce11e8d3127.png) - - * 2.进入费用中心 - 算力券 ， 点击右上角“兑换” - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/848d61e6-a1c3-4805-896a-e01fc7a6dde7.png) - -3. **租用算力** - - * 模力方舟算力市场链接：https://ai.gitee.com/compute - - * 选择沐曦芯片厂商，并根据项目要求选择相应的配置。 - - - ![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/6dbc8f01-11ef-4911-a042-815e951d71f6.png) - -4. **创建实例** - - -专属镜像文件： - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/2f0771d6-4417-4e41-bcbd-422e22202048.png) - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/85abeaab-bea1-423f-8a20-53a919014362.png) - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/218161bf-71a6-48c1-a2de-4e31d0c2b228.png) - -  进入算力容器，刚创建的实例默认开机状态，点击工具-lab开始项目创作。 - -**重要说明：**由于本次使用的是预装的专属镜像，环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后，无需再进行繁琐的依赖库版本验证即可直接进入测试环节。 - -### 代码准备 - -* 已获取基准测试脚本 `benchmark_kvcache.py` - -* 已进入项目目录 `/data/flashattn_baseline` - - ---- - -## 五、知识预备 - -### 名词解释 - -| 术语 | 说明 | -| --- | --- | -| **KV-Cache** | Key-Value Cache，Transformer 推理时缓存历史 token 的 Key 和 Value 向量，避免重复计算 | -| **Paged KV-Cache** | 将 KV-Cache 分页管理，提高显存利用率，类似操作系统的虚拟内存分页机制 | -| **flash\_attn\_with\_kvcache** | FlashAttention 提供的带 KV-Cache 支持的注意力计算核函数 | -| **batch\_size** | 批大小，一次处理的样本数量 | -| **seq\_len\_kv** | KV 序列长度，KV-Cache 中缓存的历史 token 数量 | -| **headdim** | Head Dimension，注意力头的维度 | -| **带宽 (Bandwidth)** | 显存带宽，单位 GB/s，衡量 GPU 读写显存的速度 | - -### 核心概念详解 - -#### 什么是正确性测试与性能测试 - -* **正确性测试 (Correctness Testing)：**解决“算得对不对”的问题。它的目标是验证当前算子的输出结果，在数学精度上是否与标准参考实现完全一致。这是所有测试的绝对前提底线。 - -* **性能测试 (Performance Testing)：**解决“跑得快不快”的问题。它的目标是在验证正确性的基础上，测量算子在特定硬件上的执行耗时、吞吐量和有效带宽利用率。本教程执行的 Benchmark 脚本，正是一个纯粹的性能测试。 - -* **二者区别：** - - -| 维度 | 性能测试 | 正确性测试 | -| --- | --- | --- | -| 测试目标 | 测量速度、带宽 | 验证输出结果 | -| 关注输出 | 否 | 是 | -| 关注效率 | 是 | 否 | -| 是否需要 Baseline | 是 | 不一定 | -| 是否受实现不同而影响 | 大 | 是（精度不同可能影响） | - -**二者联系：** - -在实际开发中： - -正确性测试（先） - -      ↓ - -建立 Baseline - -      ↓ - -性能分析 - -      ↓ - -优化实现 - -      ↓ - -性能测试对比 Baseline - -      ↓ - -回归正确性验证（保证没变坏） - -在算子优化迭代中，每一次修改底层代码，都必须**先通过正确性测试**确立功能基准，**再运行性能测试**对比性能基准 Baseline，确保速度的提升绝不是以牺牲结果正确性为代价。 - -#### 什么是 Benchmark（基准测试） - -Benchmark 是一种标准化的性能测量方法，通过在固定条件下反复运行同一任务，获取可重复、可对比的性能指标。在 GPU 算子优化场景中，benchmark 的作用是： - -* **建立性能基线**：在优化前记录原始性能数据，作为后续对比的参照 - -* **量化优化效果**：优化后运行同样的 benchmark，直接对比时间/带宽变化 - -* **发现性能瓶颈**：通过不同参数组合的测试结果，定位性能拐点 - - -> 参考：[MLPerf Benchmark 介绍](https://mlcommons.org/benchmarks/) - -#### 什么是 batch\_size、seq\_len、headdim - -这三个参数共同决定了注意力计算的**工作量**和**显存占用**： - -* **batch\_size（批大小）**：一次推理同时处理的样本数量。batch\_size 越大，GPU 并行度越高，但显存占用也线性增长。在 KV-Cache 场景中，batch\_size 对应同时服务的请求数。 - -* **seq\_len / seq\_len\_kv（序列长度）**：序列中 token 的数量。seq\_len\_kv 特指 KV-Cache 中已缓存的历史 token 数量。序列越长，注意力计算的计算量呈 O(n²) 增长（但 FlashAttention 将其优化为 O(n) 显存），KV-Cache 的显存占用则呈 O(n) 线性增长。 - -* **headdim（注意力头维度）**：每个注意力头的向量维度。常见的有 64、128、256。headdim 越大，单个 token 的 Key/Value 向量越宽，KV-Cache 的显存占用与 headdim 成正比。 - - -三者与显存占用的关系： - -```Plain -KV-Cache 显存 ≈ batch_size × seq_len_kv × num_heads_k × headdim × 2(K+V) × bytes_per_elem -``` - -> 参考：[Attention Is All You Need (Vaswani et al., 2017)](https://arxiv.org/abs/1706.03762) - -#### 什么是 Kernel 执行时间 - -Kernel（核函数）是运行在 GPU 上的并行计算函数。Kernel 执行时间指从 GPU 开始执行该核函数到执行完毕所花费的时间，通常以**毫秒 (ms)** 为单位。 - -测量方式有两种： - -* **CPU 端计时**：使用 `torch.cuda.synchronize()` + `time.time()`，包含 GPU 调度开销，时间偏大 - -* **GPU 端计时**：使用 CUDA Event 或 profiler，精度更高，直接测量 GPU 上的实际执行时间 - - -本教程使用 GPU 端同步计时。 - -> 参考：[PyTorch CUDA Semantics](https://pytorch.org/docs/stable/notes/cuda.html) - -#### 什么是有效带宽 - -有效带宽（Effective Bandwidth）是衡量 kernel 实际利用显存带宽效率的指标，计算公式为： - -```Plain -有效带宽 (GB/s) = 数据传输量 (GB) / kernel 执行时间 (s) - -``` - -GPU 显存带宽是有限的（例如沐曦 C500 的理论峰值带宽），有效带宽越接近理论峰值，说明 kernel 对显存带宽的利用率越高。对于**访存密集型**算子（如 KV-Cache 注意力），有效带宽是衡量优化效果的核心指标。 - -* 有效带宽 **接近理论峰值** → kernel 已接近最优，优化空间有限 - -* 有效带宽 **远低于理论峰值** → 存在优化空间（如内存访问不合并、bank conflict 等） - - -> 参考：[CUDA C++ Programming Guide - Performance Guidelines](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#performance-guidelines) - -#### 为什么要 Warmup / Repeat - -GPU 程序的首次运行往往比后续运行慢，原因包括： - -* **JIT 编译**：部分框架会延迟编译 kernel 代码 - -* **缓存冷启动**：GPU L2 Cache、TLB 等初始状态为空 - -* **频率爬升**：GPU 需要时间从低功耗状态切换到高频率状态 - - -因此，benchmark 流程通常分为两步： - -1. **Warmup（预热）**：先运行若干次（如 10 次），不记录时间，让 GPU 进入稳定状态 - -2. **Repeat（重复测量）**：正式运行多次（如 100 次），记录每次时间，取统计值（均值/中位数） - - -重复测量可以消除随机波动，获得更可靠的性能数据。次数越多，结果越稳定，但耗时也越长。 - -> 参考：[PyTorch Benchmark Utils](https://pytorch.org/tutorials/recipes/recipes/benchmark.html) - -#### CUDA Stream 与同步 - -CUDA 采用异步执行模型，CPU 提交 kernel 到 GPU 后不等待完成就继续执行。`torch.cuda.synchronize()` 会阻塞 CPU 直到 GPU 上所有已提交的任务完成，这是精确计时的前提。 - -> 参考：[CUDA Streams](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#asynchronous-concurrent-execution) - -#### 数据类型（dtype）对性能的影响 - -不同数据类型占用的字节数不同，直接影响显存带宽需求和计算吞吐： - -| 数据类型 | 字节数 | 说明 | -| --- | --- | --- | -| float32 | 4 | 单精度浮点，精度最高 | -| float16 | 2 | 半精度浮点，精度足够且带宽减半 | -| bfloat16 | 2 | Brain Float 16，动态范围与 float32 相同，训练/推理常用 | - -本教程使用 `bfloat16`，在精度和性能之间取得平衡。 - -> 参考：[Mixed Precision Training (Micikevicius et al., 2018)](https://arxiv.org/abs/1710.03740) - -#### Paged KV-Cache 与 Block Table - -传统 KV-Cache 为每个请求预分配连续显存，容易造成碎片和浪费。Paged KV-Cache（灵感来自操作系统虚拟内存）将显存分成固定大小的 page/block，通过 **block\_table** 映射逻辑位置到物理位置： - -* **page\_block\_size**：每个 block 包含的 token 数量（本教程默认 16） - -* **block\_table**：索引张量，记录每个 batch 的 KV-Cache 页面映射关系 - -* **优势**：减少显存碎片，支持动态分配，提高多请求并发效率 - - -> 参考：[Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023)](https://arxiv.org/abs/2309.06180) - -#### OOM（Out of Memory） - -OOM 表示 GPU 显存不足，无法完成当前计算。常见原因： - -* batch\_size 或 seq\_len\_kv 过大，超出显存容量 - -* 同时存在多个占用显存的进程 - -* 未释放的中间变量占用显存 - - -应对策略：减小 batch\_size/seq\_len\_kv、使用更小的 dtype（如 bfloat16 替代 float32）、使用梯度检查点等。 - -> 参考：[PyTorch CUDA Memory Management](https://pytorch.org/docs/stable/notes/cuda.html#memory-management) - -#### Tensor Core 与矩阵乘法加速 - -现代 GPU（包括沐曦 C500）配备 Tensor Core 单元，专门加速矩阵乘法运算。Attention 计算中的 Q×K^T 和 Attn×V 都是矩阵乘法，能够受益于 Tensor Core 加速。Tensor Core 对数据类型和矩阵维度有对齐要求（通常要求维度为 8 或 16 的倍数），这也是 headdim 通常取 64/128/256 的原因之一。 - -> 参考：[NVIDIA Tensor Core Technology](https://developer.nvidia.com/tensor-cores) - -### 相关链接 - -* [FlashAttention 官方仓库](https://github.com/Dao-AILab/flash-attention) - -* [FlashAttention API 文档](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_interface.py) - -* [FlashAttention 论文 (Dao et al., 2022)](https://arxiv.org/abs/2205.14135) - -* [FlashAttention-2 论文 (Dao, 2023)](https://arxiv.org/abs/2307.08691) - -* [PyTorch CUDA 编程最佳实践](https://pytorch.org/docs/stable/notes/cuda.html) - - ---- - -## 六、项目实践：FlashAttention KV-Cache Benchmark - -### 项目目标 - -对 FlashAttention 的 paged KV-cache 推理核函数（`flash_attn_with_kvcache`）进行自动化性能基准测试，覆盖多种 `batch_size × seq_len_kv` 组合，输出执行时间和有效显存带宽。 - -基准测试脚本 `benchmark_kvcache.py` 的最小闭环包括： - -1. 准备 paged KV-cache 张量和 block table。 - -2. 通过 `torch.profiler` 对 kernel 进行计时。 - -3. 计算有效显存带宽（GB/s）。 - -4. 将结果写入带时间戳的 CSV 文件。 - - ---- - -### 步骤 0：进入创建的实例环境 - -模力方舟链接：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) - -### 步骤 1：检查运行环境 - -**目标：** 确认当前环境满足本模块运行要求。 - -在JupyterLab Terminal中检查运行环境的配置。 - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/339ac3f8-e31d-43c6-992b-07e20e35ef95.png) - -**操作：** 检查 GPU 状态、Python 版本和依赖版本。 - -**命令示例：** - -```Bash -# 检查沐曦 GPU 状态 -mx-smi - -# 检查 Python 版本 -python --version - -# 检查 PyTorch 是否能识别 GPU -python -c "import torch; print(f'GPU available: {torch.cuda.is_available()}'); print(f'GPU count: {torch.cuda.device_count()}')" - -# 检查依赖版本 -python -c "import torch; print(f'PyTorch {torch.__version__}')" -python -c "import flash_attn; print(f'flash-attn {flash_attn.__version__}')" -python -c "import einops; print('einops OK')" - -``` - -**预期结果：** - -* `mx-smi` 显示沐曦 GPU 信息 - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/016119b5-960b-478d-9119-af27b9f3c727.png) - -* Python 版本 = 3.8 - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/568886e4-85a7-41c6-a94c-89a09913e0f7.png) - -* `torch.cuda.is_available()` 返回 `True` - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/d1c97030-d40f-416f-89b3-b81caff88966.png) - -* 所有依赖版本符合要求 - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/352b41ee-8a87-405e-991b-1c3714bdeff0.png) - -**常见问题：** - -| 问题 | 解决方法 | -| --- | --- | -| `torch.cuda.is_available()` 返回 `False` | 检查 MXMACA 环境变量是否正确配置 | -| `ModuleNotFoundError: No module named 'flash_attn '` | 确认 flash-attn 已安装且版本= 2.6 | -| `mx-smi` 命令不存在 | 确认已配置沐曦 GPU 驱动环境 | - ---- - -### 步骤 2：进入项目目录 - -**目标：** 进入本模块所需的项目目录。 - -1. 将flashattn\_baseline压缩文件进行解压，并将解压文件加入JupyterLab。 - - -![image](https://alidocs.oss-cn-zhangjiakou.aliyuncs.com/res/YdgOk2bRrmLe7q4B/img/3eb4d4fe-e948-405e-8d11-6301b4d8f16e.png) - -**操作：** 切换到基准测试脚本所在目录。 - -**命令示例：** - -```Bash -cd flashattn_baseline -ls -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：配置基准测试参数 - -**目标：** 根据测试需求配置基准测试参数。 - -**操作：** 在 `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 推理） | -| `dtype` | `torch.bfloat16` | 数据类型 | -| `causal` | `False` | 是否启用 causal mask | -| `warmup` | `10` | 预热迭代次数 | -| `repeat` | `100` | 正式 profiling 迭代次数 | - -**配置示例：** - -如需测试 `headdim=128`，修改对应列表： - -```Python -headdims = [128] - -``` - -如需启用 causal mask： - -```Python -causal = True - -``` ---- - -### 步骤 4：运行基准测试 - -**目标：** 运行基准测试脚本，收集性能数据。 - -**操作：** 执行基准测试脚本并观察输出。 - -**命令示例：** - -```Bash -cd flashattn_baseline -python benchmark_kvcache.py - -``` - -**预期结果：** - -脚本运行时会在终端实时打印结果表格： - -```Plain -batch_size seq_len_kv heads headdim time_ms bandwidth_GB_s ---------------------------------------------------------------------------- - 1 512 8 256 0.0948 44.27 - 2 512 8 256 0.0988 84.95 - ... - -``` - -同时生成带时间戳的 CSV 文件，命名格式为 `benchmark_kvcache_YYYYMMDD_HHMMSS.csv`。 - -**常见问题：** - -| 问题 | 解决方法 | -| --- | --- | -| 出现 `OOM` 标记 | 该配置超出 GPU 显存容量，可减小 batch\_size 或 seq\_len\_kv | -| 脚本运行缓慢 | 减少 `repeat` 次数或缩小扫描范围 | - ---- - -### 步骤 5：查看与分析结果 - -**目标：** 理解输出结果格式，分析性能数据。 - -**CSV 输出格式：** - -| 列名 | 说明 | -| --- | --- | -| `batch_size` | 批大小 | -| `seq_len_kv` | KV 序列长度 | -| `heads` | head 数量 | -| `headdim` | head 维度 | -| `time_ms` | kernel 执行时间（毫秒） | -| `bandwidth_GB_s` | 有效显存带宽（GB/s） | - -如果某个配置因显存不足而失败，对应的 `time_ms` 和 `bandwidth_GB_s` 列会标记为 `OOM`。 - -**带宽计算公式：** - -```Plain -total_bytes = q_bytes + kv_bytes -q_bytes = batch_size × seqlen_q × num_heads × headdim × bytes_per_elem -kv_bytes = batch_size × seqlen_k × num_heads_k × headdim × bytes_per_elem × 2 -bandwidth = (total_bytes / 1e9) / (time_ms / 1e3) [GB/s] - -``` - -其中 `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 - - -**测试结果示例：** - -#### headdim=128（2026-05-26） - -所有 48 个配置均成功运行，峰值带宽约 1251 GB/s。 - -| batch\_size | seq\_len\_kv | time\_ms | bandwidth\_GB\_s | -| --- | --- | --- | --- | -| 1 | 512 | 0.0322 | 65.27 | -| 128 | 512 | 0.2453 | 1095.50 | -| 1 | 16384 | 0.8356 | 80.32 | -| 128 | 16384 | 6.8668 | 1250.98 | - -#### headdim=256（2026-05-27） - -48 个配置中有 3 个因显存不足（OOM）而失败，峰值带宽约 807 GB/s。 - -| batch\_size | seq\_len\_kv | time\_ms | bandwidth\_GB\_s | -| --- | --- | --- | --- | -| 1 | 512 | 0.0948 | 44.27 | -| 128 | 4096 | 5.3221 | 807.10 | -| 128 | 8192 | OOM | OOM | -| 64 | 16384 | OOM | OOM | -| 128 | 16384 | OOM | OOM | - ---- - -### 步骤 6：自定义扩展（可选） - -如需进一步测试，可参考以下扩展方法： - -**修改 headdim：** - -```Python -headdims = [128, 256] - -``` - -**启用 causal mask：** - -```Python -causal = True - -``` - -**调整 profiling 精度：** - -```Python -warmup = 20 -repeat = 200 - -``` - -**启用详细 profiler 输出：** - -```Python -ms = run_with_profiler(run_fn, warmup=warmup, reps=repeat, print_result=True, target_kernels=["flash"]) - -``` ---- - -## 七、Agent 使用说明 - -在本模块中，Agent 可用于以下场景： - -### Prompt 模板 - -**环境验证：** - -```Plain -请帮我验证当前环境是否满足 FlashAttention KV-Cache Benchmark 的运行要求，包括： -1. 沐曦 GPU 是否可见 -2. PyTorch 版本和 CUDA 支持 -3. flash-attn 和 einops 是否已安装 - -``` - -**参数配置建议：** - -```Plain -我需要测试 headdim=128 和 headdim=256 的性能差异，请帮我推荐合适的 batch_sizes 和 seq_lens_kv 扫描范围。 - -``` - -**结果分析：** - -```Plain -请帮我分析这份 benchmark 结果 CSV 文件，找出峰值带宽配置和 OOM 边界。 - -``` ---- - -## 八、常见问题 - -### Q1: 运行时提示 `ModuleNotFoundError: No module named 'flash_attn '` - -**原因：** flash-attn 未安装或版本不兼容。 **解决：** 确认已安装 flash-attn = 2.6，可使用 `pip show flash-attn` 检查。 - -### Q2: 所有配置都显示 OOM - -**原因：** GPU 显存不足或 batch\\_size 设置过大。 **解决：** 减小 `batch_sizes` 和 `seq_lens_kv` 的范围重新测试。 - -### Q3: `mx-smi` 命令无输出或报错 - -**原因：** 沐曦 GPU 驱动未正确安装或环境变量未配置。 **解决：** 确认已正确配置 MXMACA 环境，检查 `/usr/local/maca` 目录是否存在。 - -### Q4: 带宽数值异常低 - -**原因：** 可能是 warmup 不足或 GPU 未达到稳态。 **解决：** 增加 `warmup` 次数，如 `warmup = 20`。 - ---- - -## 九、下一步学习建议 - -完成本模块后，建议继续学习以下内容： - -1. **算子优化基础** — 了解如何分析 kernel 性能瓶颈 - -2. **FlashAttention 源码解析** — 深入理解 `flash_attn_with_kvcache` 的实现原理 - -3. **自定义 kernel 开发** — 学习如何编写和优化沐曦 GPU 上的算子 - -4. **性能对比分析** — 将 baseline 结果与优化后结果进行对比 - - ---- \ No newline at end of file