PagedAttention 与 RadixAttention：上游机制与海光实现对照

概述

大语言模型（LLM）推理服务的吞吐与延迟瓶颈本质上受制于显存容量与显存带宽：

vLLM 的 PagedAttention 通过操作系统分页思想将连续的 KV Cache 离散化为物理块池，借助 block_table 虚拟映射消除了连续显存预分配带来的外部碎片与绝大部分内部碎片，并将跨序列 KV Cache 共享（如并行采样、束搜索）的显存开销降为零；
SGLang 的 RadixAttention 则跳出单请求视角，采用基数树（Radix Tree）以 Token 序列作为动态键管理物理 KV Cache 槽位，结合引用计数 lock_ref 与 LRU 驱逐策略，解决了多轮对话与树搜索中高频前缀的跨请求复用问题；
针对前缀命中后 Q 与 KV 长度不对称的计算场景，SGLang 提出了 Extend Attention（前缀 Prefill 算子），推导了平移因果掩码 $j \le i + (L_k - L_q)$ 并解耦为全可见前缀矩形区与下三角因果区；
海光 DCU 版 FlashAttention 针对底层硬件差异（Wavefront 64、MMAC 矩阵乘指令、LDS 架构等）对这两套机制进行了全方位的本土化重构：在 CUTLASS 侧将 PagedAttention 的 Grid 驱动维度从 Q Head 转向 KV Head 并引入 REUSE_KV_TIMES 模板特化以复用 KV 加载；在 HIPC 侧则采用原生 C++/HIP 替代上游 Triton 重写 Extend Prefill 算子，支持 SGLang 的 bshd 内存池布局并引入 k_delta / v_delta 硬件寄存器差分寻址。

核心要点速览

vLLM PagedAttention 将显存浪费从 60%~80% 压降至 4% 以下：传统静态预分配（按 max_seq_len 如 2048 连续切块）因输出长度不可预测引发严重的内部碎片、内存预留浪费与外部碎片；PagedAttention 通过以固定大小物理块（如 16/32 tokens）按需分配，使显存浪费仅存在于序列末尾块的内部碎片（数学期望为 $B/2$ 个 token）。
Copy-on-Write (CoW) 机制使分支采样在 Fork 阶段显存增量严格为 0：vLLM 在 BlockSpaceManagerV1.fork() 中直接复用父序列的物理块并增加 PhysicalTokenBlock.ref_count；仅在子序列 Decode 写入未满且共享的尾块时（ref_count > 1），才在 append_slots() 触发深拷贝分配。
SGLang 的 RadixAttention 实现了 Token 粒度的跨请求前缀自动合并与精准驱逐：相比 vLLM 早期的块级哈希缓存（要求 Token 数必须严格对齐 block_size，未满整块无法复用），SGLang 的 RadixCache 采用压缩前缀树（Radix Tree），通过 _split_node 在任意 Token 边界分裂节点；驱逐仅作用于 lock_ref == 0 的叶子节点（evictable_leaves），绝不会连带子节点强行驱逐内部节点。
Extend Attention 的因果掩码数学本质是平移不等式 $j \le i + (L_k - L_q)$ ：当利用前缀缓存时，输入只有新增的 $L_q = L_{extend}$ 个 query，而 KV 包含历史缓存与新增 Token，长度为 $L_k = L_{prefix} + L_{extend}$ 。这打破了标准 Prefill（ $L_q = L_k$ ）与 Decode（ $L_q = 1$ ）的假设。
海光 DCU PagedAttention 将 Grid 驱动维度重构为 KV Head，以 GQA 展开消除 DRAM 冗余带宽：上游 vLLM 的 PagedAttention v1/v2 以 (num_heads, num_seqs) 为 Grid 划分，GQA 下多个 Q Head 重复从物理显存加载相同的 KV 块；海光在 csrc/flash_attn_cutlass/src/paged_attention.cu:559-591 中将 Grid 划分为 grid.x = num_kv_heads，由一个 Thread Block 一次性将 KV 加载至 LDS/寄存器，通过模板参数 REUSE_KV_TIMES（支持 4~64 倍）流水计算该组所有 Q Head。
海光 DCU Prefix Prefill 采用硬件缓冲描述符差分寻址（k_delta / v_delta）优化寻址开销：在遍历跨 Page 的 KV Cache 时，上游代码在每个循环步重新从页表计算 64 位全局地址；海光 HIPC 算子计算页表差值 table_diff 与偏移差值 offset_diff，直接向 buffer resource 寄存器指针做增量累加（*(int64_t*)&gK += k_delta * sizeof(Element)），显著缩减了 VGPR 占用与指令条数。

第一部分：vLLM 的 PagedAttention

1.1 问题定义与显存浪费剖析

在 Transformer 解码（Decode）阶段，自注意力机制需要访问自回归生成的所有历史 Token 的 Key 和 Value 向量。若当前序列长度为 $t$ ，第 $t$ 步的计算公式为： $\text{Attention}(q_t, K_{1:t}, V_{1:t}) = \text{softmax}\left(\frac{q_t K_{1:t}^T}{\sqrt{d_k}}\right) V_{1:t}$ 为了避免每一步都对历史 Token 重新计算投影，服务系统必须将历史生成的 Key 和 Value 保存在 GPU 显存中，这就是 KV Cache。对于一个参数规模为 $13\text{B}$ 的典型模型（如 LLaMA-13B，层数 $L=40$ ，Hidden Size $H=5120$ ，FP16 存储），每个 Token 在每一层占用字节数为： $\text{Bytes per token} = 2 \times 2 \times H \times L = 4 \times 5120 \times 40 = 819,200 \text{ Bytes} \approx 800 \text{ KB}$ 当并发请求达到数十甚至数百时，KV Cache 的显存开销迅速超越模型权重本身，成为系统吞吐最核心的制约瓶颈。

传统做法的痛点与显存浪费三要素

在 vLLM（SOSP 2023）出现之前，主流推理系统（如 FasterTransformer、Orca）要求张量在物理显存中保持行优先连续排列。由于自回归生成过程中模型何时输出终止符 [EOS] 是未知的，系统在接收到请求的瞬间，必须根据模型的最大上下文长度（ $max\_seq\_len$ ，例如 2048 或 4096）一次性开辟一块物理上完全连续的显存空间。

这种静态连续分配模式造成了极度严重的显存浪费，具体分为三种形式：

+-----------------------------------------------------------------------------------------------+
|  已处理 Prompt (如 100 tokens)  |  当前已生成 (如 50 tokens)  |        未使用的预留槽位       |
+-----------------------------------------------------------------------------------------------+
|<---------------- 实际占用显存 (150 tokens) ----------------->|<--- 内存预留浪费 / 内部碎片 --->|
|<--------------------------------- max_seq_len (2048 tokens) --------------------------------->|

内存预留浪费（Reservation Waste / Over-provisioning）：
- 机制：系统为最坏情况（请求一路生成至 $max\_seq\_len$ ）预留显存，但这部分显存在当前时间步并未写入任何数据。
- 具体算例：设 $max\_seq\_len = 2048$ 。一个请求的 Prompt 长度为 100，当前生成了第 50 个 token。该请求实际只使用了 150 个 token 的空间，但系统必须锁定 2048 个 token 的显存。在此刻，有 1898 个 token（占分配量的 $92.7\%$ ）的显存虽然被空置，却无法被其他并发请求使用。
内部碎片（Internal Fragmentation）：
- 机制：当该请求在第 200 个 token 生成了 [EOS] 结束解码时，系统释放这块空间。但在请求运行的整个生命周期中，分配的 2048 槽位中只有 200 槽位实际生效，剩下的 1848 槽位就是内部碎片。
- 具体算例：单次生命周期内显存有效利用率仅为 $200 / 2048 \approx 9.77\%$ ，内部碎片浪费率高达 $90.23\%$ 。
外部碎片（External Fragmentation）：
- 机制：由于每个请求的实际生成长度不同，生命周期也不同，连续内存分配器（如 cudaMalloc 底层的伙伴系统或 Best-fit 堆）在频繁的分配与释放后，显存池会被切割得千疮百孔，形成大量分散的空闲内存孔隙（Holes）。
- 具体算例：假设显存中剩余总空闲容量为 $8\text{ GB}$ ，但由 4 段分散的 $2\text{ GB}$ 非连续区间构成。此时系统进驻一个需要连续占用 $3\text{ GB}$ 的新请求，连续内存分配器将因找不到足够大的单一连续空间而直接报 OOM（Out Of Memory），即便总空闲容量远超请求需求。
无法共享显存（Inability to Share Memory）：
- 机制：在并行采样（Parallel Sampling，单 Prompt 生成 $N$ 个输出）或束搜索（Beam Search）中，各分支在 Prompt 阶段生成的 KV Cache 是一模一样的。但由于传统系统要求每个序列独占连续空间，必须将 Prompt 的 KV Cache 硬拷贝 $N$ 份。
- 具体算例：设 Prompt 长度为 1024 tokens，采用 Beam Search 宽度 $N=4$ 。传统做法必须分配 $4 \times 1024 = 4096$ 个 token 的 KV 显存；而理想情况下只需保存一份 1024 tokens 的只读前缀。

