执行摘要

• 一句话：提取层索引搭建逻辑至独立模块 layer_setup.py
• 推荐动作：建议合入前务必解决机器人提出的向后兼容性问题：在 ModelRunner 中保留 self.start_layer/self.end_layer/self.num_effective_layers 作为 self.layer_info 的委托属性，或至少确保所有外部访问都已更新。同时应修复 adjust_hybrid_swa_layer_ids 中的 off-by-one 错误。建议补充针对提取后函数的单元测试。

功能与动机

PR body 指出这是一个提取重构系列中的一环，目的是将 ModelRunner 中与层索引搭建相关的辅助函数逐步分离到独立模块，以减少 ModelRunner 的膨胀并提高可维护性。该 PR 属于系列中的 'mrc-layer-setup' 步骤，专注于层索引设置。

实现拆解

1. 创建 layer_setup.py 模块：在 model_executor/model_runner_components/ 下新增文件，定义 AttentionAndMoeLayers NamedTuple、_PPLayerRange 和 ModelLayerInfo msgspec Struct，以及 compute_attention_and_moe_layers、_compute_model_num_layers、_resolve_pp_layer_range、_assert_pp_mtp_compat、_adjust_hybrid_swa_layer_ids 和 resolve_layer_indices 函数。其中 resolve_layer_indices 作为外部入口，统一封装层数计算、PP 范围、loop_num 乘数和 MTP 断言。
2. 改写 ModelRunner.initialize()：删除原有分散的层索引计算代码（约 50 行），替换为对 resolve_layer_indices 的单次调用，结果存入 self.layer_info。同时将 adjust_hybrid_swa_layers_for_pp 方法参数化为纯函数并移至 layer_setup.py。
3. 更新 model_runner_kv_cache_mixin.py：将原来直接引用 self.start_layer 和 self.end_layer 的地方改为 self.layer_info.start_layer 和 self.layer_info.end_layer，共 8 处。
4. 更新 MLX model_runner_stub.py：导入 ModelLayerInfo，在 initialize 中使用 ModelLayerInfo 封装之前手动赋值的 start_layer/end_layer/num_effective_layers。
5. 移除无用导入和全局变量：在 model_runner.py 中删除不再需要的 is_hip 导入和 _is_hip 全局变量，因为相关逻辑已移到 layer_setup.py（内部保留自己的 _is_hip）。

关键文件：

• python/sglang/srt/model_executor/model_runner_components/layer_setup.py（模块 层设置；类别 source；类型 data-contract；符号 AttentionAndMoeLayers, compute_attention_and_moe_layers, _PPLayerRange, ModelLayerInfo）: 新增的核心模块，包含所有层索引搭建逻辑和数据结构，是重构的主要产出。
• python/sglang/srt/model_executor/model_runner.py（模块 模型执行器；类别 source；类型 core-logic；符号 adjust_hybrid_swa_layers_for_pp）: 核心被重构的文件，删除了超过130行层索引逻辑，改为调用 resolve_layer_indices，是变更的主要影响对象。
• python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py（模块 缓存管理；类别 source；类型 data-contract）: 需要适配新的层信息访问方式，将 self.start_layer/self.end_layer 改为 self.layer_info.start_layer/self.layer_info.end_layer，共16行变更。
• python/sglang/srt/hardware_backend/mlx/model_runner_stub.py（模块 MLX 后端；类别 source；类型 data-contract）: MLX 后端的 ModelRunner 简化版需要同步修改，使用 ModelLayerInfo 结构替代手动属性赋值。

关键符号：resolve_layer_indices, _compute_model_num_layers, _resolve_pp_layer_range, _assert_pp_mtp_compat, adjust_hybrid_swa_layer_ids, compute_attention_and_moe_layers

关键源码片段

python/sglang/srt/model_executor/model_runner_components/layer_setup.py

新增的核心模块，包含所有层索引搭建逻辑和数据结构，是重构的主要产出。

