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

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

项目地址：https://github.com/AUXStar/RWKV-Server

一、引言

如果你做过大模型推理服务，大概率接触过 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 两个阶段的计算特性：

关键发现：prefill 只占总 token 处理量的 3%。这意味着即使你把 prefill 的吞吐提升 10 倍，对整体吞吐的影响也只有 0.3%。真正决定 RWKV 推理性能的是 decode 阶段的效率，而 decode 效率完全取决于 kernel 是否能走 B=1, T=1 的专用路径。

实测吞吐数据

注意 decode 峰值吞吐远高于 prefill 吞吐。这不是因为 decode 计算量更大，而是因为 decode 走了高度优化的专用 kernel（详见下一节），而 prefill 只能走通用稠密路径。

三、Albatross v3a 的 Kernel 特化路径

RWKV-Server 使用的推理后端是Albatross v3a（rwkv7_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。 """rows=B*Tifrows==1:# B=1, T=1: 单样本单步 decode -- 最优路径return'CMIX_B1T1_SPARSE'elifB==1andT>1:# B=1, T=N: 单样本 prefill -- 稠密 GEMM 路径return'CMIX_DENSE'elifB>1andT==1:# B>1, T=1: 多样本 decode -- 通用稀疏路径return'CMIX_BATCHED_SPARSE'else:# B>1, T>1: 混合 batch -- 退化为稠密路径return'CMIX_DENSE'

关键在于rows = B * T这一行。整个 kernel 选择体系围绕这个乘积展开。当且仅当rows == 1（即 B=1, T=1）时，才能进入最优的CMIX_B1T1_SPARSE路径。

B=1, T=1 的专用优化链

当select_path返回CMIX_B1T1_SPARSE时，v3a 会启用一整套专用 CUDA kernel。这些 kernel 是为"单行计算"量身定制的，与通用 GEMM 路径相比有质的飞跃：

• linear_orig_row1_exact_f16_kernel— 专门处理单行矩阵乘法的 CUDA kernel，避免了通用 GEMM 的启动开销和 padding 浪费。对于 M=1 的矩阵乘法，这个 kernel 比 cuBLAS GEMM 快数倍。
• linear_f16_m1_splitk— M=1 场景下的 split-k 优化。将 K 维度切分成多个小块并行计算，充分利用 GPU 的并行度，最后做 reduction。在 decode 阶段的线性投影层中效果显著。
• fused add_layer_norm_cmix_mix_f16— T=1 专用的融合 kernel。将 layer norm、残差加法、channel mixing 的 mix 操作融合为单次 kernel launch，大幅减少 GPU kernel launch 开销和全局内存访问次数。
• LN1_TMIX_FUSE— 当B==1 and T==1时，layer norm 和 time mixing 的 mix6 操作进一步融合。这是 v3a 最深层的优化之一，将多个逐元素操作合并为一个 kernel。

以下代码展示了forward_from_x中 T==1 时的 fused kernel 分支（第 382 行附近）：

# forward_from_x 中的 T==1 融合分支ifT==1andself.ln1_tmix_fuse:# 当 B=1, T=1 时，layer_norm 和 tmix mix6 融合为单次 kernel launchx=fused_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 和 tmixx=layer_norm(x,self.ln1_w,self.ln1_b)x=tmix_forward(x,...)

而当 B=1, T=N（prefill）时，走的是完全不同的路径：

# cmix_sparse_one vs cmix_sparse_rows 的调用（第 583-585 行）ifself.cmix_mode=='CMIX_B1T1_SPARSE':# 单行专用：cmix_sparse_one 处理单个 tokenx=cmix_sparse_one(x,self.cmix_key,self.cmix_value,self.cmix_receptance,self.cmix_output)elifself.cmix_mode=='CMIX_BATCHED_SPARSE':# 多行通用：cmix_sparse_rows 处理 batchx=cmix_sparse_rows(x,self.cmix_key,self.cmix_value,self.cmix_receptance,self.cmix_output)else:# CMIX_DENSE: 退化为通用稠密矩阵乘法x=cmix_dense(x,self.cmix_key,self.cmix_value,self.cmix_receptance,self.cmix_output)

要点：v3a 的 kernel 特化深度极高。B=1, T=1 的路径不仅换了 kernel，还做了多层算子融合。这条路径的性能优势来自"单行计算"这个前提条件，任何破坏这个前提的设计（比如把 prefill 和 decode 混在同一个 batch）都会导致整条优化链失效。

四、混跑的代价：为什么不能把 prefill 和 decode 放在同一个 batch

理解了 v3a 的 kernel 特化机制，混跑的问题就一目了然了。

问题一：kernel 路径退化

假设一个 batch 中有 7 个 decode 样本（T=1）和 1 个正在 prefill 的样本（T=128）。此时rows = 8 * 128 = 1024（取最大 T），select_path会返回CMIX_DENSE。

结果：所有 8 个样本都失去 T=1 的专用优化。原本走linear_orig_row1_exact_f16_kernel的 decode 样本，现在被迫走通用 cuBLAS GEMM。性能损失可达数倍。

问题二：padding 浪费

不同长度的 prefill 拼在同一个 batch 里，需要 padding 到最长序列的长度。如果 batch 中有一个 1024 token 的 prefill 和七个 1 token 的 decode，有效计算量只有 1031，但实际计算量是 8192（8 x 1024），87% 的计算全是 padding。

问题三：kernel 切换开销

prefill 阶段使用的是大矩阵 GEMM（M=N, K=d_model, N=d_ff），decode 阶段使用的是 GEMV（M=1, K=d_model, N=d_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) │ B=1, T=N, 走 CMIX_DENSE │ │ └────────────┘ │ ▼ 状态 = READY │ ┌────────────┘ │ ▼ 调度器收集 READY 任务 │ ▼ update_batch() ──► 注入 Worker 槽位 │ ▼ engine.generate() ──► 无锁执行 │ B>1, T=1, 走 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_lock=prefill_lockifprefill_lockelseNullLock()self._status=Status.PREFILL self.prefill(prompt)# 构造时同步 prefillself.cpu()# prefill 完成后立即回 CPUdefprefill(self,prompt):self._status=Status.PREFILL prompt=self.tokenize(prompt)iflen(prompt)>=2:ifself.shift_state.device!="cuda":self.cuda()tokens=torch.tensor(prompt[:-1],dtype=torch.long,device="cpu")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_token=prompt[-1]self._status=Status.READY# prefill 完成，进入 READYdef__enter__(self):self.prepare()# cuda() 上 GPUreturnselfdef__exit__(self,*args):self.cpu()# 回 CPU

为什么这个设计是正确的

1. Prefill 不阻塞 decode：prefill_lock只保护 prefill 的forward()，decode 的generate()完全无锁。两者可以并行执行，GPU 不会被 prefill 独占。
2. 锁只防 prefill 互斥：多个任务同时 prefill 时会串行（通过prefill_lock），避免多个 prefill 同时抢占 GPU 导致资源争抢。但 decode 不受影响。
3. 状态机自然解耦：PREFILL 状态的任务不会被调度器 pick，只有 READY 状态的任务才能进入 decode batch。不需要额外的调度逻辑。
4. 串行 prefill 的开销可忽略：prefill 只占总 token 的 3%，即使 256 个任务串行 prefill，总耗时也只有约 5ms（2.9B）或 9ms（7.2B）。对整体延迟的影响微乎其微。

六、实测数据验证

以下是三组实测数据，验证了上述设计的有效性：

动态缩容过程

RWKV-Server 的 decode 调度器实现了动态 batch 缩容。当 batch 中的任务陆续完成时，batch size 从 256 逐步缩减：

缩容过程如此迅速，是因为 RWKV-Server 不需要像 vLLM 那样重新分配 KV Cache。RWKV 的状态是一个固定大小的向量（与序列长度无关），缩容只需要减少 batch 维度，不涉及内存重分配。

对比 Transformer：在 vLLM/SGLang 中，batch 缩容需要释放和重新分配 PagedAttention 的 KV Cache blocks，这个过程可能需要数百毫秒。RWKV 的固定大小状态向量让缩容变成了一个近乎零开销的操作。

七、与 Transformer 推理框架的对比

为什么 vLLM/SGLang 需要 Chunk Prefill

Transformer 的 self-attention 在 prefill 阶段的计算量是 O(n^2)。一个 4096 token 的 prompt 需要 4096 x 4096 = 16M 次注意力计算。而 decode 阶段每步只需要 1 x n 次计算。Prefill 的计算量可以占整个推理过程的 30-50%。

在这种情况下，prefill batching 的收益巨大：通过将多个 prefill 拼成一个大矩阵做 GEMM，GPU 利用率可以从 10% 提升到 80%+。即使 decode 阶段的 kernel 路径因此退化，整体吞吐仍然是提升的。

为什么 RWKV 不需要

RWKV 的线性注意力在 prefill 阶段是 O(n) 的逐 token 扫描，计算量远小于 Transformer。实测数据表明 prefill 只占总 token 的 3%。即使不做任何 batching 优化，prefill 的总耗时也只有几毫秒。

更关键的是，RWKV 的 decode 阶段依赖高度特化的 kernel（B=1, T=1 路径），这些 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 架构特性的深刻理解：

1. 与 v3a kernel 特化策略完全吻合：prefill 走 B=1, T=N 的稠密路径，decode 走 B>1, T=1 的专用稀疏路径，两条路径互不干扰。
2. 串行 prefill 的开销可忽略：3% 的 token 占比意味着即使不做任何优化，prefill 的总耗时也只有几毫秒。
3. 保护 decode 的 kernel 特化：不混跑意味着 decode 始终能走最优路径，不受 prefill 的干扰。
4. 实现简单、正确、高效：没有复杂的 chunking 逻辑，没有 prefill/decode 的调度博弈，没有 padding 浪费。

一句话总结：架构决定策略。RWKV 的线性注意力让"简单"成为最优解。不要把 Transformer 的工程经验盲目套用到 RWKV 上 – 理解你的架构，然后为它选择正确的路径。

项目地址：https://github.com/AUXStar/RWKV-Server

Albatross v3a CUDA kernels by BlinkDL