内核生成系统

相关源文件

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

目的与范围

内核生成系统负责为不同的硬件架构生成优化的查找表（lookup table）内核代码，以加速BitNet.cpp中三元（1比特）模型的推理。该系统创建了平台特定的实现，以最大限度地提升在ARM (TL1) 和x86 (TL2) 处理器上的性能。

本页重点介绍内核生成过程和矩阵分区策略。有关运行推理的详细信息，请参见运行推理；有关模型转换的信息，请参见模型转换流程。

架构概述

来源：setup_env.py144-188 utils/codegen_tl2.py683-757

矩阵分区策略

内核生成系统根据目标架构采用不同的矩阵分区策略来优化性能。

适用于ARM架构的TL1

TL1实现了三级分区策略

1. 原始权重矩阵 (M×K) 被分为 M/BM 个块，每个块的大小为 (BM×K)
2. 每个块进一步切分为 K/BK 个大小为 (BM×BK) 的块
3. 每个计算块以更小的单位 (bm×compute_num/bm) 进行处理

来源：docs/codegen.md29-38

适用于x86架构的TL2

TL2为x86采用了专门的策略

1. K维度被分成两部分：threeK（可被6整除）和twoK（余数）
2. threeK 部分使用x86特定的优化进行处理
3. twoK 部分作为备用方案使用类似TL1的方法处理

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

代码生成过程

内核生成过程始于使用模型参数执行 setup_env.py。

来源：setup_env.py144-188 utils/codegen_tl2.py532-624 utils/codegen_tl2.py626-673

生成的内核组件

bitnet-lut-kernels.h 中生成的内核代码包含几个关键组件

1. 构造函数代码：用于创建查找表和张量级量化的函数

   • per_tensor_quant: 量化浮点张量
   • three_lut_ctor/two_lut_ctor: 构造三元权重的查找表

2. 实现模板：架构特定的优化矩阵乘法

   • three_tbl_impl_*: 3值三元权重的实现
   • two_tbl_impl_*: 2值二元权重的实现

3. 顶层API：外部代码访问优化内核的函数

   • ggml_preprocessor: 为推理准备查找表
   • ggml_qgemm_lut: 优化矩阵乘法的入口点

4. 张量变换：为优化计算准备张量

   • ggml_bitnet_transform_tensor: 为BitNet处理转换标准张量

来源：utils/codegen_tl2.py5-277 utils/codegen_tl2.py279-531 utils/codegen_tl2.py532-624 utils/codegen_tl2.py626-673

模型特定参数

不同的模型需要不同的分区参数以实现最佳性能。该系统支持多种模型配置

模型	架构	矩阵形状	BM 值	BK 值	bm 值
bitnet_b1_58-large	ARM (TL1)	1536×4096, 1536×1536, 4096×1536	256,128,256	128,64,128	32,64,32
bitnet_b1_58-large	x86 (TL2)	1536×4096, 1536×1536, 4096×1536	256,128,256	96,192,96	32,32,32
bitnet_b1_58-3B	ARM (TL1)	3200×8640, 3200×3200, 8640×3200	160,320,320	64,128,64	32,64,32
bitnet_b1_58-3B	x86 (TL2)	3200×8640, 3200×3200, 8640×3200	160,320,320	96,96,96	32,32,32
Llama3-8B-1.58	ARM (TL1)	14336×4096, 4096×14336, 1024×4096, 4096×4096	256,128,256,128	128,64,128,64	32,64,32,64
Llama3-8B-1.58	x86 (TL2)	14336×4096, 4096×14336, 1024×4096, 4096×4096	256,128,256,128	96,96,96,96	32,32,32,32

来源：setup_env.py162-168 setup_env.py179-186 utils/codegen_tl2.py683-694

参数要求和约束

内核生成系统强制执行特定要求以确保最佳性能

对于TL1 (ARM)

• M % BM == 0 (M 必须可被 BM 整除)
• K % BK == 0 (K 必须可被 BK 整除)
• BM % bm == 0 (BM 必须可被 bm 整除)
• bm 必须从 [32, 64] 中选择

对于TL2 (x86)

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

系统在生成代码之前会执行验证检查

来源：utils/codegen_tl2.py726-729 docs/codegen.md34-38 docs/codegen.md45-49

与BitNet.cpp系统的集成

内核生成系统与更大的BitNet.cpp工作流相契合

来源：setup_env.py206-210 setup_env.py144-188 setup_env.py190-204 setup_env.py97-138

使用示例

要为特定模型手动生成内核，可以使用以下命令

对于ARM (TL1)

python utils/codegen_tl1.py --model bitnet_b1_58-large --BM 256,128,256 --BK 128,64,128 --bm 32,64,32

对于x86 (TL2)

python utils/codegen_tl2.py --model bitnet_b1_58-large --BM 256,128,256 --BK 96,192,96 --bm 32,32,32

通常，这些命令会由 setup_env.py 根据检测到的架构和正在使用的模型自动执行。

来源：setup_env.py162-168 setup_env.py179-186 docs/codegen.md22-27