使用 GemLite、TorchAO 和 SGLang 加速LLM推理

由 PyTorch、Mobius Labs 和 SGLang 团队开发

大型语言模型（LLMs）通常非常资源密集，需要大量的内存、计算和电力才能有效运行。量化通过将权重和激活从 16 位浮点数降低到更低的比特率（例如 8 位、4 位、2 位）提供了解决方案，实现了显著的加速和内存节省，同时也支持更大的批量大小。

现有的低精度推理解决方案对于小批量大小效果良好，但存在以下问题：

• 当我们增加批量大小时，性能会下降
• 量化类型限制，例如，某些内核仅支持对称量化，这可能会对模型在低比特下的精度产生影响
• 量化、序列化和张量并行（TP）之间的相互作用使得加载量化模型变得困难，并需要修改用户模型

为了解决这些挑战，我们创建了一个端到端、高性能、模块化和可扩展的低精度推理解决方案，集成了以下库：

• GemLite，一个 Triton 内核库，解决了大批次大小和量化类型限制的性能瓶颈
• TorchAO，一个 PyTorch 原生库，为量化、稀疏性和张量并行（使用 DTensor）提供了一体化的体验
• SGLang，一个快速、高效且可定制的服务框架，适用于大型语言模型（LLM）和视觉语言模型（VLM），具有广泛的支持模型

如果您有兴趣在 SGLang 中尝试，请遵循以下复现说明。在本博客的其余部分，我们将详细介绍 GemLite、TorchAO 和 SGlang 的相关细节，包括库本身的设计以及解决上述问题的集成，最后我们将展示在 Llama 3.1-8B 模型上针对不同批量和张量并行大小的基准测试结果。

1. 结果预告

以下是在 Llama 3.1-8B 机器上使用 8xH100 解码结果的总结。所有实验的基线是使用 bfloat16 torch.compile 模型的：

我们的解决方案支持 NVIDIA GPU，包括 H100 和 A100，并在批大小和 TP 大小方面，对于仅 int4 权重（从 1.14 倍到 1.95 倍）和 float8 动态量化（从 1.10 倍到 1.28 倍）都实现了比编译的 bfloat16 基线更高的速度提升。请注意，量化可能会对精度产生轻微影响，这超出了本博客文章的范围。我们的 int4 权重仅量化与保持精度的技术（如 HQQ）兼容。请参阅 TorchAO 的 README、此基准和此博客以获取更多信息。

2. GemLite：内核开发

这些内核是作为 GemLite 项目的一部分开发的，该项目致力于优化低比特矩阵乘法内核。使用 Triton 开发，GemLite 在各种激活、比特率和硬件上提供高度灵活和高效的解决方案。简而言之，内核提供：

• 支持各种激活数据类型：fp16、int8 和 fp8
• 兼容性：与未打包格式（例如，int8、fp8）和打包格式（例如，uint4、uint2、uint1）无缝兼容
• 性能优化：包含优化内核和自动调优工具，以实现不同硬件和批处理大小的高性能
• 集成：与 torch.compile 和 CUDA 图兼容，确保支持张量并行等高级功能

内核选择

优化大型语言模型（LLM）生成时的内核选择，需要满足不同批量大小的特定需求。LLM工作负载涉及计算密集型和内存密集型迭代的混合：较小的批量大小的内存密集型，而较大的批量大小的计算密集型。GemLite 内核旨在适应这些不同的需求，确保每种场景下的最佳执行。

在内存密集型场景中，数据传输是限制因素，处理器通常等待数据被检索，导致计算资源利用率低下。对于批量大大小为 1 的情况，GEMV 内核表现最佳，而对于较大的批量大大小，GEMM 内核则更高效。对于 2 到 64 之间的批量大大小，当矩阵“瘦”时，使用 GEMM-SPLITK 内核以实现更好的 GPU 利用率（arXiv）。

GemLite 包括以下针对这些场景优化的内核：

单样本推理

对于单样本推理，我们使用 GEMV 核。然而，非对称量化方法需要为每个块加载额外的元数据，例如缩放和零点。这可能导致内存传输增加，因此需要谨慎处理。

具体来说，对于打包数据，我们的实验表明，每两个连续块只加载一次缩放和零点可以最小化冗余操作。由于这些块共享相同的元数据，这种方法可以实现：

• 与默认的 GEMV 核相比，端到端推理速度提高 5-8%
• 比传统的 Split-K 方法提高 30-40%

这款新的内核/算法 GEMV_REVSPLITK 现已提供。

对于非打包数据，采用 GEMV_SPLITK 算法。此算法通过迭代 k 维度来计算点积，而不依赖于 Triton 的 tl.dot。

