Merge pull request #96 from songhexiang/adjust_kNumThreads_of_notify_dispatch

Adjust kNumThreads of notify_dispatch

Merge pull request #96 from songhexiang/adjust_kNumThreads_of_notify_dispatch
Adjust kNumThreads of notify_dispatch
dcf46f1c · Chenggang Zhao · GitHub · e130cc6e · 4dd1e68a · dcf46f1c
Unverified Commit dcf46f1c authored Mar 28, 2025 by Chenggang Zhao Committed by GitHub Mar 28, 2025
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

csrc/kernels/internode.cu csrc/kernels/internode.cu +1 -1

No files found.
--- a/csrc/kernels/internode.cu
+++ b/csrc/kernels/internode.cu
@@ -428,7 +428,7 @@ void notify_dispatch(const int* num_tokens_per_rank, int* moe_recv_counter_mappe
                  buffer_ptrs, task_fifo_ptrs, head, rank, \
                  cpu_rdma_team); } break

-    constexpr int kNumThreads = 256;
+    constexpr int kNumThreads = std::max(256, 32 * num_channels);
    const auto num_rdma_ranks = num_ranks / NUM_MAX_NVL_PEERS;

    // Get clean meta