MixLLM: LLM Quantization with Global Mixed-precision between Output-features and Highly-efficient System Design

作者/机构: Zhen Zheng*, Xiaonan Song, Chuanjie Liu (Microsoft)

A1 主要贡献

本文针对现有大型语言模型（LLM）量化方案在精度、内存消耗和系统效率这三个方面（有效性三角形）存在的不足，提出了一种名为MixLLM的综合性解决方案。现有方案通常在某一方面有所侧重，但难以兼顾三者：纯权重（Weight-only）量化方案在处理大批量任务时系统效率会下降，且在Llama 3等新模型上4-bit量化精度损失明显；权重-激活（Weight-activation）量化方案虽然能提升计算效率，但激活量化带来的精度损失更大，且反量化开销会拖累系统性能；离群点分离和混合精度技术虽然能提升精度，但前者因稀疏张量处理效率低而影响系统性能，后者在低比特量化下精度下降仍然不可忽视。

为解决这些问题，MixLLM做出了以下核心贡献：

  1. 提出基于全局显著性的输出特征间混合精度权重方案，实现高精度与低内存消耗：

    • 核心思想：模型的不同神经元（输出特征/输出通道）对最终输出的贡献不同。因此，MixLLM为不同显著性的输出特征分配不同位宽的权重，对高显著性特征使用8-bit，其他使用4-bit（如图1所示）。
    • 全局显著性识别：与现有方法【50, 【Atom: Low-bit quantization for efficient and accurate LLM serving，2024，MLSys】】在单层内根据局部损失评估显著性不同，MixLLM从全局视角出发，根据每个输出特征对模型最终损失的影响来识别其显著性。这种方法承认了不同层对模型的重要性不同，从而更有效地将高位宽分配给最需要的特征。
    • 系统友好性：在输出特征间应用混合精度比在输入特征间更容易实现系统优化，因为不同输出特征的计算是相互独立的子问题。

  2. 通过算法-系统协同设计，实现高精度与高系统效率：

    • 量化配置优化：MixLLM找到了一个兼顾精度和效率的“甜蜜点”配置。它对激活值使用8-bit对称量化，因为激活值对精度的影响较大，且其张量尺寸相对权重较小，减少其位宽对计算密集度的提升有限。对权重则采用分组量化，其中8-bit部分为对称量化，4-bit部分为非对称量化，以最大化精度。
    • 高效GPU核优化：针对上述混合对称/非对称配置难以利用int8 Tensor Core的挑战，MixLLM设计了创新的系统优化策略：
      • 两步反量化（Two-step dequantization）：该方法使得混合配置能量化后的计算能有效利用快速的int8 Tensor Core。
      • 快速数据类型转换：设计了快速的整数到浮点数转换方法，显著降低了反量化开销。
      • 软件流水线：精心设计了软件流水线，以最大限度地重叠内存访问、反量化和矩阵乘法（MatMul）计算。

通过这些贡献，MixLLM仅用额外10%的比特（即W4.4A8配置），就能将Llama 3.1 70B模型的困惑度（PPL）增量从SOTA方案的约0.5降低到0.2以内，并在三个流行模型的MMLU-Pro测试上平均比SOTA提升0.93分，同时达到了业界领先的系统效率。

图1：输出特征间混合精度量化和核执行的图示。
图1：输出特征间混合精度量化和核执行的图示。

A2 背景知识、相关工作与讨论

2.1 量化背景

量化原理。量化通过仿射变换将张量X映射到具有较小位宽表示的目标范围：$X_q = \text{clamp}(\lfloor \frac{X}{s} \rceil + z, \text{range})$，其中s是缩放因子，z是零点。通过$X' = (X_q - z) \times s$可以恢复（即反量化）数值。$X'$被映射到离散的区间而非原始值，因此会产生精度损失。位宽对量化精度至关重要，因为它决定了量化值的区间数量（$2^{\text{bit width}}$）。例如，将位宽从4位增加到5位可以使区间数量翻倍，因此5-bit的RTN（Round-to-Nearest）量化可以轻松超越采用先进技术的4-bit量化（如表1所示）。

