自定义内核配置

相关源文件

• docs/codegen.md
• setup_env.py
• utils/codegen_tl2.py

本页面解释了如何在 BitNet.cpp 中为不同的模型和硬件架构自定义和优化内核参数。内核配置是实现最佳性能的关键组成部分，因为它决定了矩阵操作如何在目标硬件上分区和执行。

有关通用性能调优的信息，请参阅性能调优。

1. 内核配置概述

BitNet.cpp 使用优化的内核实现来加速三元权重模型的推理。这些内核是根据模型架构和目标硬件平台动态生成的。

• TL1 内核：针对 ARM64 架构优化
• TL2 内核：针对 x86_64 架构优化

内核生成过程通过将矩阵乘法操作划分为更小的块来创建优化的实现，这些块能更好地利用硬件能力，包括缓存层次结构和 SIMD 指令。

来源：setup_env.py141-170 docs/codegen.md1-7

2. 关键配置参数

BitNet.cpp 的性能高度依赖于三个关键参数，这些参数决定了矩阵如何分区。

参数	描述	约束
BM	沿 M 维度切割权重矩阵的块大小	必须满足 M % BM = 0
BK	沿 K 维度切割权重矩阵的块大小	对于 TL1：K % BK = 0 对于 TL2：K % BK % 32 = 0
bm	计算块的块大小	对于 TL1：从 [32, 64] 中选择 对于 TL2：必须为 32

必须针对每个模型架构和硬件平台仔细调整这些参数，以实现最佳性能。

来源：docs/codegen.md16-49 utils/codegen_tl2.py726-730

3. 矩阵分区策略

内核生成过程为 TL1 和 TL2 内核实现了不同的分区策略。

3.1 TL1 (ARM) 分区策略

对于 TL1 内核，权重矩阵被划分为大小为 BM×K 的块，然后进一步细分为 BM×BK 块。每个 BM×BK 块都在大小为 bm 的计算块中处理。

来源：docs/codegen.md29-39

3.2 TL2 (x86) 分区策略

对于 TL2 内核，方法更为复杂。K 维度被拆分为“threeK”（用 TL2 处理）和“twoK”（用备用方法处理）。TL2 实现要求 BK 必须能被 6 整除，以实现最佳 AVX2 指令利用。

来源：docs/codegen.md40-49 utils/codegen_tl2.py676-681

4. 配置自定义内核

在 BitNet.cpp 中配置自定义内核有两种方法：

4.1 使用 --use-pretuned 选项

BitNet.cpp 包含针对受支持模型的预调优内核配置。要使用这些配置：

此方法会从 preset_kernels 目录复制预调优内核配置。

来源：setup_env.py150-160 setup_env.py172-178

4.2 自定义配置

对于自定义配置，直接指定 BM、BK 和 bm 参数。

脚本将根据模型架构和目标平台自动选择合适的内核参数。

来源：setup_env.py212-221 setup_env.py147-188

5. 模型特定配置

BitNet.cpp 包含针对不同模型架构的优化配置。以下是受支持模型使用的默认配置：

5.1 ARM64 (TL1) 配置

模型	BM 值	BK 值	bm 值
bitnet_b1_58-large	256,128,256	128,64,128	32,64,32
bitnet_b1_58-3B	160,320,320	64,128,64	32,64,32
BitNet-b1.58-2B-4T	160,320,320	64,128,64	32,64,32
Llama3/Falcon3	256,128,256,128	128,64,128,64	32,64,32,64

5.2 x86_64 (TL2) 配置

模型	BM 值	BK 值	bm 值
bitnet_b1_58-large	256,128,256	96,192,96	32,32,32
bitnet_b1_58-3B	160,320,320	96,96,96	32,32,32
BitNet-b1.58-2B-4T	160,320,320	96,96,96	32,32,32
Llama3/Falcon3	256,128,256,128	96,96,96,96	32,32,32,32

来源：setup_env.py161-169 setup_env.py179-188

6. 高级自定义配置

对于希望手动调优内核的高级用户，您可以直接使用 codegen 脚本。

此方法提供了对模型中每个矩阵形状的内核参数的完全控制。

来源：docs/codegen.md18-27 utils/codegen_tl2.py696-705

7. 配置存储

内核配置存储在两个文件中：

1. bitnet-lut-kernels.h：包含生成的内核代码
2. kernel_config.ini：包含配置参数

内核配置文件以 INI 格式存储，其中包含模型中每个内核形状的节。

来源：utils/codegen_tl2.py736-756 setup_env.py202-204

8. 性能考量

自定义内核配置可以显著影响性能。

• 最佳 BM 值应平衡 SIMD 指令宽度和缓存利用率。
• 最佳 BK 值取决于架构的 SIMD 能力（ARM NEON 与 x86 AVX2）。
• 最佳 bm 值取决于可用寄存器数量和指令延迟。

寻找最佳配置可能需要对新的模型架构或硬件平台进行实验。

来源：docs/codegen.md29-49