forked from metax-maca/op_optimization
新增代码
This commit is contained in:
Ke Xiao (i26293)
parent
397632b4c4
commit
82c7b01da8
|
|
@ -36,35 +36,20 @@
|
|||
|
||||
### GPU准备
|
||||
|
||||
* 步骤1：获取算力券
|
||||
|
||||
[https://developer.metax-tech.com/activities/6](https://developer.metax-tech.com/activities/6)
|
||||
|
||||
* 登录平台
|
||||
|
||||
|
||||
|
||||
|
||||
* 首次登录需要先进行注册（使用邮箱或者手机号进行注册）
|
||||
|
||||
|
||||
|
||||
|
||||
* 登录成功后进行第二步-邮箱验证，填入自己的邮箱。
|
||||
|
||||
|
||||
|
||||
|
||||
* 第三步，提交申请。
|
||||
|
||||
|
||||
|
||||
|
||||
* 获得兑换码
|
||||
|
||||
|
||||
|
||||
|
||||
* 步骤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 @@
|
|||
* 进入费用中心 - 算力券 ，点击右上角 “兑换”。
|
||||
|
||||
|
||||
|
||||
|
||||
* 步骤3：租用算力
|
||||
|
||||
|
||||
进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图：
|
||||
|
||||
|
||||
|
||||
- 进入算力容器，选择沐曦，租用算力，建议优先选16G显存/32G显存，如下图：
|
||||
* 步骤4：创建实例
|
||||
|
||||
|
||||
基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
- 基础镜像：maca-pytorch:3.7.1.5-torch2.8-py312-ubuntu24.04-amd64
|
||||
* 步骤5：选择工具-lab进入实例环境
|
||||
|
||||
|
||||
|
||||
|
||||
* 步骤6：在JupyterLab Terminal中检查运行环境的配置。
|
||||
|
||||
|
||||
|
||||
* 确认沐曦 GPU 可见--可以使用`mx-smi`命令查看
|
||||
|
||||
* 步骤6：在JupyterLab Terminal中检查运行环境的配置，确认沐曦 GPU 可见--可以使用`mx-smi`命令查看
|
||||
|
||||
|
||||
|
||||
### 环境依赖准备
|
||||
|
||||
|
|
@ -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 记录映射），彻底消灭了外部碎片，新请求随时可以插空进入。
|
||||
|
||||
|
||||
|
||||
flashInfer 将分页 KV Cache 视作一个**块稀疏矩阵**，并巧妙地使用了 **CSR (Compressed Sparse Row)** 格式来建立索引。图 1 中的三个关键数组就是 CSR 格式的体现：
|
||||
|
||||
* `**kv_page_indices**` **(页索引池):** 把所有请求当前占用的物理页编号，按顺序“平铺”拼接在一起。
|
||||
|
||||
* 图中蓝、橙、绿三个请求分别使用了 `[0, 5, 8]`、`[1, 6, 7]`、`[3, 4]`。
|
||||
|
||||
* 拼接后得到：`[0, 5, 8, 1, 6, 7, 3, 4]`。
|
||||
|
||||
* `**kv_indptr**` **(索引指针):** 用于标记每个请求在 `kv_page_indices` 中的**起始和结束位置**。
|
||||
|
||||
* 数组长度固定为 `num_requests + 1`。
|
||||
|
||||
* 图中数值为 `[0, 3, 6, 8]`。这意味着：
|
||||
|
||||
* 请求 0 (蓝) 的页索引在 `kv_page_indices` 的 `0` 到 `3` 之间（即 `[0, 5, 8]`，共 3 页）。
|
||||
|
||||
* 请求 1 (橙) 的页索引在 `3` 到 `6` 之间（即 `[1, 6, 7]`，共 3 页）。
|
||||
|
||||
* 请求 2 (绿) 的页索引在 `6` 到 `8` 之间（即 `[3, 4]`，共 2 页）。
|
||||
|
||||
* `**kv_last_page_lens**` **(尾页有效长度):** 由于一个请求的 Token 总数很少能刚好被 `page_size`（每页容量，图中为 8）整除，最后一个页通常是不满的。这个数组记录了每个请求**最后一页实际存储的 Token 数量**。
|
||||
|
||||
* 图中分别为 `[6, 4, 7]`。
|
||||
|
||||
|
||||
图 1 左下角展示了在解码（Decode/Append）阶段，新生成的 Token 是如何追加到 KV Cache 中的。这里的细节非常值得注意：
|
||||
|
||||
* `**qo_indptr = [0, 4, 6, 9]**`: 这表示当前批次中，各个请求**新追加**的 Token 数量（Query/Output）。
|
||||
|
||||
* 请求 0 (蓝) 追加了 `4 - 0 = 4` 个 Token。
|
||||
|
||||
* 请求 1 (橙) 追加了 `6 - 4 = 2` 个 Token。
|
||||
|
||||
* 请求 2 (绿) 追加了 `9 - 6 = 3` 个 Token。
|
||||
|
||||
|
||||
### Flashinfer API
|
||||
|
||||
#### Ragged Tensor（不规则张量）
|
||||
|
||||
|
||||
|
||||
假设有 3 个请求，长度分别是 5, 3, 4。所以batchsize=3，用户请求的seq\_len分别为 5,3,4，因为在深度学习中由于底层的矩阵运算要求张量（Tensor）必须是**规整的矩形**（比如 `[batch_size, seq_len, hidden_dim]`），当一个 Batch 中包含**不同长度的句子时**，我们通常会按最长的那句话进行 **Padding（补零）**。而图中的核心思想是：**“拒绝 Padding，把所有 Token 拍扁拼接到一起。”**
|
||||
|
||||
这里不使用 `[3(batch_size), 5(seq_len), num_heads, head_dim]` 这样的多维规整张量，而是把所有请求的 Token 首尾相连，打包成一个长度为 12（即 5+3+4）的连续一维数组。这就是图中的 `data: (12, num_heads, head_dim)` = `data：(seq_len(5+3+4)，num_heads，head_dim)`的由来。
|
||||
|
||||
既然数据被拼接到了一起，系统怎么知道哪个 Token 属于哪个请求呢？这就是 `indptr`（Index Pointer，索引指针）发挥作用的地方。
|
||||
|
||||
* **颜色块代表不同的请求（Request）：**
|
||||
|
||||
* 🟦 蓝色：Request 0，长度为 5
|
||||
|
||||
* 🟧 橙色：Request 1，长度为 3
|
||||
|
||||
* 🟩 绿色：Request 2，长度为 4
|
||||
|
||||
* `**indptr = [0, 5, 8, 12]**`**：**
|
||||
|
||||
* 这是一个一维数组，用来记录每个请求在 `data` 数组里的**起始和结束位置**。它的长度永远是 `num_requests + 1`，且第一个元素必定是 0。
|
||||
|
||||
* Request 0 的数据在 `data[0:5]` （对应图中 `indptr[0]` 到 `indptr[1]`）
|
||||
|
||||
* Request 1 的数据在 `data[5:8]` （对应图中 `indptr[1]` 到 `indptr[2]`）
|
||||
|
||||
* Request 2 的数据在 `data[8:12]` （对应图中 `indptr[2]` 到 `indptr[3]`）
|
||||
|
||||
|
||||
**序列长度计算：** 第 $i$个请求的长度（Sequence length）可以直接通过 $indptr\[i+1\] - indptr\[i\]$算出来。
|
||||
|
||||
**总 Token 数：** `indptr` 的最后一个元素（即 `indptr[-1]`，在这里是 12），就是这个 Batch 中所有 Token 的总和。
|
||||
|
||||
**数据切片读取：** 当你需要单独提取第 $i$个请求的 $Q/K/V$矩阵时，只需要执行切片操作 `data[indptr[i]:indptr[i+1]]` 即可精准拿取，没有任何多余的 Padding 元素。
|
||||
|
||||
#### flashinfer.BatchPrefillWithRaggedKVCacheWrapper() 类
|
||||
|
||||
参考链接：https://docs.flashinfer.ai/api/attention.html#flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper
|
||||
|
||||
1. 构造函数：
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, kv_indptr_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None) → None
|
||||
```
|
||||
|
||||
2. 常用初始化参数理解：
|
||||
|
||||
1. `_**float_workspace_buffer**_``(``_torch.Tensor_``)`：用户预留的浮点工作空间缓冲区，用于在 split-k 算法中存储中间的注意力计算结果。建议大小为 128MB，其设备类型需与输入张量所在的设备保持一致。
|
||||
|
||||
2. `_**kv_layout**_`_**：**_输入 K/V 张量的显存布局格式，可以是 `NHD` 或 `HND`。默认_**'NHD'**_
|
||||
|
||||
* NHD：`(seq_len, num_heads, head_dim)`
|
||||
|
||||
* HND：`(num_heads, seq_len, head_dim)`
|
||||
|
||||
3. `_**backend**_`_**：**_底层实现引擎。可选值包括 `auto`、`fa2`、`fa3`、`cudnn`、`cutlass` 或 `cute-dsl`，默认值为 `auto`。系统会根据显卡架构自动选择最优后端。其中 `cute-dsl` 是专为最新一代 Blackwell 架构（如 SM100+ 系列）准备的算子。
|
||||
|
||||
4. `_**jit_args**_` _**&**_ `_**jit_kwargs**_`_**：**_用于即时编译（JIT, Just-In-Time）的参数列表和字典参数。如果提供，框架将会在运行时动态编译底层算子（手动实现算子）；否则，直接使用预编译好的默认算子。
|
||||
|
||||
3. **.Plan()** 方法
|
||||
|
||||
|
||||
**Plan()** 方法的作用就是“运筹帷幄”的预处理（AOT, Ahead-of-Time Setup）阶段：**任务规划：** 接收当前 Batch 中所有请求的形状、长度和硬件规格，计算出最优的 GPU 算力调度方案。**显存分配：** 在底层预先创建并缓存计算所需的辅助数据结构和临时工作空间（Workspace）。**解耦计算：** 将“准备工作”与真正的“执行工作（`run` 方法）”解耦。调用 `plan` 搭建好“脚手架”后，后续调用 `run` 时 GPU 就可以直接根据图纸极速开工，从而把 **CPU 调度开销降到最低**。`plan()` 方法必须在任何 `run()`之前
|
||||
|
||||
```Python
|
||||
plan(qo_indptr: Tensor, kv_indptr: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, window_left: int = -1, logits_soft_cap: float | None = None, sm_scale: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False, seq_lens: Tensor | None = None, seq_lens_q: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, v_indptr: Tensor | None = None, o_indptr: Tensor | None = None) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `_**qo_indptr**_`_**：**_Query/Output 张量的索引指针数组（indptr），形状为 `[batch_size + 1]`。用于在 Ragged 连续内存中定位每个请求的 Query 边界。
|
||||
|
||||
* `_**kv_indptr**_`_**：**_Key/Value 张量的索引指针数组，形状为 `[batch_size + 1]`。
|
||||
|
||||
* `_**num_qo_heads**_`_**：**_ Query 和 Output 的注意力头（Attention Heads）数量。
|
||||
|
||||
* `_**num_kv_heads**_`_**：**_Key 和 Value 的注意力头数量。
|
||||
|
||||
* `_**head_dim_qk**_`_**：**_Query 和 Key 张量中每个注意力头的维度大小。
|
||||
|
||||
* `_**head_dim_vo**_`_**：**_Value 和 Output 张量中每个头的维度大小。如果不提供，默认与 `head_dim_qk` 相同。
|
||||
|
||||
* `_**causal**_`_**：**_ 是否对注意力矩阵应用因果掩码（Causal Mask，即屏蔽未来信息）。如果在 `plan()` 中已经提供了自定义的 `mask` 参数，此选项将被忽略。
|
||||
|
||||
* `_**q_data_type**_`_**：**_Query 张量的数据类型，默认为 `torch.float16`。
|
||||
|
||||
* `_**kv_data_type**_`_**：**_Key/Value 张量的数据类型。如果不提供，默认与 `q_data_type` 一致。
|
||||
|
||||
|
||||
注意：`plan()` 方法包含复杂的 Python 层逻辑和动态显存分配，因此**不能**在 CUDA Graph 捕获环境或 `torch.compile` 环境中被追踪或调用。
|
||||
|
||||
1. **.run()** 方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q: Tensor, k: Tensor, v: Tensor, *args, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, enable_pdl: bool | None = None, kv_cache_sf: torch.Tensor | Tuple[torch.Tensor, torch.Tensor] | None = None) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `_**q**_`_**：**_Query（查询）张量。
|
||||
|
||||
* **形状:** `[qo_indptr[-1], num_qo_heads, head_dim_qk]`
|
||||
|
||||
* **解释:** 这里的 `qo_indptr[-1]` 正是我们之前提到的 Ragged Tensor 中所有 Token 数量的总和。它表示把 Batch 里所有的 Query 拍扁到了一个一维的连续维度上。
|
||||
|
||||
* `_**k**_`_**：**_Key（键）张量。
|
||||
|
||||
* **形状:** `[kv_indptr[-1], num_kv_heads, head_dim_qk]`
|
||||
|
||||
* `_**v**_`_**：**_Value（值）张量。
|
||||
|
||||
* **形状:** `[kv_indptr[-1], num_kv_heads, head_dim_vo]`
|
||||
|
||||
|
||||
#### flashinfer.BatchPrefillWithPagedKVCacheWrapper() 类
|
||||
|
||||
参考链接：https://docs.flashinfer.cn/api/attention.html#flashinfer.prefill.BatchPrefillWithPagedKVCacheWrapper
|
||||
|
||||
1. 构造函数
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, qo_indptr_buf: Tensor | None = None, paged_kv_indptr_buf: Tensor | None = None, paged_kv_indices_buf: Tensor | None = None, paged_kv_last_page_len_buf: Tensor | None = None, custom_mask_buf: Tensor | None = None, mask_indptr_buf: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None, jit_kwargs: Dict[str, Any] | None = None) → None
|
||||
```
|
||||
|
||||
参数含义同**BatchPrefillWithRaggedKVCacheWrapper()** 的构造函数参数相同
|
||||
|
||||
2. .Plan() 方法
|
||||
|
||||
|
||||
```Python
|
||||
plan(qo_indptr: Tensor, paged_kv_indptr: Tensor, paged_kv_indices: Tensor, paged_kv_last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim_qk: int, page_size: int, head_dim_vo: int | None = None, custom_mask: Tensor | None = None, packed_custom_mask: Tensor | None = None, causal: bool = False, pos_encoding_mode: str = 'NONE', use_fp16_qk_reduction: bool = False, sm_scale: float | None = None, window_left: int = -1, logits_soft_cap: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, q_data_type: str | dtype = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, non_blocking: bool = True, prefix_len_ptr: Tensor | None = None, token_pos_in_items_ptr: Tensor | None = None, token_pos_in_items_len: int = 0, max_item_len_ptr: Tensor | None = None, seq_lens: Tensor | None = None, seq_lens_q: Tensor | None = None, block_tables: Tensor | None = None, max_token_per_sequence: int | None = None, max_sequence_kv: int | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**o_indptr**``(``_torch.Tensor_``)` – 查询/输出张量的 indptr，形状：`[batch_size + 1]`。
|
||||
|
||||
* `**paged_kv_indptr**``(``_torch.Tensor_``)` – 分页 kv-cache 的 indptr，形状：`[batch_size + 1]`。
|
||||
|
||||
* `**paged_kv_indices**``(``_torch.Tensor_``)` – 分页 kv-cache 的页索引，形状：`[paged_kv_indptr[-1]]`。
|
||||
|
||||
* `**paged_kv_last_page_len**``(``_torch.Tensor_``)` – 分页 kv-cache 中每个请求的最后一页中的条目数，形状：`[batch_size]`。
|
||||
|
||||
* `**num_qo_heads**``(``_int_``)` – 查询/输出头的数量。
|
||||
|
||||
* `**num_kv_heads**``(``_int_``)` – 键/值头的数量。
|
||||
|
||||
* `**head_dim_qk**``(``_int_``)` – 查询/键头的维度。
|
||||
|
||||
* `**page_size**``(``_int_``)` – 分页 kv-cache 中每个页面的大小。
|
||||
|
||||
|
||||
1. run()方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], *args, k_scale: float | None = None, v_scale: float | None = None, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, enable_pdl: bool | None = None, window_left: int | None = None) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* **q** (_torch.Tensor_) – 查询张量，形状：`[qo_indptr[-1], num_qo_heads, head_dim]`
|
||||
|
||||
* **paged\_kv\_cache** (_Union\[torch.Tensor, Tuple\[torch.Tensor, torch.Tensor\]\]_) –存储的分页 KV 缓存，作为张量元组或单个张量
|
||||
|
||||
* 一个元组 `(k_cache, v_cache)`，包含 4D 张量，每个张量的形状为：`[max_num_pages, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。
|
||||
|
||||
* 一个 5D 张量，形状为：`[max_num_pages, 2, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, 2, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。其中 `paged_kv_cache[:, 0]` 是 key 缓存，`paged_kv_cache[:, 1]` 是 value 缓存
|
||||
|
||||
|
||||
#### flashinfer.mla.BatchMLAPagedAttentionWrapper()类
|
||||
|
||||
多头潜在注意力 (MLA) 是一种新的注意力机制，由 [++DeepSeek v2++](https://arxiv.org/abs/2405.04434) 提出，并用于后来的 DeepSeek 模型。MLA 将键缓存和值缓存统一到一个张量中，因此无需单独存储它们。与多头注意力或分组查询注意力相比，MLA 的 KV-Cache 没有 `num_heads` 维度，因此没有像 `NHD` 和 `HND` 布局这样的区别。
|
||||
|
||||
MLA 分离 RoPE（旋转位置编码）维度和其他头部维度。我们使用 `kpe`（带有位置编码的键）和 `ckv`（压缩的键/值）来命名这两个组件。用户可以将它们存储在单个 Paged KV-Cache 中
|
||||
|
||||
```Python
|
||||
head_dim_ckv = 512
|
||||
head_dim_kpe = 64
|
||||
mla_paged_kv_cache = torch.empty(max_num_pages, page_size, head_dim_ckv + head_dim_kpe, dtype=torch.bfloat16)
|
||||
ckv = mla_paged_kv_cache[:, :, :head_dim_ckv] # Slicing here does not copy or move data
|
||||
kpe = mla_paged_kv_cache[:, :, head_dim_ckv:] # Slicing here does not copy or move data
|
||||
```
|
||||
|
||||
**低秩联合压缩 (Joint Compression):** MLA 不再为每个注意力头单独存储巨大的 Key 和 Value 矩阵。相反，它将它们投影并压缩到一个共享的潜在向量（Latent Vector）中，即您代码中的 `ckv` (`head_dim_ckv = 512`)。
|
||||
|
||||
**解耦旋转位置编码 (Decoupled RoPE):** 位置信息对于注意力机制至关重要，但它很难被压缩。MLA 的巧妙之处在于将携带 RoPE 信息的维度单独剥离出来，即代码中的 `kpe` (`head_dim_kpe = 64`)。
|
||||
|
||||
**消除** `**num_heads**` **维度:** 存储的 Cache 不再区分 NHD（序列、头数、头维度）或 HND 布局。无论模型有多少个注意力头，KV-Cache 对于每个 Token 只需要存储 `head_dim_ckv + head_dim_kpe`（例如 512 + 64 = 576 个元素）。
|
||||
|
||||
1. 构造函数
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, use_cuda_graph: bool = False, qo_indptr: Tensor | None = None, kv_indptr: Tensor | None = None, kv_indices: Tensor | None = None, kv_len_arr: Tensor | None = None, backend: str = 'auto') → None
|
||||
```
|
||||
|
||||
2. .Plan()方法
|
||||
|
||||
|
||||
```Python
|
||||
plan(qo_indptr: Tensor, kv_indptr: Tensor, kv_indices: Tensor, kv_len_arr: Tensor, num_heads: int, head_dim_ckv: int, head_dim_kpe: int, page_size: int, causal: bool, sm_scale: float, q_data_type: dtype, kv_data_type: dtype, use_profiler: bool = False) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**qo_indptr**``(``_torch.IntTensor_``)` – 查询/输出张量的 indptr，形状：`[batch_size + 1]`。对于解码注意力，每个查询的长度为 1，张量的内容应为 `[0, 1, 2, ..., batch_size]`。
|
||||
|
||||
* `**kv_indptr**``(``_torch.IntTensor_``)` – 分页 kv-cache 的 indptr，形状：`[batch_size + 1]`。
|
||||
|
||||
* `**kv_indices**``(``_torch.IntTensor_``)` – 分页 kv-cache 的页面索引，形状：`[kv_indptr[-1]]` 或更大。
|
||||
|
||||
* `**kv_len_arr**``(``_torch.IntTensor_``)` – 每个请求的查询长度，形状：`[batch_size]`。
|
||||
|
||||
* `**num_heads**``(``_int_``)` – 查询/输出张量中的头数。
|
||||
|
||||
* `**head_dim_ckv**``(``_int_``)` – 压缩 kv 的头维度。
|
||||
|
||||
* `**head_dim_kpe**``(``_int_``)` – rope k-cache 的头维度。
|
||||
|
||||
* `**page_size**``(``_int_``)` – 分页 kv-cache 的页面大小。
|
||||
|
||||
* `**causal**``(``_bool_``)` – 是否使用因果注意力。
|
||||
|
||||
* `**sm_scale**``(``_float_``)` – softmax 运算的缩放因子。
|
||||
|
||||
* `**q_data_type**``(``_torch.dtype_``)` – 查询张量的数据类型。
|
||||
|
||||
* `**kv_data_type**``(``_torch.dtype_``)` – kv-cache 张量的数据类型。
|
||||
|
||||
* `**use_profiler**``(``_bool, optional_``)` – 是否启用内核内分析器，默认值为 False。
|
||||
|
||||
|
||||
1. run()方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q_nope: Tensor, q_pe: Tensor, ckv_cache: Tensor, kpe_cache: Tensor, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, profiler_buffer: Tensor | None = None, kv_len: Tensor | None = None, page_table: Tensor | None = None, return_lse_base_on_e: bool = False) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**q_nope**``(``_torch.Tensor_``)` – 不含 rope 的查询张量，形状：`[batch_size, num_heads, head_dim_ckv]`。
|
||||
|
||||
* `**q_pe**``(``_torch.Tensor_``)` – 查询张量的 rope 部分，形状：`[batch_size, num_heads, head_dim_kpe]`。
|
||||
|
||||
* `**ckv_cache**``(``_torch.Tensor_``)` – 压缩的 kv-cache 张量（不含 rope），形状：`[num_pages, page_size, head_dim_ckv]`。 `head_dim_ckv` 在 DeepSeek v2/v3 模型中为 512。
|
||||
|
||||
* `**kpe_cache**``(``_torch.Tensor_``)` – kv-cache 张量的 rope 部分，形状：`[num_pages, page_size, head_dim_kpe]`。 `head_dim_kpe` 在 DeepSeek v2/v3 模型中为 64。
|
||||
|
||||
|
||||
#### flashinfer.BatchDecodeWithPagedKVCacheWrapper()类
|
||||
|
||||
1. 构造函数
|
||||
|
||||
```Python
|
||||
__init__(float_workspace_buffer: Tensor, kv_layout: str = 'NHD', use_cuda_graph: bool = False, use_tensor_cores: bool = False, paged_kv_indptr_buffer: Tensor | None = None, paged_kv_indices_buffer: Tensor | None = None, paged_kv_last_page_len_buffer: Tensor | None = None, backend: str = 'auto', jit_args: List[Any] | None = None) → None
|
||||
```
|
||||
|
||||
2. .Plan()方法
|
||||
|
||||
|
||||
```Python
|
||||
plan(indptr: Tensor, indices: Tensor, last_page_len: Tensor, num_qo_heads: int, num_kv_heads: int, head_dim: int, page_size: int, pos_encoding_mode: str = 'NONE', window_left: int = -1, logits_soft_cap: float | None = None, q_data_type: str | dtype | None = 'float16', kv_data_type: str | dtype | None = None, o_data_type: str | dtype | None = None, data_type: str | dtype | None = None, sm_scale: float | None = None, rope_scale: float | None = None, rope_theta: float | None = None, non_blocking: bool = True, block_tables: Tensor | None = None, seq_lens: Tensor | None = None, fixed_split_size: int | None = None, disable_split_kv: bool = False) → None
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**indptr**``(``_torch.Tensor_``)` – 分页 kv 缓存的 indptr，形状：`[batch_size + 1]`，dtype：`torch.int32`
|
||||
|
||||
* `**indices**``(``_torch.Tensor_``)` – 分页 kv 缓存的页面索引，形状：`[kv_indptr[-1]]`，dtype：`torch.int32`
|
||||
|
||||
* `**last_page_len**``(``_torch.Tensor_``)` – 分页 kv 缓存中每个请求的最后一页中的条目数，形状：`[batch_size]`，dtype：`torch.int32`
|
||||
|
||||
* `**num_qo_heads**``(``_int_``)` – 查询/输出头的数量
|
||||
|
||||
* `**num_kv_heads**``(``_int_``)` – key/value 头的数量
|
||||
|
||||
* `**head_dim**``(``_int_``)` – 头部的维度
|
||||
|
||||
* `**page_size**``(``_int_``)` – 分页 kv 缓存的页面大小
|
||||
|
||||
|
||||
1. run()方法
|
||||
|
||||
|
||||
```Python
|
||||
run(q: Tensor, paged_kv_cache: Tensor | Tuple[Tensor, Tensor], *args, q_scale: float | None = None, k_scale: float | None = None, v_scale: float | None = None, out: Tensor | None = None, lse: Tensor | None = None, return_lse: Literal[False] = False, enable_pdl: bool | None = None, window_left: int | None = None) → Tensor
|
||||
```
|
||||
|
||||
* 关键参数
|
||||
|
||||
* `**q**``(``_torch.Tensor_``)` – 查询张量，形状：`[batch_size, num_qo_heads, head_dim]`
|
||||
|
||||
* `**paged_kv_cache**` (_Union\[torch.Tensor, Tuple\[torch.Tensor, torch.Tensor\]\]_) –存储的分页 KV 缓存，作为张量元组或单个张量
|
||||
|
||||
* 一个元组 `(k_cache, v_cache)`，包含 4D 张量，每个张量的形状为：`[max_num_pages, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。
|
||||
|
||||
* 一个 5D 张量，形状为：`[max_num_pages, 2, page_size, num_kv_heads, head_dim]`，如果 `kv_layout` 是 `NHD`，以及 `[max_num_pages, 2, num_kv_heads, page_size, head_dim]`，如果 `kv_layout` 是 `HND`。其中 `paged_kv_cache[:, 0]` 是 key 缓存，`paged_kv_cache[:, 1]` 是 value 缓存
|
||||
|
||||
|
||||
## 六、项目实践--FlashInfer-Baseline
|
||||
|
||||
### Step 1：检查运行环境
|
||||
|
|
@ -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 结果
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
Binary file not shown.
|
Binary file not shown.
|
Binary file not shown.
|
Binary file not shown.
|
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -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}")
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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}")
|
||||
|
|
@ -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
|
||||
|
|
@ -220,15 +220,14 @@ flashInfer 将分页 KV Cache 视作一个**块稀疏矩阵**，并巧妙地
|
|||
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）的参数列表和字典参数。如果提供，框架将会在运行时动态编译底层算子（手动实现算子）；否则，直接使用预编译好的默认算子。
|
||||
|
||||
* 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()`之前
|
||||
|
|
|
|||
Loading…
Reference in New Issue