量化方案。缩放因子和零点可以从整个通道/token向量计算得出，称为逐通道/token（per-channel/token）量化；也可以从通道/token内的一个小分组计算，称为分组（group-wise）量化。分组方案由于区间缩放尺度更小，导致的精度损失更小，但需要更复杂的GPU核设计。对称量化使用0作为零点值，这简化了计算（$X_q = \text{clamp}(\lfloor \frac{W}{s} \rceil, \text{range})$, $X' = X_q \times s$），并使许多工作【45, 【Smoothquant: Accurate and efficient post-training quantization for large language models，2023，ICML】；40, 【Tensorrt-llm，2024，https://github.com/NVIDIA/TensorRT-LLM】】能够在整个矩阵乘法（MatMul）的尾声阶段通过乘以缩放因子来进行反量化，从而设计出逐通道/逐token的量化核。然而，由于数据分布通常是不对称的，对称量化比非对称量化会导致更大的损失，尤其是在像4-bit这样的小位宽下 。

2.2 相关工作及通用量化原则讨论

本文主要关注训练后量化（Post-Training Quantization, PTQ）。

影响量化需求的系统。连续批处理（continuous batching）技术【47, 【Orca: A distributed serving system for transformer-based generative models，2022，OSDI】】能够将来自不同请求的解码任务批处理在一起，以增大LLM推理过程中矩阵乘法的批量维度。分块预填充（chunked-prefill）方法【19, 【Deepspeed-fastgen: High-throughput text generation for llms via MII and deepspeed-inference，2024，CoRR】；2, 【Taming throughput-latency tradeoff in LLM inference with sarathi-serve，2024，OSDI】；51, 【Batchllm: Optimizing large batched llm inference with global prefix sharing and throughput-oriented token batching，2024，arXiv preprint】】通过将预填充和解码任务合并到同一批次中，进一步推进了连续批处理，从而进一步增大了矩阵乘法的规模。这些技术使得许多LLM任务变为计算密集型（compute-bound），从而激发了减少计算量的需求。

纯权重量化及其局限性。为了提高纯权重量化的准确性，涌现了多种技术。GPTQ【14, 【GPTQ: accurate post-training quantization for generative pre-trained transformers，2022，CoRR】】在OBC【13, 【Optimal brain compression: A framework for accurate post-training quantization and pruning，2022，NeurIPS】】的基础上，基于OBS【18, 【Optimal brain surgeon and general network pruning，1993，ICNN】】的权重补偿方法，并引入了分块更新和重排序。AWQ【26, 【AWQ: activation-aware weight quantization for on-device LLM compression and acceleration，2024，MLSys】】提出根据激活值的特性来缩放权重。OminiQuant【37, 【Omniquant: Omnidirectionally calibrated quantization for large language models，2024，ICLR】】提出了可学习的缩放和权重裁剪因子。SpQR【11, 【Spqr: A sparse-quantized representation for near-lossless LLM weight compression，2024，ICLR】】、SqueezeLLM【21, 【Squeezellm: Dense-and-sparse quantization，2024，ICML】】和OWQ【25, 【OWQ: outlier-aware weight quantization for efficient fine-tuning and inference of large language models，2024，AAAI】】将离群点从量化中分离出来，并使用半精度处理。QuiP【8, 【Quip: 2-bit quantization of large language models with guarantees，2023，NeurIPS】】旨在通过非相干处理实现极低比特量化。ZeroQuant(4+2)【42, 【Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks，2023，CoRR】】旨在通过中等大小的FP6量化来提高准确性。纯权重量化不减少计算量，反而引入了额外的反量化操作。低比特权重将被反量化为float16，以便以float16数据类型执行矩阵乘法。当前纯权重量化面临两个挑战：1) 从准确性角度看，4-bit量化与float16模型之间仍然存在精度差距，特别是对于许多对微小精度下降敏感的真实业务场景，正如近期工作【42, 【Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks，2023，CoRR】；44, 【Quant-llm: Accelerating the serving of large language models via fp6-centric algorithm-system co-design on modern gpus，2024，USENIX ATC】】所讨论的。2) 在繁忙的服务器上，它可能导致系统效率下降，因为近期的LLM推理服务系统通常会将不同请求的处理批处理在一起，形成大型矩阵乘法。这些大型矩阵乘法是计算密集型的，会受到反量化开销的影响【27, 【Qserve: W4A8KV4 quantization and system co-design for efficient LLM serving，2024，CoRR】】。

