免费获取学习方案
ARTICLE DETAIL

资讯详情

深耕编程基础知识与建站技术分享的一线实战洞察。

如何编译 BitNet W2A8 CUDA 内核并用 test.py 运行 GEMV 加速测试?

如何编译 BitNet W2A8 CUDA 内核并用 test.py 运行 GEMV 加速测试? 如何编译 BitNet W2A8 CUDA 内核并用 test.py 运行 GEMV 加速测试【免费下载链接】BitNetOfficial inference framework for 1-bit LLMs项目地址: https://gitcode.com/GitHub_Trending/bitne/BitNetBitNet 仓库的 gpu/ 目录提供了一组针对 W2A8 推理2-bit 权重 × 8-bit 激活的自定义 CUDA GEMV 内核并配有现成的基准测试脚本。如果你的目标是验证这组内核相对 torch BF16 矩阵乘是否真的有加速流程只有两步先编译出libbitnet.so再在gpu/目录下运行python test.py。前提是一台可用的 NVIDIA GPU——compile.sh 的编译目标是compute_80Ampere 架构gpu/README.md 中的基准数据也是在 NVIDIA A100 40GB 上测得的。准备环境按 gpu/README.md 的说明用 conda 创建一个 Python 版本低于 3.13 的环境并安装 gpu/requirements.txt 中的依赖包含torch2.2.0、xformers0.0.22、transformers等conda create --name bitnet-gpu python3.13 conda activate bitnet-gpu # 以下命令默认在仓库的 gpu/ 目录下执行 pip install -r requirements.txt还有三点依赖需要注意都来自对 test.py 源码的直接检查test.py通过devicecuda生成张量并取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 -stdc17 -Xcudafe --diag_suppress177 --compiler-options -fPIC -lineinfo --shared bitnet_kernels.cu -lcuda -gencodearchcompute_80,codecompute_80 -o libbitnet.so几个直接影响能否跑通 test.py 的细节-shared加-o libbitnet.so表示编译成共享库产物输出在当前目录-gencodearchcompute_80,codecompute_80指定目标为 compute capability 8.0Ampere如 A100test.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.pyCUDAExtension 打包形式但 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 Ratio2560 × 256013.3218.321.383840 × 256014.9018.871.2713824 × 256018.7559.513.172560 × 691214.4937.782.613200 × 320014.6119.081.314800 × 320013.0921.841.673200 × 1024019.6460.793.1020480 × 320030.99112.393.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_80README 的对比表基于 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),仅供参考
返回列表