Accelerating Hybrid Inference in SGLang with KTransformers CPU Kernels

发表时间: 2025-10 · Blog post by KVCache.AI and Approaching AI (lmsys.org)

原文: https://lmsys.org/blog/2025-10-22-KTransformers

作者/机构：KVCache.AI and Approaching AI

速读

一句话结论 本文提出了 KTransformers 并将其集成到 SGLang 中，通过 AMX 专用 CPU 算子、设备协同优化与专家延迟执行机制，解决了千亿参数 MoE 模型在 CPU/GPU 混合推理时的协同开销与算力闲置问题，在单机混合部署下实现了最高 20 倍的 Prefill 阶段提速。

要解决什么问题 现代混合专家（MoE）语言模型（如 DeepSeek-V3）包含数千亿参数，但每个 token 仅激活极少部分专家。这种稀疏激活天然适合 CPU/GPU 混合推理：将需要大显存的路由专家放在 CPU 上，计算密集的注意力机制和共享专家放在高带宽 GPU 上，从而在单机部署万亿参数模型。然而，原有方案卡在设备协同开销和算力闲置上。具体而言，CPU 和 GPU 间存在严重同步等待，跨 NUMA 节点内存访问产生昂贵流量开销，且 CPU 任务与数据传输会打断 GPU 计算图捕获，导致混合架构有效吞吐量大幅受限。

怎么做的 核心思路是通过优化 CPU 算子、消除设备间同步断点，并重排计算顺序，让 CPU 和 GPU 极致重叠并行。方法由四个关键设计构成。第一，AMX 专用 CPU 算子。设计感知分块的内存布局使权重与 CPU 缓存对齐，并支持动态指令集切换：在计算密集的 Prefill（预填充）阶段用 Intel AMX 算子，在轻量级 Decode（解码）阶段切至 AVX-512。第二，NUMA 感知张量并行。将专家权重切片精准放置在各 NUMA 节点本地内存，避免跨节点传输瓶颈。第三，基于 CUDA Graph 的异步调度。将混合执行捕获为连续计算图，通过异步调度确保 CPU 任务不制造 GPU 计算图“断点”，将 GPU 算子下发开销从 20% 以上降至近零。第四，专家延迟执行（Expert Deferral）。利用 Transformers 残差连接对微小延迟的容忍性，跨层重排专家执行顺序。其更新式可抽象为：

$$H_{l+1} = H_l + \text{Attention}(H_l) + \text{Expert}(H_{l-k})$$


其中 $k$ 为延迟层数。系统在计算当前层 GPU 注意力时，异步执行被推迟到 CPU 的专家计算，实现两者完全重叠。此外，SGLang 集成引入混合专家并行，将“热点”专家留在 GPU 显存，缓解 CPU 带宽压力。

效果如何 实验在单卡（双路 Intel Xeon Platinum 8452Y 搭配单张 A100 或 RTX 4080）和多卡（双路 Xeon Gold 6454S 搭配 8 张 L20）环境下进行。对比基线为 PyTorch 原生实现（代表无底层协同优化的常规路线）。单卡测试中，得益于 AMX 优化，KTransformers 在 Prefill 阶段比基线最高提速 20 倍（单路 CPU 吞吐达 21.3 TFLOPS，为原生 3.9 倍）；Decode 阶段因协同开销降低最高提速 4 倍。多卡并发场景下（测试 int4 量化 DeepSeek-V3），8 并发时 8 卡配置比单卡吞吐量提升 264%。在 ShareGPT 数据集测试 DeepSeek-R1-0528（FP8）时，1000 个请求总吞吐达 227.85 tokens/s，中位数 token 间延迟 299.18 毫秒。代价与局限在于：专家延迟执行带来微小精度波动（准确率变化低于 0.5%）；单并发下多卡扩展收益极低（8 卡仅比单卡提升 26%）；且强依赖支持 AMX 的 Intel 处理器及 500 GB 以上 DDR5 内存。

主要贡献

现代混合专家（Mixture-of-Experts, MoE）语言模型（如DeepSeek-V3）通常包含数千亿个参数，导致其很难在单台显存受限的机器上进行部署。然而，MoE模型具有稀疏激活的特性，即在处理每个Token时仅有一小部分专家被激活。这一特性使得MoE模型非常适合采用CPU/GPU混合推理架构进行部署。

为了解决混合推理中存在的设备协调开销大、计算资源利用率低等挑战，MadSys @ Tsinghua与Approaching.AI联合开发了KTransformers项目（发表于SOSP’25）。KTransformers通过引入AMX专用的CPU内核、高效的设备协同机制以及专家延迟执行（Expert Deferral）技术，大幅提升了CPU/GPU协同推理MoE模型的效率。

目前，KTransformers已作为后端库被集成入SGLang中。这一集成将GPU的张量并行（Tensor Parallelism）与针对MoE的CPU/GPU混合专家并行（Hybrid Expert Parallelism）相结合，在多GPU高并发场景下实现了卓越的吞吐量、可扩展性与成本效益，使得在消费级或有限显存的硬件上高效运行万亿参数模型成为可能。

背景知识与设计原则

稀疏混合专家（MoE）模型的推理挑战与机遇：现代MoE语言模型（如DeepSeek-V3）包含数千亿个参数，但在处理每个Token时仅有一小部分专家被激活。这种稀疏激活模式使得MoE模型极其适合CPU/GPU混合推理架构。

混合推理的职责划分与优势：在这种混合设计中，被稀疏激活的专家模型可以高效地运行在具有大内存容量的CPU上；而密集且计算密集型的组件（如注意力机制和共享专家）则在具有更高带宽和吞吐量的GPU上执行。这种架构允许在GPU显存有限的单台机器上部署万亿参数模型，从而为研究和私有应用提供本地推理能力。
混合推理架构示意图

现有混合推理的瓶颈：尽管混合架构具有显著优势，但完全发挥CPU和GPU的潜力仍然具有挑战性。主要的限制因素在于设备间的协调开销以及计算资源的未充分利用，这些问题严重限制了系统的有效吞吐量。

核心方法细节

KTransformers引入了一系列优化技术，使得CPU/GPU在MoE推理中的协作变得更加高效。其改进主要分为以下三个核心类别：

AMX专用CPU内核设计与动态切换：
KTransformers利用Intel AMX优化的内核重新设计了CPU计算过程，并采用了一种感知分块（tiling-aware）的内存布局，使得权重存储与CPU缓存层次结构严格对齐。此外，系统支持在不同的指令集之间进行动态切换：在处理高强度的Prefill（预填充）阶段工作负载时使用AMX指令集，而在处理轻量级的Decode（解码）阶段时则切换至AVX-512指令集。在单路Xeon插槽上，经过AMX优化的内核能够达到高达21.3 TFLOPS的持续吞吐量，这一速度比PyTorch原生实现快3.9倍。这种底层计算能力的提升，直接转化为在Prefill阶段CPU端专家吞吐量的大幅提升，以及混合运行中整体Token吞吐量的显著增加。

高效的设备协同机制：
为了降低CPU与GPU之间的协调成本，KTransformers引入了感知NUMA的张量并行（NUMA-aware tensor parallelism）和基于CUDA Graph的调度机制。

  • 感知NUMA的张量并行：该机制将专家权重的切片放置在每个NUMA节点的本地内存中，从而确保计算主要在本地进行，避免了昂贵的跨NUMA内存流量。在双路服务器上，这一优化带来了高达63%的解码吞吐量提升。
  • 基于CUDA Graph的异步调度：CUDA Graph集成将混合的CPU/GPU执行过程捕获为连续的计算图。为了确保图捕获的鲁棒性，KTransformers采用了异步任务调度机制，使得CPU任务和数据传输不会在捕获的图中产生“断点”。通过这种方式捕获工作负载，成功将GPU内核启动的开销从超过20%降低到了几乎为零。这些优化共同确保了两个设备都能以最小的同步延迟协同运行。

