提出三轴量化方法,让大模型推理更省显存、更快。
TriAxialKV: Toward Extreme Low-Precision KV-Cache Quantization for Agentic Inference Tasks

- 按时间、模态、语义三维度给记忆块打标签,分层压缩
- 在相同精度下,显存用4.5倍,吞吐高30%
- 适合需要工具调用的复杂任务,如系统操作代理
智能体工作负载已成为大模型推理的主要场景,其特点包括长上下文处理、多模态输入和带工具调用的多轮结构化交互。这些特性使上下文在时间近似性、模态类型(如文本或图像)和语义角色(如用户提问、工具调用、观察结果、推理过程)三个维度上表现出不同的敏感度。现有KV缓存量化方法通常仅考虑单一维度的异质性,忽略多维交互。为此,我们提出TriAxialKV,一种新型混合精度量化方案,为每个令牌分配三轴标签,校准每类标签的敏感度,并在固定内存预算下分配INT2/INT4位宽。我们实现了一个端到端服务系统,包含校准、混合精度量化、内存管理及自定义Triton解码内核。在真实GPU系统上运行Qwen3-VL-32B-Thinking作为操作系统操作代理时,TriAxialKV在保持与SGLang BF16 KV缓存相当精度的同时,支持4.5×的缓存大小,实现30%更高的端到端吞吐量。
原文摘要 · Abstract (English)
Agentic workloads have emerged as a major workload for LLM inference. They differ significantly from chat-only workloads, requiring long-context processing, the ability to handle multimodal inputs, and structured multi-turn interactions with tool calling capabilities. As a result, their context exhibits structure that can carry different importance along three key axes: temporal recency to the current turn, modality such as text or image tokens, and semantic role such as user queries, tool calls, observations, or reasoning. These axes capture distinct token behaviors and lead to different sensitivities to KV-cache compression. However, existing KV-cache quantization methods are typically homogeneous or exploit only heterogeneity on a single dimension, such as temporal proximity or modality, overlooking the interactions among them. To this end, we introduce TriAxialKV, a novel mixed-precision KV-cache quantization scheme that assigns each token a triaxial tag, calibrates per-tag sensitivity, and allocates INT2/INT4 bitwidths under a fixed memory budget. We implement TriAxialKV as an end-to-end serving system, comprising calibration, mixed-precision quantization and memory management, and custom fused Triton decode kernels. When using Qwen3-VL-32B-Thinking as a computer-use agent operating the OSWorld, TriAxialKV matches the accuracy of SGLang with BF16 KV cache while supporting 4.5$\times$ KV cache size and achieving 30% higher end-to-end throughput, when running on real GPU systems.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。