DeepInfer 源码拆解（二十二）：InputMetadata 的构造——调度元数据如何变成 GPU 输入张量

positions：绝对位置

1
for (int32_t j = p_seq->seqlen_k - p_seq->seqlen_q; j < p_seq->seqlen_k; ++j, ++i_q)
2
    p_positions[i_q] = j;

对第 $i$ 个 query token，它的绝对位置是

其中 $j$ 是该 token 在本次 query 段内的局部下标。prompt 的 chunked prefill 下，seqlen_k - seqlen_q 就是这段 chunk 的起点；decode 下它等于上一步结束时的长度。RoPE、SWA ring 槽位、压缩 token 的边界判断全都以这个绝对位置为准，所以它是所有派生字段的锚点。

input_ids：prompt 取整段，decode 取一个加草稿

prompt 的 query 段就是原始 token 序列的尾部：

1
memcpy(p_input_ids + i_q, token_ids.data() + p_seq->seqlen_k - p_seq->seqlen_q,
2
       seqlen_q_golden * sizeof(int32_t));

decode 只取最后一个已确认 token，后面接投机草稿：

1
p_input_ids[i_q] = token_ids[seqlen_k - seqlen_q];
2
for (int j = 0; j < seqlen_q_draft; ++j)
3
    p_input_ids[i_q + 1 + j] = p_seq->mtp.draft_tokens[j];

这里有一个容易读错的地方：decode 的 seqlen_q 是 1 + 草稿数，而 seqlen_k 是 token_ids.size() - 1 + seqlen_q。也就是说，草稿 token 还没有进入 token_ids，但已经计入本步的 key 长度。验证分支会用 seqlen_q_golden = 1 单独检查主 token，再检查草稿数量。

gate_force_random：负载均衡的旁路开关

1
p_gate_force_random[i] = args.force_load_balance;

这是一个布尔张量，整批同值。它为真时，MoE gate 会用随机 logits 替换真实 logits，把路由强行摊平。它的用途不是推理正确性，而是压测与 EPLB 统计：当需要在不改变输入分布的前提下观察最坏情况下的专家负载时，打开它可以让每个专家都收到接近均等的 token。它逐 token 复制成 [total_q]，是因为下游 kernel 按 token 读取，不想在 kernel 里再传一个标量分支。

ngram_token_ids：Engram 的回看窗口

开启 Engram（max_ngram_size > 0）时，每个 query token 需要 max_ngram_size 个压缩后的历史 token。fill_engram_row 逐行填充：

1
for (int r = 0; r < max_ngram_size; ++r) {
2
    int s = pos - r;                       // 往回看 r 步
3
    if (s < 0) v = pad_token;              // 序列头，补 pad
4
    else if (s < L) v = compressed[s];     // 已压缩部分
5
    else v = lookup_table[last_next_n_token[s - L]]; // 本步新 token，现场压缩
6
}

$L$ 是 compressed_seq_data 的长度。回看窗口跨越“已压缩历史”和“本步/上步未压缩 token”两个区段，后者要通过压缩词表的 lookup_table 现场映射。遇到图像哨兵 token 时，从该位置起全部填 pad，避免 N-gram 跨模态边界。这个字段的形状是 [total_q, max_ngram_size]，是公共底座里唯一一个二维字段；其余都是一维。

valid_token_mask 与 query_indices：占位序列的两种标记

decode base 填 valid_token_mask：

1
bool is_valid = !p_seq->is_placeholder;
2
for (int j = 0; j < p_seq->seqlen_q; ++j) p_valid_token_mask[i_q++] = is_valid;

query_indices 则给每个 token 标记它属于哪条序列：

1
int indices = !p_seq->is_placeholder ? i_seq : 998244353;

正常序列用批内下标 i_seq；占位序列用一个远大于任何真实批大小的哨兵。为什么不用 -1？因为下游的 index-attention 与 core-attention 会拿它做 block_table[query_indices[i]] 一类的间接寻址，-1 在 C++ 里是合法下标，会读到错误位置；998244353 在 int32 范围内足够大，任何数组访问都会越界并被 CUDA 的非法地址检查捕获，属于“快速失败”的设计。

prompt 路径的 valid_token_mask 不是在这里填的，而是在 CP mirror 之后由 mirror_cp_index != -1 生成，原因在 CP 一节说明。

prompt：一段 query 产生多少个压缩 token

1
n_cq += std::max(1, seqlen_k / n - (seqlen_k - seqlen_q) / n);

设 query 段为 $[\text{seqlen}_k-\text{seqlen}_q,\ \text{seqlen}_k)$，压缩粒度为 $n$。这段区间覆盖的完整压缩块数是