批量推理

对于中等批量大小的场景，我们使用基于 GEMM 的 Split-K 方法（arXiv），该方法将 k 维度（权重行）分割成多个作业。通过自动调整从 1 到 16 的值来找到最优分割 SPLIT_K 参数。将 SPLIT_K 设置为 1 可以启用回退实现到 GEMM 内核，从而允许使用相同的内核代码处理从 32 和 64 开始的计算密集型批大小，具体取决于矩阵形状和设备。

最大化高性能：关键实施见解

为了实现高性能，必须仔细处理各种实施细节。以下是我们在确保高性能方面关注的几个关键方面：

1. 性能自调优

   自调优对于实现最佳内核性能至关重要。由于此过程可能耗时较长，GemLite 提供了工具来自动保存和加载所有内核的自调优结果。这确保了自调优过程在每个 GPU 设备上仅执行一次，最小化运行时间，减少重复开销，并保持运行的一致性。

2. 确保内核正确性

   确保在不同量化和配置设置下内核的正确性至关重要。Triton 的早期配置剪枝在这一过程中发挥着关键作用。例如，在 Split-K 调优期间，只有当 K 能被 BLOCK_SIZE_K × SPLIT_K 整除时，才会选择配置，而 BLOCKS_SIZE_K 会根据组大小值进一步剪枝。这种方法确保了内核操作的高效性和正确性。

3. 克服位解包瓶颈

   在部署到数据中心级 GPU（如 NVIDIA 的 A100 和 H100）时，观察到与位解包相关的性能瓶颈。为了缓解这些问题，探索了各种位打包配置，包括按列打包与按行打包以及尝试不同的位打包宽度（例如，8 位与 32 位）。值得注意的是，从 32 位打包过渡到 8 位打包在 A100 上实现了高达 18%的性能提升，在 H100 上实现了 6%的性能提升。

4. torch.compile 兼容性

   为确保与 PyTorch 的 torch.compile 无缝兼容，内核调用被封装在 custom_op 中。这种集成允许高级功能，如预钩子和早期配置剪枝正确运行，在不牺牲性能的情况下提供准确的结果。虽然 PyTorch 中的一些功能尚未完全支持，但 custom_op 的实现有效地弥合了这一差距，确保了平滑的集成和高效性能。

3. TorchAO

TorchAO 是一个用于训练和推理的 PyTorch 原生量化与稀疏性库，具有简单的用户 API 以训练、量化和部署低精度模型，并且与其他 PyTorch 功能（如分布式推理和 torch.compile）具有可组合性。

PyTorch 默认不支持低精度数据类型或不同的打包格式。通过 Tensor 子类，我们扩展了 PyTorch 原生的 Tensor 抽象和模型量化，作为数据类型转换，而针对自定义内核的不同打包格式则通过布局来处理。例如，我们支持使用 int4 权重的量化线性运算，以 Tensor Core 友好的布局打包，使用 tinygemm 或 GemLite 内核实现。更多详情请见此处。

除了为开发者提供更多 PyTorch 原生抽象之外，我们还想强调该设计对建模用户带来的两个好处。

1. 序列化：像浮点模型一样保存和加载量化权重到 state_dict 中，无需在加载量化权重之前将浮点模型转换为量化模型。这减少了分发和部署量化模型时的摩擦。

2. 可组合性：无缝集成下游功能，如张量并行，使用户能够专注于建模，无需担心与张量并行、torch.compile 和其他 PyTorch 功能的兼容性。由于这些功能采用张量级抽象实现，用户在大多数情况下无需修改模型即可进行量化和小批量推理。

GemLite 内核集成

为了实现上述 GemLite 内核的益处，我们将 GemLite 集成到 TorchAO 中。这种集成利用了 GemLite 广泛的兼容性和灵活性，允许在 4 位和 8 位下进行仅权重量化，支持非对称和对称量化方案，以及 32 位和 8 位打包大小，以及分组和无分组量化。我们通过 quantize_ API 启用此集成，该 API 可以与 GemLite 构造函数一起使用，如下所示

quantize_(model, gemlite_uintx_weight_only(group_size, bit_width, packing_bitwidth))

创建此集成的主要困难在于确保 TorchAO 的可组合性保证适用于 GemLite 量化内核选项的全范围。虽然主要集成相对直接，但确保每种不同的量化类型及其相关内核与张量并行工作良好并非易事。

Torch 张量并行

张量并行是一种有效加速LLM推理的方法。TP 将线性或嵌入模块的大矩阵分片到多个设备上，通常采用列向或行向风格。随着权重矩阵的分布，计算也被分解。例如，下面的列向模式允许在四个设备上同时进行矩阵-向量乘法：

