From 397632b4c445272a5842e6dfa3695524c43be4a5 Mon Sep 17 00:00:00 2001 From: "Ke Xiao (i26293)" Date: Fri, 5 Jun 2026 16:40:02 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=90=88=E5=B9=B6=E5=86=B2?= =?UTF-8?q?=E7=AA=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../赛题说明.md | 650 +++++++++++++++++- 1 file changed, 647 insertions(+), 3 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md index e9e75e7..eb737d9 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md @@ -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) + + * 登录平台 + + ![platform login](https://origin.picgo.net/2026/06/04/platform-login626620122b08424d.png) + + * 首次登录需要先进行注册（使用邮箱或者手机号进行注册） + + ![platform registration](https://origin.picgo.net/2026/06/04/platform-registrationdb267074af39bf4c.png) + + * 登录成功后进行第二步-邮箱验证，填入自己的邮箱。 + + ![email verification](https://origin.picgo.net/2026/06/04/email-verificationc3f391bb747318e0.png) + + * 第三步，提交申请。 + + ![submit application](https://origin.picgo.net/2026/06/04/submit-application3bf7ac4724e13ae8.png) + + * 获得兑换码 + + ![get redeem code](https://origin.picgo.net/2026/06/04/get-redeem-code3f6a20e5f9cbbd38.png) + +2. **在模力方舟平台兑换算力券** + + * 平台链接：[https://ai.gitee.com/](https://ai.gitee.com/) + + * 1.登录模力方舟平台 + + ![ai.gitee login](https://origin.picgo.net/2026/06/04/ai.gitee-login7d9fe2b5e35e3a92.png) + + * 2.进入费用中心 - 算力券 ， 点击右上角“兑换” + + ![redeem compute voucher](https://origin.picgo.net/2026/06/04/redeem-compute-voucher0eb15e2f3f9b7bbd.png) + +3. **租用算力** + + * 模力方舟算力市场链接：https://ai.gitee.com/compute + + * 选择沐曦芯片厂商，并根据项目要求选择相应的配置。 + + ![rent compute](https://origin.picgo.net/2026/06/04/rent-compute1197cc6d884ce429.png) + +4. **创建实例** + +专属镜像文件： + +![create instance1](https://origin.picgo.net/2026/06/04/create-instance10ab33dd1b7e14727.png) + +![create instance2](https://origin.picgo.net/2026/06/04/create-instance23666efb720fefa60.png) + +![create instance3](https://origin.picgo.net/2026/06/04/create-instance3f18b4323644d3447.png) + +  进入算力容器，刚创建的实例默认开机状态，点击工具-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 记录映射），彻底消灭了外部碎片，新请求随时可以插空进入。 + +![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。 + +##### 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 实战 @@ -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 编程最佳实践]( 教程链接：FlashAttention Baseline 入门