在 Devin 中试用 DeepWiki 处理私有仓库
DeepWiki

microsoft/BitNet

菜单

TL1 和 TL2 优化

相关源文件

本文档详细介绍了 BitNet.cpp 中实现的基于查找表的专门优化，旨在实现 1 位三元大语言模型的高效推理。TL1 (ARM) 和 TL2 (x86) 是平台特定的优化，可显著提高三元权重矩阵乘法的性能。有关整体内核生成系统的信息，请参阅内核生成系统。

1. 介绍与目的

TL1 和 TL2 是专门的矩阵乘法技术，它们通过基于查找表的计算，利用三元权重（-1, 0, 1）的特性。这些优化会根据目标平台自动选择。

  • TL1：针对 ARM 架构（移动设备、Apple Silicon 等）优化。
  • TL2：针对支持 AVX2 指令的 x86 架构优化。

这两种优化通过降低计算复杂性并利用平台特定的 SIMD（单指令多数据）能力，显著提高了推理性能。

来源：docs/codegen.md1-4

2. 性能优势

平台特定的优化提供了显著的性能和效率提升。

优化平台性能提升能耗降低
TL1ARM641.37 倍-5.07 倍的加速55.4%-70.0%
TL2x86_642.37 倍-6.17 倍的加速71.9%-82.2%

这些提升是通过精心的矩阵分区、基于查找表的计算以及平台特定的 SIMD 指令利用实现的。

3. 优化方法

3.1 通用策略

TL1 和 TL2 都采用查找表（LUT）方法，通过预先计算可能的结果并通过查找索引访问它们，从而避免昂贵的乘法操作。优化过程包括：

  1. 将权重矩阵划分为可管理的分块
  2. 创建查找表以实现高效计算
  3. 实现带有 SIMD 指令的优化内核代码

实际实现因平台而异，以充分利用可用的硬件能力。

来源：utils/codegen_tl2.py76-189 docs/codegen.md29-49

3.2 TL1 优化 (ARM)

TL1 通过系统性的分区方法优化 ARM 架构的矩阵乘法：

  1. 将权重矩阵 (M×K) 切割成 M/BM 个大小为 (BM×K) 的块
  2. 进一步将每个块划分为 (BM×BK) 子块
  3. 使用大小为 (bm×compute_num/bm) 的计算块处理每个子块

这种方法最大限度地提高了缓存局部性，并能有效利用 ARM SIMD 指令。

TL1 实现需要特定的参数约束：

  • M % BM == 0 (M 必须可被 BM 整除)
  • K % BK == 0 (K 必须可被 BK 整除)
  • BM % bm == 0 (BM 必须可被 bm 整除)
  • bm 必须为 32 或 64

来源：docs/codegen.md29-38 assets/tl1.png

3.3 TL2 优化 (x86)

TL2 优化是为支持 AVX2 的 x86 架构设计的，具有更复杂的分区策略：

  1. 由于 TL2 要求 BK 必须可被 6 整除，因此将处理拆分为 "threeK" 和 "twoK" 部分
  2. 使用 (BM×BK) 块和 x86 特定优化处理 "threeK" 部分
  3. 使用类似于 TL1 的回退方法处理剩余的 "twoK" 部分

TL2 实现需要特定的参数约束：

  • M % BM == 0 (M 必须可被 BM 整除)
  • K % BK % 32 == 0 (K 对 BK 的模数必须可被 32 整除)
  • BM % bm == 0 (BM 必须可被 bm 整除)
  • bm 必须为 32

来源：docs/codegen.md40-49 assets/tl2.png utils/codegen_tl2.py190-267

4. 代码生成系统

TL1 和 TL2 的优化内核由 codegen_tl1.py（用于 ARM）和 codegen_tl2.py（用于 x86）动态生成。这种代码生成方法使得系统能够：

  1. 生成高度优化、平台特定的代码
  2. 适应不同的模型架构和大小
  3. 在各种硬件上保持最佳性能

代码生成过程包括：

  1. 参数配置：指定模型形状 (M, K)、块大小 (BM, BK) 和计算块大小 (bm)
  2. 内核生成：根据目标平台创建优化的实现代码
  3. 配置存储：将配置保存到 kernel_config.ini 以供后续参考

