04 · 分组量化 GEMM 内核一句话一个内核同时服务 decodeM1与 prefillM32用「每线程固定几列、沿 K 扫」的映射让权重按 warp 满事务读、激活广播用 z-split 把小 M 的欠填充网格补满且只发一次 launch。前置建议先读 第 03 篇 · 权重条带缓存与融合权重。环境x86-64 NVIDIA RTX 5090sm_120· 阶段一 2B 稠密Qwen3-VL-2B q8。一、问题与结论问题做法判定一个内核怎么同时吃 decode 与 prefill同一映射M 走行、n 走列decode 时靠 z-split 补并行度实测可行权重与激活各怎么访存权重按 warp 满事务读、激活一次读全 warp 广播实测权重密集、激活广播小 M 网格欠填充怎么办沿 group 轴切 z一次 launch 完成所有 slice实测避免 launch 串行小投影该不该 offloadFLOP 门槛1e7自动把太小的留在 CPU实测见下节反直觉表二、背景要算的东西是逐块点积再乘 scaleout[m][n] Σ_g a_scale[m][g] · b_scale[n][g] · ( Σ_{k ∈ block g} Aq[m][k] · Wq[n][k] )G分组宽度Q8_0/Q4_0 是 32。prec0 int8 权重码[-127,127]Q8_01 int4[-8,7]Q4_0 nibble 已解包。整数点积必须整组完成再乘 scale——这是精确量的划分不能被拆开。头文件的 contract 与此一致include/npu/cuda/vllm_cuda.h布局与 RK3588 NPU 完全相同所以同一份 strip-cache artifact 喂两个后端而内核是「对 K 的普通归约」天然支持M1——这正是它相对 NPU cube要求M0 mod 4、无法服务 decode的结构性收益。三、核心机制3.1 两代内核int8 激活 → float 激活内核激活输入用在哪k_gw_gemmint8 Aq float asc早期 prefill offloadk_gw_gemm_f32actfloat Axfull-decode 现行 prefill两者的分块与「沿 K 扫」的结构一致只有两点不同一是激活的处理int8 码 scale vs float二是 f32 版本在M 1decode时把整块线程全映射到 n 方向避免 m 方向线程空转。第 11 篇会讲为什么最终选了 float 激活PPL Δ3.0% → 0.19%。3.2 线程映射/* 文件src/npu/cuda/vllm_cuda_kernels.cu节选 */#defineVC_NT64/* threads along n per block */#defineVC_MT4/* threads along m per block */#defineVC_VEC4/* output columns per thread *//* 文件src/npu/cuda/vllm_cuda_kernels.cuk_gw_gemm / k_gw_gemm_f32act节选 *//* 通用形状prefill / M1*/constintn0(blockIdx.x*VC_NTthreadIdx.x)*VC_VEC;/* 输出列 */constintmblockIdx.y*VC_MTthreadIdx.y;/* 输出行 *//* f32 版本在 M1decode时把线程全压到 n*/if(M1){n0(blockIdx.x*(VC_NT*VC_MT)threadIdx.y*VC_NTthreadIdx.x)*VC_VEC;m0;}每个线程负责一行m上连续的 4 个输出列[n0, n04)。内层沿 K 扫/* 文件src/npu/cuda/vllm_cuda_kernels.cuk_gw_gemm_f32act 的 float 激活路径节选 */for(intgg_begin;gg_end;g){constint8_t*wtWT(size_t)g*G*Npadn0;constfloat*xbarow(size_t)g*G;floatp00.f,p10.f,p20.f,p30.f;for(intkk0;kkG;kk){constfloatxvxb[kk];/* 整 warp 广播同一地址 */constint8_t*wwt(size_t)kk*Npad;p0(float)w[0]*xv;p1(float)w[1]*xv;p2(float)w[2]*xv;p3(float)w[3]*xv;}constfloat4 bs*reinterpret_castconstfloat4*(bsT(size_t)g*Npadn0);acc0p0*bs.x;acc1p1*bs.y;acc2p2*bs.z;acc3p3*bs.w;}int8 激活版本k_gw_gemm的内层取uchar4权重码、以int累加最后再乘a_scale与b_scale。两个访存特征值得点出权重VC_NT64线程 ×VC_VEC4字节 warp 一次读 128 字节连续数据——满事务。激活xb[kk]对同一行m的所有线程是同一个地址——一次读、全 warp 广播。这是「权重密集、激活广播」的取舍因为 decode 时瓶颈在权重带宽。3.3 z-split小 M 时怎么喂饱设备decode 时M1nblk_y 1nblk_x也不大——(x,y) 网格严重欠填充。解法是沿group 轴再切一刀用blockIdx.z/* 文件src/npu/cuda/vllm_cuda_kernels.cuk_gw_gemm节选 */constintzblockIdx.z;constintbasengroups/nsplit;constintremngroups%nsplit;constintg_beginz*base(zrem?z:rem);constintg_endg_beginbase(zrem?1:0);要点每个 slice 拥有整数个 scale block所以整数部分和仍是精确的——这是纯划分。一次 launch完成所有 slice不是循环 launch。decode 形状下nsplit可达 64逐个 launch 的开销会盖过整个 GEMM。结果落在dPart[z]最后用k_gw_reduce按固定顺序求和——固定顺序是确定性的前提。/* 文件src/npu/cuda/vllm_cuda_kernels.cuk_gw_reduce节选 */__global__voidk_gw_reduce(constfloat*__restrict part,float*__restrict dst,size_ttotal,intnsplit){constsize_ti(size_t)blockIdx.x*blockDim.xthreadIdx.x;if(itotal)return;floats0.f;for(intz0;znsplit;z)spart[(size_t)z*totali];dst[i]s;}nsplit的选择/* 文件src/npu/cuda/vllm_cuda_kernels.cuvcuda_dev_gemm_gw节选 */if(nblk_xy0nblk_xyVC_SPLIT_TARGETngroups1){nsplit(VC_SPLIT_TARGETnblk_xy-1)/nblk_xy;/* 目标 512 个 block */if(nsplitngroups)nsplitngroups;}while(nsplit1(size_t)nsplit*ou_szVC_PART_CAP_BYTES)nsplit--;/* part 缓冲上限 64 MiB */blockIdx.z 切 group权重满事务读同地址decode M1 网格欠填充nsplit 个切片VC_NT×VC_MT 线程映射每线程盯 4 输出列沿 K 扫 g_begin..g_end128 B/warp 权重xb[kk] 全 warp 广播各切片写 dPart[z]k_gw_reduce 固定序求和图 1解码/预填充共用映射下的数据流。blockIdx.x/y摊出「列 × 行」网格VC_NT64×VC_MT4每个线程盯住[n0, n04)那 4 个输出列、沿 K 扫过本 slice 的 group 区间两个访存特征正来自这套映射——权重每 warp 满事务读 128 B激活xb[kk]对同列线程是同一地址、一次读全 warp 广播。decodeM1时(x,y)网格欠填充改由blockIdx.z沿 group 轴切成nsplit片目标VC_SPLIT_TARGET512个 blocknsplit受 64 MiB 的VC_PART_CAP_BYTES上限约束各片写dPart[z]最后由k_gw_reduce按固定顺序求和保确定性。3.4 调用路径上的传输优化设备层vcuda_dev_gemm_gw()的传输形态是实测调出来的H2D 走pinned 暂存pageable 源会串行化全部cudaMemcpyAsync发在默认流只在 D2H 后做一次cudaDeviceSynchronize。/* 文件src/npu/cuda/vllm_cuda_kernels.cuvcuda_dev_gemm_gw节选 */memcpy(s-hAq,Aq,aq_sz);memcpy(s-hAsc,asc,as_sz);if(cudaMemcpyAsync(s-dAq,s-hAq,aq_sz,cudaMemcpyHostToDevice)!cudaSuccess||cudaMemcpyAsync(s-dAsc,s-hAsc,as_sz,cudaMemcpyHostToDevice)!cudaSuccess){...}.../* One async D2H into pinned staging THE single synchronize of the call. */cudaMemcpyAsync(s-hOut,s-dOut,ou_sz,cudaMemcpyDeviceToHost);cudaDeviceSynchronize();旧形态阻塞 D2H 额外一次cudaDeviceSynchronize多付了一整个往返。profile 显示对小投影传输/同步这对开销而不是内核本身才是主导。3.5 小投影为什么单独 offload 是亏的2B q8的 decode 桶实测ms/token投影单投影 offload 前单投影 offload 后结论QKV / O / K / V6.39 / 3.70 / …7.24 / 3.73 / …变慢gate/up16.967.56大幅变快down10.793.98大幅变快lm_head6.87—变快最大矩阵原因是每次调用的固定往返 ~100 µs超过小投影的计算量。所以QKV/K/V 单独 offload 是亏的 → 必须融合QKV 三合一、gateup 二合一把三次往返压成一次用 FLOP 门槛1e7把「太小不值得」的投影自动留在 CPUM1时 Q/O 是 8.4 MFLOP、K/V 4.2 MFLOP排除gate/up/down 25 MFLOP 与 lm_head 622 MFLOP通过。这条「小算子别 offload、融合才划算」的结论直接决定了 第 05 篇 里卸载集合的结构。四、实测数据口径数值判定标注QKV 单投影 offload2B q8 decode6.39 → 7.24 ms/tok变慢实测gate/up16.96 → 7.56 ms/tok大幅变快实测down10.79 → 3.98 ms/tok大幅变快实测lm_head6.87 → …变快实测每次调用固定往返~100 µs小投影的主导开销实测M1 各投影 FLOPsQ/O 8.4 M / K/V 4.2 M排除gate/up/down 25 M、lm_head 622 M通过门槛1e7干净切开实测五、边界与已知限制这里的实测是2B q8的口径小投影「负收益」的临界由固定往返~100 µs决定换设备/换后端不能照搬这条门槛需要按自己的往返成本重测。prec只支持 0/1int8/int4 码K % G ! 0一律退回 CPU。z-split 的 part 缓冲有 64 MiB 上限超限时自动减少nsplit并行度换显存。内核的数值确定性依赖k_gw_reduce的固定求和顺序以及块内整数点积整组完成改变划分方式会改变位级结果。CPU 对照迁移前基线CPU 参考kestrel-llm/src/model/vllm_safetensors.c函数st_q4_row_dot/q4_dense4_x86dyn_matvec_q8—— CPU 侧按行q4 走 4 行组逐块反量化点积。迁移要点CPU 侧逐行/逐 4 行组算 → CUDA 侧用一个映射同时吃 decode(M1) 与 prefill(M32)每线程固定 4 列沿 K 扫权重满事务读、激活广播小 M 用blockIdx.z切 groupz-split 一次 launch固定顺序 reduce 保确定性。真机验证命中 E2selftestint8/int4 × decode/prefill/非整除 N 全过。六、小结可复用结论一个映射吃两种负载每线程固定几列、沿 K 扫让权重满事务读、激活广播——这是 decode 权重带宽受限下的正确取舍。小 M 靠 z-split 补并行度且必须一次 launchdecode 形状下nsplit可达 64逐个 launch 会把 GEMM 本身盖过去。精确量要整组完成整数点积在块内做完再乘 scale才能做纯划分、用固定顺序 reduce 保确定性。传输/同步才是小投影的主导开销pinned 暂存 单次同步是实测调出来的不是默认写法。小算子不该单独 offloadFLOP 门槛加融合才是把「快」变成「更快」的正确方向。相关篇目第 03 篇 · 权重条带缓存与融合权重、第 05 篇 · 卸载集合源码与配套资源本仓库 https://gitee.com/pei-xiaoguang/kestrel-llm-cuda.gitCPU 推理源码 https://gitee.com/pei-xiaoguang/kestrel-llm