kestrel-llm 17 · 草稿模型上 GPU：双 backend 投机解码

用户36005557900
56 阅读8分钟

17 · 草稿模型上 GPU：双 backend 投机解码

一句话：把 Qwen3-0.6B 草稿模型从 CPU 搬到 GPU 是投机解码落地链条上的最后一块拼图——但草稿上 GPU 是必要条件，不是充分条件，端到端吞吐只涨了 +11%。 前置：建议先读 第 16 篇（batch verify 的数值一致性修复）。 环境：x86-64 + NVIDIA RTX 5090（CUDA 后端）· target=Qwen3-30B-A3B q4local、draft=Qwen3-0.6B · 开关 VLLM_CUDA_DECODE=1 VLLM_CUDA_INFER=1 VLLM_CUDA_SPEC=4 VLLM_CUDA_SPEC_DRAFT=1。

一、问题与结论

结论先行，下表数据均为实测：

指标CPU 草稿（改前）GPU 草稿（改后）
草稿 backend无（ist2.no_cuda=1）独立 draft_cuda（2 GiB wcache）
草稿权重预热—253 个
greedy TOKIDS（32 tok）32/3232/32 逐位一致
接受率（32 tok）30.4%34.6%
吞吐（32 tok）77.18 tok/s86.00 tok/s（+11%）
吞吐（16 tok）98.75 tok/s101.68 tok/s（+3%）

草稿上 GPU 后吞吐仍低于无 spec 基线（118.79 tok/s，实测）。瓶颈不在草稿速度，在 0.6B 对 30B 的接受率（30–50%）+ 每轮 verify/correction 的固定开销（见第四节）。

二、背景：为什么必须双 backend

主模型的 full-GPU decode 靠 device 权重缓存（wcache）hit-only 跑。wcache 的 key 是 st_cuda_wkey(layer, proj) = (layer<<32)|(proj+1)——只看 (layer, proj)，不看模型。

草稿（0.6B）和主模型（30B）的 layer 0 QKV 投影都用同一个 key 0x0000000000000001，共用一个 vllm_cuda_t 会互相覆盖权重。所以必须给草稿一个独立的 vllm_cuda_t 实例，让两套 wcache（还有两套 device KV mirror、两套 decode 几何）各归各的。

mermaid diagram

图 1：为什么必须是双 backend。device wcache 的 key st_cuda_wkey(layer, proj) = (layer<<32)|(proj+1) 只看 (layer, proj)、不看模型，于是 target（30B）与 draft（0.6B）的 layer 0 QKV 撞同一个 key，共用一个 vllm_cuda_t 就会互相覆盖权重；解法是给草稿绑独立的 draft_cuda（2 GiB 预算），两套 state 经 st_cuda_of(st) 各解析到自己的 backend，各持有自己的 wcache / device KV mirror / decode 几何。（host strip-cache 的 key 多带 (layer, proj, N, K)，草稿 dim=1024 vs 主模型 dim=2048 不碰撞，故可共用。）

一个反直觉的顺带结论：host strip-cache 可以共用。st_npu_gw_get(layer, proj, N, K, q4) 用 (layer, proj, N, K) 匹配，而草稿 dim=1024 vs 主模型 dim=2048，N/K 不同，天然不碰撞。所以 host 侧一行不用改，只要 device 侧拆开。

三、核心机制：per-state backend 参数化与接线

3.1 给 state 挂一个 backend，统一解析

现状是 st_qwen_model_forward 内部硬编码全局 g_cuda。最小改动是给 STQwenInferenceState 挂一个 backend，让所有 CUDA 路径按 st->cuda 解析。仓库里的字段定义（include/model/vllm_safetensors.h）：

/* 文件名：include/model/vllm_safetensors.h，STQwenInferenceState */
int     no_cuda;           /* 1 = force CPU decode (spec-draft model) */
struct vllm_cuda_s *cuda;  /* per-state CUDA backend（NULL = 用全局 g_cuda；
                            * spec 草稿模型绑定独立实例 draft_cuda） */
int     cuda_kv_ready;     /* per-state device KV mirror 已初始化 */
int     cuda_kv_slot;      /* device KV mirror slot（-1 = 未分配；-2 = 池满已放弃） */
int     cuda_decode_ready; /* per-state full-GPU decode 已初始化 */

统一解析函数与「逐层 GEMM offload 抑制」（src/model/vllm_safetensors.c）：

/* 文件名：src/model/vllm_safetensors.c */
/* Per-state backend resolution: the spec-draft model binds its own independent
 * vllm_cuda_t instance (draft_cuda); everything else uses the global g_cuda. */
static vllm_cuda_t *st_cuda_of(const STQwenInferenceState *st) {
    return st->cuda ? st->cuda : g_cuda;
}

/* 文件名：src/model/vllm_safetensors.c，st_qwen_model_forward() */
const int saved_cuda_suppress = g_cuda_suppress;
/* 草稿（no_cuda 或绑了非全局 backend）forward 期间，关掉走全局 g_cuda 的
 * 逐层 GEMM offload（草稿权重不在主模型 g_cuda 里）。 */
if (st->no_cuda || (st->cuda && st->cuda != g_cuda)) g_cuda_suppress = 1;

改动范围刻意收窄到「full-GPU decode + KV」两条核心路径（st_cuda_attn_on/decode_on/kv_append/try_attn_q8/decode_init、full-GPU decode 分支、st_cuda_spec_decode），全部换成 st_cuda_of(st)；全局 g_cuda_kv_ready/g_cuda_decode_ready 下沉为 per-state。

逐层 GEMM offload（dyn_matvec_*）仍走全局 g_cuda（它们没有 st 参数），但草稿走 full-GPU decode 成功就提前 return，根本不进逐层循环；失败才回退纯 CPU。

3.2 接线：创建独立 draft backend 并预热

在 main.c 里，主 backend 就绪后按 VLLM_CUDA_SPEC_DRAFT 建一个独立的 g_draft_cuda（src/main.c）：

/* 文件名：src/main.c */
/* 草稿与主模型 device wcache key 撞（(layer<<32)|proj），必须独立 vllm_cuda_t。
 * 草稿权重小（~0.6GB int8），给 2GB 预算。 */
vllm_cuda_cfg_t dcfg;
memset(&dcfg, 0, sizeof(dcfg));
dcfg.device    = cfg.device;
dcfg.verbosity = 1;
dcfg.wcache_mb = 2048;
if (vllm_cuda_init(&g_draft_cuda, &dcfg) == 0 && vllm_cuda_available(g_draft_cuda))
    printf("  CUDA draft backend: ready (wcache 2 GiB)\n");

草稿初始化时绑定 draft_cuda、显式 init KV mirror、再临时切全局句柄预热权重（src/main.c）：

/* 文件名：src/main.c */
ist2.no_cuda = 0;
ist2.cuda    = g_draft_cuda;
const int kv2  = ist2.cfg.n_kv_heads * ist2.cfg.head_dim;
const int dslot = vllm_cuda_kv_slot_alloc(g_draft_cuda);
if (dslot >= 0 &&
    vllm_cuda_kv_init(g_draft_cuda, dslot, ist2.cfg.n_layers,
                      ist2.cfg.max_seq_len, kv2, ist2.cfg.n_kv_heads) == 0) {
    ist2.cuda_kv_slot  = dslot;
    ist2.cuda_kv_ready = 1;
}
/* 预热草稿权重：临时切全局 backend，预热后切回主模型 */
vllm_cuda_t *saved = g_cuda;
st_cuda_set(g_draft_cuda);
int warmed = st_cuda_preload_all(&w2, ist2.cfg.n_layers, d2, q2,
                                 kv2, ist2.cfg.ffn_dim, ist2.cfg.vocab_size);
st_cuda_set(saved);

两个细节：

  1. KV mirror 显式 init。主模型的 KV mirror 是 prefill 时 st_cuda_kv_append 懒初始化的；草稿没有这条 CPU prefill 路径（草稿 prefill 直接走 full-GPU decode），而 vllm_cuda_decode_run 前置检查 kv_nl != 0，所以得先显式 vllm_cuda_kv_init。
  2. 预热靠切全局 backend，而不是给 st_cuda_preload_all 加参数——它内部经 st_cuda_preload_one → vllm_cuda_wcache_preload(g_cuda, ...) 用全局句柄，切一下 g_cuda 再切回来最省事（投机循环单线程，无并发）。

四、实测数据：为什么收益只有 +11%

把每轮投机解码的时间拆开（30B target 满配置，实测）：

段耗时说明
草稿自回归 4 token（GPU）~17 ms0.6B dense 的 full-GPU decode，非 ~1 ms 理想值
主模型 batch verify（M=5）~3 ms已摊销
correction forward（M=1）~8 ms全量 decode，未摊销

mermaid diagram

图 2：每轮投机解码的时间拆解（30B target 满配置）。三段依次是草稿自回归 4 token（~17ms）→ 主模型 batch verify M=5（~3ms，已摊销）→ correction forward M=1（~8ms，全量 decode 未摊销），合计 ~28ms 产出约 3 token ≈ 9 ms/token。草稿从 CPU 换 GPU 只省第一段，而第一段占比远不如「correction 每次 8ms」和「接受率只有一半」来得重——这就是端到端只涨 +11%（77.18 → 86.00 tok/s）的原因。

每轮 ~28 ms 产出 ~3 token（接受率 50%），≈ 9 ms/token。草稿从 CPU 换 GPU，省的只是第一段——而第一段在总时间里的占比，远不如「correction 每次 8 ms」和「接受率只有一半」来得重。

两个真正的天花板：

  1. 接受率：0.6B 对 30B 的 greedy 一致率只有 30–50%，仓库 wiki 早先预估的「同家族 60–80%」对 50× 参数差的 draft 明显偏乐观。
  2. correction 没摊销：每轮 correction 走一次 M=1 全量 decode（读全量权重），和普通 decode 一样贵，把投机省下的部分又吐回去。

五、边界与已知限制

  • 吞吐仍低于无 spec 基线：86.00 vs 118.79 tok/s（实测）。草稿上 GPU 没有让投机转正。
  • 收益口径绑定本机：target=Qwen3-30B-A3B + draft=Qwen3-0.6B，VLLM_CUDA_SPEC=4，RTX 5090。换模型规模/换卡需重测。
  • draft backend 预算是写死的 2 GiB：够放 0.6B int8（~0.6 GB）+ 余量；更大 draft 需调整。
  • 未做：草稿侧的 prefill 批量化、多草稿并行；本篇只覆盖单草稿、单轮 k=4。

CPU 对照（迁移前基线）

  • CPU 参考：kestrel-llm/src/model/vllm_safetensors.c:st_qwen_model_forward —— 单 backend 前向（草稿与主模型共用全局 g_cuda，无 per-state backend 概念）。
  • 迁移要点：把 backend 参数化——给 STQwenInferenceState 挂 st->cuda，草稿绑独立 vllm_cuda_t（因 device wcache key 只看 (layer,proj)）；全局 g_cuda_*_ready 下沉 per-state；host strip-cache 因 key 多带 N/K 可共用。
  • 真机验证：未独立复现：需 30B+0.6B 双 backend 与 VLLM_CUDA_SPEC=4 的接受率/吞吐。

六、小结

  1. wcache key 只看 (layer,proj)，换模型必须拆 backend（或给 key 加命名空间）。这是「双 vllm_cuda_t」的根因，也是 第 10 篇 那条「症状≠病灶」的反面——这里病因和症状都在 key 上。
  2. host strip-cache 的 key 比 device wcache 更细（多带 N/K），所以「撞不撞」要分别看，不能因为 device 撞了就假设 host 也撞。
  3. per-state ready 标志是双 backend 的代价。把全局 g_cuda_*_ready 下沉到 state，否则切 backend 会互相污染。
  4. 草稿上 GPU 不是收益的充分条件。投机解码的收益 = f(接受率, verify/correction 摊销)，草稿速度只是其中一个因子，且不是当前的主导因子。

相关篇目：第 10 篇、第 16 篇、第 18 篇 源码与配套资源：本仓库 gitee.com/pei-xiaogua…； CPU 推理源码 gitee.com/pei-xiaogua…

avatar
文章
阅读
粉丝