DeepInfer 源码拆解（三十）：KV Cache 的写入与取用——slot 映射、区块装载计划与量化布局的内核实现

KVCacheFormat：把布局写成一个可校验的描述符

DeepInfer 没有为每种模型写一套缓存代码，而是把所有差异收敛到 cache_common.py 里的 KVCacheFormat。它描述的是“一个 token 在缓存里占多少字节、这些字节怎么解释”，字段不多，但每个都不是随手加的：

1
@dataclass(frozen=True)
2
class KVCacheFormat:
3
    quantized_dtype: QuantizedFormat | None   # "e4m3" 或 "e2m1"
4
    quantized_elements: int                   # 量化后的逻辑元素数
5
    elements_per_sf: int                      # 每组多少个元素共享一个 scale
6
    bf16_elements: int                        # 历史遗留的 BF16 尾巴，现已禁止混存
7
    do_cast: bool                             # 读取时是否反量化成 BF16
8
    sf_format: ScaleFactorFormat | None       # "ue8m0" / "e4m3" / "fp32"
9
    note: str
10
    token_interleaved: bool = False

quantized_dtype 决定每个元素的位宽：e4m3 是 8 bit，e2m1 是 4 bit。quantized_elements 是逻辑元素数，不是字节数——FP4 的 512 个元素只占 256 字节。elements_per_sf 是量化分组大小，DeepInfer 的主缓存用 32，FP4 主缓存用 16，这直接对应 MXFP8 与 NVFP4 两种 microscaling 口径。do_cast 则区分“这是主缓存，读出来要反量化给注意力用”和“这是索引缓存，读出来仍保持量化值，交给索引打分内核自己处理”。

最容易被忽略的是 token_interleaved。它为真时，一条 token 的记录是 [量化数据 | scale] 紧挨着存放，整条记录作为一个整体在 block 内连续排列；为假时，数据区和 scale 区在 block 级别分开，先放整块的 data，再放整块的 scale。这个差别不是审美问题：前者让主缓存的每次 token 读写只触碰一段连续字节，后者让索引缓存的 scale 可以按块做 TMA 对齐的列主序排布。

描述符在构造时就做了互斥校验，避免运行期出现“半量化半 BF16”的缓存：

1
def __post_init__(self):
2
    if self.quantized_dtype is None:
3
        assert self.quantized_elements == 0 and self.elements_per_sf == 0
4
        assert self.sf_format is None
5
    else:
6
        assert self.quantized_elements > 0 and self.elements_per_sf > 0
7
        assert self.sf_format in ("ue8m0", "e4m3", "fp32")
8
        if self.bf16_elements:
9
            raise ValueError("Mixed quantized/BF16 KV caches are no longer supported")
10
        if self.do_cast and not self.token_interleaved:
11
            raise ValueError("Main/SWA KV caches must be token-interleaved")

“主缓存必须 token-interleaved”这条约束值得展开。主缓存的写入是逐 token 的：本轮新算出的 K/V 有多少个 token，就写多少个 slot，每个 slot 落在不同物理块的不同偏移上。如果数据和 scale 在 block 级别分区，一次写入要拆成两次访存、两个基址、两套偏移换算；交错之后，一个 token 的 528 字节是一段连续区间，slot 到地址的换算只做一次，写内核的合并访存也更好做。反过来，索引缓存是按块整体装载的，数据区与 scale 区分开更利于按块做批量搬运和 TMA 描述。

五种正式格式：不是“所有组合”，而是产品需要的组合

SUPPORTED_KV_CACHE_FORMATS 只列了五种格式，分别是：

主缓存的 512 个元素对应 MLA 的 latent 向量：DeepSeek 系列把每层的 K/V 压成一个低秩 latent，缓存只需要存这一条 latent，而不是每个 head 各存一份 K 和 V。索引缓存的 128 个元素则对应“闪电索引器”那一侧的低维 key——它只用来给 token 打分、挑 top-k，不参与最终注意力，所以可以压得更狠，连反量化都省掉。可参考 DeepInfer 源码拆解（十一）：注意力族与压缩 KV 中关于压缩 KV 与索引器共享的讨论。

bytes_per_token 由这些字段直接算出来：

1
@property
2
def quantized_bytes(self) -> int:
3
    if self.quantized_dtype is None:
4
        return 0
5
    bits_per_element = 4 if self.quantized_dtype == "e2m1" else 8
6
    return self.quantized_elements * bits_per_element // 8
7

8
@property
9
def num_scales(self) -> int:
10
    if self.quantized_dtype is None:
11
        return 0
12
    return (self.quantized_elements + self.elements_per_sf - 1) // self.elements_per_sf
13

14
@property
15
def bytes_per_token(self) -> int:
16
    return self.quantized_bytes + self.sf_bytes + self.bf16_elements * torch.bfloat16.itemsize