# 文件：python/sglang/srt/model_executor/model_runner_components/layer_setup.py
from __future__ import annotations
from typing import TYPE_CHECKING, Any, NamedTuple
import msgspec
from sglang.srt.utils import is_hip
​
if TYPE_CHECKING:
    from sglang.srt.configs.model_config import ModelConfig
    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
​
_is_hip = is_hip()
​
# --- 数据结构 ---
​
class AttentionAndMoeLayers(NamedTuple):
    """收集模型中所有注意力层、MoE 层、MoE fusion 层和 DSA indexer 的列表。"""
    attention_layers: list[Any]
    moe_layers: list[Any]
    moe_fusions: list[Any]
    dsa_indexers: list[Any]
​
​
class _PPLayerRange(msgspec.Struct, frozen=True, kw_only=True):
    """表示流水线并行（PP）的层范围（内部使用）。"""
    start_layer: int
    end_layer: int
​
​
class ModelLayerInfo(msgspec.Struct, frozen=True, kw_only=True):
    """公开的层信息，替代之前分散的 start_layer / end_layer / num_effective_layers。"""
    start_layer: int
    end_layer: int
    num_effective_layers: int
​
​
# --- 核心逻辑 ---
​
def compute_attention_and_moe_layers(layer_model: Any) -> AttentionAndMoeLayers:
    """遍历模型的所有层，识别出注意力层、MoE 专家层等并收集。"""
    attention_layers: list[Any] = []
    moe_layers: list[Any] = []
    moe_fusions: list[Any] = []
    dsa_indexers: list[Any] = []
    for layer in layer_model.layers:
        # 注意力层识别（支持多种模型结构）
        attn_layer = None
        if hasattr(layer, "self_attn"):
            if hasattr(layer.self_attn, "attn"):
                attn_layer = layer.self_attn.attn
            elif hasattr(layer.self_attn, "attn_mqa"):
                attn_layer = layer.self_attn.attn_mqa
                if _is_hip and hasattr(layer.self_attn, "attn_mha"):
                    attn_layer._pcg_mha_companion = layer.self_attn.attn_mha
        elif hasattr(layer, "attn"):
            attn_layer = layer.attn
        elif hasattr(layer, "linear_attn"):
            attn_layer = layer.linear_attn.attn if hasattr(layer.linear_attn, "attn") else layer.linear_attn
        elif hasattr(layer, "attention"):
            if hasattr(layer.attention, "attn"):
                attn_layer = layer.attention.attn
        elif hasattr(layer, "mixer"):
            if hasattr(layer.mixer, "attn"):
                attn_layer = layer.mixer.attn
            elif hasattr(layer, "_forward_mamba"):
                attn_layer = layer
        if attn_layer is not None:
            attention_layers.append(attn_layer)
        elif hasattr(layer, "mixer"):
            attention_layers.append(None)
​
        # MoE 专家层识别
        moe_block = moe_fusion = None
        if hasattr(layer, "mlp") and hasattr(layer.mlp, "experts"):
            moe_block = layer.mlp.experts
            moe_fusion = layer.mlp
        if hasattr(layer, "block_sparse_moe") and hasattr(layer.block_sparse_moe, "experts"):
            moe_block = layer.block_sparse_moe.experts
            moe_fusion = layer.block_sparse_moe
        if hasattr(layer, "moe") and hasattr(layer.moe, "experts"):
            moe_block = layer.moe.experts
            moe_fusion = layer.moe
        if hasattr(layer, "mixer") and hasattr(layer.mixer, "experts"):
            moe_block = layer.mixer.experts
            moe_fusion = layer.mixer
        moe_layers.append(moe_block)
        moe_fusions.append(moe_fusion)
​
        # DSA indexer（用于 NSA 注意力）
        dsa_indexer = None
        if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "indexer"):
            dsa_indexer = layer.self_attn.indexer
        dsa_indexers.append(dsa_indexer)
​
    return AttentionAndMoeLayers(attention_layers, moe_layers, moe_fusions, dsa_indexers)

