vLLM 源码完全拆解（十九）：TurboQuant KV Cache 量化——Hadamard 旋转、Lloyd-Max 标量量化与分页缓存内核

一条 kv-cache-dtype，两条默认的量化路线

vLLM 的缓存在服务启动时按 cache_config.cache_dtype 决定数据类型。FP8、INT8 这类方案走的是”每个元素独立存储、读取时再解量化”的简单模型；TurboQuant 走的却是”每个 KV 头在每个位置只占一段定长的 packed 字节”。这意味着它不能复用任何现有注意力后端的 KV 读写逻辑，必须同时改三处：

1. KV 缓存规格：AttentionSpec 里的 state_content_bytes 原本描述”两个向量（K 和 V）的存储字节”，TurboQuant 把它改成”一个 slot 的字节数”。
2. 注意力后端：新增 TurboQuantAttentionBackend，注册在 backend 注册表里，名字是 TURBOQUANT。
3. 写入路径：TurboQuant 注意力层把 K/V 的写缓存做成独立的 kernel，而不是依赖后端自带的 unified_kv_cache_update。

缓存规格被改写成”一个 slot”

TurboQuantAttentionBackend.customize_spec 在注意到 kv_quant_mode.is_turboquant 时，用预设参数算出 slot 大小写入 state_content_bytes：

1
@classmethod
2
def customize_spec(cls, spec: AttentionSpec) -> AttentionSpec:
3
    """TurboQuant packs K+V into one slot per head."""
4
    if spec.state_content_bytes is not None or not spec.kv_quant_mode.is_turboquant:
5
        return spec
6
    from vllm.model_executor.layers.quantization.turboquant.config import (
7
        TurboQuantConfig,
8
    )
9

10
    # KVQuantMode member names mirror the preset strings.
11
    tq = TurboQuantConfig.from_cache_dtype(
12
        spec.kv_quant_mode.name.lower(), spec.head_size
13
    )
14
    return replace(spec, state_content_bytes=tq.slot_size_aligned)

这段逻辑改写了 vLLM 里”KV Cache 的每一页由若干个 head×head_dim 的 K 和 V 组成”的默认假设。原本一页的字节数等于 block_size × num_kv_heads × head_size × 2 × element_size，而 TurboQuant 之后，K 和 V 不再各自成块，而是合并成一条”slot”。之后 KVCacheManager、页表、block table 等组件仍然按”页 → 每页若干 slot”的方式管理，只是 slot 的尺寸变了。PagedAttention 的分页结构和页粒度这篇站内文章可以对照着看：KV Cache 的页粒度：128 token 一块的账，以及五个推理框架为什么选了五个不同答案。

后端能力声明

TurboQuantAttentionBackend 声明的关键能力：

1
accept_output_buffer: bool = True
2
forward_includes_kv_cache_update: bool = False
3

4
supported_dtypes: ClassVar[list[torch.dtype]] = [
5
    torch.float16,
6
    torch.bfloat16,
7
]
8
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
9
    "turboquant_k8v4",
10
    "turboquant_4bit_nc",
11
    "turboquant_k3v4_nc",
12
    "turboquant_3bit_nc",
13
]
14

15
@classmethod
16
def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:
17
    return (KVCacheLayout.LBNHC,)
18

19
@staticmethod
20
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
21
    return [16, 32, 64, 128]
22

23
@classmethod
24
def supports_attn_type(cls, attn_type: str) -> bool:
25
    return attn_type == AttentionType.DECODER
26

27
@classmethod
28
def supports_per_head_quant_scales(cls) -> bool:
29
    return False

逐条解释这些声明的含义：

• forward_includes_kv_cache_update = False：与 FlashAttention 后端一致，KV 的写入被拆成独立一步（do_kv_cache_update），在 attention 前向之前执行。这样做的好处是当模型中某几层被 get_boundary_skip_layers 跳过时，写入端可以逐层决定是否量化，而前向端依旧走同一个后端。
• supported_kv_cache_layouts = (LBNHC,)：缓存张量的逻辑布局是 (num_blocks, block_size, num_kv_heads, slot_bytes)，不再有”最前面那个大小为 2 的 K/V 维”。代码里还额外做了一次 kv_cache.transpose(1, 2) 转成 (B, H, N, C)，让 Triton kernel 读起来像普通的 tensor layout。
• supported_kernel_block_sizes = [16, 32, 64, 128]：页大小的选择有限，因为写入 kernel 的 BLOCK_SIZE 是编译期常量，而 3-bit packing 的字节跨页有对齐问题，后文会说明。
• supports_per_head_quant_scales = False：TurboQuant 不使用 per-head 的 scale，因为每个向量自带自己的 norm 或 scale，这是它的设计核心。
• supports_attn_type = DECODER：只支持 decoder 型注意力，encoder/decoder 混合的模型不可用。

四个预设对应什么

真正控制量化方式的，是 TurboQuantConfig 里的 TQ_PRESETS：

1
# Named TQ presets: each maps to frozen config parameters.
2
# key_quant_bits: 8 = FP8 keys, 3-4 = MSE (Lloyd-Max) quantized keys.
3
# value_quant_bits: 3-4 = uniform quantized values.
4
TQ_PRESETS: dict[str, dict] = {
5
    "turboquant_k8v4": {
6
        "key_quant_bits": 8,
7
        "value_quant_bits": 4,
8
        "norm_correction": False,
9
    },
10
    "turboquant_4bit_nc": {
11
        "key_quant_bits": 4,
12
        "value_quant_bits": 4,
13
        "norm_correction": True,
14
    },
15
    "turboquant_k3v4_nc": {
16
        "key_quant_bits": 3,
17
        "value_quant_bits": 4,
18
        "norm_correction": True,
19
    },
20
    "turboquant_3bit_nc": {
21
        "key_quant_bits": 3,
22
        "value_quant_bits": 3,
23
        "norm_correction": True,
24
    },
25
}

TurboQuantConfig 的 docstring 直接给出了各预设的压缩倍数和困惑度（PPL）代价：

这张表里的四行就已经说明了整套工程取舍：K 的位宽对精度影响远大于 V；把 K 和 V 同时压到 3 bit 时 PPL 已经上涨 20%，基本只能在”显存极度紧张、又只关心能不能跑”的场景下使用。为什么 K 的量化如此敏感？因为 K 参与 softmax，误差是”在指数前”发生的；V 是”加权求和”，误差被 softmax 的凸组合稀释。稍后会给出更明确的推理。

为什么不用 QJL

TurboQuantConfig 的 docstring 里有一段值得单独拿出来讲的取舍：

QJL is intentionally omitted: community consensus (5+ independent groups) found it hurts attention quality by amplifying variance through softmax.

原论文 TurboQuant 的完整算法里还有一个 QJL（Quantized Johnson-Lindenstrauss）分量，用来在残差上进一步压缩，理论上可以做到接近无损的 1-bit 残差。但 vLLM 的实现完全没有采用它，理由是多个独立的复现团队发现 QJL 会在 softmax 之前放大方差，进而损害注意力质量。这条取舍背后有清晰的直觉：QJL 存储的是被量化之后残差的方向，如果在下游 softmax 之前没有合适的归一化，残差噪声会被直接打进 logits，再被 softmax 的指数放大。TurboQuant 在 vLLM 中实际采用的三件套是”Hadamard 旋转 + Lloyd-Max 标量量化 + norm correction”，把 K 的噪声主要限制在一个可估计的范围内。

问题的本质：非均匀维度

假设每层每个 KV 头有一个 K 向量 $k \in \mathbb{R}^d$，长度 128。如果直接对每个坐标用同一个量化区间做标量量化（比如最简单的”线性均匀量化”），效果会很差，原因有三：

1. 坐标分布不均匀：经过 RMSNorm、RoPE 之后，K 的坐标通常有一个很大的集中区（接近 0）和一条长尾，均匀量化会把大量码字浪费在尾部。
2. 坐标尺度差异：不同头、不同层的 K 数值范围不同，如果共享码本，必然有的头饱和、有的头分辨率浪费。
3. 维度间的相关性：直接量化每个坐标相当于把问题看成”独立维度”，但坐标之间有相关性，量化误差会集中到某些方向。

经典的解决办法是”每通道/每组 scale + zero”（比如 FP8 per-head scale、INT8 per-token per-head）。但那样会引入额外元数据：每个头要存一个 scale，每个组要存一个 scale。位宽越低，元数据占比越高——3 bit 时一个 fp16 的 scale 就占掉 5 个元素的存储量。

旋转的关键定理

TurboQuant 的思路是：先把这个向量乘上一个随机正交矩阵，让每个坐标近似服从同一个高斯分布。理由来自高维几何里的一个结论：

