执行摘要

• 一句话：引入 WeightLoadPlan 显式控制权重加载设备放置
• 推荐动作：该 PR 是重要的架构改进，值得所有参与 diffusion 模型加载的开发者精读。其引入的 WeightLoadPlan 抽象可用于未来扩展（如异构设备、分阶段加载）。建议 merge 后运行完整的 diffusion 测试套件以验证兼容性。

功能与动机

PR #29903 修复了在线 FP8 + dit_cpu_offload 崩溃，但修复方式是在 process_weights_after_loading 完成前将权重保留在设备上，这本质上属于加载时的设备放置策略。此 PR 将该策略显式化为 WeightLoadPlan，避免将加载决策混入运行时 server args 和组件加载器中的条件逻辑，使得加载管道更具可维护性。

实现拆解

1. 新建 WeightLoadPlan 数据类：位于 weight_load_plan.py，包含 checkpoint_load_device、weight_postprocess_device 和 defer_component_cpu_offload 三个字段，并提供 for_component 工厂方法根据是否需要设备端后处理自动设置。
2. 提取组件加载辅助方法：在 component_loader.py 中，将原 load 方法中的定制和原生加载上下文管理提取为 _load_customized_with_context 和 _load_native_with_context 两个独立方法，使得位置相关的控制流不再混入主干逻辑。
3. 重命名并扩展设备后处理检测：将 transformer_load_utils.py 中的 _requires_device_weight_processing 重命名为 _needs_device_weight_postprocess；同时从 if-else 结构改为字典映射，新增对 mxfp8 和 mxfp4_npu 的支持，以覆盖更多在线量化场景。
4. 修改 FSDP 加载入口：fsdp_load.py 中 maybe_load_fsdp_model 函数的参数从单个布尔值 defer_cpu_offload_until_after_weight_processing 替换为完整的 WeightLoadPlan 对象，并在 FSDP 模式下忽略 weight_postprocess_device 以避免不支持的设备转移。
5. 对接 transformer 和 bridge 加载器：在 transformer_loader.py 和 bridge_loader.py 中，调用 WeightLoadPlan.for_component 生成计划并传入 FSDP 加载函数，替代之前的条件逻辑。
6. 新增单元测试：在 test_transformer_quant.py 中补充了 test_weight_load_plan_defers_cpu_offload_for_device_postprocess 和扩展的 test_online_fp8_needs_device_weight_postprocess，覆盖新抽象的各种分支。

关键文件：

• python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py（模块 加载器；类别 source；类型 core-logic；符号 WeightLoadPlan, for_component）: 核心新文件，定义 WeightLoadPlan 数据类和工厂方法，是整个重构的中心抽象。
• python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py（模块 加载器；类别 source；类型 core-logic；符号 _load_customized_with_context, _load_native_with_context）: 重构了 load 方法，将定制和原生加载上下文提取为独立方法，提高了可读性和可维护性。
• python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py（模块 量化加载；类别 source；类型 core-logic；符号 _requires_device_weight_processing, _needs_device_weight_postprocess）: 重命名核心函数并扩展量化格式支持，确保在线量化正确触发设备后处理。
• python/sglang/multimodal_gen/runtime/loader/fsdp_load.py（模块 FSDP 加载；类别 source；类型 dependency-wiring）: 修改 maybe_load_fsdp_model 接受 WeightLoadPlan 替代布尔参数，并在 FSDP 模式下忽略设备 override。
• python/sglang/multimodal_gen/test/unit/test_transformer_quant.py（模块 量化测试；类别 test；类型 test-coverage；符号 _make_quant_config, test_weight_load_plan_defers_cpu_offload_for_device_postprocess, test_online_fp8_needs_device_weight_postprocess）: 新增测试用例覆盖 WeightLoadPlan 逻辑和扩展的量化格式检测。
• python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py（模块 加载器；类别 source；类型 dependency-wiring）: 使用 WeightLoadPlan.for_component 替代手动布尔条件，对接新抽象。
• python/sglang/multimodal_gen/runtime/loader/component_loaders/bridge_loader.py（模块 加载器；类别 source；类型 dependency-wiring）: 与 transformer_loader 类似，使用 WeightLoadPlan，保持一致性。