# （接上）计算模型层数、PP 范围解析、混合 SWA 调整、统一入口函数
​
def _compute_model_num_layers(*, model: Any, model_config: ModelConfig, is_draft_worker: bool) -> int:
    """计算模型实际层数，考虑 MTP draft worker 场景。"""
    _nnpl = model_config.num_nextn_predict_layers
    model_has_mtp_layers = _nnpl is not None and _nnpl > 0
    if is_draft_worker and model_has_mtp_layers:
        return getattr(model, "num_stages", _nnpl)
    # 注意：这里直接索引 architectures[0] 存在空列表风险
    if model_config.hf_config.architectures[0] in ("MiMoV2MTP", "Step3p5MTP"):
        return 1
    return max(model_config.num_hidden_layers, model_config.num_attention_layers)
​
​
def _resolve_pp_layer_range(*, model: Any, model_num_layers: int) -> _PPLayerRange:
    """从模型对象获取 PP 切分后的起始/结束层。"""
    return _PPLayerRange(
        start_layer=getattr(model, "start_layer", 0),
        end_layer=getattr(model, "end_layer", model_num_layers),
    )
​
​
def _assert_pp_mtp_compat(*, model_has_mtp_layers: bool, spec_algorithm: SpeculativeAlgorithm, num_effective_layers: int, model_num_layers: int) -> None:
    """断言 PP 与 MTP 模型兼容。"""
    if model_has_mtp_layers and not spec_algorithm.is_none() and num_effective_layers != model_num_layers:
        raise AssertionError("PP is not compatible with MTP models.")
​
​
def adjust_hybrid_swa_layer_ids(*, model_config: ModelConfig, start_layer: int, end_layer: int) -> list[int]:
    """对混合 SWA 模型，筛选出当前 PP 分片内的 full attention 和 SWA 层 ID。
    注意：end_layer 是 exclusive 上界，但此处使用了 end_layer + 1，可能引入 off-by-one。"""
    full_attention_layer_ids = [
        layer_idx
        for layer_idx in range(start_layer, end_layer + 1) # bug: 应使用 end_layer
        if hasattr(model_config, 'full_attention_layer_ids')
        and layer_idx in model_config.full_attention_layer_ids
    ]
    swa_attention_layer_ids = [
        layer_idx
        for layer_idx in range(start_layer, end_layer + 1) # 同样的 off-by-one
        if hasattr(model_config, 'swa_attention_layer_ids')
        and layer_idx in model_config.swa_attention_layer_ids
    ]
    return full_attention_layer_ids, swa_attention_layer_ids
​
​
def resolve_layer_indices(*, model: Any, model_config: ModelConfig, is_draft_worker: bool, spec_algorithm: SpeculativeAlgorithm) -> ModelLayerInfo:
    """统一入口：计算层索引信息。
    按顺序执行：计算模型层数 → 解析 PP 范围 → 处理 loop_num → 断言 MTP 兼容 → 调整混合 SWA。"""
    model_num_layers = _compute_model_num_layers(
        model=model, model_config=model_config, is_draft_worker=is_draft_worker
    )
    pp_range = _resolve_pp_layer_range(model=model, model_num_layers=model_num_layers)
    num_effective_layers = pp_range.end_layer - pp_range.start_layer
​
    loop_num = getattr(model_config.hf_config, "loop_num", 1)
    if loop_num > 1:
        num_effective_layers *= loop_num
​
    _nnpl = model_config.num_nextn_predict_layers
    model_has_mtp_layers = _nnpl is not None and _nnpl > 0
    _assert_pp_mtp_compat(
        model_has_mtp_layers=model_has_mtp_layers,
        spec_algorithm=spec_algorithm,
        num_effective_layers=num_effective_layers,
        model_num_layers=model_num_layers,
    )
​
    adjust_hybrid_swa_layer_ids(
        model_config=model_config,
        start_layer=pp_range.start_layer,
        end_layer=pp_range.end_layer,
    )
