执行摘要

• 一句话：简化扩散解耦传输层，清理过期同步 API 与重复 schema
• 推荐动作：值得快速翻阅，不必精读：作为 refactor-only 清理，它展示了在活跃子系统中收敛 schema、删除 dead API、复用已有工具函数的实践。重点看 manager.py 中同步/异步 API 的去留判断、protocol.py 中 TransferStagedMsg 的 schema 化，以及 _prefetch_transfer_ready 返回元组的收窄方式。若要在生产长期维护，建议补一个覆盖 async stage/load 的轻量单测，并说明 extra CI 失败的原因。

功能与动机

PR body 指出传输路径存在重复的 schema 组装，以及 diffusion runtime 根本不调用的 API 造成的双份维护面：“The transfer path had duplicate schema assembly and two maintenance surfaces for APIs that are not called by the diffusion runtime. Consolidating the live path makes RDMA staging and prefetch ownership easier to reason about.” 同时明确 “No public API changes are intended”，说明这是一次内部卫生清理。

实现拆解

1. schema 收敛到 transport/protocol.py：protocol.py 移除无用的 logging 与 __post_init__，TransferStagedMsg.manifest 改用 field(default_factory=dict)，并新增 scalar_fields 字段。scheduler_mixin.py 的 _disagg_encoder_transfer_stage 从手工 json.dumps(staged_data) 改为 encode_transfer_msg(TransferStagedMsg(...))，使 protocol 模块成为 staged 消息的唯一 schema。
2. 复用默认值判断逻辑：删除 scheduler_mixin.py 中 _BASE_SP_DEFAULTS 缓存与长注释，改为对每个 SamplingParams 字段调用已有 _is_default(value, f) 判断是否跳过默认值字段，避免两套“字段是否等于默认值”的实现。
3. 精简 prefetch 状态：_prefetch_transfer_ready 返回值从六元组收窄为四元组 (req, load_event, request_id, prealloc_slot_id)；_disagg_prefetch_event_loop 与 _handle_transfer_ready 同步调整解包；_disagg_denoiser_compute / _disagg_decoder_compute 移除 role_name 形参；多 rank 判断收敛为 _is_multi_rank()。
4. 删除传输层死代码：manager.py 删除同步 stage_tensors、load_tensors 及 get_receive_slot_addr / get_receive_slot_offset / get_staged_info / free_slots_count；buffer.py 的 write_tensor 去掉未使用的 name 参数，并删除 free_slots_count / get_stats 包装；allocator.py 删除 count_free_slots。
5. 测试与 CI 配套：没有新增测试文件；作者按 SGLang-Diffusion 重构指导未做本地验证。CI 中基础 PR Test 通过、extra CI 失败后仍合并，缺少对删除 API 的回归保护。

关键文件：

• python/sglang/multimodal_gen/runtime/disaggregation/transport/manager.py（模块 传输管理；类别 source；类型 core-logic；符号 stage_tensors, load_tensors, get_receive_slot_addr, get_receive_slot_offset）: 本次清理的核心文件，移除 6 个无调用方的同步/查询 API 并收窄 StagedTransfer.slot 为可空类型，直接决定传输层 API 面。
• python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py（模块 解耦调度；类别 source；类型 core-logic；符号 _disagg_denoiser_compute, _disagg_decoder_compute, _prefetch_transfer_ready, _disagg_encoder_transfer_stage）: prefetch 状态简化与 TransferStagedMsg 落地的主战场，涉及默认值判断逻辑、compute 函数签名与消息编码多处修改。
• python/sglang/multimodal_gen/runtime/disaggregation/transport/buffer.py（模块 传输缓冲；类别 source；类型 core-logic；符号 write_tensor, write_tensors_from_gpu, free_slots_count, get_stats）: 传输缓冲层的配套清理：write_tensor 移除未使用参数，并删除无人调用的统计包装方法。
• python/sglang/multimodal_gen/runtime/disaggregation/transport/allocator.py（模块 内存分配；类别 source；类型 core-logic；符号 count_free_slots）: 配合 free_slots_count 删除而移除 count_free_slots，证明该估算逻辑没有消费者。
• python/sglang/multimodal_gen/runtime/disaggregation/transport/protocol.py（模块 传输协议；类别 source；类型 core-logic；符号 TransferStagedMsg, post_init）: TransferStagedMsg 成为 staged 消息唯一 schema，新增 scalar_fields 并移除 __post_init__，是本次 schema 收敛的落点。

