执行摘要

• 一句话：修复EAGLE下静态显存分数未计入draft KV缓存导致OOM
• 推荐动作：此PR是一次重要的内存管理重构，修复了长期存在的draft模型显存预算缺失问题。设计上将初始化阶段清晰分离，提高了可维护性和可预测性。建议所有使用推测解码的团队精读，尤其是init_model_worker的编排逻辑和各worker的alloc_memory_pool/init_backends实现。

功能与动机

根据PR描述，当EAGLE推测解码启用时，--mem-fraction-static仅控制目标模型的KV缓存池。draft模型的权重和KV缓存未计入预算，导致目标KV池过分配（因内存分析在draft权重加载前进行），且draft KV池在同一max_total_num_tokens下额外分配。之前通过粗糙的4GB/6GB启发式预留缓解，但对于大型模型或用户指定的--mem-fraction-static仍然不足。

实现拆解

1. 提前获取draft模型层数（model_runner.py）：在__init__中增加eagle_draft_num_layers属性，通过_build_model_config在权重加载前读取draft配置，记录层数用于后续KV池容量核算。同时保存pre_model_load_memory快照。

2. 调度器初始化三阶段化（scheduler.py）：新增init_target_memory_pool、init_memory_pools、init_all_backends三个方法，在init_model_worker中按序调用：先加载目标权重（init_tp_model_worker），再加载draft权重（maybe_init_draft_worker），然后分配KV池（此时内存分析已见所有权重），最后初始化后端和CUDA图。

3. draft worker解耦池分配与后端初始化（多个speculative worker）：Eagle、multi-layer Eagle、DFlash、FrozenKVMTP等worker的构造器仅加载权重，不再直接分配池或初始化后端；新抽取的alloc_memory_pool和init_backends方法由调度器在适当时机调用，接收目标池的配置。

4. tp_worker通用化池分配（tp_worker.py）：新增alloc_memory_pool和init_backends方法，支持外部传入池对象或创建新池，并统一对model_runner_list中所有runner执行操作。get_worker_info后移max_req_len计算到池分配之后。

5. 移除启发式预留（server_args.py）：删除speculative_draft_load_format逻辑中针对EAGLE/STANDALONE的4GB/6GB自动mem_fraction_static覆盖，因为新机制已精确核算draft模型开销。

6. 配套测试调整：新增test_model_config_shapes.py验证ModelConfig._derive_model_shapes中的head_dim派生逻辑；更新test_eagle_worker_v2_topk1_fastpath.py适配新的后端初始化接口。

关键文件：

• python/sglang/srt/model_executor/model_runner.py（模块 模型执行器；类别 source；类型 core-logic；符号 _build_model_config, initialize, alloc_memory_pool, init_backends）: 核心变更文件：在__init__中提前加载draft模型层数用于KV池核算，保存pre_model_load_memory供内存分析使用；新增eagle_draft_num_layers属性。
• python/sglang/srt/managers/scheduler.py（模块 调度器；类别 source；类型 core-logic；符号 init_target_memory_pool, init_memory_pools, init_all_backends）: 调度器新增三个方法编排初始化阶段，是流程重构的核心调度节点。
• python/sglang/srt/speculative/eagle_worker_v2.py（模块 Eagle工作器；类别 source；类型 core-logic；符号 alloc_memory_pool, init_backends, spec_v2_attn_backends）: Eagle draft worker重构：构造仅加载权重，池分配和后端初始化延后到新方法中。
• python/sglang/srt/managers/tp_worker.py（模块 TP工作器；类别 source；类型 core-logic；符号 alloc_memory_pool, init_backends）: TP worker新增 alloc_memory_pool 和 init_backends 通用方法，支持分离调用和外部池注入。
• test/registered/unit/configs/test_model_config_shapes.py（模块 配置测试；类别 test；类型 test-coverage；符号 _make_text_config, TestModelConfigShapes, _derive_shapes, test_optional_head_dims_default_when_none）: 新增单元测试验证ModelConfig形状推导逻辑，确保head_dim等派生正确。

关键符号：_build_model_config, initialize, alloc_memory_pool, init_backends, init_target_memory_pool, init_memory_pools, init_all_backends, spec_v2_attn_backends, _derive_model_shapes, _make_text_config, _derive_shapes, test_optional_head_dims_default_when_none, test_explicit_head_dims_are_preserved

关键源码片段

python/sglang/srt/model_executor/model_runner.py

核心变更文件：在__init__中提前加载draft模型层数用于KV池核算，保存pre_model_load_memory供内存分析使用；新增eagle_draft_num_layers属性。

# python/sglang/srt/model_executor/model_runner.py
# 在 __init__ 中新增：提前获取 draft 模型层数
self.eagle_draft_num_layers = None # 新增属性，记录 draft 模型层数
if (
    (self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone())
    and not self.is_draft_worker
    and server_args.speculative_draft_model_path
):
    # 在权重加载前读取 draft 配置，获取层数用于 KV 缓存大小核算
    draft_model_config = self._build_model_config(
        server_args,
        model_path=server_args.speculative_draft_model_path,
        model_revision=server_args.speculative_draft_model_revision,
        is_draft_model=True,
    )
    num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers
    if num_nextn_predict_layers is not None:
        self.eagle_draft_num_layers = int(num_nextn_predict_layers)
    else:
        # fallback: 取 hidden_layers 和 attention_layers 的较大值
        self.eagle_draft_num_layers = int(
            max(
                draft_model_config.num_hidden_layers,
                draft_model_config.num_attention_layers,
            )
        )
​
# 保存权重加载前的可用内存快照，供 alloc_memory_pool 使用
self.pre_model_load_memory = self.init_torch_distributed()

python/sglang/srt/managers/scheduler.py

调度器新增三个方法编排初始化阶段，是流程重构的核心调度节点。

# python/sglang/srt/managers/scheduler.py
# 新增的三阶段初始化方法：
def init_target_memory_pool(self):
    """仅分配目标KV池（如尚未分配）"""
    if (
        self.tp_worker.model_runner.memory_pool_config is not None
        and self.tp_worker.model_runner.req_to_token_pool is not None
        and self.tp_worker.model_runner.token_to_kv_pool_allocator is not None
    ):
        return
    self.tp_worker.alloc_memory_pool()
​
def init_memory_pools(self):
    """分配所有worker的KV池（目标和draft）"""
    self.init_target_memory_pool()
    if self.draft_worker is not None:
        pool, allocator = self.tp_worker.get_memory_pool()
        self.draft_worker.alloc_memory_pool(
            memory_pool_config=self.tp_worker.model_runner.memory_pool_config,
            req_to_token_pool=pool,
            token_to_kv_pool_allocator=allocator,
        )
​
def init_all_backends(self):
    """初始化注意力后端和CUDA图"""
    self.tp_worker.init_backends()
    if self.draft_worker is not None:
        self.draft_worker.init_backends()
​
def init_model_worker(self):
    # 1. 加载目标权重
    self.init_tp_model_worker()
    if self.spec_algorithm.is_frozen_kv_mtp():
        # Frozen-KV MTP 需要目标 KV 池提前就绪
        self.init_target_memory_pool()
    # 2. 加载 draft 权重
    self.maybe_init_draft_worker()
    # 3. 分配所有 KV 池（此时所有权重已加载，内存分析准确）
    self.init_memory_pools()
    # 4. 初始化后端和 CUDA 图
    self.init_all_backends()
    # ... 后续逻辑不变 ...

python/sglang/srt/speculative/eagle_worker_v2.py

Eagle draft worker重构：构造仅加载权重，池分配和后端初始化延后到新方法中。