权重-激活量化及其挑战。权重-激活量化有助于利用低比特计算单元。LLM.int8()【10, 【Llm.int8(): 8-bit matrix multiplication for transformers at scale，2022，CoRR】】观察到激活离群点问题，并通过半精度将离群点从量化中分离出来。ZeroQuant【46, 【Zeroquant: Efficient and affordable post-training quantization for large-scale transformers，2022，NeurIPS】】提出了逐token激活量化和分组权重量化。SmoothQuant【45, 【Smoothquant: Accurate and efficient post-training quantization for large language models，2023，ICML】】通过平滑处理解决了激活离群点问题，而AffineQuant【29, 【Affinequant: Affine transformation quantization for large language models，2024，ICLR】】提出了通用的仿射变换量化。RPTQ【48, 【RPTQ: reorder-based post-training quantization for large language models，2023，CoRR】】对通道进行重排序，将具有相似缩放值的数值聚集在一起。SpinQuant【28, 【Spinquant: LLM quantization with learned rotations，2024，CoRR】】和QuaRot【3, 【Quarot: Outlier-free 4-bit inference in rotated llms，2024，CoRR】】利用矩阵旋转特性来缓解离群点现象。Atom【50, 【Atom: Low-bit quantization for efficient and accurate LLM serving，2024，MLSys】】在输入特征间使用混合精度来提高4-bit激活量化的准确性。QoQ【27, 【Qserve: W4A8KV4 quantization and system co-design for efficient LLM serving，2024，CoRR】】是一个整体量化解决方案，包括渐进式分组量化、注意力平滑、旋转和通道重排序。尽管权重-激活量化具有减少矩阵乘法计算量的优势（即使用更小位宽的计算单元，具有更高的计算吞吐量），但它面临着由激活量化引起的精度下降的挑战，特别是因为激活通常比权重更难量化。目前最先进的低比特权重-激活解决方案【3, 【Quarot: Outlier-free 4-bit inference in rotated llms，2024，CoRR】；28, 【Spinquant: LLM quantization with learned rotations，2024，CoRR】；27, 【Qserve: W4A8KV4 quantization and system co-design for efficient LLM serving，2024，CoRR】】与4-bit纯权重量化相比仍有差距。除了精度下降，激活量化会比纯权重量化引入更多的反量化开销，这使得设计高效的GPU核变得具有挑战性。当启用非对称量化时，$(X_q - z)$ 的结果可能会超过 $X_q$ 位宽的范围，导致难以使用相应的Tensor Core计算。像Atom【50, 【Atom: Low-bit quantization for efficient and accurate LLM serving，2024，MLSys】】这样的系统因此避免使用非对称量化，但代价是更大的精度下降。分组量化需要细粒度的整数到浮点数（I2F）转换来应用每组的缩放因子。然而，I2F指令在GPU上比普通计算指令更昂贵【1, 【Demystifying the nvidia ampere architecture through microbenchmarking and instruction-level analysis，2022，IEEE High Performance Extreme Computing Conference】】，可能导致系统性能大幅下降（在我们的实践中下降超过10%）。此外，Tensor Core的吞吐量远高于SIMT Core，int8 Tensor Core为624 TOPS，而FP32/INT32 SIMT Core为19.5 TFLOPS/TOPS。现有工作中仍缺乏一个精心设计的软件流水线来重叠Tensor Core计算和基于SIMT Core的反量化，同时实现高精度。

总结。总的来说，现有的解决方案侧重于有效性三角形的部分方面，但不能很好地覆盖所有方面。MixLLM通过探索输出特征间的混合精度与全局显著性识别，以及协同设计的量化决策和GPU核，与上述工作是正交的。

A3 方法细节

3.1 MixLLM中的量化设计与决策

为了同时覆盖有效性三角形的三个方面，我们根据第2.2节的分析，对权重和激活量化做出以下设计和决策。

3.1.1 权重的输出特征间混合精度，并采用全局显著性识别

