读 tp_worker.py：scheduler 与执行壳之间的 worker 外观#

这章解决什么问题#

前面的运行时架构已经在：

里解释了 worker 壳的职责边界，而代码导读又已经分别补了：

  • forward_batch_info.py
  • model_runner.py
  • model_loader

但如果你真的顺着源码往执行层里走，仍然会遇到一个具体问题：

  • scheduler 到底是怎样通过一个更稳定的 worker 外观去驱动 ModelRunner，而不是直接把所有控制和执行细节都压到 model_runner.py 上？

这条链最关键的文件其实是：

  • python/sglang/srt/managers/tp_worker.py

这章的目标，就是把这棵 worker façade 树正式收成一章。

为什么这棵树值得单独成章#

因为 tp_worker.py 并不是简单的“把 ModelRunner 包一层”。它至少承担了三类职责：

  1. 把配置和 rank 身份翻译成 worker 实例
  2. 把调度层的 ModelWorkerBatch 编译成 ForwardBatch
  3. 把一系列控制动作统一转发给 ModelRunner

也就是说，它不仅是执行壳的外观，也是调度层与执行壳之间的正式边界。

优秀技术书在这里应该明确告诉读者：

  • 你不是在读一个多余封装
  • 而是在读 scheduler 真正依赖的 worker 契约

一张图：TpModelWorker 在执行主线里的真实位置#

这张图解决的理解障碍是：很多读者会默认 Scheduler -> ModelRunner 是直接关系，但源码里其实明确存在一个 worker façade。

flowchart LR
    Sch["Scheduler"] --> Batch["ModelWorkerBatch"]
    Batch --> TP["TpModelWorker"]
    TP --> Fwd["ForwardBatch.init_new(...)"]
    Fwd --> MR["ModelRunner.forward() / sample()"]
    TP --> Ctrl["weights / LoRA / memory / worker info"]

图比纯文字多解释的一点是：TpModelWorker 并不是“执行以后才有的 wrapper”，它从 batch 入口开始就在参与执行链。

第一层：为什么应该先读 BaseTpWorker#

更稳的顺序，不是直接冲进 TpModelWorker.__init__()，而是先看：

  • BaseTpWorker

因为它已经很清楚地暴露了这棵树的稳定契约：

  • forward_batch_generation(...)
  • model_runner
  • get_tokens_per_layer_info()
  • get_memory_pool()
  • 各种 update_weights_from_*
  • load_lora_adapter(...)

这说明这棵树的核心不是某个具体模型，而是：

  • worker 对上层暴露哪些稳定能力

从系统阅读角度看，这非常重要，因为它让你在进入实现前，先知道 scheduler 和控制面到底期待它提供什么。

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

它们看起来像小 getter，但其实暴露了一个关键边界：

  • scheduler 或更上层有时并不需要整个 ModelRunner
  • 它只需要问 worker：“你有多少 token 容量、你的 memory pool 长什么样？”

这说明 TpModelWorker 不只是 forward façade，也是一层资源能力查询面。

这和前面调度与内存章节天然回扣，因为很多 admission / cache / profile 决策都建立在这些能力事实上。

TpModelWorker.__init__() 真正在固定什么#

这段构造函数最值得抓住的，不是每个字段，而是几个关键步骤：

  1. 保存 rank / GPU / parallelism 身份
  2. _init_model_config()
  3. _init_model_runner()
  4. 必要时初始化 multi-layer eagle runners
  5. 初始化 tokenizer / processor
  6. 同步随机种子
  7. 确定 overlap、spec、memory pool 等运行人格

也就是说，worker 自己也是一个小型人格编译器。

它并不只是拿现成 ModelRunner 来用，而是先把：

  • 这个 rank 应该成为怎样的 worker

这件事钉死。

_init_model_runner() 为什么是执行壳的真正挂载点#

这里会真正构造：

  • ModelRunner(...)

并把：

  • model_config
  • mem_fraction_static
  • rank 身份
  • req_to_token_pool
  • token_to_kv_pool_allocator
  • is_draft_worker

这些全部接进去。

这说明 TpModelWorker 并不是在运行时才找 ModelRunner，而是从初始化阶段就把它作为内核执行壳牢牢挂好。

也正因为这样，前面的 model_runner.py 章节更像“执行壳内部怎么跑”，而这章更像“执行壳怎样被上层真正接入系统”。

forward_batch_generation(...) 为什么是最值得读的主入口#

如果只读一个函数，那一定是：

  • forward_batch_generation(...)

因为它把 worker 外观的价值写得非常直接：

  1. 如果给的是 ModelWorkerBatch，先 ForwardBatch.init_new(...)
  2. 再调用 self.model_runner.forward(...)
  3. 如果当前在最后一个 PP rank，就决定是否 sample 或只算 logprobs
  4. 最后把结果包装成 GenerationBatchResult

这说明 worker façade 的真正职责不是做很多计算，而是：

  • 把“调度视图”
  • 转成“执行视图”
  • 再把结果整理成“上层还容易继续处理的返回视图”

这正是一个优秀 façade 应该做的事。

为什么 forward_batch_embedding(...) 也必须一起看#

这能防止你把 worker 理解成“只服务 token generation”。

forward_batch_embedding(...) 会：

  • 先构 ForwardBatch
  • 再 self.model_runner.forward(...)
  • 再取 embeddings

这说明 worker façade 同样服务于：

  • embedding
  • classification
  • rerank / score

也就是说，它封装的不是某一种业务，而是整个执行壳的“被上层安全调用”的方式。

为什么 TpModelWorker 里会有那么多 weights / LoRA 控制动作#

前面的运行时架构和控制面章节已经提过：

  • update_weights_from_disk
  • update_weights_from_distributed
  • update_weights_from_tensor
  • update_weights_from_ipc
  • load_lora_adapter
  • unload_lora_adapter

这些动作放回这棵树再看一次，会更清楚它们的共同点：

  • 上层不会直接改 ModelRunner
  • 它总是先通过 worker façade 进入

这说明 TpModelWorker 不只是执行 façade，也是：

  • rank-local 控制 façade

从架构设计看，这很健康，因为它避免了上层到处直接碰 ModelRunner 内部细节。

get_worker_info() 为什么很像这棵树的总结#

它一次性把：

  • token 容量
  • prefill 容量
  • 请求上限
  • memory pool 大小
  • 随机种子
  • device

这些东西打包出去。

这说明对上层来说，worker 真正代表的是：

  • 一组可执行能力和资源事实

而不只是“这里有个 forward 方法”。

因此 get_worker_info() 很适合被当成阅读这棵树后的一个自检点：如果你能理解它为什么返回这些东西，说明你已经看懂 worker 在系统里的角色了。

为什么这章和 model_runner.py 章节不重复#

它们分别站在不同层：

  • model_runner.py 讲执行壳内部如何汇合模型、KV、backend 和 sampling
  • tp_worker.py 讲这个执行壳怎样被包装成上层真正可驱动、可控制、可查询的 worker 单元

也就是说：

  • 一个偏壳内
  • 一个偏壳外

把两者都讲清楚，读者才真正能理解“为什么 scheduler 从不直接碰模型执行细节”。

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

1. 以为 TpModelWorker 只是 rank 容器#

它实际上还是调度视图到执行视图的关键翻译层。

2. 以为 weights / LoRA 控制面会直接改 ModelRunner#

上层通常先经过 worker façade。

3. 以为非生成路径绕过 worker#

forward_batch_embedding(...) 说明并没有。

如果你要顺着源码读这棵 worker façade 树，推荐顺序是什么#

建议按下面顺序：

  1. BaseTpWorker
  2. TpModelWorker.__init__()
  3. _init_model_config() / _init_model_runner()
  4. forward_batch_generation(...)
  5. forward_batch_embedding(...)
  6. 最后再看 weights / LoRA / memory 信息接口

这样读，你先知道 façade 暴露什么，再看它怎样把执行壳挂进来，最后再看控制与查询分支，最不容易被大文件打散。

小结#

这一章真正要补齐的，是代码导读里执行外壳的最后一块：

  • ModelRunner 是执行壳内部中枢
  • TpModelWorker 则是 scheduler 和控制面真正依赖的 worker façade

到这里，执行主线在代码导读层就从：

  • ForwardBatch
  • ModelRunner
  • Sampler

继续长到了：

  • 上层到底怎样安全而稳定地驱动这套执行壳

这使整条执行链在源码层更完整地闭环。