也就是“结尾处的压缩块数减去起点处的压缩块数”。但 code 外面套了 max(1, ...)：如果这段 query 太短、没有跨过任何压缩边界，仍然要分配一行。这一行不是有效压缩 token，而是占位行：positions_c 写 0，slot_mapping_store_ca 写 -1。下游的压缩器按行数启动，零行会导致 grid 为 0，某些 kernel 与 CUDA Graph 都不喜欢零长度；补一行并把 slot 标成 -1，既保持形状稳定，又让写入操作被跳过。

对每个有效压缩 token，位置与槽位分别是

其中 $j$ 是压缩 token 的全局压缩下标，$b_c = \text{block\_size}/n$ 是每个物理块能放多少个压缩 token。第一项把压缩下标还原成原序列的绝对位置（压缩 token 代表从 $j n$ 开始的 $n$ 个原始 token），第二项把压缩下标映射进物理块。这里用的是按最小压缩倍率计长的块表：fill_decode_block_tables 用 min_compress_ratio 数每条序列占了多少块；块本身的大小 block_size 对所有倍率相同，倍率更大的层只是每块能容纳的压缩 token 更少。

decode：只在跨过压缩边界的那一步写

decode 的 seqlen_q 很小，绝大多数 query token 不会产生新的压缩 token。判断条件是：

1
if ((orig_pos + 1) % n == 0) {
2
    p_pos[i_q] = (orig_pos / n) * n;
3
    p_slot[i_q] = block_table[compressed_idx / b_c] * b_c + compressed_idx % b_c;
4
} else {
5
    p_pos[i_q] = 0;
6
    p_slot[i_q] = -1;
7
}

$\text{orig\_pos}=q$ 时，$(q+1) \bmod n = 0$ 说明第 $q$ 个 token 正好是第 $q/n$ 个压缩块的最后一个原始 token；此时压缩下标是 $q/n$，它代表的压缩 token 位置是 $(q/n)\cdot n$。其余行写 slot=-1，表示“本步不写压缩缓存”。这里 positions_c 写 0 而不是 -1，是因为位置字段只被写路径读取，而写路径先看 slot 是否非负；写 0 可以让未初始化内存不进入 kernel。

一个共享的坑：block table 的右侧未初始化

fill_decode_block_tables 把每条序列的块表复制到 [total_q, max_num_blocks] 的每一行：

1
int32_t tokens_per_block = block_size / compress_ratio;
2
num_blocks = (seqlen_k / compress_ratio + tokens_per_block - 1) / tokens_per_block;
3
std::copy(block_table.begin(), block_table.begin() + num_blocks,
4
          p_block_tables + i_q * max_num_blocks);

右侧 max_num_blocks - num_blocks 列不初始化。这不是疏忽：下游 kernel 的有效长度来自 klen / offset_len，不会读越界列；而把整行清零反而要多一次写。若调试时看到块表右侧是垃圾值，先确认读取方是否用了 num_blocks 而不是 max_num_blocks 做循环上界。

写：只写窗口内的最后一段

1
if (p_seq->seqlen_q - j <= swa_block_size)
2
    p_slot_mapping_store_swa[i_q] = p_seq->state_id * swa_block_size
3
        + (p_seq->seqlen_k - p_seq->seqlen_q + j) % swa_block_size;
4
else
5
    p_slot_mapping_store_swa[i_q] = -1;

条件 seqlen_q - j <= swa_block_size 表示这个 query token 位于本段 query 的最后 swa_block_size 个之内。更早的 token 不会留在 ring 里（会被后续 token 覆盖），所以不写。槽位公式是

前一项定位到序列自己的 state 槽，后一项是 ring 内的偏移。state_id 在 prompt 阶段是每条序列一个（state_block_idx 形状 [n_seq]），在 decode 阶段展开成每个 query token 一个（形状 [total_q]），因为 decode kernel 按 token 取行。

读：把 ring 区间翻译成块粒度描述符

make_swa_block_load_plan 不生成逐 token 的索引，而是生成“装载段”：

1
history_end   = seqlen_k - seqlen_q;
2
history_begin = history_end - window_size;
3
padding_len   = clamp(swa_valid_start - history_begin, 0, window_size);
4
valid_begin   = history_begin + padding_len;
5
builder.append_zero_fill(padding_len);
6
builder.append_ring_range(state_id, valid_begin, history_end);

