KTransformers

Layerwise Prefill

Layerwise prefill 是 KT + SGLang 在长上下文 prefill 阶段的加速策略：

  • 短序列：使用 CPU+GPU 混合 prefill
  • 长序列：切换到逐层 GPU prefill（layerwise prefill）

它的目标不是“让所有场景都更快”，而是把长序列 prefill 的主瓶颈从 CPU 计算转移到更可扩展的 GPU 计算路径。

适用场景

建议在以下场景启用 layerwise prefill：

  • 业务以长输入为主（如长上下文、长文档处理、批量预填充），且 prefill 吞吐是核心目标
  • 使用原生精度链路（CPU/GPU 权重语义一致）
  • GPU 仍有可用显存，能承受额外 prefill 工作区

支持的方法：

  • RAWINT4：典型模型为 Kimi-K2-Thinking
  • FP8：典型模型为 MiniMax-M2/M2.1、Qwen3-235B-A22B
  • FP8_PERCHANNEL：典型模型为 GLM-4.7
  • BF16：典型模型为 GLM-4.7、Qwen3-235B-A22B

原理与两种瓶颈

Layerwise prefill 本质上是“按 token 数切换 prefill 路径”：

  • 当 prefill token 数 低于或等于阈值（--kt-gpu-prefill-token-threshold）：走混合路径
  • 当 prefill token 数 高于阈值：走 layerwise prefill 路径

两条路径的主要瓶颈不同：

  • 非 layerwise（混合路径）：瓶颈通常在 CPU 专家计算（即使有并行掩盖，长序列下仍可能成为主瓶颈）
  • layerwise：瓶颈通常转向权重搬运，尤其是 CPU→GPU 的 PCIe 传输与对应加载开销

这也是为什么长序列更适合 layerwise：

  • token 越长，GPU 侧矩阵计算利用率越高
  • CPU 计算占比下降，长序列扩展性更稳定
  • 代价是更高显存占用，以及对 PCIe 带宽更敏感

Layerwise 执行路径详解

进入 layerwise 后，可以把执行过程理解为“按 chunk、按层推进”的流水：

  1. 按阈值切换模式
  • 当 prefill token 数超过 --kt-gpu-prefill-token-threshold，切入 layerwise prefill。
  1. 按 chunk 切分 prefill
  • 请求会按 --chunked-prefill-size 分块。
  • chunk 数量近似为：chunk 数 ≈ ceil(prefill token 总数 / chunked-prefill-size)。
  1. 每个 chunk、每一层都要准备完整层级权重工作集
  • 当前层需要把 CPU 侧权重工作集搬到 GPU（主要经 PCIe）。
  • 这是 layerwise 路径与混合路径的关键差异，也是主要额外开销来源。
  1. 层内采用“逐专家三级流水”
  • 阶段 A：权重格式转换（写入 pinned CPU memory slot）
  • 阶段 B：PCIe 传输（从 pinned CPU memory slot 搬到 GPU）
  • 阶段 C：Marlin repack 后处理
  • 三个阶段按专家粒度流水推进，降低串行等待。
  1. DDIO + 双缓冲优化点
  • 传统方式下，格式转换写内存 + PCIe 读内存会分别占用内存写带宽和读带宽。
  • 采用逐专家紧邻流水后，单专家数据常驻于 LLC/L3；配合 DDIO，写入更倾向落在 LLC，随后 PCIe 传输也直接从 LLC 读取。
  • 结果是：显著降低对 DRAM 读写带宽的占用，把瓶颈更集中到 PCIe 与 GPU 计算本身。
  1. 当前层在 GPU 上完成该 chunk 的 MoE 计算
  • 计算完成后继续下一层，直到该 chunk 全部层处理完成。
  • 然后进入下一个 chunk，重复上述过程。

典型优化效果

以 MiniMax-M2.1 (FP8) 为例（PCIe Gen5 平台）：

  • 1 卡（RTX 5090）：prefill 吞吐至多可达 1172 tokens/s
  • 2 卡（RTX 5090）：prefill 吞吐至多可达 2879 tokens/s
  • 4 卡（RTX 5090）：prefill 吞吐至多可达 4045 tokens/s

长序列下 layerwise prefill 的上限与 PCIe 总带宽强相关。

  • PCIe 代际越高（如 Gen5 相比 Gen4）通常上限越高
  • GPU 卡数越多，可用总 PCIe 带宽通常越大，prefill 吞吐上限也更高

因此评估 layerwise 性能时，建议把“模型与参数”与“PCIe 代际 + 卡数（总带宽）”一起看。

额外显存开销

启用 layerwise prefill 后，会增加额外显存，主要来自：

  1. 单个 MoE 层的完整专家权重工作集（临时加载到 GPU）
  2. prefill 临时缓冲区（与 --chunked-prefill-size 近似线性相关）
  3. 额外计算工作区/中间缓冲

可用近似关系理解：

  • 额外显存 ≈ 单个完整 MoE 层权重工作集 + 与 chunked-prefill-size 线性相关的临时缓冲

在不同模型上，额外显存差异很大。经验上可达到数 GB 级别（例如某些模型约 3.6GB 到 9GB+）。

关键参数与调优

1）--kt-gpu-prefill-token-threshold

该参数本质是在找“路径切换平衡点”：

  • 理想阈值对应于两条路径耗时接近相等的长度：当T_hybrid(L) ≈ T_layerwise(L)时，对应的 L 就是较优阈值附近区间

其中平衡点主要由两项共同决定：

  • CPU kernel 性能（影响 hybrid 路径）
  • PCIe 总带宽（代际 × 链路 × 卡数，影响 layerwise 路径）

因此阈值没有固定最优值，必须按机器实测调优。经验上通常落在几千 token量级。

2）--chunked-prefill-size

这是 layerwise 中非常关键的参数。

  • prefill 会按 chunk 处理
  • 每个 chunk 都需要进行一次完整层级权重搬运流程（经 PCIe 从 CPU 侧进入 GPU 工作集）

因此会出现典型权衡：

  • chunk 太小：chunk 次数增多，完整层级权重搬运次数更多，PCIe 开销占比上升，吞吐受限
  • chunk 太大：单次临时显存需求上升，更容易 OOM

实操建议：

  • 在不 OOM 的前提下，尽量把 --chunked-prefill-size 调到较大
  • 与 --max-total-tokens、--kt-num-gpu-experts 联合调参，平衡 prefill 与 decode

启动示例

python -m sglang.launch_server \
  --model /path/to/model \
  --trust-remote-code \
  --kt-method FP8 \
  --kt-weight-path /path/to/model \
  --kt-cpuinfer 64 \
  --kt-threadpool-count 2 \
  --kt-num-gpu-experts 32 \
  --chunked-prefill-size 16384 \
  --kt-gpu-prefill-token-threshold 2048

常见问题

1）为什么短输入看不出收益？

短输入通常还在混合路径，或尚未进入可体现 layerwise 优势的区间。

2）为什么吞吐提升但更容易 OOM？

因为 layerwise 需要额外显存承载完整层级权重与更大的 prefill 临时工作区。

可调整以下参数：

  • 下调 --chunked-prefill-size，降低单次 prefill 临时显存占用
  • 下调 --max-total-tokens，降低 KV Cache 显存占用
  • 下调 --kt-num-gpu-experts，降低专家权重显存占用