RWKV 批量推理中 Prefill 的正确打开方式

RWKV 批量推理中 Prefill 的正确打开方式项目地址https://github.com/AUXStar/RWKV-Server文章目录RWKV 批量推理中 Prefill 的正确打开方式一、引言二、Prefill vs Decode计算特性对比实测吞吐数据三、Albatross v3a 的 Kernel 特化路径select_path 函数B1, T1 的专用优化链四、混跑的代价为什么不能把 prefill 和 decode 放在同一个 batch问题一kernel 路径退化问题二padding 浪费问题三kernel 切换开销量化分析五、RWKV-Server 的设计单线程 prefill 无锁 decode核心设计要点代码引用task.py 的 prefill 方法为什么这个设计是正确的六、实测数据验证动态缩容过程七、与 Transformer 推理框架的对比为什么 vLLM/SGLang 需要 Chunk Prefill为什么 RWKV 不需要八、结论一、引言如果你做过大模型推理服务大概率接触过 vLLM 或 SGLang。这两个框架的核心思路高度一致Chunk Prefill Dynamic Batching。把多个请求的 prefill 阶段切分成小块和 decode 阶段混合在同一个 batch 里执行以此最大化 GPU 利用率。这个思路对 Transformer 架构完全正确。Transformer 的 self-attention 是 O(n^2) 的prefill 阶段的长序列注意力计算是整个推理过程中最重的部分必须通过 batching 和 chunking 来摊薄开销。但 RWKV 不一样。核心论点RWKV 的 O(n) 线性注意力架构让 prefill 和 decode 的工程权衡与 Transformer 完全不同。在 RWKV 中prefill 只占总 token 处理量的约3%盲目套用 Transformer 的 Chunk Prefill 策略不仅不会提升吞吐反而会因为破坏 kernel 特化路径而导致性能倒退。本文将从计算特性、kernel 实现、架构设计三个层面系统分析为什么 RWKV-Server 选择了一条看似朴素的路径单线程串行 prefill batch decode以及这条路径如何与 Albatross v3a 的 CUDA kernel 特化策略完美契合。二、Prefill vs Decode计算特性对比在讨论工程策略之前先看一组实测数据。以下表格对比了 RWKV 架构中 prefill 和 decode 两个阶段的计算特性特性PrefillDecode每步 token 数整个 prompt 长度通常 20~2000 tok1 token / step计算复杂度O(n) 线性扫描稠密 GEMMO(1) 状态更新GEMV / 稀疏 kernelGPU 利用率中等矩阵维度较大极低单行计算依赖专用 kernel总 token 占比~3%实测 20,140 / 548,180~97%实测 528,040 / 548,180Kernel 特化程度通用稠密路径CMIX_DENSE高度特化SPARSE / fused / split-k能否与 decode 混跑会退化整个 batch 的 kernel 路径需要 B1, T1 才能走最优路径关键发现prefill 只占总 token 处理量的 3%。这意味着即使你把 prefill 的吞吐提升 10 倍对整体吞吐的影响也只有 0.3%。真正决定 RWKV 推理性能的是 decode 阶段的效率而 decode 效率完全取决于 kernel 是否能走 B1, T1 的专用路径。实测吞吐数据模型 GPU256 任务入队耗时Prefill 吞吐Decode 峰值吞吐Prefill 占比2.9B RTX 40904.97s~4,052 tok/s11,081 tok/s3.5%7.2B RTX 40908.64s~2,331 tok/s6,373 tok/s3.2%2.9B 5070 Ti Laptop––5,169 tok/s–注意 decode 峰值吞吐远高于 prefill 吞吐。这不是因为 decode 计算量更大而是因为 decode 走了高度优化的专用 kernel详见下一节而 prefill 只能走通用稠密路径。三、Albatross v3a 的 Kernel 特化路径RWKV-Server 使用的推理后端是Albatross v3arwkv7_fast_v3a.py。v3a 的核心设计思想是根据select_path(B, T)的返回值为不同的 batch size (B) 和 sequence length (T) 组合选择完全不同的 CUDA kernel 路径。select_path 函数以下代码摘自rwkv7_fast_v3a.py第 156-178 行。这个函数是整个 kernel 调度系统的入口defselect_path(B,T): 根据 batch size (B) 和 sequence length (T) 选择最优的 kernel 执行路径。 返回值决定了后续所有 cmix / tmix 操作使用哪套 kernel。 rowsB*Tifrows1:# B1, T1: 单样本单步 decode -- 最优路径returnCMIX_B1T1_SPARSEelifB1andT1:# B1, TN: 单样本 prefill -- 稠密 GEMM 路径returnCMIX_DENSEelifB1andT1:# B1, T1: 多样本 decode -- 通用稀疏路径returnCMIX_BATCHED_SPARSEelse:# B1, T1: 混合 batch -- 退化为稠密路径returnCMIX_DENSE关键在于rows B * T这一行。整个 kernel 选择体系围绕这个乘积展开。当且仅当rows 1即 B1, T1时才能进入最优的CMIX_B1T1_SPARSE路径。B1, T1 的专用优化链当select_path返回CMIX_B1T1_SPARSE时v3a 会启用一整套专用 CUDA kernel。这些 kernel 是为单行计算量身定制的与通用 GEMM 路径相比有质的飞跃linear_orig_row1_exact_f16_kernel— 专门处理单行矩阵乘法的 CUDA kernel避免了通用 GEMM 的启动开销和 padding 浪费。对于 M1 的矩阵乘法这个 kernel 比 cuBLAS GEMM 快数倍。linear_f16_m1_splitk— M1 场景下的 split-k 优化。将 K 维度切分成多个小块并行计算充分利用 GPU 的并行度最后做 reduction。在 decode 阶段的线性投影层中效果显著。fused add_layer_norm_cmix_mix_f16— T1 专用的融合 kernel。将 layer norm、残差加法、channel mixing 的 mix 操作融合为单次 kernel launch大幅减少 GPU kernel launch 开销和全局内存访问次数。LN1_TMIX_FUSE— 当B1 and T1时layer norm 和 time mixing 的 mix6 操作进一步融合。这是 v3a 最深层的优化之一将多个逐元素操作合并为一个 kernel。以下代码展示了forward_from_x中 T1 时的 fused kernel 分支第 382 行附近# forward_from_x 中的 T1 融合分支ifT1andself.ln1_tmix_fuse:# 当 B1, T1 时layer_norm 和 tmix mix6 融合为单次 kernel launchxfused_ln_tmix_mix(x,self.ln1_w,self.ln1_b,self.tmix_mix_w,self.tmix_mix_b,self.tmix_receptance,self.tmix_key,self.tmix_value)else:# 通用路径分步执行 layer norm 和 tmixxlayer_norm(x,self.ln1_w,self.ln1_b)xtmix_forward(x,...)而当 B1, TNprefill时走的是完全不同的路径# cmix_sparse_one vs cmix_sparse_rows 的调用第 583-585 行ifself.cmix_modeCMIX_B1T1_SPARSE:# 单行专用cmix_sparse_one 处理单个 tokenxcmix_sparse_one(x,self.cmix_key,self.cmix_value,self.cmix_receptance,self.cmix_output)elifself.cmix_modeCMIX_BATCHED_SPARSE:# 多行通用cmix_sparse_rows 处理 batchxcmix_sparse_rows(x,self.cmix_key,self.cmix_value,self.cmix_receptance,self.cmix_output)else:# CMIX_DENSE: 退化为通用稠密矩阵乘法xcmix_dense(x,self.cmix_key,self.cmix_value,self.cmix_receptance,self.cmix_output)要点v3a 的 kernel 特化深度极高。B1, T1 的路径不仅换了 kernel还做了多层算子融合。这条路径的性能优势来自单行计算这个前提条件任何破坏这个前提的设计比如把 prefill 和 decode 混在同一个 batch都会导致整条优化链失效。四、混跑的代价为什么不能把 prefill 和 decode 放在同一个 batch理解了 v3a 的 kernel 特化机制混跑的问题就一目了然了。问题一kernel 路径退化假设一个 batch 中有 7 个 decode 样本T1和 1 个正在 prefill 的样本T128。此时rows 8 * 128 1024取最大 Tselect_path会返回CMIX_DENSE。结果所有 8 个样本都失去 T1 的专用优化。原本走linear_orig_row1_exact_f16_kernel的 decode 样本现在被迫走通用 cuBLAS GEMM。性能损失可达数倍。问题二padding 浪费不同长度的 prefill 拼在同一个 batch 里需要 padding 到最长序列的长度。如果 batch 中有一个 1024 token 的 prefill 和七个 1 token 的 decode有效计算量只有 1031但实际计算量是 81928 x 102487% 的计算全是 padding。问题三kernel 切换开销prefill 阶段使用的是大矩阵 GEMMMN, Kd_model, Nd_ffdecode 阶段使用的是 GEMVM1, Kd_model, Nd_ff。两种计算的内存访问模式、并行度、寄存器使用完全不同。在同一个 batch 中交替执行会导致 GPU 的 L2 cache、shared memory、warp scheduler 频繁切换状态产生显著的切换开销。量化分析在 RWKV 中prefill 的计算量占比极低~3%。即使你通过 chunk prefill 把 prefill 的吞吐提升 2-3 倍整体吞吐的提升也只有3% * 2 6%。而混跑导致的 decode kernel 退化可能让 decode 吞吐下降 30-50%整体吞吐反而下降97% * 30% 29%。结论在 RWKV 架构下prefill 和 decode 混跑是典型的捡芝麻丢西瓜。prefill 的优化空间极小而 decode 的优化空间极大。任何破坏 decode kernel 特化的设计都是得不偿失的。五、RWKV-Server 的设计单线程 prefill 无锁 decode基于以上分析RWKV-Server 采用了状态机解耦的策略prefill 和 decode 通过任务状态自然分离prefill 串行执行并受锁保护decode 在调度器主循环中无锁 batch 执行。新任务提交 │ ▼ Task 创建 ──► prefill() ──► 状态 READY │ │ │ prefill_lock 保护 │ 多个 prefill 互斥 │ │ │ model.forward(prompt) │ B1, TN, 走 CMIX_DENSE │ │ └────────────┘ │ ▼ 状态 READY │ ┌────────────┘ │ ▼ 调度器收集 READY 任务 │ ▼ update_batch() ──► 注入 Worker 槽位 │ ▼ engine.generate() ──► 无锁执行 │ B1, T1, 走 CMIX_BATCHED_SPARSE │ decode 与 prefill 可并行 ▼ _collect() ──► 输出收集 │ ▼ 任务完成 / 状态 FINISHED核心设计要点Prefill 串行 锁保护每个新任务的 prefill 在Task.__init__()中同步完成或独立线程中通过prefill_lock保证多个 prefill 之间互斥。锁只保护model.forward()调用不锁任务列表。Decode 无锁 batch调度器主循环中的engine.generate()直接调用model.forward()没有任何锁保护。decode 和 prefill 可以同时在 GPU 上执行只要 CUDA stream 不冲突。状态机自然解耦PREFILL → READY → RUNNING → FINISHED 的四状态流转让 prefill 和 decode 通过任务状态而非锁来协调。READY 状态的任务才会被调度器 pick 进 batch。Task.__enter__ / __exit__上下文管理器自动管理 GPU/CPU 数据迁移。被调度器 pick 时cuda()上 GPU完成时cpu()回 CPU。代码引用task.py 的 prefill 方法classTask:def__init__(self,prompt,model_loader,batch_sampler,...):# ... 初始化 state ...self.prefill_lockprefill_lockifprefill_lockelseNullLock()self._statusStatus.PREFILL self.prefill(prompt)# 构造时同步 prefillself.cpu()# prefill 完成后立即回 CPUdefprefill(self,prompt):self._statusStatus.PREFILL promptself.tokenize(prompt)iflen(prompt)2:ifself.shift_state.device!cuda:self.cuda()tokenstorch.tensor(prompt[:-1],dtypetorch.long,devicecpu)self.prefill_lock.acquire()# 只锁 prefill 的 forwardself.model_loader.model.forward(tokens,(self.shift_state,self.wkv_state,self.elapsed_t))self.prefill_lock.release()self.current_tokenprompt[-1]self._statusStatus.READY# prefill 完成进入 READYdef__enter__(self):self.prepare()# cuda() 上 GPUreturnselfdef__exit__(self,*args):self.cpu()# 回 CPU为什么这个设计是正确的Prefill 不阻塞 decodeprefill_lock只保护 prefill 的forward()decode 的generate()完全无锁。两者可以并行执行GPU 不会被 prefill 独占。锁只防 prefill 互斥多个任务同时 prefill 时会串行通过prefill_lock避免多个 prefill 同时抢占 GPU 导致资源争抢。但 decode 不受影响。状态机自然解耦PREFILL 状态的任务不会被调度器 pick只有 READY 状态的任务才能进入 decode batch。不需要额外的调度逻辑。串行 prefill 的开销可忽略prefill 只占总 token 的 3%即使 256 个任务串行 prefill总耗时也只有约 5ms2.9B或 9ms7.2B。对整体延迟的影响微乎其微。六、实测数据验证以下是三组实测数据验证了上述设计的有效性配置256 任务入队耗时Prefill 吞吐Decode 峰值吞吐Prefill 占比Prefill 总耗时2.9B RTX 40904.97s~4,052 tok/s11,081 tok/s3.5%~5.0ms7.2B RTX 40908.64s~2,331 tok/s6,373 tok/s3.2%~8.6ms2.9B 5070 Ti Laptop––5,169 tok/s––动态缩容过程RWKV-Server 的 decode 调度器实现了动态 batch 缩容。当 batch 中的任务陆续完成时batch size 从 256 逐步缩减Batch Size2561286432168421缩容耗时每次缩容 15ms含 kernel 重编译–缩容过程如此迅速是因为 RWKV-Server 不需要像 vLLM 那样重新分配 KV Cache。RWKV 的状态是一个固定大小的向量与序列长度无关缩容只需要减少 batch 维度不涉及内存重分配。对比 Transformer在 vLLM/SGLang 中batch 缩容需要释放和重新分配 PagedAttention 的 KV Cache blocks这个过程可能需要数百毫秒。RWKV 的固定大小状态向量让缩容变成了一个近乎零开销的操作。七、与 Transformer 推理框架的对比维度RWKV-ServervLLM / SGLang注意力机制O(n) 线性注意力RNN-like 状态O(n^2) self-attentionPrefill 策略单线程串行B1, TNChunk Prefill batchingDecode 策略Dynamic batchingB1, T1Continuous batchingPrefill/Decode 混跑不混跑状态机解耦混跑同一 batch混跑原因混跑会退化 decode kernelPrefill 是瓶颈必须 batchingPrefill 占比~3%~30-50%取决于序列长度状态管理固定大小向量无需 KV CachePagedAttention KV CacheBatch 缩容 15ms无内存重分配~100-500msKV Cache 回收Kernel 特化极深B1,T1 有专用融合 kernel较浅FlashAttention 通用为什么 vLLM/SGLang 需要 Chunk PrefillTransformer 的 self-attention 在 prefill 阶段的计算量是 O(n^2)。一个 4096 token 的 prompt 需要 4096 x 4096 16M 次注意力计算。而 decode 阶段每步只需要 1 x n 次计算。Prefill 的计算量可以占整个推理过程的 30-50%。在这种情况下prefill batching 的收益巨大通过将多个 prefill 拼成一个大矩阵做 GEMMGPU 利用率可以从 10% 提升到 80%。即使 decode 阶段的 kernel 路径因此退化整体吞吐仍然是提升的。为什么 RWKV 不需要RWKV 的线性注意力在 prefill 阶段是 O(n) 的逐 token 扫描计算量远小于 Transformer。实测数据表明 prefill 只占总 token 的 3%。即使不做任何 batching 优化prefill 的总耗时也只有几毫秒。更关键的是RWKV 的 decode 阶段依赖高度特化的 kernelB1, T1 路径这些 kernel 的性能优势来自单行计算这个前提。如果为了 batching prefill 而破坏这个前提decode 吞吐的损失远大于 prefill batching 的收益。八、结论RWKV 的 O(n) 线性注意力架构从根本上改变了 prefill 和 decode 的工程权衡。在 Transformer 中prefill 是计算瓶颈需要通过 batching 和 chunking 来优化。在 RWKV 中prefill 只占总计算量的 3%真正的瓶颈在 decode 阶段的 kernel 效率。RWKV-Server 的设计选择 –单线程串行 prefill batch decode– 不是因为偷懒或不够先进而是基于对 RWKV 架构特性的深刻理解与 v3a kernel 特化策略完全吻合prefill 走 B1, TN 的稠密路径decode 走 B1, T1 的专用稀疏路径两条路径互不干扰。串行 prefill 的开销可忽略3% 的 token 占比意味着即使不做任何优化prefill 的总耗时也只有几毫秒。保护 decode 的 kernel 特化不混跑意味着 decode 始终能走最优路径不受 prefill 的干扰。实现简单、正确、高效没有复杂的 chunking 逻辑没有 prefill/decode 的调度博弈没有 padding 浪费。一句话总结架构决定策略。RWKV 的线性注意力让简单成为最优解。不要把 Transformer 的工程经验盲目套用到 RWKV 上 – 理解你的架构然后为它选择正确的路径。项目地址https://github.com/AUXStar/RWKV-ServerAlbatross v3a CUDA kernels by BlinkDL