输出是 [num_segments, 4] 的 int32 张量，每行 [src_block_id, src_begin, src_end, dst_offset]，其中 src_block_id=-1 表示零填充。为什么要块粒度？因为 SWA kernel 的输入布局是 [n_seq * window_size, ...] 的连续缓冲：每条序列固定占 window_size 行，历史不足 window_size 的部分要补零。装载计划把“从哪个 state 槽的哪段 ring 搬到目标缓冲的哪个偏移”描述成若干段，一次 kernel 就能完成全部序列的搬运，而不需要 CPU 逐 token 计算。

total_k_swa 恒等于 n_seq * window_size，即使真实历史更短。补零段就是为此存在。swa_valid_start 则处理 zero-state 重放：当 SWA 状态被清零后，[0, valid_start) 的槽不可信，装载计划把这段也补零，attention 的可见下界在后续 pivots 里同步 clamp 到 valid_start。两处必须一致，否则 kernel 会读到清零后的垃圾值。

混合层型意味着两套池子

V4.1-Flash 的层不是同一种注意力：一部分是全量压缩 KV，一部分是 SWA，还有 YoCo 的 consumer 层。不同层型的 KV 生命周期不同——全量 KV 要长期保留、可落盘复用；SWA ring 只保留一个窗口，且一部分层在重放时要重算。元数据因此天然分成两套池子的描述符：

• 分页池：block_table、slot_mapping_store_ca、klen、kv_ca_index_offset_len；
• 状态池：state_block_idx、slot_mapping_store_swa、kv_swa_index_offset_len、load_segments_swa。

两套池子的“块”含义也不同：分页池的块由调度器统一分配，可以跨请求复用；状态池的槽按序列独占，生命周期跟随请求。把它们混在一张块表里，会让前缀复用与 SWA 重放都变复杂。vLLM 在混合注意力模型上采用“分组”的方式处理同类问题：不同层型各自成组，组内共享块池，组间独立核算。

图片来源：vLLM 官方文档 docs/assets/design/hybrid_kv_cache_manager/overview.png。DeepInfer 没有直接照搬这套分组，但“全量层与 SWA 层分开记账”的原则是一致的；差别在于 DeepInfer 的 SWA 历史放在 StateBlock 里，而不是分页池的一个组。

压缩 KV 的装载计划

prompt 阶段还要把压缩 KV 历史从分页池搬进连续缓冲，make_paged_block_load_plan 复用同一个 BlockLoadPlanBuilder：

1
for (auto p_seq : seq_metadata_list) {
2
    num_tokens = p_seq->seqlen_k / compress_ratio;
3
    if (num_tokens == 0) builder.append_zero_fill(1);
4
    else builder.append_paged_range(p_seq->block_table, 0, num_tokens);
5
}

每段按物理块切分，total_k_compressed 记录搬完后的总长度。append_zero_fill(1) 同样是为了避免零长度序列：一条空序列也要在缓冲里占一行，保持 cu_seqlens 单调。

图片来源：vLLM 官方文档 docs/assets/design/paged_attention/key.png。DeepInfer 的 block_table 与 slot_mapping 沿用同一套“逻辑位置 → 物理块 → 块内偏移”的三级映射，只是额外叠加了压缩倍率与 SWA ring。

zigzag 的 token 归属

一条长度为 $L$ 的序列被切成 $2p$ 块（$p$ 是 CP 大小），rank $r$ 拿第 $r$ 块和第 $2p-1-r$ 块：

每个 rank 的 token 数与 cp_mirror_len 由 C++ 侧直接算出：

1
cp_mirror_len += (seqlen_q + cp_size * 2 - 1) / (cp_size * 2);
2
cp_mirror_len *= 2;

(sq + 2p - 1) / (2p) 是每块的向上取整长度，乘以 2 是两块。注意这里的 sq 是 chunk 长度，不是整条序列：chunked prefill 下，镜像切分按本次 chunk进行，而不是按历史长度，因为本 rank 只处理本次 chunk 的一部分。

重排哪些字段、保留哪些字段

prepare_for_mirror_cp 的规则是“query-aligned 的字段重排，跨边界的字段保留”：

有一处细节很值得注意：代码先判断 hasattr(self, "positions_cp") 再赋值。

1
if hasattr(self, "positions_cp"):
2
    self.positions_cp = self.positions
3
self.positions = gather(self.positions, self.mirror_cp_index)

positions_cp 保存的是重排前的原值。为什么压缩 KV 需要原位置？在 CP 下，压缩是对整条序列做的：每个 rank 只算自己的 query，但压缩 token 的边界（哪个原始位置产生一个压缩 token）取决于全局位置，不能随 mirror 重排而变。所以压缩相关的 gather 用 positions_cp，而 attention 的 query 位置用重排后的 positions。同理，prepare_cp_mirror 直接吃原始的 cu_seqlens_q/cu_seqlens_k，输出每个 mirror 行的 ks/ke。

