执行摘要

• 一句话：新增 DWDP 并行策略，消除 MoE prefill 跨 rank 同步
• 推荐动作：此 PR 实现了一个重要的分布式并行策略，设计精巧（CUDA VMM 复合 VA、double-buffered prefetch），性能收益突出。建议精读 vmm.py、layout.py 和 weight_manager.py 的设计。但需注意非 CUDA 平台的兼容性问题，建议在后续 PR 中优先修复 import 延迟。综合评分：重要度 8（重大变更），洞察价值 7。

功能与动机

当前 MoE prefill 使用 expert parallelism（EP）需要 all-to-all token 分发和 allgather 同步，带来显著通信开销。DWDP 通过预取 remote expert 权重并在本地计算，避免跨 rank 同步，从而提升 prefill 吞吐。PR 基于 DWDP 论文（arXiv 2604.01621）实现。

实现拆解

1. CUDA VMM 基础层（vmm.py）：封装 CUDA VMM API，提供 shareable handle 创建（自动探测 FABRIC/POSIX）、VA 保留/映射、DLPack 张量视图，实现零拷贝跨 rank 内存共享。
2. 布局与所有权（layout.py）：定义 DwdpExpertLayout 计算每个 rank 负责的本地 expert 区间和 peer 预取范围；WeightSpec 描述每层权重形状；PageAlignedLayout 计算页对齐复合 VA 的三段式布局（预填充 + fabric handle + 后填充）。
3. 传输层（transport.py）：DWDPTransport 的 create 方法在初始化时执行 Phase 1：将本地 expert 权重复制到 fabric handle 中，释放原始存储；Phase 2：通过 pidfd（x86_64）或 fabric（aarch64）交换 handle，导入 peer 视图。
4. 权重缓冲与管理（weight_buffer.py + weight_manager.py）：WeightBuffer 为每层权重构建复合 VA，组合本地 pool 页面和 fabric handle；DWDPWeightManager 实现 double-buffered 预取：在 layer l 计算时异步预取 layer l+1 的 remote expert 权重，利用 CUDA stream 和 event 同步。
5. 编排与模型适配（dwdp_manager.py + FusedMoE 等）：DwdpManager.setup 收集 MoE 层，验证 expert 数可被 dwdp_size 整除，创建传输和缓冲，并替换原始 expert 权重为复合 VA 张量。同时修改 fused_moe_triton/layer.py 提供 replace_expert_tensor 等接口，修改 mimo_v2.py 和 gpt_oss.py 添加 forward_dwdp 路径。自动启用 dp_attention、dp_lm_head、SCHEDULER_SKIP_ALL_GATHER。
6. 测试与集成：添加两个端到端测试 test_disaggregation_dwdp_mimo.py 和 test_disaggregation_dwdp_gpt_oss.py，在 PD disaggregation 模式下验证 GSM8K 精度。提供 --dwdp-size 命令行参数。

关键文件：

