源码精读：ModelRunner —— GPU 前向引擎

谁该读这一篇？ 想搞清"Scheduler 把 batch 交给 GPU 后到底发生了什么"的工程师；准备实现新模型 / 调 attention backend 的开发者。 前置阅读： 03-scheduler.md、02-core-concepts/05-token-attention.md。 耗时： 50 分钟 学完能： 1. 在 2100+ 行的 model_runner.py 和拆出的 model_runner_components/、runner/ 中快速定位主线； 2. 理解 ModelRunner.forward 如何在 decode/prefill CUDA Graph 与 eager runner 之间分派； 3. 知道 ForwardBatch 这个 dataclass 在调度和执行之间起的桥梁作用； 4. 解释 CUDA Graph 在哪一步被捕获和回放； 5. 看到一个 forward 时延问题能在源码里定位到瓶颈位置。

1. ModelRunner 的位置

Scheduler.run_batch
   └── self.model_worker.forward_batch_generation(batch)
          └── ModelRunner.forward(forward_batch)        ← 本章主角
                ├── decode_cuda_graph_runner.execute(...)
                ├── prefill_cuda_graph_runner.execute(...)
                └── eager_runner.execute(...)

源码：model_runner.py（当前约 2200 行；加载、graph 配置和权重更新等职责已拆到 components）。

2. 关键数据结构：ForwardBatch

源码：forward_batch_info.py:394。

class ForwardBatch:
    forward_mode: ForwardMode                       # EXTEND / DECODE / IDLE
    batch_size: int
    input_ids: torch.Tensor                         # [total_tokens]
    req_pool_indices: torch.Tensor                  # [batch_size]
    seq_lens: torch.Tensor                          # [batch_size]
    out_cache_loc: torch.Tensor                     # 新算 KV 写到哪些 slot
    token_to_kv_pool: TokenToKVPool                 # KV cache 容器
    req_to_token_pool: ReqToTokenPool               # 逻辑 token → 物理 slot 映射
    attn_backend: AttentionBackend                  # 选哪个 backend
    sampling_info: SamplingBatchInfo                # 采样参数（一 batch 多种）
    # ... 多模态、LoRA、speculative 等可选字段

每个 Scheduler 决策一步产出一个 ForwardBatch。它是 Scheduler 和 ModelRunner 的唯一接口。

ForwardMode 不止三种。常见值包括：

class ForwardMode(IntEnum):
    EXTEND = ...          # 首次或分块 prefill
    DECODE = ...          # 每请求一个 token
    MIXED = ...           # chunked prefill 中混合两者
    IDLE = ...
    TARGET_VERIFY = ...   # speculative decoding verify
    DRAFT_EXTEND_V2 = ...
    PREBUILT = ...        # disaggregated decode
    SPLIT_PREFILL = ...
    DLLM_EXTEND = ...
    ...

注意 PREFILL 和 EXTEND 在源码里大量场景被统一处理（都是"一次喂多 token"）。

3. 初始化分层（ModelRunner.__init__ 311+）

现在不是在一个 __init__ 里顺序完成全部初始化。__init__ 先配置 device、 distributed、stream 等基础状态，再调用 initialize（model_runner.py:643）加载模型和准备 MoE/LoRA；Scheduler 随后按顺序调用 memory pool、attention backend 和 CUDA Graph：

def initialize(self):
    self.load_model()
    prepare_moe_topk(...)
    self.maybe_init_lora_manager()
    self.configure_kv_cache_dtype()

# Scheduler.init_model_worker / init_memory_pools / init_all_*
self.tp_worker.alloc_memory_pool()
self.tp_worker.init_attention_backends()
self.tp_worker.init_cuda_graphs()

这个顺序保证 graph 捕获时 KV pool、attention metadata 和可选 draft worker 都已就绪。 最耗时的仍是 load_model 和 CUDA Graph capture。

4. load_model（1113+）

源码：model_runner.py:1113。

简化：