论文公布的浪费比例与出处

根据 vLLM 创始团队发表于操作系统顶会 SOSP 2023 的论文 《Efficient Memory Management for Large Language Model Serving with PagedAttention》（Woosuk Kwon et al., UC Berkeley，arXiv:2309.06180）：

在 FasterTransformer 和 Orca 等系统中，显存中实际存储有效 Token 数据的比例仅占 $20.4\% \sim 38.2\%$ ；
换言之，系统有 $61.8\% \sim 79.6\%$ 的 GPU 显存被内部碎片、预留浪费与外部碎片浪费；
引入 PagedAttention 将物理块离散化管理后，外部碎片被彻底消除为 $0$ ，内部碎片被严格局限在序列的最后一个物理块内部，整体显存浪费降至 $4\%$ 以下。在相同硬件配置下，由于有效并发 Batch Size 提升了数倍，整体服务吞吐量提升了 $2\times \sim 4\times$ 。

1.2 数据结构详解

PagedAttention 的核心是虚拟内存分页机制向 GPU 显存的映射：在逻辑上，一个序列拥有一个不断增长的连续 Token 空间（逻辑块 Logical Blocks）；在物理上，KV 向量存储在预先开辟的固定大小非连续块（物理块 Physical Blocks）中。连接二者的是页表（Block Table）。

[逻辑序列 (Logical Sequence)]
+--------------+--------------+--------------+
| Logical Blk 0| Logical Blk 1| Logical Blk 2|   (Token 0 ~ Token 47, 设 BlockSize=16)
+--------------+--------------+--------------+
        |              |              |
        | block_table 映射            |
        v              v              v
[物理显存池 (Physical KV Pool)]
+-----------+  +-----------+  +-----------+  +-----------+
| PhysBlk 0 |  | PhysBlk 7 |  | PhysBlk 3 |  | PhysBlk 12| ...
+-----------+  +-----------+  +-----------+  +-----------+
(Block 0)      (Block 1)      (Block 2)

1. KV Cache 物理池（Physical KV Cache Pool）