设 $Q \in \mathbb{R}^{d \times d}$ 是随机正交矩阵（$Q^\top Q = I$），$x$ 是任意固定的单位向量，则当 $d$ 较大时，$Q x$ 的每个坐标近似服从 $\mathcal{N}(0, 1/d)$。更准确地说，$Q x$ 在单位球面上”像随机方向一样分布”，而单位球面上的随机方向在每个坐标轴上投影的方差就是 $1/d$。

这个结论可以用两条步骤直观推导：

• 第一步：固定 $x$，随机选择正交 $Q$，等价于在单位球面上随机取一个方向 $y = Qx$。因为正交变换对所有方向一视同仁。
• 第二步：单位球面上的均匀分布 $S^{d-1}$ 在任意坐标轴上的投影，密度正比于 $(1 - t^2)^{(d-3)/2}$。当 $d \to \infty$ 时，这个分布收敛到 $\mathcal{N}(0, 1/d)$。可以用 Berry-Esseen 不等式或更直接的”球面集中”来论证：$d \geq 64$ 时误差已经很小。

所以”随机旋转 → 高斯坐标”不是巧合，而是几何必然。旋转之后的向量，每个坐标的统计特性都一样了。这时候我们只需要设计一个标量量化器，就能适用于所有坐标、所有头、所有层。这就是为什么低比特标量量化加上旋转会比每个头各自 scale 的量化更划算。

为什么 vLLM 选了 Hadamard 而不是随机正交矩阵

源码里 _build_hadamard 的注释写得非常直白：

1
def _build_hadamard(d: int, device_str: str) -> torch.Tensor:
2
    """Orthonormal Hadamard matrix (Sylvester construction), cached per (d, device).
3

4
    Precomputed D×D matrix enables matmul-based WHT — single cuBLAS GEMM
5
    instead of log2(D) butterfly kernel launches. 64KB for D=128.
6
    """
7
    return _build_hadamard_cached(d, str(torch.device(device_str)))
8

9
@functools.cache
10
def _build_hadamard_cached(d: int, device_str: str) -> torch.Tensor:
11
    H = torch.tensor([[1.0]])
12
    while H.shape[0] < d:
13
        H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0)
14
    return (H / math.sqrt(d)).to(torch.device(device_str))

Hadamard 矩阵具备三个关键性质：

1. 正交：$H^\top H = I$（归一化后乘 $1/\sqrt{d}$ 保证这一点），满足”旋转”要求。
2. 对称：$H = H^\top$，所以逆旋转等于正旋转。
3. 是”足够随机”的：Hadamard 被视为一种确定性随机旋转变体，对高斯化坐标同样有效。更重要的是它的元素只有 $\pm 1/\sqrt{d}$，没有浮点数值本身的误差。

源码 _ensure_on_device 里进一步解释了为什么不加随机符号翻转：

1
# Pure Hadamard: orthonormal + symmetric (H = H^T), enabling
2
# in-kernel butterfly fusion and trivial inverse for continuation.
3
H = _build_hadamard(D, str(device))
4
layer._tq_PiT = H
5
layer._tq_Pi = H
6
# fp16 copy for rotation in continuation prefill path
7
layer._tq_Pi_half = H.to(torch.float16)

1
# The Hadamard rotation is shared across all layers: random sign
2
# flips do not improve Lloyd-Max quantization quality because the
3
# quantizer is symmetric around zero (sign-flipping a coordinate
4
# maps it to the mirror centroid with identical distortion).

这三段注释合起来回答了四个问题：

• 为什么不用随机高斯矩阵？ 因为它需要一次真正的 GEMM，元素是浮点，需要额外存储随机种子。
• 为什么不用”Hadamard + 随机符号”这种原论文的写法？ 因为 Lloyd-Max 量化器是对称的（以 0 为中心），坐标翻转后落到对称码字上，畸变完全一样，所以随机符号对量化质量没有增益。
• 为什么共享 D×D 矩阵而不是逐层生成？ 矩阵本身和层无关，共享节省了 $L \times 64\text{KB}$ 的显存。
• 为什么不是 Fast Walsh-Hadamard 蝶形运算？ 注释写得很实在：预计算 D×D 矩阵，然后一次 matmul 比 $log_2(D)$ 个 kernel 的蝶形运算更省。对 $D=128$ 的矩阵，64KB 可以一直驻留 GPU 的一个内存块里，每次变换只花一次 cuBLAS 调用。

旋转前后的”不变”与”改变”

旋转是一个正交变换，所以：

• 不变的是：向量长度 $\|x\|_2$ 保持不变（正交变换保距）；两个向量的内积 $x^\top y$ 保持不变（正交变换保角）。这正是为什么旋转对”要算 softmax 的 $q^\top k$“没有本质损伤——如果 K 只是旋转，Q 不旋转，则 $q^\top (H k) \neq q^\top k$，所以 Q 也必须同步旋转（下文的写入和读取路径都会说明）。
• 改变的是：每个坐标的分布从”任意、不均匀”变成了”近似高斯、方差 $1/d$“；每个坐标的期望接近 0。这恰好是 Lloyd-Max 量化器所需要的假设。

一个具体的图像

设想一个维度为 128 的单元球面。任取一个单位向量 $x$，它的坐标往往有几个很大的值（比如由 RMSNorm 之后的某个高频分量主导）和大量接近零的分量。乘上 $H/\sqrt{d}$ 之后，结果是一个类似”随机方向”的单元向量，每个坐标的幅值在 $1/\sqrt{d}$ 附近波动，满足近似高斯。

这个变换把”如何构造一个能适配所有头的码本”从一个与数据强相关的问题，变成了一个与数据无关的问题：只需要知道旋转后坐标为 $\mathcal{N}(0, 1/d)$ 这一事实，就能对所有模型、所有层用同一套码本。这是 TurboQuant 的核心。

标量量化的三个对象

一旦坐标被旋转成近似 $\mathcal{N}(0, 1/d)$，问题就退化为一个一维标量量化问题。一个 $b$ bit 的标量量化器由三样东西定义：

• 码本：$2^b$ 个重建电平 $c_0 < c_1 < \dots < c_{2^b-1}$；
• 判决边界：$2^b - 1$ 个分界点 $t_0 < t_1 < \dots < t_{2^b-2}$；
• 编码规则：坐标 $x$ 落入第 $i$ 个区间 $[t_{i-1}, t_i)$ 时，编码为 $i$；
• 解码规则：拿到 $i$ 后重建为 $c_i$。

量化误差（畸变）用均方误差衡量：

其中 $t_{-1}=-\infty$，$t_{2^b-1}=+\infty$。

Lloyd-Max 的两个最优性条件

对于给定的 $b$，我们要找一组 $(c_i, t_i)$ 使 $D$ 最小。对 $D$ 分别关于 $t_i$ 和 $c_i$ 求偏导并令其为零，得到两个条件：

条件一：最近邻划分。 相邻两个重建电平之间的分界点必须取中点：

因为如果分界点偏向任意一侧，把边界移动一点都能让该区间的误差变小。这个条件与输入分布无关——无论 $p(x)$ 是什么，边界都取相邻电平的算术平均。

条件二：质心重建。 每个重建电平必须是该区间内分布的条件期望：

因为 $\mathbb{E}[(x-c)^2]$ 在 $c=\mathbb{E}[x]$ 处最小。

这两个条件互相依赖：质心依赖边界，边界依赖质心。Lloyd-Max 算法采用不动点迭代：

1. 初始化一组质心（比如在 $[-3.5\sigma, 3.5\sigma]$ 上均匀放置）；
2. 用条件一算出边界；
3. 用条件二算出新的质心；
4. 重复直到质心变化小于容差。

对于高斯分布，这个迭代一定收敛到局部最优；由于高斯是单峰的、关于 0 对称，初始值取对称均匀栅格时收敛到全局最优，且解关于 0 对称。

vLLM 的实现：数值积分而不是解析解

centroids.py 里的 solve_lloyd_max 没有用解析公式（高斯的区间期望没有初等闭式表达），而是直接做数值积分：

1
def solve_lloyd_max(
2
    d: int,
3
    bits: int,
4
    max_iter: int = 200,
5
    tol: float = 1e-10,
6
) -> tuple[torch.Tensor, torch.Tensor]:
7
    """Solve Lloyd-Max optimal quantizer for N(0, 1/d) distribution."""
8
    n_levels = 2**bits
9
    sigma2 = 1.0 / d
10
    sigma = math.sqrt(sigma2)
11

12
    def pdf(x):
13
        return _gaussian_pdf(x, sigma2)
14

15
    lo, hi = -3.5 * sigma, 3.5 * sigma
16
    centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)]
17

18
    for _ in range(max_iter):
19
        boundaries = [
20
            (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)
21
        ]
