TensorRT-LLM 源码完全拆解（六）：Attention Backend 的 prepare() 内部——块句柄如何变成块表、页表与块偏移

第四篇留下的账面结构

回顾第四篇的核心结论：TensorRT-LLM 的 KV 缓存由 KVCacheManager（Python 侧包装，内部实现是 C++ 的 KVCacheManager）管理，它把显存切成固定大小的块（page，默认 32 或 64 个 token 一块），每个块有两个属性：

• 块句柄（block id）：一个整数编号，全局唯一，是这个块在池子里的”学号”。调度器、前缀复用树、free_blocks 队列只认这个编号。
• 物理位置：这个块在显存池里的实际地址，由池的基址加上 block_id × 块字节数 得到偏移。

上层逻辑（前缀匹配、复用、回收、beam 分叉）全部在块句柄的层面工作——句柄是逻辑概念，不涉及地址；这样池可以换、布局可以变，逻辑不用改。但内核不行：内核要算出”第 3 个请求的第 17 个 KV 块在显存里的哪个位置”，它必须拿到地址偏移或者能换算成偏移的编号。

块偏移张量的四维布局

上图是 TensorRT-LLM 官方文档里分块上下文（chunked context）注意力的示意图：一条长提示被切成多段依次执行，每段产生的 KV 依次写进缓存。账本必须按”块”而不是按”步”记账，正是因为它要跨多步持续持有这些物理页——这也是本节四维张量存在的理由。图片来源：TensorRT-LLM 官方文档 media/feat_long_seq_chunked_attention.png。

这个翻译就发生在 KVCacheManager.copy_batch_block_offsets() 里。先看它填充的目标张量的形状：

1
[num_pools, max_num_requests, 2, max_blocks_per_seq]

四个维度逐一拆开：

• num_pools：池的数量。绝大多数模型只有一个池；但可变滑动窗口注意力（VSWA，Variable Sliding Window Attention）会按窗口大小分池，每种窗口一个池，这样不同层可以用不同粒度的块管理。Gemma 系列这类”层间窗口大小不同”的模型就会走这条路。
• max_num_requests：行数按最大批大小预分配，而不是按当前批大小。原因和第五篇讲过的一样：CUDA Graph 捕获的是地址，缓冲区的行数必须在捕获时就固定，运行中只能覆写前 num_seqs 行。
• 2：K 与 V。多数现代模型（MHA/GQA/MLA 的 KV 合并存储除外）把 K 和 V 分成两块独立的池，或者在同一池里按因子 2 排列，所以这个维度是 2，索引 0 拿 K、索引 1 拿 V。
• max_blocks_per_seq：每序列最多多少个块，由 max_seq_len / tokens_per_block 决定，向上取整。同样是预分配的上限值。

注意这里的 dtype：int32。块偏移不需要 64 位，因为块偏移是”第几个块”而不是”第几字节”，池的块数是 2 的 30 次方级别就够用了，int32 省一半带宽——这是个大批场景下真实影响性能的决策，因为每次前向都要把这些表从主机搬到设备。

句柄到偏移的物理含义

关键问题是：块偏移张量里存的到底是”块句柄”还是”字节偏移”？

答案是两者在数值上可以统一，而且 TensorRT-LLM 正是这么设计的。在分页 KV 池里，物理地址 = 池基址 + 块编号 × 块步长；而”每序列第 i 个块的编号”既可以直接当索引（内核自己做 base + idx * stride），也可以由主机侧预乘成字节偏移。TensorRT-LLM 选择了让张量里存块编号，由内核用块步长换算，理由是：

• 块编号是布局无关的。同一个编号，在 NHD 布局和 HND 布局里对应的字节位置不同，但如果张量里存字节偏移，换布局就得重算整张表。
• 池可以重映射。启用显存整理或跨实例 KV 传输时，块编号不变但物理位置会变；只存编号就不会出现”表里的地址已经过期”的问题。

所以 copy_batch_block_offsets 里 self.impl.copy_batch_block_offsets(...) 调用的是 C++ 实现——它在 C++ 侧直接遍历前綴复用树的节点，把每个请求持有一串块编号铺进行里。名字里的 “offsets” 指的是”这个请求的块序列在块池里的偏移序列”，而不是字节地址。

下面是一段最小化的伪代码，展示这层翻译在概念上做了什么：

1
# 输入：request_ids（本批请求 id 列表）、beam_width、num_context、num_seqs
2
for row, rid in enumerate(request_ids):
3
    # 1) 从请求状态里取出它持有的块序列（这是一个块句柄列表）
4
    blocks = blocks_held_by(rid)          # 长度 = ceil(kv_len / tokens_per_block)
5

6
    # 2) context 段（预填充请求）与 generation 段（解码请求）分开处理：
7
    #    context 段只用 1 个 beam，因为它还没有分叉
8
    #    generation 段要按 beam_width 复制行 —— beam search 时每条请求有多个 beam
9
    for b in range(1 if row < num_context else beam_width):
10
        dst_row = row_idx_for(row, b)
11
        dst[0, dst_row, 0, :len(blocks)] = blocks   # K 池
12
        dst[1, dst_row, 1, :len(blocks)] = blocks   # V 池（或同一池的 V 因子）

这段伪代码里有两个细节值得展开，它们在真实代码里都有对应的注释：

