Skip to content

[MoonEP] 跨节点 Combine 优化为单通信域、单次 Launch #111

Description

@chaowick

背景

当前跨节点 Combine 被拆成两个阶段:

  1. COMBINE_PUBLISH_ONLY
  2. Host Barrier
  3. COMBINE_CONSUME_ONLY

该方案需要两个通信域、两次 Kernel Launch 和多次 Host/Device 同步。

目标

跨节点 Combine 使用同一个 TileXR 通信域,在一次 Kernel Launch 内完成:

  1. 发布本 Rank 数据。
  2. 通知其他 Rank 数据就绪。
  3. 等待 Peer Ready Signal。
  4. 获取远端数据并完成 Reduce。
  5. 发布完成状态。

Native Combine 不再依赖 Host phase_barrier

实现要求

  • 使用 magic-tagged Signal 支持通信窗口多轮复用。
  • 保证 Payload 写入完成后才能发布 Ready Signal。
  • 支持超时和远端失败处理,不能无限等待。
  • 支持 Memory 和 UDMA 路径,UDMA 不可用时不能影响已有路径。
  • 保留现有 ABI Flag 的兼容性,但默认路径不再使用 Publish/Consume 拆分模式。

验收标准

  • 单节点和跨节点 Combine 均只执行一次 Kernel Launch。
  • Native 跨节点 Combine 不再调用 Host Phase Barrier。
  • 2 节点、16 Rank、16 卡、rank_per_dev=1 六阶段 correctness 通过。
  • 64 Rank 和 128 Rank 全流程无死锁、超时和 stale-magic 问题。
  • Combine 输出与 Reference Backend 一致。
  • 连续多轮运行可以安全复用通信窗口。
  • Kernel Launch 和 Host 同步开销低于当前双阶段方案。

相关代码

  • integrations/moonep_torch/tilexr_moonep/torch_api.py
  • src/moonep/combine/kernels/tilexr_moonep_combine_kernel.cpp
  • src/moonep/combine/host/combine_layout.cpp

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions