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_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,10 @@
 
 _ALLREDUCE = "allreduce"
 _ALLGATHER = "allgather"
+_REDUCESCATTER = "reducescatter"
+# CLI buffer-mode names mapped to the CppCollectiveBufferMode enum members exposed by
+# python/csrc/algorithm.cpp.
+_BUFFER_MODE_NAMES = {"in-place": "IN_PLACE", "out-of-place": "OUT_OF_PLACE"}
 _DEFAULT_BATCH_SIZES = (
     1,
     2,
@@ -81,6 +85,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
@@ -92,6 +104,7 @@ class BenchmarkCase:
     output: cp.ndarray
     dtype_spec: DTypeSpec
     symmetric_memory: bool = False
+    buffer_mode: str = "in-place"
 
 
 def _device_name() -> str:
@@ -190,6 +203,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 +272,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
+    multi_node = nranks > nranks_per_node
+    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 multi_node 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():
@@ -278,6 +334,9 @@ def _candidate_algorithms(comm: Comm, case: BenchmarkCase) -> list[tuple[Any, Ca
         algorithm = available.get(candidate.algorithm)
         if algorithm is None or algorithm.name in seen:
             continue
+        if not _algorithm_supports_buffer_mode(algorithm, case.buffer_mode):
+            filtered_out = True
+            continue
         seen.add(algorithm.name)
         candidates.append((algorithm, candidate))
     if candidates:
@@ -296,6 +355,34 @@ 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 _algorithm_supports_buffer_mode(algorithm: Any, buffer_mode: str) -> bool:
+    """Whether ``algorithm`` can run against buffers laid out in ``buffer_mode``.
+
+    Mirrors the ``bufferModeMatch`` test in ``src/ext/nccl/algorithm_selector.cc``: an algorithm
+    declaring ``ANY`` accepts either layout, otherwise the modes must match. The benchmark resolves
+    algorithms by name instead of going through that selector, so nothing else enforces this. Native
+    algorithms are registered without an explicit mode and therefore report ``ANY``.
+    """
+    mode = getattr(algorithm, "buffer_mode", None)
+    if mode is None:
+        return True
+    name = getattr(mode, "name", None) or str(mode).rsplit(".", 1)[-1]
+    return name in ("ANY", _BUFFER_MODE_NAMES.get(buffer_mode))
+
+
 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
@@ -332,11 +419,28 @@ def _make_case(
             output=output,
             dtype_spec=dtype_spec,
             symmetric_memory=symmetric_memory,
+            buffer_mode=buffer_mode,
+        )
+
+    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,
+            buffer_mode="in-place",
         )
 
     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
@@ -353,6 +457,7 @@ def _make_case(
         output=output,
         dtype_spec=dtype_spec,
         symmetric_memory=symmetric_memory,
+        buffer_mode=buffer_mode,
     )
 
 
@@ -441,7 +546,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 +575,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 +600,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 +645,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 +657,11 @@ 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,
+        buffer_mode=args.buffer_mode,
+        dsl_tbg=dsl_tbg,
+        dsl_tpb=dsl_tpb,
     )
     tuner = OfflineTuner(
         comm,