第一，context 与 generation 分开调用。 真实实现里是两次调用：

1
self.impl.copy_batch_block_offsets(self.host_kv_cache_block_offsets,
2
                                   request_ids[:num_context], 1, 0)
3
self.impl.copy_batch_block_offsets(self.host_kv_cache_block_offsets,
4
                                   request_ids[num_context:],
5
                                   beam_width, num_context)

第一次传 beam_width=1、num_context=0，处理预填充请求；第二次传真实的 beam_width 与 num_context，处理解码请求。为什么预填充段强制 beam_width=1？因为预填充阶段还没有分叉：一条新请求进入时只有一个序列，beam search 的分叉发生在它开始逐 token 生成之后。如果在预填充行上就按 beam 展开，会凭空多出几行指向同一个块序列的假 beam，浪费带宽也浪费内核的序列槽位。

第二，cross attention 的展开逻辑正好相反。 交叉注意力（编码器-解码器模型的 decoder 读 encoder 的 K/V）里，多条请求可能共享同一份编码器输出吗？不一定，但同一条请求的多个 beam 一定共享同一份 cross K/V——因为 cross K/V 由编码器一次性写出，解码器的每个 beam 读的是同一份。所以 cross 路径的逻辑是”每请求填一行，然后把 generation 行按 beam 复制（repeat_interleave）”：

1
gen_block_offsets = host_block_offsets[
2
    pool_idx, num_context:num_context + num_gen_requests]
3
dst_tensor[pool_idx, num_context:num_seqs].copy_(
4
    gen_block_offsets.repeat_interleave(beam_width, dim=0),
5
    non_blocking=True)

repeat_interleave(beam_width, dim=0) 把 [num_gen_requests, ...] 展开成 [num_gen_requests * beam_width, ...]，第 i 条请求的 beam_width 个 beam 拿到完全相同的一行。同一份代码里还有一条断言在守着这个语义：

1
num_gen_requests = len(request_ids) - num_context
2
expected_num_seqs = num_context + num_gen_requests * beam_width
3
assert num_seqs == expected_num_seqs, (
4
    f"Cross KV cache block offsets expected {expected_num_seqs} "
5
    f"decoder rows, got {num_seqs}.")

它把”元数据张量的行数必须等于 context 行数 + generation 行数 × beam 宽度”写成了硬性契约。这类断言在 TensorRT-LLM 里到处都是，因为行数的错位不会立刻崩溃，只会让内核读到另一个批次（或者另一条请求）的块——这类错误极难排查，所以宁可断言失败也不要静默出错。

一个真实踩过的坑：异步拷贝的源被覆写

copy_batch_block_offsets 的代码里有一段长注释，讲的是一个真实的回归（内部编号 nvbug 6293536），值得完整讲一遍，因为它是理解”重叠调度 + 异步 H2D + 持久缓冲区”三者相互作用的最好例子。

设想的朴素实现是这样的：

1
# 错误写法（曾经的真实代码）
2
self.impl.copy_batch_block_offsets(self.host_kv_cache_block_offsets,
3
                                   request_ids, beam_width, num_context)
4
dst_tensor.copy_(self.host_kv_cache_block_offsets, non_blocking=True)

主机侧先把本批的块编号铺进那块持久的主机缓冲区（它是 max_batch 大小的，跨步复用），然后异步拷到设备。看起来没问题：先写源、再拷。

问题出在重叠调度（第六篇前的第五篇提过，SGLang 侧在SGLang 源码完全拆解（六）：Overlap Scheduler 里详细讲过同一类机制）。开启重叠后，CPU 会提前一个迭代准备下一批的输入：当第 t 步的 H2D 拷贝还在流里排队时，CPU 线程已经在为第 t+1 步往同一块持久缓冲区里写数据了。于是第 t 步那次还没执行的异步拷贝，读到的源已经被第 t+1 步覆写了——注意力内核拿到了另一个批次的块表。

修法是把源改成”每次调用新分配一块 pinned 主机缓冲，把持久缓冲区的前 num_rows 行快照进去”：

1
def _stage_block_offsets_for_copy(self, num_rows, max_blocks=None):
2
    if max_blocks is None:
3
        width = self.max_blocks_per_seq
4
    else:
5
        width = min(max(max_blocks, 1), self.max_blocks_per_seq)
6
    host_block_offsets = torch.empty(self.num_pools, num_rows, 2, width,
7
                                     dtype=torch.int32,
8
                                     pin_memory=prefer_pinned(), device='cpu')
9
    host_block_offsets.copy_(self.host_kv_cache_block_offsets[:, :num_rows, :, :width])
10
    return host_block_offsets

新分配的缓冲区不会被下一迭代碰，non_blocking=True 的 H2D 就安全了。注释里还提到：“缓存式主机分配器会把它一直持有到消费它的拷贝完成”——即释放不是靠 Python 引用计数，而是靠 CUDA 缓存分配器的流语义。

而原来那块持久缓冲区不能删，因为有三类 CPU 侧读者直接读它，且依赖它”跨调用持久、行数固定为 max_batch”：

• DSA 稀疏注意力（DeepSeek Sparse Attention，站内 DSA 稀疏注意力完全拆解 讲过这个算法本身）：索引器要按块粒度打分，需要读到块表。
• 投机解码的 draft/target 缓冲区交换：草稿模型与目标模型的前向交替使用同一份元数据对象，交换时需要读原表。
• KV 缓存重定位路径：启用显存整理或弹性伸缩时，块的物理位置要搬，搬运逻辑按块编号读表。