专家延迟执行（Expert Deferral）与模型执行重叠：
KTransformers进一步引入了专家延迟机制，对跨层的专家执行顺序进行重排。部分专家被推迟到后续阶段执行，从而允许CPU端的专家计算与GPU端的注意力处理过程相互重叠。
专家延迟执行机制示意图
由于现代Transformer架构使用了残差连接（residual connections），它们天生对中间计算的微小延迟具有容忍度。因此，推迟特定的专家计算能够在仅轻微改变模型行为的代价下，极大增强调度的灵活性。该机制提高了两种设备的并发利用率，并带来了高达1.45倍的解码吞吐量提升，同时确保模型的精度变化保持在0.5%以下。

SGLang与KTransformers的系统集成设计：
SGLang目前已将KTransformers作为后端库进行集成，以支持高效的CPU/GPU混合推理。这一集成将GPU张量并行与CPU/GPU混合专家并行结合在一起。在此异构设备推理架构中，KTransformers提供高度优化的基于AMX的CPU内核，这些内核能够与GPU执行无缝协作。
虽然KTransformers本身专注于单GPU设置和高效的CPU协作，但SGLang在跨多GPU扩展方面表现出色，这在高并发场景中尤为有利。在混合设置中，多GPU可以处理更大的请求上下文并执行快速的注意力计算，而专家则在CPU和GPU之间进行智能调度——系统将常用（“热点”）专家存储在GPU上，以减轻CPU的计算和带宽压力。通过这种联合设计，使用各种硬件配置的用户都能充分利用可用资源，实现更好的吞吐量、可扩展性和成本效益。我们已经开发了概念验证（proof-of-concept）实现，并且将其完全集成到SGLang的路线图正在积极推进中。

实验环境

  • 硬件配置：

    • 单GPU + CPU 测试平台：双路 Intel Xeon Platinum 8452Y 服务器（36核 × 2，1 TB DDR5 × 2 内存），搭配 NVIDIA A100 (40 GB，用于全精度模型) 或 RTX 4080 (16 GB，用于量化模型)。
    • 多GPU + CPU 测试平台：双路 Intel Xeon Gold 6454S CPU，搭配 8× NVIDIA L20 GPU。
    • 推荐运行硬件要求：支持AMX的现代Intel Xeon处理器（如Sapphire Rapids及以上）；充足的DDR5内存以容纳所有专家权重（对于DeepSeek-V3规模模型通常需要500GB+）；一个或多个具备足够显存的GPU（用于注意力层、共享专家及部分路由专家）；双路NUMA配置以利用NUMA感知的线程池分配。

  • 软件配置：操作系统为Linux x86_64；CUDA版本 >= 12.1；编译器需 gcc, g++ >= 11；构建工具需 CMake >= 3.25；Python 3.11。依赖库包括 sglang[all] 和 KTransformers CPU 内核 (kt-kernel)。

  • 模型与数据集：测试模型包括 int4量化的 DeepSeek-V3 和 FP8量化的 DeepSeek-R1-0528 MoE模型。多GPU并发测试采用平均输入128 tokens、输出512 tokens的工作负载。ShareGPT基准测试使用了包含1000个对话请求的数据集（总计301K输入tokens，188K输出tokens）。

实验结果

单GPU + CPU 性能评估：
原生KTransformers在单GPU+CPU设置下进行了详细的性能评估，SGLang集成KTransformers后在相同配置下达到了与原生实现相当的性能。

  • Prefill（预填充）阶段：在所有Prompt长度下，KTransformers均持续优于两个基线方案。这主要得益于AMX优化的CPU内核，实现了高达20倍的加速。
    Prefill阶段性能对比

  • Decode（解码）阶段：KTransformers同样优于基线方案，其性能增益主要归功于CPU/GPU协调开销的显著减少，实现了高达4倍的加速。
    Decode阶段性能对比

多GPU + CPU 混合推理性能评估：
在8× L20 GPU和双路Intel Xeon Gold 6454S CPU系统上，对int4量化的DeepSeek-V3进行了并发吞吐量测试。

  • 并发吞吐量分析：在单并发条件下，8-GPU配置相较于1-GPU设置仅提供了有限的提升（仅增加26%）。然而，在8路并发条件下，相同的8-GPU配置获得了相较于1 GPU高达264%的吞吐量增益，展现了极佳的可用性（平均每个请求达到近20 tokens/s）。这一改进主要来自于将更多专家放置在GPU上，从而在带宽受限的情况下有效减少了CPU内存的访问压力。
    多GPU并发吞吐量对比

ShareGPT 基准测试 (NVIDIA L20 × 8 设置)：
使用 DeepSeek-R1-0528 (FP8量化MoE模型) 和 ShareGPT 数据集（1000个请求）进行了端到端测试。

  • 测试结果：总Token吞吐量达到 227.85 tok/s，输出Token吞吐量为 87.58 tok/s，请求吞吐量为 0.46 req/s。中位Token间延迟（Inter-Token Latency, ITL）为 299.18 ms（平均ITL为 431.61 ms，P99 ITL为 1935.13 ms）。
  • 结论分析：该设置证明了SGLang + KTransformers集成能够有效利用消费级GPU进行混合推理，在万亿参数MoE模型上实现了超过 220 tokens/s 的总吞吐量。相对较低的中位Token间延迟（299ms）确保了交互式应用中流畅的流式生成体验。

结论

本文介绍了通过集成KTransformers CPU内核来加速SGLang中的混合推理方案。KTransformers通过AMX专用内核、NUMA感知并行、CUDA Graph异步调度以及创新的专家延迟执行机制，彻底释放了CPU/GPU混合架构在MoE模型推理上的潜力。与SGLang的结合，进一步将这种高效的单节点混合推理能力扩展到了多GPU高并发场景下。实验证明，该方案能够在消费级或有限显存的硬件上，以极低的成本和延迟实现万亿参数模型（如DeepSeek-V3/R1）的高吞吐部署。未来，完全集成到SGLang的开发路线图将继续推进，为社区提供更强大、灵活的大模型服务基础设施。

补充细节

安装流程与前置条件：
要使用带有KTransformers混合推理的SGLang，必须同时安装SGLang和KTransformers CPU内核（kt-kernel）。

  • 步骤1：安装SGLang。使用命令 pip install "sglang[all]" 安装官方完整版SGLang。
  • 步骤2：安装KTransformers CPU内核。kt-kernel 提供了用于混合推理的AMX优化计算支持，具体编译和安装需参考官方仓库指南。

使用示例与模型准备：

  • 下载模型：针对KTransformers混合推理优化的DeepSeek-R1模型（包含GPU和CPU权重）可以从Approaching AI的ModelScope主页下载。
  • 启动服务器：可以通过以下命令启动启用了KTransformers混合推理的SGLang服务器：
python -m sglang.launch_server \
  --host 0.0.0.0 \
  --port 30000 \
  --model /path/to/gpu-weight \
  --kt-amx-weight-path /path/to/cpu-weight \
  --kt-cpuinfer 80 \
  --kt-threadpool-count 2 \
  --kt-num-gpu-experts 200 \
  --kt-amx-method AMXINT4 \
  --attention-backend triton \
  --trust-remote-code \
  --mem-fraction-static 0.98 \
  --chunked-prefill-size 4096 \
  --max-running-requests 37 \
  --max-total-tokens 37000 \
  --served-model-name DeepSeek-R1-0528-FP8 \
  --enable-mixed-chunk \
  --tensor-parallel-size 8 \
  --enable-p2p-check \
  --disable-shared-experts-fusion

核心启动参数解析：

  • --kt-amx-weight-path：指向针对CPU优化的模型权重路径。这些权重已经过预量化处理并格式化，以支持高效的AMX计算。
  • --kt-cpuinfer：专用于专家推理的CPU核心数量（例如，在双路服务器上可设置为80个核心）。
  • --kt-threadpool-count：用于并行CPU执行的线程池数量。对于双路NUMA配置，通常设置为2。
  • --kt-num-gpu-experts：保留在GPU上的“热点”专家数量。将更多的专家放置在GPU上可以减少CPU的计算压力，但这需要消耗额外的GPU显存。用户需根据显存容量和工作负载模式进行动态调整。
  • --kt-amx-method：CPU内核优化方法。对于int4量化的模型，应使用 AMXINT4 以充分利用Intel AMX指令集实现最大吞吐量。
    启动服务器后，用户可以通过兼容OpenAI的API端点 http://0.0.0.0:30000 发送推理请求。

参考文献与资源汇总