22
        edges = [lo * 3] + boundaries + [hi * 3]
23
        new_centroids = []
24
        for i in range(n_levels):
25
            a, b = edges[i], edges[i + 1]
26
            num = _trapz(lambda x: x * pdf(x), a, b)
27
            den = _trapz(pdf, a, b)
28
            new_centroids.append(num / den if den > 1e-15 else centroids[i])
29

30
        if max(abs(new_centroids[i] - centroids[i]) for i in range(n_levels)) < tol:
31
            break
32
        centroids = new_centroids
33

34
    boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)]
35
    return (
36
        torch.tensor(centroids, dtype=torch.float32),
37
        torch.tensor(boundaries, dtype=torch.float32),
38
    )

几个实现细节值得解释：

• $\sigma^2 = 1/d$：这是旋转后每个坐标的方差。注意它是由维度决定的，而不是从数据里估计出来的。这意味着码本与具体模型无关，所有层、所有头、所有请求共享同一套中心点。
• 积分区间 $[-3.5\sigma, 3.5\sigma]$，外层再扩到 $\pm 10.5\sigma$：高斯在 $3.5\sigma$ 之外的尾部质量约 $4.7\times10^{-4}$，对 KV 向量来说可以忽略；把外层边界写成 lo*3 与 hi*3 只是数值积分的端点，避免最外侧区间积分不完整。
• _trapz 固定 200 个梯形：中心点数量最多 16，每个区间 200 点积分足以把量化器求到 $10^{-10}$ 的稳定性。它是纯 Python 循环，只在模型初始化时跑一次，不进入推理热路径。
• 最多 200 次迭代：大多数情况下十几次就收敛；200 是安全上限。

求解结果用 @lru_cache(maxsize=32) 缓存，键是 (d, bits)。一个服务进程里，无论多少层调用 get_centroids(128, 3)，都只计算一次。

3-bit 和 4-bit 的码本长什么样

把上述结果算出来，会看到典型的高斯量化码本。以 $d=128$、$\sigma=1/\sqrt{128}\approx0.0884$ 为例：

• 4 bit（16 电平）：最外侧两个质心大约在 $\pm 0.242$ 附近，相邻电平间距在中心区域约 $0.023$，在尾部约 $0.059$。越靠近尾部，电平间距越大——因为高斯尾部的概率密度低，把码字放在那里不划算。
• 3 bit（8 电平）：最外侧质心约在 $\pm 0.190$，中心间距约 $0.045$，尾部间距约 $0.071$。

这与均匀量化形成鲜明对比：均匀量化会在 $\pm 3.5\sigma$ 的整个区间等距撒点，中心区域的分辨率被尾部”浪费”掉；Lloyd-Max 把码字集中到概率质量大的地方，因此在同样的 bit 数下 MSE 更小。对于高斯分布，最优固定速率标量量化的高码率近似为

代入 $\sigma^2=1/d$，得到整条单位向量的归一化畸变（NMSE）：

• 4 bit：$\approx 2.72/256 = 1.06\%$
• 3 bit：$\approx 2.72/64 = 4.25\%$

这两个数字与后面精度表的态度一致：4-bit 的 K 只引入约 1% 的方向误差，3-bit 则接近 4%，后者的 PPL 代价明显更高。注意这是”方向”的误差，向量的长度另存为 fp16，不参与这个估算。

熵编码视角：为什么固定速率不是最优但更实用

上面的式子是固定速率标量量化的结果。如果允许熵编码（不同码字用不同比特数），高斯分布的最优标量量化畸变可以降到 $\frac{\pi e}{6}\sigma^2 2^{-2R}\approx1.42\sigma^2 2^{-2R}$，比固定速率好约一倍。但 KV Cache 的读取是随机访问的：每读一个 token 的 K，都要在常量时间内定位到它的编码。变长编码会引入前缀解析、偏移表，并且让分页寻址变得复杂。TurboQuant 选择固定速率 + 定长 slot，牺牲了约 2 倍的率失真增益，换来的是”每个 token 的 K 恰好占固定字节数”这一工程性质。这是整个方案能在 vLLM 里落地的关键前提。

量化误差怎样进入注意力分数

设原始 K 向量为 $k=\|k\|\cdot\hat{k}$，$\hat{k}$ 是单位方向。量化重建为

其中 $\tilde{c}$ 是由 $d$ 个质心分量拼成的向量（可能再做 norm correction）。注意力分数为

误差来自 $\tilde{c}$ 与 $\hat{k}$ 的偏差。设 $\delta = \tilde{c} - \hat{k}$，则

如果 $q$ 的每个坐标量级为 $O(1/\sqrt{d})$（经过训练后通常近似如此），$\|\delta\|_2$ 的量级为 $\sqrt{\text{NMSE}}$，则分数误差的标准差大约为 $\|k\|\cdot\|q\|\cdot\sqrt{\text{NMSE}}$。4-bit 下 $\sqrt{1.06\%}\approx10\%$，3-bit 下约 $20\%$——这已经是相当可观的相对误差。softmax 会把 logits 的差异指数化，所以 K 的位宽必须谨慎；而 V 的误差不经过 softmax 的指数放大，只被权重加权，宽容得多。这就解释了为什么 k3v4（K 3 bit、V 4 bit）的 PPL 代价（+10.63%）明显低于 3bit（K 和 V 都 3 bit，+20.59%）：把有限的比特优先给 K，收益更大。

论文本身区分了两种量化目标：TurboQuantmse 以均方误差为目标，量化重建更接近原向量，但用它做内积估计会引入偏置；TurboQuantprod 在 MSE 量化器的残差上再加一层 1-bit QJL，换来无偏的内积估计。vLLM 采用的前者（MSE 码本 + norm correction），没有采用 QJL，原因下一节会讲。下面两张图是论文在 1536 维 DBpedia 向量上对这两种量化器的实测：

图源：Zandieh 等，TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate（ICLR 2026，arXiv:2504.19874）图 1。上排为 TurboQuantprod，下排为 TurboQuantmse；横轴是内积畸变，四列对应 1 到 4 bit。

图源：同论文图 2。2-bit 下 TurboQuantprod 的内积误差方差在不同平均内积下保持稳定，TurboQuantmse 的偏置随平均内积增大而增大。

norm correction：修正”量化会缩短向量”

Lloyd-Max 重建有一个系统性偏差：对单位向量逐坐标取质心，重建出来的向量长度通常小于 1。原因很直观——质心是条件期望，条件期望会把极端值往中心拉，相当于对每个坐标做了一次收缩。结果 $\|\tilde{c}\| < \|\hat{k}\| = 1$。如果不修正，解码得到的键长度变成 $\|k\|\cdot\|\tilde{c}\|$，所有注意力分数被整体缩小，softmax 的温度被无意中调高，分布变平。

vLLM 的 norm_correction 就是把 $\tilde{c}$ 重新归一化：

这样 $\|\tilde{c}_{\text{nc}}\|=1$，键的长度严格等于存储的 $\|k\|$。源码的 docstring 给出的数据是：在 4-bit 下，norm correction 把 PPL 再降低约 0.8%。这个数字不大，但它是”免费”的——只需要在解码时多做一次 $d$ 维的平方和开方（后来 vLLM 甚至把它折叠到了写入端，见下一节）。

与”每头一个 scale”的对比

传统 INT4/FP8 KV 量化通常为每个头（或每个 token 的每个头）存一个 scale 和一个 zero。以 head_dim=128、4 bit 为例，数据占 64 字节，加上两个 fp16 的 scale/zero 共 68 字节，元数据占 6%。到 3 bit 时数据 48 字节，元数据仍是 4 字节，占比升到 8%。TurboQuant 的 K 侧元数据只有一个 fp16 的向量 norm（2 字节），3-bit 时数据 48 字节，元数据占 4%；4-bit 时数据 64 字节，元数据占 3%。更关键的是，它不需要为每个头校准 scale——码本由维度和 bit 数决定，是全局常量。这既省了元数据，也省掉了校准时对数据的依赖。

代价是：TurboQuant 假定旋转后的坐标严格服从 $\mathcal{N}(0,1/d)$。如果某个模型的 K 分布严重偏离这个假设（比如维度很小、或某些层有异常大的离群坐标），量化误差会比理论值大。vLLM 用”首尾层跳过量化”来对冲这一点，后文会讲到。

从 slot_size 的三项算术开始

TurboQuantConfig 把每个头每个位置的存储拆成两部分。K 部分（key_packed_size）取决于 K 走 FP8 还是 MSE：

1
@property
2
def key_packed_size(self) -> int:
3
    """Packed bytes for a single KEY vector."""
4
    if self.key_fp8:
5
        return self.head_dim  # 1 byte per element
6
    mse_bytes = math.ceil(self.head_dim * self.key_mse_bits / 8)
