diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/FlashInfer 迁移 Baseline 实战.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/FlashInfer 迁移 Baseline 实战.md deleted file mode 100644 index e9c3681..0000000 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/FlashInfer 迁移 Baseline 实战.md +++ /dev/null @@ -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) - - * 登录平台 - - - ![metax developer login 1](https://origin.picgo.net/2026/06/04/metax-developer-login-17ae7b237580ceb5a.png) - - * 首次登录需要先进行注册（使用邮箱或者手机号进行注册） - - - ![metax developer login 2](https://origin.picgo.net/2026/06/04/metax-developer-login-26db93f87f79ced1b.png) - - * 登录成功后进行第二步-邮箱验证，填入自己的邮箱。 - - - ![metax developer login 3](https://origin.picgo.net/2026/06/04/metax-developer-login-34fd23fa8432bdff1.png) - - * 第三步，提交申请。 - - - ![metax developer login 4](https://origin.picgo.net/2026/06/04/metax-developer-login-4b5200e9f6c00c2bf.png) - - * 获得兑换码 - - - ![metax developer login 5](https://origin.picgo.net/2026/06/04/metax-developer-login-554e19476b3419540.png) - -* 步骤2：兑换算力和登陆平台 - - * 访问模力方舟官网：[https://ai.gitee.com/](https://ai.gitee.com/) - - * 进入费用中心 - 算力券 ，点击右上角 “兑换”。 - - -![giteeai 兑换 6](https://origin.picgo.net/2026/06/04/giteeai--6e369bba40d2be65c.png) - -* 步骤3：租用算力 - - -进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图： - -![giteeai 租用 7](https://origin.picgo.net/2026/06/04/giteeai--7a35e27970245a913.png) - -* 步骤4：创建实例 - - -基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64 - -![giteeai 实例 8](https://origin.picgo.net/2026/06/04/giteeai--8840408058665bced.png) - -![giteeai 实例 9](https://origin.picgo.net/2026/06/04/giteeai--958d645e340553967.png) - -![giteeai 实例 10](https://origin.picgo.net/2026/06/04/giteeai--10a0ef0e5bcd186ba3.png) - -* 步骤5：选择工具-lab进入实例环境 - - -![giteeai 实例 11](https://origin.picgo.net/2026/06/04/giteeai--1127a207a2ab9763db.png) - -* 步骤6：在JupyterLab Terminal中检查运行环境的配置。 - - ![giteeai 实例 12](https://origin.picgo.net/2026/06/04/giteeai--12c1772b12867f6be0.png) - - * 确认沐曦 GPU 可见--可以使用`mx-smi`命令查看 - - -![giteeai 实例 13](https://origin.picgo.net/2026/06/04/giteeai--1374ac2ed7c86248b2.png) - -### 环境依赖准备 - -```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 记录映射），彻底消灭了外部碎片，新请求随时可以插空进入。 - -![page kvcache](https://origin.picgo.net/2026/06/04/page-kvcachedd5f9b5d0f1f1f4b.png) - -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（不规则张量） - -![ragged tensor](https://origin.picgo.net/2026/06/04/ragged-tensordb0121cf56ef98f4.png) - -假设有 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、编译工具和依赖版本。 - -![giteeai 实例 12](https://origin.picgo.net/2026/06/04/giteeai--12c1772b12867f6be0.png) - -**常见问题：** - -| 问题 | 解决方法 | -| --- | --- | -| `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 脚本，对比前后性能变化。 | | | | | \ No newline at end of file