forked from metax-maca/op_optimization
添加AI Agent相关算子优化及baselines文件
This commit is contained in:
xiao
parent
d12739c5d9
commit
999c413112
|
|
@ -1,42 +1,43 @@
|
|||
# Agent 使用说明
|
||||
本说明涵盖三个算子的 Agent 辅助能力（FlashInfer、FlashAttention、Fused MoE）以及使用 Agent 时的推荐方式和注意事项。
|
||||
|
||||
## 一、FlashInfer：
|
||||
## 一、模块功能速查
|
||||
### 1.FlashInferFlashInfer：
|
||||
在本模块中，Agent 可以帮助你完成以下任务：
|
||||
|
||||
### 环境检查
|
||||
#### 环境检查
|
||||
|
||||
```plaintext
|
||||
请帮我检查当前环境是否满足 FlashInfer 运行要求，包括 GPU、Python、PyTorch 和 flashinfer 依赖。
|
||||
```
|
||||
|
||||
### 运行测试
|
||||
#### 运行测试
|
||||
|
||||
```plaintext
|
||||
请帮我运行 bench_batch_decode.py 脚本，执行 BatchDecode 的基准测试。
|
||||
```
|
||||
|
||||
### 分析结果
|
||||
#### 分析结果
|
||||
|
||||
```plaintext
|
||||
请帮我读取最新的 CSV 结果文件，分析各参数配置下的性能表现，找出带宽最高和 TFLOPs 最高的配置。
|
||||
```
|
||||
|
||||
### 问题排查
|
||||
#### 问题排查
|
||||
|
||||
```plaintext
|
||||
运行时报错 out of memory，请帮我分析原因并给出解决方案。
|
||||
```
|
||||
|
||||
### 代码理解
|
||||
#### 代码理解
|
||||
|
||||
```plaintext
|
||||
请帮我解释 bench_common.py 中 run_with_profiler 函数的工作原理。
|
||||
```
|
||||
|
||||
## 二、FlashAttention：
|
||||
### 2. FlashAttention：
|
||||
在本模块中，Agent 可用于以下场景：
|
||||
|
||||
### Prompt 模板
|
||||
#### Prompt 模板
|
||||
|
||||
**环境验证：**
|
||||
|
||||
|
|
@ -62,10 +63,10 @@
|
|||
|
||||
```
|
||||
---
|
||||
## 三、Fused MoE：
|
||||
### 3. Fused MoE：
|
||||
在本模块中，Agent可以帮助你完成以下任务：
|
||||
|
||||
### 环境检查
|
||||
#### 环境检查
|
||||
|
||||
```plaintext
|
||||
我正在算力平台部署 fused_moe_baseline 源码。
|
||||
|
|
@ -80,13 +81,13 @@
|
|||
2. 是否有潜在的不兼容风险（如 Python 与 libpython 版本）？
|
||||
```
|
||||
|
||||
### 运行测试
|
||||
#### 运行测试
|
||||
|
||||
```plaintext
|
||||
请帮我运行 scripts/run_fused_moe_i8_tn_pybind_test.sh 脚本
|
||||
```
|
||||
|
||||
### 分析结果
|
||||
#### 分析结果
|
||||
|
||||
```plaintext
|
||||
这是性能测试结果：
|
||||
|
|
@ -100,7 +101,7 @@
|
|||
3. 当前结果是否已经具备提交价值？
|
||||
```
|
||||
|
||||
### 报错检查
|
||||
#### 报错检查
|
||||
|
||||
```plaintext
|
||||
编译 pybind 时出现以下错误：
|
||||
|
|
@ -116,15 +117,34 @@
|
|||
3. 如何在 build_fused_moe_i8_tn_pybind.sh 中正确指定路径？
|
||||
```
|
||||
|
||||
### 代码理解
|
||||
#### 代码理解
|
||||
|
||||
```plaintext
|
||||
请帮我梳理释 benchmark_fused_moe_i8_tn.py 代码整体框架
|
||||
```
|
||||
|
||||
### KernelSwift 系统搜索算子
|
||||
#### KernelSwift 系统搜索算子
|
||||
|
||||
|
||||
```plaintext
|
||||
请帮我在算子广场检索 fused_moe 算子
|
||||
```
|
||||
```
|
||||
## 二、推荐工作方式
|
||||
使用 Agent 时，不要一开始就让它“直接优化到最快”。推荐节奏是：
|
||||
|
||||
先让它检查目录。
|
||||
|
||||
再让它跑通原始 baseline。
|
||||
|
||||
然后让它只做一处小修改。
|
||||
|
||||
每次修改后必须 build、test、benchmark。
|
||||
|
||||
每轮都记录结果。
|
||||
|
||||
## 三、注意事项
|
||||
使用 Agent 不要相信没有命令输出支持的环境判断。
|
||||
|
||||
不要让 Agent 同时修改很多文件。
|
||||
|
||||
不要让 Agent 修改测试标准。
|
||||
|
|
@ -163,79 +163,30 @@
|
|||
---
|
||||
|
||||
## 四、前置准备
|
||||
|
||||
开始实战前，请确认你已经完成以下准备：
|
||||
|
||||
### 环境准备
|
||||
* **设置领取与兑换算力券**
|
||||
|
||||
* 前往沐曦开发者社区注册账号并完成邮箱验证，申请并获取 MACA 算力代金券兑换码。
|
||||
|
||||
|
||||
链接：[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6)
|
||||
|
||||
* 登录 模力方舟平台 (Gitee AI)，在“费用中心 -> 算力券”页面输入兑换码完成充值。
|
||||
|
||||
|
||||
链接：[https://ai.gitee.com/](https://ai.gitee.com/)
|
||||
|
||||
* **创建并启动实例**
|
||||
|
||||
* 进入 算力市场，筛选“沐曦”芯片厂商，选择合适的 GPU 规格（推荐 曦云 C500 节点）。
|
||||
|
||||
* **关键配置：** 在预装镜像处，务必选择专属开发镜像（`PyTorch Agent/2.8.0/Python 3.12/maca 3.7.2.1`）。
|
||||
|
||||
* 创建完成后，进入算力容器，点击“工具-lab”即可打开 JupyterLab 终端开始项目创作。
|
||||
|
||||
|
||||
#### 开发环境设置
|
||||
|
||||
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 等所有依赖包。因此在启动实例后，无需再进行繁琐的依赖库版本验证即可直接进入测试环节。
|
||||
**说明：**由于本次使用的是预装的专属镜像，环境中已经默认安装并配置好了 PyTorch、FlashAttention、einops 等所有依赖包。因此在启动实例后，无需再进行繁琐的依赖库版本验证即可直接进入测试环节。
|
||||
|
||||
### 代码准备
|
||||
|
||||
|
|
@ -246,219 +197,52 @@
|
|||
|
||||
---
|
||||
|
||||
## 五、知识预备
|
||||
## 五、知识速览
|
||||
|
||||
### 名词解释
|
||||
### 关键术语
|
||||
|
||||
| 术语 | 说明 |
|
||||
| --- | --- |
|
||||
| **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)：**解决“算得对不对”的问题。它的目标是验证当前算子的输出结果，在数学精度上是否与标准参考实现完全一致。这是所有测试的绝对前提底线。
|
||||
* **KV-Cache**：缓存历史 Token 的 Key/Value 向量，避免 Transformer 推理时重复计算。
|
||||
|
||||
* **性能测试 (Performance Testing)：**解决“跑得快不快”的问题。它的目标是在验证正确性的基础上，测量算子在特定硬件上的执行耗时、吞吐量和有效带宽利用率。本教程执行的 Benchmark 脚本，正是一个纯粹的性能测试。
|
||||
* **Paged KV-Cache**：将 KV-Cache 分页管理，减少显存碎片，提高利用率。
|
||||
|
||||
* **二者区别：**
|
||||
* **Batch Size**：一次处理的样本数，越大并行度越高，但显存占用越大。
|
||||
|
||||
* **seq\_len\_kv**：KV-Cache 中已缓存的历史 Token 数量。
|
||||
|
||||
* **headdim**：注意力头维度，常见为 64/128/256。
|
||||
|
||||
|
||||
| 维度 | 性能测试 | 正确性测试 |
|
||||
| --- | --- | --- |
|
||||
| 测试目标 | 测量速度、带宽 | 验证输出结果 |
|
||||
| 关注输出 | 否 | 是 |
|
||||
| 关注效率 | 是 | 否 |
|
||||
| 是否需要 Baseline | 是 | 不一定 |
|
||||
| 是否受实现不同而影响 | 大 | 是（精度不同可能影响） |
|
||||
### 核心知识
|
||||
|
||||
**二者联系：**
|
||||
|
||||
在实际开发中：
|
||||
|
||||
正确性测试（先）
|
||||
|
||||
↓
|
||||
|
||||
建立 Baseline
|
||||
|
||||
↓
|
||||
|
||||
性能分析
|
||||
|
||||
↓
|
||||
|
||||
优化实现
|
||||
|
||||
↓
|
||||
|
||||
性能测试对比 Baseline
|
||||
|
||||
↓
|
||||
|
||||
回归正确性验证（保证没变坏）
|
||||
|
||||
在算子优化迭代中，每一次修改底层代码，都必须**先通过正确性测试**确立功能基准，**再运行性能测试**对比性能基准 Baseline，确保速度的提升绝不是以牺牲结果正确性为代价。
|
||||
|
||||
#### 什么是 Benchmark（基准测试）
|
||||
|
||||
Benchmark 是一种标准化的性能测量方法，通过在固定条件下反复运行同一任务，获取可重复、可对比的性能指标。在 GPU 算子优化场景中，benchmark 的作用是：
|
||||
|
||||
* **建立性能基线**：在优化前记录原始性能数据，作为后续对比的参照
|
||||
* **正确性测试 ：**验证算子输出结果的数学精度是否与标准实现一致，这是**绝对底线**。
|
||||
|
||||
* **量化优化效果**：优化后运行同样的 benchmark，直接对比时间/带宽变化
|
||||
|
||||
* **发现性能瓶颈**：通过不同参数组合的测试结果，定位性能拐点
|
||||
* **性能测试 ：**在正确的前提下，测算速度与吞吐。通过建立 **Baseline（基线），**才能量化后续每次代码修改带来的真实收益（加速比）。
|
||||
|
||||
|
||||
> 参考：[MLPerf Benchmark 介绍](https://mlcommons.org/benchmarks/)
|
||||
|
||||
#### 什么是 batch\_size、seq\_len、headdim
|
||||
|
||||
这三个参数共同决定了注意力计算的**工作量**和**显存占用**：
|
||||
|
||||
* **batch\_size（批大小）**：一次推理同时处理的样本数量。batch\_size 越大，GPU 并行度越高，但显存占用也线性增长。在 KV-Cache 场景中，batch\_size 对应同时服务的请求数。
|
||||
* **Benchmark （基准测试）：**在固定条件下反复运行同一任务，获取可重复的性能指标，用于建立基线、量化优化效果和定位瓶颈。
|
||||
|
||||
* **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 成正比。
|
||||
* **Kernel 执行时间**：GPU 核函数运行耗时（ms），使用 GPU 端同步计时获得。
|
||||
|
||||
* **有效带宽：**`数据传输量 (GB) ÷ Kernel 时间 (s)`，越接近理论峰值说明显存带宽利用越充分。
|
||||
|
||||
|
||||
三者与显存占用的关系：
|
||||
### 其他要点
|
||||
|
||||
```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 调度开销，时间偏大
|
||||
* **Warmup**：预热若干次（不记录），使 GPU 进入稳定状态。
|
||||
|
||||
* **GPU 端计时**：使用 CUDA Event 或 profiler，精度更高，直接测量 GPU 上的实际执行时间
|
||||
* **Repeat**：正式运行多次，取平均值或中位数以消除波动。
|
||||
|
||||
|
||||
本教程使用 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 已接近最优，优化空间有限
|
||||
* **同步**：调用 `torch.cuda.synchronize()` 确保精确计时。
|
||||
|
||||
* 有效带宽 **远低于理论峰值** → 存在优化空间（如内存访问不合并、bank conflict 等）
|
||||
* **数据类型**：本教程使用 `bfloat16`，在精度和性能取得平衡。
|
||||
|
||||
|
||||
> 参考：[CUDA C++ Programming Guide - Performance Guidelines](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#performance-guidelines)
|
||||
|
||||
#### 为什么要 Warmup / Repeat
|
||||
|
||||
GPU 程序的首次运行往往比后续运行慢，原因包括：
|
||||
|
||||
* **JIT 编译**：部分框架会延迟编译 kernel 代码
|
||||
* **显存占用估算**：`KV-Cache ≈ batch × seq_len_kv × num_heads_k × headdim × 2(K+V) × 字节数`。
|
||||
|
||||
* **缓存冷启动**：GPU L2 Cache、TLB 等初始状态为空
|
||||
* **OOM 应对**：减小 batch/seq\_len\_kv、使用更小 dtype 或释放中间变量。
|
||||
|
||||
* **频率爬升**：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)
|
||||
* **Tensor Core**：现代 GPU（含沐曦 C500）的矩阵乘法专用单元，要求维度对齐为 8 或 16 的倍数。
|
||||
|
||||
|
||||
---
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
|
@ -0,0 +1,146 @@
|
|||
from flash_attn.flash_attn_interface import flash_attn_with_kvcache
|
||||
import torch
|
||||
import math
|
||||
from einops import rearrange
|
||||
from datetime import datetime
|
||||
import csv
|
||||
|
||||
|
||||
def run_with_profiler(fn, warmup=10, reps=100, print_result=False, target_kernels=None):
|
||||
"""Run function with torch.profiler and return sum of specific kernel times in ms"""
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
with torch.profiler.profile(
|
||||
activities=[torch.profiler.ProfilerActivity.CUDA],
|
||||
record_shapes=False,
|
||||
profile_memory=False,
|
||||
with_stack=False,
|
||||
) as prof:
|
||||
for _ in range(reps):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if print_result:
|
||||
print(prof.key_averages().table(sort_by="device_time", row_limit=20))
|
||||
|
||||
if target_kernels is None:
|
||||
target_kernels = []
|
||||
|
||||
kernel_times_us = 0.0
|
||||
for evt in prof.key_averages():
|
||||
if any(k in evt.key for k in target_kernels):
|
||||
kernel_times_us += evt.device_time
|
||||
|
||||
ms = kernel_times_us / 1e3
|
||||
return ms
|
||||
|
||||
|
||||
def calc_bandwidth(batch_size, seqlen_q, seqlen_k, num_heads, num_heads_k, headdim, dtype, ms):
|
||||
"""Calculate bandwidth in GB/s"""
|
||||
bytes_per_elem = 2 if dtype == torch.bfloat16 else 4
|
||||
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
|
||||
total_bytes = q_bytes + kv_bytes
|
||||
bw_gb_s = (total_bytes / 1e9) / (ms / 1e3)
|
||||
return bw_gb_s
|
||||
|
||||
|
||||
def benchmark_kvcache(batch_size, seqlen_k, seqlen_q, num_heads, num_heads_k, headdim, page_block_size, device, dtype=torch.bfloat16, causal=False):
|
||||
num_blocks = math.ceil(seqlen_k / page_block_size) * batch_size * 3
|
||||
num_blocks = max(1024, num_blocks)
|
||||
paged_kv_block_size = page_block_size
|
||||
|
||||
nheads = num_heads
|
||||
nheads_k = num_heads_k
|
||||
d = headdim
|
||||
|
||||
torch.random.manual_seed(0)
|
||||
window_size = (-1, -1)
|
||||
|
||||
q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype)
|
||||
|
||||
k_cache_paged = torch.randn(
|
||||
num_blocks, paged_kv_block_size, nheads_k, d, device=device, dtype=dtype
|
||||
)
|
||||
v_cache_paged = torch.randn(
|
||||
num_blocks, paged_kv_block_size, nheads_k, d, device=device, dtype=dtype
|
||||
)
|
||||
block_table = rearrange(
|
||||
torch.randperm(num_blocks, dtype=torch.int32, device=device),
|
||||
"(b nblocks) -> b nblocks",
|
||||
b=batch_size,
|
||||
)
|
||||
|
||||
cache_seqlens = torch.full((batch_size,), seqlen_k, dtype=torch.int32, device=device)
|
||||
|
||||
def run_fn():
|
||||
flash_attn_with_kvcache(
|
||||
q, k_cache_paged, v_cache_paged, None, None,
|
||||
cache_seqlens=cache_seqlens,
|
||||
cache_batch_idx=None,
|
||||
block_table=block_table,
|
||||
causal=causal,
|
||||
window_size=window_size,
|
||||
rotary_interleaved=False,
|
||||
alibi_slopes=None,
|
||||
num_splits=1,
|
||||
)
|
||||
|
||||
return run_fn
|
||||
|
||||
|
||||
def main():
|
||||
headdims = [256]
|
||||
page_block_size = 16
|
||||
batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128]
|
||||
seq_lens_kv = [512, 1024, 2048, 4096, 8192, 16384]
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.bfloat16
|
||||
causal = False
|
||||
warmup = 10
|
||||
repeat = 100
|
||||
|
||||
num_heads = 8
|
||||
num_heads_k = 8
|
||||
seqlen_q = 1
|
||||
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
csv_path = f"benchmark_kvcache_{timestamp}.csv"
|
||||
|
||||
with open(csv_path, "w", newline="") as f:
|
||||
writer = csv.writer(f)
|
||||
writer.writerow(["batch_size", "seq_len_kv", "heads", "headdim", "time_ms", "bandwidth_GB_s"])
|
||||
print(f"{'batch_size':>10} {'seq_len_kv':>12} {'heads':>6} {'headdim':>8} {'time_ms':>10} {'bandwidth_GB_s':>15}")
|
||||
print("-" * 75)
|
||||
for headdim in headdims:
|
||||
for seqlen_k in seq_lens_kv:
|
||||
for batch_size in batch_sizes:
|
||||
try:
|
||||
run_fn = benchmark_kvcache(
|
||||
batch_size=batch_size,
|
||||
seqlen_k=seqlen_k,
|
||||
seqlen_q=seqlen_q,
|
||||
num_heads=num_heads,
|
||||
num_heads_k=num_heads_k,
|
||||
headdim=headdim,
|
||||
page_block_size=page_block_size,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
causal=causal,
|
||||
)
|
||||
ms = run_with_profiler(run_fn, warmup=warmup, reps=repeat, target_kernels=["flash"])
|
||||
bw = calc_bandwidth(batch_size, seqlen_q, seqlen_k, num_heads, num_heads_k, headdim, dtype, ms)
|
||||
writer.writerow([batch_size, seqlen_k, num_heads, headdim, f"{ms:.4f}", f"{bw:.2f}"])
|
||||
print(f"{batch_size:>10} {seqlen_k:>12} {num_heads:>6} {headdim:>8} {ms:>10.4f} {bw:>15.2f}")
|
||||
except Exception as e:
|
||||
writer.writerow([batch_size, seqlen_k, num_heads, headdim, "OOM", "OOM"])
|
||||
print(f"{batch_size:>10} {seqlen_k:>12} {num_heads:>6} {headdim:>8} {'OOM':>10} {'OOM':>15} # {e}")
|
||||
|
||||
print(f"\nResults saved to {csv_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
batch_size,seq_len_kv,heads,headdim,time_ms,bandwidth_GB_s
|
||||
1,512,8,128,0.0322,65.27
|
||||
2,512,8,128,0.0324,129.61
|
||||
4,512,8,128,0.0332,253.01
|
||||
8,512,8,128,0.0355,472.66
|
||||
16,512,8,128,0.0546,614.98
|
||||
32,512,8,128,0.0817,822.60
|
||||
64,512,8,128,0.1297,1035.62
|
||||
128,512,8,128,0.2453,1095.50
|
||||
1,1024,8,128,0.0578,72.55
|
||||
2,1024,8,128,0.0586,143.14
|
||||
4,1024,8,128,0.0597,281.24
|
||||
8,1024,8,128,0.0625,536.80
|
||||
16,1024,8,128,0.0982,683.59
|
||||
32,1024,8,128,0.1493,899.64
|
||||
64,1024,8,128,0.2403,1117.62
|
||||
128,1024,8,128,0.4594,1169.27
|
||||
1,2048,8,128,0.1101,76.24
|
||||
2,2048,8,128,0.1107,151.64
|
||||
4,2048,8,128,0.1119,299.88
|
||||
8,2048,8,128,0.1159,578.98
|
||||
16,2048,8,128,0.1849,726.23
|
||||
32,2048,8,128,0.2843,944.47
|
||||
64,2048,8,128,0.4607,1165.56
|
||||
128,2048,8,128,0.8868,1211.07
|
||||
1,4096,8,128,0.2139,78.46
|
||||
2,4096,8,128,0.2151,156.01
|
||||
4,4096,8,128,0.2163,310.36
|
||||
8,4096,8,128,0.2227,602.81
|
||||
16,4096,8,128,0.3574,751.13
|
||||
32,4096,8,128,0.5540,969.23
|
||||
64,4096,8,128,0.9016,1191.07
|
||||
128,4096,8,128,1.7414,1233.34
|
||||
1,8192,8,128,0.4215,79.61
|
||||
2,8192,8,128,0.4226,158.81
|
||||
4,8192,8,128,0.4242,316.39
|
||||
8,8192,8,128,0.4362,615.46
|
||||
16,8192,8,128,0.7035,763.14
|
||||
32,8192,8,128,1.0934,982.11
|
||||
64,8192,8,128,1.7814,1205.57
|
||||
128,8192,8,128,3.4505,1244.82
|
||||
1,16384,8,128,0.8356,80.32
|
||||
2,16384,8,128,0.8377,160.23
|
||||
4,16384,8,128,0.8407,319.30
|
||||
8,16384,8,128,0.8625,622.51
|
||||
16,16384,8,128,1.3934,770.60
|
||||
32,16384,8,128,2.1695,989.88
|
||||
64,16384,8,128,3.5397,1213.41
|
||||
128,16384,8,128,6.8668,1250.98
|
||||
|
|
|
@ -83,10 +83,10 @@ PyTorch-Agent / 2.8.0 / Python 3.12 / maca 3.7.2.1
|
|||
PyTorch-Agent / 2.8.0 / Python 3.12 / maca 3.7.2.1
|
||||
```
|
||||
|
||||
4. 推荐机器规格：
|
||||
|
||||
4. 日常开发与调试：选用沐曦 GPU，显存 16–32 GB 即可满足需求。
|
||||
|
||||
**注意：运行性能跑分（Benchmark）时，请务必使用整张单卡（64 GB）。**
|
||||
|
||||
沐曦GPU VRAM 16-32 GB即可。
|
||||
|
||||
### 4.2 购买模型资源包
|
||||
|
||||
|
|
@ -564,33 +564,8 @@ opencode
|
|||
请通过真实 shell 命令检查当前目录、Python 版本、mxcc 路径、MACA_PATH 和 baseline 脚本入口。然后检查当前目录结构，并找出 fusedmoe_v2.1 的 build、test、benchmark 入口。先不要修改代码。
|
||||
```
|
||||
|
||||
## 七、Agent 使用说明
|
||||
|
||||
### 7.1 推荐工作方式
|
||||
|
||||
使用 Agent 时，不要一开始就让它“直接优化到最快”。推荐节奏是：
|
||||
|
||||
1. 先让它检查目录。
|
||||
|
||||
2. 再让它跑通原始 baseline。
|
||||
|
||||
3. 然后让它只做一处小修改。
|
||||
|
||||
4. 每次修改后必须 build、test、benchmark。
|
||||
|
||||
5. 每轮都记录结果。
|
||||
|
||||
|
||||
### 7.2 使用 Agent 的注意事项
|
||||
|
||||
1. 不要相信没有命令输出支持的环境判断。
|
||||
|
||||
2. 不要让 Agent 同时修改很多文件。
|
||||
|
||||
3. 不要让 Agent 修改测试标准。
|
||||
|
||||
|
||||
## 八、常见问题
|
||||
## 七、常见问题
|
||||
|
||||
### 问题 1：OpenCode 报 Unauthorized
|
||||
|
||||
Loading…
Reference in New Issue