混合精度量化。众所周知，权重的不同元素在量化时对网络损失的显著性不同【21, 【Squeezellm: Dense-and-sparse quantization，2024，ICML】；11, 【Spqr: A sparse-quantized representation for near-lossless LLM weight compression，2024，ICLR】】。离群点分离方法可以通过使用float16存储高显著性元素来提高精度，但会受到低效的稀疏矩阵乘法的影响。我们观察到，在许多LLM的大多数线性层中，高显著性元素倾向于沿输出通道分布。基于此观察，我们可以为高显著性的输出通道分配较大的位宽，为其他通道分配较小的位宽，形成结构化的混合精度量化。通过实验，我们得出了与现有工作【21, 【Squeezellm: Dense-and-sparse quantization，2024，ICML】；11, 【Spqr: A sparse-quantized representation for near-lossless LLM weight compression，2024，ICLR】】相同的结论，即只有一小部分高显著性元素对模型的精度下降有显著贡献。因此，我们只需要为一小部分输出通道分配大位宽，即可实现良好的精度，并同时保持较小的内存消耗。

系统友好性。不同输出通道之间的结构化混合精度对系统效率和核开发很友好，因为在矩阵乘法中，不同输出特征是分离的，它们的计算是不同的子问题。图1展示了线性层如何利用输出特征间的混合精度进行计算。它将线性层划分为独立的子问题，并最终将子问题的输出收集在一起形成最终结果。这种优化空间与现有的量化优化（例如GPTQ【14, 【GPTQ: accurate post-training quantization for generative pre-trained transformers，2022，CoRR】】）是正交的，可以与它们一同应用。

全局显著性识别。一个关键问题是如何识别模型中的高显著性输出通道。固定阈值【11, 【Spqr: A sparse-quantized representation for near-lossless LLM weight compression，2024，ICLR】】或根据层的局部损失计算的固定数量/比例【50, 【Atom: Low-bit quantization for efficient and accurate LLM serving，2024，MLSys】；25, 【OWQ: outlier-aware weight quantization for efficient fine-tuning and inference of large language models，2024，AAAI】】的高显著性元素对端到端模型来说可能是次优的，因为不同层对模型的最终输出可能有不同的重要性【16, 【The unreasonable ineffectiveness of the deeper layers，2024，CoRR】；30, 【Shortgpt: Layers in large language models are more redundant than you expect，2024，CoRR】；12, 【HAWQ: hessian aware quantization of neural networks with mixed-precision，2019，ICCV】】。相对于一个层而言的高显著性通道，可能不是端到端模型的高显著性通道。在MixLLM中，我们根据它们对模型最终损失的影响，全局地计算高显著性通道（第3.2节）。因此，不同层将有不同数量的高显著性通道。图2显示了Llama 3.1 8B中前10%高显著性输出特征的分布，显示出不同线性层之间的巨大差异。

图2：Llama 3.1 8B模型中，根据每个特征对量化到4-bit后最终损失的贡献，全局前10%高显著性输出特征在每个线性层内的百分比。每个解码器层按顺序包含q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj和down_proj。
图2：Llama 3.1 8B模型中，根据每个特征对量化到4-bit后最终损失的贡献，全局前10%高显著性输出特征在每个线性层内的百分比。每个解码器层按顺序包含q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj和down_proj。

与Atom的区别。需要注意的是，此设计与Atom【50, 【Atom: Low-bit quantization for efficient and accurate LLM serving，2024，MLSys】】中的混合精度在两个方面有所不同。1) MixLLM首先解决的是全局而非局部地识别高显著性通道的问题。2) MixLLM在输出特征而非输入特征之间应用混合精度，这在系统性能和算法灵活性方面更具优势，因为输出特征是天然分离的。

3.1.2 结合算法-系统考量的量化决策：8-bit对称激活和4-bit非对称权重分组量化

激活量化决策。MixLLM在激活量化上与QoQ【27, 【Qserve: W4A8KV4 quantization and system co-design for efficient LLM serving，2024，CoRR】】做出了相同的决策，即使用8-bit，因为4-bit激活会导致较大的精度下降，但并不会带来显著的系统效率提升，原因是矩阵乘法执行更倾向于受限于较大的权重张量，而非较小的激活张量。这可以从线性层的计算强度中部分看出。给定token数量为$M$，输入/输出特征为$K/N$，计算强度$I = \frac{2MNK}{MKB_{act} + KNB_{weight}}$。$B_{act}$和$B_{weight}$分别是激活和权重的每个元素的字节数。当$M=512$，$N=K=4096$时，将$B_{weight}$从8减少到4将导致$I$增加80%，而将$B_{act}$从8减少到4只会带来5.88%的增幅。