​
    return ModelLayerInfo(
        start_layer=pp_range.start_layer,
        end_layer=pp_range.end_layer,
        num_effective_layers=num_effective_layers,
    )

评论区精华

Gemini Code Assist 机器人提出了 4 条评论：

• 关键问题：删除 self.start_layer、self.end_layer、self.num_effective_layers 将会导致 AttributeError，因为这些属性在 model_runner_kv_cache_mixin.py 和其他地方仍被广泛引用。建议恢复这些属性作为向后兼容的桥接。
• 高优先级：MLX stub 也应同步设置 self.start_layer 等属性以避免继承方法中的错误。
• 中优先级：在 _compute_model_num_layers 中直接索引 architectures[0] 可能引发 IndexError，建议先检查列表是否非空。
• 中优先级：adjust_hybrid_swa_layer_ids 中使用了 end_layer + 1 作为 range 上限，而 end_layer 是 exclusive 的，这会导致越界。

所有评论均为机器人自动检查，没有人工 reviewer 参与讨论或给出结论。

• 删除 self.start_layer/end_layer/num_effective_layers 导致 AttributeError (correctness): 未解决，作者未回复，PR 已合并但风险仍存在。
• MLX stub 缺少桥接属性 (correctness): 未解决，作者未采纳。
• architectures[0] 直接索引可能 IndexError (correctness): 未解决，作者未采纳。
• end_layer + 1 导致 off-by-one 错误 (correctness): 未解决，作者未回复。

风险与影响

• 风险：
  1. 向后兼容风险（高）：直接删除 self.start_layer/end_layer/num_effective_layers 会导致外部代码（尤其是 mixin 和测试）出现 AttributeError。虽然该 PR 修改了 model_runner_kv_cache_mixin.py 中的引用，但可能还有其他地方未触及（如 model_runner.py 内剩余方法或其他模块）。机器人评论已明确指出此风险。
  2. 边界条件风险（中）：新的 adjust_hybrid_swa_layer_ids 函数中 range(start_layer, end_layer + 1) 使用了 +1，这与 end_layer 通常作为 exclusive 上限的语义冲突，可能引入 off-by-one 错误，影响混合 SWA 模型的层过滤。
  3. 空架构名风险（低）：_compute_model_num_layers 中 architectures[0] 的访问在 architectures 为空或 None 时会崩溃，虽然实际中很少出现，但防御性编码更安全。
  4. 无新增测试：该 PR 没有增加对应的单元测试来验证提取后的逻辑是否与重构前一致。
• 影响：
  • 对用户：无直接功能变化，推理行为应与重构前一致。但如有 bug 引入（如 off-by-one），可能影响特定模型（混合 SWA、MTP）的推理正确性。
  • 对系统：ModelRunner 初始化逻辑更加模块化，layer_setup.py 成为层索引配置的中心模块，便于未来的维护和扩展。
  • 对团队：重构系列的一部分，后续 PR 可依赖 layer_info 和 resolve_layer_indices 接口。需要警惕向后兼容性桥接的缺失。
  • 风险标记：向后兼容性缺失, off-by-one 错误, 空架构名崩溃风险, 无测试覆盖

关联脉络

• PR #31169 Split initialize() into orchestration helpers: 同属 ModelRunner 初始化逻辑拆分系列，是对 initialize() 整体解耦的后续步骤。
• PR #31168 Extract cuda-graph setup into a module: 同一系列重构，将 CUDA graph 配置提取到独立模块。
• PR #31167 Extract attention-backend setup into a module: 同一系列重构，将注意力后端设置提取到独立模块。
• PR #31166 Narrow component dependencies to injected fields instead of ModelRunner: 同一系列重构，解耦组件对 ModelRunner 的依赖，使用依赖注入。
• PR #31165 Drop ModelRunner's duplicated parallel-degree fields and read them via self.ps: 同一系列重构，统一并行度字段访问。