KV 量化

主线位置：一个请求的一生里前向那一步：注意力层把新算的 K、V 写进缓存时量化，读历史 KV 时反量化；配置在启动参数里，模型加载时定下 scale。两个引擎这一段都是 Python 编排加 CUDA / Triton kernel，没有 Rust。

版本：vLLM v0.32.0、SGLang v0.5.22

结论先说：

• 两边的开关都叫 --kv-cache-dtype，默认都是 auto：用模型的 dtype；如果 checkpoint 的量化配置声明了 FP8 KV，就用 FP8。
• 最常用的 FP8 KV，两边存法一样：数据张量的 dtype 是 uint8（按 FP8 位模式解释），scale 是每层一个 K scale、一个 V scale（per-tensor），挂在注意力层对象上，不进 KV 缓存。写缓存前除以 scale，读出来乘回去。
• scale 从 checkpoint 读；读不到就是 1.0，两边都会打警告。SGLang 还能从 --quantization-param-path 指定的 JSON 读。tag 上两边都没有运行时校准 per-tensor scale 的参数。
• 更细的格式各走各的：vLLM 的 per-token-head 模式和 NVFP4 把 scale 和数据放在同一页里；SGLang 的 MXFP8、NVFP4 把块 scale 放进独立的缓冲。
• 量化后能多放多少块、吞吐和精度变化多少，待实测（C3-5）。

能选什么

vLLM 的完整取值见 CacheDType，每个取值属于哪种量化方式由 get_kv_quant_mode 判定。并不是每个注意力后端都支持每种取值，后端各自声明 supported_kv_cache_dtypes，比如 Triton 后端支持 per-token-head 三种（triton_attn.py）。

auto 的解析：

• vLLM：resolve_kv_cache_dtype_string 看 checkpoint 的 quantization_config 里有没有 KV 量化算法，有就用它，否则保持 auto，即模型 dtype。
• SGLang：configure_kv_cache_dtype 看量化配置的 kv_cache_quant_algo 是不是 FP8，是就用 float8_e4m3fn（ROCm 上是 fnuz 变体），否则用模型 dtype。

FP8 per-tensor：数据在缓存里，scale 在层上

数据：uint8

vLLM 把 fp8、fp8_e4m3、fp8_e5m2 都映射成 torch.uint8 分配缓存（torch_utils.py），kernel 里再 view 成 FP8 类型（flash_attn.py）。SGLang 同样用 uint8 存，注释给了原因：

python/sglang/srt/mem_cache/memory_pool.py

        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn, torch.float8_e4m3fnuz):
            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
            self.store_dtype = torch.uint8
        else:
            self.store_dtype = dtype

缓存的形状不变，只是每个元素从 2 字节变 1 字节：vLLM 一块的字节数公式里 dtype 字节数从 2 变成 1（见 分页注意力），SGLang 每层 K、V 两个 (槽位数 + page_size, 头数, head_dim) 张量也一样。

scale：每层一对，挂在注意力层上

两边的 FP8 scale 加载代码几乎一样，SGLang 那份文件头注明改编自 vLLM（kv_cache.py）。vLLM 的版本：

vllm/model_executor/layers/quantization/kv_cache.py

        if is_quantized_kv_cache(layer.kv_cache_dtype):
            if layer.k_scale > 0.0 and layer.v_scale > 0.0:
                # We prefer to use separate k_scale and v_scale if present
                k_scale = layer.k_scale.to("cpu").tolist()
                v_scale = layer.v_scale.to("cpu").tolist()
                # ...
            elif layer.k_scale < 0.0 and layer.v_scale < 0.0:
                # If no scales were loaded (both scales are invalid negative
                # values), use the default value of 1.0
                k_scale = 1.0
                v_scale = 1.0
            else:
                # ...
                scale_to_duplicate = max(layer.k_scale, layer.v_scale)
                # ...
            if not isinstance(k_scale, float) or not isinstance(v_scale, float):
                raise ValueError(
                    "Only support per-tensor scaling factor for fp8 KV cache"
                )
            # ...
            layer._k_scale.copy_(k_scale)
            layer._v_scale.copy_(v_scale)

• scale 参数初始化成 −1（非法值），checkpoint 里有 k_scale / v_scale 就被覆盖（L68）。
• 两个都有就各用各的；只有一个（老格式的单个 kv_scale）就 K、V 共用；都没有就用 1.0，并警告 Using KV cache scaling factor 1.0 for fp8_e4m3。
• 只接受标量，per-channel 的 scale 直接报错。
• 模型本身没有量化配置时，根本不走这段：注意力层初始化就把 scale 设成 1.0（set_default_quant_scales）。所以对一个普通 bf16 checkpoint 加 --kv-cache-dtype fp8，scale 就是 1.0。

SGLang 多一条路：--quantization-param-path 指向一个 JSON，kv_cache_dtype 为 fp8_e4m3 时由模型自己的 load_kv_cache_scales 读进来；没给就警告 Defaulting to scaling factors of 1.0（入口 load_kv_cache_scales）。以 Llama 为例，JSON 里每层一个值，同时赋给 K 和 V（llama.py）。

写和读

写缓存时，scale 作为参数传给 kernel：vLLM 的 FlashAttention 后端调 reshape_and_cache_flash(..., kv_cache_dtype, layer._k_scale, layer._v_scale)；读时把同一个 scale 扩成每个序列、每个 KV 头一份，作为 k_descale / v_descale 交给注意力 kernel。SGLang 的普通 MHA 池在 Python 里先除再转：

python/sglang/srt/mem_cache/memory_pool.py

        if cache_k.dtype != self.dtype:
            if k_scale is not None:
                cache_k.div_(k_scale)
            if v_scale is not None:
                cache_v.div_(v_scale)
            cache_k = cache_k.to(self.dtype)
            cache_v = cache_v.to(self.dtype)

        if self.store_dtype != self.dtype:
            cache_k = cache_k.view(self.store_dtype)
            cache_v = cache_v.view(self.store_dtype)

flowchart LR
  ckpt["checkpoint 的 k_scale / v_scale"] --> layer["注意力层：_k_scale、_v_scale（每层一对标量）"]
  json["SGLang：--quantization-param-path"] --> layer
  dflt["都没有：1.0 加警告"] --> layer
  kv["新算的 K、V（bf16）"] --> q["除以 scale，转 FP8"]
  layer --> q
  q --> cache["KV 缓存（uint8 存 FP8 位模式）"]
  cache --> attn["注意力 kernel：按块表读，乘 descale"]
  layer --> attn

scale 和数据放在一起：vLLM 的 per-token-head 与 NVFP4

per-tensor scale 是离线定的，一层只有一个数。vLLM 另有三种「per-token-head」模式（int8_per_token_head、fp8_per_token_head、int4_per_token_head），每个 token、每个 KV 头在写缓存时动态算一个 scale。这时层上的 scale 被固定成 1.0（kv_cache.py L100），真正的 scale 是 fp32，嵌在页里每个头的数据后面：

vllm/v1/attention/backends/triton_attn.py

    def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
        """Per-token-head modes pack inline fp32 scales after each head's
        data, so the content is (data + one scale) per K/V side."""
        # ...
        hs_k, hs_v = spec.head_size, spec.head_size_v
        if mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
            hs_k, hs_v = hs_k // 2, hs_v // 2
        scale_bytes = get_dtype_size(torch.float32)
        content = (hs_k + hs_v) * get_dtype_size(spec.dtype) + 2 * scale_bytes
        return replace(spec, state_content_bytes=content)

NVFP4 也是同一页里放数据和 scale：每个头一行，前面是 head_size / 2 字节的 4 bit 数据，后面是 head_size / 16 字节的 FP8 块 scale（每 16 个元素一个，nvfp4_kv_cache_full_dim）。

SGLang 的 NVFP4 是另一种布局：数据缓冲 (槽位, 头, head_dim / 2)，块 scale 单独一个缓冲 (槽位, 头, head_dim / 16)，外加每层一个 fp32 全局 scale（NVFP4KVCacheMethod、create_buffers）。这两种 FP4 方法由 --kv-cache-dtype 的取值在一个登记表里选（L1018），池子通过 _create_quantized_buffers 让方法自己决定数据和 scale 的形状。

每个 token、每个 KV 头占多少字节

按上面的源码公式，head_size = 128 时（K 和 V 合计）：

这是存储公式算出的字节数，不是实测。实际能多放多少块，还要看页是否按对齐要求补齐、激活和 CUDA graph 占多少，待实测（C3-1）；对吞吐和输出质量的影响待实测（C3-5）。

小结

• 开关都是 --kv-cache-dtype，默认 auto：checkpoint 声明了 FP8 KV 就用 FP8，否则用模型 dtype。
• FP8 per-tensor 两边做法相同：数据用 uint8 存，每层一对标量 scale 挂在注意力层上，写时除、读时乘。scale 来自 checkpoint（SGLang 还可以来自 JSON），缺了就是 1.0 加一条警告。
• vLLM 的 per-token-head 模式在写缓存时动态算 scale，fp32 scale 嵌在页里；vLLM 的 NVFP4 也是页内布局。SGLang 的 NVFP4 把块 scale 放在独立缓冲，另有每层全局 scale。
• FP8 把每个 token 的 KV 字节数减半，4 bit 格式约减到 0.28。效果待实测（C3-5）。