Vramkit Vibe 项目

项目简介

Vramkit 是一个针对大语言模型推理阶段显存占用做优化的开源工具集，它盯上了推理时吃掉显存的"两座大山"——模型权重与键值缓存（KV Cache），并为这两者分别给出了一套可复现的压缩方案。项目最特别的地方在于它的验证方式：全部代码与基准测试都在作者的一台 MacBook（Apple Silicon 芯片、纯 CPU 版 PyTorch）上开发并测量完成，不依赖数据中心级 GPU，普通开发者就能在自己机器上复现出文档里写的每一个数字。权重部分采用从零手写的 32 位区间编码器配合二维自适应上下文模型，官方宣称在真实 Qwen 权重量化数据上比 zstd、xz 更省且完全无损；KV 缓存部分则给出了一个反直觉却经真实困惑度检验的结论：键应该按通道量化，值才适合按 token 量化。项目采用 MIT 协议开源，目前仍属可运行的研究原型，官方也坦承 GPU 解量化内核与更大模型的验证尚未完成。

核心功能

1. 无损权重编解码器：自研 32 位区间编码器加二维自适应上下文模型，每个 4 位权重都以它左邻居与上邻居为条件建模，能看到字节序压缩器（如 zstd、xz）看不到的结构性关联。
2. 按通道量化的键值缓存压缩：这是项目最核心的发现。键向量里存在少量持续存在的通道级离群值，用按 token 的缩放系数会让这些离群值压垮其余全部数值；改成按通道（每个特征维度一个系数，跨 token 共享）即可修正。
3. 可复现的真实模型基准脚本：仓库自带困惑度验证脚本与真实 safetensors 权重验证脚本，所有质量结论都建立在真实模型的端到端困惑度上，而非余弦相似度之类的代理指标。
4. 设备无关的即插即用模块：核心模块是纯 PyTorch 实现，可运行在 AMD ROCm、NVIDIA、苹果 MPS 与纯 CPU 环境，不绑定任何特定厂商的算子栈。
5. 坦诚的自我修正记录：项目早期版本曾用余弦相似度和重复性语料报告 KV 质量，官方在文档中主动披露并纠正了这两处问题，把修正过程与结论一并公开。

技术实现

权重压缩的核心是一个完全自研的 32 位区间编码器，配合二维自适应上下文模型——每个 4 位权重的取值都以它左边和上边的邻居为条件进行概率建模，从而捕获量化权重张量在空间上的局部相关性，这是纯字节序压缩器结构上无法看到的。在真实 Qwen2.5-0.5B 权重的 self_attn.o_proj（int4）上实测：原始 4 位打包为每权重 4.000 位，zstd -19 为 2.717 位（较原始下降 32.1%），xz -9e 为 2.708 位（下降 32.3%），而 vramkit 的二维上下文模型为 2.676 位（下降 33.1%），且经逐位比对确认完全无损；在该模型全部权重张量类型上平均为每权重 2.92 位、较原始 4 位下降 27%，其中 4 类张量优于 zstd-19。KV 缓存侧则通过前向钩子把压缩施加到每一层的键与值上，在四段互不重叠的真实散文窗口上测真实困惑度（基线约 18）：按 token 量化键值的 int8 方案在 2.0 倍压缩下困惑度上升 4.6%，而按通道量化键、按 token 量化值的 int8 方案在同样 2.0 倍压缩下仅上升 0.10%（键在旋转位置编码前）与 0.13%（标准缓存中存储的旋转后键），基本无损；键改用 int4 时无论哪种方案都会明显劣化（上升 25.6% 以上），说明 int4 只适合用来压值。项目用纯 PyTorch 与 numpy 实现，不依赖任何定制算子，可直接嵌入 ROCm 技术栈；官方特别指出，长上下文推理在 ROCm 上超过约 64K 后注意力路径会明显退化，而 FP8 的 KV 缓存支持仍在完善中，因此把 KV 缓存体积压下来是提升每张加速卡并发请求数与上下文长度的直接杠杆。流式解码场景下，由于按通道的缩放系数需要 token 轴信息，作者采用类似 KIVI 的做法：按最近 G 个 token 分组量化键，最新的少数 token 保留半精度。

快速上手

1. 克隆仓库：git clone https://github.com/SRKRZ23/vramkit，项目不依赖数据中心 GPU，MacBook 上即可完整运行。
2. 运行自带自检：执行 python3 vramkit/kvquant.py，脚本会直接验证"在带离群值的键上按通道量化优于按 token 量化"这一核心结论是否成立。
3. 接入压缩与还原：在代码中引入 from vramkit.kvquant import quantize_kv, dequantize_kv，对键执行 quantize_kv(key, bits=8, per="channel")、对值执行 quantize_kv(value, bits=8, per="token")，在注意力计算前用 dequantize_kv 还原即可。
4. 验证真实模型困惑度：用 python3 bench/validate_ppl_mac.py 在 Qwen2.5-0.5B 上测出压缩带来的困惑度变化，用 python3 bench/validate_real_weights.py 在真实权重上验证权重编解码器的体积收益与逐位一致性。
5. 对照实验：用 bench/bench_weights.py 与 bench/bench_kvcache.py 在合成数据上跑出理论熵底噪，与真实权重结果对照，判断某个收益是否只是单一张量的偶然表现。

适用人群

正在为大模型推理服务做成本与并发优化的算法与基础设施工程师；受限于显存无法把更大上下文或更多并发请求塞进同一张卡的推理服务团队；希望在 AMD ROCm 或苹果 MPS 等非主流硬件栈上验证 KV 缓存压缩收益的研究者；以及希望有一份数字全部可复现、且愿意读原始基准脚本而不是只看结论摘要的推理方向工程师。

官网链接

源码仓库：https://github.com/SRKRZ23/vramkit ｜ 权重编解码器：vramkit/rangecoder.py ｜ 键值缓存压缩模块：vramkit/kvquant.py ｜ 真实困惑度基准：bench/validate_ppl_mac.py ｜ 协议：MIT ｜ 状态：GPU 解量化内核与 7B 以上模型验证为官方列出的下一步待办