执行摘要

• 一句话：H3 MPS 内存优化：权重分阶段卸载、前向分块流式
• 推荐动作：值得精读，尤其 _forward_mps_streamed_attention 是“分块流式计算规避统一内存峰值”的典型范例，LayerwiseOffloadManager 对同步/异步双路径的抽象也有借鉴价值。阅读时建议对照 commit 历史中多次 revert 的 chunk size 调参过程，理解 MPS 内存-吞吐权衡的敏感性；但由于 Draft 合并且 CI Extra 失败，合入后应优先补齐端到端精度验证，再在真实 MPS 设备上推广。

功能与动机

MiniMax H3 的 packed 长序列和大组件工作集在 MPS 统一内存上压力极大。PR body 的动机原文是："MiniMax-H3's long packed sequence and large component working set need stricter placement on MPS unified memory." 具体而言，融合 QKV 投影输出在 768px 分辨率下约 1.45 GiB，与 packed residual、SDPA workspace 同时驻留易挤占系统内存；同时 MPS 没有 CUDA 式 pinned memory 与独立 copy stream，原 layerwise offload 的异步预取与 Event 同步模型并不适用。

实现拆解

实现按五步展开，全部落在 multimodal_gen 运行时，CUDA 路径仅做兼容性收口：

1. LayerwiseOffloadManager 的 MPS 化改造（python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py）：__init__ 新增 initialize 参数使初始化可延迟；识别 _synchronous_mps 后禁用 pinned memory 与 copy stream，prefetch_size 显式保持 0 避免下一层与当前层同时驻留；新增 _mps_cpu_weights 字典保留每层原始 CPU tensor，_initialize_mps_cpu_weights 把权重迁回 CPU 并用共享空 tensor 占位；prefetch_layer 在 MPS 分支做同步搬移，release_layer 强制 torch.mps.synchronize() + empty_cache()，sync_layer_to_cpu 直接跳过（推理不改参数）。同时新增 _capture_mps_cpu_non_layer_weights/materialize_mps_non_layer_weights/release_mps_non_layer_weights/restore_mps_cpu_non_layer_weights 管理条件投影、时间嵌入等非逐层权重。
2. 加载链路先 CPU 后设备（transformer_loader.py、vae_loader.py、text_encoder_loader.py、fsdp_load.py）：三个 loader 新增 customized_load_kwargs_for_component，在 MPS 且组件被配置为 layerwise offload 时返回 cpu_offload_flag=True；load_customized 接受该参数，checkpoint_load_device 强制为 CPU，WeightLoadPlan.for_component 增加 mps_layerwise_cpu_staging 标志；VAE loader 用 load_state_dict(..., assign=True) 直接复用 CPU 原始 tensor，避免加载后再拷一份。
3. DiT 前向分块流式化（python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py、sdpa.py）：_forward_mps_streamed_attention 第一阶段按 128 token chunk 生成完整 K/V 且不保留 Q，第二阶段逐序列逐 chunk 投影 Q、立即过 attention 与 out_proj 回写；MLP 与最终投影（video_out/audio_out）同样按 128 chunk 分块；_apply_rope 从 _apply_rope_qk 中抽取 eager 路径；qkv 权重标记 mps_zero_copy_unsafe=True 表明 checkpoint 行交错需重排、不可零拷贝；forward 入口以 x.device.type == "mps" and not ulysses_active 切到流式路径。
4. 准入与平台参数收紧（pipeline_configs/minimax_h3.py、server_args.py）：validate_server_args 在 MPS 下强制 transformer/text_encoder/video_vae/audio_vae 全部为 LAYERWISE_OFFLOAD 并拒绝 torch.compile；_adjust_platform_specific 收窄为仅允许 num_gpus=1、residency 只支持 RESIDENT 或 LAYERWISE_OFFLOAD，移除原先“MPS 强制关闭 layerwise offload”的逻辑。
5. 平台杂项与配套测试：gpu_worker.py 改 capture_memory_snapshot() 统计峰值（max_memory_reserved() 在 MPS 不可用）并跳过 set_device 与 reset_peak_memory_stats；vae_vit.py 与 decoding.py 用 _cuda_autocast_disabled/nullcontext 让 autocast 只在 CUDA 生效；hf_diffusers_utils.py 为 client has been closed 增加指数退避重试；test_minimax_h3_admission.py 新增 MPS 准入用例 test_mps_admission_requires_layerwise_residency_for_every_h3_component。

关键文件：

