TriAxialKV: Toward Extreme Low-Precision KV-Cache Quantization for Agentic Inference Tasks¶

会议: NeurIPS 2026
arXiv: 2605.17170
论文: arXiv 2605.17170
领域: 模型压缩
关键词: KV Cache量化, Agent推理, 混合精度量化, 语义结构敏感度, 融合解码算子

一句话总结¶

针对 Agent 复杂长上下文交互中单维度量化破坏关键 Token 表征的瓶颈，TriAxialKV 提出联合建模时序、模态与语义角色的三轴混合精度量化体系，在 Qwen3-VL 运行 OSWorld 时无损维持 BF16 精度的同时容纳 4.5 倍 KV Cache 并实现最高 1.52 倍吞吐加速。

研究背景与动机¶

随着大语言模型从单轮问答演进到计算机操作（Computer Use）、网络导航和复杂代码工程等自主智能体（Agent）工作流，模型推理面临前所未有的上下文膨胀挑战。在长周期多轮交互中，模型不仅需要自主拆解规划、调用工具，还需在上下文内持续累积环境反馈（如截图观察、执行输出）与长链条推理思维（Chain-of-Thought）。当在 OSWorld 这类基准上运行 LLaMA-3-70B 或 Qwen3-VL-32B 时，上下文长度常突破 100K Token，在 FP16 精度下仅单个批次的 KV Cache 就吞噬超过 30 GB 显存，使高带宽显存（HBM）容量成为限制系统吞吐和并发批大小的绝对瓶颈。

现有 KV Cache 压缩方案主要聚焦于均匀低比特量化（如 KIVI、FP4）、注意力稀疏驱逐或单一维度的混合精度策略（例如仅依据时序衰减降比特的 PM-KVQ，或仅区分图文模态的 VL-Cache）。然而，这些方法严重忽视了 Agent 上下文内在的高度异构性：系统提示与工具参数定义、不同轮次的对话历史、图像输入、思考过程与外部环境观察，其量化误差对最终注意力和决策精度的敏感度相差超过一个数量级。若单纯采用同质化极低精度（如统一切入 2-bit），像工具 Schema 和系统指令这类高度脆弱的结构元数据将被破坏，引发参数名拼写错误并导致任务直接失败；而只沿单一时序或模态维度压缩，又无法区分同轮次或同模态内截然不同的语义敏感度差异。

本文的研究切入点是突破单维度压缩假设，联合建模 Token 在交互系统中所处的三维结构特性。研究团队揭示出时序远近、模态来源与语义角色这三个正交维度几乎完全解释了上下文 Token 的量化敏感度方差。核心 idea：通过纯 Chat-Template 正则快速打标，将 Agent 提示切分为“时序-模态-语义”三维正交标签空间，在离线注意力失真表驱动下按边际误差增益贪心分配 INT2/INT4 位宽，并结合统一虚拟地址的双精度内存池与 Flash-Decoding 融合解码算子，实现兼具高保真度与硬件级吞吐收益的极端低精度 KV Cache 压缩。

方法详解¶

整体框架¶

TriAxialKV 包含一套端到端协同的混合精度量化与推理服务流水线。在离线校准阶段，框架从目标工作负载中抽取小规模提示，通过前向钩子采集各层注意力输入，计算逐个三维标签量化为 2-bit 或 4-bit 时相对于 FP16 真实注意力输出的均方误差（MSE），构建全局敏感度查询表。在在线服务调度阶段，请求预处理器无需模型推理即可完成上下文单遍打标，查表确定各 Token 的 INT2 或 INT4 位宽，并将 Token 路由至混合精度内存池中紧凑存放；在解码阶段，系统通过自定义的 Triton 融合 Flash-Decoding 算子，按位宽均匀的分组并行完成在线反量化与注意力计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入请求预处理<br/>对话模板单遍解析标记"] --> B["三轴标签生成与位宽分配<br/>时序 × 模态 × 语义 正交联合空间"]
    B --> C["混合精度内存池管理<br/>统一虚拟地址空间 (INT2 / INT4 Paged)"]
    C --> D["分组非对称量化存储<br/>Key/Value 定制量化轴与残差兜底"]
    D --> E["双阶段页表重排与融合解码<br/>Flash-Decoding 分组无缝在线反量化"]
    E --> F["高吞吐低显存推理输出<br/>维持 BF16 精度 / 4.5× 容量"]

关键设计¶

1. 三轴结构化正交标签体系与敏捷打标：联合时序、模态与语义解构上下文

针对 Agent 上下文异构性，TriAxialKV 将 Token 空间划分为三个正交维度的笛卡尔积空间 $\mathcal{S} = \mathcal{A}_{\mathrm{temporal}} \times \mathcal{A}_{\mathrm{modal}} \times \mathcal{A}_{\mathrm{semantic}}$。其中时序轴 $\mathcal{A}_{\mathrm{temporal}}$ 划分为当前轮次、前一轮、前两轮及更早历史（$\mathtt{current}, \mathtt{turn\_m1}, \mathtt{turn\_m2}, \mathtt{older}$），以捕捉注意力衰减动态；模态轴 $\mathcal{A}_{\mathrm{modal}}$ 区分文本与视觉图像（$\mathtt{text}, \mathtt{image}$），反映视觉编码器与文本 Embedding 的统计方差差异；语义角色轴 $\mathcal{A}_{\mathrm{semantic}}$ 则划分出系统指令/工具规范（$\mathtt{inst}$）、用户提示（$\mathtt{user}$）、助手文本（$\mathtt{assistant}$）、思维链（$\mathtt{reasoning}$）、工具调用（$\mathtt{tool\_call}$）、环境观测（$\mathtt{obs}$）以及模板分隔符（$\mathtt{delim}$）等 7 种角色。所有标签信息完全由 Tokenizer 的特殊标记（如 <think>, <tool_call> 等边界标记）通过单遍正则直接提取，无需任何模型前向开销即可为每个 Token 绑定标签 $s(t) \in \mathcal{S}$。

2. 注意力输出失真驱动的敏感度校准与贪心位宽分配：精准保护脆弱关键信息

不同 Token 对量化误差的容忍度存在量级差异，且原始 KV 误差与注意力最终输出非线性相关。为了使优化目标贴合实际生成质量，TriAxialKV 采用一阶泰勒展开将整体注意力均方误差分解为各标签量化失真的加权叠加： $$ \min_{\mathbf{b}\in{2,4}^{|\mathcal{S}|}}\;\;\sum_{k\in\mathcal{S}}D_{k}(b_{k})\qquad\text{s.t.}\qquad\sum_{k\in\mathcal{S}}N_{k}\,b_{k}\;\leq\;B\sum_{k\in\mathcal{S}}N_{k} $$ 其中 $D_k(b_k)$ 为仅对标签 $k$ 采用 $b_k \in \{2, 4\}$ 比特量化时产生的注意力输出 MSE，通过离线校准集中跨注意力头取最大（防木桶短板）、跨层累加计算得到。分配算法首先将所有标签初始化为 2-bit，计算将标签 $k$ 从 2-bit 升级到 4-bit 所能带来的单位显存误差降幅： $$ \rho_{k}\;\triangleq\;\frac{D_{k}(2)-D_{k}(4)}{2\,N_{k}} $$ 算法按 $\rho_k$ 降序排列，优先将显存预算分配给敏感度最高、边际收益最大的标签（例如承载函数签名的 $\mathtt{inst}$ 模块），直至达到目标平均位宽 $B$。

3. 定制化混合精度缓存布局与残差兜底：兼顾通道敏感度与显存对齐

在物理量化细节上，系统统一采用组大小 $G=32$ 的非对称分组量化。针对 INT4，Key 与 Value 均采用 per-token 分组量化；而在极低精度的 INT2 路径下，由于 Key 向量极易受到个别通道异常高幅值特征的污染，系统对 INT2 Key 采用 per-channel 分组量化（在 $G$ 个 Token 维度共享缩放比例与零点），对 INT2 Value 则保持 per-token 分组量化。对于无法凑满 $G=32$ 个 Token 的尾部残差（Residual Tokens），TriAxialKV 创新性地摒弃了 KIVI 设置专用全精度缓存的臃肿方案，直接将残差 Token 无缝路由至 INT4 per-token 路径，极大精简了内存管理，消除了短序列请求下全精度残差缓冲区对显存收益的侵蚀。

4. 统一地址空间双池管理与 Flash-Decoding 融合解码算子：消除分支开销的硬件加速

为了避免混合精度调度带来跨池碎片化与控制流发散，服务系统将 INT2 与 INT4 物理内存池置于单一连续虚拟地址空间，仅通过预设的基址偏移即可单指令判定 Token 精度。在自回归解码前（Phase 1），系统对页表指针进行局部排序，使每个请求中所有的 INT2 Token 物理指针紧凑前置于 INT4 之前；而在实际解码计算时（Phase 2），自定义 Triton 算子遵循 Flash-Decoding 并行架构，将序列划分为多个物理切片（Splits）。指针重排保证了前序切片为纯净的 INT2 计算路径，后续切片为纯净的 INT4 路径，各切片内部无分支快速执行位移解包与反量化乘加，最后通过全局在线 Log-Sum-Exp 规约融合注意力输出，新生成的自回归 Token 则直接以 INT4 写入尾部。

实验关键数据¶

主实验¶

论文在两类主流复杂 Agent 评估基准上系统评测了 TriAxialKV：涵盖文本结构化工具调用的 BFCL Memory，以及多模态计算机图形界面操作的 OSWorld。对比基线包括全精度 SGLang BF16、均匀低精度 SGLang FP4 以及经典非对称 2-bit 量化方案 KIVI。

任务 / 评估模型	指标 (任务成功率)	SGLang BF16 (参考基线)	SGLang FP4	KIVI (2-bit)	TriAxialKV Mixed (本文)
BFCL Memory (Falcon3-10B-Instruct)	准确率 (%)	24.00	28.22	22.89	24.00 ($\pm 0.00$)
BFCL Memory (Qwen3-14B)	准确率 (%)	25.11	18.00	20.67	24.22 ($-0.89$)
BFCL Memory (Qwen3-32B)	准确率 (%)	25.78	20.22	21.56	25.11 ($-0.67$)
BFCL Memory (Qwen3-235B-A22B)	准确率 (%)	23.78	23.56	23.33	22.67 ($-1.11$)
OSWorld (Qwen3-VL-8B-Thinking)	成功率 (%)	34.02	29.52	34.61	35.78 ($+1.76$)
OSWorld (Qwen3-VL-32B-Thinking)	成功率 (%)	39.20	40.54	41.49	40.59 ($+1.39$)
OSWorld (InternVL3.5-38B)	成功率 (%)	12.22	12.13	12.26	12.77 ($+0.55$)

在端到端系统吞吐方面，在搭载 80GB HBM3 的 NVIDIA H100 GPU 上服务 Qwen3-VL-32B 时，TriAxialKV 相比 BF16 实现了 1.52× 的端到端吞吐提升；在单张 180GB HBM3e 的 NVIDIA B200 GPU 上，TriAxialKV 分别在 Qwen3-VL-8B 和 32B 上取得 1.26× 与 1.32× 的吞吐加速，支持的并发请求数由 BF16 的 29.79 跃升至 118.82（提升近 4 倍）。

消融实验¶

为了量化三维标签体系中各正交轴的作用及压缩位宽极限，研究人员在 BFCL Memory 基准上对 Qwen3-14B 和 Qwen3-32B 进行了剥离评测，并对平均位宽预算进行了敏感度扫描：

