ForwardBatch 与 ModelRunner#

第五章讲的是请求怎样成 batch，第六章开始要回答另一件事：这个 batch 最后怎样真正落到执行层。站在调度器视角看，手里还是 Req 和 ScheduleBatch；站在执行层看，需要的却已经是更贴近前向计算的对象。这条边界正是 ForwardBatch 和 ModelRunner 共同承担的。

这一节解决什么问题#

这一节主要回答三件事：

  1. ScheduleBatch 为什么还不等于执行输入；
  2. ForwardBatch 到底把哪些调度信息压成了前向输入；
  3. ModelRunner 为什么被设计成执行壳，而不是直接把模型暴露给 scheduler。

一张图先看执行前的最后两层#

flowchart LR
    A["Req / ScheduleBatch"] --> B["ForwardBatch"]
    B --> C["ModelRunner"]
    C --> D["模型前向 / logits / hidden states"]

这张图最重要的一点是：执行层接手 batch 之前，还有一层专门把调度对象翻译成前向对象的边界。

ForwardBatch 解决的是“执行层到底要吃什么”#

ForwardBatch 从字段定义就能看出，它已经明显偏执行层：

  • input_ids
  • req_pool_indices
  • seq_lens
  • out_cache_loc
  • positions
  • sampling_info
  • req_to_token_pool
  • token_to_kv_pool

这说明 ForwardBatch 不再关心 waiting queue 里还有谁、哪条请求为什么先入队，它只关心这一轮前向真正需要的输入、位置、缓存映射和采样附带信息。

所以 ForwardBatch 的职责不是“再包一层数据”，而是把调度对象压成执行层可直接消费的形式。

如果把这层桥接再压成更短的一段代码，可以直接看到“调度对象怎样被翻译成执行对象”：

return cls(
    reqs=reqs,
    req_to_token_pool=req_to_token_pool,
    token_to_kv_pool_allocator=token_to_kv_pool_allocator,
    tree_cache=tree_cache,
    model_config=model_config,
    enable_overlap=enable_overlap,
)

这段代码值得看的不是参数列表，而是这些参数的来源：前半部分仍然是调度器关心的 request 和 cache，后半部分已经开始偏向执行器关心的运行配置。

为什么 ScheduleBatch 不能直接交给 ModelRunner#

如果没有 ForwardBatch，ModelRunner 就必须直接理解：

  • request queue 语义；
  • prefix / batch shaping 语义；
  • execution 需要的张量输入；
  • cache 位置和采样附带信息。

这样执行层就会被调度层污染。当前设计保留 ForwardBatch，意味着：

  • scheduler 继续掌握“这一轮为什么这样排”；
  • ModelRunner 只接更靠近前向计算的对象。

从全书结构看，这也是为什么第五章和第六章必须拆开的原因。前者讲“为什么这一轮是它们”，后者讲“这一轮怎样真正算出来”。

ModelRunner 是执行壳，不是调度器的附属函数#

ModelRunner 的初始化字段已经说明，它站在很明确的执行层位置：

  • model config
  • tp / pp / dp 等并行状态
  • memory pool config
  • req_to_token_pool
  • token_to_kv_pool_allocator
  • speculative algorithm

这些字段说明两件事：

  1. 它显然不是一个“简单调用模型”的小包装；
  2. 它也不是调度器的延长线，而是一个同时理解模型、并行拓扑和 KV 映射的执行壳。

所以 ModelRunner 负责的不是“选谁先跑”，而是“这一轮已经选好了，该怎样在当前拓扑下正确执行”。

它的初始化字段本身也很能说明这件事：

self.tp_rank = tp_rank
self.tp_size = tp_size
self.model_config = model_config
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.spec_algorithm = SpeculativeAlgorithm.from_string(
    server_args.speculative_algorithm
)

这些字段已经明显在执行层，而不再是调度层的排队信息。

执行层真正接手的边界在哪里#

如果把这一节的主线再压缩一下，可以得到一个很直接的判断：

  • ScheduleBatch 里仍然带着很多调度语义；
  • ForwardBatch 开始主要携带前向所需输入；
  • ModelRunner 负责拿这些输入和当前模型 / 并行环境真正执行前向。

这条边界如果不先钉住，后面读 logits、采样和输出处理时就会很容易混淆：有些逻辑为什么还在调度层，有些却已经在执行层。

调试执行入口时先看哪里#

如果你看到的现象是：

  • batch 已经成形，但前向输入不对；
  • cache 映射看起来正常，但执行侧拿到的序列长度不对；
  • 某些并行模式下只有执行侧才出现异常；

更稳的顺序通常是：

  1. 先看 ScheduleBatch 里这一轮到底选了谁；
  2. 再看 ForwardBatch 是否正确携带了 input_ids / seq_lens / out_cache_loc；
  3. 最后再看 ModelRunner 在当前并行状态下怎样消费这些输入。

这样做的好处是，你能把“调度选错了”和“执行层解释错了”分开看。

Prefill 与 Decode 的执行本质区别#

在真正理解 ForwardBatch 怎样驱动前向计算之前，有一件事必须先搞清楚：prefill 和 decode 是两种计算量和并行度都完全不同的执行模式。这个区别不是实现细节，而是理解为什么 KV cache 存在、为什么 batch 组合方式如此重要的根本原因。

Prefill：一次处理整个 prompt#