• python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py（模块 权重卸载；类别 source；类型 dependency-wiring；符号 initialize, _initialize_mps_cpu_weights, _capture_mps_cpu_non_layer_weights, _matches_mps_weight_prefix）: 整个 PR 的核心：LayerwiseOffloadManager 从纯 CUDA 异步模型扩展为同步 MPS 与异步 CUDA 双路径，新增 _mps_cpu_weights 结构、MPS 初始化/取层/释放/非层权重生命周期管理，是所有 MPS 组件驻留行为的地基。
• python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py（模块 H3 模型；类别 source；类型 data-contract；符号 _apply_rope, _forward_mps_streamed_attention）: DiT 前向的 MPS 流式执行主体：新增 _forward_mps_streamed_attention、MLP 与最终投影的 128 token 分块路径、_apply_rope 抽取和 mps_zero_copy_unsafe 标记，直接决定内存峰值下降幅度。
• python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py（模块 组件加载；类别 source；类型 core-logic；符号 customized_load_kwargs_for_component）: DiT 组件在 MPS 下先加载到 CPU 的入口：customized_load_kwargs_for_component 下发 cpu_offload_flag，load_customized 据此强制 checkpoint_load_device 为 CPU 并配置 mps_layerwise_cpu_staging。
• python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py（模块 组件加载；类别 source；类型 core-logic；符号 customized_load_kwargs_for_component）: VAE 组件同路径改造，关键差异是 load_state_dict 使用 assign=True 以复用 CPU 原始 tensor，避免 MPS 下额外拷贝。
• python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py（模块 配置校验；类别 source；类型 dependency-wiring）: MPS 准入校验：强制 transformer/text_encoder/video_vae/audio_vae 全部 LAYERWISE_OFFLOAD，并拒绝 torch.compile，防止未经验证的组合在 MPS 上运行。
• python/sglang/multimodal_gen/runtime/server_args/server_args.py（模块 启动参数；类别 source；类型 core-logic）: 解除 MPS 对 layerwise offload 的旧限制并新增 num_gpus=1 约束，是平台参数层的策略翻转，支撑本 PR 的 MPS 驻留方案。
• python/sglang/multimodal_gen/runtime/managers/gpu_worker.py（模块 工作进程；类别 source；类型 dependency-wiring）: MPS 适配的必要杂项：跳过 set_device 与 reset_peak_memory_stats，内存峰值改用 capture_memory_snapshot()，因为 MPS 下 max_memory_reserved() 不可用。
• python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py（模块 VAE 模型；类别 source；类型 data-contract；符号 _cuda_autocast_disabled）: autocast 平台适配：_cuda_autocast_disabled 让 CUDA 之外的设备（含 MPS）走 nullcontext，避免 MPS 上触发不受支持的 autocast。
• python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py（模块 解码阶段；类别 source；类型 data-contract）: 解码阶段的 autocast 同样改为 CUDA 专属，MPS 下用 nullcontext 包裹 VAE decode。
• python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py（模块 准入测试；类别 test；类型 test-coverage；符号 test_mps_admission_requires_layerwise_residency_for_every_h3_component）: 新增 MPS 准入测试，锁定 MiniMax H3 在 MPS 下必须全组件 layerwise offload 的契约，防止后续参数调整破坏该约束。

关键符号：LayerwiseOffloadManager.initialize, LayerwiseOffloadManager._initialize_mps_cpu_weights, LayerwiseOffloadManager.prefetch_layer, LayerwiseOffloadManager.release_layer, LayerwiseOffloadManager.sync_layer_to_cpu, LayerwiseOffloadManager.materialize_mps_non_layer_weights, LayerwiseOffloadManager.release_mps_non_layer_weights, MiniMaxH3DiTModel._forward_mps_streamed_attention, _apply_rope, TransformerLoader.customized_load_kwargs_for_component, VAELoader.customized_load_kwargs_for_component

关键源码片段

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

整个 PR 的核心：LayerwiseOffloadManager 从纯 CUDA 异步模型扩展为同步 MPS 与异步 CUDA 双路径，新增 _mps_cpu_weights 结构、MPS 初始化/取层/释放/非层权重生命周期管理，是所有 MPS 组件驻留行为的地基。

# layerwise_offload.py 节选：MPS 同步层权重的卸下与取回
# 前置状态：__init__ 识别到 current_platform.is_mps() 时，
# 会设置 _synchronous_mps=True、pin_cpu_memory=False、copy_stream=None，
# 并用 _mps_cpu_weights 字典存放每层原始 CPU tensor。
​
@torch.compiler.disable
def _initialize_mps_cpu_weights(self) -> None:
    """把每层权重迁回 CPU，并在设备上替换为共享空 tensor 占位。
​
    CUDA 路径会把同层同 dtype 的权重合并成扁平化的 pinned CPU 副本；
    MPS 与主机共享物理内存，直接保留原始 CPU tensor 即可，避免再驻留
    一份主机拷贝。
    """
    for name, tensor in self._named_parameters.items():
        layer_idx = self._match_layer_idx(name)
        if layer_idx is None or layer_idx >= self.num_layers:
            continue
        local_tensor = self._to_local_tensor(tensor).detach()
        cpu_tensor = (
            local_tensor
            if local_tensor.device.type == "cpu"
            else local_tensor.to("cpu")
        )
        self._mps_cpu_weights.setdefault(layer_idx, {})[name] = cpu_tensor
        self._weight_metadata.setdefault(layer_idx, {})[name] = {
            "dtype": cpu_tensor.dtype,
        }
        # 用共享的 1 元素占位 tensor 顶替原参数，释放设备侧大块显存
        tensor.data = self._get_shared_empty_tensor_for_target(
            tensor, cpu_tensor.dtype
        )