7
    norm_bytes = 2  # vec_norm fp16
8
    return mse_bytes + norm_bytes
9

10
@property
11
def value_packed_size(self) -> int:
12
    """Packed bytes for a single VALUE vector."""
13
    data_bytes = math.ceil(self.head_dim * self.value_quant_bits / 8)
14
    return data_bytes + 4  # +2 scale(fp16) +2 zero(fp16)
15

16
@property
17
def slot_size(self) -> int:
18
    return self.key_packed_size + self.value_packed_size
19

20
@property
21
def slot_size_aligned(self) -> int:
22
    """Slot size rounded up to next even number."""
23
    s = self.slot_size
24
    return s + (s % 2)

这里有两个容易忽略的细节：

• K 的 MSE 索引按 ceil(d * bits / 8) 字节打包，再加 2 字节的 fp16 向量范数。这个范数是整条 K 向量的 $\|k\|_2$，不是每个坐标的 scale。
• V 的数据按 ceil(d * bits / 8) 打包，再加 2 字节 fp16 scale 和 2 字节 fp16 zero。V 采用”整条向量的 min/max 线性量化”，所有坐标共享一组 scale/zero。
• slot_size_aligned 向上取偶，因为 vLLM 的通用 KV 缓存规格在某些路径上会把 state_content_bytes 折半当作”等效 head_size”处理，要求它能被 2 整除。

把四个预设代入 $d=128$，可以得到源码测试里锁定的那一组具体数字：

这组数字来自 tests/quantization/test_turboquant.py 的 PRESET_EXPECTED，它是对外承诺的布局契约，任何改动都要同步更新测试。

一个 slot 的字节图

以 turboquant_k3v4_nc、$d=128$ 为例，一个 slot 的 118 字节是这样排的：

1
offset  0                      48 50                        114     118
2
        | K indices (3bit×128) | K norm | V indices (4bit×128) | V scale | V zero |
3
        |<----- 48 B --------->| 2 B    |<----- 64 B ---------->| 2 B     | 2 B    |

• 偏移 0..47：128 个 3-bit 质心下标，每 8 个下标打包成 3 字节（$8 \times 3=24$ bit），共 16 组、48 字节。
• 偏移 48..49：K 向量的 fp16 范数 $\|k\|_2$。
• 偏移 50..113：128 个 4-bit 的 V 下标，每两个打包成一个字节，共 64 字节。
• 偏移 114..115：V 的 fp16 scale。
• 偏移 116..117：V 的 fp16 zero。

布局是 AoS（Array of Structures）：同一个 token 的 K 和 V 紧挨在一起，同一页内按 token → head → slot 的顺序排列。这样做的好处是 decode 读一个 token 的 K 和 V 时只碰相邻的两段内存，坏处是 scale/zero 这类”每 token 只读一次”的标量也夹在数据流里，会给向量化访存带来额外的事务。后文的 SoA 优化正是针对这一点。

一个容易踩的坑在 turboquant_attn.py 开头的模块 docstring：它写着

1
Per-head per-position slot layout:
2
  [key_packed (kps bytes) | value_fp16 (D*2 bytes)]
3
  For turboquant_k3v4_nc head_dim=256: [100 bytes key | 512 bytes value] = 612

这段注释停留在”V 保持 FP16”的早期版本，与当前 TurboQuantConfig、store kernel、decode kernel 以及测试里的 118 字节布局都不一致。如果照着注释去算 slot 大小，会多算出一倍多的显存。源码注释会随着迭代过期，布局这类”直接影响正确性”的数字应当以 config 与测试为准，这也是读框架源码时值得养成的习惯。

页内寻址：一页多少字节，slot 落在哪里

state_content_bytes = slot_size_aligned 之后，FullAttentionSpec.page_size_bytes 变成

例如 $n_{kv}=8$、block_size=16、slot=118 时，一页 $16 \times 8 \times 118 = 15104$ 字节。这个页大小比 FP16 的 $16 \times 8 \times 128 \times 2 \times 2 = 65536$ 字节小了 4 倍多，直接决定了同样的显存能装下多少 token。

Triton kernel 里计算 slot 地址的方式是

1
page_idx = kv_offs // BLOCK_SIZE
2
page_off = kv_offs % BLOCK_SIZE
3
block_nums = tl.load(Block_table_ptr + bt_base + page_idx, ...)
4
slot_bases = (
5
    block_nums * stride_cache_block
6
    + page_off.to(tl.int64) * stride_cache_pos
7
    + tl.cast(kv_head, tl.int64) * stride_cache_head
8
)

这里的 stride_cache_pos 是”同一页内相邻 token 的字节跨度”，在 AoS 布局下等于 num_kv_heads * slot_size_aligned；stride_cache_head 等于 slot_size_aligned。注意 block_table 存的是物理页号，与 PagedAttention 的页表语义一致；不同请求的同一个逻辑位置可以映射到不同物理页，这就是分页 KV 缓存的基础。站内 PagedAttention 完全拆解 和 vLLM 源码完全拆解（四）：KVCacheManager 已经把这一层讲透，这里不再重复。

为什么 block_size 只有 16/32/64/128

get_supported_kernel_block_sizes 返回四个值，不是随意选的：

1. 16 是 SoA 优化与 FlyDSL decode 的硬性要求。FlyDSL 的 eligibility 检查里写着 kv_cache.shape[1] in (16, 32)；SoA 布局里 metadata 区按 [H, fields, block_size] 排列，block_size 太小会让 K norm、V scale、V zero 三个数组的交错粒度太细，收益下降。
2. 128 是 vLLM 的默认页大小，能保证大多数模型的调度器按惯常的页粒度工作。
3. 3-bit 打包的字节数必须是 3 的倍数。128 个坐标的 3-bit 索引正好 48 字节，而如果 block_size 不是 8 的倍数，最后一个 token 的 3-bit 组会跨页，写入 kernel 需要处理跨页写，复杂度陡增。16、32、64、128 都是 8 的倍数，避开了这个问题。
4. CUDA Graph 的静态形状：decode kernel 的 BLOCK_SIZE 是编译期常量，页大小一旦确定就不再变化，才能把整张图捕获下来。

AoS 与 SoA：两种页内布局

vLLM 里存在两套 TurboQuant 页内布局：

• AoS（默认，Triton store + Triton decode）：每个 slot 内依次是 K 数据、K norm、V 数据、V scale、V zero。整个页按 token 顺序排列。
• SoA（FlyDSL 路径必需，Triton store + Triton/SoA decode）：把一页拆成”数据区”和”元数据区”两块：

1
block start
2
| data region: [block_size, H, key_data_bytes + val_data_bytes] |
3
| meta region: [H, num_soa_fields, block_size] of fp16             |

数据区只放 packed K 索引和 packed V 数据；元数据区按 head 分组，每个 head 存 block_size 个 K norm、block_size 个 V scale、block_size 个 V zero。源码注释里的说法是 data region [bs, H, KEY_DATA_BYTES+VAL_DATA_BYTES], followed by metadata region [H, NUM_SOA_FIELDS, bs] of fp16。

SoA 的动机是访存局部性：decode kernel 每一轮要读 BLOCK_KV 个 token 的 K norm 和 V scale/zero，AoS 下这些值分散在每 token 的 slot 尾部，需要跨 stride 的小粒度 gather；SoA 下同一个 head 的这些标量在页内是连续的 block_size 个 fp16，可以一次 u16 load 拿到一段。源码注释给出的量化收益是：把 norm correction 的 1/||c_t|| 折叠进存储的 K norm 之后，gpt-oss 在 4096 上下文下 TQ 开销降低约 26%。

SoA 的另一项收益：把 norm correction 折进写入端

AoS 路径的 decode 每次都要对每个 token 的质心向量算一次 $\|\tilde{c}\|$，这需要 $d$ 个平方、一次归约、一次开方、一次除法，并且每个 tile 都要重复。SoA 的 store kernel 把它挪到了写入端：

1
# When norm_correction=True, the kernel folds 1/||c_t|| into the
2
# stored per-token K-norm so the load kernel can skip per-tile norm
3
# math (Opt#1). centroids is only dereferenced inside the kernel when
4
# NORM_CORRECTION=1; pass midpoints as a harmless placeholder otherwise.

1
vn_f32 = tl.load(Norms_ptr + pid)  # ||k_t||, fp32
2

3
if NORM_CORRECTION:
4
    safe_idx = tl.minimum(tl.maximum(idx, 0), N_CENTROIDS - 1)
5
    c_vals = tl.load(Centroids_ptr + safe_idx, mask=d_mask, other=0.0)
6
    c_norm_sq = tl.sum(tl.where(d_mask, c_vals * c_vals, 0.0), axis=0)
7
    c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