PyTorch 通过将常规张量（例如矩阵 A）转换为 DTensor 来实现 TP：

dtensor = _shard_tensor(mA, device_mesh, (Shard(0),))

由于 DTensor 存储了关于分片的元信息，因此它知道如何在需要时重建完整结果。以 Transformers 的前馈模块为例，下投影和上投影分别使用列式和行式分片，当它们移动到下一个操作时，DTensor 将自动对所有 rank 的结果执行 all-reduce。这种自动化允许模型作者专注于计算，无需担心分布式执行所需的通信。

张量并行和量化顺序

由于 DTensor 和量化都是张量级别的转换，确保工作流程可以在不同的设置上一般工作，应用顺序很重要。我们有两个观察结果：(i) 检查点通常以量化格式保存，以在每次运行之前节省量化开销；(ii) TP 可能根据资源限制或服务协议在不同的设备上运行不同的数量。因此，我们首先对原始张量进行量化，根据是否需要重用将其保存到磁盘。在服务启动时，我们加载量化检查点，并在将张量加载到模型时即时将其分片为 DTensor。

TorchAO 中的张量并行支持

首先对模型进行量化，然后分配 Tensor，我们将得到 DTensor(QuantizedTensor(weight)) ，其中 DTensor 表示分布式 Tensor 类， QuantizedTensor 表示 TorchAO 中的量化张量类。 QuantizedTensor 应支持在构建 DTensor 时调用的操作，包括切片和视图操作。为了确保整体执行效率，0 和 1 维度上切片的打包权重应与先切片再打包（打包和切片操作应可交换）的结果匹配，否则打包格式与张量并行性不兼容。

4. SGLang

SGLang 是一个用于大型语言模型和视觉语言模型的快速服务框架。它以其几乎零开销的批量调度和快速约束解码而闻名。它主要用 Python 实现，轻量级，易于修改。它也是第一个集成 torch.compile 的框架之一。

SGLang 中的 TorchAO 集成

我们将支持仅 int4 权重量化的特定类型量化 API quantize_ 集成到支持 int4 权重量化（包括 tinygemm 和 GemLite 版本）、float8 动态量化以及其他几种量化类型的 SGLang 中。用户可以通过在基准测试脚本中添加 --torchao-config 参数来启用量化。当前启用的选项还支持通过与启用 --tp-size 选项的 DTensor 的组合来实现张量并行性。

SGLang 中的 Torch 原生张量并行支持

SGLang 中现有的模型定义使用与张量并行风格耦合的特殊线性模块，例如： MergedColumnParallelLinear 、 QKVParallelLinear 和 RowParallelLinear 。为了解耦模型定义和张量并行化风格，我们定义了一个使用 PyTorch 的普通 nn.Linear 模块的原生 PyTorch 模型，并依赖于 PyTorch 张量并行 API 进行并行化以及 torch.compile 进行加速。在相关的模块层次结构中，我们添加了一个描述子模块应该如何并行化的字典。例如，在 class LlamaAttention 中，我们定义：

_tp_plan = {
    "qkv_proj": "Colwise_Sharded",
    "o_proj": "Rowwise",
}

其中 "qkv_proj" 和 "o_proj" 是 wqkv 和 wo 投影的 FQNs，值是它们的 TP 风格。

我们在 model_parallel.py 中定义了一个 TP 引擎。它递归地在模型中搜索 _tp_plan ，并使用 PyTorch 的 parallelize_module API 将指示的 TP 样式应用于子模块。

5. 结果

评估主要集中在 H100 机器上两种流行的量化技术：int4 权重仅量化以及 float8 动态量化。这些方法因其广泛用于优化 H100 机器上的内存效率和计算性能而被选中，成为各种工作负载基准测试的理想候选。

• int4 权重仅量化：这种方法显著减少了内存占用并加速了内存受限工作负载的解码，在计算密集型场景（如预填充或更大的批大小）中对性能的影响最小。以下列出了 bf16、GemLite 和 tinygemm 内核在不同批大小和张量并行配置下的结果。
• float8 动态量化：虽然提供的内存节省较少，但这种方法通常提供更高的精度和平衡的速度提升，适用于内存密集型和计算密集型任务。借助 Hopper 级硬件和本机 fp8 支持，AO 使用的 cutlass/cuBLAS 内核效率高，有助于显著加速

下面的图表显示了不同 tp 大小下的解码令牌/秒，每个图表显示了不同批处理大小和不同量化类型的测试结果：

