forked from metax-maca/op_optimization
Add FlashAttention baseline guide
This commit is contained in:
何沐川
parent
2c1367e0c1
commit
6c579ad038
|
|
@ -0,0 +1,829 @@
|
|||
# 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)
|
||||
|
||||
* 登录平台
|
||||
|
||||
|
||||
|
||||
|
||||
* 首次登录需要先进行注册（使用邮箱或者手机号进行注册）
|
||||
|
||||
|
||||
|
||||
|
||||
* 登录成功后进行第二步-邮箱验证，填入自己的邮箱。
|
||||
|
||||
|
||||
|
||||
|
||||
* 第三步，提交申请。
|
||||
|
||||
|
||||
|
||||
|
||||
* 获得兑换码
|
||||
|
||||
|
||||
|
||||
|
||||
2. **在模力方舟平台兑换算力券**
|
||||
|
||||
* 平台链接：[https://ai.gitee.com/](https://ai.gitee.com/)
|
||||
|
||||
* 1.登录模力方舟平台
|
||||
|
||||
|
||||
|
||||
|
||||
* 2.进入费用中心 - 算力券 ， 点击右上角“兑换”
|
||||
|
||||
|
||||
|
||||
|
||||
3. **租用算力**
|
||||
|
||||
* 模力方舟算力市场链接：https://ai.gitee.com/compute
|
||||
|
||||
* 选择沐曦芯片厂商，并根据项目要求选择相应的配置。
|
||||
|
||||
|
||||
|
||||
|
||||
4. **创建实例**
|
||||
|
||||
|
||||
专属镜像文件：
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
进入算力容器，刚创建的实例默认开机状态，点击工具-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进入实例环境
|
||||
|
||||
|
||||
|
||||
### 步骤 1：检查运行环境
|
||||
|
||||
**目标：** 确认当前环境满足本模块运行要求。
|
||||
|
||||
在JupyterLab Terminal中检查运行环境的配置。
|
||||
|
||||
|
||||
|
||||
**操作：** 检查 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 信息
|
||||
|
||||
|
||||
|
||||
|
||||
* Python 版本 = 3.8
|
||||
|
||||
|
||||
|
||||
|
||||
* `torch.cuda.is_available()` 返回 `True`
|
||||
|
||||
|
||||
|
||||
|
||||
* 所有依赖版本符合要求
|
||||
|
||||
|
||||
|
||||
|
||||
**常见问题：**
|
||||
|
||||
| 问题 | 解决方法 |
|
||||
| --- | --- |
|
||||
| `torch.cuda.is_available()` 返回 `False` | 检查 MXMACA 环境变量是否正确配置 |
|
||||
| `ModuleNotFoundError: No module named 'flash_attn '` | 确认 flash-attn 已安装且版本= 2.6 |
|
||||
| `mx-smi` 命令不存在 | 确认已配置沐曦 GPU 驱动环境 |
|
||||
|
||||
---
|
||||
|
||||
### 步骤 2：进入项目目录
|
||||
|
||||
**目标：** 进入本模块所需的项目目录。
|
||||
|
||||
1. 将flashattn\_baseline压缩文件进行解压，并将解压文件加入JupyterLab。
|
||||
|
||||
|
||||
|
||||
|
||||
**操作：** 切换到基准测试脚本所在目录。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```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 结果与优化后结果进行对比
|
||||
|
||||
|
||||
---
|
||||
Loading…
Reference in New Issue