DeepInfer 源码拆解（三十二）：norm、RoPE、量化与写 KV 合成一个 kernel——fused_transform 的五段流水与槽位分发

这份 KV 到底长什么样：MLA 的 K/V 是同一份潜在向量

要理解这个 kernel 为什么能同时服务「主 KV」与「index KV」，得先回到 DeepSeek 的 MLA（Multi-head Latent Attention）。DeepSeek-V2 于 2024 年提出的 MLA 把 K 与 V 压缩成一份共享的低秩潜在向量，缓存里每个 token 只存这一份压缩后的 KV，而不是像标准多头注意力那样为 K、V 各存一份完整的 head 维度 —— 部署口径下 KV 缓存下降超过 90%。

（图源：DeepSeek-V2 论文，arXiv:2405.04434）

这张图里的两处关键结构，直接决定了本 kernel 的形态：

• 被缓存的那份「潜在向量」，每个 token 只有一行。这意味着写缓存的路径天然是「一个 token 一行、一行一个 head」，上一节的 num_heads == 1 断言正是这个结构在接口上的投影；而同一个 kernel 服务 Q 预处理时，输入变成「一个 token 多个 head」，它就不写缓存、只返回 bf16。
• 还原 K 与 V 的上投影发生在读取侧（注意力 kernel 内），而不是写入侧。写入侧只需要把「潜在向量」归一化、旋转、量化、落下——这恰好就是本 kernel 四段的来源。K/V 上投影带来的额外算术如果放进写入侧，会变成每层每 token 多几次矩阵乘，没有任何收益，因为读取侧无论如何都要做。

表里的形态切换

这张表里藏着几次形态切换，每一种都对应签名里的一项：

• 输出目标从「返回张量」变成「写缓存」：out_args 是 KVCacheOutArgs 时，结果直接散写进分页缓存，函数返回 None（或可选的 bf16 中间量）。这个分支要求 num_heads == 1——因为缓存里每个 token 一行，而多 head 的输入是「一个 token 多行」，必须多对一才对得上。
• 精度从 bf16 变成「量化值 + scale」：cast_args 且不写缓存时，返回一个二元组，值可能是 FP8 也可能是打包的 FP4。
• 需要中间态：KVCacheOutArgs.out 不为空时，函数除了写缓存还能顺带返回某个阶段的 bf16。out_stage 取 "norm"、"transform" 或 "cast_back"，分别对应归一化后、旋转/变换后、量化再反量化回来三种状态。这个设计给了一条「带检查点的融合」路径：需要中间结果的调用点（后文会看到压缩 KV 为什么要 norm 之后的 latent）不用为此另跑一遍。

五个参数对象分别约束什么

每个参数对象都不是随便包的，它们在构造时就把「哪些组合合法」写死在断言里。按 kernel 的校验顺序读：

1
if do_rope:
2
    assert rotary_dim > 0 and rotary_dim <= head_dim and rotary_dim % 2 == 0
3
    assert rope_offset + rotary_dim == head_dim
4
if do_hadamard:
5
    assert head_dim == 128

RoPE 只作用于 head 的尾部连续一段：offset 是起点，默认取 head_dim - rotary_dim，断言保证「起点 + 长度 = 结尾」。对 DeepSeek 的 MLA 来说，每个 head 由 qk_nope_head_dim 个不旋转的维和 qk_rope_head_dim 个旋转的维拼成，所以 offset 恰好落在 qk_head_dim - qk_rope_head_dim，见上一张表里所有 attn 调用点。Hadamard 则被限制在 head_dim == 128，原因在第三段再讲——蝶形运算要求每一行恰好是一条完整的 warp 级归约规模。

量化这一段有三组互斥规则：

1
if do_cast:
2
    assert sf_format in ("ue8m0", "e4m3", "fp32")
3
    assert num_per_channels in (16, 32, 64, 128) and head_dim % num_per_channels == 0
4
    assert sf_format != "ue8m0" or round_sf
5
    assert sf_format != "e4m3" or not round_sf
6
    if cast_fmt == "e2m1":
7
        assert (num_per_channels, round_sf, sf_format) in ((16, False, "e4m3"), (32, True, "ue8m0"))

• FP8（e4m3）：分组可以是 16/32/64/128；scale 格式可以是 ue8m0（2 的幂）或 fp32（任意实数）。
• FP4（e2m1）：只有两种组合被支持——每 16 个元素一组、scale 存成 e4m3（即 NVFP4 系的「每 16 元素一个 FP8 scale」），或者每 32 个元素一组、scale 是 UE8M0（MXFP4 系）。这两套正好是主流 FP4 存储格式的两族。
• round_sf 与 sf_format 的绑定：UE8M0 只能配「2 的幂取整」的 scale（它本身只能表达 2 的幂），e4m3 只能配「不取整」的连续 scale。这条断言把一个很容易混淆的点固化下来：scale 的数据类型决定了它能不能承载非 2 的幂的数值。

