RotaryQuant: Fitting 120B MoE Models on Consumer Hardware via Fused Compressed-Space Attention¶
会议: NeurIPS 2026
arXiv: 2608.08081
领域: 模型压缩
关键词: MoE模型推理、KV Cache量化、结构化旋转、算子融合、端侧部署
一句话总结¶
RotaryQuant 提出三轴协同压缩系统(角色化权重量化、LRU 专家卸载以及基于 WHT 与 SO(4) 块旋转的 IsoQuant 3-bit KV 缓存压缩),通过 Metal 融合算子实现直接在 3-bit 压缩域计算注意力,在无需物化 FP16 张量的前提下将 120B MoE 模型装入 32GB 消费级硬件并实现交互式流畅推理。
研究背景与动机¶
百亿至千亿参数级稀疏混合专家(MoE)架构通过稀疏激活将前向计算成本维持在可处理区间,但将 26B–120B 的大规模 MoE 部署至仅具备 16–32 GB 统一内存的消费级终端时,依然面临三维并发内存重压:常驻的密集与专家权重矩阵、随上下文长度线性暴涨的键值(KV)缓存状态,以及数十个需要按需调度的专家子层。以 30B MoE 为例,即使在 4-bit 量化下模型权重已耗费约 15 GB 显存,一段 8K 上下文又将额外引入 4–8 GB 的 KV 缓存,瞬间突破消费级终端的物理内存极限。
现有 KV 缓存量化方法(如 KIVI、KVQuant、TurboQuant 等)普遍遵循“先反量化重建再计算”(reconstruct-then-compute)的传统范式:保存在显存中的低比特压缩向量必须先解压还原为 FP16 浮点张量,再调用标准密集 GEMM 算子进行注意力计算。这种瞬时物化(materialization)机制在解码期间大幅抵消了显存压缩红利,并引入与序列长度严格成正比的显存带宽瓶颈。此外,基于旋转的消除激活离群点方法多采用密集随机或可学习正交矩阵,带来高达 \(\mathcal{O}(d^2)\) 的乘加计算复杂度与不可忽略的旋转矩阵参数存储开销。
本文旨在彻底推翻物化计算流程,转而建立直接在压缩张量上执行计算的底层执行模型(compute-in-compressed-space)。核心 idea:通过 Walsh–Hadamard 变换与分块四元数 SO(4) 旋转的乘积构造 \(\mathcal{O}(d \log d)\) 极轻量各向同性变换(IsoQuant),彻底消除离群特征后实施 3-bit 标量量化,并配合四内核融合 Metal 算子直接在打包数据上完成注意力内积与加权累加,仅在最终输出时执行单次逆旋转。
方法详解¶
整体框架¶
RotaryQuant 针对消费级硬件上的 MoE 推理设计了三轴独立但级联协同的压缩体系:以非交换形式依次编排权重量化(\(Q_W\))、KV 缓存压缩(\(C_{KV}\))与专家按需卸载(\(O_E\))。在推理执行侧,核心引擎 IsoQuant 彻底重塑了解码流:输入 KV 向量在预填充阶段维持 FP16 精度以避免误差跨位置级联(延迟预填充),而在解码阶段自回归生成的 KV 向量首先经由模长归一化与快速正交结构化旋转变换至各向同性空间,通过 Lloyd–Max 最优码本量化并打包为 3-bit 紧凑张量。
在解码注意力计算时,Metal 定制流水线调度四个融合内核:直接在 3-bit 打包数据上执行 Query 与 Key 的点积,过 Softmax 后直接在 3-bit 打包 Value 上执行加权累加,最后仅对聚合得到的单个输出向量施加一次逆旋转变换,全程完全不物化全精度中间张量。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入查询与键值向量<br/>Query q, Key k, Value v"] --> B["三轴压缩协同调度<br/>角色化权重量化 + LRU专家卸载"]
B --> C["IsoQuant 结构化正交各向同性化<br/>WHT全局混叠 ∘ SO(4)块旋转"]
C --> D["Lloyd-Max 最优码本与 3-bit 打包<br/>d=128 压缩至 48 字节"]
D --> E["压缩域四内核融合注意力流水线<br/>直接计算打包QK点积与V累加"]
E --> F["单次输出逆旋转恢复<br/>仅对聚合输出执行 R^T 还原"]
关键设计¶
1. 三轴协同压缩体系:按架构角色解耦多维显存压力
消费级终端无法依靠单一手段解决全系统内存承载力问题,RotaryQuant 建立了严谨的三轴复合压缩算子 \(M_{\text{total}} = O_E \circ C_{KV} \circ Q_W\)。其核心机制在于认识到三者具有非交换性:权重矩阵位宽直接塑造了激活层输出的方差与峰度,而专家卸载边界则严格受制于权重与缓存压缩后剩余的显存空间。系统对不同角色子网络分配异构精度:密集层统一采用标准 4-bit(GPTQ/AWQ);路由稀疏专家(routed experts)被激进压缩至 2-bit;而针对全局共享专家(shared expert),由于其激活峰度高达 \(\kappa=10.10\)(远高于专门专家的 \(\kappa=0.41\)),若强行低比特化将引发剧烈性能崩溃,因而特意为其保留 Q8_0(8-bit)精度以锚定模型表达底座。对于剩余超出显存配额的专家层,引入 LRU 专家卸载机制在真实显存承压时换入换出至固态硬盘。
2. IsoQuant 结构化旋转变换:各向同性化与 O(d log d) 超低开销
高维激活向量中少数异常特征维度(outliers)是低比特量化精度雪崩的主因。以往方法(如 TurboQuant)使用密集正交矩阵实现各向同性旋转,但在头维度 \(d=128\) 下需存储 \(128 \times 128 = 16,384\) 个浮点参数,每次旋转消耗 16,384 次 FMA 乘加运算。IsoQuant 提出将正交变换严格分解为全局与局部的复合算子 \(R = R_{SO(4)} \circ H_d\)。其中 \(H_d\) 为归一化 Walsh–Hadamard 变换(WHT),利用蝶式算法以 \(\mathcal{O}(d \log d)\)(在 \(d=128\) 下仅需 896 次 FMA)实现全局能量均匀扩散,使各坐标方差收敛至 \(1/d\);紧随其后的 \(R_{SO(4)}\) 对相邻 4 维连续块施加独立紧凑旋转,采用双四元数配对表示 \(v \mapsto q_L v \bar{q}_R\)(\(q_L, q_R \in S^3\)),每块严格覆盖完整的 6 个旋转自由度(单四元数共轭仅覆盖 3 个自由度的 SO(3),会导致 25% 维度完全不旋转)。该分块结构仅需 \((d/4) \times 16 = 512\) 次 FMA,整套复合旋转仅需 1,408 次 FMA(比密集旋转计算量下降 11.6 倍),且每注意力头仅需存储 256 个参数(参数量骤减 64 倍)。经旋转后的向量通过 Hanson–Wright 不等式保证坐标指数级集中,随后送入 Lloyd–Max 标量码本紧凑打包为 3-bit(128 维仅耗 48 字节,实现 5.3 倍显存压缩)。
3. 压缩域四内核融合流水线:彻底消除张量物化与单次逆旋转机制
传统量化解码过程必须分配庞大的临时 FP16 缓冲区来展开全部历史 KV,不仅产生可观的分配与带宽开销,更无法利用高速片上缓存。IsoQuant 构建了四个定制 Metal 内核流水线(Kernel A: fused_qk_dot、Kernel B: softmax、Kernel C: fused_value_accum、Kernel D: rotate_inverse)。在计算过程中,Kernel A 与 Kernel C 直接加载 3-bit 整数打包张量并在片上寄存器完成即时点积与带权聚合。数学上,得益于 \(R\) 的正交严格保内积性:
$\(\langle Rq, Rk \rangle = q^\top R^\top R k = q^\top k\)$
在解码端将当前查询向量前向旋转 \(\tilde{q} = Rq\) 后,与所有已旋转并压缩的历史键值计算得到的注意力分布 \(\alpha_j\) 完全等价于全精度未旋转内积。聚合输出为:
$\(\tilde{o} = \sum_j \alpha_j \tilde{v}_j = R \sum_j \alpha_j v_j = R \cdot o\)$
因此,整个序列无需逐步反变换还原 Value 向量,仅需在注意力计算彻底收敛后,由 Kernel D 对汇总结果向量 \(\tilde{o}\) 执行一次单次逆旋转 \(o = R^\top \tilde{o}\) 即可完美还原最终输出。针对性能分析发现 Kernel C(Value 累加)因需遍历所有历史标记而成为耗时占 0.79ms 的瓶颈节点,设计了双策略调度机制:在序列长度 \(T < 512\) 时采用词级并行(word-parallel),在 \(T \ge 512\) 时自适应切换至维度级并行(dimension-parallel),保障在各类序列长度下的吞吐稳定性。
实验关键数据¶
主实验¶
评测在配备 128 GB 统一内存的单一 Apple M4 Max 设备上展开,验证了在 2,048 长度上下文下的困惑度退化(\(\Delta\text{PPL}\),越低越好),以及端到端多模型在不同硬件预算目标下的实际吞吐与峰值显存占用:
| 模型架构 | 激活参数量 | 默认 PPL | TurboQuant (3-bit) \(\Delta\text{PPL}\) | IsoQuant (3-bit) \(\Delta\text{PPL}\) | 端到端吞吐 (tok/s) | 峰值显存 | 目标部署预算 |
|---|---|---|---|---|---|---|---|
| Gemma 4-26B-A4B | 4B (Dense MoE) | 1.3483 | +0.0622 | +0.0000 | 12.85 | 5.4 GB | 16 GB |
| Qwen3-30B-A3B / 3.6-35B | 3B (Sparse MoE) | 1.0844 | +0.0405 | +0.0009 | 15.60 | 6.8 GB | 16 GB |
| Nemotron-H 120B | 混合 Mamba+MoE | 1.0866 | +0.0039 | +0.0012 | 14.85 | 17.2 GB | 32 GB |
注:Gemma 4 因仅有 5 个全局注意力层触发 KV 压缩(其余为局部滑动窗口),故 \(\Delta\text{PPL}\) 呈现极高无损性;Nemotron-H 120B 在 17.2 GB 峰值显存下平稳运行,通过了全部 12 项严格生成质量门禁与 2 小时持续高压稳定性测试(Soak pass)。
消融实验¶
针对专家卸载与 KV 压缩交互效应(Gemma 4-26B-A4B),以及不同结构化旋转机制设计的严格消融对比:
| 评估维度 | 实验配置 / 策略 | 关键指标 | 结果与机制说明 |
|---|---|---|---|
| 专家卸载 × KV模式 | 无卸载 + 默认 FP16 KV | 109.8 tok/s (10,649 MB) | 模型直接驻留显存基线,高吞吐 |
| 专家卸载 × KV模式 | 无卸载 + IsoQuant 3-bit | 20.6 tok/s (10,748 MB) | 打包计算开销在显存充裕时无加速优势(降速 5.3x) |
| 专家卸载 × KV模式 | 开启卸载 + 默认 FP16 KV | 1.01 tok/s (2,445 MB) | 频繁 SSD 读写使 I/O 成为绝对瓶颈 |
| 专家卸载 × KV模式 | 开启卸载 + IsoQuant 3-bit | 1.05 tok/s (2,804 MB) | 显存极度紧绷下吞吐相当,但显存变异系数由 5.4% 降至 0.3% |
| 旋转算子设计消融 | v1: 单四元数 SO(3) 旋转 | 质量门禁: 0/5 | 存在不变实轴,导致 25% 的维度未被混叠消除离群 |
| 旋转算子设计消融 | v2: 独立分块 SO(4) 旋转 | 质量门禁: 1/5 | 缺乏全局跨块混叠,无法处理大跨度特征相关性 |
| 旋转算子设计消融 | v3: WHT + SO(4) 复合旋转 | 质量门禁: 5/5 | 完美融合全局能量均分与块内精细旋转,无死角消除离群点 |
在 llama.cpp 跨框架移植实验中(GGML_TYPE_ISOQUANT3_0),融合内核相比 TurboQuant 仅有 −3.2% 细微吞吐差异(96.98 vs 100.15 tok/s),而未融合路径由于多调度 280 次内核发射,导致生成吞吐断崖式下跌 18%(81.92 tok/s),提示算子融合是压缩域注意力的必要先决条件。
关键发现¶
- KV 压缩价值的明确判定准则:解码时间归因剖析表明,传统 Transformer MoE(Gemma 4、Qwen3)中 KV 注意力耗时超过总解码时间的 50%(分别为 51% 和 45%),此时引入 IsoQuant 极具收益;而在 Nemotron-H 这类混合 Mamba+MoE 架构中,SSM 状态空间层完全绕过了 KV 缓存,KV 注意力耗时骤降至 14%,使得缓存压缩收益显著放缓。作者提炼出明确工程判据:仅当模型的 KV 注意力时间占比超过约 20% 时,投入压缩域 KV 缓存才具备显著正向投资回报。
- 长文本检索与精度零损耗:在 WikiText-103 的 32K 上下文评估中,Qwen3.6 开启 IsoQuant 3-bit 后 PPL 为 5.624(相比全精度 5.625 差异仅 0.024%);在 4K–32K 的多键“大海捞针”(Needle-in-a-Haystack)测试中,3-bit 状态下检索精度实现 100% 保持,证明结构化各向同性旋转能近乎无损地保留高频位置敏感特征。
亮点与洞察¶
- 单次逆旋转的数学转化巧思:巧妙利用注意力机制中 Softmax 加权和的线性可加性与正交矩阵的线性保真性,将本需在每个 Key/Value 向量上执行的反量化与逆变换推迟至最终标量求和之后,仅对单个输出向量执行一次 \(\mathcal{O}(d \log d)\) 逆旋转,直接击碎了传统量化方案随序列长度暴涨的带宽与算力瓶颈。
- WHT 与双四元数 SO(4) 的极致极简组合:相较于常规深度学习对密集正交矩阵盲目可学习微调的依赖,通过数学理论严格指出 SO(4) 的通用万有双重覆盖性质(Spin(4) \(\cong S^3 \times S^3\)),以纯代数推导出的 256 参数与 1,408 次 FMA 取代了 16,384 参数的密集矩阵,极大压低了推理框架的静态存储与访存负担。
- 客观详实的消极实验报告(Negative Results):诚实公开了注意力残差预测器(AttnRes)用于跨层专家预取遭遇失败的细节(在统一内存架构下,CPU 端预取信号的计算与 GPU 命令缓冲区产生总线带宽争抢,导致 10.6%–11.2% 的吞吐衰退且命中率零提升),对未来端侧系统软硬件协同设计极具警示价值。
局限与展望¶
- MLA 与 RoPE 架构的硬性结构阻断:对于采用多头潜在注意力(MLA,如 DeepSeek-V2/V3)的模型,其 KV 显式解耦为内容隐空间和携带旋转位置编码(RoPE)的相位子空间。RoPE 维度严禁受到二次旋转或有损量化(否则相位被污染破坏绝对位置感),当前 IsoQuant 缺少子块拆分隔离机制,尚无法直接支持此类主流架构。
- Value 累加内核(Kernel C)仍具访存优化空间:尽管引入了词级与维度级自适应并行,Kernel C 依然占据单步解码 0.79ms 的绝对大头,未来仍需探索低开销寄存器重排与 Mojo 等新一代异构语言内核优化。
- 显存非受压场景下的降速反噬:在硬件物理内存足够容纳全量模型权重与 KV 的工况下,3-bit 打包数据的拆包点积运算反而造成 5.3 倍的推理吞吐降速;该架构本质为“内存紧缩应急救生筏”,不应作为显存充裕时的通用推理加速器。
相关工作与启发¶
- vs TurboQuant (ICLR 2026):TurboQuant 依赖全密集随机正交矩阵与在线向量量化,需要存储 \(d^2=16,384\) 个浮点参数并付出 \(\mathcal{O}(d^2)\) 计算,且依旧依赖张量物化;IsoQuant 提出 WHT+SO(4) 分解使参数降低 64 倍、运算降低 11.6 倍,并完全消除了物化开销。
- vs KIVI / KVQuant:KIVI 与 KVQuant 采用非对称或非均匀量化标量映射,但计算阶段必须先反量化回 FP16 参与 GEMM;IsoQuant 则是端到端的计算范式跃迁,建立直接在 3-bit 压缩张量上运行点积与累加的融合计算流水线。
- vs FlexGen (ICML 2023):FlexGen 确立了单卡 CPU/磁盘卸载吞吐模型;RotaryQuant 将其针对 MoE 专家结构重构为按需 LRU 卸载,并作为第三轴与架构感知权重量化、压缩域 KV 注意力紧密复合。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [提出直接在压缩域进行融合计算的全新执行模型,配合 WHT 与 SO(4) 双四元数正交设计,数学优雅且体系完整。]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 4 款百亿到千亿级主流 MoE 模型、端到端吞吐、PPL 退化、长文本 NIAH 以及消极实验剖析,数据极为扎实。]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑严丝合缝,系统三轴与内核剖析清晰,技术局限与边界定义坦诚明确。]
- 价值: ⭐⭐⭐⭐⭐ [首次在 32GB 消费级个人电脑上实现 120B MoE 模型的流畅可交互本地离线运行,对端侧本地部署与隐私计算具有里程碑意义。]