权重和激活量化方案。MixLLM没有使用逐token激活量化，而是使用了分组RTN方法。一方面，表1显示，简单的分组RTN量化比逐token平滑方法表现更好。另一方面，MixLLM中的权重已经是分组的，分组激活不会在系统中导致显著增加的反量化开销。我们观察到，对于8-bit激活，对称量化已经足够（参考表1中的MixLLM W8A8），而对于4-bit权重，非对称量化至关重要。采用非对称的分组方法可能会导致核函数难以利用int8 Tensor Core，为此QoQ【27, 【Qserve: W4A8KV4 quantization and system co-design for efficient LLM serving，2024，CoRR】】引入了两步量化方法。相反，我们利用对称和非对称混合的特性，设计了一种两步反量化方法（第3.3节）。

3.2 全局精度搜索算法

显著性计算。如3.1节所述，MixLLM全局而非局部地确定所有层中所有输出特征的精度。它根据这些特征对模型最终损失的显著性来识别它们，并为导致更大损失的特征分配更大的位宽。具体来说，它将一个通道c的显著性S计算为：

$$S_c = |l(c_q) - l(c_0)|$$

这个值是量化该单通道与不量化该通道时模型损失的距离。在公式1中，$l()$是模型关于单个通道的损失函数，$c_q$是该通道的量化权重，$c_0$是原始权重。请注意，它将除c以外的其他神经元在$l()$中视为常数。

损失函数估计。我们使用泰勒展开法来估计损失函数$l(c)$（与现有量化工作类似，忽略高阶项）：

$$l(c) \approx l(c_0) + g^T(c - c_0) + \frac{1}{2}(c - c_0)^T H(c - c_0)$$

其中$g = E[\frac{\partial}{\partial c}l(c)]$是损失关于通道的梯度，$H = E[\frac{\partial^2}{\partial c^2}l(c)]$是关于通道的二阶梯度（即Hessian矩阵）。

Hessian矩阵近似。计算Hessian矩阵是不可行的，因为它成本太高。我们用校准数据集D上的（经验）Fisher信息矩阵F来近似Hessian矩阵H：

$$H \approx F = \frac{1}{|D|} \sum_{d \in D} g_{d} g_{d}^{T}$$

请注意，F是关于一个通道的，这与近期工作中忽略任何神经元间交互的对角Fisher信息矩阵不同【23, 【A fast post-training pruning framework for transformers，2022，NeurIPS】；21, 【Squeezellm: Dense-and-sparse quantization，2024，ICML】】。

显著性最终计算。基于这个近似，二阶损失因子$\frac{1}{2}(c-c_0)^T(\sum g_d g_d^T)(c-c_0)$可以进一步简化为$\frac{1}{2}(\sum g_d^T(c-c_0))^2$，将昂贵的链式矩阵乘法简化为单个向量积。最终，显著性可以计算为：

$$S_{c}=\frac{1}{|D|}\sum_{d\in D}|g_{d}^{T}(c_{q}-c_{0})+\frac{1}{2}(g_{d}^{T}(c_{q}-c_{0}))^{2}|$$

一阶信息保留。我们在计算过程中不忽略一阶信息，这与OBD【24, 【Optimal brain damage，1989，NIPS】】和许多近期的量化工作【14, 【GPTQ: accurate post-training quantization for generative pre-trained transformers，2022，CoRR】；11, 【Spqr: A sparse-quantized representation for near-lossless LLM weight compression，2024，ICLR】；21, 【Squeezellm: Dense-and-sparse quantization，2024，ICML】】不同。这是因为在公式4的估计中，一阶因子可能更显著，因为估计的二阶因子是每个样本一阶因子的平方除以二。考虑到对于训练良好的模型，$g$可能非常小，且量化权重的变化量通常不大，一阶因子可能比二阶因子更大。此外，我们需要的是损失本身，而不是损失函数的参数，因此我们不需要为了简化参数计算而忽略一阶因子。

算法流程。算法1展示了全局精度搜索的过程。它计算所有线性层的所有输出通道的显著性，并全局按降序排序。给定全局阈值$N_{\text{largebit}}$作为高位宽精度通道的数量，前$N_{\text{largebit}}$个通道将被量化为8-bit，其他通道将被量化为更小的位宽（本文中为4-bit）。任何量化方法（如GPTQ、裁剪搜索）都可以独立地应用于这两个不相交的通道部分。注意，我们是一次性计算通道的显著性，而不是以较小的步长迭代识别高显著性部分，因为我们观察到单次方法与迭代方法显示出相似的结果，并且比后者节省了大量的计算开销。

