add dsl algorithms support to bench collective

.azure-pipelines/multi-nodes-test.yml

@@ -87,12 +87,54 @@ jobs:
 
   - template: templates/run-remote-task.yml
     parameters:
-      name: RunMscclppTest
-      displayName: Run multi-nodes mscclpp-test
+      name: RunBenchCollectiveAllreduceTest
+      displayName: Run multi-nodes bench_collective DSL allreduce
       continueOnError: true
       runRemoteArgs: '--hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --host ${{ parameters.vmssName }}000000 --user azureuser'
       remoteScript: |
-        bash /root/mscclpp/test/deploy/run_tests.sh mscclpp-test
+        set -e
+        mpirun --allow-run-as-root --bind-to numa -hostfile /root/mscclpp/test/deploy/hostfile_mpi \
+          -mca btl_tcp_if_include eth0 -np 16 -npernode 8 \
+          -x MSCCLPP_DEBUG=WARN -x MSCCLPP_SOCKET_IFNAME=eth0 -x MSCCLPP_HOME=/root/mscclpp \
+          -x LD_LIBRARY_PATH=/root/mscclpp/build/lib:$LD_LIBRARY_PATH -x PATH=$PATH \
+          -x PYTHONPATH=/root/mscclpp/python \
+          python3 -m mscclpp_benchmark.bench_collective \
+            --collective allreduce --autotune --enable-dsl --buffer-mode in-place \
+            --d-model 5120 --batch-sizes 8,16,32 --dsl-tbg 1,2 --dsl-tpb 512,1024
+
+  - template: templates/run-remote-task.yml
+    parameters:
+      name: RunBenchCollectiveAllgatherTest
+      displayName: Run multi-nodes bench_collective DSL allgather
+      continueOnError: true
+      runRemoteArgs: '--hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --host ${{ parameters.vmssName }}000000 --user azureuser'
+      remoteScript: |
+        set -e
+        mpirun --allow-run-as-root --bind-to numa -hostfile /root/mscclpp/test/deploy/hostfile_mpi \
+          -mca btl_tcp_if_include eth0 -np 16 -npernode 8 \
+          -x MSCCLPP_DEBUG=WARN -x MSCCLPP_SOCKET_IFNAME=eth0 -x MSCCLPP_HOME=/root/mscclpp \
+          -x LD_LIBRARY_PATH=/root/mscclpp/build/lib:$LD_LIBRARY_PATH -x PATH=$PATH \
+          -x PYTHONPATH=/root/mscclpp/python \
+          python3 -m mscclpp_benchmark.bench_collective \
+            --collective allgather --autotune --enable-dsl --buffer-mode in-place \
+            --d-model 5120 --batch-sizes 8,16,32 --dsl-tbg 1,2 --dsl-tpb 512,1024
+
+  - template: templates/run-remote-task.yml
+    parameters:
+      name: RunBenchCollectiveReduceScatterTest
+      displayName: Run multi-nodes bench_collective DSL reducescatter
+      continueOnError: true
+      runRemoteArgs: '--hostfile $(System.DefaultWorkingDirectory)/test/deploy/hostfile --host ${{ parameters.vmssName }}000000 --user azureuser'
+      remoteScript: |
+        set -e
+        mpirun --allow-run-as-root --bind-to numa -hostfile /root/mscclpp/test/deploy/hostfile_mpi \
+          -mca btl_tcp_if_include eth0 -np 16 -npernode 8 \
+          -x MSCCLPP_DEBUG=WARN -x MSCCLPP_SOCKET_IFNAME=eth0 -x MSCCLPP_HOME=/root/mscclpp \
+          -x LD_LIBRARY_PATH=/root/mscclpp/build/lib:$LD_LIBRARY_PATH -x PATH=$PATH \
+          -x PYTHONPATH=/root/mscclpp/python \
+          python3 -m mscclpp_benchmark.bench_collective \
+            --collective reducescatter --autotune --enable-dsl --buffer-mode in-place \
+            --d-model 5120 --batch-sizes 8,16,32 --dsl-tbg 1,2 --dsl-tpb 512,1024
 
   - template: templates/run-remote-task.yml
     parameters:

