KTransformers AMD CPU 加速指南:基于 BLIS/AOCL LPGEMM 的 MOE_INT8 Prefill 优化 📅 发布时间:2026/9/13 15:20:08 👁 浏览次数: KTransformers AMD CPU 加速指南基于 BLIS/AOCL LPGEMM 的 MOE_INT8 Prefill 优化【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers本文是 KTransformersktransformers项目中 AMD 平台 CPU 推理加速的专项指南。核心内容源自 doc/en/kt-kernel/amd_blis.md并结合作业区kt-kernel/下真实源码CMake 构建逻辑、AOCL 内核封装、安装脚本与 Python 封装层进行纵深讲解。读者学完后可以独立完成 BLIS 库的构建安装、环境变量配置与MOE_INT8推理链路启用并理解这套加速方案在 prefill预填充阶段的底层原理。一、功能定位为什么 AMD CPU 需要 BLISKTransformers 的 MoEMixture-of-Experts内核默认针对 Intel 平台的 AMX / AVX512 指令集做了深度优化。而 AMD 服务器 CPU如 EPYC 系列尤其是 Zen 4虽然也支持 AVX512但缺少 AMX 指令直接复用 Intel 优化路径无法发挥硬件潜力。amd_blis.md明确说明了该特性的动机Motivation为了加速 AMD 平台的 prefill预填充速度方案是引入 AMD 官方维护的 BLIS 库并启用其中的LPGEMMLow-Precision GEMM扩展能力。BLISBLAS-like Library Instantiation Software Framework是 AMD 参与维护的下一代 BLAS 实现其 LPGEMM 接口专为低精度INT8 等矩阵乘设计能够在 AMD CPU 上获得接近理论峰值的 int8 矩阵乘算力正好匹配 MoE 专家计算中s8s8s32int8 输入 × int8 权重 → int32 累加的典型数据流。从仓库的 kt-kernel/README.md 也可以看到AMD BLIS 后端被明确列为受支持的加速路径对应 int8 prefill 与 decode并且 README 中单独提示 AMD BLIS 后端用户需要参考 AMD 特定的安装配置。二、使用前提与三步开启流程原文档给出的USAGE非常简单清晰共三个步骤使用MOE_INT8量化方法即运行时指定--kt-method MOE_INT8构建并安装正确的 AMD BLIS 库必须参考下方 Motivation 一节说明的构建方式编译带 LPGEMMaocl_gemmadd-on支持的 BLIS 版本安装前导出环境变量执行export CPUINFER_ENABLE_BLISON来启用该后端。2.1 第一步--kt-method MOE_INT8MOE_INT8是 KTransformers 提供的通用 MoE 量化方法之一。在 Python 封装层 kt-kernel/python/utils/moe_kernel.py 中GeneralMoEWrapper明确支持MOE_INT4和MOE_INT8两种方法method: str MOE_INT8,该封装在初始化时会对后端可用性做硬校验——如果编译产物中不包含 int8 内核传入MOE_INT8会直接抛出RuntimeError提示重新编译if not _HAS_INT8_SUPPORT and method MOE_INT8: raise RuntimeError( MoE_INT8 backend not available. kt_kernel_ext was not compiled with int8 support.\n Please recompile with int8 enabled. )其中Int8_KERNEL_MOE从kt_kernel_ext.moe导入即由 C/CUDA 扩展在编译期根据开关生成。这从侧面印证BLIS 加速与 MOE_INT8 是强绑定的int8 是 LPGEMM 的数据类型基础。2.2 第二步构建带 LPGEMM 的 AMD BLIS 库原文档指出参考 AMD 的 BLIS 仓库构建时需要启用 LPGEMM 支持对应 AMD AOCL BLAS 用户指南中描述的lpgemm以及 BLIS 构建系统文档中的aocl_gemmadd-on。核心要点BLIS 的 CMake 构建系统提供了aocl_gemmadd-on开关打开后才会编译出aocl_gemm_*系列低精度 GEMM 接口构建完成后库中应能看到类似aocl_gemm_s8s8s32os32、aocl_reorder_s8s8s32os32、aocl_get_reorder_buf_size_s8s8s32os32这类符号它们是 KTransformers AOCL 内核直接依赖的 API。说明由于本文只引用当前仓库证据具体的 BLIS 源码获取与完整编译参数请以 BLIS 官方构建文档为准仓库侧对 BLIS 的查找与链接逻辑见下文第三节可作为安装后验证的对照。2.3 第三步export CPUINFER_ENABLE_BLISON该环境变量是整个开关链路的总闸必须在安装编译 kt-kernel 扩展之前导出。它被 kt-kernel/setup.py 读取并转发为 CMake 选项同时触发一系列连锁的自动配置详见第三节。安装脚本 kt-kernel/install.sh 中CPUINFER_ENABLE_BLIS在 ARM 自动检测分支里默认被置为OFF:${CPUINFER_ENABLE_BLIS:OFF}即默认不开启需要用户显式导出为ON这与原文档的提示完全一致。三、源码视角CPUINFER_ENABLE_BLISON之后的完整链路3.1 环境变量 → CMake 选项的转发在 kt-kernel/setup.py 中CPUINFER_ENABLE_BLIS被映射到 CMake 开关KTRANSFORMERS_CPU_MOE_AMD_forward_bool_env(cmake_args, CPUINFER_ENABLE_BLIS, KTRANSFORMERS_CPU_MOE_AMD)更关键的是setup.py 还会自动联动开启 MoE 内核开关当CPUINFER_ENABLE_BLIS或CPUINFER_ENABLE_KML为 ON 时自动追加-DKTRANSFORMERS_CPU_MOE_KERNELONamd_env _env_get_bool(CPUINFER_ENABLE_BLIS, None) kml_env _env_get_bool(CPUINFER_ENABLE_KML, None) if amd_env or kml_env: auto_moe_kernel_ True ... if not already_set and auto_moe_kernel_: cmake_args.append(-DKTRANSFORMERS_CPU_MOE_KERNELON)也就是说导出CPUINFER_ENABLE_BLISON一个变量会同时激活「AMD MoE 后端」和「通用 MoE 内核」两个编译维度无需手动再开第二个开关。3.2 CMake 侧的 BLIS 查找与链接kt-kernel/CMakeLists.txt 在KTRANSFORMERS_CPU_MOE_AMD开启时执行 BLIS 探测通过BLIS_ROOT缓存变量指定 BLIS 安装根目录默认搜索路径为/usr/local与/usr用find_path查找blis.h头文件、find_library查找 BLIS 动态/静态库若找不到会输出警告BLIS not found; set BLIS_ROOT or specify BLIS_INCLUDE_DIR/BLIS_LIBRARY找到后记录_KT_BLIS_INCLUDE_DIR与_KT_BLIS_LIBRARY在最终的目标${PROJECT_NAME}创建完成后将头文件目录与库链接到扩展产物上target_include_directories(${PROJECT_NAME} PRIVATE ${_KT_BLIS_INCLUDE_DIR}) target_link_libraries(${PROJECT_NAME} PRIVATE ${_KT_BLIS_LIBRARY})排查提示如果编译时出现BLIS not found警告说明库没有被 CMake 找到。此时应确认 BLIS 是否安装在/usr/local、/usr下或通过-DBLIS_ROOT/path/to/blissetup.py 层面对应CPUINFER_BLIS_ROOT显式指定安装位置。同时KTRANSFORMERS_CPU_MOE_AMD开启后CMake 会把 AOCL 内核目录编入工程并定义宏if(KTRANSFORMERS_CPU_MOE_AMD) aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/operators/moe_kernel/mat_kernel/aocl_kernel SOURCE_DIR7_KERNEL) add_compile_definitions(USE_MOE_KERNEL_AMD1) endif()3.3 AOCL 内核封装LPGEMM 是如何被调用的KTransformers 的 MoE 内核通过统一接口 kt-kernel/operators/moe_kernel/mat_kernel/batch_gemm_api.hpp 声明了五个核心函数函数用途decode_cblas_gemm_s8s8s32decode逐 token 生成阶段的 int8 GEMMprefill_cblas_gemm_s8s8s32prefill预填充阶段的 int8 GEMMdecode_int4_cblas_gemm_s8s8s32decode 阶段 int4 路径prefill_int4_cblas_gemm_s8s8s32prefill 阶段 int4 路径reorder_B_gemm/get_reorder_B_size权重 B 矩阵的重排pack与缓冲大小查询AMD 的实现位于 kt-kernel/operators/moe_kernel/mat_kernel/aocl_kernel/kernel.cpp其中prefill_cblas_gemm_s8s8s32与decode_cblas_gemm_s8s8s32最终都调用同一个底层函数aocl_gemm_s8s8s32os32(order, op_a, op_b, m, n, k, alpha, a, lda, n, b, ldb, r, beta, c, ldc, nullptr);这正是 BLISaocl_gemmadd-on 提供的低精度 GEMM 接口s8s8s32offset 为标量 os32 形式。调用前布局与转置标志被映射为 AOCL 约定的小写字符r/c表示行/列主序n/t表示是否转置。值得注意的两个实现细节int4 路径不支持prefill_int4_cblas_gemm_s8s8s32与decode_int4_cblas_gemm_s8s8s32直接抛出std::runtime_error(int4 not support prefill)/(int4 not support decode)说明AOCL/BLIS 后端仅面向 int8MOE_INT8与文档「必须使用 MOE_INT8 方法」的要求严格对应B 矩阵重排reorder_B_gemm调用aocl_reorder_s8s8s32os32对权重 B 做布局重排get_reorder_B_size查询所需缓冲大小——这是 BLIS 低精度内核典型的「先 pack 权重、后批量 GEMM」优化模式权重重排可提前完成prefill 期间直接复用重排后的数据。3.4 独立验证程序simple_test_aocl仓库提供了一个独立的 AOCL 功能验证程序 kt-kernel/demo/simple_test_aocl.cpp直接#include blis.h其流程完整演示了 BLIS LPGEMM 的使用方式用bli_malloc_user分配输入矩阵 A、B 与输出 C将 B 从列主序转换为行主序通过aocl_get_reorder_buf_size_s8s8s32os32(c,n,B, kK, kN)查询重排缓冲大小并分配调用aocl_reorder_s8s8s32os32(c,n,B, ...)重排 B分别用重排后的 B与未重排的 B调用aocl_gemm_s8s8s32os32与纯标量三重循环的参考实现对比输出AOCL GEMM output matches reference.之类的结果。该 demo 中的维度kM3, kK7168, kN2048与真实 MoE 专家网络形状如 7168 维中间层一致可当作安装 BLIS 后验证「库是否可用、LPGEMM 是否生效」的最小测试程序。从 kt-kernel/CMakeLists.txt 可以看到该 demo 编译依赖 BLIS 头文件与库只有 BLIS 探测成功才会被构建。四、构建安装速查整合步骤综合原文档与仓库安装逻辑完整的开启流程如下# 1. 导出启用开关必须在编译安装 kt-kernel 之前 export CPUINFER_ENABLE_BLISON # 2.可选如果 BLIS 安装位置非默认指定根目录 # export CPUINFER_BLIS_ROOT/path/to/your/blis # 对应 CMake 的 BLIS_ROOT # 3. 初始化子模块并创建环境 git submodule update --init --recursive conda create -n kt-kernel python3.11 -y conda activate kt-kernel # 4. 构建安装 kt-kernel自动探测 CPU 并链接 BLIS ./install.sh # 或分步./install.sh deps ./install.sh build安装后从 Python 侧确认 int8 后端可用并验证扩展加载python -c from kt_kernel import KTMoEWrapper; print(✓ kt-kernel installed successfully)如果安装脚本输出BLIS not found相关警告请回到第三节 3.2 的排查提示检查库搜索路径。最终运行模型时确保推理侧以--kt-method MOE_INT8启动使 MoE 专家计算走 BLIS int8 路径。五、适用前提与边界说明仅面向 AMD CPU该后端面向 AMD 平台编译期KTRANSFORMERS_CPU_MOE_AMD、运行期USE_MOE_KERNEL_AMD宏。Intel 平台建议继续使用 AMX / AVX512 路径仅支持 int8AOCL 内核的 int4 分支为显式抛错实现使用 BLIS 加速时必须搭配MOE_INT8环境变量时机敏感CPUINFER_ENABLE_BLISON必须在编译安装阶段生效若仅修改推理参数而未重新编译不会启用该后端依赖外部库本方案强依赖系统安装的、带aocl_gemmLPGEMM支持的 AMD BLIS 库仓库本身不内置该库只负责查找、链接与调用构建产物与 CPU 绑定按 kt-kernel/README.md 的说明默认构建针对本机 CPU 优化-marchnative分发或迁移到其他机器时需注意兼容性。六、小结AMD BLIS 加速方案是 KTransformers 为 AMD CPU 补齐 prefill 性能短板的关键路径以MOE_INT8为数据契约以 BLIS 的 LPGEMMaocl_gemm_s8s8s32os32为计算核心通过CPUINFER_ENABLE_BLISON一个环境变量串联起 CMake 探测、AOCL 内核编译、权重重排与批量 GEMM 的完整链路。开发者只需三步——选对方法、装对库、开对开关即可让 AMD 平台上的 MoE 专家计算获得低精度矩阵乘的硬件级加速相关源码实现可在 aocl_kernel/kernel.cpp 与 simple_test_aocl.cpp 中进一步研读。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考