以 V41_FORMAT 为例：512 个 FP8 元素是 512 字节，512/32 共 16 个 scale，UE8M0 每个 1 字节，总计 528 字节。对比 BF16 下 512 个元素要 1024 字节，压缩比接近 2 倍；V41_FP4_FORMAT 则是 256 字节数据加 32 个 e4m3 scale，共 288 字节，相比 BF16 约 3.6 倍。这些数字不是宣传口径，而是后面所有对齐、分块、带宽估算的起点。

对齐：为什么 block stride 要补齐，token stride 却只要求 gcd

alloc_kv_cache 负责把一块连续的 uint8 存储切成逻辑上的 [num_blocks, block_size, 1, bytes_per_token] 视图：

1
alignment = math.lcm(32, layout.bytes_per_token) if layout.token_interleaved else 32
2
block_bytes = block_size * layout.bytes_per_token
3
block_stride = (block_bytes + alignment - 1) // alignment * alignment
4
storage = torch.empty((num_blocks, block_stride), dtype=torch.uint8, device=device)
5
if layout.token_interleaved:
6
    return storage.as_strided(
7
        (num_blocks, block_size, 1, layout.bytes_per_token),
8
        (block_stride, layout.bytes_per_token, layout.bytes_per_token, 1),
9
    )
10
return storage[:, :block_bytes].view(num_blocks, block_size, 1, layout.bytes_per_token)

这里有两层对齐。第一层是 block 级：每个块的起始地址必须对齐到 32 字节，并且对交错布局还要对齐到 lcm(32, bytes_per_token)。528 和 32 的最小公倍数是 1056，也就是说 V4.1 主缓存的每个 block 占用的字节数会被补到 1056 的整数倍。这样做的原因是 TileLang 内核需要给出 value_block_stride 的假设（T.assume(value_block_stride % 32 == 0)），编译器才能安全地做向量化与地址折叠；如果每个块的跨度是奇数倍的 528，某些块的首地址就会落到 16 字节甚至 8 字节边界上。

第二层是 token 级：stride_alignment = gcd(32, bytes_per_token)。528 与 32 的最大公约数是 16，于是 token 跨度只保证 16 字节对齐。这是一个刻意的折中：如果把每条 token 记录都补齐到 32 字节，528 会变成 544，缓存体积凭空多出 3%；而 16 字节对齐已经足够让内核用 16 字节向量读写数据段，scale 段则按 1 字节粒度单独处理。读取内核里 vector_size = min(group_size, 32, value_token_alignment) 正是吃这个约束：V4.1 的 group 是 32，但 token 对齐只有 16，所以一次向量化搬 16 个元素，一个 32 元素的量化组拆成两个 channel chunk。

FP4 主缓存还有一处特殊处理。split_kv_cache 在交错布局下返回的 quantized_data 是 int8 视图，因为 PyTorch 没有原生的 4 bit 张量；内核工厂通过 fp4_cache_kernel_kwargs 把逻辑的 value/scale stride 传进去：

1
return {
2
    "fp4_value_block_stride": quantized_data.stride(0) * 2,
3
    "fp4_value_stride": quantized_data.stride(1) * 2,
4
    "fp4_scale_stride": sf_data.stride(1),
5
}

* 2 把“每字节两个 FP4 元素”换算成逻辑元素跨度。TileLang 侧用 T.float4_e2m1fn 声明张量，编译器据此生成按 nibble 寻址的代码。这里如果直接把 int8 的 stride 传下去，内核会按一个元素一字节来算地址，所有 FP4 数据都会错位——这是量化缓存里最典型的一类“看起来能跑、结果全错”的 bug。

入口契约：slot 是逻辑页号乘页大小加页内偏移

写入 API 的签名很短，但每个参数的约定都很硬：

1
def store_kv_cache(
2
    kv: torch.Tensor,              # [n, ..., num_elements], bf16
3
    kv_cache: torch.Tensor,        # 分页缓存
4
    slot_mapping: torch.Tensor,    # [n], int32
5
    *,
6
    layout: KVCacheFormat | None = None,
7
) -> None:

kv 的第一维是 token 数，其余维度全部 flatten 成一条缓存记录；slot_mapping 的每个元素是 block_id * block_size + block_offset，-1 表示这个 token 是填充、不写入。把“块号 + 块内偏移”压成一个 int32 的 slot，是分页缓存最常用的做法：调度器只需要维护一张扁平数组，内核里一次除法、一次取模就能还原出物理地址，不需要再查块表。

入口先做格式推断和形状校验：

1
layout = resolve_kv_cache_layout(kv_cache, layout)
2
assert kv.device == kv_cache.device == slot_mapping.device
3
n = slot_mapping.shape[0]
4
kv = _flatten_kv(kv, n, layout)
5
if n == 0:
6
    return

