diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer 迁移 Baseline 实战.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer 迁移 Baseline 实战.md index e9c3681..9ead2af 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer 迁移 Baseline 实战.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer 迁移 Baseline 实战.md @@ -36,35 +36,20 @@ ### 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) - +* 步骤1：获取算力券[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6) + + * 首次登录需要先进行注册（使用邮箱或者手机号进行注册） + + + * 登录成功后进行第二步-邮箱验证，填入自己的邮箱。 + + + * 第三步，提交申请。 + + + * 获得兑换码 + + * 步骤2：兑换算力和登陆平台 * 访问模力方舟官网：[https://ai.gitee.com/](https://ai.gitee.com/) @@ -72,39 +57,17 @@ * 进入费用中心 - 算力券 ，点击右上角 “兑换”。 -![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) - + - 进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图： * 步骤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) + - 基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64 * 步骤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`命令查看 - +* 步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用`mx-smi`命令查看 -![giteeai 实例 13](https://origin.picgo.net/2026/06/04/giteeai--1374ac2ed7c86248b2.png) ### 环境依赖准备 @@ -143,27 +106,6 @@ opencode ## 五、知识预备 -### 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 的阶段 @@ -183,376 +125,6 @@ https://github.com/flashinfer-ai/flashinfer 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：检查运行环境 @@ -573,7 +145,7 @@ run(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], *args, q_scale: f ### Step 2：进入项目目录 -**目标：** 进入本模块所需的源码目录。 +**目标：** 进入本模块所需的源码目录[flashinfer_baseline](baselines/flashinfer_baseline)。 **操作：** 切换到 FlashInfer Baseline 项目目录。 @@ -693,41 +265,9 @@ if csv_files: ``` -## 七、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 函数的工作原理。 -``` - -## 八、常见问题 +## 七、常见问题 ### 环境相关问题 @@ -759,7 +299,7 @@ if csv_files: | TFLOPs 数值异常低 | 工作负载过小，kernel 启动开销占比大 | 增大 `batch_size` 或 `seq_len` | | 带宽数值异常低 | 数据未正确加载到 GPU | 检查 Tensor 是否在 CUDA 设备上 | -## 九、下一步学习建议 +## 八、下一步学习建议 ### 1. 保存你的 Baseline 结果 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv new file mode 100644 index 0000000..ae909a7 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv @@ -0,0 +1,145 @@ +api,batch_size,seq_len_q,seq_len_kv,num_qo_heads,num_kv_heads,head_dim,time_ms,bandwidth_GB_s,tflops +BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,8,64,0.02042879999999998,51.528822055137894,0.8212531328320811 +BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,4,128,0.02333952000000001,45.27805199078642,0.718832949435121 +BatchDecodeWithPagedKVCacheWrapper,1,1,512,32,4,256,0.0319488,66.15384615384615,1.0502564102564103 +BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,8,64,0.023262719999999973,90.32684054143292,1.4424122372620245 +BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,4,128,0.025041919999999992,84.07278675117566,1.3399304845634845 +BatchDecodeWithPagedKVCacheWrapper,1,1,1024,32,4,256,0.033387520000000004,126.11562643766291,2.0099984664928687 +BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,8,64,0.028298240000000037,148.36258368011562,2.371485435136599 +BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,4,128,0.027745280000000008,151.4670603432367,2.418748846650673 +BatchDecodeWithPagedKVCacheWrapper,1,1,2048,32,4,256,0.03723775999999999,225.7115358174069,3.604344837068611 +BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,8,64,0.03886591999999997,215.93992886312756,3.4533526544592306 +BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,4,128,0.03426815999999998,245.03212311370103,3.916689078141344 +BatchDecodeWithPagedKVCacheWrapper,1,1,4096,32,4,256,0.066048,254.26356589147287,4.064248062015504 +BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,8,64,0.052495359999999984,319.6722910367698,5.1135082414902975 +BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,4,128,0.04628480000000001,362.6548672566371,5.799646017699114 +BatchDecodeWithPagedKVCacheWrapper,1,1,8192,32,4,256,0.08975359999999999,374.0330861380491,5.981608670849972 +BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,8,64,0.08625152000000001,389.0775258221536,6.224480588863825 +BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,4,128,0.0638464,525.6776263031276,8.408789093825181 +BatchDecodeWithPagedKVCacheWrapper,1,1,16384,32,4,256,0.13059071999999994,514.0123892417473,8.222190857053247 +BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,8,64,0.02342912,89.86013986013987,1.4321678321678322 +BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,4,128,0.02486784,84.99073502161829,1.3493102738315832 +BatchDecodeWithPagedKVCacheWrapper,2,1,512,32,4,256,0.03340287999999998,126.54812998160644,2.009074187614961 +BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,8,64,0.02839040000000001,148.02524797114512,2.3637871956717755 +BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,4,128,0.028165120000000012,149.5000908925649,2.382694055626249 +BatchDecodeWithPagedKVCacheWrapper,2,1,1024,32,4,256,0.03740160000000001,225.16084873374396,3.5885557837097872 +BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,8,64,0.03881984000000001,216.30176734370872,3.457451859667633 +BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,4,128,0.03601408000000001,233.38072220642587,3.7268126243957904 +BatchDecodeWithPagedKVCacheWrapper,2,1,2048,32,4,256,0.06728704000000002,249.82498858621207,3.9894080048698815 +BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,8,64,0.052490240000000014,319.7815060476004,5.114007023019897 +BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,4,128,0.04626431999999999,362.9924745462595,5.802213368747235 +BatchDecodeWithPagedKVCacheWrapper,2,1,4096,32,4,256,0.08993791999999999,373.44870773084375,5.969349880450872 +BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,8,64,0.08536063999999999,393.18618042226495,6.289443378119003 +BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,4,128,0.0630784,532.2077922077922,8.51116883116883 +BatchDecodeWithPagedKVCacheWrapper,2,1,8192,32,4,256,0.12952576,518.3650881492608,8.289793659577834 +BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,8,64,0.15207424000000003,441.34401723789637,7.0606423809844445 +BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,4,128,0.10330112,649.8017446471055,10.394290245836638 +BatchDecodeWithPagedKVCacheWrapper,2,1,16384,32,4,256,0.2281984,588.3060354498541,9.410599057662106 +BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,8,64,0.0283904,148.3137962128043,2.3637871956717764 +BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,4,128,0.028078080000000036,150.5470459518598,2.3900802334062696 +BatchDecodeWithPagedKVCacheWrapper,4,1,512,32,4,256,0.03707903999999999,228.00331400165706,3.619773543220106 +BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,8,64,0.03844096000000004,218.64677677144357,3.4915290356952546 +BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,4,128,0.03641856000000004,231.23857725291697,3.6854210600309254 +BatchDecodeWithPagedKVCacheWrapper,4,1,1024,32,4,256,0.06640640000000002,253.63145720894363,4.04231303006939 +BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,8,64,0.059007999999999984,284.5986984815619,4.5491366594360105 +BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,4,128,0.04641792000000003,362.1442753143611,5.783013456871825 +BatchDecodeWithPagedKVCacheWrapper,4,1,2048,32,4,256,0.08961023999999998,375.1799794309223,5.991178151068451 +BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,8,64,0.09185279999999997,365.484949832776,5.84490523968785 +BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,4,128,0.06349823999999998,528.9469440412838,8.454894371875506 +BatchDecodeWithPagedKVCacheWrapper,4,1,4096,32,4,256,0.1303347200000001,515.3991200502825,8.238340666247638 +BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,8,64,0.16568319999999992,405.1421508034613,6.480692212608161 +BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,4,128,0.10290176,652.4828341128471,10.43463031147378 +BatchDecodeWithPagedKVCacheWrapper,4,1,8192,32,4,256,0.22947840000000008,585.1673360107093,9.358107987505575 +BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,8,64,0.30601215999999987,438.65613706331163,7.017641547316293 +BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,4,128,0.18384895999999992,730.2216776205863,11.680695109724857 +BatchDecodeWithPagedKVCacheWrapper,4,1,16384,32,4,256,0.4362026666666668,615.5418398787107,9.846265564630505 +BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,8,64,0.038655999999999975,217.85430463576174,3.472105960264903 +BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,4,128,0.03645951999999999,231.87754528858312,3.6812807190001418 +BatchDecodeWithPagedKVCacheWrapper,8,1,512,32,4,256,0.06676480000000001,253.25153374233125,4.020613496932515 +BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,8,64,0.05858303999999996,286.9428421604617,4.582135990211505 +BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,4,128,0.04676608000000001,360.14889424129615,5.73996058681848 +BatchDecodeWithPagedKVCacheWrapper,8,1,1024,32,4,256,0.08992768000000002,374.58437713504884,5.970029606012297 +BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,8,64,0.092416,363.43490304709144,5.8092853185595565 +BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,4,128,0.07130112000000002,471.5208961654457,7.5296280338934345 +BatchDecodeWithPagedKVCacheWrapper,8,1,2048,32,4,256,0.14862335999999993,452.41835469202175,7.224583161085851 +BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,8,64,0.16396288000000003,409.4928803397451,6.548688483637271 +BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,4,128,0.11201536000000002,599.6891854831337,9.585665965810401 +BatchDecodeWithPagedKVCacheWrapper,8,1,4096,32,4,256,0.24935424000000006,538.7869081351894,8.61218019793848 +BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,8,64,0.3056947200000001,439.16524302415155,7.024928817874248 +BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,4,128,0.20128768000000002,667.1211273337741,10.668728697156228 +BatchDecodeWithPagedKVCacheWrapper,8,1,8192,32,4,256,0.46690133333333317,575.2104541716168,9.198875628255509 +BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,8,64,0.5866495999999998,457.6296037702917,7.321179961598886 +BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,4,128,0.37337600000000015,719.1169009256082,11.503062050051419 +BatchDecodeWithPagedKVCacheWrapper,8,1,16384,32,4,256,0.8934826666666666,601.0211546726517,9.613991308915525 +BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,8,64,0.0698112,241.26145947928126,3.845163182984965 +BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,4,128,0.04724735999999999,357.8673602080625,5.681491114000869 +BatchDecodeWithPagedKVCacheWrapper,16,1,512,32,4,256,0.08954879999999998,377.6329331046313,5.995288736420813 +BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,8,64,0.12070911999999998,278.52052935188334,4.447641669494402 +BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,4,128,0.07076864000000004,475.9947909130369,7.586282737664589 +BatchDecodeWithPagedKVCacheWrapper,16,1,1024,32,4,256,0.14710784000000002,457.9702074342197,7.2990115550605585 +BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,8,64,0.22239232000000014,302.05359609540454,4.8281425545630325 +BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,4,128,0.11209728000000002,599.8355713894217,9.578660820316067 +BatchDecodeWithPagedKVCacheWrapper,16,1,2048,32,4,256,0.2504192,537.0190145164587,8.575555101206296 +BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,8,64,0.42098688000000006,318.97256275539985,5.101070247129791 +BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,4,128,0.2027008,662.7936347562515,10.594352109118466 +BatchDecodeWithPagedKVCacheWrapper,16,1,4096,32,4,256,0.46432,578.6905582356995,9.250015713301172 +BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,8,64,0.8234496000000004,326.06851955480926,5.215822918609709 +BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,4,128,0.3726506666666667,720.6924662239522,11.525451797572705 +BatchDecodeWithPagedKVCacheWrapper,16,1,8192,32,4,256,0.8939733333333334,600.8378952392316,9.608714568667223 +BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,8,64,1.6324906666666663,328.9062896122735,5.261858317107276 +BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,4,128,0.7114879999999999,754.7590177206082,12.073196725735361 +BatchDecodeWithPagedKVCacheWrapper,16,1,16384,32,4,256,1.742272,616.4387466480549,9.860612570253094 +BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,8,64,0.08406016,400.730905104154,6.386746254111341 +BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,4,128,0.08498175999999999,397.92746113989637,6.317484034220991 +BatchDecodeWithPagedKVCacheWrapper,32,1,512,32,4,256,0.1808896,373.8918765921313,5.935895839230116 +BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,8,64,0.14712832,457.0155902004454,7.297995545657015 +BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,4,128,0.14887935999999996,452.5208061077104,7.212160396175805 +BatchDecodeWithPagedKVCacheWrapper,32,1,1024,32,4,256,0.3279462400000001,410.8661712358707,6.548279522887651 +BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,8,64,0.27223039999999993,493.51137859695325,7.888478465299983 +BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,4,128,0.27833343999999993,483.1610316029581,7.715507155733786 +BatchDecodeWithPagedKVCacheWrapper,32,1,2048,32,4,256,0.6300373333333331,426.89493109403054,6.817004435715981 +BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,8,64,0.52494336,511.6104868913858,8.181772784019977 +BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,4,128,0.5080533333333332,528.8767583455806,8.453772496325847 +BatchDecodeWithPagedKVCacheWrapper,32,1,4096,32,4,256,1.2449493333333332,431.66029782202656,6.899826653186422 +BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,8,64,1.0273706666666667,522.6954607749491,8.361086091615102 +BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,4,128,1.0078719999999999,532.9377698755399,8.522842773685548 +BatchDecodeWithPagedKVCacheWrapper,32,1,8192,32,4,256,2.446784,439.05228741073995,7.021408176610604 +BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,8,64,2.0322986666666663,528.4030903594223,8.453417534430637 +BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,4,128,2.018026666666667,532.2050425498176,8.513202262276018 +BatchDecodeWithPagedKVCacheWrapper,32,1,16384,32,4,256,4.847957333333333,443.0748433429558,7.087467154826447 +BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,8,64,0.13077504,515.1671756322919,8.210602145485865 +BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,4,128,0.14377984000000002,470.39384659212305,7.4679581226408365 +BatchDecodeWithPagedKVCacheWrapper,64,1,512,32,4,256,0.27039743999999993,500.2499431947286,7.9419525865333656 +BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,8,64,0.231424,581.0973451327434,9.279433628318584 +BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,4,128,0.25729023999999995,523.6965692907746,8.34654143118682 +BatchDecodeWithPagedKVCacheWrapper,64,1,1024,32,4,256,0.502016,536.8036715961244,8.555439061703213 +BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,8,64,0.4335923199999999,619.7010131544766,9.905542828802874 +BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,4,128,0.47517866666666664,566.018137739068,9.038636616683128 +BatchDecodeWithPagedKVCacheWrapper,64,1,2048,32,4,256,0.9693866666666666,554.9070422535212,8.861205633802816 +BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,8,64,0.8388479999999999,640.3222705424583,10.240156252384224 +BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,4,128,0.9261013333333336,580.2768883462716,9.275372232844207 +BatchDecodeWithPagedKVCacheWrapper,64,1,4096,32,4,256,1.906474666666667,563.7580287805205,9.011328335161021 +BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,8,64,1.6543999999999999,649.1803481624759,10.384350328820116 +BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,4,128,1.8147413333333327,591.9665201137942,9.466841840453299 +BatchDecodeWithPagedKVCacheWrapper,64,1,8192,32,4,256,3.774634666666667,569.2026947596871,9.10279839037844 +BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,8,64,3.2680746666666667,657.1899393567508,10.513755612274872 +BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,4,128,3.5912106666666666,598.129192458031,9.567731207451676 +BatchDecodeWithPagedKVCacheWrapper,64,1,16384,32,4,256,7.526272,570.802632697835,9.130612969608327 +BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,8,64,0.2176000000000001,619.2188235294115,9.86895058823529 +BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,4,128,0.21536768,628.0715100798782,9.971243818942565 +BatchDecodeWithPagedKVCacheWrapper,128,1,512,32,4,256,0.45757866666666663,591.2264441232692,9.386292694298103 +BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,8,64,0.39856127999999985,674.826576229382,10.776177997019683 +BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,4,128,0.39381333333333335,684.2938244853738,10.906099241603465 +BatchDecodeWithPagedKVCacheWrapper,128,1,1024,32,4,256,0.8577493333333336,628.3514810853829,10.014504539010616 +BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,8,64,0.7606186666666664,706.5238121949853,11.293352330734283 +BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,4,128,0.7354026666666665,731.4625203063357,11.680586679043863 +BatchDecodeWithPagedKVCacheWrapper,128,1,2048,32,4,256,1.6673066666666665,645.2556074467406,10.30396478792144 +BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,8,64,1.4816639999999999,725.0403006349618,11.594983197270098 +BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,4,128,1.4354773333333333,748.7338009749138,11.968053263583403 +BatchDecodeWithPagedKVCacheWrapper,128,1,4096,32,4,256,3.2697173333333325,657.4209880731792,10.508473627893627 +BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,8,64,2.9226666666666676,734.9479708029195,11.75629734306569 +BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,4,128,2.825301333333333,760.4612643087983,12.161442024827087 +BatchDecodeWithPagedKVCacheWrapper,128,1,8192,32,4,256,6.484309333333334,662.6865294520187,10.5978097594357 +BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,8,64,5.794901333333332,741.2536188134122,11.858610316747416 +BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,4,128,5.61536,765.0472760428539,12.237768680191476 +BatchDecodeWithPagedKVCacheWrapper,128,1,16384,32,4,256,12.908458666666668,665.6125232199165,10.647200957222468 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchMLAPagedAttentionWrapper_20260525_074200.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchMLAPagedAttentionWrapper_20260525_074200.csv new file mode 100644 index 0000000..432acdb --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchMLAPagedAttentionWrapper_20260525_074200.csv @@ -0,0 +1,33 @@ +api,batch_size,seq_len,num_heads,head_dim_ckv,head_dim_kpe,time_ms,bandwidth_GB_s,tflops +BatchMLAPagedAttentionWrapper,1,1024,64,512,64,0.035975679999999996,34.83953604212624,3.963964989681919 +BatchMLAPagedAttentionWrapper,1,4096,64,512,64,0.05349631999999998,89.58223668469162,10.662889409963158 +BatchMLAPagedAttentionWrapper,1,8192,64,512,64,0.06174719999999999,154.0298507462687,18.47615257048093 +BatchMLAPagedAttentionWrapper,1,16384,64,512,64,0.08995584000000004,210.63775292410136,25.3646831156265 +BatchMLAPagedAttentionWrapper,4,1024,64,512,64,0.05086207999999998,98.5705657338434,11.215139923495071 +BatchMLAPagedAttentionWrapper,4,4096,64,512,64,0.08034559999999999,238.58531145451653,28.39858531145452 +BatchMLAPagedAttentionWrapper,4,8192,64,512,64,0.10866687999999997,350.0942329438373,41.99442140972485 +BatchMLAPagedAttentionWrapper,4,16384,64,512,64,0.16821760000000002,450.56155836250184,54.2559488662304 +BatchMLAPagedAttentionWrapper,16,1024,64,512,64,0.06735359999999997,297.7423033067276,33.87645762067656 +BatchMLAPagedAttentionWrapper,16,4096,64,512,64,0.14288383999999996,536.6395528003728,63.87570143691549 +BatchMLAPagedAttentionWrapper,16,8192,64,512,64,0.21618431999999987,703.9113289992544,84.43540682321462 +BatchMLAPagedAttentionWrapper,16,16384,64,512,64,0.39363328000000025,770.1826837405613,92.74424666532254 +BatchMLAPagedAttentionWrapper,64,1024,64,512,64,0.15278592,525.0226198853926,59.73590697362689 +BatchMLAPagedAttentionWrapper,64,4096,64,512,64,0.4850483199999999,632.3256206721838,75.26512413443676 +BatchMLAPagedAttentionWrapper,64,8192,64,512,64,0.9133465600000001,666.4484158127227,79.94166423750474 +BatchMLAPagedAttentionWrapper,64,16384,64,512,64,1.7720038399999998,684.3541287134007,82.40890045926764 +BatchMLAPagedAttentionWrapper,1,1024,128,512,64,0.04499968000000001,29.491409716691315,6.338104448742746 +BatchMLAPagedAttentionWrapper,1,4096,128,512,64,0.05375743999999999,90.51859612362495,21.222191532930147 +BatchMLAPagedAttentionWrapper,1,8192,128,512,64,0.08302080000000002,115.44865864939868,27.48349059512796 +BatchMLAPagedAttentionWrapper,1,16384,128,512,64,0.11321343999999998,168.01736613603475,40.30795947901592 +BatchMLAPagedAttentionWrapper,4,1024,128,512,64,0.05178880000000003,102.50123578843295,22.028907563025196 +BatchMLAPagedAttentionWrapper,4,4096,128,512,64,0.11032576,176.4247261926861,41.36298496380175 +BatchMLAPagedAttentionWrapper,4,8192,128,512,64,0.1688268800000001,227.08800873415404,54.06014435615937 +BatchMLAPagedAttentionWrapper,4,16384,128,512,64,0.30781695999999986,247.18357299091002,59.30021207408457 +BatchMLAPagedAttentionWrapper,16,1024,128,512,64,0.10527487999999995,201.69734698344004,43.34749896651511 +BatchMLAPagedAttentionWrapper,16,4096,128,512,64,0.2629478400000002,296.0920614521874,69.41913273750409 +BatchMLAPagedAttentionWrapper,16,8192,128,512,64,0.3962367999999998,387.02674764181444,92.13485980100793 +BatchMLAPagedAttentionWrapper,16,16384,128,512,64,0.7528985599999998,404.23663979381246,96.97779742333418 +BatchMLAPagedAttentionWrapper,64,1024,128,512,64,0.3242547199999998,261.9380714026308,56.29404872811108 +BatchMLAPagedAttentionWrapper,64,4096,128,512,64,1.1793126399999994,264.07507342582215,61.91271216426548 +BatchMLAPagedAttentionWrapper,64,8192,128,512,64,2.3186406399999986,264.55887532446616,62.98038839860932 +BatchMLAPagedAttentionWrapper,64,16384,128,512,64,4.6020608,264.53295358462015,63.462389746784744 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv new file mode 100644 index 0000000..3dfd620 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv @@ -0,0 +1,33 @@ +api,batch_size,seq_len,num_qo_heads,num_kv_heads,head_dim,time_ms,bandwidth_GB_s,tflops +BatchPrefillWithPagedKVCacheWrapper,1,1024,32,4,128,0.3529011200000001,29.71302556364796,24.34091054174041 +BatchPrefillWithPagedKVCacheWrapper,1,4096,32,4,128,4.62532608,9.068126068205768,29.714435500296666 +BatchPrefillWithPagedKVCacheWrapper,1,8192,32,4,128,18.113853439999996,4.631045529757804,30.350019983820744 +BatchPrefillWithPagedKVCacheWrapper,1,16384,32,4,128,71.05519616000001,2.36115258372119,30.948099145350383 +BatchPrefillWithPagedKVCacheWrapper,4,1024,32,4,128,1.2374374399999997,33.89507917264893,27.766848858234006 +BatchPrefillWithPagedKVCacheWrapper,4,4096,32,4,128,17.896878079999997,9.374381344614939,30.71797279003423 +BatchPrefillWithPagedKVCacheWrapper,4,8192,32,4,128,71.25501952,4.709062214288198,30.861310127559136 +BatchPrefillWithPagedKVCacheWrapper,4,16384,32,4,128,283.27072767999994,2.3690716139159393,31.051895457919 +BatchPrefillWithPagedKVCacheWrapper,16,1024,32,4,128,4.752537600000002,35.301595509733566,28.919067041573737 +BatchPrefillWithPagedKVCacheWrapper,16,4096,32,4,128,70.51405312000001,9.517090711803915,31.185602844439067 +BatchPrefillWithPagedKVCacheWrapper,16,8192,32,4,128,284.16772266666663,4.723186952426669,30.953878011423416 +BatchPrefillWithPagedKVCacheWrapper,16,16384,32,4,128,1129.139136,2.377346134250013,31.160351250841774 +BatchPrefillWithPagedKVCacheWrapper,64,1024,32,4,128,18.757478399999997,35.77712449878125,29.3086203894016 +BatchPrefillWithPagedKVCacheWrapper,64,4096,32,4,128,281.4907093333333,9.536210151864244,31.248253425628754 +BatchPrefillWithPagedKVCacheWrapper,64,8192,32,4,128,1134.7048106666668,4.731370722616177,31.007511167737377 +BatchPrefillWithPagedKVCacheWrapper,64,16384,32,4,128,4514.139178666666,2.378619226173592,31.177037921302507 +BatchPrefillWithPagedKVCacheWrapper,1,1024,32,4,256,0.7928422399999997,26.4510629504301,21.668710768992337 +BatchPrefillWithPagedKVCacheWrapper,1,4096,32,4,256,12.533002240000002,6.69321511267838,21.932327281224513 +BatchPrefillWithPagedKVCacheWrapper,1,8192,32,4,256,49.81321727999999,3.368024977325858,22.072688491402744 +BatchPrefillWithPagedKVCacheWrapper,1,16384,32,4,256,190.01136128,1.765917141688929,23.14622915954513 +BatchPrefillWithPagedKVCacheWrapper,4,1024,32,4,256,3.111116800000001,26.963333552761494,22.088362846422218 +BatchPrefillWithPagedKVCacheWrapper,4,4096,32,4,256,47.738091520000026,7.02885912101079,23.032165567728153 +BatchPrefillWithPagedKVCacheWrapper,4,8192,32,4,256,190.14286336,3.529391680241077,23.130221315627924 +BatchPrefillWithPagedKVCacheWrapper,4,16384,32,4,256,759.6848640000004,1.76675532658763,23.157215416649382 +BatchPrefillWithPagedKVCacheWrapper,16,1024,32,4,256,12.28442624,27.31461066593534,22.376129057534232 +BatchPrefillWithPagedKVCacheWrapper,16,4096,32,4,256,191.34602666666663,7.014398487291994,22.984780963158407 +BatchPrefillWithPagedKVCacheWrapper,16,8192,32,4,256,759.7649706666668,3.5331380935403933,23.15477380982632 +BatchPrefillWithPagedKVCacheWrapper,16,16384,32,4,256,3028.668266666667,1.77263029400997,23.234219789647476 +BatchPrefillWithPagedKVCacheWrapper,64,1024,32,4,256,49.26948266666667,27.241554149868346,22.316281159572153 +BatchPrefillWithPagedKVCacheWrapper,64,4096,32,4,256,763.6229333333335,7.030576067909256,23.037791659325052 +BatchPrefillWithPagedKVCacheWrapper,64,8192,32,4,256,3037.7449386666663,3.534667477616765,23.16479678130923 +BatchPrefillWithPagedKVCacheWrapper,64,16384,32,4,256,12110.653866666667,1.7732185822854112,23.241930601731337 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv new file mode 100644 index 0000000..4f20821 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv @@ -0,0 +1,49 @@ +api,batch_size,seq_len,num_qo_heads,num_kv_heads,head_dim_qk,head_dim_vo,time_ms,bandwidth_GB_s,tflops +BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,128,128,0.031580159999999996,66.66666666666667,272.00415045395596 +BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,128,128,0.0424448,197.82870928829917,3238.0634016887816 +BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,128,128,0.057313279999999994,292.871180989816,9592.119206717885 +BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,128,128,0.06972416000000001,481.36290204141574,31538.89922161844 +BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,128,128,0.04327423999999998,194.60482725982024,793.9998106956938 +BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,128,128,0.06579199999999998,510.5058365758757,8355.967501945528 +BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,128,128,0.09618432000000002,698.0517406579366,22862.596060896405 +BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,128,128,0.15411199999999997,871.12292358804,57075.97735548174 +BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,128,128,0.07452671999999999,451.99230557845567,1844.1567463588901 +BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,128,128,0.1668906666666667,805.0108653969065,13176.43041083983 +BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,128,128,0.2874026666666667,934.46080760095,30605.46766745843 +BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,128,128,0.5342506666666667,1005.1498622369525,65857.42289917343 +BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,128,128,0.15733333333333333,856.4111186440679,3494.2106814915255 +BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,128,128,0.5614719999999999,957.1184315513509,15666.129428017784 +BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,128,128,1.1031466666666667,973.8198414233224,31894.55505055115 +BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,128,128,2.1813759999999998,984.7032038493136,64517.75776176505 +BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,192,128,0.03564544000000001,73.88681413386956,301.22838264866414 +BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,192,128,0.04922368,213.27231121281466,3490.1635115456625 +BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,192,128,0.061327359999999984,342.16062781766584,11205.353815328106 +BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,192,128,0.08377343999999999,500.8189707859675,32812.059161471705 +BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,192,128,0.049623040000000056,212.29880313660726,865.5187783739157 +BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,192,128,0.08634367999999998,486.3377609108161,7958.831119544594 +BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,192,128,0.13644799999999999,615.1444652908068,20145.249981238278 +BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,192,128,0.2321706666666666,722.8359827253516,47357.904577781876 +BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,192,128,0.09042944,465.99479107688825,1899.8093081191257 +BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,192,128,0.3087573333333334,544.0154770952807,8902.716705589717 +BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,192,128,0.5995946666666665,559.9464882943145,18337.58185156195 +BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,192,128,1.1809706666666668,568.4182232017052,37240.94624227753 +BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,192,128,0.2555306666666667,659.6413424611787,2689.2849156787443 +BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,192,128,0.9085866666666667,739.472740079831,12101.340115520075 +BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,192,128,1.7810773333333334,754.017631276351,24693.1810808739 +BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,192,128,3.5260586666666662,761.5134193267346,49891.92667360423 +BatchPrefillWithRaggedKVCacheWrapper,1,1024,32,4,256,256,0.044037119999999964,95.61678874549479,390.12245087780525 +BatchPrefillWithRaggedKVCacheWrapper,1,4096,32,4,256,256,0.08118271999999997,206.86175580222005,3385.916448032292 +BatchPrefillWithRaggedKVCacheWrapper,1,8192,32,4,256,256,0.11204607999999996,299.6161579235972,9813.030743922503 +BatchPrefillWithRaggedKVCacheWrapper,1,16384,32,4,256,256,0.14619648000000002,459.1440778875113,30083.12177628353 +BatchPrefillWithRaggedKVCacheWrapper,4,1024,32,4,256,256,0.07792639999999999,216.1366622864652,881.8510381077531 +BatchPrefillWithRaggedKVCacheWrapper,4,4096,32,4,256,256,0.13784064000000001,487.3337790654483,7976.686902904687 +BatchPrefillWithRaggedKVCacheWrapper,4,8192,32,4,256,256,0.22408533333333336,599.2505712109672,19626.65938766184 +BatchPrefillWithRaggedKVCacheWrapper,4,16384,32,4,256,256,0.3959893333333334,678.0510720827496,44425.908890852275 +BatchPrefillWithRaggedKVCacheWrapper,16,1024,32,4,256,256,0.15150079999999996,444.6907739101049,1814.366042581954 +BatchPrefillWithRaggedKVCacheWrapper,16,4096,32,4,256,256,0.4274346666666664,628.6284687562392,10289.400589339195 +BatchPrefillWithRaggedKVCacheWrapper,16,8192,32,4,256,256,0.7913173333333334,678.7833823093305,22231.518637802277 +BatchPrefillWithRaggedKVCacheWrapper,16,16384,32,4,256,256,1.5360853333333337,699.1824898616742,45810.43946625186 +BatchPrefillWithRaggedKVCacheWrapper,64,1024,32,4,256,256,0.43906133333333336,613.773091686507,2504.2324256352945 +BatchPrefillWithRaggedKVCacheWrapper,64,4096,32,4,256,256,1.6363946666666664,656.8039006075145,10750.576497692491 +BatchPrefillWithRaggedKVCacheWrapper,64,8192,32,4,256,256,3.234005333333333,664.3564257160574,21759.006842803803 +BatchPrefillWithRaggedKVCacheWrapper,64,16384,32,4,256,256,6.420821333333334,669.0757535484556,43837.84598543405 diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchDecodeWithPagedKVCacheWrapper_20260525_093143.csv differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchMLAPagedAttentionWrapper_20260525_074200.csv differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithPagedKVCacheWrapper_20260526_150417.csv differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._BatchPrefillWithRaggedKVCacheWrapper_20260526_145454.csv differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_decode.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_decode.py new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_decode.py differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_mla.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_mla.py new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_mla.py differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_paged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_paged.py new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_paged.py differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_ragged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_ragged.py new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_batch_prefill_ragged.py differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_common.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_common.py new file mode 100644 index 0000000..53c3adb Binary files /dev/null and b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/__MACOSX/._bench_common.py differ diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py new file mode 100644 index 0000000..32861ca --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_decode.py @@ -0,0 +1,101 @@ +""" +Benchmark script for BatchDecodeWithPagedKVCacheWrapper +seq_len_q=1 (decode mode), seq_len_kv from 1K to 16K +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import dtype, page_block_size, setup_workspace, setup_paged_kv_indptr, run_with_profiler, get_csv_path, compute_reps + +target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] + + +def bench_batch_decode( + batch_size, + seq_len_kv, + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, +): + """Benchmark BatchDecodeWithPagedKVCacheWrapper""" + seq_lens = [seq_len_kv] * batch_size + + kv_indptr, last_page_len, num_blocks = setup_paged_kv_indptr(batch_size, seq_lens) + + q = torch.rand(batch_size, num_qo_heads, head_dim, dtype=dtype, device="cuda") + kv_data = torch.randn(num_blocks, 2, page_block_size, num_kv_heads, head_dim, dtype=dtype, device="cuda") + + workspace_buffer = setup_workspace() + wrapper = flashinfer.BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, kv_layout="NHD", use_tensor_cores=True + ) + wrapper.plan( + kv_indptr.to("cuda"), + torch.arange(num_blocks, dtype=torch.int32, device="cuda"), + last_page_len.to("cuda"), + num_qo_heads, + num_kv_heads, + head_dim, + page_block_size, + data_type=dtype, + q_data_type=dtype, + ) + + reps = compute_reps(batch_size, seq_len_kv, head_dim, base_reps=100) + ms = run_with_profiler(lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps) + + io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() + flops = 2 * batch_size * seq_len_kv * num_qo_heads * num_kv_heads * head_dim + return ms, io, flops + + +def run_benchmark(): + records = [] + + batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128] + head_dims = [64, 128, 256] + seq_lens_kv = [512, 1024, 2048, 4096, 8192, 16384] + + api_name = "BatchDecodeWithPagedKVCacheWrapper" + test_cases = list(itertools.product(batch_sizes, seq_lens_kv, head_dims)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + print(f" seq_len_q=1 (decode mode), causal=False") + for idx, (bs, sl_kv, hd) in enumerate(test_cases, 1): + num_qo_heads = 32 + num_kv_heads = 8 if hd == 64 else 4 + ms, io, flops = bench_batch_decode(bs, sl_kv, num_qo_heads, num_kv_heads, hd, page_block_size) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append({ + "api": api_name, + "batch_size": bs, + "seq_len_q": 1, + "seq_len_kv": sl_kv, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim": hd, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + }) + print(f" [{idx}/{total_cases}] bs={bs}, kv_len={sl_kv}, hd={hd}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + + return records + + +if __name__ == "__main__": + import numpy as np + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchDecodeWithPagedKVCacheWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py new file mode 100644 index 0000000..5c48bc9 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_mla.py @@ -0,0 +1,110 @@ +""" +Benchmark script for BatchMLAPagedAttentionWrapper +headdim: ckv=512, kpe=64 (DeepSeek MLA configuration) +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import dtype, page_block_size, setup_workspace, run_with_profiler, get_csv_path, compute_reps + +target_kernels = ["BatchMLAPagedAttentionKernel"] + + +def bench_batch_mla_paged_attention( + batch_size, + seq_len, + num_heads, + head_dim_ckv, + head_dim_kpe, +): + """Benchmark BatchMLAPagedAttentionWrapper for DeepSeek MLA""" + # MLA decode mode: q has length 1, not seq_len + q_nope = torch.randn(batch_size, num_heads, head_dim_ckv, dtype=dtype, device="cuda") + q_pe = torch.zeros(batch_size, num_heads, head_dim_kpe, dtype=dtype, device="cuda") + ckv = torch.randn(batch_size * seq_len, 1, head_dim_ckv, dtype=dtype, device="cuda") + kpe = torch.zeros(batch_size * seq_len, 1, head_dim_kpe, dtype=dtype, device="cuda") + + sm_scale = 1.0 / ((head_dim_ckv + head_dim_kpe) ** 0.5) + + # q_indptr for decode: each query has length 1 + q_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + kv_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * seq_len + kv_indices = torch.arange(0, batch_size * seq_len, dtype=torch.int32, device="cuda") + kv_lens = torch.full((batch_size,), seq_len, dtype=torch.int32, device="cuda") + + page_size = 1 # MLA uses page_size=1 + + workspace_buffer = setup_workspace() + wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(workspace_buffer, backend="auto") + wrapper.plan( + q_indptr, + kv_indptr, + kv_indices, + kv_lens, + num_heads, + head_dim_ckv, + head_dim_kpe, + page_size, + False, # causal + sm_scale, + q_nope.dtype, + ckv.dtype, + ) + + reps = compute_reps(batch_size, seq_len, head_dim_ckv + head_dim_kpe, base_reps=100) + ms = run_with_profiler(lambda: wrapper.run(q_nope, q_pe, ckv, kpe, return_lse=False), target_kernels=target_kernels, reps=reps) + + io = sum([t.numel() * t.element_size() for t in [q_nope, q_pe, ckv, kpe]]) + # MLA FLOPs: 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len + flops = 2 * batch_size * num_heads * (2 * head_dim_ckv + head_dim_kpe) * seq_len + return ms, io, flops + + +def run_benchmark(): + records = [] + + # MLA configuration - same as DeepSeek + head_dim_ckv = 512 + head_dim_kpe = 64 + batch_sizes = [1, 4, 16, 64] + seq_lens = [1024, 4096, 8192, 16384] + num_heads_list = [64, 128] + + api_name = "BatchMLAPagedAttentionWrapper" + test_cases = list(itertools.product(num_heads_list, batch_sizes, seq_lens)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + for idx, (num_heads, bs, sl) in enumerate(test_cases, 1): + ms, io, flops = bench_batch_mla_paged_attention(bs, sl, num_heads, head_dim_ckv, head_dim_kpe) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append({ + "api": api_name, + "batch_size": bs, + "seq_len": sl, + "num_heads": num_heads, + "head_dim_ckv": head_dim_ckv, + "head_dim_kpe": head_dim_kpe, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + }) + print(f" [{idx}/{total_cases}] bs={bs}, sl={sl}, num_heads={num_heads}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs") + + return records + + +if __name__ == "__main__": + import numpy as np + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchMLAPagedAttentionWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py new file mode 100644 index 0000000..1dee3d0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_paged.py @@ -0,0 +1,133 @@ +""" +Benchmark script for BatchPrefillWithPagedKVCacheWrapper +headdim: 64/128/256 +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import ( + dtype, + setup_workspace, + setup_paged_kv_indptr, + run_with_profiler, + get_csv_path, + compute_reps, +) + +target_kernels = ["BatchPrefillWithPagedKVCacheKernel"] + + +def bench_batch_prefill_with_paged_kv_cache( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim, + causal=True, +): + """Benchmark BatchPrefillWithPagedKVCacheWrapper""" + q_lens = [seq_len] * batch_size + kv_lens = [seq_len] * batch_size + + qo_indptr = torch.cat( + [torch.tensor([0]), torch.cumsum(torch.tensor(q_lens), 0)], dim=0 + ).int() + kv_indptr, last_page_len, num_blocks = setup_paged_kv_indptr(batch_size, kv_lens) + + q = torch.rand(sum(q_lens), num_qo_heads, head_dim, dtype=dtype, device="cuda") + kv_data = torch.randn( + num_blocks, 2, 16, num_kv_heads, head_dim, dtype=dtype, device="cuda" + ) + + workspace_buffer = setup_workspace() + wrapper = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace_buffer, kv_layout="NHD", backend="auto" + ) + wrapper.plan( + qo_indptr, + kv_indptr, + torch.arange(num_blocks, dtype=torch.int32, device="cuda"), + last_page_len, + num_qo_heads, + num_kv_heads, + head_dim, + 16, + q_data_type=dtype, + kv_data_type=dtype, + ) + + reps = compute_reps(batch_size, seq_len, head_dim, base_reps=100) + ms = run_with_profiler( + lambda: wrapper.run(q, kv_data), target_kernels=target_kernels, reps=reps + ) + + io = q.numel() * q.element_size() + kv_data.numel() * kv_data.element_size() + # Attention FLOPs calculation: + # - causal=True: triangular pattern + # - causal=False: full attention + flops = ( + 2 + * batch_size + * seq_len + * seq_len + * num_qo_heads + * head_dim + * (1 if causal else 2) + ) + return ms, io, flops + + +def run_benchmark(): + records = [] + + batch_sizes = [1, 4, 16, 64] + seq_lens = [1024, 4096, 8192, 16384] + head_dims = [128, 256] + + api_name = "BatchPrefillWithPagedKVCacheWrapper" + test_cases = list(itertools.product(head_dims, batch_sizes, seq_lens)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + for idx, (head_dim, bs, sl) in enumerate(test_cases, 1): + num_qo_heads = 32 + num_kv_heads = 8 if head_dim == 64 else 4 + ms, io, flops = bench_batch_prefill_with_paged_kv_cache( + bs, sl, num_qo_heads, num_kv_heads, head_dim + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append( + { + "api": api_name, + "batch_size": bs, + "seq_len": sl, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + } + ) + print( + f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd={head_dim}: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" + ) + + return records + + +if __name__ == "__main__": + import numpy as np + + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchPrefillWithPagedKVCacheWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py new file mode 100644 index 0000000..43aa4e0 --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_batch_prefill_ragged.py @@ -0,0 +1,136 @@ +""" +Benchmark script for BatchPrefillWithRaggedKVCacheWrapper +headdim configurations: [64,64], [128,128], [192,128], [256,256] +""" + +import itertools +import pandas as pd +import torch + +import flashinfer +from bench_common import ( + dtype, + setup_workspace, + run_with_profiler, + get_csv_path, + compute_reps, +) + +target_kernels = [ + "BatchPrefillWithRaggedKVCacheKernel", + "PersistentVariableLengthMergeStates", +] + + +def bench_batch_prefill_with_ragged_kv_cache( + batch_size, + seq_len, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal=True, +): + """Benchmark BatchPrefillWithRaggedKVCacheWrapper for MLA""" + qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + kv_indptr = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * seq_len + ) + + q = torch.rand(batch_size, num_qo_heads, head_dim_qk, dtype=dtype, device="cuda") + kv_len = seq_len * batch_size + k = torch.rand(kv_len, num_kv_heads, head_dim_qk, dtype=dtype, device="cuda") + v = torch.rand(kv_len, num_kv_heads, head_dim_vo, dtype=dtype, device="cuda") + + workspace_buffer = setup_workspace() + wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper( + workspace_buffer, kv_layout="NHD", backend="auto" + ) + wrapper.plan( + qo_indptr, + kv_indptr, + num_qo_heads, + num_kv_heads, + head_dim_qk, + head_dim_vo, + causal=causal, + q_data_type=dtype, + kv_data_type=dtype, + ) + + reps = compute_reps(batch_size, seq_len, head_dim_qk + head_dim_vo, base_reps=100) + ms = run_with_profiler( + lambda: wrapper.run(q, k, v), target_kernels=target_kernels, reps=reps + ) + + io = ( + q.numel() * q.element_size() + + k.numel() * k.element_size() + + v.numel() * v.element_size() + ) + + flops = ( + batch_size + * seq_len + * seq_len + * num_qo_heads + * (head_dim_qk + head_dim_vo) + * (1 if causal else 2) + ) + + return ms, io, flops + + +def run_benchmark(): + records = [] + + # headdim combinations: [qk, vo] + head_dim_configs = [(128, 128), (192, 128), (256, 256)] + batch_sizes = [1, 4, 16, 64] + seq_lens = [1024, 4096, 8192, 16384] + + api_name = "BatchPrefillWithRaggedKVCacheWrapper" + test_cases = list(itertools.product(head_dim_configs, batch_sizes, seq_lens)) + total_cases = len(test_cases) + + print(f"[{api_name}] Starting benchmark, total cases: {total_cases}") + for idx, ((head_dim_qk, head_dim_vo), bs, sl) in enumerate(test_cases, 1): + num_qo_heads = 32 + num_kv_heads = 4 + ms, io, flops = bench_batch_prefill_with_ragged_kv_cache( + bs, sl, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo + ) + bw = io / ms / 1e6 + tflops = flops / ms / 1e9 + records.append( + { + "api": api_name, + "batch_size": bs, + "seq_len": sl, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "head_dim_qk": head_dim_qk, + "head_dim_vo": head_dim_vo, + "time_ms": ms, + "bandwidth_GB_s": bw, + "tflops": tflops, + } + ) + print( + f" [{idx}/{total_cases}] bs={bs}, sl={sl}, hd=[{head_dim_qk},{head_dim_vo}]: {ms:.3f}ms, {bw:.2f} GB/s, {tflops:.2f} TFLOPs" + ) + + return records + + +if __name__ == "__main__": + import numpy as np + + np.random.seed(42) + torch.random.manual_seed(42) + + records = run_benchmark() + df = pd.DataFrame(records) + csv_path = get_csv_path("BatchPrefillWithRaggedKVCacheWrapper") + df.to_csv(csv_path, index=False) + print(f"\nResults saved to {csv_path}") diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py new file mode 100644 index 0000000..3a0563a --- /dev/null +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/baselines/flashinfer_baseline/bench_common.py @@ -0,0 +1,91 @@ +""" +Common utilities for FlashInfer benchmarks +""" + +import os +import random +from datetime import datetime +import numpy as np +import torch + +page_block_size = 16 +dtype = torch.bfloat16 + + +def get_timestamp(): + return datetime.now().strftime("%Y%m%d_%H%M%S") + + +def get_csv_path(prefix): + """Generate CSV path in current execution directory with timestamp""" + return f"{prefix}_{get_timestamp()}.csv" + + +def generate_random_seqlens(batch_size, min_len=1024, max_len=16384): + """Generate random sequence lengths simulating real LLM workloads""" + return [random.randint(min_len, max_len) for _ in range(batch_size)] + + +def setup_workspace(): + """Create workspace buffer for FlashInfer""" + return torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda") + + +def setup_paged_kv_indptr(batch_size, seq_lens): + """Setup paged KV cache indptr and last_page_len""" + seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32) + seq_lens_blocks = torch.ceil(seq_lens_tensor / page_block_size).int() + kv_indptr = torch.cat([torch.tensor([0]), torch.cumsum(seq_lens_blocks, 0)], dim=0).int() + num_blocks = kv_indptr[-1].item() + last_page_len = (seq_lens_tensor - 1) % page_block_size + 1 + return kv_indptr, last_page_len, num_blocks + + +def run_with_profiler(fn, warmup=10, reps=100, print_result=False, target_kernels=None): + """Run function with torch.profiler and return sum of specific kernel times in ms""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA], + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as prof: + for _ in range(reps): + fn() + torch.cuda.synchronize() + + if print_result: + print(prof.key_averages().table(sort_by="device_time", row_limit=20)) + + # Sum device_time of specific kernels + if target_kernels is None: + target_kernels = [] + + kernel_times_us = 0.0 + for evt in prof.key_averages(): + if any(k in evt.key for k in target_kernels): + kernel_times_us += evt.device_time + + ms = kernel_times_us / 1e3 + return ms + + +def compute_reps(batch_size, seq_len, head_dim, base_reps=100): + """Dynamically compute repetition count based on workload size""" + # Estimate workload: batch_size * seq_len * head_dim + workload = batch_size * seq_len * head_dim + if workload < 1e5: # tiny workload + return base_reps + elif workload < 1e6: # small workload + return base_reps // 2 + elif workload < 1e7: # medium workload + return base_reps // 4 + elif workload < 1e8: # large workload + return base_reps // 8 + elif workload < 1e9: # very large workload + return base_reps // 16 + else: # huge workload + return base_reps // 32 \ No newline at end of file diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/赛题说明.md index e9e75e7..07f0029 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,387 @@ 面向 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 +460,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 入门