跳转到主要内容
博客

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

作者： 2025 年 1 月 21 日2025 年 5 月 5 日暂无评论

大型语言模型 (LLM) 通常是资源密集型的，需要大量的内存、计算和电力才能有效运行。量化提供了一种解决方案，它将权重和激活从 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.compiled 模型。

  bfloat16 w/ torch.compile int4 仅权重量化，组大小 64 float8 每行动态量化
批处理大小 1，TP 大小 1 131 tokens/秒 255 tokens/秒 (1.95 倍加速) 166 tokens/秒 (1.27 倍加速)
批处理大小 32，TP 大小 1 2799 tokens/秒 3241 tokens/秒 (1.16 倍加速) 3586 tokens/秒 (1.28 倍加速)
批处理大小 32，TP 大小 4 5575 tokens/秒 6334 tokens/秒 (1.14 倍加速) 6159 tokens/秒 (1.10 倍加速)

我们的解决方案支持 NVIDIA GPU，包括 H100 和 A100，并且在所有批处理大小和 TP 大小下，对于 int4 仅权重（1.14 倍至 1.95 倍）和 float8 动态量化（1.10 倍至 1.28 倍），都比编译的 bfloat16 基线实现了加速。请注意，量化可能会对准确性产生很小的影响，这不在本博客的讨论范围之内。我们的 int4 仅权重量化与 HQQ 等精度保持技术兼容。有关更多信息，请参阅TorchAO 的 README、此基准和此博客。

2. GemLite：内核开发

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

  • 支持各种激活数据类型：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 可在此处找到：https://github.com/mobiusml/gemlite/blob/master/gemlite/triton_kernels/gemv_revsplitK_A16fWnO16f_int32packing.py。

对于非打包数据，采用 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. 克服比特解包瓶颈在部署到 NVIDIA A100 和 H100 等数据中心级 GPU 时，观察到与比特解包相关的性能瓶颈。为了缓解这些问题，探索了各种比特打包配置，包括按列打包与按行打包，以及尝试不同的比特打包宽度（例如，8 位与 32 位）。值得注意的是，从 32 位打包到 8 位打包在 A100 上带来了高达 18% 的性能提升，在 H100 上带来了 6% 的性能提升。
  4. torch.compile 兼容性为了确保与 PyTorch 的 torch.compile 无缝兼容，内核调用被包装在自定义操作中。这种集成允许高级功能（例如预钩子和早期配置剪枝）正常工作，在不牺牲性能的情况下提供准确的结果。尽管 PyTorch 尚未完全支持其中一些功能，但 custom_op 实现有效地弥合了差距，确保了平滑集成和高性能。

3. TorchAO

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

PyTorch 默认不支持低精度 dtype 或不同的打包格式。通过 Tensor Subclass，我们扩展了 PyTorch 原生 Tensor 抽象，并将模型量化作为 dtype 转换，而自定义内核的不同打包格式通过布局处理。例如，我们支持使用 int4 权重（以 Tensor Core 友好布局打包）的量化线性操作，并使用 tinygemm 或 GemLite 内核实现。更多详细信息可在此处找到：https://pytorch.ac.cn/ao/stable/contributor_guide.html。

flow diagram

除了为开发人员提供更原生的 PyTorch 抽象之外，我们还希望强调这种设计对建模用户的两个好处。

  1. 序列化：像浮点模型一样将量化权重保存和加载到 state_dict 中，无需在加载量化权重之前将浮点模型转换为量化模型。这减少了分发和部署量化模型的摩擦。
  2. 可组合性：与张量并行等下游功能无缝集成，允许用户专注于建模，而无需担心与张量并行、torch.compile 和其他 PyTorch 功能的兼容性。由于这些功能是通过张量级抽象实现的，因此用户可以在大多数情况下进行量化和分布式推理，而无需更改模型。

GemLite 内核集成

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

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

创建此集成的主要困难是确保 GemLite 量化内核选项的整个范围都满足 TorchAO 的可组合性保证。虽然主要集成相对简单，但确保每种不同的量化类型及其相关内核与张量并行良好地配合并非易事。

Torch 张量并行

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

equation

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

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

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

张量并行和量化顺序

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

TorchAO 中的张量并行支持

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

4. SGLang

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

SGLang 中的 TorchAO 集成

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

SGLang 中 PyTorch 原生张量并行支持

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

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

其中 `qkv_proj` 和 `o_proj` 是 `wqkv` 和 `wo` 投影的 FQN，值是它们的 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 大小的解码 tokens/秒，每个图表显示了不同批处理大小和不同量化类型的结果。

  • BF16 是我们的 bfloat16，经过 torch.compile 编译的基线
  • 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` 量化，激活和权重均使用每行比例进行量化
bar chart
bar chart
bar chart

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

Float8 动态量化在张量并行大小为 1 的情况下，在不同批处理大小下持续比 bfloat16 加速 1.3 倍，在更大的张量并行大小下加速 1.1 倍至 1.2 倍。随着张量并行大小的增加，整体加速会降低，这是由于矩阵乘法大小的减小而预期的。请注意，我们也确实期望在预填充方面获得加速，但由于我们依赖 torch.compile 来加速，而 SGLang 中尚未启用预填充编译，因此我们将此留待未来工作。

复现说明

我们使用 GemLite 0.4.1、从 feb2b76 提交构建的 SGLang、TorchAO nightly 0.8.0.dev20241223+cu124 和 PyTorch 2.5.1 在 8xH100 机器上进行了基准测试。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，我们展示了针对 int4 和 float8 的快速端到端量化推理，涵盖了不同批处理大小和张量并行大小，同时提供了简单且可组合的用户 API，以降低 LLM 的资源需求。此次集成是我们满足跨不同模型、工作负载、精度和硬件的快速推理需求的第一步，我们期待继续推动端到端混合和低精度 LLM 推理的最新进展。

我们当前的未来工作重点如下：

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

我们还邀请社区积极测试、提供反馈并为塑造快速高效的 LLM 推理的未来做出贡献。