Merge some updates#674
Conversation
| if not is_cached_dispatch: | ||
| self.cached_recv_src_metadata_before_sort = self.recv_src_metadata.clone() | ||
| assert self.cached_recv_src_metadata_before_sort is not None | ||
| sort_keys = self.cached_recv_src_metadata_before_sort[:, 0] |
There was a problem hiding this comment.
🟡 warning: deterministic_sort is now correctness-critical (runs in the dispatch hot path under deterministic=True) and its expand-mode path is non-trivial: it uses a magic constant 1e10 as the high sort key, relies on src_token_global_index_max being < 5e9, and mutates recv_src_metadata[:, 2:] in place while leaving recv_src_metadata[:, :2] unchanged. The interplay between cached/non-cached and expand/non-expand (and the unaligned-vs-aligned counts) is hard to verify by inspection. Suggest adding a dedicated unit test (weave it into tests/elastic/test_ep.py or as a standalone ref impl) and documenting the 1e10 assumption.
🤖 v4
| assert self.event is not None | ||
| self.event.current_stream_wait() | ||
|
|
||
| if self.hook_after_wait is not None: |
There was a problem hiding this comment.
🟡 warning: The hook fires only on the first current_stream_wait() and is then cleared. If a caller never calls current_stream_wait() (or calls it multiple times), the deterministic sort either silently never runs or re-runs inconsistently. Combined with release_handle=True the ordering is subtle. Consider documenting that exactly one current_stream_wait() is required when async_with_compute_stream=True, or making the sort more defensive.
🤖 v4
| num_dispatch_warps = std::min<int>(std::min<int>( | ||
| (num_smem_bytes - num_notify_smem_bytes) / token_layout.get_num_bytes<true>(), 32 - num_notify_warps), | ||
| math::ceil_div(512, num_sms)); | ||
| num_dispatch_warps = std::min<int>( |
There was a problem hiding this comment.
🔵 suggestion: The math::ceil_div(512, num_sms) cap on num_dispatch_warps was removed (the 'Too many warps may cause performance degrade' comment was dropped too). This lets dispatch use more warps on high-SM GPUs but can increase register pressure. Worth a benchmark across the workload range.
🤖 v4
| int kNumChannels = kNumScaleoutWarps * kNumSMs, | ||
| int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, kNumChannels), | ||
| int kScaleoutUpdateInterval = 3, | ||
| int kScaleoutUpdateInterval = 6, |
There was a problem hiding this comment.
🔵 suggestion: kScaleoutUpdateInterval changed from 3 to 6. This is a performance tuning that alters how often scale-out updates are issued; please confirm the before/after throughput justifies it on your target configurations.
🤖 v4
| num_bytes_per_entry = align(hidden * dtype.itemsize + num_sf_packs * 4, 32) | ||
| num_bytes_per_entry = align(hidden * dtype.itemsize, 32) | ||
| num_gpu_bytes = align(num_bytes_per_entry * num_max_tokens_per_rank, buffer_alignment) | ||
| num_cpu_bytes = align(num_bytes_per_entry * num_entries, buffer_alignment) |
There was a problem hiding this comment.
🔵 suggestion: For FP8 engram, scaling factors are now fully replicated across ranks rather than transported via RDMA, which makes the per-entry size no longer reserve SF space (correct). Note however that the replicated SF table is num_ranks * num_entries * ceil(hidden/32) * 4 bytes — ~1.25x the FP8 data — and lives for the whole window lifetime, so the per-rank memory cost can grow with num_ranks.
🤖 v4
| ptx::tma_store_commit(); | ||
|
|
||
| // Prefetch: wait until this stage's buffer is safe to reuse, then issue next load | ||
| // Prefetch: wait until this stage's store is completed, then issue next load |
There was a problem hiding this comment.
🔵 suggestion: The TMA copy prefetch now fully drains the store pipeline before re-issuing the next load (tma_store_wait() instead of the old tma_store_wait()). This correctly avoids a WAR hazard when recycling the stage buffer (the intended data-race fix), but it serializes store->load, so prefetch overlap is reduced. Please verify PP throughput didn't regress.
🤖 v4
| self.nccl_comm_handle = get_nccl_comm_handle(group) | ||
| self.deterministic = deterministic | ||
|
|
||
| if os.environ.get('NCCL_GIN_CROSS_NIC') == '0': |
There was a problem hiding this comment.
🟡 warning: 行为变化需回归确认：原先 csrc/kernels/backend/symmetric.hpp 中在 hybrid 模式下会无条件 setenv('NCCL_SYM_REUSE_SYSMEM_HANDLES','1',overwrite=0)，现改为仅当 os.environ.get('NCCL_GIN_CROSS_NIC') == '0' 时在 Python 侧 setdefault。提交信息标注为 'Fix Set Env Bug'（应是修复原 setenv 相对 comm 创建时机过晚的问题），属有意为之。建议确认：非 NCCL_GIN_CROSS_NIC=0 的 hybrid（多平面）场景是否仍依赖该变量，如是则需针对该场景做回归验证。
🤖 v3
| self.cached_recv_src_metadata_before_sort = self.recv_src_metadata.clone() | ||
| assert self.cached_recv_src_metadata_before_sort is not None | ||
| sort_keys = self.cached_recv_src_metadata_before_sort[:, 0] | ||
|
|
There was a problem hiding this comment.
🔵 suggestion: 可读性建议：deterministic_sort 中 num_recv_tokens 在 not do_cpu_sync 时为 0-dim tensor（psum_num_recv_tokens_per_scaleup_rank[-1]），在 do_cpu_sync 时为 Python int。后续用于 tensor 比较与 torch.arange 比较均可正常工作，但类型混用降低可读性，建议统一 .item() 或补充注释说明。非缺陷。
🤖 v3
| const int expert_idx = i * 32 + lane_idx; | ||
| num_experts_per_lane[i] = expert_idx < kNumExpertsPerRank ? | ||
| num_unaligned_recv_tokens_per_expert[expert_idx] : 0; | ||
| } |
There was a problem hiding this comment.
🔵 suggestion: 潜在编译器告警：zero-padding 循环中 wave_num_experts_per_lane 未初始化即通过 unroll 循环条件赋值（依赖 i == wave_idx 必命中，由外层 while 的 wave_idx < kNumExpertsPerLane 保证）。逻辑正确，但编译器可能告 'may be used uninitialized'。建议初始化为 0 以消除告警。
🤖 v3
|
|
||
| // TODO: once NCCL supports ncclCoopWarp gin.get, drop the elect_one_sync and let the whole warp | ||
| // gather SF packs in parallel. | ||
| if constexpr (kNumSFPacks > 0) { |
There was a problem hiding this comment.
🔵 suggestion: 格式问题：issue RDMA get 后新增代码引入了尾随空格（空行含空格）。建议运行 format.sh 清理。此外 elastic.py（self.deterministic = deterministic 后、若干处）与 test_ep.py（lambda 后）也存在尾随空格，一并清理。
🤖 v3
🤖 ds-review-bot Code Reviewv4This review covers the cumulative changes between origin/main and HEAD (PRs ~#165-#177, titled 'Merge some updates', ~26 files). The headline changes are: (1) FP8 Engram — engram_write/fetch now support torch.float8_e4m3fn with globally-replicated scaling factors (SF moved off the RDMA path); (2) deterministic dispatch refactored from a dedicated C++ kernel into a Python EPHandle.deterministic_sort() post-hook registered via EventOverlap; (3) cached-mode dispatch now supports expanding layout plus a TMA-based do_zero_padding pass that clears expert-alignment gaps, tracked via a new num_unaligned_recv_tokens_per_expert written by the notify warps; (4) combine expand backward — kDoExpandedSend now reads 1D topk_weights by slot when allow_multiple_reduction. Supporting changes include NCCL_WIN_STRICT_ORDERING, a sequential barrier mode, GIN flush-depth mechanism, PP TMA pipeline rewrite, removal of the 512-warp cap, and kScaleoutUpdateInterval 3->6. v3这是一次从 af9a040 到 8d8b70f 的合并（26 个文件），囊括多个功能与修复提交，整体质量高，接口改动前后一致，测试覆盖充分。主要变更：(1) FP8 Engram 支持：engram_write/engram_fetch 支持 torch.float8_e4m3fn 存储 + 全局复制的 scaling factor 本地聚集，索引从 [num_tokens] 扩展为 [num_tokens, num_entries_per_token]，支持 TMA 对齐列主序 SF 布局。(2) 重构确定性 dispatch：删除 GPU 端 dispatch_deterministic_prologue.cuh（143 行）及相关 C++ 启动逻辑，改为 Python 侧 EPHandle.deterministic_sort 后处理排序（通过 event.register_hook_after_wait 挂钩），移除 deterministic 构造参数并简化内核路径。(3) Zero padding + cached expand 模式：dispatch 支持 do_expand 的 cached 模式并新增 do_zero_padding 清零专家对齐间隙，新增 num_unaligned_recv_tokens_per_expert。(4) 性能/正确性调优：NCCL MR 改用 NCCL_WIN_STRICT_ORDERING；本地 rank/TagRail 原子改用 GPU-scoped release；Engram fetch 周期性 flush 避免 hang；PP TMA copy 修复数据竞争；新增 sequential barrier。一致性核验全部通过：构造函数与 pybind init 均为 14 参数，dispatch 返回元组三处（C++/Python/_unpack_handle）对应更新，red_add_rel_gpu 与 red_add_rel_sys 重载集一致，已删除文件无残留引用，zero-padding SF 布局与主循环一致，combine 内核模板参数定义完整。结论：逻辑自洽，无阻塞性缺陷，建议合并；合并前建议确认多平面场景回归并清理尾随空格。 Files reviewed: 25 |
Merge 8 upstream commits including: - fence.proxy.async.shared::cta for TMA load (deepseek-ai#642) - Fix single-node intra EP init (deepseek-ai#630) - Fix internode dispatch args (deepseek-ai#641) - Symmetric memory refactor with CPU buffer support - NCCL SO name resolution fix for pip wheels (deepseek-ai#627) - Various other updates (deepseek-ai#637, deepseek-ai#674) Conflicts resolved: - buffer.hpp: merged cpu_comm/cpu_buffer params with has_nvlink param - nccl.cu: merged symmetric alloc refactor with has_nvlink and debug prints - elastic.py: merged NCCL comm handle refactor with NVLink detection logic - setup.py: adopted upstream dynamic NCCL lib name resolution Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
3 participants
Bug Fixes
Features
do_expandfor cache mode.