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)