From 02c65cfbcda1b3e4fd782e20d996f595134142a3 Mon Sep 17 00:00:00 2001 From: Binyang Li Date: Tue, 14 Jul 2026 16:33:25 -0700 Subject: [PATCH] Support hidden size 6656 in low-latency EP (#835) Instantiate low-latency dispatch and combine kernels for hidden size 6656 and expose the shape through the functional and benchmark entry points. Reuse syncNamedBarrier for scheduler named barriers instead of inline PTX. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/ext/ep/README.md | 8 ++++---- src/ext/ep/low_latency/combine.cu | 5 ++++- src/ext/ep/low_latency/dispatch.cu | 15 ++++++++------- test/python/ep/ep_bench_ll.py | 6 +++--- test/python/ep/mscclpp_ep_bench.cu | 4 ++-- test/python/ep/run_ep_bench.py | 2 +- test/python/ep/test_low_latency_multirank.py | 2 +- 7 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/ext/ep/README.md b/src/ext/ep/README.md index 6feccb35..3a4fd1b9 100644 --- a/src/ext/ep/README.md +++ b/src/ext/ep/README.md @@ -529,10 +529,10 @@ has completed, the sender publishes readiness through a lightweight `BaseMemoryChannel` signal; the receiver waits on that signal before copying data into expert-major output. -The optimized kernels are instantiated for hidden sizes `4096`, `7168`, -`8192`, and `9216`; other hidden sizes are rejected. FP8 E4M3 currently fixes -the scale block at 128. A future scale layout must use a distinct -`DispatchDataType`. +The optimized kernels are instantiated for hidden sizes `4096`, `6656`, +`7168`, `8192`, and `9216`; other hidden sizes are rejected. FP8 E4M3 +currently fixes the scale block at 128. A future scale layout must use a +distinct `DispatchDataType`. ### Current H100 performance diff --git a/src/ext/ep/low_latency/combine.cu b/src/ext/ep/low_latency/combine.cu index a9dd8794..450f545c 100644 --- a/src/ext/ep/low_latency/combine.cu +++ b/src/ext/ep/low_latency/combine.cu @@ -403,7 +403,7 @@ inline void combineHiddenMode(void* output, const void* expertOutput, const int6 const low_latency::Workload& workload, void* recvBuffer, void* dispatchRecvBuffer, const low_latency::CommContext& comm, void* workspace, int numBlocks, cudaStream_t stream) { - static_assert(Hidden == 4096 || Hidden == 7168 || Hidden == 8192 || Hidden == 9216); + static_assert(Hidden == 4096 || Hidden == 6656 || Hidden == 7168 || Hidden == 8192 || Hidden == 9216); const int nExperts = workload.numExperts_; const int nRanks = comm.numRanks_; const int nLocalExperts = nExperts / nRanks; @@ -492,6 +492,9 @@ inline void combine(void* output, const void* expertOutput, const int64_t* topkI case 4096: return combineHidden<4096>(output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, workload, recvBuffer, dispatchRecvBuffer, comm, workspace, numBlocks, mode, stream); + case 6656: + return combineHidden<6656>(output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, workload, + recvBuffer, dispatchRecvBuffer, comm, workspace, numBlocks, mode, stream); case 7168: return combineHidden<7168>(output, expertOutput, topkIndices, topkWeights, srcInfo, layoutRange, workload, recvBuffer, dispatchRecvBuffer, comm, workspace, numBlocks, mode, stream); diff --git a/src/ext/ep/low_latency/dispatch.cu b/src/ext/ep/low_latency/dispatch.cu index e1968298..5d0c6e05 100644 --- a/src/ext/ep/low_latency/dispatch.cu +++ b/src/ext/ep/low_latency/dispatch.cu @@ -304,7 +304,7 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvScheduler(int64_t* outputLayout, int* out sharedMem[warpId] = rankTokenPrefix; sharedMem[nRankWarps + warpId] = activeRankPrefix; } - asm volatile("bar.sync %0, %1;" ::"r"(DispatchSchedulerPrefixBarrier), "r"(nRankWarps * WARP_SIZE) : "memory"); + syncNamedBarrier(DispatchSchedulerPrefixBarrier, nRankWarps * WARP_SIZE); if (warpId == 0) { const int tokenTotal = laneId < nRankWarps ? sharedMem[laneId] : 0; @@ -320,7 +320,7 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvScheduler(int64_t* outputLayout, int* out sharedMem[2 * nRankWarps + 1] = activePrefix; } } - asm volatile("bar.sync %0, %1;" ::"r"(DispatchSchedulerPrefixBarrier), "r"(nRankWarps * WARP_SIZE) : "memory"); + syncNamedBarrier(DispatchSchedulerPrefixBarrier, nRankWarps * WARP_SIZE); rankTokenPrefix += sharedMem[warpId]; activeRankPrefix += sharedMem[nRankWarps + warpId]; @@ -351,8 +351,7 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvScheduler(int64_t* outputLayout, int* out } if (threadId == 0) *workspaceView.dispatchNumRecvTasks_ = nTasks; - asm volatile("bar.sync %0, %1;" ::"r"(DispatchSchedulerReadyBarrier), "r"((nRankWarps + nLayoutWarps) * WARP_SIZE) - : "memory"); + syncNamedBarrier(DispatchSchedulerReadyBarrier, (nRankWarps + nLayoutWarps) * WARP_SIZE); if (threadId == 0) { mscclpp::atomicStore(workspaceView.dispatchTasksReadyEpoch_, dispatchEpoch, mscclpp::memoryOrderRelease); @@ -381,8 +380,7 @@ MSCCLPP_DEVICE_INLINE void dispatchRecvScheduler(int64_t* outputLayout, int* out } outputCount[localExpertIdx] = outputOffset; } - asm volatile("bar.sync %0, %1;" ::"r"(DispatchSchedulerReadyBarrier), "r"((nRankWarps + nLayoutWarps) * WARP_SIZE) - : "memory"); + syncNamedBarrier(DispatchSchedulerReadyBarrier, (nRankWarps + nLayoutWarps) * WARP_SIZE); } } @@ -554,7 +552,7 @@ inline void dispatchHiddenMode(void* output, float* outputScales, int* outputSrc const low_latency::Workload& workload, void* recvBuffer, const low_latency::CommContext& comm, void* workspace, int numBlocks, cudaStream_t stream) { - static_assert(Hidden == 4096 || Hidden == 7168 || Hidden == 8192 || Hidden == 9216); + static_assert(Hidden == 4096 || Hidden == 6656 || Hidden == 7168 || Hidden == 8192 || Hidden == 9216); using OutputType = DispatchElementType; constexpr int NRecvTmaWorkers = tmaWorkerCount(); static_assert(NRecvTmaWorkers > 0); @@ -629,6 +627,9 @@ inline void dispatch(void* output, float* outputScales, int* outputSrcInfo, int6 case 4096: return dispatchHidden<4096>(output, outputScales, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream); + case 6656: + return dispatchHidden<6656>(output, outputScales, outputSrcInfo, outputLayout, outputCount, input, topkIdx, + topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream); case 7168: return dispatchHidden<7168>(output, outputScales, outputSrcInfo, outputLayout, outputCount, input, topkIdx, topkWeights, workload, recvBuffer, comm, workspace, numBlocks, stream); diff --git a/test/python/ep/ep_bench_ll.py b/test/python/ep/ep_bench_ll.py index 07ff1a8f..e23e6e45 100644 --- a/test/python/ep/ep_bench_ll.py +++ b/test/python/ep/ep_bench_ll.py @@ -93,7 +93,7 @@ def parse_args() -> argparse.Namespace: "--hidden", type=int, default=int(os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "7168")), - choices=(4096, 7168, 8192, 9216), + choices=(4096, 6656, 7168, 8192, 9216), help="hidden dimension", ) p.add_argument( @@ -165,8 +165,8 @@ def parse_args() -> argparse.Namespace: ) p.add_argument("--seed", type=int, default=0xB3C4, help="per-rank RNG seed base") args = p.parse_args() - if args.hidden not in (4096, 7168, 8192, 9216): - p.error("--hidden must be one of 4096, 7168, 8192, 9216") + if args.hidden not in (4096, 6656, 7168, 8192, 9216): + p.error("--hidden must be one of 4096, 6656, 7168, 8192, 9216") if not 1 <= args.num_topk <= 9: p.error("--num-topk must be in [1, 9]") if args.num_tokens <= 0 or args.num_experts <= 0: diff --git a/test/python/ep/mscclpp_ep_bench.cu b/test/python/ep/mscclpp_ep_bench.cu index 66e8253b..f2334c19 100644 --- a/test/python/ep/mscclpp_ep_bench.cu +++ b/test/python/ep/mscclpp_ep_bench.cu @@ -205,8 +205,8 @@ int main(int argc, char** argv) { if (rank == 0) fprintf(stderr, "tokens, experts, and iters must be positive; warmup must be non-negative\n"); MPI_Abort(MPI_COMM_WORLD, 1); } - if (H != 4096 && H != 7168 && H != 8192 && H != 9216) { - if (rank == 0) fprintf(stderr, "hidden must be one of 4096, 7168, 8192, 9216\n"); + if (H != 4096 && H != 6656 && H != 7168 && H != 8192 && H != 9216) { + if (rank == 0) fprintf(stderr, "hidden must be one of 4096, 6656, 7168, 8192, 9216\n"); MPI_Abort(MPI_COMM_WORLD, 1); } if (K <= 0 || K > 9) { diff --git a/test/python/ep/run_ep_bench.py b/test/python/ep/run_ep_bench.py index 872efc0a..7ed2248c 100644 --- a/test/python/ep/run_ep_bench.py +++ b/test/python/ep/run_ep_bench.py @@ -109,7 +109,7 @@ def parse_args() -> argparse.Namespace: "--hidden", type=int, default=7168, - choices=(4096, 7168, 8192, 9216), + choices=(4096, 6656, 7168, 8192, 9216), help="hidden dimension", ) p.add_argument("-k", "--num-topk", type=int, default=8, choices=range(1, 10), help="top-k experts per token") diff --git a/test/python/ep/test_low_latency_multirank.py b/test/python/ep/test_low_latency_multirank.py index 29315802..1230eab7 100644 --- a/test/python/ep/test_low_latency_multirank.py +++ b/test/python/ep/test_low_latency_multirank.py @@ -56,7 +56,7 @@ def parse_args(): "--hidden", type=int, default=7168, - choices=(4096, 7168, 8192, 9216), + choices=(4096, 6656, 7168, 8192, 9216), help="BF16 hidden size compiled into the optimized low-latency kernels", ) parser.add_argument("--num-topk", type=int, default=8)