注释里那句 “Reassigning this attribute to a per-call, num_seqs-sized buffer broke those readers (nvbug 6293536 regression)” 说明一开始的修法可能是把属性直接指向新缓冲——那会把这些读者的视野从 max_batch 行缩成 num_seqs 行，同样是错的。正确修法是”两份缓冲并存”：持久的那份给同步的 CPU 读者，新的那份给异步的 H2D。 这是一个很典型的工程结论：当一个数据结构同时被”需要稳定视图的读者”和”需要瞬时快照的读者”使用时，共用一个实例迟早出事。

三类字段

把它的几十个字段按用途归类，会看得清楚得多：

第一类：容量上限（构造时确定，终身不变）

1
max_num_requests: int      # 单批最多多少条请求
2
max_num_tokens: int        # 单批最多多少 token
3
max_num_sequences: int     # 单批最多多少序列（beam 展开后）

这三个数字是预分配的依据。所有按行索引的缓冲区——seq_lens、kv_lens、块偏移表、块表——都按这三个上限之一来分配行数，运行中只填前 num_seqs 行。第五篇讲过：这是 CUDA Graph 的硬性要求，因为图里记录的是设备地址，缓冲区不能换。

第二类：本步的实际形状（每步都变）

1
seq_lens: torch.Tensor      # (batch_size)，本步每条序列的 query 长度
2
seq_lens_kv: torch.Tensor   # 交叉注意力专用，KV 侧长度
3
num_contexts: int           # 本批里有多少条是 context（预填充）请求

seq_lens 有一个很关键的细节——它有两个副本：

1
_seq_lens: torch.Tensor       # 固定在 CPU（pinned）
2
_seq_lens_cuda: torch.Tensor  # 同一份数据的设备副本

为什么要有两份？代码注释解释了：seq_lens_cuda “被 logits processor 使用”，而且”保留两份副本可以避免 FlashInfer 的 prepare() 实现里大量多余的拷贝”。真实的因果链是这样：采样阶段要在设备上做 logits 处理，需要 per-sequence 长度；注意力后端的内核也需要它。 如果只有一份设备张量，采样侧每次想读就要 D2H 同步——而 D2H 同步会打断重叠调度（第五篇讲的”CPU 提前跑一个迭代”正是靠避免同步才能成立的）。两份副本的分工是：CPU 侧那份给调度与元数据构造用（CPU 需要直接算数，不需要同步），设备侧那份给内核与采样用。

seq_lens 的 setter 里还有一行很值得看的保护：

1
if self.is_cuda_graph and self._seq_lens_cuda is not None:
2
    # Very important: do not reallocate if we are using CUDA graphs.
3
    self._seq_lens_cuda.copy_(self._seq_lens, non_blocking=True)
4
else:
5
    self._seq_lens_cuda = self._seq_lens.cuda(non_blocking=True)

注释里那句 “Very important” 是重点：CUDA Graph 模式下必须原地 copy_，绝不重新分配。 一旦走了 else 分支（新建张量），_seq_lens_cuda 的地址就变了，而捕获的图里写死了旧地址——重放时会读到一块已经释放或被复用的显存。

顺带还能看到一个反直觉的性质：CUDA Graph 模式下批大小不变。注释说 “This copy is safe because the batch size is guaranteed to not change in the CUDA graph case”。因为解码图是按批大小分桶捕获的，某个桶对应的图只在那个批大小下重放，所以 _seq_lens_cuda 的 shape 恒定，原地拷贝永远装得下。但序列长度本身可以变（投机解码时草稿 token 数不同），所以值要每步更新——这才需要 copy_。这正好呼应了第五篇讲的结论：TensorRT-LLM 用”形状固定”换”稳态确定性”，而序列长度这种”形状不变但值变”的东西就落在可更新的那一侧。

第三类：KV 账本接口（跨越到第四篇的世界）

1
kv_cache_manager: KVCacheManager | KVCacheManagerV2 | None
2
kv_cache_params: KVCacheParams          # 含 num_cached_tokens_per_seq
3
kv_lens: torch.Tensor                   # 本步之后每条序列的 KV 长度
4
request_ids: List[int]                  # 本批请求 id，用于向 manager 查块

这里最关键的是 request_ids——元数据本身不持有块表，它只持有”向谁要块表”的钥匙。真正的块在第四篇的 KVCacheManager 里，元数据每次 prepare() 时拿着 request_ids 去问一次。这个设计的好处是：块表天然跟着账本走，不存在”元数据里的块表已经和账本不一致”的窗口。

还有一个隐藏的强约束：is_cross 属性。

1
@property
2
def is_cross(self) -> bool:
3
    return self.seq_lens is not self.seq_lens_kv

交叉注意力的判定用的是 is not——对象身份比较，不是值比较。 自注意力时 seq_lens_kv 的 getter 直接返回 seq_lens（同一个对象），所以 is not 为假；交叉注意力时两者被显式设为不同张量，is not 为真。为什么不用值比较？因为一行 if 判断省下的开销无所谓，但这个约定把”两者是不是同一份数据”这件事变成了对象身份问题，于是任何重新分配 seq_lens 的代码都会立刻破坏判定。这跟第五篇讲的”元数据必须是同一对象（is 而非 ==）“是同一种设计语言——用对象身份编码语义，因为身份不会撒谎。