写缓存分支的约束更紧：

1
assert num_heads == 1 and do_cast
2
if not cache_layout.do_cast:
3
    assert head_dim == 128
4
    assert sf_format == cache_layout.sf_format
5
else:
6
    assert (head_dim, num_per_channels, round_sf, sf_format) == (
7
        512, cache_layout.elements_per_sf, cache_layout.sf_format != "e4m3", cache_layout.sf_format,
8
    )

「写缓存必须量化」是设计决定而不是限制——主缓存与 index 缓存本来就是量化存储的，bf16 缓存格式 do_cast=False 的存在只是为了兼容读取侧，写入路径不接受它。而主缓存只有 512 维、index 缓存只有 128 维这两条，加上上面 Hadamard 只在 128 维可用，把「哪条路径能用什么旋钮」限定得很死。

这段在算什么

RMSNorm 对每一行（一个 token 的一个 head）计算：

其中 $d$ 是 head_dim（512 或 128），$w_i$ 是可学习的逐元素权重，$\epsilon$ 是防除零的常数。计算上要求先做一次全行归约（平方和），再用归约结果缩放每个元素。融合的第一段就要解决：这条归约的通信范围有多大。

按「每线程持有多大一块」反推线程分工

kernel 里没有把线程数写死，而是按 head_dim 推导，注释写明要考虑四个目标：最大化向量宽度、最小化补零、最大化占用率、最大化每线程处理的元素数。核心的几行：

