TriAxialKV：面向代理推理任务的极端低精度 KV 缓存量化
机器学习
2026-05-19 v1
摘要
代理工作负载已成为 LLM 推理的主要工作负载。它们与仅聊天的工作负载有显著不同，需要长上下文处理、能够处理多模态输入，以及结构化的多轮交互以支持 tool calling 功能。因此，其 context 具有可以在三个关键轴向上携带不同重要性的结构：时间上的 recency（当前轮次的时序新近性）、模态性（如文本或图像 token）以及语义角色（如用户查询、tool 调用、观察或推理）。这些轴向捕获了不同 token 行为，导致对 KV 缓存压缩有不同的灵敏度。然而，现有的 KV 缓存量化方法通常是均匀的，或仅利用单一维度的异构性，如时间接近性或模态性，忽略了它们之间的相互作用。为此，我们引入 TriAxialKV，一种新的混合精度 KV 缓存量化方案，为每个 token 分配 triaxial tag，校准 per-tag 灵敏度，并在固定内存预算下分配 INT2/INT4 位宽。我们将 TriAxialKV 实现为一个端到端的 serving 系统，包括 calibration（校准）、混合精度量化和 memory management（内存管理），以及自定义融合 Triton 解码 kernel。当使用 Qwen3-VL-32B-Thinking 作为运行 OSWorld 的 computer-use agent 时，TriAxialKV 在 KV 缓存为 BF16 时匹配 SGLang 的准确性，同时支持 4.5 倍的 KV 缓存大小，并在真实 GPU 系统上实现 30% 的端到端吞吐量提升。
引用
@article{arxiv.2605.17170,
title = {TriAxialKV: Toward Extreme Low-Precision KV-Cache Quantization for Agentic Inference Tasks},
author = {Hanzhang Shen and Haoran Wu and Yiren Zhao and Robert Mullins},
journal= {arXiv preprint arXiv:2605.17170},
year = {2026}
}