TpWorker、ModelRunner 与 draft worker 的边界#

这章解决什么问题#

运行时架构前面已经讲了 scheduler 的模式树、typed IPC、data parallel controller 和配置编译层,但还有一个非常关键的执行边界没有单独讲:Scheduler 自己并不直接 forward 模型,它是通过 TpModelWorkerModelRunner,以及在 speculative 场景下的 draft_worker 来完成真正的执行。

这一章的目标,就是把这三个层的边界讲清楚。

为什么这层值得单独成章#

很多读者在第一次读 scheduler.py 时会看到:

  • self.tp_worker
  • self.draft_worker
  • self.model_worker
  • self.tp_worker.model_runner

如果不把这层分工讲明白,就很容易把它们误读成重复包装。事实上它们承担的是不同粒度的职责:

  • Scheduler 负责调度和状态推进
  • TpModelWorker 负责 worker 级执行接口与资源拥有权
  • ModelRunner 负责更接近 forward/sample/graph/backend 的执行细节
  • draft_worker 只在 speculative 路径下引入另一条执行人格

一张图:scheduler 并不是直接把 batch 丢给模型,而是经过 worker / runner 边界#

这张图解决的理解障碍是:很多人默认 Scheduler -> ModelRunner 是直接关系,但中间实际上隔了一层 worker 边界。

flowchart LR
    Sch["Scheduler"] --> TP["TpModelWorker"]
    TP --> MR["ModelRunner"]
    Sch --> Draft["draft_worker (spec only)"]
    Draft --> DMR["draft ModelRunner"]

这张图比纯文字多解释的一点是:执行层不只是一个 ModelRunner,而是存在 worker 边界和 speculative 分支边界。

TpModelWorker 真正在包什么#

tp_worker.py 看,TpModelWorker 至少负责:

  • 拥有 ModelRunner
  • 提供 get_memory_pool()
  • 提供 get_tokens_per_layer_info()
  • 代理权重更新 / LoRA 更新 / load/unload
  • 组织 ForwardBatch.init_new(...)
  • 在某些路径下直接调 model_runner.forward(...)sample(...)

这说明它不是简单转发器,而更像“执行资源拥有者 + runner 的运行时壳”。

为什么 TpModelWorker 不等于 ModelRunner#

因为 TpModelWorker 还承担:

  • 分布式 rank / GPU / NCCL 相关初始化边界
  • memory pool 暴露边界
  • draft worker / multi-layer eagle 相关 runner 组织

ModelRunner 则更偏:

  • 真正的 forward 路径
  • attention backend / graph runner / kv cache 执行细节
  • sample / logprobs / embeddings 输出

这就是典型的“worker 边界”和“runner 边界”分工。

model_worker 为什么有时是 tp_worker,有时是 draft_worker#

scheduler.py 里在 init_model_worker() 之后会根据 speculative 算法决定:

  • 没有 speculative:self.model_worker = self.tp_worker
  • 有 speculative:self.model_worker = self.draft_worker

这说明 scheduler 真正依赖的不是“某个固定类”,而是“当前模式下负责本轮执行人格的 worker”。这是一种很重要的架构抽象,因为它让 speculative 路径可以在不推翻 scheduler 结构的情况下插入另一条执行分支。

maybe_init_draft_worker() 揭示了什么#

这段初始化逻辑说明:

  • draft worker 不是总存在
  • 它由 speculative algorithm 决定是否需要
  • 它也会带自己的 worker kwargs、target_worker=self.tp_worker

这说明 speculative 并不是在 sampler 里突然做点小技巧,而是会在运行时架构上显式长出一条额外执行分支。

ModelRunner 这一层最适合怎么理解#

前面的 execution 章节已经讲了很多 ModelRunner 内部细节。放回运行时架构里,更好的理解方式是:

  • 它是 worker 层向执行细节下钻的主要边界
  • 它承接真正的模型、graph runner、attention backend、KV cache、sample
  • 它并不负责 manager 间通信,也不直接负责请求排队

这让它在整本书里的位置更清楚:它属于“执行核心”,而不是“控制面核心”。

为什么这章和执行模型章节互补#

执行模型章节会告诉你 ModelRunner.forward(...)sample(...)SamplingBatchInfo 怎样运作;这一章则告诉你:

  • 这些执行细节是被哪一层对象拥有和调起的
  • speculative 路径为什么会多出一条 draft worker 分支

换句话说:

  • execution 章节讲“怎么跑”
  • 这一章讲“是谁在跑”

get_memory_pool()get_tokens_per_layer_info() 为什么值得特别提#

这两个接口很像小细节,但它们特别能说明 worker 层的职责:

  • scheduler 不直接向 ModelRunner 裸取所有细节
  • 它通过 worker 暴露更稳定的资源/能力接口

这说明 worker 层在架构上也承担了一部分“稳定 façade”作用,而不是把 scheduler 直接绑死在 runner 内部实现上。

这一层最容易出现的误判#

1. 把 TpModelWorker 当成完全多余的包装#

它实际上承担了分布式、资源和 speculative 组织边界。

2. 把 draft_worker 理解成 sampler 内部细节#

它已经长成了运行时架构上的独立执行分支。

3. 以为 scheduler 直接调用模型#

实际调用链更接近:Scheduler -> worker -> runner -> model/backend

如果你要排查“执行到底卡在哪一层”,先怎么走#

建议按这个顺序:

  1. 先确认当前 batch 走的是 tp_worker 还是 draft_worker
  2. 再确认问题更像发生在 worker 边界,还是 runner 内部
  3. 如果进入 runner 内部,再回 execution model 章节继续往下钻

小结#

这一章真正要补齐的,是运行时架构里经常被隐含掉的一层执行边界:

  • Scheduler 不是直接执行模型
  • TpModelWorkerModelRunner 之间存在稳定边界
  • speculative 场景下又会长出 draft_worker 这条额外执行分支

到这里,运行时架构对“执行是谁在跑”这件事也有了更完整的解释。