def load_model(self):
    self.load_config = build_load_config(...)
    maybe_enable_ipc_weight_cache(load_config=self.load_config, ...)
    self.model = load_model_with_memory_saver(
        model_config=self.model_config,
        load_config=self.load_config,
        ...,
    )
    dist_barrier_after_load(...)

srt/models/<name>.py 里每个模型都遵循同样的接口：

• 继承 nn.Module
• 暴露 forward(input_ids, positions, forward_batch, ...) -> hidden_states
• 暴露 compute_logits(hidden_states) -> logits

要加新模型就写一个 <name>.py，在 model_loader/ 注册。

4.1 已驻留权重如何影响 KV 容量

普通加载先记录模型加载前的空闲显存，再根据加载后的剩余量计算 KV pool。IPC weight cache、startup preload 或 speculative worker 可能在这次采样前就让权重驻留 设备；若仍把它们当作“加载前就不存在的占用”，会高估可供 KV cache 使用的显存。

ModelRunner.preloaded_weights_bytes 从 loader 取得已驻留字节数， account_preloaded_weights() 将其加回 pre_model_load_memory，再交给 KV cache configurator。零值保持原路径，非法的负数或非整数会直接报错。这里修正的是容量 核算基线，不会重复分配或复制权重。源码： model_runner.py:847。

5. 主路径：forward（1563+）

def forward(self, forward_batch, ...):
    # 统一入口：观测、上下文和最终执行都在这里协调。
    return self._forward_raw(forward_batch, ...)

真正的分派在 _forward_raw：它先检查当前 mode 和静态 shape 是否允许 decode graph；否则准备 DP/attention batch，再尝试 prefill graph，最后由 EagerRunner.execute 处理 decode、extend、mixed、idle 及特殊模式。

6. Prefill 与 eager 路径

源码：model_runner.py:1563。

简化伪代码：

def _forward_raw(self, forward_batch, ...):
    if forward_batch.forward_mode.is_split_prefill():
        return self.forward_split_prefill(...)
    if self.prefill_cuda_graph_runner.can_run_graph(forward_batch):
        return self.prefill_cuda_graph_runner.execute(forward_batch, ...)
    return self.eager_runner.execute(forward_batch, ...)

要点：

• 普通 extend 走 EagerRunner；固定形状且满足约束的 extend/target verify 可以走 prefill_cuda_graph_runner（piecewise prefill graph）。

• 算最后一个 token 的 logits：prefill 阶段我们只关心"prompt 末尾下一步该出啥 token"，不算中间 token 的 logits。

• attention backend 已经在 init_forward_metadata 里建好 wrapper（FlashInfer 模式下 wrapper 对每个请求建索引）。

7. Decode CUDA Graph

can_run_graph = (
    forward_batch.forward_mode.is_cuda_graph()
    and self.decode_cuda_graph_runner is not None
    and self.decode_cuda_graph_runner.can_run_graph(forward_batch)
)
if can_run_graph:
    ret = self.decode_cuda_graph_runner.execute(forward_batch, ...)
else:
    ret = self.eager_runner.execute(forward_batch, ...)

Decode 通常是每请求一个 token，形状规整，因而最容易命中 graph；但 TARGET_VERIFY、IDLE 和部分特殊 mode 也被标记为可捕获，最终是否命中仍由 runner 的 shape/状态检查决定。execute 会准备静态 buffer 并 replay 已捕获的图。

adaptive speculative decoding 会在运行时切换 draft steps。共享 logits buffer 不能只按 启动时那一个 steps 值分配；max_decode_logits_rows（model_runner.py:824）会遍历 adaptive config 中所有候选档位，取最大的 capture_bs × tokens_per_req，避免切到更大档位时越界或启动崩溃。

8. Idle 与特殊模式

IDLE、TARGET_VERIFY、DRAFT_EXTEND_V2、PREBUILT、SPLIT_PREFILL 等状态由 ForwardMode 的谓词方法描述。idle batch 主要服务 DP/并行协同， 而不是一个独立的 forward_idle API；实际执行仍由 graph runner 或 EagerRunner 选择。

9. CUDA Graph 捕获

