如何编译 BitNet W2A8 CUDA 内核并用 test.py 运行 GEMV 加速测试？
【免费下载链接】BitNetOfficial inference framework for 1-bit LLMs项目地址: https://gitcode.com/GitHub_Trending/bitne/BitNet
BitNet 仓库的 gpu/ 目录提供了一组针对 W2A8 推理（2-bit 权重 × 8-bit 激活）的自定义 CUDA GEMV 内核，并配有现成的基准测试脚本。如果你的目标是验证这组内核相对 torch BF16 矩阵乘是否真的有加速，流程只有两步：先编译出libbitnet.so，再在gpu/目录下运行python test.py。前提是一台可用的 NVIDIA GPU——compile.sh 的编译目标是compute_80（Ampere 架构），gpu/README.md 中的基准数据也是在 NVIDIA A100 40GB 上测得的。
准备环境
按 gpu/README.md 的说明，用 conda 创建一个 Python 版本低于 3.13 的环境，并安装 gpu/requirements.txt 中的依赖（包含torch>=2.2.0、xformers>=0.0.22、transformers等）：
conda create --name bitnet-gpu "python<3.13" conda activate bitnet-gpu # 以下命令默认在仓库的 gpu/ 目录下执行 pip install -r requirements.txt还有三点依赖需要注意，都来自对 test.py 源码的直接检查：
test.py通过device='cuda'生成张量并取torch.cuda.current_stream()，因此需要可用的 NVIDIA GPU 和对应 CUDA 环境的 PyTorch；- 除 requirements.txt 中的包外，
test.py还import numpy，环境里缺 numpy 时需要先补装； test.py依赖同目录下的本地模块 pack_weight.py，它负责权重的 16×32 块重排、2-bit 压缩打包和位交织，不要移动该文件。
编译 W2A8 CUDA 内核
cd bitnet_kernels bash compile.sh cd ..compile.sh 的内容就是单条 nvcc 命令：
nvcc -std=c++17 -Xcudafe --diag_suppress=177 --compiler-options -fPIC -lineinfo --shared bitnet_kernels.cu -lcuda -gencode=arch=compute_80,code=compute_80 -o libbitnet.so几个直接影响能否跑通 test.py 的细节：
-shared加-o libbitnet.so表示编译成共享库，产物输出在当前目录；-gencode=arch=compute_80,code=compute_80指定目标为 compute capability 8.0（Ampere，如 A100）；- test.py 通过
ctypes.CDLL('bitnet_kernels/libbitnet.so')加载库，相对路径是bitnet_kernels/libbitnet.so——所以 compile.sh 必须在bitnet_kernels/目录内执行、且test.py必须从gpu/目录启动，否则找不到编译产物。
内核源码 bitnet_kernels.cu 的入口bitlinear_int8xint2按 (N, K) 形状硬编码了各形状专用的ladder_int8xint2_kernel模板实例，且只处理M == 1的 GEMV 调用；遇到未特化的形状，它只向屏幕打印required ladder gemm kernel: M ... N ... K ...提示，不会执行计算。bitnet_kernels/下另有一个 setup.py（CUDAExtension 打包形式），但 README 给出的主路径是 compile.sh。
README 的 “Optimizations” 一节说明了这组内核相对常规实现的三个优化点：权重矩阵按 16×32 块重排以优化访存、16 个 2-bit 值按交错模式打包进一个 32-bit 整数以加速解码，以及使用dp4a指令完成低精度点积累加。
运行 test.py 做 GEMV 基准测试
回到gpu/目录执行：
python test.py脚本内部固定了 8 组测试形状（N×K）：(2560, 2560)、(3840, 2560)、(13824, 2560)、(2560, 6912)、(3200, 3200)、(4800, 3200)、(3200, 10240)、(20480, 3200)，每一组都分两个阶段：
- 正确性检查：生成随机 int8 激活和权重，权重经
pack_weight.convert_weight_int8_to_int2压缩打包后调用内核，结果与 numpy 的 int32 矩阵乘逐元素比较，打印一行custom == np True或custom == np False； - 计时对比：用
torch.utils.benchmark.Timer分别对 W2A8 内核和torch.matmul(input0_bf16, weight_bf16)各计时 50 次（timeit(50)），按Shape(N, K), W2A8: <均值>us, torch BF16: <均值>us的格式打印。
这 8 组形状与 bitnet_kernels.cu 中的内核特化一一对应，因此按文档流程执行时每个形状都会真正走内核计算，而不会落到“unsupported shape”的提示分支。
如何判断结果
运行输出的两类行分别对应两种判断：
- 正确性行应当是
custom == np True，即内核输出与 CPU 上的 numpy 整数矩阵乘完全一致； - 性能行给出 W2A8 内核与 torch BF16 基线的延迟均值，例如输出形如：
Shape(3200, 10240), W2A8: 19.64us, torch BF16: 60.79us注意这是文档中展示的结果格式，读者机器上跑出的具体数值取决于硬件和驱动环境。gpu/README.md 的 “Kernel Benchmarks” 表格给出了文档在 NVIDIA A100 40GB 上测得的参考数据：
| Shape (N×K) | W2A8 Latency (us) | BF16 Latency (us) | Speedup Ratio |
|---|---|---|---|
| 2560 × 2560 | 13.32 | 18.32 | 1.38 |
| 3840 × 2560 | 14.90 | 18.87 | 1.27 |
| 13824 × 2560 | 18.75 | 59.51 | 3.17 |
| 2560 × 6912 | 14.49 | 37.78 | 2.61 |
| 3200 × 3200 | 14.61 | 19.08 | 1.31 |
| 4800 × 3200 | 13.09 | 21.84 | 1.67 |
| 3200 × 10240 | 19.64 | 60.79 | 3.10 |
| 20480 × 3200 | 30.99 | 112.39 | 3.63 |
表中加速比从 1.27 到 3.63 不等，K 维度越大（即权重矩阵越宽的行数越多）加速越明显——这个规律可以直接用你自己跑出的两列数值核对。
限制与注意事项
- 该内核只覆盖
M == 1的 GEMV 场景，对应解码阶段的逐 token 生成计算；批量矩阵乘（M > 1）不在 bitnet_kernels.cu 的处理范围内，会落入提示分支； - 只有源码中硬编码的 (N, K) 形状会被实际计算，除 test.py 使用的 8 组外，源码还特化了
(5120, 27648)和(55296, 5120)两组； - 编译目标固定为
compute_80，README 的对比表基于 A100 40GB；在其他 GPU 上执行 test.py 得到的绝对延迟和加速比都会与表中不同，判断时以正确性行全为 True、再结合自己环境下的两列延迟为准。
如果内核验证通过，gpu/README.md 还给出了端到端路径：用huggingface-cli下载 BitNet-b1.58-2B 权重，经convert_safetensors.py和convert_checkpoint.py转换后用generate.py --interactive --chat_format做交互式推理，那一步会真正用上这里编译的内核。
【免费下载链接】BitNetOfficial inference framework for 1-bit LLMs项目地址: https://gitcode.com/GitHub_Trending/bitne/BitNet
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