InputMetadataV4VL 覆写了一件事：图像掩码 images_seq_mask 在重排前另存为 images_seq_mask_full。YoCo compact 视图需要按全局下标 gather 图像掩码，而已经重排过的版本无法反推全局位置。这是“跨边界字段必须保留原件”原则的一个具体例子。

mirror 与 YoCo 的叠加顺序

当两种机制同时开启时，顺序不能反：

1. 先在 full 布局上算截断（fill_yoco_truncation_metadata）；
2. 再在 compact 布局上做一次 mirror（get_consumer_view 里的 prepare_cp_mirror）。

代码注释把这称为“先 compact 再 CP”。反过来的话，mirror 后的行已经打乱，无法按连续的保留区间计算 clip_prefix。Compact 之后重新 mirror 的长度由 fill_yoco_consumer_cp_scalars 预先算出：

1
cp_mirror_len_c += 2 * ((keep_len + 2 * cp - 1) / (2 * cp));

它的另一项输出 yoco_consumer_cp_contrib_pad_len 是“每个 rank 对 compact 结果的贡献最大值”：对 rank $r$，只数那些落在保留区间内的块：

1
lo = max(ps * b, clip);  hi = min(ps * (b + 1), q_len);
2
kept_per_rank[r] += max(hi - lo, 0);

取所有 rank 的最大值，作为 all-gather 后每个 rank 的连续缓冲长度。为什么取 max 而不是和？因为每个 rank 只把自己的贡献发给汇聚点，汇聚后的总长等于所有贡献之和，但单 rank 的发送缓冲只需要放下自己那一份；取最大值保证每个 rank 用同一块大小的缓冲，避免在 CUDA Graph 里出现按 rank 变化的形状。

ks/ke：一个 query 行给一个连续区间

对每个稀疏倍率 $r$，prompt 的 ks/ke 按序列顺序累加：

1
int curr_ks = 0;
2
for (int j = seqlen_k - seqlen_q; j < seqlen_k; ++j) {
3
    p_ks[i_q] = curr_ks;
4
    p_ke[i_q] = curr_ks + (j + 1) / r;
5
}
6
curr_ks += std::max(1, seqlen_k / r);

这里的下标不是全局压缩下标，而是搬进连续缓冲之后的下标：前面序列的压缩 key 先占位，本序列的 key 从 curr_ks 开始。第 $j$ 个位置的因果注意力只能看到位置 $0..j$ 的 key，其中压缩 key 的数量是

于是区间是 $[\text{curr\_ks},\ \text{curr\_ks} + k_{\text{len}})$。为什么 curr_ks 要加 max(1, seqlen_k / r) 而不是直接用 seqlen_k / r？因为装载计划里空序列会补一行，这里的步进必须与装载后的实际长度一致，否则下一序列的起点会错位。两处都出现了同一个 max(1, ...)，它们必须成对修改。

ks/ke 的形状是 [total_q]，每个 query 行两个 int32。稀疏索引层拿它做什么？top-k 选出来的压缩 key 只给出一组“候选位置”，真正的注意力计算仍然需要知道每个 query 行的候选在缓冲里的绝对区间；ks/ke 就是这个区间的左闭右开边界。与 index_attn.py 里的索引器实现相配合，这一节只讨论元数据侧的输入。更早讨论压缩 KV 与索引器跨层共享的文章见 DeepInfer 源码拆解（十一）：注意力族与压缩 KV。

pivots：一个 query 行给两段

SWA 的可见 key 来源是两段：ring 里的历史（section 1）和本步 chunk 里新算的 key（section 2）。它们不在同一块内存里，无法用一个连续区间表达。DeepInfer 的做法是让 CPU 只写四元组，让 kernel 展开：

其中 $X$ 是该 query 行最终的 kv_swa_full_index（长度 window_size），$W$ 是窗口大小，$(-1)$ 为 mask。四元组的产生：

1
abs_lo = max(pos - window_size + 1, valid_start);   // 可见下界（绝对位置）
2
if (abs_lo < 0) abs_lo = 0;
3

4
sec1_lo_abs = min(abs_lo, q_start);
5
start_slot  = sec1_lo_abs - (q_start - window_size);
6
P0 = sec1_start + start_slot;              // section1 起点（缓冲下标）
7
P1 = window_size - start_slot;             // section1 长度
8

9
k_lo = max(0, abs_lo - q_start);
10
P2 = sec2_start + k_lo;                    // section2 起点
11
P3 = j - k_lo + 1;                         // section2 长度

逐步解释：