关键符号：stage_tensors, load_tensors, get_receive_slot_addr, get_receive_slot_offset, get_staged_info, free_slots_count, count_free_slots, get_stats, _disagg_denoiser_compute, _disagg_decoder_compute, _prefetch_transfer_ready, _disagg_encoder_transfer_stage, write_tensor

关键源码片段

python/sglang/multimodal_gen/runtime/disaggregation/scheduler_mixin.py

prefetch 状态简化与 TransferStagedMsg 落地的主战场，涉及默认值判断逻辑、compute 函数签名与消息编码多处修改。

    def _prefetch_transfer_ready(self: Scheduler, msg: dict) -> tuple:
        """Load tensors from transfer buffer and build Req for a transfer_ready message.
​
        Called from the recv prefetch thread. Loads on _transfer_stream and builds the Req,
        so the main thread can start compute immediately. Returns (req, load_event, request_id, prealloc_slot_id).
        """
        request_id = msg['request_id']
        manifest = msg.get('manifest', {})
        scalar_fields = msg.get('scalar_fields', {})
​
        if self._disagg_metrics:
            self._disagg_metrics.record_request_start(request_id)
​
        # 处理预分配 slot：若发送方已分配好接收槽，直接注册进 transfer manager
        prealloc_slot_id = scalar_fields.pop('_prealloc_slot_id', None)
        if (prealloc_slot_id is not None and prealloc_slot_id in self._preallocated_slots):
            slot = self._preallocated_slots[prealloc_slot_id]
            self._transfer_manager.register_prealloc_as_receive(request_id, slot)
​
        # 在 transfer_stream 上非阻塞加载张量，CPU 侧可同时构建 Req
        local_device = f'{current_platform.device_type}:{self.worker.local_rank}'
        tensors, load_event = self._transfer_manager.load_tensors_async(
            request_id, manifest, device=local_device, stream=self._transfer_stream,
        )
​
        # NOTE: 不要在这里释放 receive slot——异步 load 仍在进行，
        # slot 必须保持有效直到主线程等待 load_event；释放发生在 _disagg_prefetch_event_loop。
​
        # CPU 工作与 GPU load 重叠：先构建 Req 再交给主线程 compute
        req = self._build_disagg_req(scalar_fields, tensors)
​
        # NOTE: 不要在这里调用 scheduler_mod.set_timesteps()！
        # 该调用修改共享 scheduler 状态（self.sigmas），会破坏主线程正在运行的
        # denoising loop，应推迟到主线程的 _disagg_prefetch_event_loop 中执行。
​
        # 返回四元组：role_name 与 scalar_fields 已无消费者，不再向下游传递
        return req, load_event, request_id, prealloc_slot_id

python/sglang/multimodal_gen/runtime/disaggregation/transport/protocol.py

TransferStagedMsg 成为 staged 消息唯一 schema，新增 scalar_fields 并移除 __post_init__，是本次 schema 收敛的落点。

# SPDX-License-Identifier: Apache-2.0
"""Transfer protocol messages for disaggregated diffusion.
​
All messages are sent as ZMQ multipart with a b'__transfer__' discriminator
in frame[0] and JSON payload in frame[1].
"""
​
import json
from dataclasses import asdict, dataclass, field
from typing import Any
​
TRANSFER_MAGIC = b'__transfer__'
​
​
class TransferMsgType:
    # Instance -> DiffusionServer
    STAGED = 'transfer_staged'
    ALLOCATED = 'transfer_allocated'
    PUSHED = 'transfer_pushed'
    DONE = 'transfer_done'