源码：runner/decode_cuda_graph_runner.py 与 model_runner_components/cuda_graph_setup.py。

启动时 init_cuda_graphs：

def init_cuda_graphs(self):
    # 给一组常见 batch size 都录一份 graph
    capture_batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128, 256]   # 可配
    self.decode_cuda_graph_runner.capture()

运行时 decode_cuda_graph_runner.can_run_graph 检查：

• batch size 是否在 pool 里。
• attention backend 元数据是否兼容（FlashInfer 要 page table 形状一致）。
• decode graph 只接受其 runner 支持的 mode/shape；prefill 由独立的 prefill_cuda_graph_runner 判断，普通变长 extend 通常回退 eager。

不匹配就 fallback 到普通 forward。

Capture 阶段是启动慢的主因（数十秒）。

10. Overlap 模式下的 forward

enable_overlap=True 时 Scheduler 在 GPU 还在跑时就准备下一批；ModelRunner 这边：

def forward_batch_generation(self, batch, on_publish=None):
    with self.forward_stream_ctx:
        # 在 forward_stream 上跑 model
        output = self.forward(batch)
    if on_publish:
        on_publish(...)                                # spec_v2 模式中途回调
    return output

要点：

• 用 separate forward_stream 避免和 schedule_stream 抢；
• on_publish 是 spec_v2（投机解码 v2）在 verify 后 / draft 前的回调点。

11. 一次 forward 的 GPU 时间组成

粗略（H100、Llama-3-8B、batch=32 decode）：

CUDA Graph 省的是 kernel launch 开销，约 1 ms / forward。

12. 加新模型的流程

1. 写 srt/models/my_model.py，模仿 llama.py 结构。
2. 在 model_loader/__init__.py 注册映射。
3. 如果 attention 形状特殊（如 MLA），在 attention backend 加分支。
4. 单元测试 + 跑一个 prompt 验证。

PR 历史里有大量"add support for X model"，照抄即可。

13. 权重热更新（WeightUpdater）

ModelRunner.init_weight_updater 注入模型 getter、字段更新回调和 graph 重捕获回调； 四条实现已集中到 weight_updater.py:55，不是 ModelRunner 自身的方法：

def update_weights_from_disk(self, recv_req):
    ...

def update_weights_from_distributed(self, recv_req):
    ...

def update_weights_from_tensor(self, recv_req):
    ...

def update_weights_from_ipc(self, recv_req):
    ...

四条更新路径，对应 RL 训练 / online 微调场景。

• from_disk：从模型文件重新加载（最慢）。
• from_distributed：从训练侧 NCCL 广播。
• from_tensor：直接传 tensor。
• from_ipc：CUDA IPC handle（最快）。

详见 08-production-deployment/09-versioning.md 的"模型权重切换"小节。

14. 关键源码索引

15. 小结

• ModelRunner 是 Scheduler 和 GPU 之间的薄包装层，主入口 forward 按 mode 分派。
• forward 是统一入口；_forward_raw 按 mode 和 shape 在 decode graph、prefill graph、split prefill 与 EagerRunner 之间选择。
• ForwardBatch 是接口，包含 input_ids、KV pool 引用、attention backend、采样信息等。
• CUDA Graph 在启动时按多个 batch size 预捕获，运行时 replay。
• 加新模型只动 srt/models/ 和 model_loader/ 注册。

16. 自检

1. prefill 何时能用 CUDA Graph？

2. CUDA Graph 在哪个 batch size 下命中？哪种情况会 fallback？

3. ForwardBatch.forward_mode 有几种值？分别表示什么？

4. 加新模型 PR 需要改哪些文件？

5. 看到 forward 时延突增，应该先看哪几个数？

17. 下一步

• 05-radix-tree.md — KV cache 索引细节。
• 06-attention-backends.md — Attention kernel 选择。
• 04-optimizations/04-cuda-graph.md — CUDA Graph 进阶。
• 源码：model_runner.py、forward_batch_info.py。

上游源码：sglang/python/sglang/srt/model_executor/。