来源：utils/codegen_tl2.py682-757 docs/codegen.md17-27

4.1 生成的函数

生成代码包括几种关键函数类型：

  1. 查找表构造函数:

    • three_lut_ctor<act_k>：为 3 值权重创建查找表
    • two_lut_ctor<act_k>：为 2 值权重创建查找表

  2. 基于表的实现函数:

    • three_tbl_impl_{pre}<batch_size, K3>：使用 3 值表处理矩阵
    • two_tbl_impl{pre}<batch_size, K2>：使用 2 值表处理矩阵

  3. QGEMM 封装函数:

    • three_qgemm_lut_{pre}<BATCH_SIZE>：3 值计算的高级封装函数
    • two_qgemm_lut_{pre}<BATCH_SIZE>：2 值计算的高级封装函数

  4. 顶层 API 函数:

    • ggml_preprocessor：为计算准备查找表
    • ggml_qgemm_lut：使用查找表的主矩阵乘法函数
    • ggml_bitnet_transform_tensor：转换张量以与优化内核配合使用

来源：utils/codegen_tl2.py107-189 utils/codegen_tl2.py279-528 utils/codegen_tl2.py532-624 utils/codegen_tl2.py626-672

5. 实现细节

5.1 查找表构建

查找表构建是 TL1 和 TL2 优化的关键组成部分。这些表根据权重模式预先计算所有可能的输出值。

特别是对于 TL2，查找表构建包括：

  1. 使用比例因子将权重值量化为整数
  2. 计算三元权重（-1, 0, 1）的所有可能结果组合
  3. 组织和转置表以实现高效的 SIMD 访问

来源：utils/codegen_tl2.py76-96 utils/codegen_tl2.py106-267

5.2 SIMD 优化计算

TL2 计算路径利用 AVX2 SIMD 指令进行高效并行处理。

主要的 AVX2 操作包括：

  • _mm256_loadu_si256：加载 256 位向量
  • _mm256_shuffle_epi8：执行表查找
  • _mm256_add_epi16：并行添加 16 位整数
  • _mm256_cvtepi16_epi32：将 16 位整数转换为 32 位整数
  • _mm256_storeu_si256：存储 256 位向量

来源：utils/codegen_tl2.py286-422 utils/codegen_tl2.py432-486

6. 在 BitNet.cpp 中的使用

TL1 和 TL2 优化会根据目标架构由 BitNet.cpp 框架自动利用。典型的工作流程是：

  1. 环境设置：setup_env.py 检测平台并生成相应的内核
  2. 代码生成：codegen_tl1.py 或 codegen_tl2.py 使用适当的参数创建优化内核
  3. 配置存储：参数保存到 kernel_config.ini
  4. 模型准备：模型被转换并为推理做好准备
  5. 推理：优化内核在模型推理过程中使用

6.1 参数示例

不同模型的参数示例：

  1. BitNet-b1.58-large (0.7B):

    • TL1：--BM 256,128,256 --BK 128,64,128 --bm 32,64,32
    • TL2：--BM 256,128,256 --BK 96,192,96 --bm 32,32,32

  2. BitNet-b1.58-3B (3.3B):

    • 矩阵形状：[3200, 8640], [3200, 3200], [8640, 3200]
    • 参数需要满足 TL1 或 TL2 的约束。

  3. Llama3-8B-1.58 (8B):

    • 矩阵形状：[14336, 4096], [4096, 14336], [1024, 4096], [4096, 4096]
    • 参数需要满足 TL1 或 TL2 的约束。

来源：docs/codegen.md18-27 utils/codegen_tl2.py683-694

7. 结论

TL1 和 TL2 优化是 BitNet.cpp 的核心组件，分别在 ARM 和 x86 平台上实现了三元权重模型的高效推理。通过使用基于查找表的计算和平台特定的 SIMD 指令，这些优化与传统方法相比，实现了显著的性能提升（高达 6.17 倍加速）和能效提升（高达 82.2% 的降低）。

动态代码生成系统使这些优化能够适应不同的模型架构，同时保持最佳性能。矩阵分区策略、查找表技术和 SIMD 加速的这种结合，使得 BitNet.cpp 成为在消费级硬件上运行 1 位三元语言模型的高效框架。