TensorRT-LLM 源码完全拆解（八）：KV Cache 的写入与读取路径——一次前向里 K/V 如何落进分页池、又如何被内核取回

为什么写入要”合并”

一个自然的疑问是：加 bias、施加 RoPE、量化、写池，这四件事为什么要塞进同一个内核，而不是拆成四个小内核顺序执行？

答案在 kernel launch 与访存的账上：

• 中间张量的显存流量。如果拆成四个内核，每一步都要把整个 QKV 张量读一遍、写一遍。以 num_tokens = 4096、hidden = 8192、BF16 为例，qkv_input 是 $4096 \times 8192 \times 2 = 64$ MB，读 + 写一轮就是 128 MB 的流量。四步串行就多出三倍的显存往返。
• kernel launch 开销。每个内核的启动本身有微秒级开销，解码阶段单步的 GPU 计算时间本身就只有几百微秒，几个额外内核的启动开销占比可观。
• 融合之后中间值留在寄存器里。bias 加完的 q、RoPE 旋转后的 k，全程在寄存器中传递，只有最终的 Q（写回）和 K/V（写池）落地。

这三条理由对预填充阶段（token 多、张量大）是显存带宽账，对解码阶段（token 少、单步时间短）是启动开销账。TensorRT-LLM 把两件事用一个内核覆盖，靠的是模板参数区分：同样的 kernel 代码，编译出”预填充变体”和”解码变体”两份二进制。

线程组织：一个 block 负责一个头

内核的并行划分是三维的：

1
int const head_idx = blockIdx.y;                       // 一个 block 负责一个头
2
int const tidx = threadIdx.x;                          // x 维：head_dim 方向（固定 32 个线程）
3
int local_token_idx = blockIdx.x * blockDim.y + threadIdx.y;  // y 维：token 方向
4
// blockIdx.z 维：batch（以 gridDim.z 为步长循环跨越）

核心设计是 x 维固定 32 个线程，每个线程负责 VEC_SIZE = size_per_head / 32 个元素。以 head_dim = 128、BF16 为例：VEC_SIZE = 4，每个线程处理 4 个 BF16 元素，打包成一个 uint16_t（VecType<half> 的 GPTNeoXEltType = uint16_t，VecType<__nv_bfloat16> 的 Type = mmha::bf16_8_t）。这个 32 的来源有实际理由：

• RoPE 的配对语义需要跨线程通信。GPT_NEOX 的旋转是”前半维与后半维配对”，如果 head_dim 维的线程数是 8 或 16，配对的两半可能落在同一个 warp 的不同 lane 上，就得走 shuffle。固定 32 则让配对关系与 warp 边界对齐，两个线程各自持有配对的两半，在寄存器层面完成旋转，必要时才退到 shared memory。
• warp 内的 32 个 lane 恰好铺满 head_dim。对 GQA 场景（qheads_per_kv_head = 4 之类），多个 Q 头共享一个 KV 头，写入时只有”该 KV 头对应的第一个 Q 头”那个 block 才执行写池（见下面的代码），其余 block 白跑一趟——这是 GQA 在写入侧的固有浪费，也是为什么写入侧会做”head_num == kv_head_num 就直接写”的判断来加速。

token 到池位置的映射：三步算术

这是整个写入内核里最值得逐行读的部分，因为它把”逻辑 token 序号”翻译成了”池里的物理位置”：

1
// The index of the token in the batch.
2
int const global_token_idx = local_token_idx
3
    + ((variable_sequence_length && params.remove_padding)
4
           ? params.cu_seq_lens[batch_idx]        // 变长 + 去 padding：用累积偏移
5
           : batch_idx * params.max_input_seq_len); // 否则：按最大长度整齐排布
6

7
int const cache_seq_len = params.cache_lens[batch_idx];   // 写完之后池里应该有多少个 token
8
int const actual_seq_len = variable_sequence_length ? params.seq_lens[batch_idx]
9
                                                    : params.max_input_seq_len;
10
int const past_seq_len = (cache_seq_len - actual_seq_len);  // 这一步之前池里已有的 token 数
11
int const token_idx_in_seq = past_seq_len + local_token_idx;
12
bool const valid_token = token_idx_in_seq < cache_seq_len;

三个量必须分清，它们代表了三个不同的”长度”：

• actual_seq_len：本次前向新输入的 token 数（预填充时是提示长度，解码时是 1）。
• past_seq_len：这次之前池里已经有的 token 数。它等于 cache_seq_len - actual_seq_len——注意这个式子本身就是”池里有多少”的定义，而不是从别处查来的。
• cache_seq_len：写完之后池里应该有的 token 数，即 past + actual。它是调度器算好、由 cache_lens 传进来的。

为什么 past_seq_len 要这么算？因为分块预填充（chunked prefill）会把一条长提示切成几段分批送入。第一批处理 token 0…N-1，第二批处理 N…2N-1，此时第二批的 past_seq_len = N，新 token 应该写进池的第 N 个位置往后。没有这行减法，第二批就会从头覆写第一批的 KV——这是分块预填充最经典的实现错误，也是”一个不报错但算错”的典型。

RoPE 的施加位置：为什么不是”每 token 查表”

在拿到 token_idx_in_seq 之后，内核算的是旋转位置，而不是直接用它：

1
int const rotary_position
2
    = (params.helix_position_offsets != nullptr ? params.helix_position_offsets[global_token_idx]
3
       : params.spec_decoding_position_offsets != nullptr
4
           ? (params.spec_decoding_position_offsets[local_token_idx + batch_idx * params.max_input_seq_len]
5
              + past_seq_len)
6
           : token_idx_in_seq)
7
    + (params.mrope_position_deltas != nullptr ? params.mrope_position_deltas[batch_idx] : 0);

三层来源，每一层都对应一种特殊场景：

1. helix_position_offsets（Helix 并行）：当模型按 token 维切分到多个 rank 时，本 rank 看到的”局部 token 序号”不等于全局位置，必须查一张全局位置表。
2. spec_decoding_position_offsets（投机解码）：一次前向要验证 $d$ 个草稿 token，它们的 Q 各自处在不同的未来位置上。这个偏移数组给出每个草稿位置的相对偏移，加上 past_seq_len 得到绝对位置。
3. 默认：就是 token_idx_in_seq，即”池里的第几个 token”，对普通自回归解码而言这正是它的绝对位置。

而 mrope_position_deltas 是多模态位置修正（M-RoPE）：图像 token 占据的是二维/三维位置网格，文本位置需要一个 delta 来补偿，否则图文混合序列的位置会错位。

关键设计点：旋转系数是”逐位置”而非”逐 token”缓存的。 代码里有一个 float2 rotary_coef_cache[ROTARY_COEF_VEC_SIZE] 的寄存器数组，配合 cached_rotary_position：

1
if (DYNAMIC_ROTARY_SCALING || rotary_position != cached_rotary_position) {
2
    apply_rotary_embedding_gptj<...>(q, k, rotary_coef_cache, ..., rotary_position);
3
    cached_rotary_position = rotary_position;
4
} else {
5
    apply_rotary_embedding_gptj<..., /*计算系数=*/false>(q, k, rotary_coef_cache, ..., rotary_position);
6
}

