microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-18 01:37:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support hidden size 6656 in low-latency EP (#835)

Browse Source 
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>
This commit is contained in:
Binyang Li
2026-07-14 16:33:25 -07:00
committed by GitHub
parent 47c3255825
commit 02c65cfbcd
7 changed files with 23 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/ext/ep/README.md
View File

@@ -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

5
src/ext/ep/low_latency/combine.cu
View File

@@ -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);

15
src/ext/ep/low_latency/dispatch.cu
View File

@@ -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<uint32_t, mscclpp::scopeDevice>(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<DataType>;
constexpr int NRecvTmaWorkers = tmaWorkerCount<Hidden, OutputType, DispatchMaxNRecvTmaWorkers>();
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);

6
test/python/ep/ep_bench_ll.py
View File

@@ -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:

4
test/python/ep/mscclpp_ep_bench.cu
View File

@@ -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) {

2
test/python/ep/run_ep_bench.py
View File

@@ -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")

2
test/python/ep/test_low_latency_multirank.py
View File

@@ -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)