prepare() 的两个入口

基类里的 prepare() 短得几乎不像个钩子：

1
def prepare(self):
2
    """Hook to be called before the forward step of the model."""
3
    self._prepare_mamba_metadata()

它只做一件事：如果这个模型带 Mamba 层（混合架构，如 Jamba、Mamba2 混合模型），顺便准备一下 Mamba 的元数据。真正的注意力元数据准备全部由子类覆写。这是一条很清晰的分层：基类负责跨注意力机制共用的部分，子类负责把 KV 账本翻译成自己内核要的形状。

此外还有一个专为编码器设计的旁路：

1
def prepare_encoder_only(self) -> None:
2
    """Hook for encoder-only (no-KV-cache) forward setup.
3

4
    Defaults to the full ``prepare()``; backends with a leaner encoder-only
5
    path override this."""
6
    self.prepare()

编码器（比如 BERT 类或 Whisper 的编码器部分）不需要 KV 缓存，块表、块偏移、前缀复用全都不存在。如果硬走完整的 prepare()，会白白做一堆”向 manager 要块表”的工作——而 manager 是 None，只能靠分支绕开。所以给了个专门入口，让后端提供一个”什么账都不用算”的快路径。

TRT-LLM 原生后端的 prepare_encoder_only 实现里有一段注释，讲了一个具体的坑：

1
# `host_request_types` is allocated with `torch.empty_like` (not zeroed),
2
# and the full `prepare()` is what normally fills it. This fast path skips
3
# `prepare()`, so explicitly mark every request as a context request (type 0)
4
# here -- otherwise the FMHA reads whatever garbage the buffer happened to
5
# hold and may treat a context segment as a generation step.
6
n = self.num_seqs
7
self.host_request_types[:n].fill_(0)

host_request_types 是一个 per-sequence 的类型标记数组（0 = context，1 = generation），它的用途是告诉内核”这一行是预填充还是解码”。它用 torch.empty_like 分配——未清零。完整 prepare() 会把它填满，而快路径跳过了 prepare()，于是这个数组里全是上一次留下的垃圾值。垃圾值恰好是 1 的话，内核会把一条预填充请求当成解码请求处理：只对最后一个 token 算注意力输出，其余 token 的输出全是垃圾。这就是”快路径”的典型代价——你省掉的每一步初始化，都得自己补回来。

一个被刻意”提前算好”的数字

基类里有一组字段，注释很直白：

1
# These fields are set when changing seq_lens and _num_contexts to avoid computation
2
# during execution. If the calculation happens during execution, torch compile treats it
3
# as DDS and fails to compile.
4
_num_generations: int = field(init=False, default=0, repr=False)
5
_num_ctx_tokens: int = field(init=False, default=0, repr=False)
6
_num_tokens: int = field(init=False, default=0, repr=False)

DDS 是 Dynamo 的 “data-dependent shape” 缩写。含义是：如果一个张量的形状依赖运行时的数据值，torch.compile 就无法为它生成静态形状的代码，只能放弃编译（graph break）。 num_contexts 是每步从调度器传进来的一个 Python int，seq_lens 是张量；如果在前向函数体里现算 seq_lens[:num_contexts].sum().item()，就产生了”形状依赖数据”。

解法是在 setter 里算好：

1
@seq_lens.setter
2
def seq_lens(self, value):
3
    ...
4
    self.on_update()
5

6
def on_update(self):
7
    if (self._seq_lens is not None
8
            and self._seq_lens.shape[0] >= self.num_contexts
9
            and self.num_contexts >= 0):
10
        self._num_ctx_tokens = self._seq_lens[:self.num_contexts].sum().item()
11
        self._num_generations = self._seq_lens.shape[0] - self.num_contexts
12
    if self._seq_lens_kv is not None:
13
        self._num_tokens = self._seq_lens_kv.sum().item()
14
    elif self._seq_lens is not None:
15
        self._num_tokens = self._seq_lens.sum().item()

.item() 是同步操作——它强制 CPU 等 GPU 上的 sum() 算完。把它放在 setter 里（调度阶段，CPU 本来就闲着）而不是前向体里（GPU 正在跑，CPU 等就是浪费），既避开了 DDS，又把同步藏到了不心疼的地方。注意 on_update 在 seq_lens、num_contexts、seq_lens_kv 三个 setter 里都被调用，因为这些量互相依赖——改任何一个，三个派生值都要重算。 这种”在 setter 里维护不变式”的写法在 TensorRT-LLM 里反复出现，第五篇的 InputBuffers 也是同一套路数。

稀疏后端的旁路入口

契约里还有一个稀疏注意力的挂载点：

1
# Sparse settings for metadata allocation/update; dense metadata leaves it None.
2
sparse_metadata_params: Optional[SparseMetadataParams] = None

默认是 None，稠密注意力完全不受影响。稀疏算法（第 next 节详述）通过它注入自己的元数据参数——比如”这次的稀疏粒度是 token 还是 block""索引器的 Top-K 是多少”。这种”默认 None + 可选协议”的写法，是 TensorRT-LLM 给同一个元数据类挂载十来种注意力变体（MLA、MHA、稀疏、VSWA、FP4）的通用手法：契约不膨胀，扩展靠旁路。

