KTransformers

AMX 后端

KT-Kernel 的 AMX 后端面向支持 AMX 指令集的 Intel CPU，专门用于加速 CPU 侧 MoE 专家计算，并与 SGLang 配合实现 CPU-GPU 异构推理。

适用场景

建议在以下场景使用 AMX 后端：

  • CPU 支持 amx-bf16 / amx-int8 / amx-tile
  • 需要比通用 AVX 路径更高的 CPU 侧 MoE 吞吐
  • 使用 SGLang 做异构部署，希望 CPU 高效承载冷专家

支持的方法及选择建议：

  • BF16：全精度路径，适合优先保证精度的场景
  • AMXINT8：自定义量化权重格式，吞吐更高，精度可能轻微下降
  • AMXINT4：自定义量化权重格式，吞吐最高，但部分模型可能存在较明显精度下降

说明：AMXINT8 / AMXINT4 使用 KT-Kernel 的自定义量化权重格式（需先转换）。

AMX 架构要点

相较传统向量化内核，AMX 使用 Tile 矩阵计算模型，其理论吞吐可达 AVX-512 指令的 8 倍，因此特别适合计算密集型任务；KT-Kernel 在此基础上做了面向 MoE 的优化：

1）面向 Tiling 的内存布局

  • 加载阶段将专家权重预重排为 Tile 友好的子矩阵
  • 子矩阵按 64 字节边界对齐
  • 内存布局按计算访问顺序组织，提升缓存命中和预取效率

这样可以减少数据搬运开销，避免推理时的动态转置/重排。

2）Cache 友好的计算流水

  • 任务按列切分并分配到线程
  • 线程内 block 大小按缓存层级（L2/L3）调优
  • 使用 Tile 寄存器进行累加，减少中间结果回写

核心目标是减少 DRAM 往返，提升持续吞吐。

3）AMX / AVX512 动态路径选择

  • AMX 在大矩阵场景（常见于 prefill）优势明显
  • 算术强度较低的场景（常见于 decode）使用轻量 AVX512 核心可能延迟更优

KT-Kernel 会根据运行时负载特征在不同路径间切换，以平衡吞吐与延迟。

4）MoE 任务融合与动态调度

  • 融合专家 GEMM 任务，减少调度与同步开销
  • 将任务细粒度拆分并在线程间动态均衡
  • prefill 阶段通过任务窃取缓解专家激活分布倾斜

这对真实路由不均衡下的稳定性能非常关键。

典型优化效果

在单个 Xeon4 CPU 上，面向 MoE 场景专门设计的 AMX 后端可带来显著收益：

  • BF16 算子吞吐超过 20 TFLOPS
  • Int8/Int4 算子吞吐超过 37 TOPS

在端到端推理上，双路 Xeon4 CPU + 单张 RTX 4090 可实现超过 500 tokens/s 的 DeepSeek-V3 prefill 吞吐。

这些结果说明，面向 MoE 优化的 AMX Kernel 能在低成本的 CPU 上提供更高吞吐，显著降低大模型部署门槛。

系统要求

CPU 与平台

  • Intel Sapphire Rapids（第四代 Xeon）及更新平台
  • Linux x86-64
  • Python 3.10 / 3.11 / 3.12

检查 AMX 能力

lscpu | grep -i amx

典型输出包含：

amx-bf16 amx-int8 amx-tile

若没有 AMX 标志：

  • 检查 CPU 代际是否满足要求
  • 在 BIOS 中开启 AMX 相关选项
  • 确认操作系统/内核支持 AMX 状态管理

AMX 权重准备

不同 method 的 CPU 侧权重要求不同：

BF16：不需要转换

BF16 方法不需要运行 convert_cpu_weights.py。直接使用原始 BF16 模型权重目录作为 --kt-weight-path 即可（通常与 --model 指向同一目录）。

AMXINT8 / AMXINT4：需要转换

这两种方法使用 KT-Kernel 自定义量化格式，需要先把 CPU 侧专家权重转换为 AMX 友好格式。

python scripts/convert_cpu_weights.py \
  --input-path /path/to/model \
  --input-type bf16 \
  --output /path/to/cpu-weights-int8 \
  --quant-method int8

常用参数：

  • --input-type：fp8、fp16 或 bf16
  • --quant-method：int8 或 int4

转换完成后，将输出目录作为 SGLang 的 --kt-weight-path。

通过 SGLang 启动

当前 AMX 后端应通过 SGLang 启动流程使用。

python -m sglang.launch_server \
  --model /path/to/model \
  --trust-remote-code \
  --kt-method BF16 \
  --kt-weight-path /path/to/model \
  --kt-cpuinfer 64 \
  --kt-threadpool-count 2 \
  --kt-num-gpu-experts 32

关键参数：

  • --kt-method：BF16、AMXINT8 或 AMXINT4
  • --kt-weight-path：
    • BF16：填写原始 BF16 模型权重目录
    • AMXINT8/AMXINT4：填写转换后的 CPU 权重目录

完整安装和端到端示例请参考 KT-Kernel 安装指南。

常见问题

出现 Illegal instruction 或未走 AMX 路径

  • 使用 lscpu 确认 AMX 标志
  • 确认 BIOS 已启用 AMX
  • 环境变更后重新构建/安装 KT-Kernel

AMX 权重转换耗时较长

大模型转换耗时较长是正常现象。该步骤通常对同一组权重只需执行一次。

decode 场景吞吐无提升

decode 阶段常见算术强度较低，性能特征与 prefill 不同，属于正常现象；系统会在该阶段更多利用非 AMX 路径以优化延迟。