8
    vn_f32 = vn_f32 * c_inv_norm
9

10
vn_u16 = vn_f32.to(tl.float16).to(tl.uint16, bitcast=True)
11
tl.store(KV_cache_u16_ptr + knorm_u16_addr, vn_u16)

这段代码的含义是：写入时算出这条 K 的质心向量 $\tilde{c}$，求它的范数 $\|\tilde{c}\|$，把 $\|k\|/\|\tilde{c}\|$ 当作新的”存储范数”。读取时只需要做

就等价于先归一化 $\tilde{c}$ 再乘以 $\|k\|$。每次 decode 只多一次乘加，省掉了每 tile 的归约与开方。注意写入端需要先拿到完整的 idx 数组，因此这项优化只能在 store kernel 内部做，不能放在外层 Python 里——这也是 SoA store 把 bucketize 和 norm folding 融进同一个 kernel 的原因。

元数据为什么用 fp16 而不是 fp32

K norm、V scale、V zero 都存成 fp16。理由有三：

• 带宽：每个 token 每头要读这三个标量，fp32 会让元数据带宽翻倍。decode 是带宽受限的，元数据的每字节都值钱。
• 精度足够：norm 和 scale 本身是从 fp16 的 K/V 算出来的，存成 fp16 不会引入比原始数据更大的误差；zero 也类似。
• 与 SoA 的 u16 寻址对齐：kv_cache.view(torch.uint16) 要求每个 fp16 落在偶数字节上。AoS 布局里 K norm 的偏移是 MSE_BYTES（48 或 64，均为偶数），V scale 的偏移是 KPS + VAL_DATA_BYTES（50+64=114，偶数），保证了 2 字节对齐；这也是 slot_size_aligned 取偶的原因之一。

布局层面的取舍总结

写入发生在 attention 之前

TurboQuant 的 attention 后端声明 forward_includes_kv_cache_update = False，因此模型层在执行注意力前，会先调用 do_kv_cache_update：

1
def do_kv_cache_update(self, layer, key, value, kv_cache, slot_mapping):
2
    N = slot_mapping.shape[0]
3
    if N <= 0:
4
        return
5
    device = key.device
6
    self._ensure_on_device(layer, device)
7

8
    k = key[:N].view(N, self.num_kv_heads, self.head_size)
9
    v = value[:N].view(N, self.num_kv_heads, self.head_size)
10
    # (B, H, N, C) -> (B, N, H, C) for TQ kernels
11
    kv_cache = kv_cache.transpose(1, 2)
12
    self._store_kv(k, v, kv_cache, slot_mapping, layer)

这段代码包含三个关键约定：

1. 只处理 N 个实际 token：key 和 value 是批内全部 token 的 QKV 投影结果，slot_mapping 被调用方裁到 num_actual_tokens，避免把 padding token 写进缓存。kernel 内部还会再次检查 slot < 0 并直接返回，这是给 CUDA Graph 的 padding 槽位留的保险。
2. K/V 的视图是 (N, H_kv, D)：vLLM 的投影输出通常把 head 维拼在最后一维，这里按 KV 头数重新解释；head_size 是模型真实的 K 维度，不是 slot 的等效 head size。
3. 缓存转置成 (B, N, H, C)：vLLM 的缓存物理布局是 (num_blocks, block_size, num_kv_heads, slot_bytes)，也就是逻辑上的 (B, N, H, C)；但有些上层路径传进来的是 (B, H, N, C) 视图，这里统一转置，后续 kernel 的 stride 计算才一致。

_store_kv 再按 _soa_store 开关分派到两套 Triton store：AoS 版本和 SoA 版本。参数完全一致，差别只在页内偏移的计算方式。

K 的归一化与旋转为什么放在 kernel 外面

MSE 路径的 Python 侧代码非常短：

1
# Normalize + rotation GEMM externally (cuBLAS is faster than in-kernel)
2
k_flat = key.float().reshape(NH, D)
3
norms = k_flat.norm(dim=1, keepdim=True)
4
x_hat = k_flat / (norms + 1e-8)
5
y = x_hat @ PiT
6

7
v_flat = value.float().reshape(NH, D)

它做了四件事：

• 把 K 转成 fp32，避免 fp16 归一化时的下溢；
• 求每个头向量的 L2 范数 norms；
• 归一化成单位向量 x_hat，分母加 1e-8 防止零向量；
• 右乘 PiT（Hadamard 矩阵）得到旋转后的单位向量 y。

为什么不把归一化和旋转也塞进 Triton kernel？源码注释给出的理由是 cuBLAS 更快：x_hat @ PiT 是一个 $(N H_kv) \times D$ 乘 $D \times D$ 的 GEMM，对于 $D=128$、prefill 几千 token 的规模，cuBLAS 的 tiling 和 Tensor Core 利用率远高于一个手写 Triton 的逐元素归约。把旋转放在外面，kernel 只负责”给定旋转后的单位向量，做查表和打包”，每个 program 的计算量小、访存规整，反而更容易打满带宽。

代价是多了一次显存的往返：y 要物化成 fp32 张量，k_flat 也是 fp32，prefill 8K token、8 个 KV 头时约 32 MiB 级别，可以接受。对 decode 这种每步只有几十个 token 的场景，这个 GEMM 的 launch 开销占比更高，但 decode 的瓶颈本来就在 KV 读取而不是写入，所以问题不大。

融合 kernel 的骨架：一个 program 处理一个 (token, head)

store kernel 的 grid 是 (NH,)，即每个 token 的每个 KV 头一个 program。pid 同时编码 token 和 head：

1
pid = tl.program_id(0)
2
token_idx = pid // H
3
head_idx = pid % H
4

5
slot = tl.load(Slot_mapping_ptr + token_idx)
6
if slot < 0:
7
    return