prepare() 逐段拆解

FlashInfer 的 prepare() 是整个文件里最长的一段，按顺序拆开看：

步骤一：qo_indptr（query 侧的 ragged 边界）

1
assert self.seq_lens_cuda is not None
2
torch.cumsum(self.seq_lens_cuda, dim=0, dtype=torch.int32,
3
             out=self._qo_indptr[1:self.seq_lens_cuda.size(0) + 1])

注意这是设备上的 cumsum，且 out= 指向预分配缓冲。两个细节各有原因：设备上算是因为要用 seq_lens_cuda；用 out= 是因为 CUDA Graph 要求地址稳定。cumsum 天然给出 prefix sum，配合 _qo_indptr[0] = 0（预分配时已置零），indptr 就成形了。

步骤二：无 KV 缓存的早退分支

1
if self.kv_cache_manager is None:
2
    assert self.request_ids is not None
3
    assert self.num_generations == 0, (
4
        "FlashInfer without a KV cache manager only supports context-only "
5
        "batches (num_generations == 0) in TRT-LLM.")
6
    ...
7
    self.kv_cache_params = KVCacheParams(use_cache=False)
8
    n = self.num_seqs
9
    self._cached_token_lens[:n].zero_()
10
    self.num_ctx_cached_tokens = 0
11
    self._clean_cached_plans(defer_plan=False)
12
    return

没有 KV 缓存的 FlashInfer 只支持 context-only 批。原因很直白：解码阶段的每步都要读写 KV 缓存，没有缓存就没有东西可读，解码无从进行。所以这个分支必须断言 num_generations == 0，并且把”已缓存 token 数”全部清零——不清零的话，内核会以为这些序列已经有一堆历史 KV（实际没有），把后续 token 写到错误位置。

步骤三：从账本取缓存长度，算出本步 KV 长度

