6.8 计算示例 · Infra Notes

KV Cache 量化到 INT8 的存储收益（head_dim = 128，per-(token, KV head) 对称量化，scale 存 FP16）：

BF16 : 128 × 2 bytes                = 256 B
INT8 : 128 × 1 byte + 1 FP16 scale  = 130 B
比值 : 130 / 256                     ≈ 50.8%

码 6-1 · KV Cache INT8 量化的 store kernel 片段

@triton.jit
def store_kvcache_int8(k_ptr, k_cache_ptr, k_scale_ptr, slot_ptr,
                       BLOCK_D: tl.constexpr):
    token, head = tl.program_id(0), tl.program_id(1)   # 一个 program = 一个 head vector
    offs = tl.arange(0, BLOCK_D)
    slot = tl.load(slot_ptr + token)                   # 逻辑位置 → 物理槽位（见码 3-2）

    key = tl.load(k_ptr + token * BLOCK_D + offs)
    # 除以 127 而不是 128：最大元素恰好映射到 ±127，加 0.5 后向零截断仍回 127
    # → 永不溢出，整个 kernel 不需要 clamp
    k_scale = tl.maximum(tl.max(tl.abs(key), axis=0) / 127.0, 1.0e-8)
    qk = key / k_scale
    # round-half-away-from-zero：先 ±0.5 再向零截断，正负号对称
    qk = tl.where(qk >= 0, qk + 0.5, qk - 0.5).to(tl.int8)

    tl.store(k_cache_ptr + slot * BLOCK_D + offs, qk)
    tl.store(k_scale_ptr + slot, k_scale.to(tl.float16))   # scale 单独存，粒度 = per (token, KV head)

两个设计选择值得记住： ① /127 而不是 /128：最大元素映射到 ±127，+0.5 后是 127.5，向零截断回 127，因此永不溢出、无需 clamp，省掉整个 kernel 的一遍逐元素判断。 ② 舍入是 round-half-away-from-zero，不是 PyTorch round() 的 half-to-even。写法 where(q >= 0, q + 0.5, q - 0.5) 再向零截断，对正负号处理一致。

存储账（head_dim = 128）：BF16 占 128 × 2 = 256 B，INT8 占 128 × 1 + 1 × 2 = 130 B，比值 50.78%——注意 scale 也是要占空间的，算收益时不能漏掉它。

一个容易被问倒的表述细节：KV Cache 通常按「剩余显存」定容，所以量化后显存占用总量几乎不变，涨的是可用块数/上下文容量（约 1.95×）。要准确表述为「用吞吐换上下文容量」，而不是「省了显存」。