forked from metax-maca/op_optimization
Delete 基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/FlashInfer 迁移 Baseline 实战.md
This commit is contained in:
parent
409d474254
commit
daf335c155
|
|
@ -1,809 +0,0 @@
|
|||
# FlashInfer 迁移 Baseline 实战
|
||||
|
||||
## 一、教程定位
|
||||
|
||||
本教程是参赛训练课程的 **FlashInfer Baseline 入门**模块，主要帮助用户快速跑通 FlashInfer 的最小可运行流程。完成本教程后，用户应能够完成源码编译、API 调用、正确性测试和 Benchmark 测试，并记录一份 baseline 性能结果，为后续算子优化提供对比基准。
|
||||
|
||||
## 二、学习目标
|
||||
|
||||
1. 理解 FlashInfer Attention Kernel 的基本作用与适用场景；
|
||||
2. 完成 FlashInfer 环境、工具链的准备与源码编译；
|
||||
3. 跑通 BatchDecode、BatchPrefill、MLA 等典型算子的 API 调用示例；
|
||||
4. 完成各算子在不同参数配置下的 Benchmark 测试；
|
||||
5. 输出各算子的 Baseline 性能结果记录表，为后续算子优化提供对比基准。
|
||||
|
||||
|
||||
## 三、适用对象
|
||||
|
||||
适合人群
|
||||
|
||||
* 参赛选手：需要完成 Baseline 入门模块，为后续算子优化做准备
|
||||
* 软件开发者和Vibe Coding开发：希望从事AI相关行业开发，以及用智能体方式来做开发工作
|
||||
* LLM 推理开发者：希望了解 FlashInfer Attention Kernel 的性能表现
|
||||
* 算子优化工程师：希望基于MXMACA软件栈在沐曦国产 GPU 上做算子迁移和优化
|
||||
|
||||
前置基础
|
||||
|
||||
* Python 基础：能够运行和修改 Python 脚本
|
||||
* PyTorch 基础：了解MXMACA软化栈的使用
|
||||
* Linux 命令行：能够使用终端执行命令
|
||||
* 了解 Attention 机制：理解 Q/K/V、KV Cache 等基本概念
|
||||
|
||||
|
||||
## 四、前置准备
|
||||
|
||||
开始实战前，请确认你已经完成以下准备：
|
||||
|
||||
### 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/)
|
||||
|
||||
* 进入费用中心 - 算力券 ，点击右上角 “兑换”。
|
||||
|
||||
|
||||
|
||||
|
||||
* 步骤3：租用算力
|
||||
|
||||
|
||||
进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图：
|
||||
|
||||
|
||||
|
||||
* 步骤4：创建实例
|
||||
|
||||
|
||||
基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
* 步骤5：选择工具-lab进入实例环境
|
||||
|
||||
|
||||
|
||||
|
||||
* 步骤6：在JupyterLab Terminal中检查运行环境的配置。
|
||||
|
||||
|
||||
|
||||
* 确认沐曦 GPU 可见--可以使用`mx-smi`命令查看
|
||||
|
||||
|
||||
|
||||
|
||||
### 环境依赖准备
|
||||
|
||||
```bash
|
||||
pip install flashinfer torch pandas numpy
|
||||
```
|
||||
|
||||
### opencode Agent安装
|
||||
|
||||
```shell
|
||||
curl -fsSL https://opencode.ai/install | bash
|
||||
```
|
||||
|
||||
说明：
|
||||
|
||||
自动下载最新版本 自动配置环境 适用于 macOS / Linux / WSL 安装完成后，还会提示你怎么用，OpenCode 包含免费模式，使用方式：
|
||||
|
||||
```shell
|
||||
cd # 进入项目目录
|
||||
opencode
|
||||
```
|
||||
|
||||
### 代码准备
|
||||
|
||||
* 已进入指定项目目录；
|
||||
|
||||
* 准备测试脚本和 Benchmark 脚本。
|
||||
|
||||
|
||||
| 命令 | 说明 |
|
||||
| --- | --- |
|
||||
| `python bench_batch_decode.py` | 运行 Batch Decode 基准测试 |
|
||||
| `python bench_batch_prefill_paged.py` | 运行 Batch Prefill (Paged KV Cache) 基准测试 |
|
||||
| `python bench_batch_prefill_ragged.py` | 运行 Batch Prefill (Ragged KV Cache) 基准测试 |
|
||||
| `python bench_batch_mla.py` | 运行 MLA (Multi-head Latent Attention) 基准测试 |
|
||||
|
||||
## 五、知识预备
|
||||
|
||||
### Flashinfer基础知识
|
||||
|
||||
FlashInfer 是一个用于推理的库和内核生成器，能够在多种 GPU 架构上实现最先进的性能。它为注意力、GEMM和MoE操作提供统一API，支持包括FlashAttention-2/3、cuDNN、CUTLASS和TensorRT-LLM在内的多种后端实现。
|
||||
|
||||
https://github.com/flashinfer-ai/flashinfer
|
||||
|
||||
#### Attention Kernels：
|
||||
|
||||
* Paged and Ragged KV-Cache: Efficient memory management for dynamic batch serving
|
||||
|
||||
* Decode, Prefill, and Append: Optimized kernels for all attention phases
|
||||
|
||||
* MLA Attention: Native support for DeepSeek's Multi-Latent Attention
|
||||
|
||||
* Cascade Attention: Memory-efficient hierarchical KV-Cache for shared prefixes
|
||||
|
||||
* Sparse Attention: Block-sparse and variable block-sparse patterns
|
||||
|
||||
* POD-Attention: Fused prefill+decode for mixed batching
|
||||
|
||||
|
||||
#### LLM推理阶段重要概念：
|
||||
|
||||
* **Prefill阶段**：prefill 阶段是指处理输入 prompt 的阶段
|
||||
|
||||
* 输入：用户一次性给出的完整 prompt，长度为 seq\_len
|
||||
|
||||
* 计算：对 prompt 中的每个 token 并行计算注意力，生成第一个输出 token 及 KV cache
|
||||
|
||||
* 特点：这是**计算密集型（compute-bound）**阶段，因为需要做完整的 seq\_len x seq\_len 注意力矩阵乘法
|
||||
|
||||
* **decode阶段**：
|
||||
|
||||
* 每次只生成 1 个 token，利用 prefill 阶段填充好的 KV cache 做自回归生成
|
||||
|
||||
* **显存带宽密集型（memory-bound**），瓶颈在从显存读取 KV cache 而非计算
|
||||
|
||||
|
||||
prefill = 并行处理用户输入，decode = 逐个生成回答 token
|
||||
|
||||
#### KV Cache原理
|
||||
|
||||
标准的多头自注意力（Multi-Head Self-Attention）：
|
||||
|
||||
$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d\_k}}\right)V$
|
||||
|
||||
其中 $Q$、$K$、$V$ 分别是查询（Query）、键（Key）和值（Value）矩阵，$d\_k$是每个注意力头的维度。在训练阶段，由于有 causal mask（因果掩码），整个序列的 $Q$、$K$、$V$可以并行计算。但在推理的自回归生成阶段，当我们生成第 $t$个 token 时：
|
||||
|
||||
位置 1 到 $t-1$的 $K$、$V$向量在生成第 $t$个 token 时就已经计算过了。如果不做任何缓存，每一步都要重新从头计算所有历史 token 的 $K$和 $V$，这意味着生成第 $t$ 个 token 的复杂度是 $O(t)$，整个序列生成的总复杂度是 $O(n^2)$，在长序列下极其低效。KV Cache 的思路非常直观：把已经计算过的 Key 和 Value 向量缓存起来，下一步直接拿来用。
|
||||
|
||||
#### Ragged KV-Cache（非分页缓存）
|
||||
|
||||
传统的实现中，显存分配要求是**物理连续**的。
|
||||
|
||||
* **分配逻辑：** 由于不知道用户最终会生成多少个 Token，系统只能“往大了猜”，按照模型允许的最大长度（例如 2048）为每个新请求一次性预留一长条连续的显存。
|
||||
|
||||
* **问题：**这样的分配方法容易造成**内部显存碎片化，显存利用率不高**，很容易引起“gpu显存利用不足”的问题，进而影响模型推理时的吞吐量。
|
||||
|
||||
|
||||
#### Paged KV-Cache（分页缓存）
|
||||
|
||||
借用操作系统虚拟内存的思想，将物理显存切分成大小固定的“页/块”（Blocks）。
|
||||
|
||||
* **分配逻辑：** 新请求到来时，不再预留连续大空间，而是先只分配一个 Block（比如 16 个 Token 的大小）。随着模型的逐字解码（Decode），当这个 Block 填满时，再向系统动态申请下一个 Block
|
||||
|
||||
* **按需分配：** 生成几个 Token 就用几个位置，几乎消灭了内部碎片。
|
||||
|
||||
* **物理离散，逻辑连续：** 不同的 Block 在物理显存上完全可以是分散的（通过 Block Table 记录映射），彻底消灭了外部碎片，新请求随时可以插空进入。
|
||||
|
||||
|
||||
|
||||
flashInfer 将分页 KV Cache 视作一个**块稀疏矩阵**，并巧妙地使用了 **CSR (Compressed Sparse Row)** 格式来建立索引。图 1 中的三个关键数组就是 CSR 格式的体现：
|
||||
|
||||
* `**kv_page_indices**` **(页索引池):** 把所有请求当前占用的物理页编号，按顺序“平铺”拼接在一起。
|
||||
|
||||
* 图中蓝、橙、绿三个请求分别使用了 `[0, 5, 8]`、`[1, 6, 7]`、`[3, 4]`。
|
||||
|
||||
* 拼接后得到：`[0, 5, 8, 1, 6, 7, 3, 4]`。
|
||||
|
||||
* `**kv_indptr**` **(索引指针):** 用于标记每个请求在 `kv_page_indices` 中的**起始和结束位置**。
|
||||
|
||||
* 数组长度固定为 `num_requests + 1`。
|
||||
|
||||
* 图中数值为 `[0, 3, 6, 8]`。这意味着：
|
||||
|
||||
* 请求 0 (蓝) 的页索引在 `kv_page_indices` 的 `0` 到 `3` 之间（即 `[0, 5, 8]`，共 3 页）。
|
||||
|
||||
* 请求 1 (橙) 的页索引在 `3` 到 `6` 之间（即 `[1, 6, 7]`，共 3 页）。
|
||||
|
||||
* 请求 2 (绿) 的页索引在 `6` 到 `8` 之间（即 `[3, 4]`，共 2 页）。
|
||||
|
||||
* `**kv_last_page_lens**` **(尾页有效长度):** 由于一个请求的 Token 总数很少能刚好被 `page_size`（每页容量，图中为 8）整除，最后一个页通常是不满的。这个数组记录了每个请求**最后一页实际存储的 Token 数量**。
|
||||
|
||||
* 图中分别为 `[6, 4, 7]`。
|
||||
|
||||
|
||||
图 1 左下角展示了在解码（Decode/Append）阶段，新生成的 Token 是如何追加到 KV Cache 中的。这里的细节非常值得注意：
|
||||
|
||||
* `**qo_indptr = [0, 4, 6, 9]**`: 这表示当前批次中，各个请求**新追加**的 Token 数量（Query/Output）。
|
||||
|
||||
* 请求 0 (蓝) 追加了 `4 - 0 = 4` 个 Token。
|
||||
|
||||
* 请求 1 (橙) 追加了 `6 - 4 = 2` 个 Token。
|
||||
|
||||
* 请求 2 (绿) 追加了 `9 - 6 = 3` 个 Token。
|
||||
|
||||
|
||||
### Flashinfer API
|
||||
|
||||
#### Ragged Tensor（不规则张量）
|
||||
|
||||
|
||||
|
||||
假设有 3 个请求，长度分别是 5, 3, 4。所以batchsize=3，用户请求的seq\_len分别为 5,3,4，因为在深度学习中由于底层的矩阵运算要求张量（Tensor）必须是**规整的矩形**（比如 `[batch_size, seq_len, hidden_dim]`），当一个 Batch 中包含**不同长度的句子时**，我们通常会按最长的那句话进行 **Padding（补零）**。而图中的核心思想是：**“拒绝 Padding，把所有 Token 拍扁拼接到一起。”**
|
||||
|
||||
这里不使用 `[3(batch_size), 5(seq_len), num_heads, head_dim]` 这样的多维规整张量，而是把所有请求的 Token 首尾相连，打包成一个长度为 12（即 5+3+4）的连续一维数组。这就是图中的 `data: (12, num_heads, head_dim)` = `data：(seq_len(5+3+4)，num_heads，head_dim)`的由来。
|
||||
|
||||
既然数据被拼接到了一起，系统怎么知道哪个 Token 属于哪个请求呢？这就是 `indptr`（Index Pointer，索引指针）发挥作用的地方。
|
||||
|
||||
* **颜色块代表不同的请求（Request）：**
|
||||
|
||||
* 🟦 蓝色：Request 0，长度为 5
|
||||
|
||||
* 🟧 橙色：Request 1，长度为 3
|
||||
|
||||
* 🟩 绿色：Request 2，长度为 4
|
||||
|
||||
* `**indptr = [0, 5, 8, 12]**`**：**
|
||||
|
||||
* 这是一个一维数组，用来记录每个请求在 `data` 数组里的**起始和结束位置**。它的长度永远是 `num_requests + 1`，且第一个元素必定是 0。
|
||||
|
||||
* Request 0 的数据在 `data[0:5]` （对应图中 `indptr[0]` 到 `indptr[1]`）
|
||||
|
||||
* Request 1 的数据在 `data[5:8]` （对应图中 `indptr[1]` 到 `indptr[2]`）
|
||||
|
||||
* Request 2 的数据在 `data[8:12]` （对应图中 `indptr[2]` 到 `indptr[3]`）
|
||||
|
||||
|
||||
**序列长度计算：** 第 $i$个请求的长度（Sequence length）可以直接通过 $indptr\[i+1\] - indptr\[i\]$算出来。
|
||||
|
||||
**总 Token 数：** `indptr` 的最后一个元素（即 `indptr[-1]`，在这里是 12），就是这个 Batch 中所有 Token 的总和。
|
||||
|
||||
**数据切片读取：** 当你需要单独提取第 $i$个请求的 $Q/K/V$矩阵时，只需要执行切片操作 `data[indptr[i]:indptr[i+1]]` 即可精准拿取，没有任何多余的 Padding 元素。
|
||||
|
||||
#### flashinfer.BatchPrefillWithRaggedKVCacheWrapper() 类
|
||||
|
||||
参考链接：https://docs.flashinfer.ai/api/attention.html#flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper
|
||||
|
||||
1. 构造函数：
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, kv_indptr_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None) → None
|
||||
```
|
||||
|
||||
2. 常用初始化参数理解：
|
||||
|
||||
1. `_**float_workspace_buffer**_``(``_torch.Tensor_``)`：用户预留的浮点工作空间缓冲区，用于在 split-k 算法中存储中间的注意力计算结果。建议大小为 128MB，其设备类型需与输入张量所在的设备保持一致。
|
||||
|
||||
2. `_**kv_layout**_`_**：**_输入 K/V 张量的显存布局格式，可以是 `NHD` 或 `HND`。默认_**'NHD'**_
|
||||
|
||||
* NHD：`(seq_len, num_heads, head_dim)`
|
||||
|
||||
* HND：`(num_heads, seq_len, head_dim)`
|
||||
|
||||
3. `_**backend**_`_**：**_底层实现引擎。可选值包括 `auto`、`fa2`、`fa3`、`cudnn`、`cutlass` 或 `cute-dsl`，默认值为 `auto`。系统会根据显卡架构自动选择最优后端。其中 `cute-dsl` 是专为最新一代 Blackwell 架构（如 SM100+ 系列）准备的算子。
|
||||
|
||||
4. `_**jit_args**_` _**&**_ `_**jit_kwargs**_`_**：**_用于即时编译（JIT, Just-In-Time）的参数列表和字典参数。如果提供，框架将会在运行时动态编译底层算子（手动实现算子）；否则，直接使用预编译好的默认算子。
|
||||
|
||||
3. **.Plan()** 方法
|
||||
|
||||
|
||||
**Plan()** 方法的作用就是“运筹帷幄”的预处理（AOT, Ahead-of-Time Setup）阶段：**任务规划：** 接收当前 Batch 中所有请求的形状、长度和硬件规格，计算出最优的 GPU 算力调度方案。**显存分配：** 在底层预先创建并缓存计算所需的辅助数据结构和临时工作空间（Workspace）。**解耦计算：** 将“准备工作”与真正的“执行工作（`run` 方法）”解耦。调用 `plan` 搭建好“脚手架”后，后续调用 `run` 时 GPU 就可以直接根据图纸极速开工，从而把 **CPU 调度开销降到最低**。`plan()` 方法必须在任何 `run()`之前
|
||||
|
||||
```Python
|
||||
plan(qo_indptr: Tensor, kv_indptr: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, window_left: int = -1, logits_soft_cap: float | None = None, sm_scale: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False, seq_lens: Tensor | None = None, seq_lens_q: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, v_indptr: Tensor | None = None, o_indptr: Tensor | None = None) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `_**qo_indptr**_`_**：**_Query/Output 张量的索引指针数组（indptr），形状为 `[batch_size + 1]`。用于在 Ragged 连续内存中定位每个请求的 Query 边界。
|
||||
|
||||
* `_**kv_indptr**_`_**：**_Key/Value 张量的索引指针数组，形状为 `[batch_size + 1]`。
|
||||
|
||||
* `_**num_qo_heads**_`_**：**_ Query 和 Output 的注意力头（Attention Heads）数量。
|
||||
|
||||
* `_**num_kv_heads**_`_**：**_Key 和 Value 的注意力头数量。
|
||||
|
||||
* `_**head_dim_qk**_`_**：**_Query 和 Key 张量中每个注意力头的维度大小。
|
||||
|
||||
* `_**head_dim_vo**_`_**：**_Value 和 Output 张量中每个头的维度大小。如果不提供，默认与 `head_dim_qk` 相同。
|
||||
|
||||
* `_**causal**_`_**：**_ 是否对注意力矩阵应用因果掩码（Causal Mask，即屏蔽未来信息）。如果在 `plan()` 中已经提供了自定义的 `mask` 参数，此选项将被忽略。
|
||||
|
||||
* `_**q_data_type**_`_**：**_Query 张量的数据类型，默认为 `torch.float16`。
|
||||
|
||||
* `_**kv_data_type**_`_**：**_Key/Value 张量的数据类型。如果不提供，默认与 `q_data_type` 一致。
|
||||
|
||||
|
||||
注意：`plan()` 方法包含复杂的 Python 层逻辑和动态显存分配，因此**不能**在 CUDA Graph 捕获环境或 `torch.compile` 环境中被追踪或调用。
|
||||
|
||||
1. **.run()** 方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q: Tensor, k: Tensor, v: Tensor, *args, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, enable_pdl: bool | None = None, kv_cache_sf: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `_**q**_`_**：**_Query（查询）张量。
|
||||
|
||||
* **形状:** `[qo_indptr[-1], num_qo_heads, head_dim_qk]`
|
||||
|
||||
* **解释:** 这里的 `qo_indptr[-1]` 正是我们之前提到的 Ragged Tensor 中所有 Token 数量的总和。它表示把 Batch 里所有的 Query 拍扁到了一个一维的连续维度上。
|
||||
|
||||
* `_**k**_`_**：**_Key（键）张量。
|
||||
|
||||
* **形状:** `[kv_indptr[-1], num_kv_heads, head_dim_qk]`
|
||||
|
||||
* `_**v**_`_**：**_Value（值）张量。
|
||||
|
||||
* **形状:** `[kv_indptr[-1], num_kv_heads, head_dim_vo]`
|
||||
|
||||
|
||||
#### flashinfer.BatchPrefillWithPagedKVCacheWrapper() 类
|
||||
|
||||
参考链接：https://docs.flashinfer.cn/api/attention.html#flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper
|
||||
|
||||
1. 构造函数
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None) → None
|
||||
```
|
||||
|
||||
参数含义同**BatchPrefillWithRaggedKVCacheWrapper()** 的构造函数参数相同
|
||||
|
||||
2. .Plan() 方法
|
||||
|
||||
|
||||
```Python
|
||||
plan(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, page_size: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, sm_scale: float | None = None, window_left: int = -1, logits_soft_cap: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, seq_lens: Tensor | None = None, seq_lens_q: Tensor | None = None, block_tables: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**o_indptr**``(``_torch.Tensor_``)` – 查询/输出张量的 indptr，形状：`[batch_size + 1]`。
|
||||
|
||||
* `**paged_kv_indptr**``(``_torch.Tensor_``)` – 分页 kv-cache 的 indptr，形状：`[batch_size + 1]`。
|
||||
|
||||
* `**paged_kv_indices**``(``_torch.Tensor_``)` – 分页 kv-cache 的页索引，形状：`[paged_kv_indptr[-1]]`。
|
||||
|
||||
* `**paged_kv_last_page_len**``(``_torch.Tensor_``)` – 分页 kv-cache 中每个请求的最后一页中的条目数，形状：`[batch_size]`。
|
||||
|
||||
* `**num_qo_heads**``(``_int_``)` – 查询/输出头的数量。
|
||||
|
||||
* `**num_kv_heads**``(``_int_``)` – 键/值头的数量。
|
||||
|
||||
* `**head_dim_qk**``(``_int_``)` – 查询/键头的维度。
|
||||
|
||||
* `**page_size**``(``_int_``)` – 分页 kv-cache 中每个页面的大小。
|
||||
|
||||
|
||||
1. run()方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], *args, k_scale: float | None = None, v_scale: float | None = None, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, enable_pdl: bool | None = None, window_left: int | None = None) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* **q** (_torch.Tensor_) – 查询张量，形状：`[qo_indptr[-1], num_qo_heads, head_dim]`
|
||||
|
||||
* **paged\_kv\_cache** (_Union\[torch.Tensor, Tuple\[torch.Tensor, torch.Tensor\]\]_) –存储的分页 KV 缓存，作为张量元组或单个张量
|
||||
|
||||
* 一个元组 `(k_cache, v_cache)`，包含 4D 张量，每个张量的形状为：`[max_num_pages, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。
|
||||
|
||||
* 一个 5D 张量，形状为：`[max_num_pages, 2, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, 2, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。其中 `paged_kv_cache[:, 0]` 是 key 缓存，`paged_kv_cache[:, 1]` 是 value 缓存
|
||||
|
||||
|
||||
#### flashinfer.mla.BatchMLAPagedAttentionWrapper()类
|
||||
|
||||
多头潜在注意力 (MLA) 是一种新的注意力机制，由 [++DeepSeek v2++](https://arxiv.org/abs/2405.04434) 提出，并用于后来的 DeepSeek 模型。MLA 将键缓存和值缓存统一到一个张量中，因此无需单独存储它们。与多头注意力或分组查询注意力相比，MLA 的 KV-Cache 没有 `num_heads` 维度，因此没有像 `NHD` 和 `HND` 布局这样的区别。
|
||||
|
||||
MLA 分离 RoPE（旋转位置编码）维度和其他头部维度。我们使用 `kpe`（带有位置编码的键）和 `ckv`（压缩的键/值）来命名这两个组件。用户可以将它们存储在单个 Paged KV-Cache 中
|
||||
|
||||
```Python
|
||||
head_dim_ckv = 512
|
||||
head_dim_kpe = 64
|
||||
mla_paged_kv_cache = torch.empty(max_num_pages, page_size, head_dim_ckv + head_dim_kpe, dtype=torch.bfloat16)
|
||||
ckv = mla_paged_kv_cache[:, :, :head_dim_ckv] # Slicing here does not copy or move data
|
||||
kpe = mla_paged_kv_cache[:, :, head_dim_ckv:] # Slicing here does not copy or move data
|
||||
```
|
||||
|
||||
**低秩联合压缩 (Joint Compression):** MLA 不再为每个注意力头单独存储巨大的 Key 和 Value 矩阵。相反，它将它们投影并压缩到一个共享的潜在向量（Latent Vector）中，即您代码中的 `ckv` (`head_dim_ckv = 512`)。
|
||||
|
||||
**解耦旋转位置编码 (Decoupled RoPE):** 位置信息对于注意力机制至关重要，但它很难被压缩。MLA 的巧妙之处在于将携带 RoPE 信息的维度单独剥离出来，即代码中的 `kpe` (`head_dim_kpe = 64`)。
|
||||
|
||||
**消除** `**num_heads**` **维度:** 存储的 Cache 不再区分 NHD（序列、头数、头维度）或 HND 布局。无论模型有多少个注意力头，KV-Cache 对于每个 Token 只需要存储 `head_dim_ckv + head_dim_kpe`（例如 512 + 64 = 576 个元素）。
|
||||
|
||||
1. 构造函数
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, use_cuda_graph: bool = False, qo_indptr: Tensor | None = None, kv_indptr: Tensor | None = None, kv_indices: Tensor | None = None, kv_len_arr: Tensor | None = None, backend: str = 'auto') → None
|
||||
```
|
||||
|
||||
2. .Plan()方法
|
||||
|
||||
|
||||
```Python
|
||||
plan(qo_indptr: Tensor, kv_indptr: Tensor, kv_indices: Tensor, kv_len_arr: Tensor, num_heads: int, head_dim_ckv: int, head_dim_kpe: int, page_size: int, causal: bool, sm_scale: float, q_data_type: dtype, kv_data_type: dtype, use_profiler: bool = False) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**qo_indptr**``(``_torch.IntTensor_``)` – 查询/输出张量的 indptr，形状：`[batch_size + 1]`。对于解码注意力，每个查询的长度为 1，张量的内容应为 `[0, 1, 2, ..., batch_size]`。
|
||||
|
||||
* `**kv_indptr**``(``_torch.IntTensor_``)` – 分页 kv-cache 的 indptr，形状：`[batch_size + 1]`。
|
||||
|
||||
* `**kv_indices**``(``_torch.IntTensor_``)` – 分页 kv-cache 的页面索引，形状：`[kv_indptr[-1]]` 或更大。
|
||||
|
||||
* `**kv_len_arr**``(``_torch.IntTensor_``)` – 每个请求的查询长度，形状：`[batch_size]`。
|
||||
|
||||
* `**num_heads**``(``_int_``)` – 查询/输出张量中的头数。
|
||||
|
||||
* `**head_dim_ckv**``(``_int_``)` – 压缩 kv 的头维度。
|
||||
|
||||
* `**head_dim_kpe**``(``_int_``)` – rope k-cache 的头维度。
|
||||
|
||||
* `**page_size**``(``_int_``)` – 分页 kv-cache 的页面大小。
|
||||
|
||||
* `**causal**``(``_bool_``)` – 是否使用因果注意力。
|
||||
|
||||
* `**sm_scale**``(``_float_``)` – softmax 运算的缩放因子。
|
||||
|
||||
* `**q_data_type**``(``_torch.dtype_``)` – 查询张量的数据类型。
|
||||
|
||||
* `**kv_data_type**``(``_torch.dtype_``)` – kv-cache 张量的数据类型。
|
||||
|
||||
* `**use_profiler**``(``_bool, optional_``)` – 是否启用内核内分析器，默认值为 False。
|
||||
|
||||
|
||||
1. run()方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q_nope: Tensor, q_pe: Tensor, ckv_cache: Tensor, kpe_cache: Tensor, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, profiler_buffer: Tensor | None = None, kv_len: Tensor | None = None, page_table: Tensor | None = None, return_lse_base_on_e: bool = False) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**q_nope**``(``_torch.Tensor_``)` – 不含 rope 的查询张量，形状：`[batch_size, num_heads, head_dim_ckv]`。
|
||||
|
||||
* `**q_pe**``(``_torch.Tensor_``)` – 查询张量的 rope 部分，形状：`[batch_size, num_heads, head_dim_kpe]`。
|
||||
|
||||
* `**ckv_cache**``(``_torch.Tensor_``)` – 压缩的 kv-cache 张量（不含 rope），形状：`[num_pages, page_size, head_dim_ckv]`。 `head_dim_ckv` 在 DeepSeek v2/v3 模型中为 512。
|
||||
|
||||
* `**kpe_cache**``(``_torch.Tensor_``)` – kv-cache 张量的 rope 部分，形状：`[num_pages, page_size, head_dim_kpe]`。 `head_dim_kpe` 在 DeepSeek v2/v3 模型中为 64。
|
||||
|
||||
|
||||
#### flashinfer.BatchDecodeWithPagedKVCacheWrapper()类
|
||||
|
||||
1. 构造函数
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, use_tensor_cores: bool = False, paged_kv_indptr_buffer: Tensor | None = None, paged_kv_indices_buffer: Tensor | None = None, paged_kv_last_page_len_buffer: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None) → None
|
||||
```
|
||||
|
||||
2. .Plan()方法
|
||||
|
||||
|
||||
```Python
|
||||
plan(indptr: Tensor, indices: Tensor, last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim: int, page_size: int, pos_encoding_mode: str = 'NONE', window_left: int = -1, logits_soft_cap: float | None = None, q_data_type: str | dtype | None = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, data_type: str | dtype | None = None, sm_scale: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, non_blocking: bool = True, block_tables: Tensor | None = None, seq_lens: Tensor | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**indptr**``(``_torch.Tensor_``)` – 分页 kv 缓存的 indptr，形状：`[batch_size + 1]`，dtype：`torch.int32`
|
||||
|
||||
* `**indices**``(``_torch.Tensor_``)` – 分页 kv 缓存的页面索引，形状：`[kv_indptr[-1]]`，dtype：`torch.int32`
|
||||
|
||||
* `**last_page_len**``(``_torch.Tensor_``)` – 分页 kv 缓存中每个请求的最后一页中的条目数，形状：`[batch_size]`，dtype：`torch.int32`
|
||||
|
||||
* `**num_qo_heads**``(``_int_``)` – 查询/输出头的数量
|
||||
|
||||
* `**num_kv_heads**``(``_int_``)` – key/value 头的数量
|
||||
|
||||
* `**head_dim**``(``_int_``)` – 头部的维度
|
||||
|
||||
* `**page_size**``(``_int_``)` – 分页 kv 缓存的页面大小
|
||||
|
||||
|
||||
1. run()方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], *args, q_scale: float | None = None, k_scale: float | None = None, v_scale: float | None = None, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, enable_pdl: bool | None = None, window_left: int | None = None) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**q**``(``_torch.Tensor_``)` – 查询张量，形状：`[batch_size, num_qo_heads, head_dim]`
|
||||
|
||||
* `**paged_kv_cache**` (_Union\[torch.Tensor, Tuple\[torch.Tensor, torch.Tensor\]\]_) –存储的分页 KV 缓存，作为张量元组或单个张量
|
||||
|
||||
* 一个元组 `(k_cache, v_cache)`，包含 4D 张量，每个张量的形状为：`[max_num_pages, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。
|
||||
|
||||
* 一个 5D 张量，形状为：`[max_num_pages, 2, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, 2, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。其中 `paged_kv_cache[:, 0]` 是 key 缓存，`paged_kv_cache[:, 1]` 是 value 缓存
|
||||
|
||||
|
||||
## 六、项目实践--FlashInfer-Baseline
|
||||
|
||||
### Step 1：检查运行环境
|
||||
|
||||
**目标：** 进入terminal确认当前环境满足本模块运行要求。
|
||||
|
||||
**操作：** 检查 GPU、Python、编译工具和依赖版本。
|
||||
|
||||
|
||||
|
||||
**常见问题：**
|
||||
|
||||
| 问题 | 解决方法 |
|
||||
| --- | --- |
|
||||
| `mx-smi: command not found` | 使用 `python -c "import torch; print(torch.cuda.get_device_name(0))"` |
|
||||
| `No GPUs are available` | 检查 MXMACA 驱动是否正确安装 |
|
||||
| `ModuleNotFoundError: No module named 'flashinfer'` | 执行 `pip install flashinfer` 安装 |
|
||||
|
||||
### Step 2：进入项目目录
|
||||
|
||||
**目标：** 进入本模块所需的源码目录。
|
||||
|
||||
**操作：** 切换到 FlashInfer Baseline 项目目录。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
cd /data/flashinfer_baseline
|
||||
ls -la
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```plaintext
|
||||
bench_common.py
|
||||
bench_batch_decode.py
|
||||
bench_batch_prefill_paged.py
|
||||
bench_batch_prefill_ragged.py
|
||||
bench_batch_mla.py
|
||||
README.md
|
||||
...
|
||||
|
||||
```
|
||||
|
||||
### Step 3：验证项目脚本
|
||||
|
||||
**目标：** 确认所有基准测试脚本可正常执行。
|
||||
|
||||
**操作：** 检查脚本文件是否存在且可读。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
# 检查脚本文件
|
||||
python -c "import os; scripts = ['bench_common.py', 'bench_batch_decode.py', 'bench_batch_prefill_paged.py', 'bench_batch_prefill_ragged.py', 'bench_batch_mla.py']; [print(f'✓ {s}') if os.path.exists(s) else print(f'✗ {s} missing') for s in scripts]"
|
||||
|
||||
# 测试脚本导入
|
||||
python -c "from bench_common import setup_workspace, get_csv_path; print('脚本导入正常')"
|
||||
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```plaintext
|
||||
✓ bench_common.py
|
||||
✓ bench_batch_decode.py
|
||||
✓ bench_batch_prefill_paged.py
|
||||
✓ bench_batch_prefill_ragged.py
|
||||
✓ bench_batch_mla.py
|
||||
脚本导入正常
|
||||
|
||||
```
|
||||
|
||||
### Step 4：运行单算子 Benchmark
|
||||
|
||||
**目标：** 执行基准测试，获取 Baseline 性能数据。
|
||||
|
||||
**操作：** 运行 BatchDecode 基准测试脚本。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
python bench_batch_decode.py
|
||||
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```plaintext
|
||||
[BatchDecodeWithPagedKVCacheWrapper] Starting benchmark, total cases: 144
|
||||
[1/144] bs=1, kv_len=512, hd=64: 0.020ms, 51.53 GB/s, 0.82 TFLOPs
|
||||
[2/144] bs=1, kv_len=512, hd=128: 0.023ms, 45.28 GB/s, 0.72 TFLOPs
|
||||
...
|
||||
|
||||
Results saved to BatchDecodeWithPagedKVCacheWrapper_20260525_xxxxxx.csv
|
||||
|
||||
```
|
||||
|
||||
**常见问题：**
|
||||
|
||||
| 问题 | 解决方法 |
|
||||
| --- | --- |
|
||||
| `out of memory` | 减小 batch\_size 或 seq\_len 参数 |
|
||||
| 运行时间过长 | 脚本会自动调整重复次数，耐心等待 |
|
||||
|
||||
### Step 5：查看测试结果
|
||||
|
||||
**目标：** 查看并分析 Benchmark 输出结果。
|
||||
|
||||
**操作：** 读取生成的 CSV 结果文件。
|
||||
|
||||
**命令示例：**
|
||||
|
||||
```bash
|
||||
# 查看最新生成的 CSV 文件
|
||||
ls -lt *.csv | head -1
|
||||
|
||||
# 使用 Python 查看结果
|
||||
python -c "
|
||||
import pandas as pd
|
||||
import glob
|
||||
csv_files = glob.glob('BatchDecodeWithPagedKVCacheWrapper_*.csv')
|
||||
if csv_files:
|
||||
df = pd.read_csv(max(csv_files))
|
||||
print(df[['batch_size', 'seq_len_kv', 'head_dim', 'time_ms', 'bandwidth_GB_s', 'tflops']].head(10))
|
||||
"
|
||||
|
||||
```
|
||||
|
||||
**预期结果：**
|
||||
|
||||
```plaintext
|
||||
batch_size seq_len_kv head_dim time_ms bandwidth_GB_s tflops
|
||||
0 1 512 64 0.020429 51.528822 0.821253
|
||||
1 1 512 128 0.023340 45.278052 0.718833
|
||||
2 1 512 256 0.031949 66.153846 1.050256
|
||||
...
|
||||
|
||||
```
|
||||
|
||||
## 七、Agent 使用说明
|
||||
|
||||
在本模块中，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 函数的工作原理。
|
||||
```
|
||||
|
||||
## 八、常见问题
|
||||
|
||||
### 环境相关问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| `No GPUs are available` | MXMACA 驱动未安装或 GPU 不可见 | 检查驱动安装，运行 `python -c "import torch; print(torch.cuda.device_count())"` 验证 |
|
||||
| `ModuleNotFoundError: No module named 'flashinfer'` | flashinfer 未安装 | 执行 `pip install flashinfer` |
|
||||
| `out of memory` | GPU 显存不足 | 减小 `batch_size` 或 `seq_len` 参数 |
|
||||
|
||||
### 运行相关问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| Benchmark 运行时间过长 | 参数组合过多， workload 较大 | 耐心等待，脚本会自动调整重复次数 |
|
||||
| `KeyError: 'BatchPrefillWithPagedKVCacheKernel'` | profiler 未捕获目标 kernel | 检查 `target_kernels` 配置是否正确 |
|
||||
| CSV 文件为空 | 测试未正常完成 | 检查 GPU 显存是否充足，重新运行 |
|
||||
|
||||
### 代码相关问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| `ImportError: cannot import name 'xxx' from 'bench_common'` | 函数名拼写错误 | 检查 `bench_common.py` 中的函数名 |
|
||||
| `RuntimeError: error: device-side assert triggered` | 输入参数超出范围 | 检查 `num_qo_heads`、`num_kv_heads`、`head_dim` 配置 |
|
||||
|
||||
### 性能相关问题
|
||||
|
||||
| 问题 | 原因 | 解决办法 |
|
||||
| --- | --- | --- |
|
||||
| TFLOPs 数值异常低 | 工作负载过小，kernel 启动开销占比大 | 增大 `batch_size` 或 `seq_len` |
|
||||
| 带宽数值异常低 | 数据未正确加载到 GPU | 检查 Tensor 是否在 CUDA 设备上 |
|
||||
|
||||
## 九、下一步学习建议
|
||||
|
||||
### 1. 保存你的 Baseline 结果
|
||||
|
||||
将本次运行生成的 CSV 文件妥善保存，后续优化时需要以此作为对比基准。
|
||||
|
||||
```bash
|
||||
# 建议创建 results 目录保存
|
||||
mkdir -p results
|
||||
mv *.csv results/
|
||||
|
||||
```
|
||||
|
||||
### 2. 深入理解 FlashInfer 核心概念
|
||||
|
||||
* 阅读 FlashInfer 官方文档，理解 Paged KV Cache、Ragged KV Cache 的设计理念
|
||||
|
||||
* 学习 MLA (Multi-head Latent Attention) 的原理，了解 DeepSeek 的注意力优化方案
|
||||
|
||||
* 理解 `plan()` 和 `run()` 两阶段设计的作用
|
||||
|
||||
|
||||
### 3. 进入算子优化模块
|
||||
|
||||
参考后续优化模块，学习以下优化技术：
|
||||
|
||||
* **Kernel Tuning**：调整 Block Size、Thread Count 等参数
|
||||
|
||||
* **Memory Optimization**：减少显存占用、优化数据搬运
|
||||
|
||||
* **Compute Optimization**：提升计算效率
|
||||
|
||||
|
||||
### 4. 参考资源
|
||||
|
||||
* FlashInfer 官方仓库：https://github.com/flashinfer-ai/flashinfer
|
||||
|
||||
* FlashInfer 文档：https://flashinfer.ai
|
||||
|
||||
|
||||
### 5. 记录优化过程
|
||||
|
||||
建议维护一份优化日志，记录每次优化的改动和性能变化：
|
||||
|
||||
| 优化项 | 改动内容 | Baseline | 优化后 | 提升比例 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 例：调整 block\_size | 16 → 32 | xx ms | xx ms | xx% |
|
||||
| 完成优化后，再次运行本模块的 Benchmark 脚本，对比前后性能变化。 | | | | |
|
||||
Loading…
Reference in New Issue