1
num_cached_tokens_per_seq = self.kv_cache_params.num_cached_tokens_per_seq
2
...
3
kv_lens_host = np.asarray(num_cached_tokens_per_seq, dtype=np.int64) + self.seq_lens_kv.numpy()
4
logical_num_blocks = ((kv_lens_host + self.page_size - 1) // self.page_size)

这一段有两条设计要点。

其一，全程在主机侧用 NumPy 算。 注释写得很清楚：

1
# Kept on the host: every consumer below needs host values, so a device-side computation
2
# would force a sync per step.

每次算完如果留在设备上，下游要用就得 D2H 一次；而下游有十来个消费者，就是十几次同步。一次性在主机上算完，只做一次 H2D（拷贝设备侧要用的那几份），总同步次数降到零。

其二，logical_num_blocks 是”逻辑块数”，不是”已分配块数”。 注释强调：

1
# These logical block counts describe committed history, not
2
# all pages that the KV manager has reserved for fixed-width speculative appends.

区别在哪？投机解码时，KV manager 会预先把草稿阶段的块也分配好（因为草稿长度是固定的，可以提前占位）。但这些块在草稿 token 被接受之前，里面的数据是未提交的。如果按”已分配的块数”来算，内核就会去读这些未提交的槽位——读到的是上一轮的残值，等于把无效 KV 混进注意力计算。所以必须用”已经真正写入的历史长度”来算逻辑块数。

步骤四：预留宽度与”全页表”

1
use_full_generation_page_table = (
2
    self.kv_cache_params.use_full_generation_page_table)
3
self._uses_full_generation_page_table = use_full_generation_page_table
4
if use_full_generation_page_table:
5
    # Generation rows need every reserved page to keep a speculative append addressable
6
    # across a page boundary. Treating that wider table as the logical sequence would make
7
    # FlashInfer attend to uncommitted draft slots, so plans and positions use the separate
8
    # device logical lengths populated below.
9
    assert self.request_ids is not None
10
    num_blocks = np.asarray(
11
        _get_page_table_num_blocks(
12
            self.kv_cache_manager, self.request_ids,
13
            logical_num_blocks.tolist(), self.num_contexts),
14
        dtype=np.int64)
15
self.num_blocks = num_blocks.tolist()

这里出现了两个不同的块数：logical_num_blocks（真实历史）和 num_blocks（页表宽度）。为什么要分开？

_get_page_table_num_blocks 的实现回答了这个问题：

1
def _get_page_table_num_blocks(kv_cache_manager, request_ids,
2
                               logical_num_blocks: list[int],
3
                               num_contexts: int) -> list[int]:
4
    """Keep context rows logical and expose every reserved generation page."""
5
    generation_request_ids = request_ids[num_contexts:]
6
    reserved_num_blocks = [
7
        len(block_ids) for block_ids in
8
        kv_cache_manager.get_batch_cache_indices(generation_request_ids)
9
    ]
10
    return list(logical_num_blocks[:num_contexts]) + [
11
        max(logical, reserved) for logical, reserved in zip(
12
            logical_num_blocks[num_contexts:], reserved_num_blocks)
13
    ]

context 行保持逻辑块数，generation 行取”逻辑块数”与”预留块数”的较大者。

对 context（预填充）行：它们每一步都在追加新 token，没有”预留给还没生成的 token”这回事，所以逻辑块数就是真实块数。

对 generation（解码）行：投机解码要跨页边界可寻址。设某条序列已经写了 33 个 token，页大小 32，那么逻辑上它占了 2 页（第 2 页只用了 1 个 slot）。现在草稿阶段要一次追加 8 个 token。如果页表只暴露 2 页，第 3 页根本不在表里——草稿 token 写到一半需要跨页时，内核会去查一个不存在的页，结果是写到别处或直接越界。所以页表必须暴露预留的全部页。

但页表变宽又带来新问题，代码注释点破了：

1
Treating that wider table as the logical sequence would make FlashInfer attend to uncommitted draft slots

如果内核把”页表覆盖的全部页”当成”序列的逻辑长度”，它就会去 attend 那些还没被接受、内容是残值的草稿槽位——等于把垃圾 KV 参加进注意力计算，输出直接错。 所以解法是把两件事拆开：页表用宽版（保证可寻址），而位置计算与 plan 用另一份单独维护的”逻辑长度”张量：

1
if use_full_generation_page_table:
2
    logical_kv_lens = _to_int32_tensor(kv_lens_host)
3
    self._logical_kv_lens[:logical_kv_lens.numel()].copy_(
4
        logical_kv_lens, non_blocking=True)
5
    position_kv_lens = self._logical_kv_lens[:self.num_seqs]
6
else:
7
    position_kv_lens = flashinfer.get_seq_lens(
8
        self.paged_kv_indptr, self.paged_kv_last_page_len, self.page_size)

这是一条很值得记住的工程原则：当”可寻址范围”和”有效数据范围”必须不一致时，用两个显式的量分别表示它们，绝不用一个量兼任两个含义。 同一个数字承担两重语义，迟早会在某个边界条件下出错——而且这种错误往往表现为”结果略微不对”，极难定位。

步骤五：块索引与 page table indptr

1
paged_kv_indices = self.kv_cache_manager.get_batch_cache_indices_flat(
2
    self.request_ids, self.num_blocks, layer_idx=primary_layer_idx)
3
self._paged_kv_indices[:paged_kv_indices.size(0)].copy_(
4
    paged_kv_indices, non_blocking=True)

get_batch_cache_indices_flat 的实现（第四篇讲过 get_batch_cache_indices，这是它的扁平版）：

1
def get_batch_cache_indices_flat(self, request_ids, num_blocks, layer_idx=None):
2
    """Concatenated per-request block tables, trimmed to real widths.
3

4
    Equivalent to concatenating
5
    ``get_batch_cache_indices(request_ids, layer_idx)[i][:num_blocks[i]]``
6
    over all requests into one CPU int32 tensor; matches the interface
7
    of ``KVCacheManagerV2.get_batch_cache_indices_flat``."""
8
    block_ids_per_seq = self.get_batch_cache_indices(
9
        request_ids, layer_idx=layer_idx, num_blocks_per_seq=num_blocks)
10
    indices_list = []
11
    for block_ids, n in zip(block_ids_per_seq, num_blocks):
12
        indices_list.extend(block_ids[:n])
13
    return torch.tensor(indices_list, dtype=torch.int32)

注意 num_blocks_per_seq=num_blocks 这个参数——它把每条序列的块数裁到我们上一步算出的宽度。裁掉的部分是”该请求持有但本次不需要暴露的块”。这就解释了为什么不能像 Vanilla 那样用 len(block_ids) 当块数，第四篇的注释也提过同一件事：“len(block_ids) 会给你一个太大的数，因为 chunked prefill / KV 复用会提前把所有块都分配好。”

紧接着构造 indptr：

1
paged_kv_indptr_decode = _to_int32_tensor(
2
    np.concatenate([[0], np.cumsum(num_blocks[self.num_contexts:])]))
3
self.paged_kv_indptr_decode[:...].copy_(paged_kv_indptr_decode, non_blocking=True)
4
self._host_paged_kv_indptr_decode = paged_kv_indptr_decode

解码段和预填充段各自一个 indptr，因为两段用的内核不同（decode 走 split-K 或 tensor-core 解码内核，prefill 走 paged-prefill 内核），各自的 varlen 边界不一样。

而两段混合时（mixed batch）又要造第三个：

1
else:
2
    assert not self.is_cuda_graph, "Cannot mix decode/prefill with CUDA graphs"
3
    # Accumulate on the host and stage through pinned memory: .cuda()
4
    # on an unpinned tensor is a synchronous H2D that stalls
5
    # the executor thread behind in-flight kernels on every mixed step.
6
    self.paged_kv_indptr = _to_int32_tensor(
7
        np.concatenate([[0], np.cumsum(num_blocks)])).to(device='cuda', non_blocking=True)

两个信息：其一，混合批不能配 CUDA Graph。 因为图必须捕获静态形状，而”混了多少 context、多少 generation”是每步变的，形状定不下来。其二，非 pinned 张量的 .cuda() 是同步的。 这是一个容易被忽略的性能细节：PyTorch 的 .cuda() 在源不是 pinned 内存时，会先隐式做一次同步 H2D，执行线程得等在那里；而 pinned 内存允许 true async。所以这里刻意绕道 _to_int32_tensor（内部用 maybe_pin_memory）。

步骤六：位置与 batch 索引

1
batch_indices, positions = flashinfer.get_batch_indices_positions(
2
    self.kv_indptr, position_kv_lens, self.num_tokens)

这两个数组是 FlashInfer 的 append 内核要的：新算出的 K/V 要写进分页池，需要知道”第 i 个新 token 属于哪条序列（batch_index）""它在该序列里是第几个位置（position）“。因为新 token 在批里是紧凑排列的，而它们在各自序列里的位置是跳着的，这个”从紧凑索引到逻辑位置”的映射必须显式算出来。

步骤七：plan 的重新规划

1
# Re-plan MLA wrappers outside of forward/capture using the params
2
# cached by prior warmup forwards.
3
self._mla_ragged_planned = False
4
self._mla_context_planned = False
5
self._mla_decode_planned = False
6
if (self.num_contexts > 0 and self._mla_ragged_plan_params is not None
7
        and self._ragged_prefill_wrapper is not None):
8
    ragged_indptr = self.qo_indptr[:self.num_contexts + 1]
9
    self._do_plan_ragged(ragged_indptr, ragged_indptr, self._mla_ragged_plan_params)

FlashInfer 的 wrapper 有一个”plan”阶段（选内核、算 tile 划分），它本身可能触发同步或设备操作。在 prepare() 里做（图捕获之前），而不是在前向体里做——因为前向体被捕获进 CUDA Graph 之后，任何不该固定的东西都会被冻进图里。这段代码的三次 _do_plan_* 分别对应 ragged prefill、MLA context、MLA decode 三种 wrapper。

步骤八：VSWA 的多池索引切换

1
if self._vswa_layer_to_pool is not None:
2
    unique_pools = set(self._vswa_layer_to_pool.values())
3
    ...
4
    for pool_id in unique_pools:
5
        if pool_id == primary_pool_id:
6
            continue
7
        rep_layer = self._vswa_pool_to_rep_layer[pool_id]
8
        pool_indices = self.kv_cache_manager.get_batch_cache_indices_flat(
9
            self.request_ids, self.num_blocks, layer_idx=rep_layer)
10
        ...
11
        buf = getattr(self, f'_vswa_pool_buf_{pool_id}')
12
        buf[:pool_indices.size(0)].copy_(pool_indices, non_blocking=True)

可变滑动窗口（VSWA）时，不同层用不同窗口，也就落在不同的块池里。每层要用属于自己池的那份块索引，所以 prepare() 为每个池各准备一份索引张量，存在 _vswa_pool_indices_cache 里；前向时每层按 layer_to_pool 映射去取对应那份，并把它拷回 _paged_kv_indices（forward_impl 里的 copy_()，代码注释说明这是为了 device-to-device 且可被 CUDA Graph 捕获）。

步骤九：自检

1
self._debug_validate_page_table(kv_lens_host, logical_num_blocks, num_blocks)

把”逻辑块数""页表块数""真实 KV 长度”三者交叉核对一遍。这类自检默认只报有限次数：

1
_page_table_check_budget = _PAGE_TABLE_CHECK_REPORTS
2
...
3
if cls._page_table_check_budget <= 0:

用一个类级别的预算计数器，报告若干次后就不再报告，避免在长期运行的服务里把日志刷爆。这是个很务实的细节：调试信息的价值随时间衰减，但成本（字符串拼接、D2H 取值）是持续的。

prepare_flash_mla：块表怎么喂给 FlashMLA

启用 FlashMLA（DeepSeek 系 MLA 模型的高性能解码内核）时，块偏移还不够，要另外准备块 id 表：

1
def prepare_flash_mla(self) -> None:
2
    self._flash_mla_metadata_valid = False
3
    # Request-specific fills and H2D copies must happen before replay, not
4
    # become fixed operations in the captured graph.
5
    if torch.cuda.is_current_stream_capturing():
6
        return
7

8
    block_ids_per_seq = maybe_pin_memory(
9
        self.kv_cache_manager.get_block_ids_per_seq(self.request_ids))
10
    num_blocks = block_ids_per_seq.shape[1]
11
    self.kv_block_ids_per_seq.fill_(0)
12
    self.kv_block_ids_per_seq[:self.num_seqs, :num_blocks].copy_(
13
        block_ids_per_seq, non_blocking=True)
14
    self.block_ids_per_seq.fill_(0)
15
    self.block_ids_per_seq[:self.num_generations, :num_blocks].copy_(
16
        block_ids_per_seq[self.num_contexts:], non_blocking=True)

get_block_ids_per_seq 是二维表版本：

1
def get_block_ids_per_seq(self, request_ids: List[int]) -> torch.Tensor:
2
    block_ids_per_seq = self.get_batch_cache_indices(request_ids)
3
    block_ids_per_seq_tensors = [
4
        torch.tensor(sublist, dtype=torch.int) for sublist in block_ids_per_seq
5
    ]
6
    padded_tensor = torch.nn.utils.rnn.pad_sequence(
7
        block_ids_per_seq_tensors, batch_first=True, padding_value=0)
8
    return padded_tensor

注意 padding_value=0——短的那些行用块编号 0 填充。所以块编号 0 必须是”可以安全读到”的（通常是一个永远预留的空块），否则被 padding 的行会让内核访存到脏数据。这是一个隐含的全局约定：编号 0 的块是”垃圾但合法”的哨兵块。

这个实现里有三处值得记的设计：

其一，if torch.cuda.is_current_stream_capturing(): return。 这是整个文件里最直白的一行防御：如果当前正在捕获 CUDA Graph，就不要做这些带 H2D 的准备工作。 原因写在注释里——“请求相关的填充和 H2D 拷贝必须发生在重放之前，不能变成捕获进图里的固定操作”。如果被捕获，那么图里就固定了”拷贝某一次特定批次的块表”这个动作，重放时永远拷同一份数据，后面的批次全用错表。

其二，fill_(0) 再 copy_ 的固定套路。 先清零整个缓冲区，再覆写前 num_blocks 列。清零是为了让”本步没用到的列”保持哨兵值 0——如果不清零，那些列里留着上一步的块编号，而内核在某些分支下可能访问到它们（比如 padding 行），就会读到上一批的块。

其三，两份块表 + 精确的行切片。

1
self.kv_block_ids_per_seq[:self.num_seqs, :num_blocks]              # 全部序列
2
self.block_ids_per_seq[:self.num_generations, :num_blocks]          # 只解码序列

kv_block_ids_per_seq 装全批（context + generation），block_ids_per_seq 只装 generation 段——用 block_ids_per_seq[self.num_contexts:] 切片，然后写进第 0 行开始的位置。为什么？因为MLA 的注意力计算只发生在解码序列上（预填充的 attention 走另一条路径）。让块表的第 0 行对应第一条解码请求，可以让内核直接按 generation 索引，不必再做 - num_contexts 的偏移修正。把偏移藏在准备阶段，而不是让内核每层都算一次——这是在多层的模型里能省下可观指令数的选择。

六种稀疏算法与它们的粒度

文档列出的公开算法，按”选择粒度”排序很有启发：

粒度从细到粗：token → 块 → 压缩条目 → 内核 tile。粒度越细，稀疏率可以越高，选择和索引的开销越大。token 级要维护一个 per-token 的打分与 Top-K，索引数组很长；块级（128 token 一格）的选择粒度换来的是索引数组短了 128 倍、内核更好向量化，代价是最小稀疏单位变大——即使一个块里只有一个重要 token，整块也不能丢。

而 skip_softmax 是另一类：它不选任何东西，而是在 FMHA 内核内部判断”这一块 tile 的 softmax 贡献是否可忽略”，可忽略就整块跳过计算。它的”选择”发生在最内层，连索引数组都不需要。这也解释了它为什么不需要 model-trained selector——它不依赖模型学出来的重要性，只依赖 softmax 本身的数值分布。

上图是 TensorRT-LLM 官方文档的滑动窗口注意力示意。deepseek_v4、csa2 这类算法的”滑窗分支”就是它：最近的 raw token 走稠密滑窗，更远的历史交给压缩后的稀疏/稠密分支。它对块表的要求与全注意力不同——窗口滑走后，早期块可能整块不再需要暴露在页表里，这也是第四篇讲的 VSWA 会按窗口大小分池的原因。图片来源：TensorRT-LLM 官方文档 media/feat_long_seq_sliding_win_attn.png。

GVR：Top-K 的三级替换路径

DSA 的文档里还提到一个值得注意的实现细节：

On supported Blackwell configurations, Guess-Verify-Refine (GVR) can replace the regular decode Top-K dispatcher. The current implementation accepts index_topk values 512, 1024, and 2048, and indexer compression ratios 1 and 4. Unsupported combinations fall back to the production insertion/radix Top-K path.

这段对应 sparse/params.py 里的一个函数：

1
def use_self_sampling_gvr(*, enable_heuristic_topk, use_self_sampling_topk,
2
                          index_topk, compress_ratio, is_cute_dsl_available,
3
                          sm_version) -> bool:
4
    """Return whether the two-level dispatch picks the self-sampling engine.
5

6
    Shared by the indexer (per-layer TopK construction) and the attention
7
    metadata (prior-state allocation and warmup) so both sides of the
8
    dispatch agree."""
9
    return (enable_heuristic_topk
10
            and use_self_sampling_topk
11
            and is_cute_dsl_available
12
            and sm_version in (100, 103)
13
            and index_topk in (512, 1024, 2048)
14
            and compress_ratio in (1, 4))

这个函数的存在意义比它的逻辑更重要。注释说明它”被索引器与注意力元数据共享，让 dispatch 的两侧达成一致”。

这是个很容易踩的坑：索引器要在每层构造 Top-K，需要知道走哪条路径；注意力元数据要做”prior-state 分配与 warmup”，也需要知道走哪条路径（因为两条路径需要的辅助缓冲不同）。如果两侧各自实现一遍判定条件，那么一旦条件写得不一样（比如一边漏了 sm_version 判断），就会出现”索引器按 GVR 路径构造了索引，而元数据按常规路径分配了缓冲”——缓冲大小对不上，直接越界。

所以作者把判定抽成一个函数，两侧调用同一个。“让两侧独立推导出同一结论”是这个代码库里反复出现的模式——第十八篇讲弹性 EP 时，joiner 与主部署各自算 join_effective_ep_size 然后核对，也是同一个思路。

注意 use_self_sampling_gvr 里还有一个隐藏的一致性风险，被显式点名了：sm_version（GPU 架构版本）与 is_cute_dsl_available（CuTe DSL 是否可用）都是环境相关的量。同一台机器上、同一个进程里它们不变，所以”两侧算出同一结论”是安全的；但如果把判定结果跨进程或跨实例共享（比如缓存到某个配置文件里），环境一变就会失效。这类”环境派生量”必须在使用点现算，不能缓存。