Kestrel_VLLM - 15 · decode 优化四连击：19ms → 8ms

15 · decode 优化四连击：19ms → 8ms

一句话 ：从 19 ms/tok 一路打到 8.4 ms/tok（2.25× ），靠的不是一个大招，而是用 stub 探针定位真瓶颈 + 四个位级安全/近似安全的小优化 ；本篇讲方法、讲数据，也讲两条走错的岔路。 前置 ：建议先读 第 13 篇 · Tiling 与 Tensor Core 判否、第 14 篇 · router 判否与 decode 真瓶颈。 环境：x86-64 + NVIDIA RTX 5090（sm_120）· 阶段三 30B-A3B MoE（Qwen3-30B-A3B q4：128 专家 / top-8 / dim=2048 / moe_ffn=768 / 48 层）。

一、问题与结论

口径：64 token decode 的 tok/s（精确值） 为权威；加速比一律用 tok/s 算，不用取整 ms。 graph off（拆解各优化单步收益）与 graph on（最终交付态）两条主线分开。

阶段	口径	tok/s	ms/tok	单步加速
基线（无优化）	graph off	52.7	19.0	1.00×
+ attention 转置归约	graph off	60.6	16.5	1.15×
+ router 并行化	graph off	79.2	12.6	1.31×
+ M=1 GEMV 去空闲	graph off	95.5	10.5	1.21×
+ CUDA Graph（叠加）	graph on	118.8	8.4	1.24×

• 总加速 = 118.8 / 52.7 ≈ 2.25×（此前的 2.4× 是拿取整 ms 估算，已订正）。
• CUDA Graph 单独收益（无其他优化时）1.08×；叠加后 1.24×------launch 占比变了，收益非恒定。

greedy 全程 16/16 逐位一致（含关键点 t=16、t=27、t=57、t=64）。

二、方法论：先用 stub 探针把瓶颈拆到底

直觉一直说「MoE FFN 是 decode 瓶颈」。但加了细粒度 stub（VLLM_CUDA_STUB_{ATTN,ROUTER,GU,LMHEAD}） 逐段跳过，19ms 被拆成四块，推翻了直觉：

组件	原成本	真相	标注
attention（QKV/O + attention 核）	8 ms（42%）	一直被当小头忽略，实为最大	实测
router	7 ms（37%）	f32 延迟受限	实测
gate/up/down	2 ms	非瓶颈	实测
lm_head + norm + launch	2 ms	小	实测

教训：排序必须由探针定，不由直觉定 （与 第 14 篇 同源）。

图 1 ：stub 探针把 19ms 拆成四块的定位链。用 VLLM_CUDA_STUB_{ATTN,ROUTER,GU,LMHEAD} 逐段跳过后取差，结果推翻直觉------一直被当小头忽略的 attention 实为最大（8ms，42%），router 次之（7ms，37%），而直觉认为的 MoE FFN（gate/up/down）只有 2ms，lm_head + norm + launch 也是 2ms。

三、四连击

图 2 ：四连击的累计收益链（52.7 → 118.8 tok/s）。① attention 转置归约（1.15×）、② router 并行化（1.31×）、③ M=1 GEMV 去空闲（1.21×）三档在 graph off 下量；④ CUDA Graph（1.24×）在 graph on 下叠加，单独收益只有 1.08×------因为它压的是剩余的 launch 占比，非恒定。总加速用精确 tok/s 算得 2.25×，不是拿取整 ms 估的 2.4×。

① attention 转置归约（8→5ms，归约树变化，greedy 未翻）

k_attn_decode_f32q pass 3 原来对每个输出元素 i<hd=128 各做一次 butterfly 归约 （~1408 次 barrier/层）。 改成转置 V-sum：每线程一个 (i, t-stride)，128 次串行归约塌成一次 combine。 归约树变了（属「数学等价非位级等价」），但实测 greedy 未翻：

/* 文件：src/npu/cuda/vllm_cuda_kernels.cu（k_attn_decode_f32q 的 Pass 3，节选） */
/* Pass 3 (transposed V sum): one thread per (output element i, t-stride), so
 * the hd sequential block-reductions collapse into a single combine. The
 * reduction tree over t changes (reduction_tree_sensitivity) -> greedy/PPL
 * must be re-verified. */
{
    const int ntg = VCUDA_ATT_T / hd;      /* t-strides per output element */
    const int i   = threadIdx.x % hd;      /* output element */
    const int tg  = threadIdx.x / hd;      /* t-stride id */
    float acc = 0.0f;
    for (int t = tg; t < seq_len; t += ntg) {
        const int8_t *vt = vb + (size_t)t * kv_dim;
        const float vss = vs[(size_t)t * nkv + kh] * (1.0f / 127.0f);
        acc += (float)vt[i] * scores[t] * vss;
    }
    red[threadIdx.x] = acc;
    __syncthreads();
    if (tg == 0) {
        float s = 0.0f;
        for (int g = 0; g < ntg; g++) s += red[(size_t)g * hd + i];
        outp[i] = s * invS;
    }
}

② router 并行化（7→4ms，放开位级一致）

router 是 [128][2048] f32 矩阵向量乘，原 128 线程（4 warps）× 2048 串行 = 延迟受限。 拆成 k_moe_router_gemm（32 block × 256 线程，4 专家 × 64 d-chunk 树归约 = 256 warps，64× 并行度）

• k_moe_router_topk：

/* 文件：src/npu/cuda/vllm_cuda_kernels.cu（k_moe_router_gemm，节选） */
const int e  = blockIdx.x * 4 + (threadIdx.x >> 6);  /* 0..127 */
const int c  = threadIdx.x & 63;                     /* 0..63 */
const int chunk = d >> 6;                            /* 32 */
const float *rw = router + (size_t)e * d + (size_t)c * chunk;
const float *xc = x + (size_t)c * chunk;
float s = 0.0f;
for (int i = 0; i < chunk; i++) s += rw[i] * xc[i];
__shared__ float lg[4][64];
lg[threadIdx.x >> 6][c] = s;
__syncthreads();
for (int off = 32; off > 0; off >>= 1) {
    if (c < off) lg[threadIdx.x >> 6][c] += lg[threadIdx.x >> 6][c + off];
    __syncthreads();
}
if (c == 0) logits[e] = lg[threadIdx.x >> 6][0];

这一步在「位级一致」硬约束下是判否的（拆 i 维破坏保序，见 第 14 篇）， 放开后落地，且本 prompt 未翻转任何 token。

③ M=1 GEMV 去空闲（13→10ms，位级等价）

k_gw_gemm_f32act 在 M==1（decode）时 block 的 4 个 m 线程只有 1 个干活（M=1 只有一行），75% 线程空转 。 改成 M==1 时把 m 线程映射到 n 维。位级等价（每列点积仍单线程计算，仅换线程分配）， greedy 16/16 保持：

/* 文件：src/npu/cuda/vllm_cuda_kernels.cu（k_gw_gemm_f32act 的线程映射，节选） */
int n0, m;
if (M == 1) {
    /* GEMV (decode): map all VC_NT*VC_MT threads onto n so no m-thread is
     * idle (the general shape leaves 3/4 of the block idle when M==1). */
    n0 = (blockIdx.x * (VC_NT * VC_MT) + threadIdx.y * VC_NT + threadIdx.x) * VC_VEC;
    m  = 0;
} else {
    n0 = (blockIdx.x * VC_NT + threadIdx.x) * VC_VEC;
    m  =  blockIdx.y * VC_MT + threadIdx.y;
}

④ CUDA Graph 压缩 launch（叠加再省 2ms）

把 decode 整段 ~870 次内核 launch 捕获成图、每 token replay 一次。零数值风险，opt-in：

/* 文件：src/npu/cuda/vllm_cuda_kernels.cu（vc_decode_graph_capture，节选） */
if (cudaStreamBeginCapture(s->stream, cudaStreamCaptureModeThreadLocal) != cudaSuccess) {
    cudaGetLastError(); return -1;
}
if (vc_decode_body(s, s->stream) != 0) {
    cudaStreamEndCapture(s->stream, &s->graph);
    return -1;
}
if (cudaStreamEndCapture(s->stream, &s->graph) != cudaSuccess) {
    cudaGetLastError(); return -1;
}
if (cudaGraphInstantiate(&s->gexec, s->graph, 0) != cudaSuccess) {
    cudaGetLastError(); return -1;
}
s->graph_ready = 1;

四、两条走错的岔路（判否）

方向	判定	卡点
MoE FFN tiling	❌	25%→100% 占用，吞吐 1.02×（访存延迟受限，非占用）
Tensor Core	❌	float 激活保 greedy 冲突 + 非计算受限
router fp16 量化	❌	greedy 保持但吞吐不变（延迟受限非带宽，减字节不减延迟）
router 并行化（位级约束下）	❌	拆 i 维破坏保序，greedy t=27 发散

详见 第 13 篇、第 14 篇。

五、最终 breakdown（graph off，10ms）

组件	成本	剩余瓶颈	标注
attention（QKV/O + attention 核）	4 ms	int8 权重读延迟 + rope/kv 小核	实测
router	2 ms	已并行化，剩余归约/launch	实测
gate/up/down	2 ms	标量 q4 点积（stride 72B tile 访存延迟）	实测
lm_head + norm + launch	2 ms	小核 launch + norm	实测

六、边界与已知限制

• 口径混用要分开：单步收益在 graph off 下量，最终交付（2.25×）是 graph on；CUDA Graph 的增益随 launch 占比变化（单独 1.08×，叠加后 1.24×），不可当常数外推。
• 基线口径备注 ：52.7 tok/s 来自 16-token 轮（moe_final16=52.69 / moe_graphoff16=53.53），无 64-token 纯基线 log；如需严格对齐可补一次 --stream-n 64 graph-off 基线复测。
• greedy 一致性 ：全程 16/16（含 t=16/27/57/64），但 ① 的 attention 转置归约属「归约树变化」，② 的 router 并行化属「放开位级一致」------两者都是近似安全而非位级等价，换 prompt 需重新对拍。
• ③ 是位级等价（每列点积仍单线程计算），④ 是零数值风险（同一批内核 replay）；① ② 不是。
• 剩余四项都在 2--4ms 且收益递减，本轮收口于此，未继续压。

CPU 对照（迁移前基线）

• CPU 参考 ：kestrel-llm/src/model/vllm_safetensors.c（函数 flash_attn_single_q_q8_neon）------ CPU 侧的归约序参考。
• 迁移要点：CPU 侧的归约序作为位级/近似判据的参照 → CUDA 侧先 stub 探针定位真瓶颈（attention 42%），再做四个针对性小优化：转置归约 / router 并行化 / M=1 线程映射 / CUDA Graph；合计 2.25×。
• 真机验证 ：未独立复现：需 30B + --stream-n 64 精确 tok/s（graph off/on 两档）。

七、小结（可复用结论）

认知	内容
探针先于排序	stub 逐段拆，才看清 attention 42% > FFN
位级等价 vs 位级一致	M=1 GEMV 是「位级等价」的免费午餐；attention/router 是「位级不一致但 greedy 未翻」的近似收益
放开约束解锁收益	router 并行化在「位级一致」下判否，放开后 7→4ms
延迟受限 vs 带宽受限	router/attention 都是延迟受限，减字节/加占用都不如加并行度

1. 大的收益是拼出来的：2.25× 不是一个大招，而是「探针定位真瓶颈 + 四个针对性小优化」的叠加（1.15× · 1.31× · 1.21× · 1.24×）。
2. 区分「位级等价」与「位级一致」：前者（M=1 GEMV）是免费午餐，后者（router 并行化）要付放开约束的代价，出手前先分清楚自己动的是哪一种。
3. 加速比用精确 tok/s 算：拿取整 ms 估算会把 2.25× 抬成 2.4×，口径不一致就会自我误导。
4. 收益随占比变化：CUDA Graph 单独 1.08×、叠加后 1.24×，因为它压的是「剩余的 launch 占比」，不是固定时间。
5. 收口也是一种结论：剩余项都落在 2--4ms 且收益递减，阶段收口于此，比继续挤 1% 更诚实。

一句话（收束）：2.25× 来自「探针定位真瓶颈」+「四个针对性小优化」， 而不是一个想象中的大优化；剩下四项都在 2--4ms 且收益递减，本轮收口于此。

---

相关篇目：第 13 篇 · Tiling 与 Tensor Core 判否、第 14 篇 · router 判否与 decode 真瓶颈、第 09 篇 · full-GPU decode 源码与配套资源：本仓库 gitee.com/pei-xiaogua...； CPU 推理源码 gitee.com/pei-xiaogua...