关键符号：WeightLoadPlan.init, WeightLoadPlan.for_component, ComponentLoader._load_customized_with_context, ComponentLoader._load_native_with_context, _needs_device_weight_postprocess, maybe_load_fsdp_model

关键源码片段

python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py

核心新文件，定义 WeightLoadPlan 数据类和工厂方法，是整个重构的中心抽象。

from dataclasses import dataclass
​
import torch
​
​
@dataclass(frozen=True)
class WeightLoadPlan:
    """Device plan for checkpoint loading, before runtime residency takes over."""
​
    # Device used while materializing checkpoint tensors from files.
    checkpoint_load_device: torch.device
    # Device required while running process_weights_after_loading;
    # None means use checkpoint_load_device (unchanged).
    weight_postprocess_device: torch.device | None = None
    # Delay non-FSDP component CPU offload until after weight postprocessing.
    defer_component_cpu_offload: bool = False
​
    @classmethod
    def for_component(
        cls,
        *,
        checkpoint_load_device: torch.device,
        needs_device_weight_postprocess: bool,
        component_cpu_offload: bool,
    ) -> "WeightLoadPlan":
        # If on-device weight postprocessing is required, load directly to
        # device to speed up loading; otherwise keep postprocess device as None.
        weight_postprocess_device = (
            checkpoint_load_device if needs_device_weight_postprocess else None
        )
        # Defer CPU offload only when both conditions are true:
        # on-device postprocessing required AND component CPU offload enabled.
        return cls(
            checkpoint_load_device=checkpoint_load_device,
            weight_postprocess_device=weight_postprocess_device,
            defer_component_cpu_offload=(
                needs_device_weight_postprocess and component_cpu_offload
            ),
        )

评论区精华

Reviewer (gemini-code-assist[bot]) 提出三个关键意见：

• _needs_device_weight_postprocess 应支持 mxfp8 和 mxfp4_npu 格式，而非仅 fp8 和 mxfp4。
• 当 use_fsdp=True 时，应忽略 weight_postprocess_device override，因为 FSDP 模型不能直接调用 model.to(device)。

• 需增加单元测试覆盖新增的量化格式。
  所有意见均在后续提交中解决，最终版本已包含相应改动和测试。

• 扩展 _needs_device_weight_postprocess 支持 mxfp8 和 mxfp4_npu (correctness): 作者采纳建议，将函数改为字典映射，新增 mxfp8 和 mxfp4_npu 条目，并统一使用 is_checkpoint_fp8_serialized 作为 mxfp8 的序列化标志。

• FSDP 模式下忽略 weight_postprocess_device (correctness): 作者在 fsdp_load.py 中添加逻辑：如果 use_fsdp 且 weight_postprocess_device 非 None，则记录警告并置为 None。
• 为 mxfp8 和 mxfp4_npu 添加单元测试 (testing): 作者在 test_transformer_quant.py 中增加了 test_online_fp8_needs_device_weight_postprocess 用例，使用 _make_quant_config 构造模拟配置进行多格式测试。

风险与影响

• 风险：主要风险在于新引入的 WeightLoadPlan 可能与其他未迁移的加载路径不兼容，尤其是当某些 loader 仍然使用旧的布尔参数时。但本次变更已将所有已知路径（transformer、bridge、FSDP）统一使用计划对象，因此风险较低。潜在的边界情况是：如果组件加载器（如 native fallback）未通过 WeightLoadPlan，则可能沿用旧的隐式行为。此外，FSDP 路径中忽略 weight_postprocess_device 可能掩盖某些配置错误。最后，新增的量化格式支持依赖于 is_checkpoint_mxfp4_npu_serialized 等属性是否在量化配置中一致设置。
• 影响：影响范围限于 diffusion 模型加载流程，特别是使用在线量化（FP8/MXFP4/MXFP8）搭配 CPU offload 的场景。对无该配置的用户透明，加载行为不变。对开发者而言，加载设备放置策略变得显式、可测试，减少了 server_args 的条件耦合，有助于未来支持更多量化格式和硬件。测试覆盖显著提升。
• 风险标记：核心路径变更, 新抽象引入, FSDP 兼容性处理, 量化格式扩展

关联脉络

• PR #29903 [diffusion] fix: fix z-Image online fp8 quantization crash with dit_cpu_offload: 本 PR 直接解决 #29903 留下的隐式加载决策问题，将其显式化为 WeightLoadPlan。