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

Adjust token major layout (#844)

Browse Source 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Binyang Li
2026-07-16 18:56:10 -07:00
committed by GitHub
parent 39813bfff1
commit dd9fa728ad
15 changed files with 190 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -142,6 +142,11 @@ def parse_args() -> argparse.Namespace:
default="expert_major",
help="low-latency dispatch output layout",
)
p.add_argument(
"--token-major-init-padding",
action="store_true",
help="initialize unused token-major top-k IDs and weights for fixed-capacity kernels",
)
p.add_argument("--num-blocks", type=int, default=130, help="total low-latency dispatch blocks")
p.add_argument(
"--no-kernel-timing",
@@ -396,6 +401,7 @@ def main() -> None:
low_latency_num_blocks=args.num_blocks,
low_latency_combine_mode=combine_mode,
output_layout=output_layout,
token_major_init_padding=args.token_major_init_padding,
quant=dispatch_quant,
)
assert moe_comm.is_available()

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

@@ -134,6 +134,11 @@ def parse_args() -> argparse.Namespace:
default="expert_major",
help="MSCCL++ Python low-latency output layout",
)
p.add_argument(
"--token-major-init-padding",
action="store_true",
help="initialize token-major padding metadata for fixed-capacity kernels",
)
p.add_argument("--num-blocks", type=int, default=130, help="MSCCL++ low-latency dispatch blocks")
# Launch / fabric.
@@ -333,6 +338,8 @@ def build_mscclpp_cmd(args: argparse.Namespace) -> str:
f"--dispatch-dtype {args.dispatch_dtype} --combine-mode {args.combine_mode} "
f"--output-layout {args.output_layout} --num-blocks {args.num_blocks}"
)
if args.token_major_init_padding:
bench_flags += " --token-major-init-padding"
cupti_build = ""
extra_exports = ""
if args.cupti_inproc or args.kernel_only:

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

@@ -86,6 +86,11 @@ def parse_args():
default="expert_major",
help="Low-latency dispatch output layout",
)
parser.add_argument(
"--token-major-init-padding",
action="store_true",
help="Initialize unused token-major top-k IDs to -1 and weights to zero",
)
parser.add_argument("--bench", action="store_true", help="Run dispatch/combine benchmark after correctness")
parser.add_argument(
"--cuda-graph",
@@ -249,6 +254,7 @@ def validate_token_major_dispatch(
all_topk_weights,
all_x,
expected_scales,
initialize_padding,
):
assert all_x is not None
assert dispatch_out.topk_ids is not None
@@ -258,11 +264,23 @@ def validate_token_major_dispatch(
assert dispatch_out.weights.shape == (num_ranks * num_tokens, num_topk)
source_token_ids = handle.combine_context.source_token_ids
assert source_token_ids.shape == (num_ranks * num_tokens,)
rank_offsets = handle.combine_context.rank_offsets
assert rank_offsets.shape == (num_ranks + 1,)
assert dispatch_out.layout.offsets is rank_offsets
assert int(rank_offsets[0].item()) == 0
total_recv_tokens = int(rank_offsets[-1].item())
assert total_recv_tokens == int(packed_recv_count.sum().item())
if initialize_padding:
assert torch.all(dispatch_out.topk_ids[total_recv_tokens:] == -1)
assert torch.all(dispatch_out.weights[total_recv_tokens:] == 0)
local_expert_begin = rank * num_local_experts
local_expert_end = local_expert_begin + num_local_experts
for source_rank in range(num_ranks):
recv_count = int(packed_recv_count[source_rank].item())
row_begin = int(rank_offsets[source_rank].item())
row_end = int(rank_offsets[source_rank + 1].item())
assert row_end - row_begin == recv_count
source_routing = all_topk_idx[source_rank]
expected_source_tokens = (
((source_routing >= local_expert_begin) & (source_routing < local_expert_end))
@@ -274,8 +292,6 @@ def validate_token_major_dispatch(
if recv_count == 0:
continue
row_begin = source_rank * num_tokens
row_end = row_begin + recv_count
source_tokens = source_token_ids[row_begin:row_end].long()
assert torch.equal(torch.sort(source_tokens).values, expected_source_tokens)
@@ -361,6 +377,7 @@ def reconstruct_token_major_reference(
group,
):
source_token_ids = handle.combine_context.source_token_ids
rank_offsets = handle.combine_context.rank_offsets
destination_ranks = torch.where(
all_topk_idx >= 0,
all_topk_idx // num_local_experts,
@@ -372,8 +389,9 @@ def reconstruct_token_major_reference(
source_count = int(packed_recv_count[source_rank].item())
if source_count == 0:
continue
row_begin = source_rank * num_tokens
row_end = row_begin + source_count
row_begin = int(rank_offsets[source_rank].item())
row_end = int(rank_offsets[source_rank + 1].item())
assert row_end - row_begin == source_count
source_tokens = source_token_ids[row_begin:row_end].long()
selected = first_destination_rank[source_rank, source_tokens] == rank
dispatched_reference_x[source_rank, source_tokens[selected]] = dequantized_x[row_begin:row_end][selected]
@@ -473,6 +491,7 @@ def main():
low_latency_num_blocks=args.num_blocks,
low_latency_combine_mode=combine_mode,
output_layout=output_layout,
token_major_init_padding=args.token_major_init_padding,
quant=dispatch_quant,
)
if rank == 0:
@@ -519,7 +538,8 @@ def main():
torch.cuda.synchronize()
print(f"[rank {rank}] post-dispatch", flush=True)
# expert-major: packed_recv_x [num_local_experts, num_ranks * max_tokens, hidden]
# token-major: packed_recv_x [num_ranks * max_tokens, hidden], rank-grouped
# token-major: packed_recv_x has worst-case capacity, with valid rows compacted
# into [0 : layout.offsets[-1]).
# Reference: gather source tokens, routing IDs, and weights from all ranks.
all_topk_idx = torch.empty((num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="cuda")
@@ -579,6 +599,7 @@ def main():
all_topk_weights=all_topk_weights,
all_x=all_x,
expected_scales=expected_scales,
initialize_padding=args.token_major_init_padding,
)
if rank == 0: