用 GAAC 在本地做无损 KV 压缩：长上下文实测

长上下文推理中，KV cache 占用的显存会随序列长度线性增长；序列从 8K 增加到 32K 时，KV cache 通常会超过模型权重，成为单个请求的显存大头。本教程在本地用 GAAC 对 KV cache 做无损压缩，完整走通“接入、调参、测量”三步，最终得到一张显存、带宽、吞吐、输出一致性的对照表，帮助你判断在自己的机器上是否值得启用。

这篇能做出什么

完成后你会得到四个具体产物：

1. 一个能量化 KV cache 占用与理论下限的小脚本，用来判断“该不该压缩”。

2. 一个接入骨架，在不改动模型权重的前提下，把 KV 从“全量 fp16 常驻”换成“分块压缩、按需解压”。

3. 一组参数扫描结果：block 大小、压缩级别、是否保留近期明文窗口、是否只压 V 等组合下的显存节省率和吞吐变化。

4. 一套无损性验证流程：逐元素比对解压前后张量，而不是只看“输出看着差不多”。

先说清一个容易忽略的前提：无损压缩降低的是显存占用和 HBM 读取字节数，不会减少注意力计算量本身。decode 阶段本来就受显存带宽限制，KV 读取量下降才有可能带来吞吐提升；但这部分收益要扣掉压缩、解压的算力和 kernel 启动开销，小 batch、短上下文时常常净亏。是否启用，要靠实测决定。

前置条件清单

安装示意（具体包名与入口以仓库 README 为准）：

```bash

python -m venv .venv && source .venv/bin/activate

pip install --upgrade pip

pip install torch transformers numpy

git clone <GAAC 官方仓库地址> # 地址以项目官方页面为准

cd gaac && pip install -e .

python -c "import gaac; print(gaac.__file__)" # 确认能导入

```

如果最后一行报错，先解决安装，不要急着往下跑。

步骤一：先算清 KV 的显存账

动手前先算一遍，避免做完发现量级不对。对于常见的 GQA 模型，单个请求的 KV cache 字节数为：

```

bytes = 2 × layers × kv_heads × head_dim × seq_len × batch × dtype_bytes

```

其中 2 代表 K 和 V 各一份。举例：32 层、8 个 KV 头、head_dim 128、fp16（2 字节）、序列 32768：

• 每 token：2 × 32 × 8 × 128 = 65536 个元素
• 乘上 32768 个 token，得到 2^31 个元素
• 乘 2 字节，恰好 4 GiB

也就是说，仅仅一条 32K 的请求就吃掉 4 GiB，batch=4 就是 16 GiB。这就是长上下文在个人设备上难以运行的根本原因。

把这段写成脚本，方便随时改参数：

```python

def kv_gib(layers, kv_heads, head_dim, seq_len, batch=1, dtype_bytes=2):

elements = 2 * layers * kv_heads * head_dim * seq_len * batch

return elements * dtype_bytes / 2**30

print(kv_gib(32, 8, 128, 32768)) # 4.0

print(kv_gib(32, 8, 128, 32768, 4)) # 16.0

```

字段名在模型的 config.json 里：num_hidden_layers、num_key_value_heads、hidden_size / num_attention_heads。

步骤二：跑通不压缩的基线

任何对比都必须有基线。下面这段脚本输出三个数：首 token 延迟（TTFT）、decode 吞吐（token/s）、峰值显存（GiB）。用贪心解码，保证可复现。

```python

import os, time, torch

from transformers import AutoModelForCausalLM, AutoAI 词典：Token">Tokenizer

MODEL = os.environ["MODEL_PATH"]

tok = AutoTokenizer.from_pretrained(MODEL)

model = AutoModelForCausalLM.from_pretrained(

MODEL, torch_dtype=torch.float16, device_map="cuda"

).eval()

def build_ids(n_tokens: int):

filler = "请逐条解释下面这段运维手册的操作步骤，并指出每一步可能的失败点。\n"

s = ""

while len(tok(s).input_ids) < n_tokens:

s += filler

return tok(s, return_tensors="pt").input_ids[:, :n_tokens].cuda()

@torch.no_grad()

def bench(ids, new_tokens=128):

torch.cuda.empty_cache()

torch.cuda.reset_peak_memory_stats()

t0 = time.perf_counter()

out = model(ids, use_cache=True)

torch.cuda.synchronize()

ttft = time.perf_counter() - t0

past = out.past_key_values

nxt = out.logits[:, -1:].argmax(-1)

del out

t1 = time.perf_counter()

for _ in range(new_tokens - 1):

out = model(nxt, past_key_values=past, use_cache=True)

past = out.past_key_values

nxt = out.logits[:, -1:].argmax(-1)

del out

torch.cuda.synchronize()

dt = time.perf_counter() - t1

return {

"ttft_s": round(ttft, 3),

"decode_tok_s": round((new_tokens - 1) / dt, 2),

"peak_gib": round(torch.cuda.max_memory_allocated() / 2**30, 2),

}

for n in (8192, 16384, 32768):

print(n, bench(build_ids(n)))

```

把这三行结果记下来，后面每次改配置都对同一组长度、同一段 prompt、同样的 new_tokens 重跑。

步骤三：把 GAAC 接进推理栈

无损 KV 压缩有三个可接入的层次，按改动成本从低到高排列：

1. 缓存对象层：替换 transformers 里传给 past_key_values 的 Cache 对象。改动最小，适合先验证。

2. 注意力后端层：在推理引擎（vLLM 等）的 attention backend 里对 block 做压缩。性能最好，但需要改引擎代码。

3. 离线层：把长文档的 KV 预先算好、压缩落盘，用时加载。适合固定知识库场景。

先做第 1 层。核心是包一层编解码器，把“写满的块”压掉，只留一个尾部明文块。GAAC 实际提供的编解码入口名称与参数以仓库文档为准，这里给出的是适配骨架：

```python

import torch

class KVCodec:

"""把 GAAC 的无损编解码包成 encode/decode 两个方法。"""

def __init__(self, level: int = 3, block_tokens: int = 256):

self.level = level

self.block_tokens = block_tokens

def encode(self, tensor: torch.Tensor):

输入: [batch, heads, tokens, dim] 的连续 fp16 张量

调用 GAAC 的压缩函数，返回可存的字节串或张量；以官方 API 为准

raw = tensor.contiguous().view(torch.uint8).cpu().numpy().tobytes()

return gaac.compress(raw, level=self.level) # 名称以仓库为准

def decode(self, blob, shape, dtype=torch.float16):

raw = gaac.decompress(blob) # 名称以仓库为准

buf = torch.frombuffer(bytearray(raw), dtype=torch.uint8)

return buf.view(dtype).view(shape)

class BlockCompressedCache:

"""最小骨架：明文 tail + 已压缩的历史块。"""

def __init__(self, codec: KVCodec):

self.codec = codec

self.plain = {} # layer_idx -> (k, v)，未满一块的尾部

self.blocks = {} # layer_idx -> list[(k_blob, v_blob, shape)]

def append(self, layer_idx, k, v):

pk, pv = self.plain.get(layer_idx, (None, None))

k = k if pk is None else torch.cat([pk, k], dim=-2)

v = v if pv is None else torch.cat([pv, v], dim=-2)

n = self.codec.block_tokens

if k.shape[-2] >= n:

cut = (k.shape[-2] // n) * n

head_k, head_v = k[..., :cut, :], v[..., :cut, :]

self.blocks.setdefault(layer_idx, []).append(

(self.codec.encode(head_k), self.codec.encode(head_v), head_k.shape)

)

k, v = k[..., cut:, :], v[..., cut:, :]

self.plain[layer_idx] = (k, v)

def read_all(self, layer_idx):

parts_k, parts_v = [], []

for kb, vb, shape in self.blocks.get(layer_idx, []):

parts_k.append(self.codec.decode(kb, shape))

parts_v.append(self.codec.decode(vb, shape))

pk, pv = self.plain.get(layer_idx, (None, None))

if pk is not None:

parts_k.append(pk); parts_v.append(pv)

return torch.cat(parts_k, dim=-2), torch.cat(parts_v, dim=-2)

```

接进 transformers 时，把它接到模型读取 KV 的位置；缓存接口的方法名在不同版本间有变化，以官方文档当前版本的 Cache 定义为准。如果暂时不想改引擎，也可以先只用 encode、decode 做离线验证：把一段长 prompt 跑完，取出 KV，压缩再解压，逐元素比对，先确认“无损”成立。

步骤四：压缩率调参

通常需要扫的四个维度：

• block_tokens：每块包含多少 token。太小则每块的头信息占比高、kernel 调用多；太大则尾部明文窗口浪费显存、首次解压延迟高。从 128、256、512 开始扫。
• level / 压缩级别：级别越高通常压缩率越好、CPU/GPU 开销越大。无损场景下收益会很快饱和，扫 1、3、5 这类小集合就够。
• 压缩范围：只压 K、只压 V、两者都压。有的模型 V 的数值分布更集中，可压缩性更好，但必须自己测。
• 近期窗口：最近 N 个 token 保持明文（示例里就是那个 plain 尾部）。N 越大，解码越省事，省下的显存越少。

扫描脚本的关键是“只改一个变量并做完整重跑”，别在一次运行里改两个参数：

```python

import itertools, json

results = []

for block, level in itertools.product((128, 256, 512), (1, 3, 5)):

codec = KVCodec(level=level, block_tokens=block)

for n in (8192, 16384, 32768):

r = bench(build_ids(n), new_tokens=128) # 换成挂了压缩缓存的版本

r.update({"block": block, "level": level, "seq": n})

results.append(r)

print(r)

json.dump(results, open("sweep.json", "w"), ensure_ascii=False, indent=2)

```

判断标准只有一条：在吞吐可接受的配置里，挑显存节省最多的那组。如果所有配置的 decode 吞吐都低于基线，说明当前负载（小 batch、短序列）不适合压缩，直接放弃是最合理的选择。

步骤五：验证“无损”

这一步不能省，而且不能只看输出像不像。三层验证，从强到弱：

第一层，张量级比对（判据）：解压后的 KV 与原始 KV 必须逐元素完全相等。

```python

k_orig = torch.randn(1, 8, 1024, 128, dtype=torch.float16, device="cuda")

codec = KVCodec(level=3, block_tokens=256)

blob = codec.encode(k_orig)

k_back = codec.decode(blob, k_orig.shape).to(k_orig.device)

print(torch.equal(k_orig, k_back)) # 必须为 True

```

对真实 KV 也要做一遍：跑一次 prefill 拿到 past_key_values，对每一层、K 和 V 都做上述比对。只要有一层为 False，就说明这个配置不是无损的，例如内部悄悄走了量化或截断，需要回退参数或换级别。

第二层，logits 比对：在同一个后端、同一个 batch 配置下比较最后一位 logits。由于解压后的 kernel 可能走不同的分块方式，浮点归约顺序会有差异，通常会出现极小量级的差；如果差到肉眼可见，那就不是无损问题了。

第三层，输出比对：固定 prompt、贪心解码、do_sample=False，比较生成的 token id 序列。这一层只能作为参考，不能作为无损的证明，因为小误差完全可能不改变 argmax。

步骤六：测带宽与吞吐

decode 阶段每生成一个 token，都要把整个 KV 读一遍。所以有一个粗略的理论下限：

```

每步耗时下限 ≈ KV 字节数 / 显存带宽

```

显存带宽取你显卡官方规格里的 HBM 带宽（不同卡差异很大，以厂商规格为准）。用这个下限可以判断“还有多少优化空间”：比如 4 GiB 的 KV、带宽 N GB/s，理论下限就是 4 × 1024 / N 秒。

实测侧，三个数最有价值：

多 batch 时，把 batch 尺寸也扫一遍（1、2、4、8）。常见的现象是：batch=1 时压缩是净亏，batch 上去以后因为 KV 读取被摊薄，压缩才开始有正收益，但收益也会被压缩开销吃掉一部分。这个拐点在哪，只有你自己的机器能回答。

常见坑与排错

显存看起来没降。 nvidia-smi 显示的是 PyTorch 缓存分配器持有的总量，压缩后旧张量释放了，但 block 未必归还驱动。看 torch.cuda.max_memory_allocated()（真实用量）和 memory_reserved()（含缓存）两个数，必要时在每轮之间 torch.cuda.empty_cache()，但注意它会拖慢测量，标记清楚哪些轮次调用了它。

OOM 发生在 prefill 而不是 decode。 prefill 本身要算 attention 矩阵，激活值峰值可能远大于 KV。此时压缩 KV 帮助有限，先降 max_batch_tokens 或分段 prefill。

吞吐反而下降。 三个常见原因：block 太小导致 kernel 启动次数暴涨；压缩在同步路径上阻塞了下一个 token；只压了很少的 token（上下文本来就短）。对策分别是调大 block、改成异步或流水线、直接不用。

输出和基线不一致。 先确认不是非确定性引起的：attention 后端的分块策略、CUDA Graph、多流并行都会改变浮点归约顺序。对比时锁死同一后端、同一 batch 配置、同一次进程运行。再做步骤五的逐元素比对，定位到底是压缩环节还是算子环节。

块边界对齐出错。 拼接时 dim=-2 用错（拼到 head 维度上），或者最后一块不足 block_tokens 时被当成完整块处理，会得到形状正确但内容错位的结果。形状检查能通过，但逐元素比对会立刻暴露。

prefix cache 被重复压缩。 多轮对话或共享前缀场景下，同一段 KV 可能被压多次、解压多次。在缓存层加一个“已压缩块”的标记，读取时跳过。

版本不匹配。 transformers 的 Cache 接口在版本迭代中调整过方法签名。报 unexpected keyword argument 一类错误时，第一件事是核对官方文档当前版本的接口定义，不要照搬旧博客的代码。

下一步建议

• 叠加而不是替代：无损压缩和权重量化解决的是不同问题（KV 显存 vs 权重显存），两者可以同时用。先做好无损压缩，再在它之上评估量化，能够更清楚地评估各手段的贡献。
• 把前缀缓存写入磁盘：长系统提示词、固定知识库这类重复内容，压缩后写入磁盘，第二次请求直接加载，省掉一整次 prefill。
• 关注多序列共享：multi-turn 与 agent 场景里，多条请求共享同一段长前缀。压缩后的前缀块天然适合在请求间共享，这也是显存节省容易被低估的地方。
• 建立自己的基准集：准备 5~10 条覆盖真实业务的长 prompt（代码库问答、长文档摘要、日志排查），固定下来，每次改配置都跑同一组。比起追逐某个“通用压缩率”，这套基准能告诉你配置在这个业务上到底值不值。

最后提醒一句：无损压缩的收益高度依赖数据分布、硬件带宽和批处理形态，任何来自其他机器或其他模型的数字都只能当参考。把上面六个步骤跑完，你手上的那张表才是决策依据。