python/mscclpp/default_algos/reducescatter_multi_nodes.py

@@ -188,3 +188,37 @@ def reducescatter_multi_nodes(
             )
 
     return prog
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--name", type=str, help="name of the program")
+    parser.add_argument("--num_gpus", type=int, help="total number of gpus")
+    parser.add_argument("--gpus_per_node", type=int, help="number of gpus per node")
+    parser.add_argument("--tbg", type=int, default=1, help="thread block group size")
+    parser.add_argument("--num_threads_per_block", type=int, default=1024, help="number of threads per block")
+    parser.add_argument("--min_message_size", type=int, default=0, help="minimum message size")
+    parser.add_argument("--max_message_size", type=int, default=2**64 - 1, help="maximum message size")
+
+    args = parser.parse_args()
+
+    spec = AlgoSpec(
+        name=args.name,
+        collective=ReduceScatter(args.num_gpus, 1, True),
+        nranks_per_node=args.gpus_per_node,
+        world_size=args.num_gpus,
+        in_place=True,
+        instances=1,
+        protocol="LL",
+        auto_sync=False,
+        num_threads_per_block=args.num_threads_per_block,
+        reuse_resources=True,
+        use_double_scratch_buffer=True,
+        min_message_size=args.min_message_size,
+        max_message_size=args.max_message_size,
+    )
+
+    prog = reducescatter_multi_nodes(spec, args.tbg)
+    print(prog.to_json())

python/mscclpp/language/collectives.py

@@ -190,10 +190,17 @@ def __init__(self, num_ranks, chunk_factor, inplace):
             num_ranks (int): The number of ranks participating in the ReduceScatter.
             chunk_factor (int): The size factor for data chunks.
             inplace (bool): Whether the operation should be performed in-place.
+                ReduceScatter only supports in-place mode.
+
+        Raises:
+            ValueError: If ``inplace`` is False, since out-of-place ReduceScatter
+                is not supported.
 
         Example:
-            >>> reduce_scatter = ReduceScatter(num_ranks=4, chunk_factor=1, inplace=False)
+            >>> reduce_scatter = ReduceScatter(num_ranks=4, chunk_factor=1, inplace=True)
         """
+        if not inplace:
+            raise ValueError("ReduceScatter only supports in-place mode.")
         Collective.__init__(self, num_ranks, chunk_factor, inplace)
         self.name = "reducescatter"
 

python/mscclpp_benchmark/bench_collective.py

@@ -12,7 +12,7 @@
 
 _mscclpp_module = None
 
-from mscclpp_benchmark.comm import Comm
+from mscclpp_benchmark.comm import DEFAULT_DSL_TBG, DEFAULT_DSL_TPB, Comm
 from mscclpp_benchmark.correctness import (
     CorrectnessStats,
     check_correctness as _check_correctness,
@@ -24,6 +24,7 @@
 
 _ALLREDUCE = "allreduce"
 _ALLGATHER = "allgather"
+_REDUCESCATTER = "reducescatter"
 _DEFAULT_BATCH_SIZES = (
     1,
     2,
@@ -81,6 +82,14 @@ class CandidateSpec:
     supported_skus: tuple[str, ...] | None = None
     requires_nvls: bool = False
     requires_symmetric_memory: bool = False
+    # Native algorithms use all-pairs CUDA-IPC and only work within a single node; they hang if
+    # tuned on a multi-node job. Only algorithms that explicitly opt in are considered when the
+    # world spans more than one node.
+    supports_multi_node: bool = False
+    # None means "use the tuner's global sweep"; an explicit tuple overrides it, which DSL
+    # algorithms need since they bake their launch geometry into the plan and ignore nblocks/nthreads.
+    candidate_nblocks: tuple[int, ...] | None = None
+    candidate_nthreads: tuple[int, ...] | None = None
 
 
 @dataclass
@@ -190,6 +199,10 @@ def _parse_int_list(raw: str | None, default: tuple[int, ...]) -> tuple[int, ...
 
 
 def _candidate_specs(collective: str, *, symmetric_memory: bool = False) -> tuple[CandidateSpec, ...]:
+    if collective == _REDUCESCATTER:
+        # There are no native reducescatter algorithms in the default collection, so the compiled
+        # DSL variants are the only candidates.
+        return ()
     if collective == _ALLGATHER:
         allgather_candidates = (
             CandidateSpec("default_allgather_fullmesh2", max_nblocks=64, supported_skus=("MI300X",)),
@@ -255,18 +268,57 @@ def _candidate_specs(collective: str, *, symmetric_memory: bool = False) -> tupl
     return candidates
 
 
+def _dsl_candidate_specs(comm: Comm, collective: str) -> tuple[CandidateSpec, ...]:
+    """Synthesize candidate specs for the DSL algorithms compiled by Comm.
+
+    Unlike the native algorithms, DSL variant names are only known at runtime, so their specs cannot
+    be listed statically. Each variant reports its own applicable message size range, which the usual
+    message size filter then applies against the effective (whole buffer) size, and pins the tuner to
+    a single launch config because the plan already bakes in its launch geometry.
+    """
+    available = comm.algorithms.get(collective, {})
+    specs: list[CandidateSpec] = []
+    for name in sorted(getattr(comm, "dsl_algorithms", ())):
+        algorithm = available.get(name)
+        if algorithm is None:
+            continue
+        min_message_size, max_message_size = algorithm.message_size_range
+        specs.append(
+            CandidateSpec(
+                name,
+                min_message_size=min_message_size,
+                max_message_size=max_message_size,
+                candidate_nblocks=(0,),
+                candidate_nthreads=(0,),
+                supports_multi_node=True,
+            )
+        )
+    return tuple(specs)
+
+
 def _candidate_algorithms(comm: Comm, case: BenchmarkCase) -> list[tuple[Any, CandidateSpec]]:
     available = comm.algorithms.get(case.collective, {})
     candidates: list[tuple[Any, CandidateSpec]] = []
     seen: set[str] = set()
     symmetric_memory = case.symmetric_memory
     profile = getattr(comm, "hardware_profile", None)
+    comm_group = comm.comm_group
+    nranks = getattr(comm_group, "nranks", 1) or 1
+    nranks_per_node = getattr(comm_group, "nranks_per_node", nranks) or nranks
+    n_nodes = nranks // nranks_per_node if nranks_per_node else 1
+    effective_message_size = _effective_message_size(case.collective, case.message_size, nranks)
     filtered_out = False
-    for candidate in _candidate_specs(case.collective, symmetric_memory=symmetric_memory):
+    for candidate in (
+        *_candidate_specs(case.collective, symmetric_memory=symmetric_memory),
+        *_dsl_candidate_specs(comm, case.collective),
+    ):
+        if n_nodes > 1 and not candidate.supports_multi_node:
+            filtered_out = True
+            continue
         if not _candidate_supports_profile(candidate, profile):
             filtered_out = True
             continue
-        if not _candidate_supports_message_size(candidate, case.message_size):
+        if not _candidate_supports_message_size(candidate, effective_message_size):
             filtered_out = True
             continue
         if candidate.requires_nvls and not _mscclpp().is_nvls_supported():
@@ -296,6 +348,19 @@ def _candidate_supports_profile(candidate: CandidateSpec, profile: HardwareProfi
     return sku in candidate.supported_skus
 
 
+def _effective_message_size(collective: str, message_size: int, nranks: int) -> int:
+    """Return the buffer size the executor validates against an algorithm's message size range.
+
+    Mirrors ``matchExecutionPlan`` in ``src/ext/nccl/algorithm_selector.cc`` and
+    ``ExecutionPlan::Impl::checkMessageSize``, which both compare the range against the whole buffer:
+    the output size for allgather and the input size otherwise. For allgather and reducescatter that
+    spans every rank, whereas ``case.message_size`` is only this rank's chunk.
+    """
+    if collective in (_ALLGATHER, _REDUCESCATTER):
+        return message_size * nranks
+    return message_size
+
+
 def _candidate_supports_message_size(candidate: CandidateSpec, message_size: int) -> bool:
     if candidate.min_message_size is not None and message_size < candidate.min_message_size:
         return False
@@ -334,9 +399,24 @@ def _make_case(
             symmetric_memory=symmetric_memory,
         )
 
+    if collective == _REDUCESCATTER:
+        # The DSL reducescatter is compiled in-place, so the per-rank output chunk always aliases the
+        # matching slice of the full input buffer (mirrors python/test/executor_test.py build_bufs).
+        input_buffer = _mscclpp().GpuBuffer(nelems * comm_group.nranks, dtype=dtype_spec.cupy_dtype)
+        start = comm_group.my_rank * nelems
+        output = input_buffer[start : start + nelems]
+        return BenchmarkCase(
+            collective=collective,
+            message_size=output.nbytes,
+            total_size=input_buffer.nbytes,
+            input=input_buffer,
+            output=output,
+            dtype_spec=dtype_spec,
+            symmetric_memory=symmetric_memory,
+        )
+
     if collective != _ALLGATHER:
         raise ValueError(f"Unsupported collective: {collective}")
-
     if buffer_mode == "in-place":
         output = _mscclpp().GpuBuffer(nelems * comm_group.nranks, dtype=dtype_spec.cupy_dtype)
         start = comm_group.my_rank * nelems
@@ -441,7 +521,7 @@ def _busbw_factor(collective: str, nranks: int) -> float:
         return 1.0
     if collective == _ALLREDUCE:
         return 2 * (nranks - 1) / nranks
-    if collective == _ALLGATHER:
+    if collective in (_ALLGATHER, _REDUCESCATTER):
         return (nranks - 1) / nranks
     raise ValueError(f"Unsupported collective: {collective}")
 
@@ -470,7 +550,7 @@ def _format_mismatches(stats: CorrectnessStats | None) -> str:
 
 def _build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(description="Benchmark MSCCL++ collectives without PyTorch dependencies")
-    parser.add_argument("--collective", choices=(_ALLREDUCE, _ALLGATHER), default=_ALLREDUCE)
+    parser.add_argument("--collective", choices=(_ALLREDUCE, _ALLGATHER, _REDUCESCATTER), default=_ALLREDUCE)
     parser.add_argument("--d-model", type=int, default=5120)
     parser.add_argument("--dtype", default="float16")
     parser.add_argument("--accum-type", help="Accumulation type for reductions: native, float16, or float32")
@@ -495,6 +575,13 @@ def _build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--tune-iterations", type=int, default=20)
     parser.add_argument("--candidate-nblocks", help="Comma-separated nblocks tuning candidates")
     parser.add_argument("--candidate-nthreads", help="Comma-separated nthreads tuning candidates")
+    parser.add_argument(
+        "--enable-dsl",
+        action="store_true",
+        help="Compile DSL algorithm variants and tune them alongside the native algorithms",
+    )
+    parser.add_argument("--dsl-tbg", help="Comma-separated DSL thread_block_group_size candidates")
+    parser.add_argument("--dsl-tpb", help="Comma-separated DSL num_threads_per_block candidates")
     parser.add_argument("--symmetric-memory", action="store_true")
     return parser
 
@@ -533,6 +620,8 @@ def main(argv: list[str] | None = None) -> None:
     batch_sizes = _parse_int_list(args.batch_sizes, _DEFAULT_BATCH_SIZES)
     candidate_nblocks = _parse_int_list(args.candidate_nblocks, _DEFAULT_CANDIDATE_NBLOCKS)
     candidate_nthreads = _parse_int_list(args.candidate_nthreads, _DEFAULT_CANDIDATE_NTHREADS)
+    dsl_tbg = _parse_int_list(args.dsl_tbg, DEFAULT_DSL_TBG)
+    dsl_tpb = _parse_int_list(args.dsl_tpb, DEFAULT_DSL_TPB)
 
     comm_group = _mscclpp().CommGroup(MPI.COMM_WORLD)
     setattr(comm_group, "_mpi_comm", MPI.COMM_WORLD)
@@ -543,6 +632,10 @@ def main(argv: list[str] | None = None) -> None:
         config_store=config_store,
         hardware_profile=hardware_profile,
         scratch_buffer_size=args.scratch_buffer_size,
+        collective=args.collective,
+        enable_dsl=args.enable_dsl,
+        dsl_tbg=dsl_tbg,
+        dsl_tpb=dsl_tpb,
     )
     tuner = OfflineTuner(
         comm,