在解码阶段同一个 token 的所有头共享同一个位置，所以同一线程连续处理不同 head 时，三角函数只需算一次，后续复用寄存器里的 cos/sin。这是一个”以寄存器换三角函数调用”的优化：__sincosf 本身不便宜，而解码阶段每步要调用的次数是 kv_head_num × half_rotary_dim / VEC_SIZE 级别。

三类精度分支：同一个 store 语句的三副面孔

拿到旋转过的 k 与原始 v 之后，写入池的分支长这样：

1
if (valid_kv_cache_pos && !helix_inactive) {
2
    if constexpr (ENABLE_8BITS_CACHE) {
3
        inBlockIdx = inBlockIdx * VEC_SIZE;                 // 元素索引 → 字节索引
4
        mmha::store_8bits_vec(kDst, k_to_cache, inBlockIdx, scaleOrigQuant);
5
        mmha::store_8bits_vec(vDst, v, inBlockIdx, scaleOrigQuant);
6
    } else if constexpr (ENABLE_4BITS_CACHE) {
7
        auto* kBlockScales = reinterpret_cast<uint8_t*>(
8
            params.kv_cache_block_scales_buffer.getKBlockPtr(batch_idx, token_idx_in_kv_cache));
9
        auto* vBlockScales = ...;
10
        float kSecondLevelSF = params.qkv_scale_orig_quant[1];   // K 的第二级 scale
11
        float vSecondLevelSF = params.qkv_scale_orig_quant[2];   // V 的第二级 scale
12
        quantizeAndWriteFP4KVCache<T, VECS_PER_HEAD>(kBlockScales, vBlockScales,
13
            reinterpret_cast<uint32_t*>(kDst), reinterpret_cast<uint32_t*>(vDst),
14
            kSecondLevelSF, vSecondLevelSF, inBlockIdx, kPacked, vPacked);
15
    } else {
16
        reinterpret_cast<VecType*>(kDst)[inBlockIdx] = k_to_cache;   // 直接向量写
17
        reinterpret_cast<VecType*>(vDst)[inBlockIdx] = v;
18
    }
19
}

三个分支对应第七篇讲的三种存储类型，注意它们在索引单位上的差异：

• FP8/INT8 路径：inBlockIdx * VEC_SIZE——第七篇提过，量化是逐元素（每元素一字节），而 getKVLocalIdx 返回的是按 sizePerHeadDivX 为元素单位算的索引（sizePerHeadDivX = size_per_head / VEC_SIZE）。所以这里要乘回 VEC_SIZE 才能从”向量槽位号”变成”字节偏移”。
• FP4 路径：走 quantizeAndWriteFP4KVCache 辅助函数，同时写主池与 scale 池。注意 kSecondLevelSF 取自 qkv_scale_orig_quant[1]、vSecondLevelSF 取自 [2]——K 和 V 各有自己的第二级 scale，因为 K 与 V 的数值分布完全不同（K 有 RoPE 旋转、V 是原始投影值），共用一个全局 scale 会让其中一方精度塌陷。这正是第七篇”两级缩放”在写入侧的落地。
• BF16/FP16 路径：直接向量写，inBlockIdx 就是向量下标，不需要换算。

FP8_OUTPUT 分支里还有一处细节：mmha::store_8bits_vec 在 STORE_QKV 场景下也会把 Q 与 K/V 一起量化写回 quantized_qkv_output。这意味着 FP8 路径下 Q 也走 8 位存储，但Q 走的是输出缓冲而不是池——两条通路共用了同一个”量化 + 存储”辅助函数，但目的地不同。

zeroPagedVCacheTail：为什么要主动清零

在写入分支的最后，有一行不起眼但极其重要的调用：

1
zeroPagedVCacheTail<T, TCache, VecType, GEN_PHASE>(params, batch_idx, token_idx_in_kv_cache,
2
    cache_seq_len, kv_head_idx, sizePerHeadDivX, channelIdx, /*isUniqueTokenWriter=*/true);

它的实现（节选）与注释直说了动机：

1
// Context owns the last physical page exclusively. Clear unused V rows because paged FMHA loads full pages and
2
// masked zero probabilities can still propagate NaNs through BMM2.
3
if constexpr (!GEN_PHASE && std::is_same_v<KVCacheBuffer, KVBlockArray> && std::is_same_v<T, TCache>) {
4
    ...
5
    int const validTokensInLastBlock = params.kv_cache_buffer.getLocalIdx(cacheSeqLen - 1) + 1;
6
    if (validTokensInLastBlock >= tokensPerBlock) return;
7
    auto* vBlock = ...getVBlockPtr(batchIdx, cacheSeqLen - 1);
8
    TVec const zero{};
9
    for (int tokenInBlock = validTokensInLastBlock; tokenInBlock < tokensPerBlock; ++tokenInBlock) {
10
        int const vIdx = kvHeadIdx * tokensPerBlock * vecsPerHead + tokenInBlock * vecsPerHead + channelIdx;
11
        vBlockVec[vIdx] = zero;
12
    }
13
}

条件是预填充阶段 + 分页池 + 池与输入同精度（量化路径另有处理）。做的是：把最后一个物理块里超出有效长度的那几行 V 清零。

为什么要清零？注释给的理由是”paged FMHA loads full pages and masked zero probabilities can still propagate NaNs through BMM2”。展开讲：

1. 分页注意力内核按整块加载 V（页大小通常 32/64/128），不会只加载有效行——因为按块加载才能连续访存，逐行加载会退化成非合并访问。
2. 注意力权重的掩码把无效行对应的 softmax 概率压成 0。数学上 $0 \times \text{NaN} = \text{NaN}$——如果那些无效行里恰好是上一次请求残留的脏数据或者未初始化的 NaN，乘 0 之后仍然会把整行的输出污染成 NaN。
3. 所以必须在写入时主动把这些位置清零，而不是依赖”反正权重是 0”。

第 1 条同时解释了为什么是”清零 V 而不清零 K”：K 参与的是 $QK^T$，无效行会被掩码加成 $-\infty$，softmax 后为 0，它自身取什么值不影响结果（只要不是 NaN/Inf 引起溢出）；而 V 是加权求和的对象，会直接乘进输出。实现里也确实只处理 V 块。

再看条件 !GEN_PHASE：解码阶段不清理。原因是解码时每个块都已被写入过至少一次（块是按需分配、用完就走的），不会有”新建但没写满”的块；而预填充阶段一条请求的最后一个块通常是半满的，必须清。另外 isUniqueTokenWriter 这个参数提示了一种并行风险：如果多个 block 同时写同一个块的不同行，清理者必须唯一，否则会互相清掉对方刚写的数据。

一个 kernel 同时覆盖预填充与解码

AppendPagedKVCacheKernel 的全部循环体只有十来行：

1
__global__ void AppendPagedKVCacheKernel(paged_kv_t<DType, IdType> paged_kv,
2
                                         DType* __restrict__ append_key,
3
                                         DType* __restrict__ append_value,
4
                                         IdType* __restrict__ batch_indices,
5
                                         IdType* __restrict__ positions, uint32_t nnz,
6
                                         size_t append_k_stride_n, size_t append_k_stride_h,
7
                                         size_t append_v_stride_n, size_t append_v_stride_h) {
8
  uint32_t tx = threadIdx.x, ty = threadIdx.y;
9
  uint32_t num_heads = paged_kv.num_heads;
10
  uint32_t head_idx = ty;
11
  uint32_t cta_id = blockIdx.x;
12
  uint32_t num_ctas = gridDim.x;
13

14
#pragma unroll 4
15
  for (uint32_t i = cta_id; i < nnz; i += num_ctas) {
16
    uint32_t page_iter, entry_idx;
17
    paged_kv.page_size.divmod(paged_kv.indptr[batch_indices[i]] * paged_kv.page_size + positions[i],
18
                              page_iter, entry_idx);
19
    DType* k_ptr = paged_kv.get_k_ptr(page_iter, head_idx, entry_idx, tx * vec_size);
20
    DType* v_ptr = paged_kv.get_v_ptr(page_iter, head_idx, entry_idx, tx * vec_size);
21
    vec_t<DType, vec_size>::memcpy(
22
        k_ptr, append_key + i * append_k_stride_n + head_idx * append_k_stride_h + tx * vec_size);
23
    vec_t<DType, vec_size>::memcpy(
24
        v_ptr, append_value + i * append_v_stride_n + head_idx * append_v_stride_h + tx * vec_size);
25
  }
26
}

逐个参数理解：

• nnz：要写入的 (token, 序列) 对总数，即”总共有多少行要落盘”。预填充时是所有请求的提示长度之和，解码时是批大小。
• batch_indices[i]：第 i 行属于哪条请求。
• positions[i]：第 i 行是该请求内的第几个 token（绝对位置，从 0 数）。
• indptr：这条请求在页表里的起始页表项下标（也就是第六篇讲的 kv_indptr）。indptr[b] 到 indptr[b+1] 是请求 b 在 kv_indices 里占的区间。

核心算术：divmod 一行拆解

整个分页写入的灵魂是这一行：

1
paged_kv.page_size.divmod(paged_kv.indptr[batch_indices[i]] * paged_kv.page_size + positions[i],
2
                          page_iter, entry_idx);

拆开看，左边这个量是：请求 b 的第一页之前的全部 token 数（indptr[b] * page_size）加上本 token 在请求内的位置（positions[i]）= 该 token 在”逻辑全局 token 坐标系”里的下标。然后对这个下标做除以页大小的 divmod：

• page_iter = 逻辑下标 / page_size → 这是”逻辑页号”，即该 token 落在这条请求的第几个页里。
• entry_idx = 逻辑下标 % page_size → 这是页内偏移，即该 token 在这个页里的第几个槽位。

为什么 indptr[b] * page_size 这个写法成立？ 因为它隐含了一个约定：每条请求在页表里占用的页数，正好等于它已分配 token 数除以页大小向上取整。于是”前面请求分掉的逻辑页数”乘以页大小，就是前面请求占用的逻辑 token 总数——而本请求的逻辑 token 编号从这里接着往下续。

这个约定有个重要后果：逻辑页号是连续的，物理页号不一定连续。page_iter 是一个纯粹的逻辑量，只用来在页表里定位；真正要访问的物理页要靠 paged_kv.indices[page_iter] 这一跳：

1
__device__ __forceinline__ DType* get_k_ptr(IdType page_iter, uint32_t head_idx,
2
                                            uint32_t entry_idx, uint32_t feat_idx) const {
3
    return k_data + get_elem_offset(__ldg(indices + page_iter), head_idx, entry_idx, feat_idx);
4
}

注意 indices + page_iter 传进去的是 page_iter，但取出来的是 indices[page_iter]——一个物理页号，再乘 stride_page。这就是分页的核心一跳：逻辑页号 → 物理页号 → 字节偏移。前两跳靠这张索引表，第三跳靠第七篇讲的 stride。

get_elem_offset 把三跳合成一个式子：

1
__host__ __device__ __forceinline__ size_t get_elem_offset(size_t page_idx, size_t head_idx,
2
                                                           size_t entry_idx,
3
                                                           size_t feat_idx) const {
4
    return page_idx * stride_page + head_idx * stride_h + entry_idx * stride_n + feat_idx;
5
}

四个 stride 全部来自第七篇的池形状：stride_page 是页间步长，stride_h 是头间步长，stride_n 是页内 token 间步长，feat_idx 是 head_dim 内元素偏移（连续，所以没有 stride 变量）。NHD 与 HND 的区别就体现在 stride_h 与 stride_n 的相对大小上：NHD 下 token 维在外、head 维在内，stride_n = head_dim × num_heads、stride_h = head_dim；HND 下正好反过来。内核代码一行没变，换的是三个 stride 的取值——这就是第七篇说”维序是契约字段”在代码层面的样子。

vec_size 与 page_size.divmod

两个实现细节值得单独讲，因为它们都是”看起来是小事、算错了全盘崩”的类型。

第一，vec_size 怎么来的。

1
constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32);
2
uint32_t bdx = HEAD_DIM / vec_size;   // x 维线程数
3
uint32_t bdy = num_heads;             // y 维线程数

vec_size 取两个量的较大者：

• 16 / sizeof(DType)：一个 128 位访问能装多少个元素。BF16 是 16/2 = 8，FP8 是 16/1 = 16。这是”至少要有 16 字节向量化”的下限——不足 16 字节的访问在 GPU 上浪费 cache line。
• HEAD_DIM / 32：让 x 维恰好 32 个线程。head_dim = 128 时是 4，head_dim = 256 时是 8。

取 max 的效果是：head_dim 小的时候由 16 字节下限主导（保证向量宽度），head_dim 大的时候由 32 线程主导（保证线程数不超 32，warp 内对齐）。以 head_dim = 128、BF16 为例 vec_size = 8，bdx = 16、bdy = num_heads；以 head_dim = 64、FP8 为例 vec_size = 16，bdx = 4。这个算式是”访存效率”与”线程组织”之间的一个折中，也是为什么 HEAD_DIM 必须走 DISPATCH_HEAD_DIM 宏在编译期确定——vec_size 是 constexpr。

第二，divmod 不是两条指令。

1
paged_kv.page_size.divmod(linear_index, page_iter, entry_idx);

page_size 在 FlashInfer 里被包装成一个带 divmod 方法的类型，当 page_size 是 2 的幂时它退化成移位与掩码，当不是 2 的幂时才走通用整数除法。为什么页大小默认是 2 的幂（16/32/64/128）？ 因为除法在 GPU 上代价高（几十个周期，且会占用特殊功能单元），而写入内核是每个 token 都要做一次这个除法。页大小取 2 的幂，这一步就变成一条移位加一条掩码——和 TensorRT-LLM 的 getBlockPtr 里那个 offsets[tokenIdx >> mTokensPerBlockLog2] 完全同源：

1
inline void* getBlockPtr(DataType const* offsets, int32_t tokenIdx) const {
2
    auto const offset = offsets[tokenIdx >> mTokensPerBlockLog2];
3
    return reinterpret_cast<void*>(
4
        reinterpret_cast<char*>(getPoolPtr(offset)) + offset.get() * static_cast<uint64_t>(mBytesPerBlock));
5
}

注意这里用了两个不同的信息：tokenIdx >> mTokensPerBlockLog2 只取”页号”（这是自带的连续页表，不含物理页跳转），返回的 offset 则是一个带 isPrimary() 标志的封装值（第七篇的两级池），既给出物理页号又给出它在哪个池里。而页内偏移在这个函数里根本没算——它由调用方的 getKVLocalIdx 单独处理：

1
return headIdx * mTokensPerBlock * dimsPerHead + getLocalIdx(globalTokenIdx) * dimsPerHead + channelIdx;
2
// 其中 getLocalIdx(globalIdx) = globalIdx & ((1 << mTokensPerBlockLog2) - 1)

& (page_size - 1) 就是 % page_size 在 2 的幂下的形态。TensorRT-LLM 把”页号”与”页内偏移”拆成两个 helper，FlashInfer 用一个 divmod 一起返回——两种写法等价，差别只在代码组织。

写入侧的并行划分：为什么用 grid-stride loop

回看内核的循环头：

1
for (uint32_t i = cta_id; i < nnz; i += num_ctas) {

这是一个 grid-stride loop：每个 block 处理第 cta_id、cta_id + num_ctas、cta_id + 2*num_ctas… 行，直到超过 nnz。CTAs 的数量由主机端算出：

1
uint32_t num_threads = bdx * bdy;
2
cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, num_threads, smem_size);
3
num_blocks_per_sm = min(num_blocks_per_sm, ceil_div(int(nnz), num_sms));
4
dim3 nblks(num_blocks_per_sm * num_sms);

三行代码的逻辑是：

1. 问硬件”每个 SM 最多能同时跑几个这样的 block”（cudaOccupancyMaxActiveBlocksPerMultiprocessor）。
2. 取它与 ceil(nnz / num_sms) 的较小者——保证每个 SM 平均至少有一行活干，不会出现大量空转的 block。
3. 乘上 SM 总数，得到最终 grid 大小。

这套算法的价值在于 nnz 的巨大动态范围：预填充一条 32K 的提示，nnz = 32768；解码一个 batch=1 的请求，nnz = 1。如果 grid 大小写死成 nnz（像 AppendPagedKVCacheDecode 那样 dim3 nblks(batch_size)），解码时只有 1 个 block 在跑，GPU 上其余 SM 全部空闲——解码阶段的写入是典型的小批量问题，需要靠更细的线程组织而不是更多 block 来填满机器。反过来预填充时 nnz 很大，用 grid-stride 保证不会启动几万个 block 把启动开销撑爆（FlashInfer 在 TensorRT-LLM 里还有一层额外的 CTA 上限切分，见下面）。

头数太多怎么办：CTA 上限切分

TensorRT-LLM 的 FlashInfer 后端在调用前做了一次包装（_append_paged_kv_cache）：