1
threads = 128
2
max_elements_per_thread = 16 if small_batch else 32
3
alignment = max(tilelang.next_power_of_2(ceil_div(head_dim, 16)), threads)
4
blk_d = min(tilelang.next_power_of_2(head_dim), align(head_dim, alignment))
5
max_vec_size = math.gcd(32 // T.bfloat16.bytes, head_dim)

• max_vec_size 是「一个 32 字节的向量事务能装下多少 bf16」，即 16 个元素。这是向量化读取的上限。
• blk_d 是每个 program 一次处理的行宽（向上取到 2 的幂，但不无限放大，因为超过 head_dim 的部分要补零）。对 512 维，ceil_div(512,16)=32 取成 32，alignment=max(32,128)=128，align(512,128)=512，最终 blk_d = min(512, 512)=512；对 128 维，blk_d = 128。换句话说，当前支持的两个维度都恰好让 blk_d == head_dim，尾部补零只在理论上的非 2 幂维度出现。
• chunks_per_row = blk_d // thread_vec_size 决定「一行的 head_dim 个元素要分给多少个线程」：512 维是 32 个 16 元素块，128 维是 8 个块。

再乘上 blk_m（一个 program 同时处理的行数，由 max_elements_per_thread * threads // blk_d 推得），512 维时 blk_m=8：8 行 × 32 块 = 256 个块，正好分给 128 个线程各 2 块，所以每个线程持有 32 个元素（来自两行），而每一行由一个 warp（32 线程）完整覆盖。128 维时 blk_m=32：32 行 × 8 块 = 256 个块，128 个线程各 2 块，每个 warp 覆盖 4 行。两种 shape 的 norm_threads_per_row 分别是 32 和 8，都落在单个 warp 内：

• norm_warps_per_row == 1：当前两个维度的情形，每个线程先算自己那块的平方和，再在 warp 内用 shfl_xor 做蝶形归约。
• norm_warps_per_row > 1：一行跨多个 warp（理论情形），各 warp 的部分和先写进 shared memory，再汇总。

1
for r in T.unroll(norm_warp_round):
2
    partial = partial + T.shfl_xor(partial, 1 << (norm_warp_round - 1 - r))
3
rstd_local[local_norm_row] = partial
4
...
5
rstd_local[local_norm_row] = T.rsqrt(rstd_local[local_norm_row] / head_dim + norm_eps)

norm_warp_round = log2(min(norm_threads_per_row, 32))。归约只在一个 warp 内完成，前提是同一行的元素被分到同一个 warp 的连续 lane 上——这一点由 chunks_per_row <= warp_size 与 (blk_m * chunks_per_row) % threads == 0 两条断言保证。当一行需要超过一个 warp 时，代码会先把各 warp 的部分和写进 shared memory，再由每个 warp 的 lane 0 读出求和：

1
if norm_warps_per_row > 1:
2
    ...
3
    norm_partial_shared[local_row, w] = ...
4
...
5
total = ...
6
for w in T.unroll(norm_warps_per_row):
7
    total = total + norm_partial_local[w]
8
rstd_local[local_norm_row] = total

这里的同步由 program 内部其余同步点承担。代码注释特别强调 Keep bitwise aligned with the fragment version：两条归约路径必须产生逐位相同的结果，否则同一次前向在不同 shape 下会有微小的数值差异，进而破坏 batch-invariant 的承诺。这也是 rsqrt 只做一次、平方和按固定顺序两两相加的原因。

为什么是「先算平方和，再统一缩放」

代码先把整行的平方和算完再对所有元素乘缩放，而不是边读边乘。原因是 scale 依赖整行；如果按块读进来、立刻乘，就只能等归约结果回来再重读一遍——这在 warp 里可以靠寄存器缓住，但实现上要保证「同一行的所有块都还在这批线程的寄存器里」。代码用 x_fp32_local（每线程 elements_per_thread 个 fp32）把整块数据缓存住：

1
for chunk in T.unroll(chunks_per_thread):
2
    for k in T.vectorized(thread_vec_size):
3
        x_fp32_local[idx] = T.cast(x_bf16_local[idx], T.float32)
4
        norm_sq_local[...] = x_fp32_local[idx] * x_fp32_local[idx]

bf16 输入先转 fp32 再平方，是因为 bf16 只有 8 位尾数，直接乘平方会丢精度，而归约结果要参与后面所有元素的缩放。转成 fp32 后，归约用 fp32、缩放用 fp32、最后才 T.cast(..., T.bfloat16) 写回 bf16——这与上一篇里 DropoutAddRMSNorm 的 fp32 累加规范一致。

交换对：NEOX 与交错两种排布

RoPE 对尾部 rotary_dim 个维度按位置 $m$ 旋转。记第 $j$ 对元素为 $(x_0, x_1)$，第 $i$ 个频率的角为 $\theta_i m$，标准公式是：

关键在于「$x_0$ 与 $x_1$ 在下标上相隔多远」有两种约定：

• NEOX 排布：前半段与后半段配对，$x_0$ 在下标 $j$，$x_1$ 在下标 $j + \text{half}$。
• 交错排布（GPT-J 风格，代码里 is_neox=False）：相邻两维配对，$x_0$ 在下标 $2j$，$x_1$ 在下标 $2j+1$。

1
if is_neox:
2
    x0_idx = j
3
    x1_idx = half_rotary_dim + j
4
else:
5
    x0_idx = j * 2
6
    x1_idx = j * 2 + 1

两种排布下，一对元素都可能被不同的线程持有。NEOX 尤其明显：$j$ 与 $j + 256$ 相隔一整段，在「一行分给一个 warp」时落在不同的 lane 上。所以 RoPE 必然要在线程之间搬运数据，这就是为什么融合里必须引入一块 shared memory。

用 shared memory 做「行内转置」

代码的写法是把尾部旋转段从寄存器搬到 shared memory rope_shared，在 shared 里按对的布局读取、计算、写回，再搬回寄存器：

1
寄存器(本线程块) --写--> rope_shared[row, col - offset] --按对读+旋转--> 写回 rope_shared
2
--读--> 寄存器(本线程块)

之所以用 shared memory 而不是直接 shfl，是因为一对元素的 lane 距离不固定（NEOX 是 16 个 lane，交错是 0/1 个 chunk 的偏移），统一用一个可寻址的暂存区最省事。代价是两次额外的往返，但它发生在片上，而不是全局内存里。

shared memory 的排布有一个细节：rope_shared_dim = rotary_dim + (2 * chunks_per_row - rotary_dim) % (2 * warp_size)。这是错开行首的 bank 冲突处理——每行的起点按「chunk 数 × 2」个元素错开，使得相邻行写入时落在不同的 bank 上。注释直接写了理由：Stagger row banks。如果 chunks_per_row 恰好让错开量为 0（例如 64 个块与 warp 宽度成倍数），这项就没效果，但这是代价不高的启发式。

只同步 warp：为什么可以

共享内存在一般 kernel 里需要 __syncthreads() 级别的同步，而这个 kernel 在 RoPE 段只用了 T.sync_warp()：

1
# rope_layout preserves each row's input warp, so warp-level synchronization is sufficient.
2
T.sync_warp()

注释里的 rope_layout 是一个 T.Fragment：

1
rope_layout = T.Fragment(
2
    (blk_m, rope_pairs),
3
    forward_thread_fn=lambda i, j: (i * chunks_per_row + j % chunks_per_row) % threads,
4
)

它规定「行 $i$ 的第 $j$ 对元素由哪个线程处理」，取模运算保证同一行的所有块落在同一个 warp 内，而且与写入 shared 时使用的线程映射一致：写进 shared 的线程、读出计算的线程、再把结果读回的线程是同一批人。于是行内通信的闭包完全在一个 warp 里，sync_warp 就够了。前提断言写得很清楚：

1
assert chunks_per_row <= warp_size and warp_size % chunks_per_row == 0
2
assert (blk_m * chunks_per_row) % threads == 0

循环读 cos/sin 表

RoPE 要在每行读取 cos[position] 与 sin[position]。decode 时每行对应一个位置，表是 [max_position, rotary_dim] 的 fp32。当 num_heads 很大（比如 Q 预处理里有 64 个头），同一个 token 的 64 行共享同一份 cos/sin；代码在 do_loop 模式下先把这次 program 会用到的那几行 cos/sin 预取进 shared：

1
for cache_row, j in T.Parallel(rope_cache_rows, half_rotary_dim):
2
    segment_row = bx * blk_m * head_blocks_per_program + cache_row * rope_rows_per_cache
3
    ...
4
    cos_shared[cache_row, j] = cos[position, j]
5
    sin_shared[cache_row, j] = sin[position, j]
6
T.sync_threads()

这里用 T.sync_threads()（不是 warp 级），因为预取时读 cos_shared 的线程与写入它的线程未必在同一 warp。rope_rows_per_cache = gcd(blk_m, num_heads) 是「一个 head_block 里出现了多少个共享同一位置的连续行」，用它把表行数压到最少。do_loop 只在 num_heads % blk_m == 0 and num_heads > blk_m 时开启——也就是只有「多 head 且能被整块整除」的场景（如 Q 的 64 个 head）才值得预取。

conjugate：逆旋转在同一段里

conjugate=True 时交换的是符号：

1
if conjugate:
2
    rope_shared[i, x0_idx] = T.cast(T.ieee_fmaf(x0, cos_val, x1 * sin_val), T.bfloat16)
3
    rope_shared[i, x1_idx] = T.cast(T.ieee_fmaf(x0, -sin_val, x1 * cos_val), T.bfloat16)
4
else:
5
    rope_shared[i, x0_idx] = T.cast(T.ieee_fmaf(x0, cos_val, -(x1 * sin_val)), T.bfloat16)
6
    rope_shared[i, x1_idx] = T.cast(T.ieee_fmaf(x0, sin_val, x1 * cos_val), T.bfloat16)

两处都用 T.ieee_fmaf 而不是先乘后加。IEEE 精度的 FMA 不给编译器把乘加拆开做重关联的自由，从而保证同一份代码在 CUDA 与 Ascend 两侧、以及不同编译选项下逐位一致。这与上一段「两条归约路径必须逐位对齐」是同一个主题：融合把多个数值步骤压进一个 kernel，任何一处引入非确定性的浮点差异，都会让「同一输入、同一 batch 组成 ⇒ 同一输出」的承诺失效。

conjugate 在 MLA 的 context 逆旋转里有用：注意力输出的尾部 64 维带着与 K 相同的位置旋转，为了让 O 投影看到原始基下的 V，需要先反旋一次。上一篇 Linear、RMSNorm 与 RoPE 已经把这层动机讲透，此处不再展开。

分组与并行归约

量化按 num_per_channels（16/32/64/128）一组做 per-group 对称量化：

其中正则是 FP8 e4m3 的最大有限值（448）或 FP4 e2m1 的最大值（6）。AMAX_MARGIN 是一个极小的下界，防止整组为零时除零。

代码先做 amax 归约。一个组的元素可能跨多个线程，所以归约分两级：线程内先求自己的 amax，再在同一位置的多个线程之间用 shfl_xor 做 warp 归约；若一组跨多个 warp，则部分结果落到 shared：

1
partial_float = T.cast(T.max(pair_max_local[0], pair_max_local[1]), T.float32)
2
for r in T.unroll(int(math.log2(quant_warp_width))):
3
    partial_float = T.max(partial_float, T.shfl_xor(partial_float, 1 << r))
4
scale_f32_local[local_group] = partial_float

注意 pair_max_local 是 2 个 bf16 槽位：一个线程一次处理 16 个元素，把它拆成奇偶两个流各自 max，最后合并。这样向量化读取后可以直接在 bf16 域内比较（组内最大值的精度要求低于归约），再把结果转 fp32 参与 warp 归约。

UE8M0：只买 2 的幂的 scale

UE8M0 是一个 8 位指数、0 位尾数的格式——它的值只能是 $2^{e-127}$。用这种 scale 量化，等于把每个组的缩放限制在「最近的 2 的幂」。实现里最关键的一行：

1
exp_scale = -ue8m0_exp_scale(amax * (1.0 / finfo_max))

其中

1
def ue8m0_exp_scale(scale):
2
    scale_bits = T.reinterpret("uint32", scale)
3
    return -T.reinterpret("int32", ((scale_bits - 1) >> 23) + 1 - 127)

这行位运算做的事：把 fp32 的 $s$ 的指数部分取出来，并且向上取整到 2 的幂（减 1 再右移是「floor(log2 s) 除非 s 恰是 2 的幂，否则取 ceil」）。返回的 exp_scale 是「量化乘子 $2^{-\lceil\log_2 s\rceil}$ 的指数」，所以 dequant 用的 scale 是 $2^{\lceil\log_2 s\rceil}$，恰好是「不小于 amax/max 的最小 2 的幂」——这样量化不会溢出，代价是最多 2 倍的精度损失。UE8M0 的存储字节就是 exp_scale + 127，一个字节一个组。

对量化乘子，代码没有用 exp2 计算，而是直接造了一个 bf16：

1
scale_bf16_local[local_group] = T.reinterpret(T.cast((-exp_scale + 127) << 7, T.uint16), T.bfloat16)

bf16 由 1 符号 + 8 指数 + 7 尾数组成，(e+127) << 7 恰好把指数摆到正确的位置、尾数置零，得到的正是 $2^{e-127}$。这一步把「求一个 2 的幂」变成一条移位指令，避免了 exp2 的延迟与精度问题。

三个分支的对比

代码里按 sf_format 分三条路：

ue8m0 分支用 bf16 相乘是有意为之：乘子本身是 2 的幂，bf16 能精确表示；而「造一个 bf16 的 2 的幂」只需要一条位运算（上一段的 (e+127) << 7），乘法本身又能在 bf16 域直接做，整条量化的关键路径上没有 exp2、没有 fp32 转换。e4m3 分支的乘子不是 2 的幂，必须在 fp32 里算，否则丢精度后 amax 附近的元素会被量化到错误的档位。fp32 分支最慢但最准确，只在「scale 段预算宽松、精度优先」的旧格式上出现。

为什么选块状缩放而不是张量级缩放

DeepSeek-V3 于 2024 年底提出的 FP8 训练框架把 fp8 GEMM 写成「细粒度块状缩放 + 高精度累加」：权重按较小的块、激活按更大的块各自求 scale，部分和在 CUDA Core 里用 FP32 累加并周期性提升回高精度寄存器。训练侧需要这套口径，是因为 FP8 的尾数太短，长 K 维上整体一个 scale 会让误差在累加末端放大。

推理侧的 KV 量化把同一个思路换了一个切分方向：不是沿 GEMM 的输出块切部分和，而是沿每个 token 的 head 维度按 16/32/128 个元素一组求 scale。动机是同一个——局部统计出来的 scale 比全局 scale 更贴合实际数值分布，量化点的相对误差更小；代价是每组要多存一个 scale 字节。一个 512 维的 FP8 主缓存记录因此是 512 字节数据 + 20 字节 scale（每 32 元素一个 UE8M0 字节），bytes_per_token = 532，正是前面 stride_alignment 里那个 gcd(32, 532) = 4 的来源。

分组粒度也不是只有一种：V41 的 FP8 主缓存是每 32 元素一个 scale，FP4 主缓存是每 16 元素一个 FP8 scale（NVFP4 口径），旧版 V4 index 缓存则是每 128 元素一个 fp32 scale。粒度越细，scale 开销越大、但分布拟合越好——这三档正好覆盖了「省字节」到「保精度」的取舍。

压缩与写槽位：一次地址算术

量化完，最后一个动作是写。return_quant 时写到返回张量；store_cache 时按槽位写进分页缓存。写槽位前先把 slot_mapping 洗一遍：

1
if thread_in_cache_group == 0:
2
    slot = T.if_then_else(row < rows, slot_mapping[row], -1)
3
    slot = T.Select(slot < num_blocks * block_size, slot, -1)
4
slot = T.shfl_sync(slot, 0, cache_slot_broadcast_width)

三步：

1. row < rows 之外的行（补零产生的 padding 行）不写，槽位置 -1。
2. 越界值（slot >= num_blocks * block_size，例如 -1 被当成无符号大数传进来，或调度器给了非法槽位）也置 -1。
3. 由组内 0 号线程读一次，然后广播给同组所有线程——同组线程写同一行，共用一个槽位，广播省掉了每个线程一次全局读。

之后 slot_block = slot // block_size、slot_offset = slot % block_size 完成从「逻辑槽位」到「页号 + 页内偏移」的翻译。

同一个 4D 张量，两种含义

分页 KV 缓存的公开形状是 [num_blocks, block_size, 1, bytes_per_token]（uint8）。同样是这个形状，主缓存与 index 缓存的字节排布不一样：

• 主缓存（token-interleaved）：每个 token 的一条记录是 [数据段 | scale 段]，token 之间没有填充，records = cache[:, :, 0, :]，然后 data = records[..., :data_bytes]、sf = records[..., data_bytes:]。
• index 缓存（block-wide）：一个 block 里先把所有 token 的数据段连续放、再把所有 token 的 scale 段连续放，block_data = cache.view(num_blocks, block_size * bytes_per_token)。

1
if layout.token_interleaved:
2
    # V41 main/SWA: each record is [data | scale], with no token padding.
3
    records = kv_cache[:, :, 0, :]
4
    data = records[..., :data_bytes]
5
    sf_data = records[..., data_bytes:]
6
else:
7
    # Index caches keep block-wide data/scale sections.
8
    block_data = kv_cache.view(num_blocks, block_size * layout.bytes_per_token)
9
    data = block_data[:, : block_size * data_bytes].view(num_blocks, block_size, data_bytes)
10
    sf_data = block_data[:, block_size * data_bytes :].view(num_blocks, block_size, layout.sf_bytes)

为什么差异这么大？主缓存的读取 kernel（FlashMLA 系的 decode）是逐 token 取 K 的，让一个 token 的数据与它的 scale 在内存里相邻，gather 时一次向量化读就能拿到全部所需；index 缓存的读取是以大块 topk 候选为单位扫描的，按段组织能让同一段内的访问完全连续、便于大块搬运。两种布局服务两种访问模式，kernel 里用一个布尔量 token_interleaved 区分。

主缓存的偏移量怎么算

主缓存里每个 slot 写入时，数据段的地址是：

1
cache_quant_stride = (
2
    stride_cache_sf_block * sf_element_scale * value_element_scale,
3
    stride_cache_quant,
4
    1,
5
)

块步长从 scale 张量 的 block stride 推出来，经过两次单位换算：

1
数据段块步长（以值的元素数计）
2
  = scale 张量的块步长（以 scale 元素计）
3
  × sf_element_scale（一个 scale 元素占几字节）
4
  × value_element_scale（一个字节装几个值）

• sf_element_scale：UE8M0 与 e4m3 scale 都是 1 字节，fp32 是 4 字节。
• value_element_scale：FP8 一个字节一个值，乘 1；FP4 一个字节装两个值，乘 2。
• 对 token-interleaved 的主缓存，scale 张量的块步长本身就以「每 token 记录（数据 + scale）」为单位计算，换算出字节数后再换算成值的元素数，刚好等于量化张量在该布局下的块步长（FP8 是 532 字节，FP4 的主缓存是 288 字节）。
• 对 index 缓存，数据段与 scale 段在 block 内分别连续，stride_cache_quant 是数据段的单行步长，stride_cache_sf_block 主要用于块级对齐的校验与基址换算。

注释把为什么从 scale 张量推解释了：

1
Bind the physical block stride through cache_sf, whose element size is
2
layout-dependent. Derive the value stride from that directly bound variable;
3
expressions such as ``stride // 2`` cannot themselves be recovered as packed-API
4
arguments by TileLang.

TileLang 的参数绑定机制要求 stride 以「实际传入张量的 stride」形式出现；如果直接给 stride // 2 这样的表达式，它无法回溯到某个实参的 stride，jit 时就无法把它当作 packed API 的参数处理。于是代码绕一步：从 scale 张量拿到块 stride，再乘以元素宽度得到数据段 stride。这是融合算子与普通 kernel 的差异——普通 kernel 里 stride/4 随便写，融合后要为多后端 + jit 参数化的约束让路。

对齐：保证「数据段」是连续的一块

stride_alignment 的计算也值得看：

1
return math.gcd(32, self.bytes_per_token) if self.token_interleaved else 32

主缓存每个 token 的记录长度是 532 字节（512 字节数据 + 20 字节 scale），最大公约数是 4，所以块 stride 至少要 4 字节对齐；index 缓存按段组织，要求 32 字节对齐。上层还有一组更严的断言：

1
assert cache_quant.stride(1) % cache_layout.stride_alignment == 0
2
assert cache_quant.stride(0) % 32 == 0

kernel 内部再用 T.assume 把这两个对齐条件告诉编译器，从而选用向量化事务。前面第一篇里讲的「Linear 输出按 32 字节对齐」是同一个主题：写出去的行如果不对齐，向量化写就被拆成若干个标量写，融合的收益就没了。

Q 预处理：只旋不写

attn.py 里 Q 的预处理调用只传了 norm_args（可选）与 rope_args，不传 cast_args、out_args：

1
result = fused_norm_rope_hadamard_cast_store_kvcache(
2
    q,
3
    norm_args=(NormArgs(self.q_b_norm.weight, ...) if self.use_q_b_norm else None),
4
    rope_args=RopeArgs(
5
        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
)

这里 num_heads 是 64（MLA 的 head 数），head_dim 是 512。前面按四个目标推出来的 blk_m 对大 batch 是 8（32 × 128 / 512），64 能被它整除且远大于它，再加上不在 small_batch 分支，do_loop 条件成立：一个 program 连续处理 head_blocks_per_program = 64 / 8 = 8 个头块，cos/sin 表按位置预取进 shared。小 batch 时 max_elements_per_thread 降到 16，blk_m 变成 4，同时 do_loop 被 not small_batch 关掉——省下的 shared 预取开销比多出来的重复读更贵，这是融合 kernel 里典型的「按数据量切策略」。结果以 bf16 返回，直接送去注意力。这一路没有量化、没有写缓存，但它和写 KV 用的是同一个 kernel——这是接口统一的直接收益：Q 与 K 的旋转逻辑只写一份，测试矩阵一份覆盖两种 shape。

SWA 写入：写缓存 + 可选返回中间量

1
if use_cp:
2
    cast_args = CastArgs("e4m3", self.swa_cache_num_per_channels, round_sf=True, sf_format="ue8m0", cast_back=True)
3
    out_args = out   # 一个 bf16 张量，接收量化再反量化的结果
4
else:
5
    cast_args = CastArgs("e4m3", ..., round_sf=True, sf_format="ue8m0")
6
    out_args = KVCacheOutArgs(swa_cache, input_metadata.slot_mapping_store_swa)

CP 开启时，K 要在各 rank 之间 all-gather（上一篇 上下文并行 CP 全链路 涉及过这一步）。为什么要先量化再反量化回 bf16？因为 all-gather 走的是通信缓冲，缓冲里的数据精度是写定的 bf16；如果不先在本地量化再反量回到与远端相同精度的 bf16，拼接后的结果与「先 all-gather 再整体量化」不一致——同一份数据，量化粒度可能不同，最终写进缓存的字节就不一样，这又破坏了「不同并行度下输出一致」。所以这里用 cast_back=True 让 kernel 在本地做一次「量化 → 反量化」的往返，输出 bf16 给通信层，代价是这份数据经过了有损压缩。非 CP 时则直接写缓存，不需要往返。

压缩 KV：先写缓存，再把 norm 后的 latent 拿回来

compress.py 的调用最特别：

1
fused_norm_rope_hadamard_cast_store_kvcache(
2
    kv_compressed,
3
    norm_args=NormArgs(self.norm.weight, ...),
4
    rope_args=...,
5
    cast_args=CastArgs.from_cache_format(cache_format),
6
    out_args=KVCacheOutArgs(
7
        kv_cache,
8
        input_metadata.slot_mapping_store_ca[self.ratio_key],
9
        out=pre_rope_kv,
10
        out_stage="norm",
11
        layout=cache_format,
12
    ),
13
)

out_stage="norm" 意味着：在写进缓存的同时，把 norm 之后、RoPE 之前的那个 bf16 张量拷贝到 pre_rope_kv。为什么要 norm 后的 latent？indexer 的 key 投影 KProject 的输入是「压缩后的 latent 在 norm 之后、RoPE 之前」——它自己还要做一次线性投影、再做 norm 和 RoPE。如果融合只写缓存不返回中间量，这条路径就得再重跑一遍 norm，白白多花一次全量读。

这条设计让 kernel 在流水中间开了一个「出口」：在计算路径上，norm 之后立刻 copy_cache_out("norm", work_3d)，然后继续 RoPE、量化、写缓存。参考实现里这段代码很短：

1
copy_cache_out("norm", work_3d)
2
if rope_args is not None:
3
    apply_rotary(work_3d[..., offset : offset + rotary_dim], ...)
4
...
5
copy_cache_out("transform", work_3d)

out_stage 可以取 "transform"（旋转/Hadamard 之后）或 "cast_back"（量化再反量化后），对应三种中间态。这三个出口共享同一条流水，区别只是在哪一步多拷一份出来。

indexer Q：返回量化值本身

index_attn.py::_quantize_q 传 rope_args 与 cast_args，不传 out_args：

1
result = fused_norm_rope_hadamard_cast_store_kvcache(
2
    q_b,
3
    rope_args=RopeArgs(...),
4
    cast_args=CastArgs.from_cache_format(self.cache_format),
5
)
6
assert isinstance(result, tuple)

result 是 (quantized, sf)：Q 被量化成与 K 相同的 index 格式（FP4 或 MXFP8），然后交给稀疏注意力的打分 kernel。量化粒度与 scale 格式必须与 K 一致——两边都用 CastArgs.from_cache_format 从同一个 KVCacheFormat 派生，保证两边不会漂移。Hadamard 不在 KVCacheFormat 里，若某一侧启用旋转，另一侧必须由调用点显式传相同的 HadamardArgs，否则 Q 与 K 落在不同基上，点积失去意义。这也是把「格式」做成 KVCacheFormat 对象而不是散落的上百个参数的意义：它既是缓存文件的元数据，又是量化算子的配置，一处定义、两处引用。

收益：一次 launch、一次往返、一次调度

从调用方看，融合带来的最直接变化是每次写 KV 只发一个 kernel。在 decode 这种每步几十到上千个 token 的场景里，GPU 的空闲常常发生在 kernel 之间的 gap 里；四个小 kernel 的串行会把 gap 累积成可观的启动开销。融合后，Q/K 的旋转、量化、写页全部在一个 grid 里完成，前一 kernel 的输出就是后一 kernel 的输入，没有「等待上一 kernel 全部完成再启动下一个」的 barrier。

第二个收益是中间张量不落地。不融时 bf16 中间结果会写回显存再读回，总量大约几 MB 到几十 MB（取决于 batch 与维度）。这个量级在 L2 里可以命中（kernel 里有一个 small_batch 判定：x.numel() * x.element_size() <= _L2_CACHE_SIZE，即 126 MB），但命中与否依赖调度；融合把它彻底变成寄存器与 shared 里的数据，与 L2 无关。

第三个收益是数值路径统一。前文的 batch-invariant、ieee_fmaf、Keep bitwise aligned 都是为这个服务的：如果归一化在 attn 里做、RoPE 在独立的 rope 算子里做、量化在量化库里做，三处的数值约定（fp32 还是 bf16 中间值、FMA 还是分两步、归约顺序）可能各不一样，最终「同一 batch 组成输出一致」就很难保证。把它们收进同一个 kernel，数值约定只需要在一处维护。

代价：可调试性、组合爆炸与形状约束

融合的代价同样真实：

• 可调试性。一个 600 行的 TileLang kernel 出错时，错误位置只有「某个 lane 的某个寄存器」这类线索。代码的对策是把参考实现 fused_..._ref 写成「同样的流水但不融合」的朴素 PyTorch 版本，逐段调用既有算子（norm、apply_rotary、hadamard_transform_ref、quant.per_token_cast、store_kv_cache_ref），kernel 只和它对照。这是一个非常值得借鉴的做法：融合算子的正确性标准是「与不融合的流水逐字节一致」，而不是「和某个数学公式一致」。
• 组合爆炸。五段流水 × 三种输出 × 五种缓存格式 × 若干维度，合法组合不多但组合方式很多。代码的对策是把组合收敛到参数对象（NormArgs/RopeArgs/HadamardArgs/CastArgs/KVCacheOutArgs）和一组断言里，而不是为每种组合写一个分支。
• 形状约束。head_dim 暂时锁在 128/512；Hadamard 只在 128；写主缓存必须 512 维；RoPE 只作用于尾部连续段；输入 stride 必须 16 元素对齐。这些约束在模型升级时会成为阻碍，但它们的代价由断言「早失败」承担，而不是运行中出现错误结果。

融合边界的画法

代码把「哪些事可以融」的边界画得很明确：

• 能融的：都是逐行独立、行内依赖有限（有界的 warp 内通信）的变换。归一化、RoPE、Hadamard、量化都属于此类。
• 不能融的：需要看全序列的操作（如 attention 本身、topk 选择）、跨 token 归约（如 RoPE 之外的某种统计）、以及需要全局同步的步骤。store_kv_cache 单独作为一个 kernel 存在（store_kv_cache.py），承接那些「数据已经做过 norm/rope、只剩搬运」的写入：prefill 的 SWA 分支就是典型——它先把历史 SWA 段 load_paged_kv_cache 进连续缓冲、与新算出的 KV 拼好，再在缓存读取之后用 store_kv_cache 落页；此时若再走融合路径，就会把同一批 KV 重复归一化和旋转一遍。

把边界画在这里有两个好处：一是 kernel 规模可控，每个 program 只处理固定行数、不跨 token；二是 CUDA Graph 捕获容易，固定形状下 grid 和 block 尺寸都是编译期常量。