• python/sglang/srt/layers/moe/dwdp/vmm.py（模块 VMM 层；类别 source；类型 dependency-wiring；符号 align_up, align_down, _make_prop, shareable_handle_types）: CUDA VMM 基础抽象，提供 handle 创建、VA 管理、DLPack 视图。核心依赖，自动探测 FABRIC/POSIX。
• python/sglang/srt/layers/moe/dwdp/layout.py（模块 布局；类别 source；类型 dependency-wiring；符号 DwdpExpertLayout, init, WeightSpec, expert_bytes）: 定义 DWDP 布局和所有权计算，核心数据模型。
• python/sglang/srt/layers/moe/dwdp/transport.py（模块 传输层；类别 source；类型 dependency-wiring；符号 _close_fds, _copy_local_weights_to_handles, DWDPTransport, init）: 实现权重跨 rank 共享 handle 交换和 peer 视图导入。
• python/sglang/srt/layers/moe/dwdp/weight_buffer.py（模块 权重缓冲；类别 source；类型 core-logic；符号 WeightBuffer, init, create, _setup_layer）: 构建复合 VA 视图，组合本地 pool 和 fabric handle，提供连续 full-expert 张量。
• python/sglang/srt/layers/moe/dwdp/weight_manager.py（模块 权重管理；类别 source；类型 core-logic；符号 DWDPWeightManager, init, weight_buffer, is_moe_layer）: 实现 double-buffered 预取逻辑，在计算当前层时异步预取下一层远程 expert。
• python/sglang/srt/layers/moe/dwdp/dwdp_manager.py（模块 编排器；类别 source；类型 core-logic；符号 DwdpManager, init, setup, prefetch_first_layers）: 全局编排器，组装各组件并集成到模型 forward 流程。
• python/sglang/srt/layers/moe/fused_moe_triton/layer.py（模块 MoE 层；类别 source；类型 core-logic；符号 num_global_routed_experts, bind_full_expert_weights, named_per_expert_tensors, replace_expert_tensor）: MoE 层适配，添加 replace_expert_tensor 等接口支撑 DWDP 权重替换。
• test/registered/disaggregation/test_disaggregation_dwdp_mimo.py（模块 测试；类别 test；类型 test-coverage；符号 TestDisaggregationDWDPMiMo, setUpClass, start_prefill, start_decode）: 测试 DWDP 在 MiMo 模型上的 PD disaggregation 端到端功能。

关键符号：shareable_handle_types, _copy_local_weights_to_handles, DWDPTransport.create, WeightBuffer.create, DwdpManager.setup, DWDPWeightManager.prefetch_layer, DwdpExpertLayout.init, PageAlignedLayout.compute, FusedMoE.replace_expert_tensor

评论区精华

核心讨论集中在 vmm.py 中 from cuda.bindings import driver as cuda 会破坏非 CUDA 设备（如 ROCm、NPU）的服务启动。Reviewer silencejade 建议仅在功能启用时才导入该模块。另一位 reviewer OrangeRedeng 表示已准备修复 PR 解决此问题。当前状态：评论已记录，但 PR 合并时未包含该修复，后续需跟进。

• cuda.bindings import 在非 CUDA 平台导致崩溃 (other): OrangeRedeng 表示已准备修复 PR。但当前 PR 未包含该修复。

风险与影响

• 风险：此变更存在以下风险：1）非 CUDA 设备兼容性风险：vmm.py 直接导入 cuda.bindings，在 AMD ROCm、Intel XPU、NPU 等平台会立即失败，需按 review 建议做延迟导入。2）CUDA VMM 平台依赖：shareable_handle_types 自动探测依赖 GPU 驱动和 NVLink fabric，在虚拟化环境或旧驱动上可能回退到 POSIX fd，但 fabric 不可用时可能导致异常。3）MoE 模型适配局限：目前仅显式适配 mimo_v2.py 和 gpt_oss.py，其他 MoE 模型（如 DeepSeek、LLaDA）需额外修改。4）内存占用增加：double-buffered 预取需要额外 GPU 内存保存两层权重，可能在高并发时加剧 OOM。
• 影响：
  1) 用户影响：需要指定 --dwdp-size 启用，无默认值，需配合 --tp 使用。PD disaggregation 模式推荐使用。性能提升显著（1.3x-1.9x）。
  2) 系统影响：引入了复杂的 CUDA VMM 层，增加了维护成本和故障排查难度。
  3) 团队影响：核心代码由一人提交，需要更多 reviewer 理解 VMM 细节。
  4) 兼容性：非 CUDA 平台（AMD/NPU/XPU）将因 import 错误无法运行 sglang，需紧急修复。
  • 风险标记：非 CUDA 平台兼容性, CUDA VMM 系统级依赖, MoE 模型适配有限, double-buffered 内存压力

关联脉络

• 暂无明显关联 PR