MRV2 学习笔记

MRV2 在 GPU 执行 step N 时，CPU 已经开始调度并准备 step N+1。问题在于，step N 的结果可能使旧 request 结束，而新的 request 也可能在这段时间到达。

MRV2 的解决方式不是“提前准确知道 N+1 的 batch”，而是：

有界的乐观调度 + GPU 内部传递真实依赖 + 结果按顺序提交；猜错的工作丢弃或回滚。

CPU: schedule N ── schedule N+1（使用 placeholder）── process output N ── schedule N+2
GPU:             execute N ── sample N ──────────── execute N+1
                                │
                                └─ 把真实 token 写入 GPU request state

1. N+1 不需要 CPU 知道 N 的 token 值

Scheduler 只需要知道“这个 request 大概率会产生一个 token”，所以在 CPU 状态中加入 num_output_placeholders，并乐观增加：

num_computed_tokens
num_in_flight_tokens

代码入口：

真正的 token 值不必回传 CPU。step N sampling 完成后，GPU 把它写入 persistent request state 的 last_sampled_tokens；step N+1 的 input-prep kernel 再直接读取这个值：

这些 kernel 在同一 CUDA stream 上按顺序执行。因此，即使 CPU 正在并行准备 N+1，GPU 真正执行 N+1 时读取的仍然是 N 产生的真实 token。

2. 如果 N 的 token 使 request 结束

例如 N 输出 EOS 或 stop token，但 N+1 已经把这个 request 排进 batch：

N+1 通常不会被取消，因为相关工作可能已经 enqueue 到 GPU。
N+1 可能多计算一次“EOS 后面的 token”。
CPU 处理 N 的输出时，把 request 标记为 finished。
N+1 的结果返回时，如果 request 已不存在或已经 finished，直接忽略结果。

对应逻辑位于 Scheduler.update_from_output：

if request is None or request.is_finished():
    continue

所以这种情况损失的是少量多余计算，不会把多算的 token 返回给用户。

对于 max_tokens 这种可以从计数提前判断的结束条件，Scheduler 会结合 placeholder 避免再调度一个 step。EOS 和自定义 stop token 是数据依赖的，无法提前预测，因此允许一次 speculative overrun。

3. 新 request 如何进入 batch

分为三种情况：

新 request 在 schedule N+1 前已经到达，并且有容量：可以直接进入 N+1。
新 request 在 schedule N+1 后才到达：最早进入 N+2。
只有等 N 的结果让旧 request 结束后才腾出容量：新 request 也要等到 N+2 才能进入。

也就是说，MRV2 接受“新 request 最多晚一个 pipeline window 被 admit”，以换取 CPU scheduling/input prep 与 GPU execution 的重叠。

finished request ID 会被放入下一份 SchedulerOutput。ModelRunner 收到后先删除旧 request 的永久 slot，再添加新 request：

4. 乐观假设影响 token 数时如何修正

Speculative decoding 可能拒绝一部分 draft token。step N 的结果返回后，Scheduler 会回滚：

num_computed_tokens
num_output_placeholders

GPU 侧同时维护真实的 accepted/rejected 状态，所以 N+1 真正执行时使用的 GPU 状态仍然正确。CPU 的 num_computed_tokens_np 明确只是 optimistic upper bound，而不是最终事实。

5. Structured output：提前执行 N+1 forward，延迟 N+1 sampling

图：MRV2 对 structured output 的 deferred sampling。CPU/GPU lane 中的实线块表示实际工作，底部 lane 表示 grammar 状态而非独立执行线程；红色暂停标记表示 CPU 尚未调用 sample_tokens()。

Structured output 维护一个 CPU 侧 grammar/FSM 状态。假设 step N 采样出 token t_N，那么 step N+1 的合法 token mask 是：

G_N --accept(t_N)--> G_{N+1} --生成 bitmask--> mask_{N+1}

在 t_N 从 GPU 返回 CPU、并由 grammar.accept_tokens(t_N) 推进 FSM 之前，CPU 无法正确生成 mask_{N+1}。这里不能用普通 placeholder 代替，因为不同的 t_N 会把 grammar 推进到不同状态，并产生不同的合法 token 集合。

不过，grammar mask 不参与 Transformer forward，只在 forward 之后限制 logits 并进行 sampling。因此 N+1 被拆成两部分：

可以提前执行：
  schedule N+1
    → GPU input prep（直接读取 GPU 上的真实 t_N）
    → model forward N+1
    → 保存 N+1 hidden states

必须等待：
  step N 的 token D2H 完成
    → Scheduler.update_from_output(N)
    → grammar.accept_tokens(t_N)，得到 G_{N+1}
    → 计算 mask_{N+1}
    → sample_tokens(N+1, mask_{N+1})

具体控制流程如下。

AsyncScheduler._update_after_schedule 发现某个 structured-output request 仍有上一轮尚未提交的 num_output_placeholders，于是设置：
```
scheduler_output.pending_structured_output_tokens = True
```
用来阻止立刻mask N+1。
EngineCore.step_with_batch_queue 仍然立即调用：
```
execute_model(scheduler_output, non_block=True)
```
因此 N+1 的 input prep 和 model forward 可以与 N 的输出回传/处理重叠。
但 EngineCore 暂时不调用 N+1 的 sample_tokens，而是把 N+1 的 scheduler_output 保存在 deferred_scheduler_output 中。后续grammar mask的时候需要知道在N+1 batch中每个request是如何被调度的信息，结合accept结果得到mask。
EngineCore 取回并提交 N 的结果。Scheduler.update_from_output 将 t_N 追加到 request，并调用：
```
grammar.accept_tokens(req_id, [t_N])
```
此时 grammar 才从 G_N 推进到 G_{N+1}。
EngineCore 使用更新后的 grammar 生成 mask_{N+1}，然后恢复之前延迟的：
```
sample_tokens(mask_{N+1}, non_block=True)
```

MRV1 vs MRV2

理解了MRV2的执行也就不难理解MRV2相对MRV1的优势在哪里了。

MRV1是要对request table进行大量操作的，一方面这占用了CPU资源且在critical path上，另一方面我们很难安排optimistic guess和request table的调整。