1
head_dim = append_key.shape[-1]
2
vec_size = max(16 // append_key.element_size(), head_dim // 32)
3
threads_per_head = head_dim // vec_size
4
max_heads_per_launch = _MAX_CUDA_THREADS_PER_BLOCK // threads_per_head
5

6
num_kv_heads = append_key.shape[1]
7
if num_kv_heads <= max_heads_per_launch:
8
    flashinfer.page.append_paged_kv_cache(...)
9
    return
10

11
for start in range(0, num_kv_heads, max_heads_per_launch):
12
    end = min(start + max_heads_per_launch, num_kv_heads)
13
    flashinfer.page.append_paged_kv_cache(
14
        append_key=append_key[:, start:end], append_value=append_value[:, start:end],
15
        paged_kv_cache=_slice_paged_kv_cache_heads(paged_kv_cache, start, end, kv_layout), ...)

这段代码解决一个很具体的约束：内核把 head 放在 threadIdx.y 上，所以线程总数是 bdx × num_heads，而 CUDA 的 block 线程上限是 1024。以 head_dim = 128、BF16、num_kv_heads = 64（DeepSeek 级模型的 KV 头数）为例：vec_size = max(8, 4) = 8，bdx = 16，threads_per_head = 16，max_heads_per_launch = 1024 / 16 = 64——恰好卡在上限。如果 head_dim 更小（比如 64，bdx = 4，即 1024/4 = 256 头才超）或者 KV 头数更多（MLA 的某些配置），就必须切分。

切分时有两个细节：

• append_key[:, start:end] 是切片而非拷贝。head_idx 在内核里从 0 数起，所以切片后的第 0 个 head 对应原张量的第 start 个 head——这对吗？对的，因为 K/V 的 head 维在池里也按同样的顺序排布，切片只是”这次只处理池里的第 start 到 end 个头”。
• paged_kv_cache 也要同步切片（_slice_paged_kv_cache_heads），否则内核算出的 head_idx * stride_h 会指向错误的头。

这里有一个通用教训：当一个 kernel 用线程维去铺某个”逻辑维”时，那个逻辑维的上界就被 CUDA 的 1024 线程上限卡住了。解决方法不是改内核，而是在调用侧切分——前提是池的布局允许按该维切片（即该维在池里是一段连续或规则步长的区间）。这正是第七篇说”池是六维规则张量”的回报：规则性让切分成为可能。

ragged：变长序列怎么分段

上面假设了每行 token 都能独立算位置。但”i 属于哪条请求”这件事本身是需要查表的：

1
paged_kv.indptr[batch_indices[i]] * paged_kv.page_size + positions[i]

batch_indices[i] 把”总行号”映射到”请求号”，positions[i] 把”总行号”映射到”请求内位置”。这两个数组就是 ragged 批处理的全部——它们不需要 indptr 之外的结构，因为 batch_indices 与 positions 已经把每条请求的 token 段”拉平”了。

生成这两个数组的方式，正是第五篇讲的扁平化：cu_seq_lens（累积序列长度）在变长场景下给出每一行属于哪条请求；positions 则由每行的全局位置减去该请求的起点得到。第七篇提过的 kv_indptr 与 kv_last_page_len 则是同一个 ragged 结构在页表维的镜像：

1
__host__ __device__ __forceinline__ uint32_t get_length(uint32_t batch_idx) const {
2
    if (indptr[batch_idx + 1] == indptr[batch_idx]) return 0;
3
    return (indptr[batch_idx + 1] - indptr[batch_idx] - 1) * page_size + last_page_len[batch_idx];
4
}

get_length 算的是”这条请求已经占用的 token 数”，式子是整页数 × 页大小 + 最后一页的有效长度。这个式子与写入侧的 divmod 是互逆的：写入时从”逻辑 token 下标”正推 (页, 页内偏移)，读取时从 (页区间的长度, 最后一页长度) 反推”总 token 数”。两者必须严格一致，否则写入位置和读取范围就会错位——而错位不会报错，只会读出一段看似合理但内容错乱的上下文。

MLA 与 FP4：写入侧的两个额外分支

分页写入在两种模型结构上会走特殊路径。

第一，MLA（Multi-head Latent Attention）。 第七篇讲过 MLA 的 latent 池是 kvFactor=1、numKvHeads=1、sizePerHead=d_latent——一份张量同时当 K 和 V，且每个 token 只有一条 latent 向量加一条 RoPE 分量。写入侧的差别在于：

• 没有”K 和 V 分头写”这回事，只有一条 ckv（compressed KV）加一条 kpe（RoPE 部分），FlashInfer 里对应 append_paged_mla_kv_cache(append_ckv, append_kpe, ...)。
• 池布局变成 [page, entry, ckv_dim] 与 [page, entry, kpe_dim] 两段（或同一段的不同偏移，见 paged_kv_mla_t::get_elem_offset_ckv）。
• 上投影发生在读取侧：写进池的是压缩后的 latent，注意力内核读出来之后才做上投影恢复成 K/V。所以 MLA 的写入更省带宽（写得更少），读取更贵（要算投影）——这是第七篇”收益是缓存压到约三成、代价是读出来之后要做上投影”在通路上的具体体现。

第二，NVFP4 写入。 前面提过 quantizeAndWriteFP4KVCache 同时写主池与 scale 池。FlashInfer 侧的对应实现是 nvfp4_append_quantize_block：

1
__device__ __forceinline__ void nvfp4_append_quantize_block(
2
    const DType* __restrict__ input, const float global_scale, const size_t input_base,
3
    const uint32_t dim_base, uint8_t* __restrict__ packed_out, uint8_t* __restrict__ sf_out) {
4
  float values[16];
5
  float amax = 0.0f;
6
#pragma unroll
7
  for (uint32_t i = 0; i < 16; ++i) {
8
    const float value = nvfp4_append_to_float(input[input_base + dim_base + i]);
9
    values[i] = value;
10
    amax = fmaxf(amax, fabsf(value));
11
  }
12

13
  float sf_value = 0.0f;
14
  if (amax > 0.0f && global_scale > 0.0f) {
15
    sf_value = precise_rounding ? amax * __frcp_rn(6.0f * global_scale) : amax / (6.0f * global_scale);
16
  }
17
  __nv_fp8_e4m3 sf_fp8 = __nv_fp8_e4m3(sf_value);
18
  *sf_out = sf_fp8.__x;
19

20
  const float sf_rounded = static_cast<float>(sf_fp8);
21
  const float output_scale = (amax > 0.0f && sf_rounded > 0.0f && global_scale > 0.0f)
22
                                 ? (precise_rounding ? __frcp_rn(sf_rounded * global_scale)
23
                                                     : (1.0f / (sf_rounded * global_scale)))
24
                                 : 0.0f;
25

26
#pragma unroll
27
  for (uint32_t i = 0; i < 8; ++i) {
28
    const uint8_t lo = nvfp4_append_quantize_e2m1(values[i * 2] * output_scale);
29
    const uint8_t hi = nvfp4_append_quantize_e2m1(values[i * 2 + 1] * output_scale);
30
    packed_out[i] = lo | (hi << 4);
31
  }
32
}

逐段读：

1. 16 个一组：NVFP4 的块内 scale 覆盖 16 个元素（第七篇的 head_size / N，N = 16）。所以一次处理 16 个元素、输出 1 个 scale 字节。
2. amax 规约：块内最大绝对值。这就是第七篇说的”写入路径上的一次规约”——逐 token（更准确说是逐 16 元素块）缩放，而不是整池共用 scale。
3. sf_value = amax / (6 * global_scale)：这里有个必须解释的常数 6。E2M1 的可表示范围是 $[-6, 6]$，所以”块内最大值的映射目标”就是 6；除以 global_scale 是因为第二级 scale 是全局的——第一级 scale 存的应该是在”全局 scale 已生效”的坐标系下的相对缩放因子。
4. sf_fp8 = __nv_fp8_e4m3(sf_value)，然后再读回来 sf_rounded：这一步”先量化再反量化”是故意的。第一级 scale 本身以 FP8 存储，它也会丢精度；如果反量化时用的 output_scale 是按未量化的 sf_value 算的，那么 16 个元素的反量化结果，与读取侧拿 FP8 版本的 scale 去反量化的结果就不一致了。用 sf_rounded（FP8 往返后的值）算 output_scale，保证写入的码位在读取侧能精确还原。
5. packed_out[i] = lo | (hi << 4)：两个 4 位码位挤进一个字节容器。这正是第七篇 getNumEltsPerContainer() 返回 2 的东西——在写入侧就必须按容器打包，因为读取侧只会按容器粒度取数。

第 4 点是整个函数里最容易出错的地方，也是最值得记住的工程教训：当一个缩放因子本身也要被量化存储时，所有依赖它的计算都必须用”量化后的值”，否则编码与解码两侧的假设不一致。这类不一致不会导致崩溃，只会让数值悄悄变差。

索引器 K 缓存的写入：完全按 stride 走的 scatter

第七篇提过 DSA（DeepSeek 稀疏注意力）的索引器有自己独立的 K 缓存池（containsIndexerKCache）。它的写入路径与主池不同，走的是 invokeIndexerKCacheScatter：

1
__device__ __forceinline__ int64_t flatIndexToMemoryOffset(
2
    int64_t flat_idx, int32_t d0, int32_t d1, int32_t d2, int32_t d3, int64_t s0, int64_t s1, int64_t s2, int64_t s3)
3
{
4
    int32_t i3 = flat_idx % d3; flat_idx /= d3;
5
    int32_t i2 = flat_idx % d2; flat_idx /= d2;
6
    int32_t i1 = flat_idx % d1; flat_idx /= d1;
7
    int32_t i0 = flat_idx;
8
    return i0 * s0 + i1 * s1 + i2 * s2 + i3 * s3;
9
}

这个函数的名字说明了一切：它把”扁平元素下标”翻译成”考虑了 stride 的内存偏移”。注意注释里的一句话：

1
// Convert flat index to memory offset using strides (k cache pool from cpp kv cache manager is non-contiguous)

non-contiguous 是关键——索引器 K 缓存的池可能是不连续的（跨多个分配段），所以不能像主池那样用一条统一公式算地址，而是必须老老实实按四维 shape + stride 逐维展开。四维的语义是 [num_blocks, block_size, 1, per_token_size]：块维、块内 token 维、一个恒为 1 的占位维、token 内字节维。

写入内核本身非常轻：

1
__global__ void indexerKCacheScatterUnifiedKernel(uint8_t const* __restrict__ k_fp8_bytes,
2
    uint8_t const* __restrict__ k_scale_bytes, uint8_t* __restrict__ k_cache,
3
    int64_t const* __restrict__ slot_mapping_fp8, int64_t const* __restrict__ slot_mapping_scale, ...)
4
{
5
    constexpr int VEC_SIZE = 4;
6
    int32_t token_idx = blockIdx.x;      // 一个 block 处理一个 token
7
    if (token_idx >= num_tokens) return;
8

9
    int64_t flat_idx_fp8_base = slot_mapping_fp8[token_idx];     // 该 token 在池里的扁平起点
10
    int64_t flat_idx_scale_base = slot_mapping_scale[token_idx];
11
    if (flat_idx_fp8_base < 0 || flat_idx_scale_base < 0) return; // 负值 = 该 token 不写（如 padding）
12

13
    int32_t head_dim_idx = threadIdx.x * VEC_SIZE;
14
    int64_t flat_idx = flat_idx_fp8_base + head_dim_idx;
15
    int64_t dst_offset = flatIndexToMemoryOffset(flat_idx, cache_dim_0, cache_dim_1, cache_dim_2, cache_dim_3,
16
        cache_stride_0, cache_stride_1, cache_stride_2, cache_stride_3);
17
    int64_t src_offset = static_cast<int64_t>(token_idx) * head_dim + head_dim_idx;
18

19
    *reinterpret_cast<uint32_t*>(&k_cache[dst_offset]) = *reinterpret_cast<uint32_t const*>(&k_fp8_bytes[src_offset]);
20

21
    if (threadIdx.x == 0) {   // 只有 thread 0 写那 4 字节 scale
22
        int64_t dst_offset_scale = flatIndexToMemoryOffset(flat_idx_scale_base, ...);
23
        int64_t src_offset_scale = static_cast<int64_t>(token_idx) * scale_size;
24
        *reinterpret_cast<uint32_t*>(&k_cache[dst_offset_scale]) = *reinterpret_cast<uint32_t const*>(&k_scale_bytes[src_offset_scale]);
25
    }
26
}

三个设计点：

• 索引由 slot_mapping 给出，而不是内核自己算。 注释说得很直白：slot_mapping_fp8 是”Flat element index for FP8 data start position”，即该 token 的数据在池里的扁平元素起点。这是谁算的？是管理侧（页表 + 块分配）算好后传进来的。这让内核完全不需要知道页结构、页大小、块编号这些概念——它只做”按 stride 展开地址 + 拷贝”。
• scale 与数据走两条独立的 slot mapping。 因为 scale 池与主池是两块不同的张量，各自有自己的 stride。这也解释了为什么函数签名里有 slot_mapping_fp8 和 slot_mapping_scale 两个数组。
• -1 作为”不写”的哨兵值。 负的 slot 直接在开头 return，这让 padding token 或不该写的位置可以统一用 -1 表示，不需要额外的 mask 分支。

把索引器 K 缓存的写入和主池的写入对比看，是很清楚的两种设计哲学：

两条路线的取舍很清楚：主池的结构是框架内部固定契约，可以让内核吃掉复杂度；索引器缓存的结构更灵活（可跨分配段），于是复杂度的成本被前移到主机侧的 slot 计算。这也解释了为什么索引器缓存的写入内核代码这么短——它的短是”复杂度搬家”的结果，不是问题更简单。

三条读取路线的取舍

TensorRT-LLM 里读取侧实际有三条并存的路线，各自服务不同的场景：

路线一（通用分页） 的形态最接近大家的直觉：每个请求对应页表里的一段，内核在序列维上按页循环，每页内按 token 连续读。它的代价是对小批量解码不够高效——页循环带来的地址计算、页号查表、边界判断在小批量下摊不开。

路线二（MMHA） 的组织方式完全不同：一个 block 处理一条请求，把整条请求的 KV 序列切成”块对齐的循环”。它的访存形式是先取块指针数组，再在块内连续读。这条路线在第一批 vLLM/TRT-LLM 的对比里是 TRT-LLM 解码延迟优势的主要来源之一，但它的适用面窄——需要 head_dim、head 数等都在编译期确定，因此仓库里靠 decoderMaskedMultiheadAttentionLaunch.h + 大量实例化文件来覆盖组合。

路线三（XQA） 是新一代做法：把解码注意力交给 Tensor Core，用 JIT 编译生成针对具体 head_dim、page_size、精度组合的 kernel，并把 KV 池的地址描述成 TensorMap（TMA 描述符）。TMA 让内核可以用一条指令搬一整块数据，而不用手写地址算术——页表在这里的角色从”内核自己算地址”变成”为 TMA 准备描述符”。

三条路线的共同点是：页表都是被内联的，没有任何一条路线会先把 KV 收集成连续张量。差异只在”谁来承担地址算术”：内核自己算（路线一）、主机侧预先摊平成块指针数组（路线二）、硬件描述符承担（路线三）。

一条更具体的读取形态：get_length 与页表三件套的配合

无论走哪条路线，读取侧起步时都要先回答一个问题：这条请求要读多少 token、从哪些页读。这两个问题的答案就是 FlashInfer 那三件套（第七篇、第六篇都出现过）：

1
# from flashinfer docs/tutorials/kv_layout.rst
2
# For each request, we keep records of its page_indices and last_page_len ...
3
# The KV sequence length of request i is
4
#     page_size * (len(page_indices[i]) - 1) + last_page_len[i]

这个式子在写入侧有严格的对偶（前面 get_length 那段）：

其中 $L_i$ 是请求 $i$ 当前的 KV 长度，$n_i$ 是它占用的页数，$\text{last\_page\_len}_i$ 是最后一页里有效 token 数。式子里那个 $-1$ 是关键：前面的 $n_i - 1$ 页是满页，最后一页按有效长度计。这三件套（kv_indices / kv_indptr / kv_last_page_len）就是第六篇讲的 AttentionMetadata.prepare() 要构造出来的东西，而它们的语义正是由这条式子定义的。

读取侧拿到这三件套之后的循环形态（伪代码）：

1
for each request r in batch:
2
    seq_len = page_size * (indptr[r+1] - indptr[r] - 1) + last_page_len[r]
3
    for page_slot in range(indptr[r], indptr[r+1]):
4
        phys_page = page_table[page_slot]          # 逻辑页 → 物理页
5
        valid_tokens = page_size if not last page else last_page_len[r]
6
        for token in range(valid_tokens):
7
            load K[phys_page][token][:] and V[phys_page][token][:]
8
            accumulate into attention output

这段伪代码里有两层循环嵌套在同一个”序列维”上（页循环 + 页内 token 循环）。这正是分页注意力内核在序列维上的标准组织方式，也是它相对连续 KV 的访存代价来源：每跨一个页边界就一次间接跳转。

n_i 越大，跳转次数越多。页大小 32 时，128K 上下文有 4096 个页，即 4096 次页号查表；页大小 128 时只有 1024 次。这是”页大小怎么选”的一个直接影响因素。

上图是 FlashInfer 官方文档给出的分页布局示意：整池被切成固定大小的页，每条请求持有自己那一串页号。注意图里”页表”的视角——它把请求到页的映射画成一张表，而读写两侧都是拿这张表在做”逻辑槽位 → 物理页”的翻译。图片来源：FlashInfer 官方文档 KV-Cache Layout 教程。

变量注意力窗口与读取侧的长度裁剪

第七篇讲过 KVBlockArray 里那组与窗口相关的字段。它们在读取侧的作用，可以从官方文档的一句话理解：

Blocks that leave the attention window are freed and placed on the radix search tree so they can be reused.

这句话描述的是有限注意力窗口（VSWA，variable sliding window attention）下的行为：块滑出窗口后就被释放、挂回前缀复用树，等待其他请求复用。而 KVBlockArray::getKVTokenIdx 正是这个机制在内核里的对应物：

1
__host__ __device__ [[nodiscard]] inline int32_t getKVTokenIdx(int32_t tokenIdx) const
2
{
3
    if (!isSinkToken(tokenIdx))
4
    {
5
        // Apply cyclic kv cache if tokenIdx >= max_attention_window_size.
6
        return mSinkTokens + mBubbleLen + (tokenIdx - mSinkTokens) % mCyclicCacheLen;
7
    }
8
    return tokenIdx;
9
}

逐段拆：

• isSinkToken(tokenIdx)：前 mSinkTokens 个 token 是 attention sink，永远保留，位置不走环形映射。这就是 StreamingLLM 那条”若干 sink + 滑窗”的路线在缓存层的落地。（本站在 JSSA 那篇里对比过”不保 sink”的另一条哲学，两条路线的分歧点就在这里。）
• mBubbleLen：气泡长度。它是为了让 sink 之后的有效区域在环形缓冲区里对齐而留的余量，构造式是：

1
auto sinkTokensInLastBlock = mSinkTokens % mTokensPerBlock;
2
mBubbleLen = sinkTokensInLastBlock == 0 ? 0 : mTokensPerBlock - sinkTokensInLastBlock;

这个”气泡”的存在理由是块对齐：环形缓冲区的起点如果不是块的起始位置，那么”滑窗的边界”就会落在块的中间，导致每个块都要额外处理”部分有效”的情况。补齐到块边界之后，窗口边界与块边界一致——读取侧就可以整块加载、整块丢弃，不必逐 token 判断有效性。这是一个纯粹为”读写粒度对齐”而存在的字段，它不产生任何数值语义。

• (tokenIdx - mSinkTokens) % mCyclicCacheLen：环形映射。超出窗口的历史 token 会覆盖掉最老的槽位，而不是新分配内存。mCyclicCacheLen 就是环的长度：

1
mCyclicCacheLen = (mEnableOneMoreBlock) ? mMaxAttentionWindow + mTokensPerBlock - mSinkTokens
2
                                        : mMaxAttentionWindow - mSinkTokens;

注意 mEnableOneMoreBlock 决定环长是”窗口 + 一整块”还是”窗口”。多留一整块的意义在于：当窗口恰好跨块边界时，不会因为环形覆盖而把正在读取的那一块中的有效部分覆盖掉。这也解释了为什么这个开关要由主机侧算好并传进来（canUseOneMoreBlock）——它取决于池里到底有几个块可用，是个容量约束，不是数学约束。

环形映射在写入与读取两侧必须完全一致。写入侧走 getKVTokenIdx 算出实际槽位再写，读取侧走同一个函数算出实际槽位再读。两个方向用同一个函数（都在 kvCacheUtils.h 里，__host__ __device__ 双端可用）是刻意的设计——如果两侧各写一份取模逻辑，任何一处偏差都会让读写指向不同槽位。

上图是 TensorRT-LLM 官方文档的滑动窗口注意力示意。窗口滑走后，早期块被释放并挂回复用树；写入侧靠 getKVTokenIdx 的环形映射把新 token 落进被腾出的槽位，读取侧靠同一映射找到窗口内的全部槽位。两侧共用同一个位置映射函数，是这套机制唯一的安全保证。图片来源：TensorRT-LLM 官方文档 media/feat_long_seq_sliding_win_attn.png。

TrtllmAttention：写池藏在注意力算子内部

tensorrt_llm/_torch/attention/backends/trtllm.py 里出现的是一个开关而不是一次调用：

1
forward_args.update_kv_cache = True            # 默认：这次前向要更新池
2
...
3
forward_args.update_kv_cache = not has_q_only  # Q-only 前向不更新
4
forward_args.update_kv_cache = not metadata.is_cross or k is not None

对应的 C++ 契约在预处理参数里有明确注释：

1
// The input contains Q only and the existing KV cache must not be updated.

也就是说：TrtllmAttention 后端把写池合进了它自己的融合算子。update_kv_cache 这个布尔值一路传到设备侧，决定融合 kernel 内部是否执行写分支。三种情况必须关掉写入：

• Q-only 前向：只算 Q 不产生新 KV（例如某些需要重算注意力的场景）。
• 交叉注意力：cross attention 的 K/V 来自编码器，不该进解码器的自回归池。
• 投机解码的验证步：草稿 token 的 KV 需要特殊处理（可能只写被接受的部分），不能走常规路径。

这个设计的收益很明显：一次 kernel launch 完成”预处理 + 写池 + 注意力计算”，中间不需要任何跨 kernel 的中间张量。代价是算子接口变复杂（要传池指针、页表、窗口参数），而且写入逻辑与注意力逻辑耦合在一起，无法单独测试或替换。

FlashInferAttention：显式的 append_paged_kv_cache

FlashInfer 后端走的是另一条路线：写池是一次独立的、显式的调用（前面的 _append_paged_kv_cache）。原因是 FlashInfer 的注意力内核（BatchDecodeWithPagedKVCache / BatchPrefillWithPagedKVCache）只负责读，不负责写——它需要一个已经填好的池。

于是流程变成：先 append_paged_kv_cache 把新 token 写进池，再启动注意力内核从池里读。这多了一次 kernel launch 和一次中间张量的读写（写入侧要先把 K/V 从 qkv_input 布局转成 append 需要的布局），换来的是职责清晰：写池与注意力各自独立、各自可替换、各自可单独测试。

VanillaAttention：教学式的 PyTorch 参考实现

vanilla.py 的 _single_request_update_kv_cache 是三个实现里唯一用 PyTorch 算子写清楚每个步骤的，它的价值在于把索引算术暴露成可读的 Python：

1
seq_len = past_seen_token + kv_len
2
tokens_per_block = kv_cache_tensor.shape[2]
3

4
if k is not None and v is not None:
5
    written = 0
6
    while written < kv_len:
7
        pos = past_seen_token + written
8
        blk = block_ids[pos // tokens_per_block]     # 块号：pos 除以页大小
9
        off = pos % tokens_per_block                 # 页内偏移：pos 取模页大小
10
        n = min(tokens_per_block - off, kv_len - written)   # 本块还能装几个
11
        assert blk != BAD_PAGE_INDEX, (
12
            f"Writing new KV into an evicted/invalid page (pos {pos}); "
13
            "block_ids/metadata are inconsistent.")
14
        dst = torch.arange(off, off + n, device=kv_cache_tensor.device)
15
        kv_cache_tensor[blk, 0].view(dtype=access_type).index_copy_(
16
            0, dst, k_selected[0, written:written + n].view(dtype=access_type))
17
        kv_cache_tensor[blk, 1].view(dtype=access_type).index_copy_(
18
            0, dst, v_selected[0, written:written + n].view(dtype=access_type))
19
        written += n

这段代码就是写入侧的数学定义，逐行对照着 CUDA 版本读会非常清楚：

• pos // tokens_per_block 与 pos % tokens_per_block：就是 CUDA 里那个 divmod，这里用 Python 的整除/取模直接写出来。注意这里直接用了 blk = block_ids[...] 一步拿到物理块号——因为 block_ids 是主机侧从 KVCacheManager 拿到的块号列表，逻辑页到物理页的翻译在外面做完了。CUDA 侧那两步（indptr[b]*page_size + pos 再 divmod）是为了在没有 block_ids 数组的场合也能算出来，两者的目标一致。
• n = min(tokens_per_block - off, kv_len - written)：这段 while 循环就是”跨页”的全部处理。它把一次写入拆成若干”页内连续段”，每段落在单个页内——前面说”跨页的是序列而不是 token”，这里就是它的实现形态。
• assert blk != BAD_PAGE_INDEX：注释写得很直白——“New tokens must land in a live page; fail loudly rather than writing into the last page via negative indexing.”。这是一个明确的工程决策：宁可让断言炸掉，也不要让负索引悄悄写到错误的位置。第八篇这个系列反复出现的主题——“不报错的错误最难查”——在这里被主动防御了。
• view(dtype=access_type)：按字节容器重解释。access_type 由 k_selected.dtype.itemsize 决定（FP8 时是 uint8、FP4 打包时是 uint32）。这是第七篇讲的”元素数 vs 容器数”在 PyTorch 侧的写法：用一个更宽的 dtype 视图去一次搬运多个窄元素。

读取侧的同款参考实现

vanilla.py 里还有 _gather_paged_kv，它是”读取侧要不要物化连续张量”这个问题的明确答案：

1
def _gather_paged_kv(kv_cache_tensor, block_ids, kv_idx, num_tokens, tokens_per_block):
2
    """Materialize the first ``num_tokens`` logical K (``kv_idx=0``) or V
3
    (``kv_idx=1``) tokens from the paged pool as a contiguous
4
    ``[num_tokens, num_kv_heads, head_dim]`` tensor (NHD view).
5

6
    A single block returns the contiguous view directly (no copy), which
7
    keeps a within-one-page request allocation-free.
8
    """
9
    if num_tokens <= 0:
10
        return kv_cache_tensor[0, kv_idx, :0]
11
    chunks = []
12
    read = 0
13
    while read < num_tokens:
14
        blk = block_ids[read // tokens_per_block]
15
        off = read % tokens_per_block
16
        n = min(tokens_per_block - off, num_tokens - read)
17
        if blk == BAD_PAGE_INDEX:
18
            chunks.append(kv_cache_tensor.new_zeros((n, *kv_cache_tensor.shape[3:])))
19
        else:
20
            chunks.append(kv_cache_tensor[blk, kv_idx, off:off + n])
21
        read += n
22
    return chunks[0] if len(chunks) == 1 else torch.cat(chunks, dim=0)

三个细节值得注意：

1. “A single block returns the contiguous view directly (no copy)“：如果一条请求的全部 token 都在一个页内，那就直接返回池里那一段的视图，零拷贝。这个优化针对的正是解码阶段的短请求——第 100 步解码时 num_tokens = 101，页大小 32 的话其实横跨 4 个页，但要到请求很长时页数才多。“能不拷贝就不拷贝”这个判断放在最前面，说明作者知道这条路径命中率高。
2. BAD_PAGE_INDEX 用零填充而不是报错（读取侧与会写入侧的断言形成对比）：写入不能容忍错误（会把数据写到别人页上），读取则可以退化成”读出一段 0”——因为注意力权重会在别处把这段 0 的贡献压掉（或者这段本来就不该被读到）。这是读写两侧错误处理策略不对称的又一个例子。
3. 返回的是”前缀 num_tokens 个 token”而不是整段：调用方用 past_seen_token 与 seq_len 两个值区分”读历史”还是”读全量”。注意 _single_request_update_kv_cache 里 _gather_paged_kv 被调用了两次，一次读 past_seen_token（历史）、一次读 seq_len（含新 token）——这正是滑动窗口与稀疏注意力需要的”过去 vs 现在”两种视图。

vanilla.py 的这套实现慢（每个请求一次 Python 循环、每个页一次 index_copy_），但它把接口语义钉死了：block_ids 是逻辑槽位到物理块的映射，BAD_PAGE_INDEX 是无效块哨兵，页内偏移按 pos % tokens_per_block 算。任何新的后端实现，只要满足这套语义，就能和调度、KV 管理、前缀复用正确对接。这是”参考实现”存在的意义：它不是性能路径，它是契约的文档。

物理页从哪来

写入内核拿到的从来不是”这块数据该写哪个物理页”的直接答案，而是一张间接表（第六篇的 kv_block_ids_per_seq、第七篇的 block_ids、FlashInfer 的 kv_indices）。内核顺着这张表跳到物理页。谁维护这张表？调度器与 KVCacheManager（第三、四篇）——它们决定了每条请求当前持有哪些块，并在请求被抢占、块被驱逐时同步更新表。

于是有一条重要的时序约束：写入必须发生在”页表已经包含本步要用的块”之后。顺序是：

1
调度器分配/复用块（更新账本与页表）
2
        │
3
        ▼
4
prepare() 构造索引数组（第六篇：块句柄 → 块表/页表/偏移）
5
        │
6
        ▼
7
写入内核执行（本步新 token 的 KV 落进池）
8
        │
9
        ▼
10
注意力内核执行（读取含新 token 在内的整段 KV）

这四步的顺序不能乱。特别是最后两步：如果写入与注意力并行执行，注意力就可能读不到刚写的 token（读到的是上一步的内容），或者更糟——读到部分写入的中间状态。TensorRT-LLM 的做法是把”写入 + 注意力”做成同一算子内部的两段（TrtllmAttention 路线），从而在算子内部保证顺序；FlashInfer 路线则靠同一个 CUDA stream 上的顺序执行来保证。

分块预填充的写入节奏

分块预填充（chunked prefill）把一条长提示切成多段依次执行，这让写入侧的节奏变得不规则：

上图是 TensorRT-LLM 官方文档的分块上下文注意力示意：长提示被切成多段，每段依次执行、产生的 KV 依次写进池。写入侧的关键在于 past_seq_len 的逐段递增（前面 cache_seq_len - actual_seq_len 那段）：第 $k$ 段的 token 必须落在池里的第 $k \cdot C$ 个位置往后（$C$ 是分块大小），而不是从头写。这也是为什么”分块预填充 + 前缀复用”的组合特别容易出错——池里的位置既要与当前请求的实际 token 序号一致，又要与前缀复用命中的那些块的物理内容对得上。图片来源：TensorRT-LLM 官方文档 media/feat_long_seq_chunked_attention.png。

在本站 SGLang 源码完全拆解（四）：PrefillAdder 如何预算 KV、切分长提示并组装 Prefill Batch 那篇里，分块预填充的接纳算法是从调度侧讲的；这一节讲的是它的写入侧后果：同一段提示被切成 $k$ 段之后，只有”这段写在池的哪个位置”这件事必须由写入路径自己算对。