算法 1 全局精度搜索过程

输入: 线性层数量 L，所有线性层的权重和梯度 (Wi ∈ R^(O,I), Gi ∈ R^(O,I) for layer i)。
输出: 全局高位宽和低位宽通道索引 (largebit_channels, smallbit_channels)。
1: S_global ← []
2: for i = 1, 2, ..., L do
3:   W_delta_i ← quantize(W_i) - W_i
4:   S_1st ← sum(G_i ⊙ W_delta_i, dim=1)  // 逐通道点积
5:   S_2nd ← 0.5 * (S_1st)^2
6:   S ← |S_1st + S_2nd|  // S ∈ R^O，O个通道的显著性
7:   for channel_id = 1, 2, .., O do  // 记录该层每个输出通道的显著性
8:     S_global.append(tuple(i, channel_id, S[channel_id]))
9: sort(S_global)  // 根据显著性降序排序
10: largebit_channels, smallbit_channels ← S_global[:N_largebit], S_global[N_largebit:]

3.3 高效的量化计算系统

利用int8 Tensor Core的两步反量化。对于W4A8计算，反量化后的权重和激活分别为$(W_q - z)s_w$和$A_q s_a$，其中$W_q$和$z$是uint4数据类型，$A_q$是int8数据类型，$s_w$和$s_a$是float16数据类型。如果在矩阵乘法计算之前直接将张量反量化为float16类型，将无法使用GPU上快速的8-bit Tensor Core。相反，MixLLM在每个组内使用两步反量化。具体来说，MixLLM首先将权重部分反量化为$(W_q - z)$，然后用8-bit Tensor Core将其与$A_q$相乘。最后，它将这个矩阵乘法的结果与每个组内的两个缩放因子相乘。注意，我们对$(W_q - z)$使用int8数据类型，因此不存在溢出问题。

图3：二进制(010010110xx...x)的浮点数和整数值，各自在一个连续的范围内。
图3：二进制(010010110xx...x)的浮点数和整数值，各自在一个连续的范围内。

快速整数到浮点数转换并部分融入Tensor Core指令。在上述两步反量化计算中，第二步是整数张量$A_q(W_q - z)$和浮点张量$s_a s_w$之间的矩阵乘法。这需要在乘法操作之前进行整数到浮点数的转换（I2F）。I2F指令在现代GPU上是昂贵的。我们利用依赖于范围的快速I2F转换，将I2F指令转换为两个加/减指令。具体来说，它基于一个事实：存在一个特定范围，其中整数值的二进制表示与相应浮点数的二进制表示相同。如图3所示，前9位为010010110的二进制数分别代表一系列连续的int32和float32值。我们可以通过给整数值加上一个偏置，使其落入这个连续范围内，然后在浮点数中减去一个相应的偏置（底层二进制相同）来恢复其浮点类型的值。我们取该范围的中间值作为偏置，以最大化可安全转换的数据范围，其十六进制数为0x4b400000（即在剩余的23位中，第一位为1，其他位为0）。这允许将一个连续范围内的$2^{23}$个int32数转换为float32。k个int8元素点积的范围是$2^{16}k$，因此上述快速I2F转换允许k值最大为128。我们使用量化组大小为128，可以安全地使用快速I2F。我们进一步将整数减法融合到Tensor Core的mma（Matrix Multiply-Accumulate）指令中。mma指令在矩阵乘法计算中计算$D = AB + D$。我们在每个量化组的矩阵乘法计算前，将累加器D初始化为偏置整数，之后只需在矩阵乘法后减去浮点偏置即可。换句话说，昂贵的I2F被转换成了一次浮点减法。上述I2F简化在A100 GPU上为512/4096/4096（M/N/K）的量化矩阵乘法计算带来了超过20 TOPS的性能提升。

图4：分组W4A8/W8A8量化矩阵乘法的GPU核软件流水线。假设完美重叠。G2S：全局内存加载到共享内存；S2R：共享内存加载到寄存器；MMA：矩阵乘法累加；I2F：整数到浮点数转换；deq：反量化；acc：累加。
图4：分组W4A8/W8A8量化矩阵乘法的GPU核软件流水线。假设完美重叠。G2S：全局内存加载到共享内存；S2R：共享内存加载到寄存器；MMA：矩阵乘法累加；I2F：整数到浮点数转换；deq：反量化；acc：累加。

