用四元数旋转压缩大模型键值缓存,速度提升超4倍且保持精度
IsoQuant: Hardware-Aligned SO(4) Isoclinic Rotations for LLM KV Cache Compression
- 基于四元数的等斜旋转,分块处理降低计算复杂度
- 在128维下将计算量从2408降至512,速度提升4.7倍
- 适合需要高效量化的大模型推理场景,尤其适配现代硬件
正交特征去相关在低比特在线向量量化中有效,但密集随机正交变换带来$O(d^2)$的存储与计算开销。RotorQuant采用分块3D Clifford旋转降低代价,但其3D划分与现代硬件不匹配且局部混合有限。本文提出IsoQuant,基于四元数代数与SO(4)等斜分解的分块旋转框架。每个4维块表示为四元数,通过闭式变换$T(v)=q_L v \overline{q_R}$实现旋转。提出两种变体:IsoQuant-Full实现完整SO(4)旋转,IsoQuant-Fast仅保留一个等斜因子以降低成本;框架还支持轻量级2D特例。在$d=128$时,IsoQuant-Full将前向旋转成本从RotorQuant的约2,408 FMAs降至1,024,IsoQuant-Fast进一步降至512。在18组融合CUDA设置($d \in \{128,256,512\}$,比特宽度${2,3,4}$,FP16/FP32执行)中,平均核级加速达4.5×–4.7×,峰值超6×,重建均方误差相当。当前验证限于合成归一化向量的量化-反量化路径,端到端KV缓存评估尚待后续工作。
原文摘要 · Abstract (English)
Orthogonal feature decorrelation is effective for low-bit online vector quantization, but dense random orthogonal transforms incur prohibitive $O(d^2)$ storage and compute. RotorQuant reduces this cost with blockwise $3$D Clifford rotors, yet the resulting $3$D partition is poorly aligned with modern hardware and offers limited local mixing. We propose \textbf{IsoQuant}, a blockwise rotation framework based on quaternion algebra and the isoclinic decomposition of $SO(4)$. It represents each $4$D block as a quaternion and applies a closed-form transform $T(v)=q_L v \overline{q_R}$. This yields two main variants: \emph{IsoQuant-Full}, which realizes the full $SO(4)$ rotation, and \emph{IsoQuant-Fast}, which keeps only one isoclinic factor for lower cost; the framework also admits a lightweight $2$D special case. At $d=128$, IsoQuant-Full reduces forward rotation cost from about $2{,}408$ FMAs in RotorQuant to $1{,}024$, while IsoQuant-Fast further reduces it to $512$. Across $18$ fused CUDA settings with $d \in {128,256,512}$, bit widths ${2,3,4}$, and FP16/FP32 execution, IsoQuant achieves mean kernel-level speedups of about $4.5\times$--$4.7\times$ over RotorQuant while maintaining comparable reconstruction MSE, with peak speedups above $6\times$. Current validation is limited to the stage-1 quantize--dequantize path on synthetic normalized vectors; end-to-end KV-cache evaluation remains future work.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。