1. 可见下界。因果注意力 + 滑窗：位置 $p$ 的 query 只能看 $[p-W+1, p]$。再与 valid_start 取 max，保证不读被清零的槽。
2. section 1 的偏移。sec1_start = i_seq * window_size 是这条序列在历史缓冲里的行起点。历史缓冲里每个序列占 window_size 行，对应绝对位置 $[q_{\text{start}}-W, q_{\text{start}})$。所以绝对位置 $a$ 在 section 1 内的偏移是 $a-(q_{\text{start}}-W)$，即 start_slot。若历史不足一整窗，start_slot 会接近 $W$，表示前部是补零或不存在。
3. section 2 的起点。section 2 存放本 chunk 新算的 KV，从 sec2_start 开始，长度等于已处理的 query 数。局部位置 $k$ 对应绝对位置 $q_{\text{start}}+k$，所以起点是 sec2_start + k_lo。
4. 长度。section 2 的长度是“当前局部下标 $j$ 减去起点 $k_{\text{lo}}$ 再加 1”，包含自身。

四元组的好处是 CPU 只写 4 个 int32，而展开成 $W$ 列的稠密索引交给 GPU 上的 expand_swa_index kernel。$W$ 通常是 128–1024 量级，若在 CPU 展开，每行要写 $W$ 个值，prompt 的 total_q 动辄数万，会直接拖慢每步准备时间。compute_attn_info 在 forward 开头调用 _expand_kv_swa_full_index 展开；展开结果走 _cp_gather，在 CP 开启时同时按 mirror 索引重排。

kv_ca_index_offset_len 用的是同样的“给区间、让 kernel 展开”思路，区别是展开结果指向分页池而不是连续缓冲。dense 层每行的公式在 InputMetadataV4.compute_attn_info 里写成注释：

1
kv_ca_index[i_q, :len] = offset[i_q] + arange(len)

decode 时的展开还要走 block table：

1
pos = offset[i_q] + arange(len[i_q])
2
kv_ca_index[i_q, :len] = block_tables[pos // tpb] * tpb + pos % tpb

tpb = block_size // n。这与压缩槽位公式是同一条映射的两种写法：一个在写侧算 slot，一个在读侧把 slot 展开成稠密索引。展开宽度必须是 128 的整数倍（(max_model_len // n + 127) // 128 * 128），因为稀疏注意力 kernel 的索引布局按 128 对齐。

为什么不在 CPU 上直接展开

可以把这看作一次“元数据的分层”：CPU 层只承担 $O(\text{total\_q})$ 的工作（每行常数个值），GPU 层承担 $O(\text{total\_q} \cdot W)$ 的展开。展开是纯数据并行，没有依赖，适合放在 kernel 里。更深一层的原因是 CPU 时间必须落在关键路径之外：CP 广播在 driver 上先发对象再发张量，跟随 rank 的空闲窗口由 driver 的 CPU 时间决定；每多一次逐 token 的 Python/C++ 循环，重叠窗口就少一点。第二篇讨论过的 ZOS 把 CPU 后处理移出关键路径，这里的设计与它是同一个思路：能在 GPU 上并行做的，就不要让 CPU 顺序做。

eager：prepare_decode_v4

在 prepare_decode_base 与 update_common_v4 之后，eager 路径补齐三组字段：

context_lens：每个 query token 的上下文长度（含自身）：

投机解码时 seqlen_q > 1，同一序列的草稿 token 共享前缀，但每个草稿位置看到的 key 数依次增加，所以是逐 token 的。

klen：稀疏层每个 query token 要看的压缩 key 数，$\lfloor (j+1)/r \rfloor$。它与 prompt 的 ke-ks 同构，只是不带偏移。

kv_swa_index_offset_len：decode 的 SWA 读窗口。与 prompt 不同，decode 的 query token 本身也已经写进了 ring（slot_mapping_store_swa 先写后读），所以窗口就是 ring 本身的一个环形区间：

1
real_window_size = min(current_offset + 1, window_size);
2
offset = state_id * swa_block_size
3
       + (current_offset - real_window_size + 1) % swa_block_size;

current_offset = seqlen_k - seqlen_q + j 是 query 的绝对位置。窗口起点 (current_offset - W + 1) % swa_block_size 落在 ring 里，长度不超过 $W$。这里出现一个 subtle 的点：decode 的 ring 大小 swa_block_size 可能略大于 window_size（留投机回退余量），但读窗口只取 $W$ 个，所以环上“多出来”的槽不会被读到。

ZOS：先讲清单，再让 kernel 算

ZOS 路径的入口在 Python 侧 prepare_decode_v4_zos。它分三个阶段：

1. prepare_decode_zos_inputs（C++）：把每条序列的 CPU 标量抽出来，装进一个 MetadataZOS 结构。抽出的东西是每序列的，不是每 token 的：seqlens_q、seqlens_k、cu_seqlens_q（前缀和）、block_tables_in（紧凑一维）、block_tables_offsets、state_block_idx、is_placeholder、input_ids_cpu、ngram_token_ids_cpu 等。
2. prepare_common_decode_buffers_zos：按 total_q 分配或复用 positions、valid_token_mask、gate_force_random、context_lens、query_indices、block_tables 六个公共缓冲区。
3. prepare_v4_decode_metadata（TileLang kernel）：一个 CTA 处理一条序列，把上面所有每 token 的字段一次性算完；block_tables_out 由第二个 kernel 展开。

为什么要把逐 token 的工作交给 kernel？对比一下工作量：一个典型的 decode 步 $N$ 条序列、$M$ 个 query 行（含投机），每行要写 positions、valid_token_mask、context_lens、query_indices、block_tables（max_num_blocks 列）、每个压缩倍率的 positions_c 与 slot_mapping_store_ca。在 CPU 上，block_tables 这一项就是 $M \times \text{max\_num\_blocks}$ 次写；在 GPU 上，同一个 kernel 用 $N$ 个 CTA 并行完成，并且与 driver 的其余准备工作重叠。

MetadataZOS 里有一个在调度论文里反复出现的字段——remap。它描述“上一步的请求在本步批里去了哪里”，用于把上一步的采样结果、草稿 token、已接受数量搬运到本步正确的位置。prepare_token_ids_zos 就是消费者：

1
gather_input_ids_zos(fresh_input_ids, prev_sampled_token_ids, prev_predict_token_ids,
2
                    remap, seqlens_q, cu_seqlens_q, input_ids, max_seqlen_q=nextn + 1)

有新请求（fresh）时，token 在 CPU 现成；延续的请求（prev）的 token 在 GPU 上由上一步产生，通过 remap 做设备端 gather。两条路分在一处，是因为“本步的 input_ids”必须是一块连续缓冲给 embedding 用，而它的来源是异构的。这个设计的背景在 DeepInfer 源码完全拆解（二）：异步调度与零开销调度 ZOS 里讨论过；本篇只关心它对元数据布局的影响。

分桶栈与零拷贝行视图

ZOS 路径里最容易被忽略的是 get_or_create_stacked。按倍率分桶的字段（klen、positions_c、slot_mapping_store_ca、kv_ca_index_offset_len）在 CUDA Graph 里不能是“每倍率一块独立张量”再加 Python 循环去填，因为 kernel 想在一次启动里写完所有倍率。做法是把它们 stack 成 [n_ratio, q_stride, ...] 的一块缓冲，kernel 按倍率与序列二维展开时，一次启动就能写完所有桶。

1
stacked = torch.empty([n, *size], ...)            # eager：新分配
2
return stacked, {k: stacked[i] for i, k in enumerate(keys)}

返回的 dict 是 stacked 的行视图。CUDA Graph 捕获时按 max_batch 捕获，回放时 total_q 更小，不能换 buffer，只能按捕获时的行间距切前缀：

1
q_stride = stride_elems // rest_numel
2
kernel_out = rows[0].as_strided([n, q_stride, *size[1:]], [stride_elems, *rows[0].stride()])

所以 kernel 写入的是一块 [n, q_stride, ...]、隔着 stride 写每行前 size[0] 个位置的视图。dict 里的值是 _cut 之后每行的前缀切片，与 kernel 的写入地址是同一块存储。代码里有一条很长的断言，检查所有行“同一 storage、等间距、连续”，失败就说明捕获流程被改动，零拷贝前提被破坏。

这个设计还有一个隐式收益：当 dense_ratios 与 sparse_ratios 不同时（比如稀疏层 4、稠密层 1），ca_ratios 是并集，同一个 token 可能出现在多个倍率的桶里。stack 让 kernel 的 grid 维度自然对应倍率，而不需要为每个倍率单独启动。

to_flatten_dict：只展开一层

1
for k, v in self.__dict__.items():
2
    if isinstance(v, dict) and v:
3
        for sub_k, sub_v in v.items():
4
            flat[f"{k}.{sub_k}"] = sub_v
5
    else:
6
        flat[k] = v

按倍率分桶的字段被展开成 klen.4、positions_c.2 这样的键。为什么不递归展开？注释给了答案：broadcast_tensor_dict 只处理一层嵌套，更深的结构在接收侧无法按路径回填。把“结构深度”限制在一层，是广播协议与数据类之间的契约。from_flatten_dict 是它的逆操作：按第一个 . 切分，把子键折回 dict。

broadcast：按路径收集非张量字段

InputMetadata.broadcast 遍历对象，把张量收进 tensor_list，把非张量按 a.b.c 路径收进 non_tensor_values；用 id(obj) 去重，防止同一个 tensor 被环形引用重复收集。接收侧按路径逐级 getattr / [key] 回填：

1
*parents, last = path.split(".")
2
obj = self
3
for part in parents:
4
    obj = obj[part] if isinstance(obj, dict) else getattr(obj, part)

这条路径式回填的好处是接收侧的 InputMetadata 对象结构可以独立构造，不需要发送侧序列化整个对象图。第二十一篇详细讲了共享内存环与 NCCL 两条通道，这里只强调一个易错点：空 dict 被当成非张量广播，因为“没有键”本身是信息；如果接收侧把空 dict 丢掉，compute_attn_info 会以为该倍率不需要展开，直接跳过。

recursive_copy_：把一份元数据搬进另一份

除了广播，InputMetadata 还提供了 recursive_copy_(other)：逐字段、逐字典值做 copy_，张量用 non_blocking=True。它解决的是“对象换了、缓冲区不变”的场景：当一份元数据的张量地址需要与另一份保持一致时（例如图捕获时为固定档位准备好的输入缓冲），重新构造对象会让 tensor 地址变化，依赖旧地址的图就会失效；recursive_copy_ 保持接收侧地址不变，只搬运内容。它与 get_or_create 的原地填充是同一目标的两条实现路径：前者用于“已有完整对象、只换值”，后者用于“字段不全、按名补齐”。

需要留意的是，recursive_copy_ 只处理 tensor 与一层 dict，标量字段被跳过。这与广播的“非张量走对象通道”互补：标量（num_reqs、total_q、is_prompt）在捕获时固定，回放时由调用方单独更新；张量则按块拷贝。若误以为标量也会被复制，图回放会用到捕获时的 num_reqs，在某些 kernel 里表现为“少算了后几个请求”。

图内图外：get_clean_list 与 compute_attn_info

InputMetadataV4.get_clean_list() 返回五个字段：

1
["kv_ca_index_len", "kv_swa_index", "kv_ca_index",
2
 "mlp_relu_schedule_metadata", "flash_mla_sched_meta_dict"]

它们都是展开结果或 kernel 调度缓存，不是调度器给出的原始信息。复用 CUDA Graph 时先把它们置 None，等 forward 开头再由 compute_attn_info 重建。为什么必须清空？因为图回放时的 total_q、block_tables、positions 都换成了本步的值，上一轮展开的索引若残留，就会指向旧位置。清空加重建把“派生字段”与“原始字段”明确分开，避免在图复用路径上出现难以排查的脏读。

compute_attn_info 还把 kv_ca_index_offset_len 展开成 kv_ca_index，把 kv_swa_index_offset_len 展开成 [total_q, window_size] 的 kv_swa_index，并为每个压缩倍率维护 klen。它的调用点有两个：普通 forward 开头，以及图回放前的清理之后。函数本身是幂等的（if self.kv_ca_index is None 才展开），所以重复调用不会重复分配。

mlp_relu_schedule_metadata 与 flash_mla_sched_meta_dict 是另一类：它们是 FlashMLA 与 DeepGEMM 的调度元数据缓存，由第一个消费该配置的层创建，后续同配置层复用。键里包含 ratio、是否 packed KV、是否用 SWA、topk 大小；这些参数相同的层可以共享同一份调度决策。把这类缓存放进 get_clean_list，说明它们也依赖本步形状，不能跨 step 复用。

cu_seqlens 的两个约定

prompt 的 cu_seqlens_q 是从 0 开始的前缀和：cu[0]=0，cu[i+1]=cu[i]+seqlen_q[i]，长度 n_seq+1，这是 FlashAttention 系列 varlen 接口的标准约定。

ZOS decode 的 cu_seqlens_q 在 prepare_v4_decode_metadata 的签名里写的是 inclusive cumsum：每个序列的结束位置，不含前导 0，长度 n_seq。prepare_decode_inputs_zos 里的 torch.cumsum(seqlens_q, dim=0, out=cu_seqlens_q) 也印证了这一点。两者同名不同义，跨路径复用代码时最容易错。判断方法很简单：看形状是 n_seq+1 还是 n_seq。

零长度与至少一行

代码里反复出现 max(1, ...) 和 append_zero_fill(1)：

• 压缩 token 数为 0 时补一行，slot=-1；
• 空序列在装载计划里补一行零；
• curr_ks 的步进用 max(1, seqlen_k / r)；
• InputMetadataV4 的注释明确写“零长度张量补到长度 1”。

这些不是防御性编程的装饰，而是形状契约：下游 kernel 的 grid、cu_seqlens 的单调性、按倍率 stack 的行数都假设“至少一行”。删掉任何一个 max(1, ...)，都可能让某个空序列在真实负载里触发越界或除零。

int32 的边界

kv_ca_index_offset_len 的 offset 是 i_q * max_num_blocks * (block_size // n)。当 $n=1$、total_q 很大、块表很长时，这个乘积可能超过 $2^{31}-1$。ZOS kernel 的 Python 包装里显式检查：

1
max_offset = total_q * max_num_blocks * (block_size // min(dense_ratios))
2
assert max_offset <= torch.iinfo(torch.int32).max

C++ 侧也有对应的 TORCH_CHECK_LE。这不是理论风险：128K 上下文、128 块大小、$n=1$ 时每行 offset 步进就是 131072，total_q 超过 16384 就会溢出。看到这条断言失败，应该先检查是不是把 dense 倍率配成了 1 而上下文又极长，而不是简单改成 int64——改成 int64 会让整个索引张量的带宽翻倍。

压缩边界与 max(1, ...) 的三次出现

max(1, seqlen_k / r) 在元数据里至少出现三次，每次的语义都不同：

1. cu_compressed_seqlens_q：给压缩器的行数定界，保证每条序列至少一行；
2. ks/ke 的 curr_ks 步进：保证下一序列在连续缓冲里的起点与装载结果一致；
3. kv_ca_index_offset_len 的 kv_ca_offset 步进：同样是连续缓冲的起点。

这三处必须严格同源。想象一条序列的 seqlen_k < r：它只产生 0 个完整压缩块，但装载计划会补一行零，压缩器也要处理一行，ks/ke 的下一序列起点要跳过这一行。如果只改了其中一处——比如为了省内存把装载计划的补零去掉——偏移量就会错位，attention 会读到上一序列的压缩 KV，表现为“偶发的首 token 错误”。修这类问题的正确做法不是加边界检查，而是找到三处并检查它们是否仍然用同一个公式；代码里三处都用 std::max(1, ...) 而不是各写各的，就是为了让这一点肉眼可查。

投机解码下 seqlens_q 的级联

开启投机后 seqlen_q = nextn + 1，元数据里有三处会随之变化：

• positions 每行是 q_start + j，草稿位置的绝对位置可以大于已确认长度；
• input_ids 的后 nextn 行来自 mtp.draft_tokens，它们还没有进入 token_ids；
• slot_mapping_store_swa 只给最后 swa_block_size 个 token 发槽位，草稿不会触发额外的写。

还有一个隐藏约束：prepare_decode_inputs_zos 里，当 DEEPINFER_ENABLE_MTP_VARLEN 打开时，cu_seqlens_q 由上一步的 metadata_zos.seqlens_q 现场 cumsum 得到，而不是从 C++ 重取：

1
torch.cumsum(metadata_zos.seqlens_q, dim=0, dtype=torch.int32, out=cu_seqlens_q)

原因是可变长投机下，每条序列本步实际要验证的草稿数取决于上一步接受了多少；这个数字只有上一轮采样后才知道，所以 seqlens_q 被“滞后一拍”地写进 MetadataZOS，下一轮元数据构造时再消费。元数据构造因此有一个回看着上一步结果的成分，这也是为什么 MetadataZOS 里同时存着 prev_sampled_token_ids、prev_predict_token_ids 与 prev_num_accepted_tokens。

空 ks 的 fallback

prepare_for_mirror_cp 里有一段：

1
kernel_ratios = tuple(sparse_ratios) or (1,)

纯 dense 模型的 ks 是空 dict，但 prepare_cp_mirror 的 kernel 至少需要一个倍率来算 ks/ke 输出。传 (1,) 让 kernel 有合法的 grid，输出随后被丢弃（if has_ks 才写回）。这类“给 kernel 一个 dummy 参数”的写法在 CP 路径上还有几处，读的时候不要误以为模型真的有一个倍率为 1 的稀疏层。

block table 的构建倍率

fill_decode_block_tables 用 min_compress_ratio 建表，而不是每个倍率各建一张。原因是块表描述的是物理块到逻辑块的映射，物理块大小固定；压缩倍率只改变“一个块能装多少个压缩 token”。因此所有倍率共享同一张表，读取时用 block_size // n 换算块内偏移。若某天有人按倍率各建一张表，kv_ca_index 的展开公式必须同步改成按倍率取表，否则会读到错误的物理块。