Prefill 阶段把 prompt 里所有 token 一起喂给模型。在每一个 Transformer 层的注意力计算里：

  • Q、K、V 的形状都是 [batch, S, d_head]，其中 S 是 prompt 的序列长度
  • 对每个注意力头，先算 Q×K^T，得到 [batch, S, S] 的分数矩阵
  • 再做 softmax，再乘 V，得到 [batch, S, d_head] 的输出
  • 把这次算出的 K 和 V 全部写入 KV cache

计算量是 O(S²)——对 S 个 token 中的每一个，都要和其他 S 个 token 做一次点积。但关键是：所有 token 可以在一次大矩阵乘法里同时算完，GPU 的并行度被充分利用。

Decode：每步只处理一个新 token#

Decode 阶段每一步只生成一个新 token。在注意力层：

  • Q 的形状是 [batch, 1, d_head]——只有一个新 token 的 query
  • K 和 V 来自 KV cache，形状是 [batch, S+t, d_head]，其中 t 是已经生成的 token 数
  • Q×K^T 的结果是 [batch, 1, S+t]——只有一行分数
  • softmax 之后乘 V，得到 [batch, 1, d_head]
  • 把这个新 token 的 K 和 V 追加到 KV cache 里

计算量是 O(S+t)——比 prefill 低得多。但问题在于：这一步的输出依赖上一步的结果，无法跨 token 并行。每一步都是独立的小矩阵乘法，GPU 利用率远低于 prefill。

两种模式的张量形状对比#

Prefill（S 个 prompt token）：
  Q  [batch, S,   d] × K^T [batch, d, S  ] = scores [batch, S,   S  ] → O(S²)
  → 把 K[batch, S, d] 和 V[batch, S, d] 写入 KV cache

Decode（生成第 t 个新 token）：
  Q  [batch, 1,   d] × K^T [batch, d, S+t] = scores [batch, 1,   S+t] → O(S+t)
  → 把 K[batch, 1, d] 和 V[batch, 1, d] 追加到 KV cache

这个形状差异直接决定了 ForwardBatch 里 out_cache_loc 字段的语义：prefill 时是 S 个 slot 的索引，decode 时只是 1 个新 slot 的索引。

ForwardBatch.forward_mode 是执行分支的开关#

ForwardBatch 有一个 forward_mode 字段，它决定执行路径走哪条分支：

  • PREFILL：第一次处理 prompt，Q/K/V 都有完整的序列长度
  • DECODE：扩展一个新 token，Q 的序列长度为 1，K/V 来自 cache
  • EXTEND：chunked prefill 的继续——prompt 被拆成多段，当前段是 prompt 的一部分
  • MIXED：用于 speculative decoding，同一个 batch 里既有 draft token 的 verify，也有正常的 decode

ModelRunner.forward() 根据这个字段走不同的代码路径：

if batch.forward_mode == ForwardMode.DECODE:
    # Q 的 seq_len=1，K/V 全部来自 KV cache
    output = self.model.forward_decode(batch)
elif batch.forward_mode == ForwardMode.PREFILL:
    # Q/K/V 都有完整的 seq_len
    output = self.model.forward_prefill(batch)

这个分支不是可选的优化，而是必须的正确性保证：prefill 和 decode 在注意力层的计算逻辑根本不同，不加区分会算出错误的结果。

为什么 forward_mode 直接影响 KV cache 的写入方式#

  • PREFILL 时：S 个 token 的 K/V 都要在这一步写入，out_cache_loc 是 S 个 slot 的地址列表
  • DECODE 时：只有 1 个新 token 的 K/V 要追加，out_cache_loc 是 1 个新 slot 的地址

这就是为什么 ForwardBatch 里会有 out_cache_loc 这个字段，而不是直接让 attention 层自己决定写在哪里——slot 分配在调度阶段已经由 token_to_kv_pool 完成，执行层只需按地址写入。

为什么 decode 比 prefill 慢得多（按 token 算）#

这里有一个让很多人直觉上觉得反常的不对称：

  • prefill 1000 个 token 需要 T 秒，平均每 token 是 0.001T 秒
  • decode 生成 1 个 token 需要约 0.3T 秒，是 prefill 平均每 token 耗时的 300 倍

原因在于计算密度的差异：

  • prefill 是一次大矩阵乘法（GEMM），形状够大，GPU 的矩阵乘法单元被充分喂满
  • decode 每步是一次极小的矩阵乘法（Q 只有 1 行），矩阵太小，GPU 的并行计算单元大量闲置，主要时间花在内存带宽上（从 KV cache 里读历史 K 和 V）

此外，随着生成长度增加，每一步 decode 需要从 KV cache 读取的数据量也在增长——序列越长，每一步越慢。

这个不对称如何影响调度策略#

理解了这个不对称，第五章里的调度决策就有了直接的物理解释：

  • decode 阶段要尽量把更多请求打包进同一个 batch：同时跑 32 个请求的一步 decode，和跑 1 个请求的一步 decode，总的 GPU 时间差不多，但吞吐量是 32 倍
  • 但 batch size 受 KV cache 容量限制：每个请求的历史 K/V 都要放在 GPU 显存里，显存满了就没法再加请求
  • decode 吞吐随序列长度退化：序列越长，每步要读的 KV cache 越大，单步延迟越高，整体吞吐下降

这就是为什么 SGLang 的调度器要精确追踪每个请求已经占用了多少 KV cache slot，而不是简单地最大化 batch size。

小结#

这一节真正要建立的是一个判断：

  • ScheduleBatch 还是调度对象；
  • ForwardBatch 是执行前的翻译层；
  • ModelRunner 是真正的执行壳。

理解了这个边界，第六章后面再讲采样和输出收口时，就不会把所有逻辑都误看成前向函数内部的细节。