GPU上量化线性核的端到端软件流水线。图4展示了量化核的软件流水线。除了基本的warp tile和block tile，我们引入了量化组tile，用于快速I2F和逐组缩放因子乘法。它在寄存器级别使用两个输出缓冲区进行输出累加，一个用于逐组累加，另一个用于全局累加。这允许在组级缓冲区上应用逐组缩放因子。我们在组tile开始时用偏置整数初始化组缓冲区，在组tile结束时减去浮点偏置。至于两步反量化，第一步在warp tile内，每个输入元素在送入MMA之前减去零点；第二步在组tile结束时，通过乘以缩放因子完成。我们使用向量化内在函数（vsub4）【27, 【Qserve: W4A8KV4 quantization and system co-design for efficient LLM serving，2024，CoRR】】在单个指令中执行四个int8减法。此外，为了提高全局内存加载效率，我们提前预打包权重张量的内存布局，以避免运行时输入元素的置换。总的来说，这个软件流水线可以最大限度地重叠内存加载和计算，以及SIMT Core的反量化计算和Tensor Core的矩阵乘法计算，并最小化了分组反量化的开销。

不同位宽子问题的并行执行。如图1所示的执行过程，MixLLM利用CUDA Graph在GPU上并行执行不同的子问题。最终，两部分的矩阵乘法执行结果会写入同一个目标张量的不同通道索引，以生成最终输出。我们通过矩阵乘法的融合尾声（fused epilogue）来实现此功能，将输出分散到相应的索引，这基本上没有成本。

A4 实验

实验环境

实验结果

4.2 困惑度（Perplexity）评估

表1和表2展示了在Wikitext2和C4数据集上的困惑度评估结果，主要结论如下：


表 1: 在wikitext2/c4上的困惑度评估（↓）（灰色为c4），序列长度2048。NA表示不支持。Abn表示值过大(> 10^5)。对于MixLLM，pn表示n%的8-bit。
表1


表 2: 与相关工作中报告数值的PPL（wikitext2）比较。
表2

4.3 系统性能

图5展示了在单个线性层上的性能评测，输入特征4096，输出特征4096/14336，token数从1到1024。主要发现如下：

图5：在A100 GPU上，两种类型的单线性层相对于torch float16基线的加速比。
图5：在A100 GPU上，两种类型的单线性层相对于torch float16基线的加速比。

4.4 下游任务评估

表3展示了在三个流行LLM上的下游任务评估结果，结论如下：


表 3: 在Llama-3.1-8B/Qwen2.5-7B/Mistral-7B-v0.3上的下游任务评估（↑）。上方是三个模型的平均值。BBH为3-shot，MMLU pro为5-shot，其他为zero-shot。
表3

4.5 详细分析

图6：Llama 3.1 8B模型在不同配置下的困惑度（wikitext2）。
图6：Llama 3.1 8B模型在不同配置下的困惑度（wikitext2）。


表 4: 在Llama 3.1 8B中，MixLLM全局10% 8-bit输出特征下，七类线性层中8-bit输出特征的平均百分比。
表4


表 5: MixLLM中全局精度搜索的开销。
表5

A5 结论

本文提出了MixLLM，通过探索一个较少被研究的优化空间——输出特征间的混合精度量化，实现了高精度、低内存消耗和高系统效率。MixLLM根据每个输出特征对全局模型损失的距离估计来识别其显著性，而非局限于局部层损失。通过为最需要高精度的特征分配更大的位宽，MixLLM以低内存消耗实现了超越SOTA的精度。不同位宽的子问题是分离的，可以在GPU上高效地并行运行。我们确定了对精度和系统效率都友好的量化配置“甜蜜点”。为了解决系统效率的挑战，我们设计了两步反量化以启用int8 Tensor Core计算，并通过快速整数-浮点转换减少反量化开销。我们还设计了端到端的软件流水线，以重叠内存访问、SIMT Core的反量化计算和Tensor Core的矩阵乘法。实验结果表明，MixLLM在精度上优于现有工作，并以低内存成本实现了SOTA的系统效率。