​
    # DiffusionServer -> Instance
    ALLOC = 'transfer_alloc'
    PUSH = 'transfer_push'
    READY = 'transfer_ready'
​
    # Registration
    REGISTER = 'transfer_register'
    REGISTER_ACK = 'transfer_register_ack'
​
​
@dataclass
class TransferStagedMsg:
    # 使用 field(default_factory=dict) 而不是可变默认值 None，
    # 避免共享同一 dict 实例，也省掉了 __post_init__ 的判空补丁
    msg_type: str = TransferMsgType.STAGED
    request_id: str = ''
    data_size: int = 0
    manifest: dict = field(default_factory=dict)
    session_id: str = ''
    pool_ptr: int = 0
    slot_offset: int = 0
    # scalar_fields 承载 encoder 阶段消息的标量载荷（如 seed、采样参数），
    # 现在由 protocol schema 统一接管，调用方不再手工拼 dict
    scalar_fields: dict = field(default_factory=dict)

评论区精华

本 PR 没有任何 review 评论。唯一交互是作者在 issue 评论中两次触发 /tag-and-rerun-ci 重跑 CI，说明关注点主要在回归验证而非设计讨论。最终基础 PR Test（Run #31085927775）通过，extra CI（Run #31085927225）失败但未阻止合并；由于没有同行评审意见，TransferStagedMsg schema 化与 API 删除的取舍只能从 PR body 推断。

• CI 重跑与 extra CI 失败 (other): 无设计争议；extra CI 失败原因未公开说明，PR 仍被合并。

风险与影响

• 风险：
  • manager.py 等删除的同步 API 虽在仓库内无调用方，但对下游分支或插件属于可见接口，若被直接引用会出现 ImportError / AttributeError。
  • _prefetch_transfer_ready 返回值是跨线程隐式契约，与旧版主线程逻辑混跑会导致解包失败，存在部署/回滚期间的兼容风险。
  • TransferStagedMsg 新增 scalar_fields 字段改变 dataclass 结构；虽然 sender 此前已手工发送该字段，但任何直接反序列化旧 schema 的外部调试工具可能不兼容。
  • 无本地验证且无新增测试，基础 CI 通过但 extra CI 失败原因未说明，存在未被发现的回归可能。
  • 改动集中在 disaggregation 传输热路径（prefetch 线程、RDMA 消息发送），异步事件处理若有疏忽可能影响 pipeline 并行下的请求调度。
  • 影响：对用户无感知，无公开 API 变更；对系统而言显著降低 transport 层维护面，消除重复 schema 组装，使 RDMA stage/load 与 prefetch 的所有权更易推理，同时减少 _BASE_SP_DEFAULTS 这类缓存的启动开销（幅度很小）。对团队而言，这为后续 diffusion disaggregation 的 prefetch/alloc 演进（如 capture-safe 传输、DP 扩展）提供了更干净的地基，但缺少测试覆盖是主要薄弱点。
  • 风险标记：无本地验证, 缺少测试覆盖, extra CI 失败, 删除传输层 API 存在外部调用风险, prefetch 线程返回值变更

关联脉络

• PR #33845 [diffusion] centralize entrypoint API hygiene: 同属 diffusion runtime 的 hygiene 重构系列，目标都是集中重复逻辑、消减维护面。
• PR #33843 [diffusion] consolidate pipeline core hygiene: 与本 PR 一脉相承的 diffusion 卫生清理，涉及多个模块的重复逻辑收敛。
• PR #33775 [diffusion] feat: capture-safe pynccl all-to-all: 同为 diffusion disaggregation/transport 基础能力；该 PR 为 RDMA 异步路径铺路，本 PR 清理传输层让 staging 与 prefetch 所有权更可推理。