resolve_kv_cache_layout 有一个容易被低估的作用：当调用方没有显式传 layout 时，它靠 kv_cache.dtype、shape[-2:] 和 bytes_per_token 反推格式。比如 uint8、[num_blocks, block_size, 1, 528] 只可能是 V4.1 主缓存；[..., 1, 296] 则可能是 FP4 主缓存或某种索引缓存，需要 sf_format 辅助区分。这种“按字节宽度反查”的设计让上层不必到处传描述符，但也意味着任何布局改动都必须同步更新 SUPPORTED_KV_CACHE_FORMATS，否则线上会出现“缓存分配成功、写入时报 unsupported shape”的问题。

NPU 与 GPU 在这里走了两条不同的路径。NPU 上直接调用 ops.store_kv_cache_tl，并把 cache 视图降成三维；GPU 上则先 split_kv_cache 拆出 data 与 scale 两个 strided 视图，再交给内核。这不是简单的平台分支，而是因为两边的 TileLang 后端对 strided tensor 的支持程度不同：GPU 侧需要在 host 上把 block/token 两级 stride 显式算好，NPU 侧的 kernel 工厂自己接受原始 4D 视图。

三块内核逻辑：FP8、FP4、纯 BF16

store_kv_cache_tl 是一个按 layout 特化的工厂函数，返回的 prim_func 结构相同，但内部有三套分支。三套分支共享同一个网格：T.Kernel(T.ceildiv(n, blk), threads=threads)，每个 CTA 处理 blk 个 token。blk 的选择来自一条经验公式：