​
    torch.mps.empty_cache()
    self.register_forward_hooks()
    self._configured = True
    logger.info(
        f"Initialized synchronous MPS layerwise offload with "
        f"{self.num_layers} layers"
    )
​
​
def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None:
    """幂等地把指定层权重搬回设备；MPS 下是同步搬移。"""
    if not self.enabled or self.device is None:
        return
    if layer_idx < 0 or layer_idx >= self.num_layers:
        return
    if layer_idx in self._gpu_layers:
        return
​
    if self._synchronous_mps:
        cpu_weights = self._mps_cpu_weights.get(layer_idx)
        if not cpu_weights:
            return
        with torch.inference_mode(False), torch.no_grad():
            for name, cpu_tensor in cpu_weights.items():
                target = self.get_target_with_name(name)
                # 统一内存下 CPU 到 MPS 的搬移没有异步语义，non_blocking 恒为 False
                target.data = self._wrap_for_target(
                    target,
                    cpu_tensor.to(device=self.device, non_blocking=False),
                )
        self._gpu_layers.add(layer_idx)
        return
​
    # 以下是 CUDA 异步路径：pinned 副本 + copy stream + Event 记录
    # MPS 分支已在上方返回，CUDA 路径行为保持不变

评论区精华

本 PR 没有任何 reviewer 评论，唯一交互是作者在 Issue 评论区触发 /tag-and-rerun-ci 重跑 CI。PR body 明确标注 Draft 状态并声明“end-to-end and accuracy validation are pending”，且 CI Extra 运行（Run #32208109060）状态为失败，说明正确性验证尚未收敛；这些未决事项直接决定了本 PR 的合入风险等级。

• 暂无高价值评论线程

风险与影响

• 风险：
  1) 正确性未验证：PR body 明说 end-to-end and accuracy validation are pending，CI Extra 失败；MPS 流式 attention 按 chunk 重算 norm 与 RoPE，其等价性依赖 _apply_rope 与 CUDA 路径逐字节一致，建议补齐精度对照后再在生产启用。
  2) 性能同步开销：每个 chunk 后强制 torch.mps.synchronize() + torch.mps.empty_cache()，同步点密集；最终 chunk size 128 来自 30 个 commit 中多次“Increase/Batch”后被 revert、再“Lower/Tune”的反复试验，说明对具体机型内存-吞吐权衡高度敏感。
  3) 全局下载重试改动：hf_diffusers_utils.py 新增 client has been closed 的捕获与退避重试，修改 snapshot_download 的错误语义，影响所有平台的模型拉取路径。
  4) 内存统计口径变化：gpu_worker.py 在 MPS 下改用 snapshot 峰值，与 CUDA 的 max_memory_reserved() 口径不同，跨平台对比数据时易产生误读。
  5) 双份权重字典契约：_mps_cpu_weights 与 _consolidated_cpu_weights 并存，update_cpu_weights 等写回路径需同时维护两套结构，后续扩展时容易漏改一侧。
  • 影响：影响面集中在 MPS（Apple Silicon）平台的 MiniMax H3 推理：所有分支均以 current_platform.is_mps() 或 x.device.type == "mps" 为条件，CUDA/ROCm 主路径行为不变；唯一全局影响是 Hugging Face 下载重试逻辑。对用户而言，MiniMax H3 在 MPS 上首次具备可运行性且可支撑更长序列；对团队而言，LayerwiseOffloadManager 从此承载同步（MPS）与异步（CUDA）两套执行模型，后续新增平台需同时考虑两条分支；由于 Draft 且 CI Extra 失败，该能力目前处于“可运行但未验证”状态，不建议直接作为发布基线。
  • 风险标记：正确性未验证（Draft）, CI Extra 失败, MPS 专用路径，CUDA 不受影响, 全局 Hugging Face 下载重试改动, 同步点密集，吞吐受限, chunk size 经验调参敏感

关联脉络

• PR #34993 [diffusion] fix: make MiniMax-H3 AdaLN cache rebuild transactional: 同改 python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py，同为 MiniMax H3 在 diffusion 运行时的稳定性补强；本 PR 引入 MPS 流式执行，34993 修复 AdaLN 缓存重建，两条链路共同保障 H3 推理正确与稳定。
• PR #35339 [diffusion] Per-request lossy accelerations: Cache-DiT, CFG gating, attention backend override: 同改 MiniMax H3 的 denoising 阶段与采样配置，说明 H3 pipeline 正在被系统性优化；本 PR 关注 MPS 平台内存，35339 关注按请求精度开关，方向互补。