张量形状：
- 上游 vLLM 原生布局：
  - key_cache: [num_blocks, num_kv_heads, head_size // x, block_size, x]，其中 $x = 16 / \text{sizeof(scalar\_t)}$ （FP16 下 $x=8$ ）。
  - value_cache: [num_blocks, num_kv_heads, head_size, block_size]。
- 海光 DCU / CUTLASS 布局（见 csrc/flash_attn_cutlass/src/paged_attention.cu:40-41）：
  - key_cache: [num_blocks, num_heads, block_size, head_size]
  - value_cache: [num_blocks, num_heads, head_size, block_size]（部分算子支持 [num_blocks, block_size, num_heads, head_size] 即 bshd 布局，详见后文）。
物理排列机理与向量化访存：为什么 value_cache 的内层维度排布为 [head_size, block_size] 而不是 [block_size, head_size]？在计算注意力矩阵乘值向量 $O = P \cdot V$ 时，同一 Head 内，多个上下文 Token（沿 block_size 方向）的值需要与对应 Token 的注意力权重标量相乘并累加。将 block_size 置于最内层连续维度，使得单个线程或线程束在读取连续多个 Token 的同一维度时，可以通过 128-bit 向量化访存（如 float4、half8）一次性完成读取，最大化 GPU 显存总线带宽利用率并规避 Bank Conflict。
生命周期与分配时机：物理池在服务启动（Server Initialization）时一次性分配。系统首先加载模型权重，测算出剩余空闲 GPU 显存，乘以预设利用率系数（如 gpu_memory_utilization=0.9），将剩余显存划分为大小恒定的 num_blocks 个物理块，生命周期与推理 Engine 进程完全一致。

2. 块表（`block_table`）

张量形状：[num_seqs, max_num_blocks_per_seq]。
数据类型：torch.int32（C++ 侧为 const int* __restrict__ block_tables）。
含义：每一行代表一个正在运行的序列，第 i 个元素表示该序列逻辑块 i 所映射到的物理块池全局索引（Physical Block Number）。
持有者与填充时机：由 CPU 侧的调度器调度核心 BlockSpaceManager 持有。每次调度循环执行 step() 时，调度器将当前 Batch 内所有就绪请求的块表收集为一个 2D Tensor，通过 PyTorch C++ 扩展传入 CUDA/HIP kernel。
空槽填充值：若某个序列的当前块数不足 max_num_blocks_per_seq，未分配的空槽填充为 0 或 -1（通常填 0，但在 Kernel 内部依据 seq_lens 计算的有效块数截止，绝不会越界访问空槽）。

3. `BlockAllocator` 与 `BlockManager`

参考 vLLM 核心源码 vllm/core/block_manager_v1.py：

核心数据结构： PhysicalTokenBlock：代表物理块实体。

class PhysicalTokenBlock:
    device: Device
    block_number: int       # 物理块全局 ID (0 <= ID < num_blocks)
    block_size: int         # 块包含的 Token 数 (如 16)
    ref_count: int          # 物理块当前的共享引用计数
    last_accessed: float    # 最后访问时间戳
    computed: bool          # 是否已完成计算 (用于 Prefix Caching)

空闲块链表维护： UncachedBlockAllocator 内部维护 free_blocks: List[PhysicalTokenBlock]。初始化时所有物理块推入栈中；分配时执行 heappop() 或 pop()；释放时，当且仅当 ref_count == 0，才执行 append() 放回空闲链表。

引用计数（ref_count）与写时复制（Copy-on-Write, CoW）：

Fork 触发：在 Beam Search 或 Parallel Sampling 中，多个子序列共享父序列历史上下文。调用 BlockSpaceManagerV1.fork(parent_seq, child_seq)（vllm/core/block_manager_v1.py:431）：
```
def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
    src_block_table = self.block_tables[parent_seq.seq_id]
    self.block_tables[child_seq.seq_id] = src_block_table.copy()
    # 引用计数增加
    for block in set(src_block_table):
        block.ref_count += 1
```
子序列直接复制父序列的 block_table，所有物理块的 ref_count 递增。此时显存开销为零，没有发生任何物理显存拷贝。

CoW 触发：当某个分支子序列生成新 Token 调用 append_slots() 时（vllm/core/block_manager_v1.py:385）：

last_block = block_table[-1]
if last_block.ref_count == 1:
    # 独占块，直接安全追加新 Token 的 KV
    return {}
else:
    # 存在分支共享！触发写时复制
    new_block = self._allocate_last_physical_block(seq)
    block_table[-1] = new_block
    self.gpu_allocator.free(last_block)  # last_block.ref_count 减 1
    # 返回拷贝指令字典，通知 GPU 执行真实显存复制
    return {last_block.block_number: [new_block.block_number]}

GPU 侧会在正式 Attention 计算前触发 copy_blocks kernel，将旧尾块的内容复制到新物理块中，之后子序列独占新尾块并写入新 Token，分支解耦完成。

4. `seq_lens`（或 `context_lens`）

张量形状：[num_seqs]，torch.int32。
与 block_table 的映射关系：对于序列 seq_idx，其当前完整的有效 Token 数为 $L = seq\_lens[seq\_idx]$ 。该序列在 block_table 中有效占据的逻辑块总数为： $num\_seq\_blocks = \left\lceil \frac{L}{block\_size} \right\rceil = \frac{L + block\_size - 1}{block\_size}$ 在 Kernel 内部，外层块循环的上界被严格限制为 num_seq_blocks，而最后一个块内有效的 Token 偏移上限为 $(L - 1) \pmod{block\_size} + 1$ 。

1.3 寻址推导

定义推导目标：给定请求编号 seq_idx、逻辑 Token 位置索引 t（ $0 \le t < seq\_lens[seq\_idx]$ ）、Head 编号 h（ $0 \le h < num\_heads$ ）、Head 内部特征维度标量偏移 d（ $0 \le d < head\_size$ ）。设物理块大小为 $B = block\_size$ ，Head 维度大小为 $D = head\_size$ 。每个元素的数据类型大小为 $S = \text{sizeof}(scalar\_t)$ （FP16/BF16 下 $S=2$ 字节）。

第一步：逻辑位置解构

逻辑块索引（Logical Block Index）： $block\_idx = \lfloor t / B \rfloor$
块内 Token 偏移（Block Offset）： $block\_offset = t \pmod B$
查询物理块号（Physical Block Number）： $P = block\_tables[seq\_idx][block\_idx]$

第二步：不同内存布局下的字节偏移推导

布局 A：`bshd` 布局（以 SGLang 内存池 / 海光 `layout=1` 为例）

张量整体逻辑形状为：[num_blocks, block_size, num_heads, head_size]

各维度 Stride 计算（以元素为单位）：
- $Stride_{head\_size} = 1$
- $Stride_{num\_heads} = D$
- $Stride_{block\_size} = H \times D$
- $Stride_{num\_blocks} = B \times H \times D$
元素线性索引公式： $\text{Index}_{bshd} = P \times (B \times H \times D) + block\_offset \times (H \times D) + h \times D + d$
物理显存绝对字节偏移： $\text{ByteOffset}_{bshd} = \left[ P \cdot B \cdot H \cdot D + block\_offset \cdot H \cdot D + h \cdot D + d \right] \times S$

布局 B：`bhsd` 布局（以标准 PagedAttention K-Cache / 海光 `layout=0` 为例）

张量整体逻辑形状为：[num_blocks, num_heads, block_size, head_size]

各维度 Stride 计算（以元素为单位）：
- $Stride_{head\_size} = 1$
- $Stride_{block\_size} = D$
- $Stride_{num\_heads} = B \times D$
- $Stride_{num\_blocks} = H \times B \times D$
元素线性索引公式： $\text{Index}_{bhsd} = P \times (H \times B \times D) + h \times (B \times D) + block\_offset \times D + d$
物理显存绝对字节偏移： $\text{ByteOffset}_{bhsd} = \left[ P \cdot H \cdot B \cdot D + h \cdot B \cdot D + block\_offset \cdot D + d \right] \times S$

布局 C：vLLM 专为向量化转置排布的 V-Cache 布局

张量整体逻辑形状为：[num_blocks, num_heads, head_size, block_size]

元素线性索引公式： $\text{Index}_{V\_trans} = P \times (H \times D \times B) + h \times (D \times B) + d \times B + block\_offset$
物理显存绝对字节偏移： $\text{ByteOffset}_{V\_trans} = \left[ P \cdot H \cdot D \cdot B + h \cdot D \cdot B + d \cdot B + block\_offset \right] \times S$

1.4 Kernel 怎么跑：v1、v2 Split-KV 与 Combine 归约

1. PagedAttention v1 的执行逻辑

Grid 划分：2D 结构 gridDim = (num_heads, num_seqs, 1)。
Block 负责范围：每个 CUDA Thread Block（通常分配 128 或 256 个线程，即 4 或 8 个 Warps）独占处理单个序列的一个 Q Head 在整个历史上下文上的注意力计算。
线程组（Thread Group）协作与点积：
- 为实现 16 字节对齐访问，多个线程组成一个 Thread Group（例如 FP16 下 2 个线程一组，每线程读取 4 个 half，合力完成 8 维点积）。
- 外层循环：遍历该序列关联的所有逻辑块 $block\_idx \in [0, num\_seq\_blocks - 1]$ 。从 block_table 读取对应的物理块号 $P$ 。
- 内层计算：加载 $Q$ （常驻寄存器或共享内存）与物理块中的 $K$ ，计算点积并乘以缩放因子 $\frac{1}{\sqrt{d_k}}$ ，存入共享内存中的 logits 数组。
Softmax 块内归约：
1. 第一阶段：每个 Warp 内部通过 Shuffle 指令（__shfl_xor_sync）做树状规约，得出 Warp 局部最大值。
2. 第二阶段：Warp 0 将所有 Warp 的局部最大值写入共享内存并规约，得出该 Thread Block 的全局最大标量 $m = \max(logits)$ 。
3. 第三阶段：广播 $m$ ，所有线程并行计算 $\exp(logits_t - m)$ 并累加求和得 $\ell = \sum \exp(logits_t - m)$ 。
4. 第四阶段：归一化更新 $logits_t \leftarrow \frac{\exp(logits_t - m)}{\ell}$ 。
乘 $V$ 并写回：遍历各物理块的 $V$ ，读取转置排布的 $V$ 矩阵与共享内存中的归一化权重做向量点积，归约后写回输出显存 out[seq_idx, head_idx, :]。
长序列下 v1 的瓶颈剖析：在推理服务的高并发场景中，往往包含长上下文（如 8K、32K、128K）。若此时并发 Batch Size 较小（例如 Batch Size = 1 或 2），整个 Grid 仅有 $BatchSize \times NumHeads$ 个 Block。例如 7B 模型 $H_q=32$ ，Batch Size=1 时只有 32 个 Block。而在现代 GPU（拥有上百个计算单元）上，大部分计算核心处于空闲状态（Under-occupancy）。与此同时，这 32 个 Block 必须在自己的 SM 上循环遍历长达数千个 Block 的 KV 缓存，无法利用其他闲置计算核心的能力与显存带宽。

2. PagedAttention v2 的 Split-KV 机制

为了在长序列小 Batch 场景下打满 GPU，vLLM 设计了 PagedAttention v2，引入序列切分（Split-KV / Partitioning）：

PARTITION_SIZE：沿上下文 Token 维度切分的步长阈值（通常设定为 256 或 512 个 Token）。
3D Grid 结构： $\text{gridDim} = (num\_heads, num\_seqs, max\_num\_partitions)$ 其中 $max\_num\_partitions = \lceil max\_seq\_len / PARTITION\_SIZE \rceil$ 。通过在 Z 轴引入分区维度，原本单个长序列的单 Head 任务被拆分给数十个 Thread Block 在不同的计算核心上并行执行，彻底释放硬件占用率。
Partition 局部计算产物：第 $p$ $p$ 个 Partition 的 Block 仅负责计算属于该区间的 Token 子集，其计算结果是局部的，写入临时显存缓冲区 out_tmp：
1. 局部最大值 $m_p \in \mathbb{R}$ ：该分区范围内的最大 Logit。
2. 局部指数和 $\ell_p \in \mathbb{R}$ ： $\ell_p = \sum_{t \in \text{partition } p} \exp(S_t - m_p)$ 。
3. 局部输出向量 $O_p \in \mathbb{R}^{head\_size}$ ： $O_p = \sum_{t \in \text{partition } p} \frac{\exp(S_t - m_p)}{\ell_p} V_t$ 。
临时缓冲区内存占用：系统在全局显存中分配临时空间：
- exp_sums: 形状 [num_seqs, num_heads, max_num_partitions]，FP32。
- max_logits: 形状 [num_seqs, num_heads, max_num_partitions]，FP32。
- tmp_out: 形状 [num_seqs, num_heads, max_num_partitions, head_size]，FP16/BF16。 (见海光实现 csrc/flash_attn_cutlass/src/paged_attention.cu:491：temp_out_size 在主机端一次性开辟 200MB~500MB 全局内存作为重用缓冲)。

3. Combine Kernel 与在线 Softmax 重缩放（Online Softmax Rescaling）推导

当所有的 Partition Block 计算完毕后，主线程启动轻量级的 paged_attention_combine kernel（见 csrc/flash_attn_cutlass/src/paged_attention.h:337-404）。其 Grid 为 2D (num_heads, num_seqs)，Block 内只需 1 个 Warp（32 或 64 线程）。

完整数学推导

设一个序列划分为 $P$ 个分区。每个分区输出了局部统计量 $(m_p, \ell_p, O_p)$ ，其中： $O_p = \frac{\sum_{j \in Part_p} \exp(S_j - m_p) V_j}{\ell_p}, \quad \ell_p = \sum_{j \in Part_p} \exp(S_j - m_p)$ 现在要恢复出真实的全局自注意力输出 $O_{global}$ ： $O_{global} = \frac{\sum_{p=1}^P \sum_{j \in Part_p} \exp(S_j - M) V_j}{L_{global}}, \quad M = \max_{1 \le k \le P} m_k, \quad L_{global} = \sum_{p=1}^P \sum_{j \in Part_p} \exp(S_j - M)$

将局部指数和转化为全局基准： $\exp(S_j - M) = \exp(S_j - m_p + m_p - M) = \exp(S_j - m_p) \cdot \exp(m_p - M)$ 对属于分区 $p$ 的所有元素求和： $\sum_{j \in Part_p} \exp(S_j - M) = \left( \sum_{j \in Part_p} \exp(S_j - m_p) \right) \cdot \exp(m_p - M) = \ell_p \cdot \exp(m_p - M)$ 定义分区重缩放因子（Rescaled Weight）： $\ell'_p = \ell_p \cdot \exp(m_p - M)$ 则全局总分母为： $L_{global} = \sum_{p=1}^P \ell'_p = \sum_{p=1}^P \ell_p \cdot \exp(m_p - M)$ 分子部分的合并： $\sum_{j \in Part_p} \exp(S_j - M) V_j = \exp(m_p - M) \sum_{j \in Part_p} \exp(S_j - m_p) V_j = \exp(m_p - M) \cdot (\ell_p O_p) = \ell'_p \cdot O_p$ 代入 $O_{global}$ 表达式中： $O_{global} = \frac{\sum_{p=1}^P \ell'_p O_p}{L_{global}} = \sum_{p=1}^P \left( \frac{\ell'_p}{L_{global}} \right) O_p$ 这就是在线 Softmax 的分区重缩放公式：各个分区的局部向量输出 $O_p$ ，只需乘以加权系数 $w_p = \frac{\ell_p \exp(m_p - M)}{L_{global}}$ ，求和后即可完全等价于全局未经切分的精确 Softmax 输出。

为什么底层硬件采用 `exp2` 而非 `exp`

在 GPU 体系结构中，通用浮点单元计算自然指数 $\exp(x) = e^x$ 需要调用多项式近似拟合，指令周期漫长。而 GPU 的特殊功能单元（SFU，Special Function Unit）内部拥有硬件级以 2 为底的指数指令（NVIDIA 的 __exp2f，海光/AMD 的 __builtin_amdgcn_exp2f），吞吐极高且单指令即可完成。利用恒等式： $e^x = 2^{x \cdot \log_2(e)}$ 其中常数 $\log_2(e) \approx 1.4426950408889634$ 。 vLLM 与海光实现巧妙地在主机端发起 Kernel 时，将 Softmax Scale 事先乘上 $\log_2(e)$ ： $\text{scale}_{internal} = \frac{1}{\sqrt{d_k}} \times 1.4426950408889634$ (见 csrc/flash_attn_cutlass/src/paged_attention.cu:60)。这样在 Kernel 内部算出的点积 $S_t$ 直接就是以 2 为底指数的幂次，后续计算直接调用快速单周期指令 exp2f(S_t - m)，大幅消减了指令延迟。

4. GQA（Grouped-Query Attention）的 Kernel 优化

在 GQA 架构中，多个 Query Head（例如 8 个）共享同一组 Key-Value Head。

朴素做法的缺陷：若 Grid 以 Q Head 展开（ $gridDim.x = num\_q\_heads$ ），则同属一个 KV 组的 8 个 Q Head 分别处于 8 个独立的 Thread Block 中。这 8 个 Block 必须各自通过显存总线从全局 DRAM 中读取完全一致的物理 KV 块，导致全局内存读取带宽放大 8 倍。
GQA 融合展开：先进实现与海光 DCU 将 Grid 的驱动维度修改为 KV Head（ $gridDim.x = num\_kv\_heads$ ）。由一个 Thread Block 统一读取物理 KV 块并放入 Shared Memory / 寄存器中，在块内部利用多组寄存器同时与多个 Query Head 计算点积（见海光实现中的 REUSE_KV_TIMES 机制），将 DRAM 的 KV 读取开销降低至理论下限。

1.5 参数权衡与选型依据

1. `block_size`（16, 32, 64, 128）的四维权衡

`block_size` 大小	内部碎片浪费	显存事务效率（Memory Transaction）	页表（`block_table`）规模	Wavefront / Warp 对齐开销
16	极低（期望浪费 8 tokens $\approx 6.4\text{ KB}$ ）	较差。序列长度较短时，物理块不连续导致无法做大跨度 Burst Read。	很大。4K 序列需要 256 个块索引，增加 CPU 调度与 TLB 负担。	在 32 线程 Warp 下勉强对齐；在 64 线程 Wavefront 下严重撕裂。
32	低（期望浪费 16 tokens $\approx 12.8\text{ KB}$ ）	平衡。完美契合 NVIDIA 32 线程 Warp（每线程读 1 token 时对齐）。	中等。vLLM 在 NVIDIA GPU 上的推荐默认值。	在 64 线程架构下仍需跨步合并。
64	中等（期望浪费 32 tokens $\approx 25.6\text{ KB}$ ）	优异。单次突发传输能抓取 64 tokens 数据，总线饱和度高。	较小。页表长度减半，调度处理快。	完美契合 AMD / 海光 DCU 的 Wavefront 64 架构。
128	偏高（期望浪费 64 tokens $\approx 51.2\text{ KB}$ ）	极高。完全满足所有矢量化加载管线，TLB Cache 命中率最高。	极小。在极长上下文（100K+）中优势显著。	适合 Prefill 算子，但在小请求并发时尾部碎片略大。

2. `PARTITION_SIZE`（256 vs 512）的权衡

取 256：
- 优势：划分粒度更细，在短序列（512~2048）下就能迅速将任务打散至大量计算单元，硬件利用率与占用率（Occupancy）迅速爬升。
- 劣势：生成的 Partition 数量翻倍，out_tmp 临时缓冲区内存占用翻倍，后置 combine kernel 需要合并的局部量增加，规约耗时上升。
取 512：
- 优势：Combine Kernel 的规约输入较少，计算密集度更高；适合超长序列（8K~32K），减少临时内存占用。
- 劣势：在 1K 上下文的中等请求中，无法提供足够的并发 Block 数来喂饱大型 GPU。

第二部分：SGLang 的 RadixAttention

2.1 问题定义：跨请求前缀复用

PagedAttention 成功解决了单请求内部动态生成的显存管理与显存共享（通过 CoW 支持单请求内的 Fork 分支），但它完全没有跨请求生命周期的记忆能力。当一个请求结束服务后，其占用的物理块全部被释放回空闲链表。

真实生产场景的痛点与算力浪费

多轮对话（Multi-turn Chat）：
- 设第 1 轮 Prompt（500 tokens），模型回复（200 tokens）。
- 第 2 轮用户提问时，发送给模型的输入必须拼接历史：[Prompt 1, Resp 1, Prompt 2]，总长 $500 + 200 + 100 = 800$ tokens。
- 痛点：前 700 个 token 的 Key 和 Value 在第 1 轮已经算过，但在第 2 轮必须从第 0 层到最后一层全部重新做一次 Prefill 投影计算！显存中也重新分配了 700 tokens 的物理块。
少样本与系统提示（Few-shot Prompting / System Prompt）：
- 系统级指令（如详尽的角色设定、API 说明文档、Code 工具示例）常达 $2000 \sim 4000$ tokens。
- 当 1000 个不同的用户向该系统并发提问时，这 3000 tokens 的固定前缀被重复 Prefill 计算了 1000 次！
- 浪费评估： $1000 \text{ reqs} \times 3000 \text{ tokens} = 3,000,000$ 冗余 Prefill Tokens，GPU 被算力密集但无意义的矩阵乘淹没，服务响应延迟急剧恶化。
Agent 树搜索与工具调用（Tree Search / ReAct）：
- 思考链（Chain of Thought）或 MCTS 探索分支时，数十个分支频繁共享前半截相同的推理轨迹。

vLLM 早期块级哈希缓存（Block-level Hash Caching）与 Radix Tree 的本质差异

vLLM 在 0.4.x 之前尝试通过块哈希（CachedBlockAllocator）实现前缀缓存：对逻辑块内的 Token 计算哈希值 $H_k = \text{hash}(H_{k-1}, \text{tokens in block } k)$ 。

硬伤 1：块边界对齐限制（Block Boundary Constraint）：若 block_size = 16，用户前缀长度为 45 tokens。块级哈希只能将前两块（32 tokens）成功缓存并哈希；剩下的 13 个 token 因不足 16 个无法形成完整物理块，系统无法为未填满的块生成确定性哈希，导致这 13 个 token 无法复用，必须重新 Prefill。
硬伤 2：扁平哈希表无法表达层次关系：扁平 Hash Map 将各个块孤立存储，丢失了前缀序列的父子层级关系，无法感知分支分叉点，导致复杂前缀树结构下的 LRU 驱逐策略难以保证整体最优。

SGLang 提出的 RadixAttention 则抛弃了块对齐哈希，采用基数树（Radix Tree / 压缩前缀树），以 Token 序列片段为边，支持任意 Token 粒度的精准匹配、分裂与物理槽位复用。

2.2 Radix Tree 数据结构详解

参考 SGLang 源码 python/sglang/srt/mem_cache/radix_cache.py：

                    [Root Node] (lock_ref=0, value=None)
                     /        \
          "Hello "  /          \  "Translate "
                   v            v
              [Node A]        [Node B]
             /        \
   "world!" /          \ "there!"
           v            v
       [Node C]       [Node D]

1. 核心节点数据结构 `TreeNode`

源码对应位置：python/sglang/srt/mem_cache/radix_cache.py:259

class TreeNode:
    def __init__(self, id: Optional[int] = None, priority: int = 0):
        self.children: Dict[RadixKey, TreeNode] = defaultdict(TreeNode)
        self.parent: Optional[TreeNode] = None

        # 键：该树边所代表的 Token ID 列表片段（RadixKey 类型，支持切片与高效匹配）
        self.key: Optional[RadixKey] = None

        # 值：该段 Token 在底层物理显存池（TokenToKVPool）中对应的全局物理槽位索引 Tensor
        self.value: Optional[torch.Tensor] = None

        # 核心保护锁：正在引用该节点的活跃请求计数（Active Request Lock Reference）
        self.lock_ref: int = 0

        # LRU 时间戳：记录该节点最后一次被匹配访问或修改的单调时间
        self.last_access_time: float = time.monotonic()

        # 命中统计与优先级
        self.hit_count: int = 0
        self.priority: int = priority
        self.id: int = TreeNode.counter

2. 为什么选择 Radix Tree 而非标准字典树（Trie）

标准 Trie 的内存与遍历开销：在标准 Trie 中，每个节点只能存一个单一 Token。一个长 2048 的 Prompt 需要在 CPU 内存中构建 2048 个独立的 Python 对象节点，产生 2048 次指针跳转（Pointer Chasing）与哈希查找，CPU 开销极大。
Radix Tree（压缩前缀树）的边合并：若一段前缀序列中没有分支，Radix Tree 会将连续的多个 Token 合并到一个节点（Edge Compression）。一个 2000 tokens 的 System Prompt 在树中只表现为单一的根子节点，其 key 为长度 2000 的 Tensor，value 为对应长度 2000 的物理槽位数组。匹配时单步即可完成千级 Token 的比对，内存极小，查找极快。

3. 节点分裂（Split）算法

源码对应位置：python/sglang/srt/mem_cache/radix_cache.py:727 当新请求的 Token 序列与树中某个已有节点 child 的前缀匹配，但在节点中间发生分叉时，调用 _split_node(key, child, split_len)：

分裂前:
  [Parent] -------- key: "ABCDEF", value: [0,1,2,3,4,5] --------> [Child]

新请求传入 "ABCXYZ" (匹配 split_len = 3, 即 "ABC"):

分裂后:
  [Parent] -------- key: "ABC", value: [0,1,2] --------> [NewNode] (承接 Parent)
                                                          /       \
               key: "DEF", value: [3,4,5]                /         \  新边插入
                                                        v           v
                                                     [Child]    [新请求剩余节点]

核心 Python 实现精析：

def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode:
    # 1. 创建新中间节点，继承子节点的 priority 与 hit_count
    new_node = TreeNode(priority=child.priority)
    new_node.hit_count = child.hit_count
    new_node.parent = child.parent
    new_node.lock_ref = child.lock_ref  # 继承活跃锁状态

    # 2. 将原 child 的前半截赋予 new_node
    new_node.key = child.key[:split_len]
    new_node.value = child.value[:split_len].clone()

    # 3. 将原 child 的后半截截断保留在 child
    child.key = child.key[split_len:]
    child.value = child.value[split_len:].clone()

    # 4. 重新建立父子拓扑连接
    child.parent = new_node
    new_node.children = {child.key.child_key(self.page_size): child}
    new_node.parent.children[key.child_key(self.page_size)] = new_node

    return new_node

2.3 匹配与分配流程

当调度器接收到一个新请求 $Req$ （Token 列表为 $T$ ）时，全流程如下：

遍历树匹配最长前缀（Prefix Matching）：从 root_node 开始执行 _match_prefix_helper(root, key)：沿子分支循环匹配 Token。若遇到完全匹配的节点，则继续向深层探索；若匹配终止于某个节点的内部，则调用 _split_node 将该节点就地分裂，分裂产生的新节点成为最长前缀终点。
计算命中与未命中切分：
- 命中长度： $L_{hit}$ （若完全未命中则 $L_{hit} = 0$ ）。
- 命中物理槽位：收集从根到终点路径上所有节点 value 拼接而成的 Tensor $V_{cached}$ 。
- 未命中部分（Extend Tokens）： $T_{extend} = T[L_{hit} : ]$ 。
加锁保护（Increment Lock Reference）：调用 self.inc_lock_ref(terminal_node)（radix_cache.py:645）：从命中终点节点开始，沿 parent 指针一路向上回溯至根节点，对路径上每一个节点的 lock_ref 统一执行 lock_ref += 1。 这确保了当其他请求在并发分配显存时，正在运行的请求所依赖的前缀节点绝对不会被后台驱逐程序释放。
未命中部分物理显存分配：从底层内存池（TokenToKVPool）开辟长度为 $len(T_{extend})$ 的空闲物理槽位数组 $V_{new}$ 。
产出元数据并送入 Kernel：此时该请求的整体 KV 物理槽位为 torch.cat([V_{cached}, V_{new}])。模型只需对 $T_{extend}$ 计算 Query，随后调用下一章详述的 Extend Attention 算子，在无需对前 $L_{hit}$ 个 token 进行重算的前提下，直接读取 $V_{cached}$ 对应的历史 KV。
请求结束时插入树中：当请求完成 Prefill 或生成结束时，新生成的 Token 及其对应的物理槽位 $V_{new}$ 被封装为一个新叶子节点插入树中，同时沿路径递减 lock_ref（dec_lock_ref）。如果该请求后续成为其他请求的前缀，该物理显存块将一直保留在树中直至触发 LRU 驱逐。

2.4 驱逐机制（Eviction）

当显存池耗尽，新的请求无法分配空闲物理块时，RadixCache 必须从树中腾退已缓存但当前处于闲置状态的历史节点。

1. 锁保护（`lock_ref > 0`）与绝对安全边界

只有 node.lock_ref == 0 的节点才被定义为“可驱逐（Evictable）”。
只要系统中仍有任何一个请求处于 Prefill 或 Decode 状态，其路径上所有祖先节点的 lock_ref \ge 1。这些节点处于物理锁定状态，驱逐扫描程序对其绝对不可见，从数据结构层面杜绝了悬垂指针（Dangling Pointer）引发显存访问越界或数据踩踏。

2. 驱逐顺序与优先队列

源码对应位置：python/sglang/srt/mem_cache/radix_cache.py:615

def evict(self, params: EvictParams) -> EvictResult:
    leaves = list(self.evictable_leaves)
    eviction_heap = [(self.eviction_strategy.get_priority(node), node) for node in leaves]
    heapq.heapify(eviction_heap)

    num_evicted = 0
    while num_evicted < num_tokens and len(eviction_heap):
        _priority, x = heapq.heappop(eviction_heap)

        # 1. 将该叶子节点持有的物理 KV 显存归还给内存池
        self.token_to_kv_pool_allocator.free_segment(x.value, start_pos=0)
        num_evicted += len(x.value)

        # 2. 从树中剔除该叶子
        self._delete_leaf(x)

        # 3. 检查其父节点：若父节点失去了所有子节点，且处于未锁定状态，则晋升为新叶子节点！
        if len(x.parent.children) == 0 and x.parent.lock_ref == 0:
            new_priority = self.eviction_strategy.get_priority(x.parent)
            heapq.heappush(eviction_heap, (new_priority, x.parent))

3. 核心机制辨析：驱逐内部节点是否连带子节点？

绝对不会！在 RadixCache 中，驱逐只能且仅能作用于叶子节点（Leaves）！

逻辑推论：若允许直接强行剔除一个内部节点，其子节点虽然持有更新的后缀，但失去了前半截前缀的物理映射，必然导致整个树逻辑出现拓扑断裂，子节点的物理槽位变为孤儿数据。
自底向上逐级收缩：驱逐必须由外向内、自叶向根。只有当一个内部节点的所有子节点全部因为久未访问被依序驱逐干净后，该内部节点“退化”为叶子节点，它才有资格被放入优先堆中，在显存依然不足时被最终驱逐。

2.5 和调度器的配合

在 Continuous Batching 调度框架中，RadixCache 不再是被动的存储组件，而是反向决定了 Batch 编排决策：

Cache-Aware Scheduling（缓存感知调度）：调度器在等待队列（Waiting Queue）中审查就绪请求时，首先调用 RadixTree 的探查接口计算每个候选请求的前缀命中率： $\text{HitRatio} = \frac{L_{hit}}{L_{total}}$ 调度器依据 HitRatio 从高到低对候选请求进行重排序（Priority Scheduling）。
批处理吞吐最大化：具有高命中率的请求只需很少的算力做 Extend Prefill，在同一个计算 Iteration 中可以塞入更多此类请求；更重要的是，将命中了相同前缀节点的多个请求聚类在同一个批次中执行，能够最大化利用 GPU L2 Cache 对同一前缀 KV 的读取复用，极大平抑了 Prefill 对 Decode 步骤造成的延迟抖动（Bubble Jitter）。

第三部分：Extend Attention（SGLang 的前缀 prefill 算子）

3.1 为什么需要一个新算子

在引入 RadixAttention 后，模型在 Prefill 阶段处理的注意力问题呈现出一种全新的几何形态：

标准 Prefill： $Q \in \mathbb{R}^{B \times L \times H \times D}, \quad K, V \in \mathbb{R}^{B \times L \times H \times D}$ $Q$ 与 $KV$ 的长度完全相等（ $L_q = L_k = L$ ）。这是一个标准的方阵自注意力，注意力掩码为一个严格下三角矩阵。
标准 Decode： $Q \in \mathbb{R}^{B \times 1 \times H \times D}, \quad K, V \in \mathbb{R}^{B \times L_{ctx} \times H \times D}$ $Q$ 的长度为 1，KV 长度为整个历史上下文。这是典型的内存受限（Memory-bound）单向量-矩阵乘，使用 PagedAttention 即可高效执行。
Prefix-Cached Prefill（Extend 阶段）：设前缀命中长度为 $L_{prefix}$ ，用户新输入的 Token 数为 $L_{extend}$ 。因为前 $L_{prefix}$ 个 Token 的 KV 已经被计算并保存在物理池中，模型在前向传播中绝不能对这部分 Token 重新运行 Q/K/V 投影。因此，传入注意力层的 Query 张量仅包含新增的 $L_{extend}$ 个 Token： $Q \in \mathbb{R}^{L_q \times H \times D}, \quad L_q = L_{extend}$ 但是，为了让这 $L_{extend}$ 个新 Token 能够与上下文建立关联，Key 和 Value 必须包含完整的历史加上新输入的 Token： $K, V \in \mathbb{R}^{L_k \times H \times D}, \quad L_k = L_{prefix} + L_{extend}$

这导致 $Q$ 和 $KV$ 的序列长度产生严重的不对称性（ $L_q \ll L_k$ 且 $L_q \ne 1$ ）：

无法使用标准 FlashAttention Prefill（如 flash_attn_varlen_func）：标准算子假定传入的 Query、Key 具有一致的累计偏移量 cu_seqlens，强制绑定 $L_q = L_k$ ；
无法使用 PagedAttention Decode：此时 $L_q$ 往往多达数百至数千，属于计算密集型（Compute-bound）GEMM 操作。若按 Decode 模式强行拆分成 $L_q$ 次逐 Token 运算，GPU 的 Tensor Core 矩阵乘单元利用率将跌落至谷底；
结论：必须专门开发一个支持高维不对称矩形因果掩码、支持跨物理页表寻址的 GEMM 算子——这就是 SGLang 的 Extend Attention（海光体系中称作 Prefix Prefill Attention）。

3.2 平移因果掩码的完整数学推导

1. 全局位置空间映射

设序列全局 Token 时间轴从 $0$ 编号至 $L_k - 1$ ：

整个序列的完整 Key/Value 集合覆盖全局区间： $\text{Pos}_K(j) = j, \quad \text{for } j \in [0, L_k - 1]$
当前输入注意力层的 Query 仅代表后半截的新增 Token，其对应的全局序列位置为： $\text{Pos}_Q(i) = L_{prefix} + i, \quad \text{for } i \in [0, L_q - 1]$

2. 自回归因果约束（Causal Constraint）

根据因果注意力定义，任何一个 Query 只能关注位置小于或等于自身全局位置的 Key： $\text{Pos}_K(j) \le \text{Pos}_Q(i)$ 将具体坐标代入不等式： $j \le L_{prefix} + i$ 由于总长度 $L_k = L_{prefix} + L_q \implies L_{prefix} = L_k - L_q$ ，代入得： $j \le i + (L_k - L_q)$ 这就是平移因果掩码（Shifted Causal Mask）的标准判定条件。对于 $L_q \times L_k$ 的注意力矩阵，任意元素 $(i, j)$ 合法的充分必要条件为： $\text{Mask}(i, j) = \begin{cases} 0 & \text{if } j \le i + (L_k - L_q) \\ -\infty & \text{if } j > i + (L_k - L_q) \end{cases}$

3. 几何分解：两段式结构分析

我们将整个 Key 维度区间 $[0, L_k - 1]$ 划分为两段：

前段：前缀区（Prefix Region） $j \in [0, L_{prefix} - 1]$ 对于任意 Query 索引 $i \ge 0$ 以及任意前缀 Key 索引 $j \le L_{prefix} - 1$ ： $j \le L_{prefix} - 1 < L_{prefix} \le L_{prefix} + i$ 不等式恒成立！ 数学结论：在前缀区，不存在任何因果遮挡，整个 $L_q \times L_{prefix}$ 区域为一个全可见的密集矩形（Dense Matrix）！
后段：扩展区（Extend Region） $j \in [L_{prefix}, L_k - 1]$ 令 $j' = j - L_{prefix} \in [0, L_q - 1]$ 代表 Key 在新增块内的相对索引。代入平移掩码不等式： $(j' + L_{prefix}) \le (i + L_{prefix}) \iff j' \le i$ 数学结论：在扩展区内部，掩码退化为一个标准的 $L_q \times L_q$ 下三角因果掩码矩阵！

4. 掩码矩阵示意图

设 $L_{prefix} = 4$ , $L_{extend} = L_q = 3$ , 则 $L_k = 7$ 。平移偏移量 $L_k - L_q = 4$ 。矩阵行代表 Query（ $i \in [0, 2]$ ），列代表 Key（ $j \in [0, 6]$ ）：

                  前缀区 (Prefix, j=0..3)          扩展区 (Extend, j=4..6)
                     j=0   j=1   j=2   j=3        j=4   j=5   j=6
              i=0  [  1     1     1     1   ]    [ 1     0     0  ]   (j <= 0+4=4)
Query (Lq=3)  i=1  [  1     1     1     1   ]    [ 1     1     0  ]   (j <= 1+4=5)
              i=2  [  1     1     1     1   ]    [ 1     1     1  ]   (j <= 2+4=6)
                   |<--- 全可见密集矩形 --->|    |<-- 下三角因果区 -->|

3.3 索引数组详解与三请求具体算例

在 Ragged Batching（变长拼批）模式下，一个 Batch 内的多个请求长短各异，且每个请求的前缀命中情况完全不同。SGLang 构建了一组精巧的 1D 元数据数组传给 Kernel。

1. 数组定义、形状与含义

extend_seq_lens: 形状 [batch_size]，int32。表示每个请求在当前批次中新输入的 Token 数量（即 Query 长度 $L_q$ ）。
seq_lens: 形状 [batch_size]，int32。表示每个请求的历史前缀加新增的总 Token 数量（即完整 KV 长度 $L_k$ ）。
extend_prefix_lens: 形状 [batch_size]，int32。每个请求命中的前缀长度：extend_prefix_lens[b] = seq_lens[b] - extend_seq_lens[b]。
qo_indptr: 形状 [batch_size + 1]，int32。Query 与 Output 紧密排列张量（q_extend, o_extend）的前缀和索引指针（Indptr）。
- 计算方式：qo_indptr[0] = 0，qo_indptr[1:] = torch.cumsum(extend_seq_lens)。
- 作用：Kernel 中确定第 $b$ 个请求的 Query 起始行地址为 qo_indptr[b]，行数为 qo_indptr[b+1] - qo_indptr[b]。
kv_indptr: 形状 [batch_size + 1]，int32。kv_indices 数组的前缀和索引指针。
- 计算方式：kv_indptr[0] = 0，kv_indptr[1:] = torch.cumsum(seq_lens)。
kv_indices: 形状 [sum(seq_lens)]，int32。保存整个批次中所有请求的物理内存槽位全局 ID。
- 作用：对于请求 $b$ ，其所有的 $L_k$ 个历史与当前 Token 在显存物理池中的存储地址直接通过切片 kv_indices[kv_indptr[b] : kv_indptr[b+1]] 查表寻址。
extend_start_loc: 形状 [batch_size]，int32。等于 qo_indptr[:-1]，标记每个请求在当前批次新增 Token 中的相对偏移。

2. 包含 3 个请求的具体数值实例

假设当前调度批次包含 3 个请求：

请求 0：多轮对话第三轮。前缀命中 $L_{prefix} = 4$ tokens，当前新发提问 $L_{extend} = 2$ tokens。总长 $L_k = 6$ 。
请求 1：全新独立请求。完全无前缀命中 $L_{prefix} = 0$ tokens，新提问 $L_{extend} = 3$ tokens。总长 $L_k = 3$ 。
请求 2：Few-shot 模板提问。命中固定前缀 $L_{prefix} = 5$ tokens，新输入 $L_{extend} = 1$ token。总长 $L_k = 6$ 。

假设底层的物理内存池在分配时返回的离散全局物理槽位 ID 如下：

请求 0 物理槽位：前缀槽位 [108, 109, 110, 111]，新增槽位 [42, 43]。
请求 1 物理槽位：无前缀，新增槽位 [201, 202, 203]。
请求 2 物理槽位：前缀槽位 [80, 81, 82, 83, 84]，新增槽位 [99]。

调度器计算出的所有张量具体数值完全展开如下：

数组名称	形状	具体数值内容	物理意义解读
`extend_seq_lens`	`[3]`	`[2, 3, 1]`	各请求新增 Q 长度分别为 2, 3, 1
`seq_lens`	`[3]`	`[6, 3, 6]`	各请求完整 KV 长度分别为 6, 3, 6
`extend_prefix_lens`	`[3]`	`[4, 0, 5]`	各请求前缀长度分别为 4, 0, 5
`qo_indptr`	`[4]`	`[0, 2, 5, 6]`	Q 总长度为 6；Req0 占 `[0,2)`，Req1 占 `[2,5)`，Req2 占 `[5,6)`
`kv_indptr`	`[4]`	`[0, 6, 9, 15]`	KV 总长为 15；Req0 占 6 项，Req1 占 3 项，Req2 占 6 项
`extend_start_loc`	`[3]`	`[0, 2, 5]`	各请求在 Q 张量中的起始行号
`kv_indices`	`[15]`	`[108, 109, 110, 111, 42, 43,` `201, 202, 203,` `80, 81, 82, 83, 84, 99]`	15 个物理槽位索引，精确对应各个请求的历史与新增物理空间

3.4 Kernel 分块（Tiling）与循环上界裁剪

在 SGLang 的 Triton 实现中（参考 python/sglang/kernels/ops/attention/extend_attention.py:327-750），计算被明确组织为两阶段流水线：

+-------------------------------------------------------------------------------+
|  阶段 1: 前缀全可见遍历循环 (Stage 1: Prefix Loop)                            |
|  for start_n in range(0, cur_seq_len_prefix, BLOCK_N_PREFIX):                 |
|      - 无因果掩码检查，直接以最大 Tile (如 BLOCK_N=64 或 128) 跑 GEMM         |
|      - 维持 Online Softmax 的局部 m 与 l                                      |
+-------------------------------------------------------------------------------+
                                      |
                                      v (累积局部 Softmax 状态)
+-------------------------------------------------------------------------------+
|  阶段 2: 扩展因果下三角循环 (Stage 2: Extend Triangular Loop)                 |
|  extend_end = min(cur_seq_len_extend, (cur_block_m + 1) * BLOCK_M)            |
|  for start_n in range(0, extend_end, BLOCK_N):                               |
|      - 施加 (cur_block_m * BLOCK_M + offs_m) >= (start_n + offs_n) 因果判断   |
|      - 早停裁剪 (Early Exit): 凡是超出 extend_end 的块直接截断跳过             |
+-------------------------------------------------------------------------------+

1. Q 的分块（M 方向）

沿 Query 维度以 BLOCK_M（如 64）为步长切块。
当前正在处理的 Q 分块索引为 cur_block_m，其负责的相对 Token 区间为： $[\text{cur\_block\_m} \times BLOCK\_M, \; (\text{cur\_block\_m} + 1) \times BLOCK\_M - 1]$

2. KV 的跨阶段分块与统一遍历

Stage 1（前缀区）： Key 遍历区间为 $[0, L_{prefix}-1]$ ，步长为 BLOCK_N_PREFIX。由于该区域内部没有任何 Token 受到因果掩码阻挡，每个 Block 内加载的全部 KV 均参与有效运算。无需计算任何逐元素掩码判断，分支预测开销为零。
Stage 2（新增因果区）： Key 遍历区间为 $[0, L_{extend}-1]$ $[0, L_{e x t e n d} - 1]$ ，步长为 BLOCK_N。此时加载的是新生成的 KV 块，并在 Tile 内部利用向量比较指令施加掩码：
```
mask_causal = (cur_block_m * BLOCK_M + offs_m[:, None]) >= (start_n + offs_n[None, :])
```

3. 循环上界裁剪与早停（Early Exit）

对于新增因果区，若某个 KV 分块满足： $start\_n \ge (cur\_block\_m + 1) \times BLOCK\_M$ 这意味着该 KV 块中的所有 Token，其物理生成时间严格晚于当前 Query 块中的所有 Token。在因果注意力机制下，整个 Tile 将全部被掩码置为 $-\infty$ ，对最终 Softmax 贡献为零。 SGLang 通过直接截断循环上界： $\text{extend\_end} = \min(L_{extend}, \; (cur\_block\_m + 1) \times BLOCK\_M)$ 所有位于对角线上方的无用 Tile 直接在循环条件中被剔除，完全不发起内存加载，算力与显存带宽节约高达 $50\%$ 。

3.5 和 Chunked Prefill 的关系辨析

在 LLM 推理系统文献中，Chunked Prefill 与 Extend Attention 常同时出现，二者处于不同的系统层级：

维度	Chunked Prefill（分块 Prefill）	Extend Attention（扩展注意力算子）
所属抽象层级	系统宏观调度策略（Scheduling Policy）	微观 GPU 执行算子（Attention Kernel）
解决的核心问题	解决长 Prompt Prefill 耗时数十毫秒甚至数秒导致同批次 Decode 步骤产生严重延迟气泡（Inter-token Latency Spike）的问题。	解决 Q 和 KV 长度不对称（ $L_q \ne L_k$ ）且带有平移因果掩码时的底层注意力张量加速问题。
调度行为	将超长 Prompt 人工切割为若干等长小块（如每块 512），在多个调度 Iteration 中分步执行。	无论分块与否，只要当前步骤中存在“已有历史 KV + 计算新增部分 Q”的几何形态，就调用该算子。
协同关系	Chunked Prefill 依赖 Extend Attention 落地：当第二块 Chunk 进入模型时，第一块 Chunk 的 KV 已经存在于显存中，此时计算该 Chunk 的注意力必须依赖 Extend Attention 算子。

第四部分：三者如何协同（端到端全流程）

下面以一个实际请求的完整生命周期，梳理 RadixCache、BlockManager / Paged Memory 与 Extend / Paged Attention 算子 在调度引擎内部的闭环协同：

sequenceDiagram
    autonumber
    participant Client as 客户端 (Client)
    participant Scheduler as 调度器 (Scheduler)
    participant Radix as 前缀基数树 (RadixCache)
    participant MemPool as 显存物理池 (Paged Token Pool)
    participant ModelRunner as 执行器 (Model Runner / Kernel)

    Client->>Scheduler: 发送提问请求 Req (Prompt 1000 tokens)
    Scheduler->>Radix: 查询最长公共前缀 match_prefix(Prompt)
    Radix-->>Scheduler: 命中前缀 L_prefix = 800, 返回已缓存物理槽位, lock_ref++

    Note over Scheduler: 算出 L_extend = 200 (未命中需计算部分)
    Scheduler->>MemPool: 为新增的 200 tokens 申请新物理块 (200 slots)
    MemPool-->>Scheduler: 返回新分配物理槽位列表

    Scheduler->>ModelRunner: 构建 metadata: qo_indptr, kv_indptr, kv_indices
    Note over ModelRunner: 执行 Prefill 阶段: 仅对 200 个 Q 做前向计算
    ModelRunner->>ModelRunner: 触发 Extend Attention 算子 (Shifted Causal Mask)
    ModelRunner-->>Scheduler: 输出 Prefill 首字 Logits, 采样出 Token 1001

    Scheduler->>Radix: 插入新计算的 200 tokens 至 Radix 树中

    loop 自回归解码循环 (Decode Loop, 逐 Token 生成)
        Scheduler->>MemPool: 为第 t 个 Token 申请 1 个槽位 (若物理块满则分配新块)
        Note over Scheduler: 构造 block_table 映射
        Scheduler->>ModelRunner: 传入 Q (长度 1), block_tables, seq_lens
        ModelRunner->>ModelRunner: 触发 PagedAttention (v1 或 v2 Split-KV)
        ModelRunner-->>Scheduler: 返回第 t 个 Token 输出, 采样下一个 Token
    end

    Client<<-Scheduler: 输出生成完毕 (EOS)
    Scheduler->>Radix: 解锁前缀路径: dec_lock_ref(terminal_node)
    Note over Radix: lock_ref 归零, 节点进入可驱逐候选池 (Evictable List)

第五部分：对照海光 DCU 实现

本节直接对照本地代码库 flash-attention-hygon 进行源码级深入分析。

5.1 双引擎架构与职责划分

代码库为了兼顾上游算子演进与国产架构硬件性能，设计了清晰的双引擎分流架构：

csrc/flash_attn_cutlass/：移植并重构自上游 CUTLASS 体系。PagedAttention 的所有实现均位于此目录（如 src/paged_attention.cu、src/paged_attention.h、src/paged_attention_bshd.cu）。它承担 Paged 场景与通用回退兜底。
csrc/flash_attn_hipc/：海光自研的 MLS（Matrix-Level Shuffler）原生汇编与手写 HIP 引擎。所有的 Prefix Prefill 与 Prefix Decode 算子均位于此目录（如 src/flash_fwd_b16_fa.h、flash_api.cpp）。它绕过任何外部编译框架，直接释放海光 DCU 的原生矩阵乘指令能力。

5.2 PagedAttention 在 CUTLASS 侧的重构

海光 DCU 版 PagedAttention 绝非简单的语法宏替换，而是针对海光 GCN/CDNA 架构硬件特性的深度重写：

1. Wavefront 64 架构适配

在 csrc/flash_attn_cutlass/src/paged_attention.h:6 中：

#define WARP_SIZE 64

改动原因：NVIDIA GPU 的基本线程调度单位 Warp 为 32 线程；而海光 DCU（类比 AMD CDNA 架构）的底层硬件执行波面（Wavefront）宽度为 64 线程。
关联影响：所有的 Warp 内部规约掩码（Mask）必须从 32 位升级为 64 位（如 paged_attention.h:297 的 mask = WARP_SIZE / 2; mask >= 1; mask /= 2，外层循环从 32 开始二分折半至 1）。

2. 源码保留的上游历史痕迹

在 csrc/flash_attn_cutlass/src/paged_attention.cu:329 中，海光代码完好保留了原版注释：

// NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
float4_t accs[Mloop][NUM_ROWS_PER_THREAD];

印证了海光 DCU PagedAttention 算子的历史基因直接继承自 vLLM 论文原版内核。

3. Grid 驱动维度重构：从 Q Head 到 KV Head

上游 vLLM 原版逻辑：gridDim.x = num_heads。每个 Block 负责一个 Q Head，GQA 下多个 Q 独立加载相同 KV。

海光 DCU 重构逻辑（见 csrc/flash_attn_cutlass/src/paged_attention.cu:559-561）：

dim3 grid;
grid.x = num_kv_heads;  // ★ 核心改动：由 KV Head 驱动 Grid！
grid.y = num_seqs;
grid.z = max_num_partitions;

REUSE_KV_TIMES 模板特化展开：在 csrc/flash_attn_cutlass/src/paged_attention.h:224-248 中，海光定义了 REUSEKV_SWITCH 宏：

#define REUSEKV_SWITCH(reusekv,...) \
[&] { \
    if (reusekv==64){ constexpr static int REUSE_KV_TIMES = 64; return __VA_ARGS__(); } \
    else if (reusekv==48){ constexpr static int REUSE_KV_TIMES = 48; return __VA_ARGS__(); } \
    ... \
    else { constexpr static int REUSE_KV_TIMES = 4; return __VA_ARGS__(); } \
}()

每个 Thread Block 加载一份物理 KV Head 数据，在寄存器内部展开并复用 REUSE_KV_TIMES 次（对应 GQA Group 比率），彻底消除了并发 Q Head 重复访存的带宽膨胀。

4. 原生硬件矩阵乘指令（MMAC）替换

上游 CUDA 代码依靠标量点积或 CUDA Tensor Core WMMA API。海光在 csrc/flash_attn_cutlass/src/paged_attention.h:314-332 中直接调用 DCU 底层内置内联函数（Intrinsics）：

template<bool is_half>
inline __device__ void builtin_amdgcn_mmac(const half4_t& reg_a, const half4_t& reg_b, float4_t& reg_c)
{
  if constexpr (is_half){
    reg_c = __builtin_hcu_mmac_f32_16x16x16_f16(reg_a, reg_b, reg_c);
  } else {
    reg_c = __builtin_hcu_mmac_f32_16x16x16_bf16(*(v4bh*)&reg_a, *(v4bh*)&reg_b, reg_c);
  }
}

通过调用 __builtin_hcu_mmac_f32_16x16x16_f16（针对 FP16）以及 __builtin_hcu_mmac_f32_16x16x32_fp8_fp8（针对 FP8），直接驱动海光硬件的 Matrix Core（矩阵加速单元），在寄存器级完成矩阵内积。

5. `PARTITION_SIZE` 运行时动态自适应切换

在 csrc/flash_attn_cutlass/src/paged_attention.cu:486-490 中：

int PARTITION_SIZE = 512;
if (real_reuse_times > 15) PARTITION_SIZE = 256;
if (PA_PARTITION_SIZE != 0) PARTITION_SIZE = PA_PARTITION_SIZE;

改动依据：当 GQA 比例极高（real_reuse_times > 15，如某些超大模型的 8:1 甚至更高比例）时，每个线程块内部需要分配给多个 Q Head 状态的寄存器和 LDS 显存急剧增加。将 PARTITION_SIZE 动态从 512 下调至 256，能降低单 Block 对 LDS 容量的压力，防止寄存器溢出（Spill）到全局显存导致性能崩塌；同时支持通过环境变量 PA_PARTITION_SIZE 强制覆盖调优。

5.3 Prefix Attention 在 HIPC 侧的重构

海光自研的 MLS 算子（HIPC 引擎）承担了前缀 Prefill 的重任：

1. 平移因果掩码的代码落地

源码对应位置：csrc/flash_attn_hipc/include/fwd/softmax.h:181 在 apply_mask_causal 函数中：

const int col_idx_limit_right = std::min(max_seqlen_k, row_idx + max_seqlen_k - max_seqlen_q);
...
tensor[mi + ni * (WARP_M / 32)][min_tile_n * 2 + min_tile_m].f32[vec_idx] =
    (col_idx > col_idx_limit_right) ? -INFINITY : tensor[...];

这行代码中的 row_idx + max_seqlen_k - max_seqlen_q，精确验证了第三部分推导出的平移因果掩码公式 $j \le i + (L_k - L_q)$ 。

2. 导出接口与对 SGLang 内存池布局的专属定制

源码对应位置：csrc/flash_attn_hipc/flash_api.cpp:1745-1822，导出 API：

hg_prefix_prefill_varlen_fwd
hg_prefix_decode_varlen_fwd

特别关注 flash_api.cpp:1800-1822 对 layout 参数的解析：

layout == 0（bhsd）： $K/V$ 形状为 [num_blocks, num_heads_k, block_size, head_dim]；
layout == 1（bshd）： $K/V$ 形状为 [num_blocks, block_size, num_heads_k, head_dim]。

layout = 1 正是海光研发团队专门为了对接 SGLang 的 TokenToKVPool 显存池布局而原生开辟的无开销直通分支。在 tests/prefix_decode_sglang_decode.py:305-325 中，SGLang 测试用例直接传入 layout=1，免去了运行时对 KV Cache 执行昂贵转置（Transpose）的显存拷贝开销。

3. `k_delta` / `v_delta` 差分寻址黑科技

在长上下文循环步进中，上游算子通常在每次 Tile 迭代时，重新执行乘法与页表寻址计算物理基地址。海光在 csrc/flash_attn_hipc/src/flash_fwd_b16_fa.h:553-568 中创造性地引入了差分步进寻址：

const int block_table_idx_cur     = n_block_loop * kBlockN / params.page_block_size;
const int block_table_offset_cur  = n_block_loop * kBlockN - block_table_idx_cur * params.page_block_size;
const int block_table_idx_next    = min(n_block_max - 1, n_block_loop + 1) * kBlockN / params.page_block_size;
const int block_table_offset_next = min(n_block_max - 1, n_block_loop + 1) * kBlockN - block_table_idx_next * params.page_block_size;

// 计算页表与偏移差值
const int table_diff              = block_table[block_table_idx_next] - block_table[block_table_idx_cur];
const int offset_diff             = block_table_offset_next - block_table_offset_cur;

// 转化为线性物理增量
const int64_t k_delta             = int64_t(table_diff) * int64_t(params.k_batch_stride) + offset_diff * int64_t(seqlen_k_stride);
const int64_t v_delta             = int64_t(table_diff) * int64_t(params.v_batch_stride) + offset_diff * int64_t(seqlen_v_stride);

// 直接原地累加硬件缓冲资源描述符 (Buffer Resource Descriptor)
*(int64_t*)&gK += k_delta * sizeof(Element);
...
*(int64_t*)&gV += v_delta * sizeof(Element);

改动原因与收益：海光 DCU 底层采用类似于 AMD 的 128-bit Buffer Resource 描述符来执行全局内存加载。通过仅计算相邻两步之间的差值 k_delta，直接将增量原子加到硬件寄存器 gK 上，省去了每轮迭代中的 64 位宽乘法运算与冗余的全局页表重新读取，极大降低了 VGPR（矢量通用寄存器）的压力，规避了寄存器溢出。

4. 原生 C++/HIP 替代 Triton 的技术必然性

上游 SGLang 与 vLLM 的 Extend/Prefix 算子大量基于 OpenAI Triton 编写（如 extend_attention.py）。而在国产硬件生态中，Triton 后端编译器对特定 ISA 指令集（如 Wavefront 64、专用 MMAC 指令、局部共享内存 LDS 布局）的自动调优与指令调度能力仍处于演进阶段，直接运行 Triton 算子容易产生大量局部内存溢出（Spill）并丧失硬件峰值性能。海光选择全手写原生 C++/HIP 模板算子，精确控制循环展开深度（#pragma unroll）、显式插入调度屏障（__builtin_amdgcn_sched_barrier）以及直接内联汇编指令，从而在硬件上释放出更高的执行效率。

未确认事项

csrc/flash_attn_cutlass/src/paged_attention.cu:479-482 中的 layout == 2：在海光 PagedAttention 的主机封装函数中，存在 if (layout == 2) 的特殊分支逻辑（block_size = num_kv_heads; num_kv_heads = key_cache.size(2)），推测是为某种特定的自定义 KV 布局（可能是多头隐藏维度在内侧的高性能交错格式）设计的，但本地测试集仅覆盖了 layout=0 与 layout=1，未在代码注释中直接确认其对应的上层推理引擎名称。
海光 DCU FP8 Prefix Prefill 在 gfx936 架构下的 scales_q_ 强制断言：在 csrc/flash_attn_hipc/flash_api.cpp:1854-1857 中，针对 arch_id == 936 的硬件定义了 use_fp8_bf16_mixed 模式，并强制要求 scales_q_、scales_k_、scales_v_ 必须传入且仅支持 bshd 布局。鉴于本地环境缺乏真实 DCU 设备无法动态挂载驱动执行，未能通过实际运行证实其在硬件层面是否为指令集硬件级混合精度限制。