diff --git a/python/mscclpp/ep/README.md b/python/mscclpp/ep/README.md index d6276e9f..4ea79cb2 100644 --- a/python/mscclpp/ep/README.md +++ b/python/mscclpp/ep/README.md @@ -152,7 +152,7 @@ specialized advanced path. The active implementation supports `mode=MoEMode.LOW_LATENCY` and `mode=MoEMode.HIGH_THROUGHPUT`. `mode` must be a `MoEMode` enum value, not a string. LL uses an expert-major output layout. HT uses a flat output layout and -supports only 2, 4, or 8 ranks within one detected GPU IPC/NVL fabric domain; +supports 2, 4, 8, or 16 ranks within one detected GPU IPC/NVL fabric domain; that domain may span multiple hosts. ```python diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 4a9f75de..d76c1dd4 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -13,7 +13,7 @@ It builds two backends: | Feature | Status | |---|---| | LL dispatch/combine | Validated on Hopper and newer GPUs | -| HT dispatch/combine | Supports 2, 4, or 8 ranks in one GPU IPC/NVL fabric domain | +| HT dispatch/combine | Supports 2, 4, 8, or 16 ranks in one GPU IPC/NVL fabric domain | | HT RDMA/IB fallback | Not supported | | Python frontend | `mscclpp.ep.MoECommunicator` selects LL or HT with `MoEMode` | | ROCm | Not supported | diff --git a/src/ext/ep/ht/config.hpp b/src/ext/ep/ht/config.hpp index 59391bb8..3819ff37 100644 --- a/src/ext/ep/ht/config.hpp +++ b/src/ext/ep/ht/config.hpp @@ -35,7 +35,7 @@ struct Config { size_t get_nvl_buffer_size_hint(size_t hidden_bytes, int num_ranks) const { EP_HOST_ASSERT(hidden_bytes > 0); - EP_HOST_ASSERT(num_ranks == 2 or num_ranks == 4 or num_ranks == 8); + EP_HOST_ASSERT(num_ranks == 2 or num_ranks == 4 or num_ranks == 8 or num_ranks == 16); const size_t ranks = static_cast(num_ranks); const size_t recv_tokens = static_cast(num_max_nvl_chunked_recv_tokens); diff --git a/src/ext/ep/ht/intranode_kernel.cu b/src/ext/ep/ht/intranode_kernel.cu index ffcfdb33..79eafc29 100644 --- a/src/ext/ep/ht/intranode_kernel.cu +++ b/src/ext/ep/ht/intranode_kernel.cu @@ -754,6 +754,7 @@ __global__ void __launch_bounds__(kNumThreads, 1) const auto send_lane_id = send_thread_id % 32; const auto send_rank_id = thread_id / num_threads_per_rank; const auto send_warp_id_in_rank = send_thread_id % num_threads_per_rank / 32; + if (send_rank_id >= kNumRanks) return; // Calculate pointers by the specific layout auto ptr = reinterpret_cast(reinterpret_cast(buffer_ptrs[send_rank_id])); @@ -1012,14 +1013,23 @@ __global__ void __launch_bounds__(kNumThreads, 1) #ifndef EP_ICMB_TMA_WARPS #define EP_ICMB_TMA_WARPS 16 // token-parallel warps per block #endif +#ifndef EP_ICMB_TMA_WARPS_WIDE +#define EP_ICMB_TMA_WARPS_WIDE 14 // low block count: more token-parallelism +#endif +#ifndef EP_ICMB_TMA_WARPS_NARROW +#define EP_ICMB_TMA_WARPS_NARROW 12 // high block count: less scheduling overhead +#endif +#ifndef EP_ICMB_TMA_WARPS_MAX_BLOCKS +#define EP_ICMB_TMA_WARPS_MAX_BLOCKS 24 +#endif -template +template __global__ void __launch_bounds__(kWarps * 32, 1) combine_intranode_gather_tma(int4* combined_x, float* combined_topk_weights, const int* send_head, int num_combined_tokens, int hidden, int num_topk, int num_ranks, void** recv_pool_ptrs, const int* ep_combine_recv_idx, int64_t recv_pool_header_bytes) { - constexpr int kMaxContrib = kNumRanks; + static_assert(kMaxContrib <= kNumRanks); constexpr int kChunkInt4 = EP_ICMB_TMA_CHUNK_INT4; constexpr int kStages = EP_ICMB_TMA_STAGES; constexpr int kChunkBytes = kChunkInt4 * static_cast(sizeof(int4)); @@ -1166,14 +1176,14 @@ bool combine_tma(cudaDataType_t type, void* combined_x, float* combined_topk_wei constexpr int kStages = EP_ICMB_TMA_STAGES; constexpr int kChunkInt4 = EP_ICMB_TMA_CHUNK_INT4; const int num_blocks = std::max(1, combine_sms); + const bool use_wide_kernel = num_blocks <= EP_ICMB_TMA_WARPS_MAX_BLOCKS; - // SMEM/block = kWarps*kStages*kMaxContrib(=ranks)*kChunkBytes + mbars. kChunkBytes=1KB, - // kStages=2. Keep it under the GB200 ~227KB opt-in cap: 16 warps for <=4 ranks (128KB), - // 12 warps for 8 ranks (192KB). -#define COMBINE_INTRANODE_TMA_LAUNCH(ranks, WARPS) \ + // Rank discovery still scans every GPU rank. Staging capacity only needs to + // cover distinct contributors, which cannot exceed top-k. +#define COMBINE_INTRANODE_TMA_LAUNCH(ranks, MAX_CONTRIB, WARPS) \ { \ - auto tma_func = combine_intranode_gather_tma; \ - const size_t tma_smem = static_cast(WARPS) * kStages * (ranks) * kChunkInt4 * sizeof(int4) + \ + auto tma_func = combine_intranode_gather_tma; \ + const size_t tma_smem = static_cast(WARPS) * kStages * (MAX_CONTRIB) * kChunkInt4 * sizeof(int4) + \ static_cast(WARPS) * kStages * sizeof(uint64_t); \ CUDA_CHECK( \ cudaFuncSetAttribute(tma_func, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(tma_smem))); \ @@ -1181,16 +1191,31 @@ bool combine_tma(cudaDataType_t type, void* combined_x, float* combined_topk_wei static_cast(num_blocks), static_cast((WARPS) * 32), tma_smem, stream, nullptr, 0}; \ LAUNCH_KERNEL(&cfg, tma_func, reinterpret_cast(combined_x), combined_topk_weights, send_head, num_tokens, \ hidden, num_topk, num_ranks, recv_pool_ptrs, ep_combine_recv_idx, recv_pool_header_bytes); \ - } \ - break + } switch (num_ranks) { case 2: - COMBINE_INTRANODE_TMA_LAUNCH(2, EP_ICMB_TMA_WARPS); + COMBINE_INTRANODE_TMA_LAUNCH(2, 2, EP_ICMB_TMA_WARPS); + break; case 4: - COMBINE_INTRANODE_TMA_LAUNCH(4, EP_ICMB_TMA_WARPS); + COMBINE_INTRANODE_TMA_LAUNCH(4, 4, EP_ICMB_TMA_WARPS); + break; case 8: - COMBINE_INTRANODE_TMA_LAUNCH(8, 12); + if (use_wide_kernel) + COMBINE_INTRANODE_TMA_LAUNCH(8, 8, EP_ICMB_TMA_WARPS_WIDE) + else + COMBINE_INTRANODE_TMA_LAUNCH(8, 8, EP_ICMB_TMA_WARPS_NARROW) + break; + case 16: + if (num_topk <= 8) { + if (use_wide_kernel) + COMBINE_INTRANODE_TMA_LAUNCH(16, 8, EP_ICMB_TMA_WARPS_WIDE) + else + COMBINE_INTRANODE_TMA_LAUNCH(16, 8, EP_ICMB_TMA_WARPS_NARROW) + } else { + COMBINE_INTRANODE_TMA_LAUNCH(16, 16, 7); + } + break; default: EP_HOST_ASSERT(false and "Unsupported ranks"); } diff --git a/src/ext/ep/ht_runtime.cc b/src/ext/ep/ht_runtime.cc index b427d402..fc8d2ff6 100644 --- a/src/ext/ep/ht_runtime.cc +++ b/src/ext/ep/ht_runtime.cc @@ -34,7 +34,8 @@ MoEHighThroughputRuntime::MoEHighThroughputRuntime(mscclpp::Communicator& commun EP_HOST_ASSERT(numNvlRanks_ > 0); EP_HOST_ASSERT(maxHiddenBytes_ > 0); - if ((numRanks_ != 2 and numRanks_ != 4 and numRanks_ != 8) or numRanksPerIpcDomain_ < numRanks_) return; + if ((numRanks_ != 2 and numRanks_ != 4 and numRanks_ != 8 and numRanks_ != 16) or numRanksPerIpcDomain_ < numRanks_) + return; ringBufferBytes_ = config_.get_nvl_buffer_size_hint(static_cast(maxHiddenBytes_), numRanks_); EP_HOST_ASSERT(ringBufferBytes_ <= static_cast(std::numeric_limits::max())); diff --git a/src/ext/ep/include/launch.cuh b/src/ext/ep/include/launch.cuh index cbae659f..d42f887d 100644 --- a/src/ext/ep/include/launch.cuh +++ b/src/ext/ep/include/launch.cuh @@ -18,6 +18,8 @@ #define LAUNCH_KERNEL(config, kernel, ...) CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__)) #endif +// HT uses the rank index as a named-barrier ID and dispatch assigns one warp +// per rank, so 16 is the architectural maximum for this launch family. #define SWITCH_RANKS(case_macro) \ do { \ switch (num_ranks) { \ @@ -27,6 +29,8 @@ case_macro(4); \ case 8: \ case_macro(8); \ + case 16: \ + case_macro(16); \ default: \ EP_HOST_ASSERT(false and "Unsupported ranks"); \ } \ @@ -41,6 +45,8 @@ case_macro(dtype, 4); \ case 8: \ case_macro(dtype, 8); \ + case 16: \ + case_macro(dtype, 16); \ default: \ EP_HOST_ASSERT(false && "Unsupported ranks"); \ } \