forked from metax-maca/op_optimization
修复合并冲突
This commit is contained in:
Ke Xiao (i26293)
parent
c3eebcfacd
commit
397632b4c4
|
|
@ -2,7 +2,67 @@
|
|||
|
||||
## 一、赛题简要说明
|
||||
|
||||
待更新
|
||||
本赛题旨在提升国产 GPU 平台（如沐曦 MACA 平台）上的大模型推理核心算子性能，通过使用或构建 AI Agent / Skill 工作流，在国产 GPU（MACA 软件栈）上完成大模型推理核心算子库的迁移适配与性能突破。
|
||||
|
||||
### 环境准备
|
||||
|
||||
#### 开发环境设置
|
||||
|
||||
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开始项目创作。
|
||||
|
||||
## 二、任务方向
|
||||
|
||||
|
|
@ -12,7 +72,388 @@
|
|||
|
||||
面向 Prefill、Decode、Paged KV Cache、MLA Attention 等推理场景，完成 MACA 平台适配与性能优化。
|
||||
|
||||
参考知识：待更新
|
||||
参考知识：
|
||||
|
||||
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。
|
||||
|
||||
##### 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 缓存
|
||||
|
||||
教程链接：<a href="https://gitlink.org.cn/metax-maca/op_optimization/tree/master/%E5%9F%BA%E4%BA%8EAI%20Agent%E5%BC%80%E5%8F%91%E8%8C%83%E5%BC%8F%E7%9A%84%E5%9B%BD%E4%BA%A7GPU%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E7%AE%97%E5%AD%90%E5%BA%93%E4%BC%98%E5%8C%96/FlashInfer%20%E8%BF%81%E7%A7%BB%20Baseline%20%E5%AE%9E%E6%88%98.md">FlashInfer 迁移 Baseline 实战</a>
|
||||
|
||||
|
|
@ -20,7 +461,210 @@
|
|||
|
||||
围绕 `flash_attn_with_kvcache` 等核心接口，提升长序列场景下的 Attention 计算性能。
|
||||
|
||||
参考知识：待更新
|
||||
参考知识：
|
||||
|
||||
#### 名词解释
|
||||
|
||||
| 术语 | 说明 |
|
||||
| ------------------------------ | ------------------------------------------------------------ |
|
||||
| **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 编程最佳实践](
|
||||
|
||||
教程链接：<a href="https://gitlink.org.cn/metax-maca/op_optimization/tree/master/%E5%9F%BA%E4%BA%8EAI%20Agent%E5%BC%80%E5%8F%91%E8%8C%83%E5%BC%8F%E7%9A%84%E5%9B%BD%E4%BA%A7GPU%E5%A4%A7%E6%A8%A1%E5%9E%8B%E6%8E%A8%E7%90%86%E7%AE%97%E5%AD%90%E5%BA%93%E4%BC%98%E5%8C%96/FlashAttention_Baseline%E5%85%A5%E9%97%A8.md">FlashAttention Baseline 入门</a>
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue