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 Decode
每步 token 数 整个 prompt 长度（通常 20~2000 tok） 1 token / step
计算复杂度 O(n) 线性扫描，稠密 GEMM O(1) 状态更新，GEMV / 稀疏 kernel
GPU 利用率 中等（矩阵维度较大） 极低（单行计算，依赖专用 kernel）
总 token 占比 ~3%（实测 20,140 / 548,180） ~97%（实测 528,040 / 548,180）
Kernel 特化程度 通用稠密路径（CMIX_DENSE） 高度特化（SPARSE / fused / split-k）
能否与 decode 混跑 会退化整个 batch 的 kernel 路径 需要 B=1, T=1 才能走最优路径

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

实测吞吐数据

模型 + GPU 256 任务入队耗时 Prefill 吞吐 Decode 峰值吞吐 Prefill 占比
2.9B + RTX 4090 4.97s ~4,052 tok/s 11,081 tok/s 3.5%
7.2B + RTX 4090 8.64s ~2,331 tok/s 6,373 tok/s 3.2%
2.9B + 5070 Ti Laptop – – 5,169 tok/s –

注意 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 调度系统的入口：

def select_path(B, T):
    """
    根据 batch size (B) 和 sequence length (T) 选择最优的 kernel 执行路径。
    返回值决定了后续所有 cmix / tmix 操作使用哪套 kernel。
    """
    rows = B * T

    if rows == 1:
        # B=1, T=1: 单样本单步 decode -- 最优路径
        return 'CMIX_B1T1_SPARSE'
    elif B == 1 and T > 1:
        # B=1, T=N: 单样本 prefill -- 稠密 GEMM 路径
        return 'CMIX_DENSE'
    elif B > 1 and T == 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 融合分支
if T == 1 and self.ln1_tmix_fuse:
    # 当 B=1, T=1 时，layer_norm 和 tmix mix6 融合为单次 kernel launch
    x = 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 和 tmix
    x = 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 行）
if self.cmix_mode == 'CMIX_B1T1_SPARSE':
    # 单行专用：cmix_sparse_one 处理单个 token
    x = cmix_sparse_one(x, self.cmix_key, self.cmix_value,
                          self.cmix_receptance, self.cmix_output)
elif self.cmix_mode == 'CMIX_BATCHED_SPARSE':
    # 多行通用：cmix_sparse_rows 处理 batch
    x = 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 方法

class Task:
    def __init__(self, prompt, model_loader, batch_sampler, ...):
        # ... 初始化 state ...
        self.prefill_lock = prefill_lock if prefill_lock else NullLock()
        self._status = Status.PREFILL
        self.prefill(prompt)   # 构造时同步 prefill
        self.cpu()             # prefill 完成后立即回 CPU

    def prefill(self, prompt):
        self._status = Status.PREFILL
        prompt = self.tokenize(prompt)
        if len(prompt) >= 2:
            if self.shift_state.device != "cuda":
                self.cuda()
            tokens = torch.tensor(prompt[:-1], dtype=torch.long, device="cpu")
            self.prefill_lock.acquire()   # 只锁 prefill 的 forward
            self.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 完成，进入 READY

    def __enter__(self):
        self.prepare()   # cuda() 上 GPU
        return self

    def __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）。对整体延迟的影响微乎其微。

六、实测数据验证

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

配置 256 任务入队耗时 Prefill 吞吐 Decode 峰值吞吐 Prefill 占比 Prefill 总耗时
2.9B + RTX 4090 4.97s ~4,052 tok/s 11,081 tok/s 3.5% ~5.0ms
7.2B + RTX 4090 8.64s ~2,331 tok/s 6,373 tok/s 3.2% ~8.6ms
2.9B + 5070 Ti Laptop – – 5,169 tok/s – –

动态缩容过程

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

Batch Size 256 128 64 32 16 8 4 2 1
缩容耗时 每次缩容 < 15ms（含 kernel 重编译） –

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

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

七、与 Transformer 推理框架的对比

维度 RWKV-Server vLLM / SGLang
注意力机制 O(n) 线性注意力（RNN-like 状态） O(n^2) self-attention
Prefill 策略 单线程串行（B=1, T=N） Chunk Prefill + batching
Decode 策略 Dynamic batching（B>1, T=1） Continuous batching
Prefill/Decode 混跑 不混跑（状态机解耦） 混跑（同一 batch）
混跑原因 混跑会退化 decode kernel Prefill 是瓶颈，必须 batching
Prefill 占比 ~3% ~30-50%（取决于序列长度）
状态管理 固定大小向量，无需 KV Cache PagedAttention KV Cache
Batch 缩容 < 15ms（无内存重分配） ~100-500ms（KV Cache 回收）
Kernel 特化 极深（B=1,T=1 有专用融合 kernel） 较浅（FlashAttention 通用）

为什么 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

# 语言模型 # 人工智能 # 深度学习
Logo
AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

cover

G-Star 精选开源项目推荐｜第二十期

avatar AtomGit开源社区
cover

AtomGit 即将亮相 WAIC 2026｜ 来赴一场属于开发者的 AI Coding Party

avatar AtomGit开源社区
cover

开源鸿蒙大学生创新大赛三大赛道全面开赛，176 万奖金池已就位

avatar AtomGit开源社区