• BF16 是我们的 bfloat16，torch.compile’d 基线
• tinygemm-4-64 在 TorchAO 中使用 int4_weight_only 量化，它是一种 4 位组量化，组大小为 64，使用 tinygemm 内核
• gemlite-4-64 正在使用 TorchAO 中的 gemlite_uintx_weight_only 量化，4 表示 4 位，64 也是组大小，使用 GemLite 内核
• fp8dq-per_row 正在使用 TorchAO 中的 float8_dynamic_activation_float8_weight 量化，激活和权重都使用每行比例进行量化

对于 int4 仅权重量化，在批大小为 1 时，tinygemm 内核实现了最佳性能。然而，其效率随着批大小的增加而下降。相反，GemLite 有效地弥合了这一差距，在更大的批大小下提供了更优的性能。与 tinygemm 相比，尽管受到 Triton 的性能优化限制，GemLite 在预填充阶段也实现了 9-10 倍的速度提升。

Float8 动态量化在不同批大小和 1 个张量并行大小下，始终比 bfloat16 快 1.3 倍，在更大的张量并行大小下，速度提升了 1.1 倍到 1.2 倍。随着张量并行大小的增加，整体速度提升会减少，这是由于矩阵乘法大小的减少而预期的。请注意，我们确实期望在预填充阶段也能获得速度提升，但由于我们依赖于 torch.compile 进行速度提升，而 SGLang 中尚未启用预填充编译，我们将将其留待未来工作。

复制说明

我们在 8xH100 机器上使用 GemLite 0.4.1、从 commit feb2b76 构建的 SGLang、TorchAO 夜间版 0.8.0.dev20241223+cu124 和 PyTorch 2.5.1 进行了基准测试。选择 Llama-3.1 Instruct 模型作为评估的架构。

BATCH_SIZE=16
# Note: gemlite is only compatible with float16
# while int4wo-64 (tinygemm-4-64 as shown in the graph) and fp8dq-per_row should use bfloat16
DTYPE=float16
# int4wo-64, fp8dq-per_tensor
TORCHAO_CONFIG=gemlite-4-64
TP_SIZE=2
# Decode performance
python3 -m sglang.bench_offline_throughput --model-path meta-llama/Llama-3.1-8B-Instruct --json-model-override-args '{"architectures": ["TorchNativeLlamaForCausalLM"]}' --dataset-name random --random-input 1024 --random-output 512 --random-range 1 --num-prompts $BATCH_SIZE --enable-torch-compile --dtype $DTYPE --torchao-config $TORCHAO_CONFIG --tp-size $TP_SIZE

# Example output
# Benchmark...
# [2024-12-20 12:42:16 TP0] Prefill batch. #new-seq: 2, #new-token: 2046, #cached-token: 4, cache hit rate: \0.06%, token usage: 0.00, #running-req: 0, #queue-req: 0
# ...
# [2024-12-20 12:45:35 TP0] Decode batch. #running-req: 16, #token: 16763, token usage: 0.01, gen throughput\ (token/s): 2.20, #queue-req: 0
# [2024-12-20 12:45:38 TP0] Decode batch. #running-req: 16, #token: 24443, token usage: 0.02, gen throughput\ (token/s): 2739.89, #queue-req: 0

# We reported the last throughput (token/s) as the performance for decode

结论

通过 GemLite 的高性能和可扩展内核、PyTorch 原生架构优化库 TorchAO 和高效推理框架 SGLang，我们展示了简单且可组合的用户 API，在不同批处理大小和张量并行大小下实现了 int4 和 float8 的快速端到端量化推理，从而降低了LLMs的资源需求。这是我们迈向满足不同模型、工作负载、精度和硬件的快速推理需求的第一步，我们期待继续推进端到端混合和低精度LLM推理的先进技术。

我们接下来的工作重点如下：

• 探索权重和激活量化的多种组合，以在速度和精度之间取得最佳平衡
• 扩展对其他 GPU 架构的支持，以拓宽可访问性
• 增强与 MoE 模型的兼容性，以满足可扩展推理不断增长的需求
• 允许在 TorchAO 中轻松集成快速自定义内核，以便它们可以轻松被 SGLang 和其他推理框架利用
• 虽然在这篇博客文章中我们没有测量准确性的影响，但我们可以在 TorchAO 中开发自动量化工具，让用户在性能和准确性之间进行权衡
• 在 SGLang 中更好地集成张量并行性，以支持运行更大的模型
• 启用 torch.compile 在 SGLang 中的预填充阶段

我们还邀请社区积极参与测试、提供反馈，并共同塑造快速高效的LLM推理的未来