1
blk, threads = min(max(4096 // input_dim, 1), 16), 64   # 量化路径
2
blk, threads = max(4096 // input_dim, 1), 256           # BF16 路径

主缓存一条记录 512 个元素，4096/512 = 8，所以一次处理 8 个 token、64 个线程；BF16 路径元素更多、寄存器压力更大，于是线程数提到 256。这个公式的本质是让每个 CTA 的工作集大致落在 4K 元素的量级，既不会因为 token 太少而让 CTA 数量爆炸，也不会因为 token 太多而把寄存器打满。

FP8 分支的核心是“按 32 元素一组求 amax、算 UE8M0 scale、量化、写回”：

1
channel = T.alloc_fragment((blk, groups, group_size), "bfloat16")
2
channel_amax = T.alloc_fragment((blk, groups), "bfloat16")
3
scale_local = T.alloc_fragment((blk, groups), "float32")
4
sf_local = T.alloc_fragment((blk, groups), sf_dtype)
5
cast_result = T.alloc_fragment((blk, groups, group_size), T.float8_e4m3fn)
6

7
for i, j, k in T.Parallel(blk, groups, group_size):
8
    if bx * blk + i < n and j < valid_groups:
9
        channel[i, j, k] = kv[bx * blk + i, j * group_size + k]
10
    else:
11
        channel[i, j, k] = 0
12
T.reduce_absmax(channel, channel_amax, dim=-1)

groups 取的是 sf_dim（scale 的存储槽数），而不是 valid_groups（有效的量化组数）。这两个值在 V4.1 主缓存里通常相等；差异出现在 scale 存储带 padding 的布局里（store 内核注释里提到的“8B padding”）。把超出 valid_groups 的槽填 0 再统一走 reduce，是为了保持 fragment 形状固定，避免 TileLang 在编译期展开出两套循环。

UE8M0 的换算值得单独讲。它不是一个普通的浮点 scale，而是“只存指数、不存尾数”的 8 位格式：数值等于 $2^{e-127}$，其中 $e$ 是存的 8 位无符号整数。DeepInfer 用位运算同时得到两个东西——给量化乘法用的 float32 scale，和写回缓存的 8 位指数：

1
scale = T.max(T.cast(amax[i, j], "float32"), FP8_AMAX_MARGIN) / FLOAT8E4M3_MAX
2
exp_scale = ue8m0_exp_scale(scale)
3
scale_local[i, j] = T.reinterpret("float32", (exp_scale + 127) << 23)
4
sf_ue8m0[i, j] = T.cast(-exp_scale + 127, "uint8")

(exp_scale + 127) << 23 直接构造出 float32 的指数位，尾数全 0，因此 scale_local 就是 $2^{exp\_scale}$。而写回的 -exp_scale + 127 是它的倒数指数：读取时内核会把存下来的字节左移 23 位再 reinterpret 成 float32，得到 $2^{-exp\_scale}$，再乘回量化值。一个存乘法尺度、一个存反量化尺度，两者互为倒数，这样写入端只需要做一次除法，读取端只需要做一次移位。

这里还有一个工程细节：FP8_AMAX_MARGIN 用来兜住全零组。如果一组 32 个元素全是 0，amax 为 0，scale = 0 / 448 = 0，取对数会得到负无穷；加上一个最小正数下界之后，scale 变成一个很小的正数，量化结果仍然是 0，但指数是有限的，不会污染后续的块统计。这不是数学上的洁癖，而是因为 UE8M0 的指数一旦落到非法值，反量化会得到 NaN，而 NaN 会沿着注意力分数传播到整行。

FP4 分支的骨架相同，但 scale 有两种口径。V41_FP4_FORMAT 用 e4m3 scale，每组 16 个元素，scale 直接存“amax / 6”的 FP8 值，读取端用 1 / sf 作为乘数；索引缓存则用 UE8M0，直接存指数。代码里用 layout.sf_format 分流：

1
if layout.sf_format == "e4m3":
2
    sf_local[i, j] = T.cast(
3
        T.min(
4
            T.max(T.cast(channel_amax[i, j], T.float32),
5
                  FLOAT8E4M3_MIN_POSITIVE * FLOAT4_E2M1_MAX) * (1.0 / FLOAT4_E2M1_MAX),
6
            FLOAT8E4M3_MAX,
7
        ),
8
        T.float8_e4m3fn,
9
    )
10
    scale_local[i, j] = 1.0 / T.cast(sf_local[i, j], T.float32)

1.0 / sf 而不是 amax / 6，是为了让写入与读取使用同一份被 FP8 舍入过的 scale。如果写入用未舍入的 amax / 6 做乘法，读取却用舍入后的 sf 做反乘，量化误差会多出一层；这里把舍入后的值再取倒数，保证两端严格互逆。

纯 BF16 分支则简单得多：没有 reduce、没有 cast，只有一次带掩码的散射：

1
for i, j in T.Parallel(blk, input_dim):
2
    index = bx * blk + i
3
    slot = slot_mapping[index] if index < n else -1
4
    T.assume(-1 <= slot < num_blocks * block_size)
5
    if slot != -1:
6
        T.assume(0 <= index < n)
7
        kv_cache_value[slot // block_size, slot % block_size, j] = kv[index, j]

注意 T.assume 的位置：它在 if slot != -1 之外，先告诉编译器 slot 的取值范围，再在分支里做实际的除法取模。这样编译器可以把 slot // block_size 和 slot % block_size 优化成一次乘加移位，而不是真正的整数除法。

-1 与 slot 0：两个看起来奇怪但必须存在的补丁

写入路径里有两段代码初读会觉得“为什么要这么绕”。第一段是 _masked_cache_store 对 slot 0 的处理：

1
safe_slots = slots.clamp_min(0).long()
2
block, offset = safe_slots // block_size, safe_slots % block_size
3
old_slot_zero = cache[0, 0].clone()
4
_slot_scatter(cache, block, offset, values)
5
slot_zero = slots == 0
6
slot_zero_valid = slot_zero.any()
7
row0_index = slot_zero.to(torch.int32).argmax().reshape(1)
8
row0 = values.index_select(0, row0_index).squeeze(0)
9
cache[0, 0] = torch.where(slot_zero_valid, row0, old_slot_zero)

这段是参考实现（store_kv_cache_ref）里的路径，不是 TileLang 内核。它解决的是“padding token 的 slot 被压成 -1 之后，clamp_min(0) 会把它们全写到 slot 0”的问题。如果多个 padding token 都落在 slot 0，简单的散射会让它们互相覆盖，最后 slot 0 里留下的是任意一个 padding token 的值。参考实现的处理是：先正常散射，再单独挑出第一个 slot 0 的 token 写回去，其余 padding 不参与。这样 slot 0 的内容是确定的，测试可复现，也不会因为 padding 顺序不同而让参考结果漂移。

第二段是 NPU 上的 _slot_scatter：

1
flat = torch.as_strided(cache, (cache.untyped_storage().nbytes() // cache.element_size(),), (1,), 0)
2
dst = cache.storage_offset() + block * cache.stride(0) + offset * cache.stride(1)
3
if cache.dim() > 2:
4
    dst = dst.unsqueeze(1) + torch.arange(cache.shape[2], device=dst.device)
5
flat.index_put_((dst.reshape(-1),), values.reshape(-1))

它绕开的是 Ascend 上 IndexPutV2 的 32 位偏移回绕：当缓存超过 2 GB 时，普通的 cache[block, offset] = values 会把 64 位地址截成 32 位。做法是把整个 storage 视为一维，自己算 64 位线性偏移，再用 index_put_ 写入。这段代码只在 PLATFORM.is_npu 下生效，GPU 路径仍然用普通的索引赋值。跨平台推理框架里，这类“某个后端有 bug、另一个后端没有”的补丁不可避免，关键是把它收在一个函数里并写清楚原因，而不是散落在业务代码中。

写入路径最后还要回答一个问题：为什么主缓存的写入常常和 norm、RoPE、Hadamard 变换融在一起？ 在 attn.py 里，SWA 分支调用的是 fused_norm_rope_hadamard_cast_store_kvcache，而不是单独的 store_kv_cache：

1
result = fused_norm_rope_hadamard_cast_store_kvcache(
2
    kv_swa,
3
    norm_args=NormArgs(self.kv_swa_norm.weight, self.kv_swa_norm.variance_epsilon),
4
    rope_args=RopeArgs(
5
        input_metadata.positions,
6
        self.rotary_emb.cos_sin_cache,
7
        self.qk_rope_head_dim,
8
        self.rotary_emb.is_neox_style,
9
        offset=self.qk_head_dim - self.qk_rope_head_dim,
10
    ),
11
    cast_args=CastArgs("e4m3", self.swa_cache_num_per_channels, round_sf=True, sf_format="ue8m0"),
12
    out_args=KVCacheOutArgs(swa_cache, input_metadata.slot_mapping_store_swa),
13
)

解码阶段没有预填充那样的“先把 KV 拼成连续缓冲区”的步骤，新 token 的 K/V 只在寄存器里存在一瞬间，最划算的做法是在同一次 kernel 里完成 RMSNorm、RoPE、量化与按 slot 写入。_prepare_swa_kv 在解码分支给出的 cast_args 就带着 CastArgs("e4m3", ..., round_sf=True, sf_format="ue8m0") 和 KVCacheOutArgs(swa_cache, input_metadata.slot_mapping_store_swa)，也就是说 norm 之后的 BF16 值直接经过 cast 与 store，中间不落盘。这样 K/V 不需要以 BF16 的形态完整写一次显存，带宽省了一整趟。至于 Hadamard 变换是否在这个融合内核里完成，取决于调用内核的版本与层配置（预填充侧的 CastArgs 还带有 cast_back=True，为的是随后把 BF16 读回给注意力），读代码时应以具体 kernel 实现为准，而不是从函数名推断。store_kv_cache 单独存在，主要是给预填充、参考实现和测试用；真正跑在线上的解码热路径，走的是融合版本。可参考 DeepInfer 源码拆解（十五）：Linear、RMSNorm 与 RoPE 中关于融合残差与量化的讨论。

为什么解码不读、预填充要读

先把两条路径的差别摆在最前面，否则后面所有设计都会显得多余。

解码阶段每个请求只前进一个 token。注意力要读的是这个请求的全部历史 KV，但这些历史在分页缓存里本来就是按块连续存放的：只要知道块表，FlashMLA 就能按“页号 + 页内偏移”直接访问，不需要把整段历史搬到一个新缓冲区。DeepInfer 在解码路径里做的是索引换算，不是数据搬运：把逻辑 token 下标转成物理 slot，然后让内核按 slot 读。

预填充阶段不一样。一个 chunk 里可能同时包含几百到几千个 query token，核心注意力需要一个 [total_kv, head_dim] 的连续 KV 缓冲区，并且每一行 query 还要通过一张索引表指定自己能看到哪些 key。这个缓冲区由三部分组成：SWA 历史（来自环形缓存）、本轮新算出的 SWA KV（还在寄存器/临时张量里）、压缩主缓存（来自分页的 ca_cache）。三者的物理位置完全不同，必须先在 GPU 上拼成一段连续内存，注意力内核才能用统一的 kv_index 去索引。

所以 load_paged_kv_cache 只会出现在预填充路径。它的任务不是“读 KV 给注意力”，而是“把分散的物理块整理成注意力能接受的连续布局”。

load_segments：把块表压缩成四元组

读取 API 的输入不是块表，而是一张装载计划：

1
def load_paged_kv_cache(
2
    kv_cache: torch.Tensor,        # [num_blocks, block_size, 1, bytes_per_token]
3
    load_segments: torch.Tensor,   # [num_segments, 4], int32
4
    total_k: int,                  # 输出逻辑 token 数
5
    out: torch.Tensor | None = None,
6
    *,
7
    layout: KVCacheFormat | None = None,
8
    sf_format: ScaleFactorFormat | None = None,
9
) -> KVCacheTensors:

每一行的四个整数是 [src_block_id, src_begin, src_end, dst_offset]：从物理块 src_block_id 的第 src_begin 到第 src_end 个 token（左闭右开），搬到输出缓冲区的第 dst_offset 个位置。src_block_id == -1 表示这一段不读缓存，直接写零。total_k 是输出总长度，所有段的目标区间必须不重叠且落在 [0, total_k) 内。

这个格式的关键在于段是块粒度的。为什么不是“每个 token 一行”的映射表？因为那会把元数据放大到和 token 数同阶：十万 token 的历史就是十万行、每行 16 字节，1.6 MB 的元数据要先传上 GPU，内核还要为每个 token 做一次独立的地址计算。改成块粒度之后，一段最多覆盖 block_size 个 token，元数据量缩小到块数量级；内核里每个 CTA 负责一段，段内地址是连续的，可以用向量化搬。

C++ 侧的 BlockLoadPlanBuilder 负责生成这张表。三个 append 方法对应三种来源：

1
void append_segment(int32_t src_block_id, int32_t src_begin, int32_t src_end) {
2
    segments_.push_back(BlockLoadSegment{src_block_id, src_begin, src_end, total_k_});
3
    total_k_ += src_end - src_begin;
4
}
5

6
void append_zero_fill(int32_t num_tokens) {
7
    while (num_tokens > 0) {
8
        const int32_t segment_len = std::min(num_tokens, source_block_size_);
9
        append_segment(-1, 0, segment_len);
10
        num_tokens -= segment_len;
11
    }
12
}
13

14
void append_paged_range(const std::vector<int> &block_table,
15
                        int32_t token_begin, int32_t token_end) {
16
    while (token_begin < token_end) {
17
        const int32_t block_index = token_begin / source_block_size_;
18
        const int32_t src_block_id = block_table[block_index];
19
        const int32_t src_begin = token_begin % source_block_size_;
20
        const int32_t segment_len = std::min(token_end - token_begin,
21
                                             source_block_size_ - src_begin);
22
        append_segment(src_block_id, src_begin, src_begin + segment_len);
23
        token_begin += segment_len;
24
    }
25
}

append_paged_range 的循环有一个隐含假设：同一个逻辑块在块表里只出现一次，且块内 token 是连续的。对压缩 KV 来说，block_table 存的是压缩 token 的块号，一个块装 block_size / compress_ratio 个压缩 token；source_block_size_ 就取这个值。循环每跨过一个块边界就切一段，因此一个逻辑区间最多被切成“块数 + 1”段，段与段之间在目标缓冲区里紧挨着，dst_offset 自动累加。

append_zero_fill 则用于“这段历史不存在但输出必须有占位”的场景。压缩装载计划里有一个很能说明问题的分支：

1
const int32_t num_tokens = p_seq->seqlen_k / compress_ratio;
2
if (num_tokens == 0)
3
    builder.append_zero_fill(1);
4
else
5
    builder.append_paged_range(p_seq->block_table, 0, num_tokens);

当序列长度还不足一个压缩组时，压缩缓存里没有任何有效 token。但后续的注意力元数据需要每个序列至少有一个压缩 token 的位置，否则 cu_compressed_seqlens 会出现零长度区间，索引计算要额外处理空序列。这里的做法是塞一个零段：输出里有一行零，长度账目连续，真正的可见性由后面的 ks/ke 索引控制，零行不会被读到。零填充是形状与确定性的补丁，不是语义上的 mask——mask 仍然由索引里的 -1 表达，这一点后面还会再遇到。

SWA 的装载计划更能体现“物理顺序不等于逻辑顺序”。SWA 缓存是环形缓冲区，一个序列的 window_size 个槽位可能跨过环的回绕点；如果历史不足一个窗口，前面还要补零。构造逻辑是：

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

append_ring_range 和 append_paged_range 的区别在于：它只有一个物理块号 state_id，所有 token 都在这个块里，但逻辑区间可能跨过 block_size 的边界，所以要按环的模数切成最多两段。加上前面的零填充，一个序列最多三段——这正是 InputMetadata 里“at most three segments per sequence”注释的来源。total_k_swa 最终严格等于 num_seqs * window_size，每个序列的窗口长度统一，后面的 kv_swa_full_index 才能用一个固定宽度的二维张量表示。

内核：按段并行，段内按 token tile 并行

load_paged_kv_cache_tl 的网格是二维的：

1
with T.Kernel(num_segments, num_token_tiles, threads=threads) as (segment_idx, token_tile_idx):
2
    T.pdl_sync()
3
    src_block_id = load_segments[segment_idx, 0]
4
    src_begin = load_segments[segment_idx, 1]
5
    src_end = load_segments[segment_idx, 2]
6
    dst_offset = load_segments[segment_idx, 3]
7
    token_base = token_tile_idx * token_tile

第一维是段，第二维是段内的 token tile。token_tile 的取值按布局分三档：索引缓存直接用整个 block_size（块内一次搬完），纯 BF16 用 min(block_size, 8)，量化主缓存用 min(block_size, 32)。这三个数字对应三种寄存器压力：索引缓存一条记录只有 128 字节，整块搬进来也不会爆寄存器；BF16 一条记录 1024 字节，8 个 token 就是 8 KB；量化主缓存一条 528 字节，32 个 token 约 16.5 KB，已经接近一个 CTA 的合理上限。

T.pdl_sync() 是 Hopper 之后的一个细节。装载计划可能是由上游的依赖 kernel（比如 CP 的 mirror 重排或 ZOS 的元数据生成）刚写出来的，PDL（Programmatic Dependent Launch）允许当前 kernel 在上一 kernel 还没完全结束时提前启动，但读元数据之前必须 pdl_sync 等到依赖满足。DeepInfer 在 store/load 两个内核里都保留了这个同步，注释写得很直白：slot_mapping may be produced by the dependency launch。这类同步点如果漏掉，表现是“大多数 step 正确、偶发读到上一轮的 slot”，非常难查。

三条分支里，最复杂的是 do_cast 路径。它要读量化值和 scale，做反量化，再写 BF16 输出：

1
if layout.do_cast:
2
    sf_float = T.alloc_local((1,), T.float32)
3
    for token_offset, group_chunk in T.Parallel(token_tile, valid_groups * channel_chunks):
4
        block_offset = token_base + token_offset
5
        if block_offset < block_size and src_begin <= block_offset and block_offset < src_end:
6
            dst = dst_offset + block_offset - src_begin
7
            group = group_chunk // channel_chunks
8
            channel_chunk = group_chunk % channel_chunks
9
            channel_base = group * group_size + channel_chunk * vector_size
10
            if src_block_id == -1:
11
                for channel in T.vectorized(vector_size):
12
                    kv[dst, channel_base + channel] = 0
13
            else:
14
                if layout.sf_format == "ue8m0":
15
                    sf_float[0] = T.reinterpret(
16
                        T.float32,
17
                        T.cast(kv_cache_sf[src_block_id, block_offset, group], T.uint32) << 23,
18
                    )
19
                else:
20
                    sf_float[0] = T.cast(kv_cache_sf[src_block_id, block_offset, group], T.float32)
21
                for channel in T.vectorized(vector_size):
22
                    kv[dst, channel_base + channel] = T.cast(
23
                        T.cast(kv_cache_value[src_block_id, block_offset, channel_base + channel], T.float32) * sf_float[0],
24
                        T.bfloat16,
25
                    )

这里把“组”和“通道块”拆成两层并行，是为了让最内层能用 T.vectorized。vector_size = min(group_size, 32, stride_alignment)：V4.1 主缓存的组是 32、token 对齐是 16，于是 vector_size = 16，channel_chunks = 2；FP4 主缓存的组是 16、token 对齐是 16，于是 vector_size = 16、channel_chunks = 1。如果直接按 32 个元素做向量化，编译器会发现起始地址只保证 16 字节对齐，退化成一堆标量读写；显式拆成 16 元素一块，反而能生成 16 字节的向量指令。

UE8M0 的读取端就是写入端的逆运算：把 8 位指数左移 23 位，reinterpret 成 float32，得到 $2^{-e}$，再乘量化值。这里没有查表、没有 pow，只有一次移位和一次类型双关。FP8 的乘加在 CUDA 上会走 FP32 中间精度，最后 cast 回 BF16；这是 DeepInfer 的通用选择——反量化结果只用于一次点积，保留 FP32 中间值不会带来额外存储，反而能减少二次舍入。

索引缓存分支（is_index）不反量化。它把量化值和 scale 原样搬到输出，scale 在 UE8M0 下按 int32 视图搬运：

1
output_sf_shape = (total_k,) if is_index and index_sf_dim == 1 else (total_k, index_sf_dim) if is_index else (total_k, output_dim)
2
...
3
index_sf_dim = layout.sf_bytes // 4 if layout.sf_format == "ue8m0" else layout.num_scales
4
index_sf_dtype = "int32" if layout.sf_format == "ue8m0" else cache_sf_dtype

sf_bytes // 4 把 4 个 UE8M0 字节打包成一个 int32。为什么读取端要打包，写入端却不打包？因为写入是逐 token 的，scale 只有 16 个字节，按字节写最简单；读取是按段搬的，int32 打包能让一个线程一次搬 4 个 scale，并且和下游索引打分内核期望的 int32 scale 布局对齐。sf_data_as_int32 在 host 侧做同样的视图转换，注释里特意说明“View packed UE8M0 scale bytes in the int32 form used by load kernels”。

纯 BF16 分支最简单：没有 scale、没有 cast，只有一次带零填充的搬运。它的 token_tile 只有 8，因为 BF16 记录宽、寄存器占用高。向量化宽度需要单独说明：stride_alignment 对 BF16 布局返回 32（token_interleaved=False 时固定为 32 字节），而 vector_size = min(group_size, 32, alignment)、group_size 在非量化分支里取 1，所以最内层实际是按元素搬运，只保证 32 字节的 stride 假设成立，并不会自动生成 16 元素的宽向量读。这与量化分支“组内可向量化、组间按 token tile 并行”的结构不同——BF16 布局本来就没有量化组，最内层的循环维度直接是通道，编译器自行决定向量宽度。

调用现场：[SWA 历史 | 本轮 SWA | 压缩主缓存] 的拼接

把内核放回 attn.py 的预填充路径，才能看清 dst_offset 的意义。代码先算三段长度，再一次性分配输出：

1
total_k_swa_cached = input_metadata.total_k_swa if self.use_swa else 0
2
total_k_swa_cur = (input_metadata.mirror_cp_reverse_index.shape[0] if use_cp else kv_swa.shape[0]) if self.use_swa else 0
3
total_k_compressed = input_metadata.total_k_compressed[self.ratio_key] if self.has_packed_kv else 0
4
total_kv = total_k_swa_cached + total_k_swa_cur + total_k_compressed
5
main_kv_offset = total_k_swa_cached + total_k_swa_cur
6
kv = self._prepare_prefill_kv(total_kv, main_kv_offset, q_b.device, cross_layer_runtime)

kv 是一个 [total_kv, 1, qk_head_dim] 的 BF16 缓冲区。SWA 历史放在最前面，本轮新 SWA 紧随其后，压缩主缓存从 main_kv_offset 开始。接下来两次 load_paged_kv_cache 分别填前段和后段：

1
load_paged_kv_cache(
2
    kv_cache.ca_cache,
3
    input_metadata.load_segments_ca[self.ratio_key],
4
    total_k_compressed,
5
    out=kv[total_k_swa_cached + total_k_swa_cur :],
6
)
7

8
load_paged_kv_cache(
9
    kv_cache.swa_cache,
10
    input_metadata.load_segments_swa,
11
    total_k_swa_cached,
12
    out=kv[:total_k_swa_cached],
13
)

注意两次调用都传了 out。load_paged_kv_cache 的默认行为是分配一块新输出，但在预填充热路径上，每层都分配一次会把 allocator 打爆，也会破坏 CUDA Graph 的固定地址假设。传入 out 之后，内核直接写进 kv 的切片，dst_offset 是相对于这个切片的局部偏移。切片本身是连续的（kv[a:b] 在第一维上切），所以内核的 dst 计算不需要知道全局偏移。

这里有一个容易写错的顺序：SWA 必须先装载、后写入。代码里专门有一行注释：

1
# The SWA store slots may overlap the ranges loaded above, so prefill must store only after the cache load.

预填充的 chunk 可能和 SWA 环形缓存的槽位重叠：本轮新 token 要写进的 slot，可能正是同一批 query 需要读的历史窗口。如果先 store 再 load，装载计划读到的就是被本轮覆盖后的数据，历史窗口里会混进未来 token。先 load 到临时 kv，再从寄存器写回缓存，才能保证“读到的历史是写之前的历史”。这类顺序约束不会在类型系统里体现，只能靠调用点的注释和测试兜住。

压缩主缓存还有一层跨层复用。MLA 的压缩 KV 在同一组的多个层之间共享，如果每层都重新装载一遍，同样的几十 MB 数据会被搬 N 次。DeepInfer 用一个 cross_layer_runtime 记录“源层已经装载好的缓冲区”：

1
if self.use_shared_kv and self.is_kv_source:
2
    cross_layer_runtime.loaded_prefill_kv = kv
3
    cross_layer_runtime.loaded_prefill_kv_main_offset = main_kv_offset

is_kv_source 的层在装载完压缩段之后，把整块 kv 和压缩段的起点记进运行时上下文；后续共享层不再重复调用 load_paged_kv_cache(ca_cache, ...)，而是直接复用这块已经拼好的缓冲区，必要时只补自己那一段 SWA。这是压缩 KV 架构带来的额外优化空间：传统 MHA 每层 K/V 都不同，没法共享；MLA 把 K/V 压成一条与层数无关的 latent 之后，跨层复用才成立。具体的共享边界在 DeepInfer 源码拆解（十一）：注意力族与压缩 KV 里有更完整的讨论。

解码侧：把逻辑下标直接换算成物理 slot

解码不装载数据，但要做一次索引换算。convert_kv_index_tl 的整个函数只有十几行：

1
for k_i in T.Parallel(k):
2
    if k_i < context_lens[bid, sid]:
3
        x = topk_idxs_fragment[k_i]
4
        y = block_tables[bid, x // block_size] * block_size + x % block_size
5
        topk_idxs_fragment[k_i] = y
6
    else:
7
        topk_idxs_fragment[k_i] = -1

topk_idxs 是索引器挑出来的逻辑 token 下标，block_tables 是这个序列的块表。换算公式就是分页的标准地址变换：逻辑下标除以块大小得到逻辑块号，查表得到物理块号，再乘块大小加页内偏移。越界的 top-k 位置写成 -1，由 FlashMLA 当作无效索引跳过。整个 kernel 原地改写 topk_idxs，不分配新张量，也不搬任何 KV 数据。

SWA 分支的解码索引则是另一条路。每个 query token 的可见窗口由 kv_swa_index_offset_len 描述成 [start_slot, valid_length]，get_kv_swa_index_decode 把它展开成 [total_q, window_size] 的 slot 表。窗口可能跨环回绕，所以展开逻辑要处理两段；不足一个窗口时剩余位置填 -1。这和预填充的 expand_swa_index 是同一个思路，只是预填充的索引指向拼接后的连续缓冲区，解码的索引指向物理 slot。

两条路径的带宽账可以直接对比。假设 128K 上下文、窗口 4K、block 64：

• 预填充若每层都装载：压缩缓存约 128K × 528 B ≈ 69 MB（约 66 MiB），SWA 约 4K × 528 B ≈ 2.1 MB，乘以层数就是 GB 级搬运；跨层共享之后，压缩缓存只搬一次。
• 解码每步只需要换算 top-k 索引（几千个 int32），KV 数据由 FlashMLA 按索引直接读，搬运量是注意力的实际读取量，没有额外的整段拷贝。

这就是为什么 DeepInfer 宁可多写一个索引换算内核，也不在解码路径复用 load_paged_kv_cache。分页缓存的全部意义就是“让内核按需读页”，如果每步都先拼成连续缓冲区，分页就退化成了“分配时麻烦、运行时更麻烦”。