Kestrel_VLLM_CUDA -09 · full-GPU decode：单 token 全设备

09 · full-GPU decode单 token 全设备一句话decode 一个 token 只付一次 H2Dembedding 一次 D2Hlogits中间所有层、所有中间量都留在设备上——把「每层一次宿主↔设备往返」换成「整条链路一次往返」。前置建议先读 第 08 篇 · MoE 专家权重显存驻留本篇要用到它的专家内核。环境x86-64 NVIDIA RTX 5090sm_120· 阶段三 30B-A3B MoEQwen3-30B-A3B q4128 专家 / top-8 / dim2048 / moe_ffn768 / 48 层。一、问题与结论指标GPU full-decodeCPU判定标注decode30B-A3B q4--stream-n 6452.0 tok/s19 ms/tok15.0 tok/s66 ms/tok约3.5×实测每 token 的传输1 次 H2Dembedding 1 次 D2Hlogits—设计目标实测二、背景为什么要全设备通用 offload 路径是「每层一次往返」宿主算好激活 → 传上去 → GEMM → 传回来 → 下一个算子。decode 时每层有好几次这样的往返而且每次都有宿主↔设备同步。full-decode 换了一个思路隐藏态、QKV、FFN 中间量、K/V 镜像全部常驻显存一个 token 的整条链路在设备上跑完。宿主只做两件事喂 token embedding、取 logits。设备 GPU宿主 CPU设备 GPU宿主 CPU开跑前全量预检所有权重 key任一缺失即返回 -1 且不改状态loop[每层 l 0..47]H2D embedding每 token 仅 1 次RMSNorm→融合QKV→QK-norm→RoPE→KV量化attentionfloat query × int8 KVFFNMoE 或稠密分支final RMSNorm→lm_headD2H logits每 token 仅 1 次图 1full-decode 的宿主↔设备往返。一个 token 只付一次 H2Dembedding 一次 D2Hlogits48 层的注意力与 FFNMoE 分支s16 量化→router→gate/up→down全在设备上跑完对比通用 offload 路径的「每层多次往返 每层同步」这就是 52.0 tok/s19 ms/tok≈3.5×的来源。开跑前做 hit-only 全量预检任一权重缺失即整体返回 -1、回落 CPU decode 循环。实测30B-A3B q4--stream-n 64GPU full-decodeCPUdecode52.0 tok/s19 ms/tok15.0 tok/s66 ms/tok约3.5×。三、核心机制3.1 驻留了什么vcuda_dev_decode_init()一次分配全部设备缓冲/* 文件src/npu/cuda/vllm_cuda_kernels.cuvcuda_dev_decode_init节选 */if(cudaMalloc((void**)s-d_x,(size_t)dim*4)!cudaSuccess||cudaMalloc((void**)s-d_normed,(size_t)dim*4)!cudaSuccess||cudaMalloc((void**)s-d_res,(size_t)dim*4)!cudaSuccess||cudaMalloc((void**)s-d_qkv,(size_t)qkv_n*4)!cudaSuccess||cudaMalloc((void**)s-d_attn_out,(size_t)nh*hd*4)!cudaSuccess||cudaMalloc((void**)s-d_gu,(size_t)ff*8)!cudaSuccess||cudaMalloc((void**)s-d_act,(size_t)ff*4)!cudaSuccess||cudaMalloc((void**)s-d_logits,(size_t)vocab*4)!cudaSuccess||cudaMalloc((void**)s-d_asc,(size_t)(aq_n/G)*4)!cudaSuccess||cudaMalloc((void**)s-d_aq,(size_t)aq_n*2)!cudaSuccess||cudaMalloc((void**)s-d_attn_norm,(size_t)nl*dim*4)!cudaSuccess||cudaMalloc((void**)s-d_ffn_norm,(size_t)nl*dim*4)!cudaSuccess||cudaMalloc((void**)s-d_final_norm,(size_t)dim*4)!cudaSuccess||(q_normcudaMalloc((void**)s-d_qnorm,(size_t)nl*hd*4)!cudaSuccess)||(k_normcudaMalloc((void**)s-d_knorm,(size_t)nl*hd*4)!cudaSuccess)){cudaGetLastError();vc_decode_free(s);return-1;}缓冲形状用途d_x / d_normed / d_res[dim]隐藏态、归一化、残差d_qkv[nh·hd 2·nkv·hd]融合 QKV 输出d_attn_out[nh·hd]注意力输出d_gu / d_act[2·ff] / [ff]稠密分支 gateup、swiglud_logits[vocab]最终 logitsd_aq / d_ascs16 激活 scaleMoE 激活d_attn_norm / d_ffn_norm / d_final_norm[nl·dim]三种 RMSNorm 权重f32不进权重缓存d_qnorm / d_knorm[nl·hd]可选 Q/K per-head normd_freqs[hd/2]RoPE 频率表MoE 专家缓冲见 第 08 篇gate/up/down q4 router3.2 单 token 的层循环vc_decode_body()的结构graph 与非 graph 两条路径共用同一段循环体差别只在提交方式H2D embedding for l in 0..nl-1: # ---- 注意力 ---- k_copy(d_res ← d_x) # 存残差 k_rms_norm(d_normed ← d_x, attn_norm[l]) vc_gemm_fused_dev(QKV) # 融合 QKVfloat 激活 k_qk_norm(Q/K per-head) # 可选Qwen3-VL k_rope(q, k, pos) k_kv_quantize(→ d_k/d_v 镜像 scale) k_attn_decode_f32q # f32 query int8 KV vc_gemm_dev(O) # float 激活 k_add(d_x ← d_res d_normed) # 注意力残差 # ---- FFN ---- k_copy(d_res ← d_x) k_rms_norm(d_normed ← d_x, ffn_norm[l]) if is_moe: k_quantize_blocks_s16(d_normed → d_aq) # 激活量化一次全专家共享 k_moe_router(router → sel/pr) k_moe_gate_up(gate/up → av) k_quantize_blocks_s16(av → avq) k_moe_down(down → y, 加权聚合) k_add(d_x ← d_res y) else: vc_gemm_fused_dev(GU); k_swiglu; vc_gemm_dev(DOWN); k_add k_rms_norm(d_normed ← d_x, final_norm) vc_gemm_dev(lm_head → d_logits) D2H logits对应的真实发射序列MoE 分支/* 文件src/npu/cuda/vllm_cuda_kernels.cuvc_decode_body 的 FFN 段节选 *//* quantize the FFN input once (shared by every experts gate/up) */k_quantize_blocks_s16dim3(d/G,1),1,0,stream(s-d_normed,s-d_aq,s-d_asc,1,d,G);/* router - top-k - softmax */k_moe_router_gemm(unsigned)(ne2),DC_BLK,0,stream(s-d_moe_routerr_loff,s-d_normed,s-d_moe_logits,ne,d);k_moe_router_topk1,1,0,stream(s-d_moe_logits,s-d_moe_sel,s-d_moe_pr,ne,tk);/* gate/up - av silu(gate) * up */k_moe_gate_updim3((unsigned)ge4,(unsigned)tk),DC_BLK,0,stream(s-d_moe_gateg_loff,s-d_moe_upg_loff,s-d_aq,s-d_asc,s-d_moe_sel,s-d_moe_av,ne,ef,d,tk,nbG,ge4,ggB,moe_probe);/* quantize av (per expert) */k_quantize_blocks_s16dim3(ef/G,tk),1,0,stream(s-d_moe_av,s-d_moe_avq,s-d_moe_avd,tk,ef,G);/* down - y (weighted aggregate in expert order) */k_moe_down(unsigned)(d2),DC_BLK,0,stream(s-d_moe_downd_loff,s-d_moe_avq,s-d_moe_avd,s-d_moe_sel,s-d_moe_pr,s-d_moe_y,ne,ef,d,tk,cbe,gdB,moe_probe);几条设计点激活量化一次共享MoE 的 FFN 输入只量化一次d_aqtop-k 所有专家的 gate/up 共用省掉 k 次重复量化。残差用d_res承接k_copy存、k_add取避免为每层单独开缓冲。专家贡献按 j 序累加与 CPU 求和顺序一致见 第 08 篇。3.3 hit-only 契约重要vcuda_dev_decode_run()是全有或全无的开跑前先把所有需要的权重 key 检查一遍任何一个缺失就返回 -1且不改变任何状态/* 文件src/npu/cuda/vllm_cuda_kernels.cuvcuda_dev_decode_run 的预检节选 */constintis_moes-dc_is_moe;constintnls-dc_nl;if(is_moes-dc_G!32)return-1;/* MoE expert tiles are 32-col grouped */for(intl0;lnl;l){if(!vc_ent_find(s,vc_wkey(l,VC_FUSED_QKV))||!vc_ent_find(s,vc_wkey(l,VC_PROJ_O)))return-1;if(!is_moe(!vc_ent_find(s,vc_wkey(l,VC_FUSED_GU))||!vc_ent_find(s,vc_wkey(l,VC_PROJ_DOWN))))return-1;}if(!vc_ent_find(s,vc_wkey(VC_LMHEAD_LAYER,VC_PROJ_LMHEAD)))return-1;为什么这个契约重要调用方引擎在拿到 -1 时会回落 CPU decode 循环。如果decode_run跑了一半才失败K/V 镜像就被写脏了CPU 回落的起点不对 → 静默错。全有或全无把这个问题消灭在源头。MoE 分支还有个细节专家权重在专用 q4 缓冲里不在 int8 权重缓存所以 GU/DOWN 的缓存项不要求存在is_moe时跳过检查。3.4 为什么 attention 用 float query有两个 query 版本的 attention内核query说明k_attn_decode_q8int8每 32 块量化快但多一层量化误差k_attn_decode_f32qfloat与 CPU int8-KV 路径同口径量化的选择不是「能省就省」int8 query 的量化误差约 0.7%。它不会在单层里显形但在48 层30B 上会复利累积最终把 greedy 的 argmax 推偏。所以 full-decode 用 float query。float query 版的头部逻辑GQA 映射kh h * nkv / nh、float query × int8-KV 的即时反量化/* 文件src/npu/cuda/vllm_cuda_kernels.cuk_attn_decode_f32q 的 Pass 1节选 */constinthblockIdx.x;constintseq_lendyn[1]1;/* n_attn visible tokens */constintkh(h*nkv)/nh;constint8_t*kbk(size_t)kh*hd;/* Pass 1: float-query scores block max. */for(inttthreadIdx.x;tseq_len;tblockDim.x){constint8_t*ktkb(size_t)t*kv_dim;constfloatkssks[(size_t)t*nkvkh]*(1.0f/127.0f);floats0.0f;for(inti0;ihd;i)sqt[i]*(float)kt[i];ss*kss*scale;scores[t]s;if(slm)lms;}两者都做「int8 存 K/V f32 scale 即时反量化」K/V 镜像用 int8 存带宽减半内核里读 int8 乘 scale 还原——与 CPU 的flash_attn_single_q_q8_neon同口径。attention 的线程组织是一格 block 一个 query headVCUDA_ATT_T256线程三次 pass打分块内 max → exp求和 → 加权 V 求和动态共享内存放scores[seq_len]。GQA 的映射是kh h * nkv / nh。3.5 KV 镜像设备侧 K/V 是扁平[nl][max_seq][kv_dim]int8 [nl][max_seq][nkv]f32 scale每层每 token 追加一行k_kv_quantize写d_k/d_v/d_ks/d_vs。这个镜像被两条路径共享VLLM_CUDA_ATTN只加速 attention与VLLM_CUDA_DECODEfull decode。四、实测数据口径数值判定标注GPU full-decode30B-A3B q4--stream-n 6452.0 tok/s19 ms/tok约 3.5×实测CPU同口径15.0 tok/s66 ms/tok基线实测每 token 传输1 次 H2D 1 次 D2H达成设计目标实测五、边界与已知限制口径是--stream-n 64这是 64-token decode 的 tok/s 精确值log 里的ms/tok为取整展示加速比一律用 tok/s 算。hit-only所有需要的权重必须先驻留任何一项缺失就整体返回 -1、回落 CPU——没有「部分 GPU」的中间态。MoE 要求 32 列分组is_moe G ! 32直接判拒。设备镜像与 CPU 的数值一致性只在前 16 token 的 greedy 上验证见 第 10 篇整模型的统计口径见 第 11 篇。这里的 3.5× 是阶段三 full-decode 的基线值后续四连击优化见 第 15 篇。CPU 对照迁移前基线CPU 参考kestrel-llm/src/model/vllm_safetensors.c函数flash_attn_single_q_q8_neon—— CPU 侧 attention 用 int8 KV f32 scale 即时反量化。迁移要点CPU 侧每层在宿主内存里算、逐层与设备往返 → CUDA 侧把往返从「每层一次」压到「每 token 一次 H2D 一次 D2H」hidden/QKV/FFN/KV 镜像全驻留attention 用 float queryint8 query ~0.7% 误差在 48 层复利会破 greedyhit-only 预检保证回退起点干净。真机验证未独立复现需 30B-A3B q4 --stream-n 64的 tok/s 与传输计数。六、小结可复用结论把往返次数从「每层」压到「每 token」隐藏态、QKV、FFN 中间量、KV 镜像全部常驻显存宿主只负责喂 embedding、取 logits收益来自消除逐层同步。共享中间量FFN 输入只量化一次top-k 所有专家共用省掉 k 次重复量化。残差用一块缓冲承接k_copy存、k_add取不用为每层单独开缓冲。全有或全无的预检失败即回落 CPU绝不在设备上留下半成品状态——否则 CPU 回落起点被污染出静默错。精度不能满足于「能省就省」query 用 float、KV 用 int8是因为 int8 query 的 ~0.7% 误差会在 48 层上复利破 greedy。七、本环境的跑法# prefill 走 CPU隔离 decode 路径只让解码上 GPU$env:VLLM_CUDA_DECODE1$env:VLLM_CUDA_INFER1$env:VLLM_CUDA_FORCE_CPU1.\build-x64\vllm_kestrel_x64.exe--cuda--model 模型目录--stream-test--stream-ctx 0--stream-n 64VLLM_CUDA_FORCE_CPU1在这里的用途是隔离——让 prefill 完全走 CPU这样 decode 的实测与对拍不被 prefill offload 的任何变量干扰第 10 篇正是靠这个搭配找出 bug 的。下一篇是这轮迁移里最值钱的一篇一个指针越界把上面所有努力都打回原形。相关篇目第 08 篇 · MoE 专家权重显存驻留、第 10 篇 · 一个指针越界引发的血案、第 15 篇 · decode 优化四连击源码与配套资源本仓库 https://gitee.com/pei-xiaoguang/kestrel-llm-cuda.gitCPU 推理源码 https://gitee.com/pei-xiaoguang/kestrel-llm