配置 / 剥离项	Qwen3-14B 准确率 (%)	Qwen3-32B 准确率 (%)	说明与影响分析
TriAxialKV Full (完整模型)	24.22	25.11	联合时序、模态与语义分配（$B \approx 2.7$ bits）
w/o Temporal (去除时序轴)	22.00	24.00	无法识别早期轮次衰减，精度下降 1.11 ~ 2.22 点
w/o Semantic (去除语义轴)	18.00	20.89	无法隔离脆弱的工具与指令，精度剧烈滑坡 4.22 ~ 6.22 点
预算扫描: $B = 2.7$ bits	24.22	-	处于稳定悬崖边界，保持接近全精度表现
预算扫描: $B = 2.6$ bits	19.56	-	显存压缩每增加 0.1 bit，准确率下降约 5 个百分点
预算扫描: $B = 2.5$ bits	16.22	-	过度压缩导致关键语义丢失，模型发生严重指令漂移

关键发现¶

语义轴占据支配地位：去除语义轴使 Qwen3-14B 准确率断崖式下跌 6.22 点，而去除时序轴仅下降 2.22 点。这表明系统提示与工具参数模式对量化极其敏感，必须由 4-bit 强力保护，而粗暴地按轮次切分无法捕捉这类跨轮次分布的元数据。
均匀低比特量化的不可靠性：SGLang FP4 在 Falcon3-10B 上虚高但在 Qwen3-14B 上暴跌 7.11 点，表明均匀浮点低比特极易与特定模型的权重分布发生病态共振；相反，TriAxialKV Mixed 在所有模型上稳定紧贴 BF16（最大偏差仅 1.1%）。
显存带宽瓶颈硬件收益更显著：在显存相对受限且访存瓶颈突出的 H100 上，KV Cache 压缩换来的并发扩展带来高达 1.52 倍吞吐；而在显存极其宽裕的 B200 上加速比为 1.26-1.32 倍。

亮点与洞察¶

将 Prompt 视为异质化结构数据而非同质文本张量：论文跳出将 KV Cache 视为单一序列张量的狭隘传统，敏锐捕捉到 Agent Prefill 具备类似编译器 AST 般的清晰语义元结构，并用无额外开销的 Chat Template 解析将其结构化提取，构思精巧高效。
残差转投 INT4 的系统级权衡：以往工作如 KIVI 为对齐残差保留 128 Token 的全精度缓冲区，导致短上下文显存优势荡然无存；TriAxialKV 果断将不足 32 Token 的残差转入 INT4 per-token 路径，以微不足道的离散误差换取了极简统一的页管理。
硬件友好的页表重排策略：在 Flash-Decoding 前通过重排让 INT2 物理切片集中在先、INT4 集中在后，使单 Kernel 内部无需复杂的动态精度分支判断，完美避开了 SIMT 线程发散和访存不对齐。

局限与展望¶

工作负载离线校准依赖：当前的敏感度失真表依赖特定任务（如 OSWorld/BFCL）5% 样本的离线前向校准，缺乏对完全未知全新分布的零样本泛化能力。
位宽选择空间受限：目前调度空间仅在 INT2 与 INT4 之间进行粗粒度二值决策，若能支持三比特（INT3）或非均匀浮点格式（如 FP2/FP3），有望在帕累托前沿上进一步挤压压缩冗余。
模板标准化强绑定：打标流程极度依赖标准化的特殊标记（Special Tokens）；若面对非标准自由对话格式或第三方工具链时，可能退化或需要人工重写打标规则。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从时序、模态与语义角色三轴正交解构 Agent KV 敏感度并构建混合量化理论。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 Qwen3/Falcon/InternVL 等 7 种主流尺寸模型，横跨文本工具调用与 GUI 真实环境，兼备精度与真实 GPU 吞吐验证。
写作质量: ⭐⭐⭐⭐⭐ 逻辑推导层次分明，公式节制且硬件系统实现细节翔实。
价值: ⭐⭐⭐⭐⭐ 为长上下文 Agent 大规模工程化部署提供了极具实用价值的显存压缩与系统加速范式。