8
blk = (slot // BLOCK_SIZE).to(tl.int64)
9
off = (slot % BLOCK_SIZE).to(tl.int64)
10
head_idx_i64 = tl.cast(head_idx, tl.int64)
11
slot_base = (
12
    blk * stride_cache_block
13
    + off * stride_cache_pos
14
    + head_idx_i64 * stride_cache_head
15
)

这个 program 的粒度选择有两层考虑：

• 一个 program 正好覆盖一条 K 向量和一条 V 向量，所有坐标都能放进寄存器/SRAM，不需要跨 program 同步；
• H 是编译期常量，token_idx 和 head_idx 的除法取模被编译器优化掉；slot_base 的三个 stride 也都是 constexpr，地址算术不会变成运行时乘法。

slot_mapping 是调度器为每个 token 分配的全局逻辑槽位，slot // BLOCK_SIZE 得到物理页号，slot % BLOCK_SIZE 得到页内偏移。这里没有直接使用 block_table，因为 store 端的 slot_mapping 已经是经过页表映射后的物理槽位；只有读取端才需要按 block_table 查页。

二分查找：3 到 4 次比较替代线性扫描

K 的量化核心是”给定一个旋转后的坐标 $y$，找到它落在哪个 Lloyd-Max 区间”。最直接的做法是线性扫描 n_centroids - 1 个边界，3-bit 要 7 次比较、4-bit 要 15 次。kernel 用了二分：

1
y_vec = tl.load(Y_ptr + base + d_offs, mask=d_mask, other=0.0)
2
lo = tl.zeros([BLOCK_D], dtype=tl.int32)
3
hi = tl.full([BLOCK_D], N_CENTROIDS - 1, dtype=tl.int32)
4
for _ in range(MSE_BITS):
5
    mid = (lo + hi) >> 1
6
    safe_mid = tl.minimum(mid, N_CENTROIDS - 2)
7
    mid_val = tl.load(Midpoints_ptr + safe_mid, mask=d_mask, other=0.0)
8
    lo = tl.where(y_vec >= mid_val, mid + 1, lo)
9
    hi = tl.where(y_vec >= mid_val, hi, mid)
10
idx = tl.minimum(lo, N_CENTROIDS - 1)

逐行解释：

• lo=0, hi=n_centroids-1，搜索区间是质心下标的闭区间。循环 MSE_BITS 次（3 或 4 次），每次把区间折半，最后 lo 就是插入位置。
• safe_mid = min(mid, n_centroids-2)：当 mid 落到最后一个质心时，midpoints 只有 n_centroids-1 个元素，直接索引会越界；把它夹到最后一个边界，比较结果仍然正确，因为已经收敛的 lane 不再改变 mid。
• tl.where(y_vec >= mid_val, mid+1, lo) 与 hi = where(..., hi, mid) 是标准的 lower_bound 写法；用 >= 而不是 >，保证等于边界时落到右侧区间，和训练时”左闭右开”的定义一致。
• mask=d_mask：BLOCK_D 是 next_power_of_2(D)，当 $D$ 不是 2 的幂时（例如 96），超出部分不参与加载；y_vec 的对应 lane 是 0，比较结果会被后面的 mask 丢掉。

为什么用二分而不是把边界做成查找表？因为 midpoints 是 7 或 15 个 fp32，常量缓存里一放，二分查找只需要 3 到 4 次常量加载；而查找表要么是 256 项的线性表（需要更多常量内存），要么是插值近似（会引入额外误差）。二分在 GPU 上还有一个好处：每个 lane 独立走同样的迭代次数，没有 warp divergence，3 到 4 次循环完全展开后就是 3 到 4 条比较指令。

位打包：3-bit 的 8 个一组与 4-bit 的 2 个一组

查表得到下标后，要把 int32 的 idx 向量压成字节流。4-bit 的打包很直观：

1
idx_pairs = tl.reshape(idx, [BLOCK_D // 2, 2])
2
shifts_4 = tl.arange(0, 2) * 4
3
packed = tl.sum((idx_pairs & 0xF) << shifts_4[None, :], axis=1).to(tl.uint8)
4
mse_offs = tl.arange(0, BLOCK_D // 2)
5
mse_mask = mse_offs < MSE_BYTES
6
tl.store(KV_cache_ptr + slot_base + mse_offs, packed, mask=mse_mask)

idx_pairs 是 (64, 2)，每个 pair 的两个 4-bit 值分别左移 0 和 4 位，按位或（用 sum 实现，因为两个 nibble 不重叠）得到一个字节。mse_offs = arange(0, 64) 是字节偏移，mse_mask 保证只写到 MSE_BYTES（$D=128$ 时为 64）个字节。

3-bit 更麻烦，因为 8 个 3-bit 值正好 24 bit，跨 3 个字节，不能按 2 的幂对齐：

1
grp_offs = tl.arange(0, BLOCK_GRP)
2
grp_mask = grp_offs < (D // 8)
3
idx_grp = tl.reshape(idx, [BLOCK_GRP, 8])
4
shifts_3 = tl.arange(0, 8) * 3
5
packed_24 = tl.sum((idx_grp & 0x7) << shifts_3[None, :], axis=1)
6
b0 = (packed_24 & 0xFF).to(tl.uint8)
7
b1 = ((packed_24 >> 8) & 0xFF).to(tl.uint8)
8
b2 = ((packed_24 >> 16) & 0xFF).to(tl.uint8)
9
tl.store(KV_cache_ptr + slot_base + grp_offs * 3, b0, mask=grp_mask)
10
tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 1, b1, mask=grp_mask)
11
tl.store(KV_cache_ptr + slot_base + grp_offs * 3 + 2, b2, mask=grp_mask)

BLOCK_GRP 是 next_power_of_2(D // 8)，$D=128$ 时为 16。每个 group 的 8 个下标被排成 (16, 8)，分别左移 0、3、6、…、21 位，累加成一个 24-bit 整数，再拆成低、中、高三个字节。这里用 sum 而不是 or 也可以，因为 3-bit 值只占 0..7，移位后不会重叠；但如果某个 lane 的 idx 越界（理论上有 min 保护，不会发生），sum 会产生进位而 or 不会，所以生产代码里 min 是必要的。

grp_offs * 3 是这个设计的核心约束：每个 3-bit 组必须从 3 的倍数偏移开始。这要求 D 是 8 的倍数（128、96、64 都满足），否则最后一组不足 8 个值，打包后的字节数不是 3 的倍数，跨 slot/跨页的边界就会错位。vLLM 的 TurboQuant 只对 head_dim 为 8 的倍数的模型开放，测试里也只覆盖 64、96、128、256 这些常见值。

向量范数的存储与 norm correction 的折叠

K 的范数在 MSE 路径里由 Python 侧算好传进来（norms），kernel 只负责存成 fp16：

1
vn_f32 = tl.load(Norms_ptr + pid)  # ||k_t||, fp32
2

3
if NORM_CORRECTION:
4
    safe_idx = tl.minimum(tl.maximum(idx, 0), N_CENTROIDS - 1)
5
    c_vals = tl.load(Centroids_ptr + safe_idx, mask=d_mask, other=0.0)
6
    c_norm_sq = tl.sum(tl.where(d_mask, c_vals * c_vals, 0.0), axis=0)
7
    c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
8
    vn_f32 = vn_f32 * c_inv_norm
9

10
vn_u16 = vn_f32.to(tl.float16).to(tl.uint16, bitcast=True)
11
tl.store(KV_cache_u16_ptr + knorm_u16_addr, vn_u16)

idx 是刚刚二分得到的质心下标向量，Centroids_ptr + safe_idx 是一次 gather：每个坐标取自己的质心值，拼成重建向量 $\tilde{c}$。c_norm_sq 是 $\|\tilde{c}\|^2$，c_inv_norm 是 $1/\|\tilde{c}\|$。当 NORM_CORRECTION=1 时，存储的范数变成 $\|k\|/\|\tilde{c}\|$，读取端就可以直接做 $\tilde{k} = \text{stored\_norm} \cdot \tilde{c}$，不必再算范数。

这段代码也解释了为什么 norm correction 几乎”免费”：写入时每个 token 已经要读一遍 idx 和 centroids，多算一次 $d$ 维平方和与开方，相对于整个 store kernel 的访存量可以忽略；而读取端每次 decode 都要对每个 tile 的每个 token 重复算，省下来的就是实打实的算力。源码注释里说，在 gpt-oss、4096 上下文下这条优化减少了约 26% 的 TQ 开销，量级是合理的。

V 的 min/max 均匀量化

V 的处理和 K 完全不同：不做旋转、不做 Lloyd-Max，而是对每条 V 向量做 min/max 线性量化。3-bit 分支：

1
val_vec = tl.load(Value_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
2
val_min = tl.min(tl.where(d_mask, val_vec, float("inf")), axis=0)
3
val_max = tl.max(tl.where(d_mask, val_vec, -float("inf")), axis=0)
4
v_scale = (val_max - val_min) / 7.0
5
v_scale = tl.where(v_scale > 1e-8, v_scale, 1e-8)
6

7
q_vals = tl.minimum(
8
    tl.maximum(((val_vec - val_min) / v_scale + 0.5).to(tl.int32), 0), 7
9
)

四个细节：

• tl.min/tl.max 用 where(d_mask, ..., ±inf) 把 padding lane 排除。如果直接对未 mask 的 0 取 min，零会被误当成最小值，量化区间就被污染了。
• v_scale = (max-min)/7：3-bit 有 8 个电平，区间宽度除以 2^b-1 而不是 2^b，保证最大值能映射到 7。
• +0.5 后转 int32 实现四舍五入；clamp(0, 7) 防止浮点误差把边界值推到 8 或 -1。
• v_scale 下限 1e-8：常数向量（max=min）时 scale=0，除零会产生 NaN；夹到极小值后所有值都量化到 0，重建为 val_min，语义正确。

4-bit 分支只是把 7.0 换成 15.0、0x7 换成 0xF，其余结构相同。打包时 3-bit 走”8 个一组、3 字节”的路径，4-bit 走”2 个一组、1 字节”的路径，与 K 的打包逻辑对称。

scale 和 zero 都存成 fp16：

1
sc_offset = val_cache_offset + VAL_DATA_BYTES
2
sc_f16 = v_scale.to(tl.float16)
3
sc_u16 = sc_f16.to(tl.uint16, bitcast=True)
4
tl.store(KV_cache_ptr + slot_base + sc_offset, (sc_u16 & 0xFF).to(tl.uint8))
5
tl.store(KV_cache_ptr + slot_base + sc_offset + 1, ((sc_u16 >> 8) & 0xFF).to(tl.uint8))

这里用两个 uint8 写一个 fp16，而不是直接 tl.store 一个 fp16 指针：KV_cache_ptr 是 uint8*，Triton 的指针类型是显式的，不能在一个 uint8* 上直接存 fp16。用 bitcast 加两次字节写，编译器会把它们合并成一条 16-bit store（如果地址对齐），即使不合并，两次相邻写也落在同一 cache line 里。

一次写入的成本账

把上面的步骤合起来，一个 program 的访存和计算大致是：

• 读 K：$D$ 个 fp16（Python 侧已经读过一次做归一化，这里读的是旋转后的 fp32 y，$4D$ 字节）；
• 读 V：$D$ 个 fp16（$2D$ 字节）；
• 读 midpoints：3 到 4 次标量常量加载；
• 写 K 索引：$D \cdot b_k / 8$ 字节；
• 写 K norm：2 字节；
• 写 V 数据：$D \cdot b_v / 8$ 字节；
• 写 V scale/zero：4 字节。

$D=128$、3-bit K、4-bit V 时，单头单 token 约读 768 字节、写 118 字节。这个比例说明 store kernel 是”读多写少”的带宽型 kernel，计算（二分、打包、min/max）相对于访存很轻。融合的主要收益不是省算力，而是省 kernel launch 和中间张量：如果分成”归一化、旋转、二分、打包、V 量化、写缓存”六个 kernel，每个 kernel 都要把中间结果写回显存再读出来，prefill 大 batch 下这部分流量会翻倍。源码注释明确写了融合消除了”每次调用 4 个 PyTorch kernel launch”（相对于先 pack、再单独写 scale 的旧实现）。

FP8 K 路径：不旋转，直接 cast

turboquant_k8v4 走另一条分支。K 不归一化、不旋转、不查表，而是在 kernel 里直接转 FP8：

1
k_vals = tl.load(Key_ptr + base + d_offs, mask=d_mask, other=0.0).to(tl.float32)
2
k_fp8 = k_vals.to(tl.float8e4b15) if FP8_E4B15 else k_vals.to(tl.float8e4nv)
3
k_bytes = k_fp8.to(tl.uint8, bitcast=True)
4
tl.store(KV_cache_ptr + slot_base + d_offs, k_bytes, mask=d_mask)

FP8_E4B15 是一个按设备能力缓存的常量：

1
def _use_fp8_e4b15(device: int = 0) -> int:
2
    """Return 1 if device needs fp8e4b15 (Ampere/Ada, SM < 8.9), else 0."""
3
    if device not in _FP8_E4B15:
4
        if current_platform.is_cuda_alike():
5
            cap = torch.cuda.get_device_capability(device)
6
            _FP8_E4B15[device] = 1 if cap < (8, 9) else 0
7
        else:
8
            _FP8_E4B15[device] = 0
9
    return _FP8_E4B15[device]

Ampere（SM 8.0/8.6）和 Ada（SM 8.9 之前）没有原生的 E4M3 FP8 指令，只能用 E4B15 这种”指数 4 位、尾数 3 位”的格式来模拟；Hopper（SM 9.0）及以上有 e4nv，精度和动态范围都更好。这个分支只影响 K 的 cast，V 仍然是 4-bit 量化，因此 k8v4 的 slot 是 128 + 68 = 196 字节。

k8v4 的意义是给”只要比 FP16 省一半、但不想承担 K 的量化误差”的场景一个中间档。它的 PPL 增量只有 +1.17%，在四个预设里最小，但压缩比也最低（2.6×）。值得注意的是，它没有 norm correction：FP8 的 per-tensor 动态范围已经足够，K 的范数信息保留在 FP8 的指数里，不需要额外的 norm 存储。

读取比写入难：分页、解包和在线 softmax 要同时处理

写入时一个 token 的 K/V 就是全部数据；读取时一条请求的历史可能横跨几十个物理页，页号由 block_table 决定，页内位置由 seq_len 决定。decode 每一步对每个请求只产生一个 query，却要读完整段历史 KV。TurboQuant 的 decode kernel 采用”split-KV + 两阶段”结构，和 vLLM 的 Triton 解码注意力一致：

• Stage 1：grid 是 (B, H_q, NUM_KV_SPLITS)，每个 program 负责一个请求、一个 query 头、一段 KV 区间。它在 tile 内部完成”解包 K → 算分 → 在线 softmax → 解包 V → 加权累加”，最后写出这一段的局部部分和（$acc$）与 log-sum-exp（$lse$）。
• Stage 2：grid 是 (B, H_q)，复用 triton_decode_attention._fwd_kernel_stage2，对 NUM_KV_SPLITS 个部分和做一次合并，得到最终输出。

NUM_KV_SPLITS 是 max_num_kv_splits，来自 attention 配置 tq_max_kv_splits_for_cuda_graph。注释里说得很明确：grid 维度必须是常量才能进 CUDA Graph，而基准测试显示固定 split 数和动态 split 在 eager 模式下没有性能回退。这是一个典型的”服务系统优先”的决策：牺牲一点动态性，换 CUDA Graph 的可捕获性。

程序入口：split 区间怎么算

1
bid = tl.program_id(0)  # batch index
2
hid = tl.program_id(1)  # q_head index
3
sid = tl.program_id(2)  # kv_split index
4

5
kv_head = hid // KV_GROUP_SIZE
6

7
seq_len = tl.load(Seq_lens_ptr + bid)
8

9
split_len = tl.cdiv(seq_len, NUM_KV_SPLITS)
10
split_start = split_len * sid
11
split_end = tl.minimum(split_start + split_len, seq_len)
12

13
if split_start >= split_end:
14
    return

这里有三个值得注意的点：

1. kv_head = hid // KV_GROUP_SIZE：GQA 下多个 query 头共享一个 KV 头，KV_GROUP_SIZE = H_q / H_kv 是编译期常量。grid 按 query 头展开（每个 query 头独立算分），而 KV 读取按 KV 头寻址，这样同一段 KV 会被 KV_GROUP_SIZE 个 program 重复读。这是 GQA 的固有代价，也是为什么 split 数不能开太大：split 越多，KV 读取的重复度越高。
2. split_len = ceil(seq_len / splits)：区间按”向上取整”均分，最后一个 split 可能比前面短。split_end 用 min 夹到 seq_len。
3. split_start >= split_end 直接返回：当序列短于 split 数时，尾部几个 split 没有数据，直接退出，避免空转。返回是安全的，因为 stage 2 会把没有写过的中间缓冲当成 lse = -inf 处理（前提是缓冲区被初始化过；实际实现里工作区预分配时已清零）。

为解包预计算位索引

K 和 V 的解包都要”从字节里取第 $i$ 个 3-bit/4-bit 值”，所以 kernel 在循环外先把位偏移、字节下标、位内偏移算好：

1
if not KEY_FP8:
2
    mse_bit_off = d_offs * MSE_BITS
3
    mse_byte_idx = mse_bit_off // 8
4
    mse_bit_shift = mse_bit_off % 8
5
    mse_mask = (1 << MSE_BITS) - 1
6

7
if VQB == 3:
8
    val_bit_off = d_offs * 3
9
    val_byte_idx = val_bit_off // 8
10
    val_bit_shift = val_bit_off % 8

d_offs = arange(0, BLOCK_D) 是坐标下标向量，mse_bit_off 是每个坐标的起始 bit。mse_byte_idx 和 mse_bit_shift 是它的字节位置和位内起始位置。这些量只依赖 d_offs，与 token、split 无关，放在循环外可以避免每个 tile 重复计算 $D$ 次除法与取模。

mse_mask = 2^b - 1 用于 (raw16 >> shift) & mask 提取。这里用 raw16 = raw0 | (raw1 << 8) 把两个相邻字节拼成 16-bit，再按 mse_bit_shift 右移、按 mse_mask 截取。为什么拼 16 位？因为一个 3-bit 值可能跨字节边界（例如 bit 6 开始的 3 位占 byte0 的 bit6-7 与 byte1 的 bit0），拼成 16 位后一次移位可以覆盖所有跨字节情况；4-bit 值同样可能跨越，用 16 位拼装比两次分支判断更简单。

分数计算：FP8 K 与 MSE K 的两条路径

FP8 K 路径直接把字节 cast 回 FP8 再乘：

1
k_addrs = slot_bases[:, None] + d_offs[None, :]
2
k_raw = tl.load(KV_cache_ptr + k_addrs, mask=kv_mask[:, None] & d_mask[None, :], other=0)
3
if FP8_E4B15:
4
    k_float = k_raw.to(tl.float8e4b15, bitcast=True).to(tl.float32)
5
else:
6
    k_float = k_raw.to(tl.float8e4nv, bitcast=True).to(tl.float32)
7
scores = tl.sum(tl.where(d_mask[None, :], q_rot[None, :] * k_float, 0.0), axis=1) * ATTN_SCALE

注意 slot_bases 的形状是 (BLOCK_KV,)，[:, None] 变成 (BLOCK_KV, 1)，加上 d_offs[None, :] 的 (1, BLOCK_D) 得到 (BLOCK_KV, BLOCK_D) 的地址矩阵，一次 load 拿到整个 tile 的 K。q_rot 是 (BLOCK_D,)，广播成 (BLOCK_KV, BLOCK_D) 后逐元素乘、按 axis=1 求和，得到每个 token 的分数。

ATTN_SCALE 是 1/sqrt(D)，在 launcher 里从 scale 参数传入。这里没有像标准的 softmax(q·k^T / sqrt(d)) 那样把 scale 放在最后的 softmax 前，而是乘在分数上，效果等价，但少一次逐 token 的乘。

MSE K 路径更复杂，需要”解包 + 查表 + 点积 + 乘范数”：

1
mse_addrs0 = slot_bases[:, None] + mse_byte_idx[None, :]
2
mse_raw0 = tl.load(KV_cache_ptr + mse_addrs0, mask=..., other=0).to(tl.int32)
3
mse_raw1 = tl.load(KV_cache_ptr + mse_addrs0 + 1, mask=..., other=0).to(tl.int32)
4
raw16 = mse_raw0 | (mse_raw1 << 8)
5
mse_idx = (raw16 >> mse_bit_shift[None, :]) & mse_mask
6

7
c_vals = tl.load(Centroids_ptr + mse_idx, mask=..., other=0.0)
8

9
term1 = tl.sum(tl.where(d_mask[None, :], q_rot[None, :] * c_vals, 0.0), axis=1)
10

11
norm_bases = slot_bases + MSE_BYTES
12
n_lo = tl.load(KV_cache_ptr + norm_bases, mask=kv_mask, other=0).to(tl.uint16)
13
n_hi = tl.load(KV_cache_ptr + norm_bases + 1, mask=kv_mask, other=0).to(tl.uint16)
14
vec_norms = (n_lo | (n_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)
15

16
scores = vec_norms * term1 * ATTN_SCALE

这条路径的核心是 c_vals = tl.load(Centroids_ptr + mse_idx, ...)：mse_idx 是一个 (BLOCK_KV, BLOCK_D) 的下标矩阵，每个元素都是 [0, 2^b) 的整数，这一步是一次”tile 内 gather”。Triton 的 gather 在底层是每个 lane 独立从常量内存或 L1 读取，centroids 只有 8 或 16 个 fp32，完全能放进常量缓存，所以 gather 的代价是每个元素的几次缓存访问，而不是 DRAM 随机访问。这也是为什么 Lloyd-Max 码本只支持 3 到 4 bit：如果码本太大（比如 8 bit 的 256 个中心），gather 的缓存命中率会下降，反而不如直接存 FP8。

term1 是 $q_{\text{rot}} \cdot \tilde{c}$，即”旋转后的 query 和质心向量的点积”。因为旋转矩阵是正交且对称的，$q_{\text{rot}}^\top \tilde{c} = (H q)^\top \tilde{c} = q^\top (H \tilde{c})$，而 $H \tilde{c}$ 正是解包后的 K（未乘范数）。所以 term1 乘上存储的范数，就是 q·k 的近似值。这个推导说明Q 必须旋转：如果 Q 不旋转，直接用原始 Q 和旋转后的质心点积，结果完全不对。

当 NORM_CORRECTION=0（AoS 路径的 k8v4 之外的某些组合）时，vec_norms 就是 $\|k\|$，读取端直接乘；而当 NORM_CORRECTION=1 且是 SoA store 时，存储的范数已经是 $\|k\|/\|\tilde{c}\|$，读取端直接乘就完成了归一化，不需要再算质心范数。AoS 路径的 NORM_CORRECTION 分支则是在 kernel 里实时算：

1
if NORM_CORRECTION:
2
    c_norm_sq = tl.sum(tl.where(d_mask[None, :], c_vals * c_vals, 0.0), axis=1)
3
    c_inv_norm = 1.0 / tl.sqrt(c_norm_sq + 1e-16)
4
    c_vals = c_vals * c_inv_norm[:, None]

这段代码在每个 tile、每个 token 上算一次 $d$ 维平方和、开方、缩放，这正是 SoA 路径把它折叠到写入端要省掉的工作。

V 的解包：3-bit 的跨字节与 4-bit 的 nibble

V 的解包与 K 的索引解包类似，但多了 scale/zero 的读取。4-bit 分支：

1
vb_idx = d_offs // 2
2
vb_shift = (d_offs % 2) * 4
3
val_addrs = val_bases[:, None] + vb_idx[None, :]
4
val_raw = tl.load(KV_cache_ptr + val_addrs, mask=..., other=0).to(tl.int32)
5
v_idx = ((val_raw >> vb_shift[None, :]) & 0xF).to(tl.float32)

每个字节装两个坐标：偶数坐标取低 4 位，奇数坐标取高 4 位。vb_idx 把 d_offs 折半，vb_shift 在 0/4 之间交替。

3-bit 分支用与 K 相同的”拼 16 位”策略：

1
val_bit_off = d_offs * 3
2
val_byte_idx = val_bit_off // 8
3
val_bit_shift = val_bit_off % 8
4
val_raw0 = tl.load(KV_cache_ptr + val_addrs0, ...)
5
val_raw1 = tl.load(KV_cache_ptr + val_addrs0 + 1, ...)
6
raw16 = val_raw0 | (val_raw1 << 8)
7
v_idx = ((raw16 >> val_bit_shift[None, :]) & 0x7)

scale 和 zero 从 slot 的 val_bases + VAL_DATA_BYTES 处各读 2 字节，反序列化成 fp16：

1
sc_lo = tl.load(KV_cache_ptr + sc_bases, mask=kv_mask, other=0).to(tl.uint16)
2
sc_hi = tl.load(KV_cache_ptr + sc_bases + 1, mask=kv_mask, other=0).to(tl.uint16)
3
v_scales = (sc_lo | (sc_hi << 8)).to(tl.float16, bitcast=True).to(tl.float32)

注意 sc_bases = val_bases + VAL_DATA_BYTES 是 (BLOCK_KV,) 形状，mask=kv_mask 是逐 token 的，所以每个 token 读一次 scale 和 zero，而不是每个坐标读一次。这在循环里只多了 2 条 load，相对于 (BLOCK_KV, BLOCK_D) 的 V 数据读取是九牛一毛。

在线 softmax：为什么要用”减法”而不是”除法”

TurboQuant 的 stage 1 用的是标准在线 softmax：

1
n_e_max = tl.maximum(tl.max(scores, 0), m_prev)
2
re_scale = tl.exp(m_prev - n_e_max)
3
p = tl.exp(scores - n_e_max)
4
acc = acc * re_scale + tl.sum(p[:, None] * values, 0)
5
l_prev = l_prev * re_scale + tl.sum(p, 0)
6
m_prev = n_e_max

逐条解释：

• scores 已经是 q·k * scale，可能很大（fp16 下 logits 超过 11 就溢出 exp）；先减去当前 tile 的最大值 n_e_max 再 exp，保证指数参数 ≤ 0，exp 结果在 $[0,1]$。
• re_scale = exp(m_prev - n_e_max) 修正之前 tile 累积的 acc 和 l：之前用的是旧的 max m_prev，现在 max 变大（或持平），之前所有项的指数都要按比例缩小。
• acc 是未归一化的加权 V 和，l 是未归一化的权重和。每个 tile 结束时 acc 和 l 的比值就是归一化后的输出。
• 最后存 acc / safe_l 和 lse = m_prev + log(l)：lse 是 log-sum-exp，stage 2 用它做跨 split 合并。

safe_l = where(l_prev > 0, l_prev, 1.0) 是为了防止空 split（split 区间全是 mask 掉的 token）时除以 0。空 split 的 acc=0, l=0，存 0/1=0、lse=-inf+log(1)=-inf，stage 2 合并时 exp(-inf)=0，不影响结果。

数值精度：累加器为什么是 fp32

acc 声明为 tl.zeros([BLOCK_D], dtype=tl.float32)，m_prev、l_prev 也是 fp32。原因有二：

1. 累加：acc 要把整个 BLOCK_KV 行的 p * v 加起来，p 在 $[0,1]$、v 是 fp16 反序列化的值，如果 acc 用 fp16，几千个 token 累加后会丢精度；fp32 的 24 位尾数足够。
2. exp 与 log：Triton 的 tl.exp/tl.log 默认返回 fp32，用 fp16 反而会引入额外转换。

输出阶段 acc / safe_l 仍然在 fp32 里算，最后在 launcher 的 stage 2 里按 out_dtype 转成 query 的 dtype（fp16/bf16）。

prefill 路径：不读缓存，直接算，再写缓存

Prefill 不走 decode kernel。forward 按批次类型分三种情况：

• 纯 decode：所有请求 max_query_len == 1，直接调 _decode_attention。
• 纯 prefill：num_decodes == 0，走 _prefill_attention。
• 混合批：decode 在前、prefill 在后（由 reorder_batch_threshold=1 保证），两段分别调，结果写进同一块输出。

_prefill_attention 里有一个重要的快速路径：

1
if _HAS_FLASH_ATTN and attn_metadata.max_query_len == attn_metadata.max_seq_len:
2
    return self._flash_attn_varlen(
3
        q=query, k=key, v=value,
4
        cu_seqlens_q=attn_metadata.query_start_loc,
5
        cu_seqlens_k=attn_metadata.query_start_loc,
6
        max_seqlen_q=attn_metadata.max_query_len,
7
        max_seqlen_k=attn_metadata.max_query_len,
8
    )