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

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:

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

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

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)