# python/sglang/srt/speculative/eagle_worker_v2.py
# 构造中仅加载 draft 模型权重，不再分配 KV 池和初始化后端
# 原 __init__ 中与池分配、后端初始化相关的代码被移除
class EagleDraftWorker:
    def __init__(self, ..., target_worker):
        # ... 前面复制参数 ...
        # 只加载 draft 模型权重（不传 pool 和 memory_pool_config）
        self.draft_worker = TpModelWorker(
            server_args=server_args,
            gpu_id=gpu_id,
            tp_rank=tp_rank,
            is_draft_worker=True,
            # 不再传入 req_to_token_pool、token_to_kv_pool_allocator、memory_pool_config
            ...
        )
        self.draft_runner = self.draft_worker.model_runner
        # ... 其他一些配置 ...
​
    # 新增：由调度器在适当时候调用，分配 KV 池
    def alloc_memory_pool(self, memory_pool_config=None, req_to_token_pool=None, token_to_kv_pool_allocator=None):
        self.req_to_token_pool = req_to_token_pool
        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
        self.draft_worker.alloc_memory_pool(
            memory_pool_config=memory_pool_config,
            req_to_token_pool=req_to_token_pool,
            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
        )
        self.init_token_map()
        self.init_lm_head()
​
    # 新增：由调度器在池分配后调用，初始化后端和 CUDA 图
    def init_backends(self):
        with self.draft_tp_context(self.draft_runner.tp_group), ...:
            self.draft_worker.init_backends(disable_cuda_graph=True)
            self.init_attention_backend()
            if check_cuda_graph_backend(...):
                self.draft_runner.init_prefill_cuda_graph(force_for_draft_worker=True)
            self.init_cuda_graphs()
        if (c := self.draft_runner.canary_manager) is not None:
            c.mark_init_finished()

评论区精华

Review中merrymercy提出了三个关键评论：

• 在model_config.py中，询问将head_dim写回到hf_text_config是否有必要，cctry确认不需要并承诺删除。
• 在model_runner_kv_cache_mixin.py中，指出注释将draft模型权重描述为“动态部分”不够准确，cctry解释这是向后兼容的临时措施，若计入draft则静态分数可能超过0.9导致测试失败，故保留注释说明。

• 在frozen_kv_mtp_worker_v2.py中，建议将init_backends拆分为更具体命名的小函数，cctry同意并拆分。

• _derive_model_shapes中写回head_dim到hf_text_config的必要性 (design): 将移除写回操作，保持派生逻辑仅存储在self.head_dim上。

• _profile_available_bytes注释描述draft权重为动态部分是否准确 (documentation): 保留注释作为临时兼容措施的说明，不修改行为。
• init_backends命名模糊，建议拆分 (design): 将init_backends拆分为多个职责单一的小函数。

风险与影响

• 风险：

  1. 初始化顺序依赖：三阶段化依赖于调度器正确调用顺序，若后续新增worker类型未按此模式实现，可能引发未初始化池访问。
  2. 移除启发式预留：此前依赖4GB/6GB自动预留的配置（即使不合理）可能因移除而OOM，需要用户重新调整--mem-fraction-static。
  3. 池共享契约变更：draft worker构造不再接收req_to_token_pool和token_to_kv_pool_allocator，而是通过alloc_memory_pool传递，若第三方扩展未更新将出错。
  4. 测试覆盖风险：虽然新增了单元测试，但多worker交互场景（如EAGLE+DFLASH组合）可能覆盖不足，引入回归。
  5. 启动时间增加：阶段化可能导致额外同步开销，但影响通常很小。
     - 影响：影响范围：所有使用EAGLE或其他推测解码（Standalone/DFLASH）的用户。启动时OOM问题消除，显存利用更精确。
     影响程度：较大。涉及核心初始化流程重构，但对外部API无破坏性变更；用户无需修改配置即可受益。
     团队影响：开发者需要理解新的三阶段初始化模式，future speculative worker应遵循此模式。

• 风险标记：核心路径变更, 移除启发式配置, 兼容性风险, 测试覆盖不足

关联脉络

• 暂无明显关联 PR