[Perf][Distributed] Add low-SM multimem reduce-scatter for SM100/SM103 (#55072)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Co-authored-by: Summer Yang <girasoleyang@gmail.com>

benchmarks/kernels/benchmark_multimem_reduce_scatter.py

@@ -0,0 +1,480 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark production MNNVL reduce-scatter backends on SM100/SM103.
+
+The multimem candidate includes the device-to-device copy into its persistent
+symmetric input buffer. Timings use pointer-distinct CUDA graphs and report the
+median of the per-sample maximum across ranks.
+
+Run the default BF16 message-size sweep with:
+
+    .venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=8 \
+        benchmarks/kernels/benchmark_multimem_reduce_scatter.py
+"""
+
+import argparse
+import json
+import os
+import statistics
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+import vllm._custom_ops as ops
+from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce
+
+_DTYPES = {
+    "bf16": torch.bfloat16,
+    "f16": torch.float16,
+    "f32": torch.float32,
+}
+
+
+@dataclass
+class Candidate:
+    name: str
+    runs: list[Callable[[], None]]
+    outputs: list[torch.Tensor]
+
+
+def parse_size(value: str) -> int:
+    suffixes = {"k": 1 << 10, "m": 1 << 20, "g": 1 << 30}
+    value = value.strip().lower()
+    multiplier = suffixes.get(value[-1], 1)
+    number = value[:-1] if multiplier != 1 else value
+    return int(float(number) * multiplier)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--message-bytes",
+        type=parse_size,
+        nargs="+",
+        default=[
+            16 << 20,
+            32 << 20,
+            64 << 20,
+            68 << 20,
+            72 << 20,
+            76 << 20,
+            96 << 20,
+        ],
+        help="Full input bytes on each rank. K/M/G suffixes are accepted.",
+    )
+    parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16")
+    parser.add_argument(
+        "--blocks",
+        type=int,
+        nargs="+",
+        default=[8],
+        help="Maximum CTAs used by each multimem candidate.",
+    )
+    parser.add_argument("--num-workspaces", type=int, default=4)
+    parser.add_argument("--warmup-replays", type=int, default=8)
+    parser.add_argument("--samples", type=int, default=31)
+    parser.add_argument(
+        "--cold-l2",
+        action="store_true",
+        help="Flush approximately twice L2 capacity before each replay.",
+    )
+    return parser.parse_args()
+
+
+def make_lamport_run(
+    comm: CustomAllreduce,
+    input_tensor: torch.Tensor,
+    output: torch.Tensor,
+) -> Callable[[], None]:
+    def run() -> None:
+        ops.mnnvl_lamport_reduce_scatter(
+            comm._ptr,
+            input_tensor,
+            output,
+            comm.mnnvl_lamport_rs_local_ptr,
+            comm.mnnvl_lamport_rs_epoch_ptr,
+            comm.mnnvl_buffer_size,
+        )
+
+    return run
+
+
+def make_multimem_run(
+    comm: CustomAllreduce,
+    input_tensor: torch.Tensor,
+    output: torch.Tensor,
+    blocks: int,
+) -> Callable[[], None]:
+    def run() -> None:
+        ops.mnnvl_multimem_reduce_scatter(
+            comm._ptr,
+            input_tensor,
+            output,
+            comm.mnnvl_multimem_rs_local_ptr,
+            comm.mnnvl_multimem_rs_multicast_ptr,
+            comm.mnnvl_multimem_rs_buffer_size,
+            blocks,
+        )
+
+    return run
+
+
+def make_nccl_run(
+    input_tensor: torch.Tensor,
+    output: torch.Tensor,
+    device_group: dist.ProcessGroup,
+) -> Callable[[], None]:
+    def run() -> None:
+        dist.reduce_scatter_single(output, input_tensor, group=device_group)
+
+    return run
+
+
+def build_candidates(
+    comm: CustomAllreduce,
+    inputs: list[torch.Tensor],
+    local_elements: int,
+    blocks_values: list[int],
+    device_group: dist.ProcessGroup,
+) -> list[Candidate]:
+    candidates: list[Candidate] = []
+
+    lamport_outputs = [
+        input_tensor.new_empty(local_elements) for input_tensor in inputs
+    ]
+    candidates.append(
+        Candidate(
+            "lamport",
+            [
+                make_lamport_run(comm, input_tensor, output)
+                for input_tensor, output in zip(inputs, lamport_outputs)
+            ],
+            lamport_outputs,
+        )
+    )
+
+    for blocks in blocks_values:
+        outputs = [input_tensor.new_empty(local_elements) for input_tensor in inputs]
+        candidates.append(
+            Candidate(
+                f"multimem_b{blocks}",
+                [
+                    make_multimem_run(comm, input_tensor, output, blocks)
+                    for input_tensor, output in zip(inputs, outputs)
+                ],
+                outputs,
+            )
+        )
+
+    nccl_outputs = [input_tensor.new_empty(local_elements) for input_tensor in inputs]
+    candidates.append(
+        Candidate(
+            "nccl",
+            [
+                make_nccl_run(input_tensor, output, device_group)
+                for input_tensor, output in zip(inputs, nccl_outputs)
+            ],
+            nccl_outputs,
+        )
+    )
+    return candidates
+
+
+def check_correctness(
+    candidates: list[Candidate],
+    inputs: list[torch.Tensor],
+    rank: int,
+    world_size: int,
+    cpu_group: dist.ProcessGroup,
+    device_group: dist.ProcessGroup,
+    graphs: dict[str, list[torch.cuda.CUDAGraph]] | None = None,
+) -> None:
+    expected_outputs = []
+    for input_tensor in inputs:
+        expected = input_tensor.float().clone()
+        dist.all_reduce(expected, group=device_group)
+        expected_outputs.append(expected.to(input_tensor.dtype).chunk(world_size)[rank])
+
+    if inputs[0].dtype == torch.float32:
+        rtol, atol = 1e-5, 1e-5
+    elif inputs[0].dtype == torch.float16:
+        rtol, atol = 1e-2, 1e-2
+    else:
+        rtol, atol = 5e-2, 1.25e-1
+
+    for candidate in candidates:
+        for workspace_index, expected in enumerate(expected_outputs):
+            dist.barrier(group=cpu_group)
+            if graphs is None:
+                candidate.runs[workspace_index]()
+            else:
+                graphs[candidate.name][workspace_index].replay()
+            torch.accelerator.synchronize()
+            error = None
+            try:
+                torch.testing.assert_close(
+                    candidate.outputs[workspace_index],
+                    expected,
+                    rtol=rtol,
+                    atol=atol,
+                )
+            except AssertionError as assertion_error:
+                error = assertion_error
+            correct = torch.tensor(
+                error is None,
+                dtype=torch.int32,
+                device=candidate.outputs[workspace_index].device,
+            )
+            dist.all_reduce(correct, op=dist.ReduceOp.MIN, group=device_group)
+            if not correct.item():
+                location = f"{candidate.name} workspace {workspace_index}"
+                detail = str(error) if error is not None else "failed on another rank"
+                raise AssertionError(
+                    f"{location} produced an incorrect shard: {detail}"
+                )
+    dist.barrier(group=cpu_group)
+
+
+def capture_graph(
+    run: Callable[[], None],
+    cpu_group: dist.ProcessGroup,
+) -> tuple[torch.cuda.CUDAGraph, torch.cuda.Stream]:
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    dist.barrier(group=cpu_group)
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            run()
+    stream.synchronize()
+    dist.barrier(group=cpu_group)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        run()
+    torch.cuda.current_stream().wait_stream(stream)
+    dist.barrier(group=cpu_group)
+    return graph, stream
+
+
+def capture_candidates(
+    candidates: list[Candidate],
+    cpu_group: dist.ProcessGroup,
+) -> tuple[dict[str, list[torch.cuda.CUDAGraph]], list[object]]:
+    graphs = {}
+    keepalive: list[object] = []
+    for candidate in candidates:
+        bundles = [capture_graph(run, cpu_group) for run in candidate.runs]
+        graphs[candidate.name] = [graph for graph, _ in bundles]
+        keepalive.extend(bundles)
+    return graphs, keepalive
+
+
+def benchmark_graphs(
+    graphs: dict[str, list[torch.cuda.CUDAGraph]],
+    warmup_replays: int,
+    samples: int,
+    device_group: dist.ProcessGroup,
+    device_barrier: Callable[[], None],
+    flush: torch.Tensor | None,
+) -> dict[str, float]:
+    names = list(graphs)
+    for round_index in range(warmup_replays):
+        for candidate_index in range(len(names)):
+            name = names[(round_index + candidate_index) % len(names)]
+            device_barrier()
+            graphs[name][round_index % len(graphs[name])].replay()
+    torch.accelerator.synchronize()
+
+    timings = {name: [] for name in names}
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    for sample_index in range(samples):
+        for candidate_index in range(len(names)):
+            name = names[(sample_index + candidate_index) % len(names)]
+            if flush is not None:
+                flush.zero_()
+            device_barrier()
+            start.record()
+            graphs[name][sample_index % len(graphs[name])].replay()
+            end.record()
+            end.synchronize()
+
+            elapsed = torch.tensor(
+                start.elapsed_time(end) * 1000,
+                dtype=torch.float64,
+                device=torch.accelerator.current_device_index(),
+            )
+            dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group)
+            timings[name].append(elapsed.item())
+
+    return {name: statistics.median(values) for name, values in timings.items()}
+
+
+def benchmark_size(
+    args: argparse.Namespace,
+    comm: CustomAllreduce,
+    message_bytes: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+    device_group: dist.ProcessGroup,
+    cpu_group: dist.ProcessGroup,
+    device_barrier: Callable[[], None],
+) -> dict[str, object]:
+    element_size = dtype.itemsize
+    if message_bytes % (element_size * world_size) != 0:
+        raise ValueError("Message bytes must divide evenly into TP shards")
+    num_elements = message_bytes // element_size
+    local_elements = num_elements // world_size
+    if local_elements * element_size % 16 != 0:
+        raise ValueError("Each output shard must be a multiple of 16 bytes")
+
+    inputs = []
+    for workspace_index in range(args.num_workspaces):
+        generator = torch.Generator(device=device)
+        generator.manual_seed(1000 + rank * 101 + workspace_index)
+        input_tensor = torch.empty(num_elements, dtype=dtype, device=device)
+        input_tensor.uniform_(-1.0, 1.0, generator=generator)
+        inputs.append(input_tensor)
+
+    candidates = build_candidates(
+        comm,
+        inputs,
+        local_elements,
+        args.blocks,
+        device_group,
+    )
+    torch.accelerator.synchronize()
+    dist.barrier(group=cpu_group)
+    check_correctness(
+        candidates,
+        inputs,
+        rank,
+        world_size,
+        cpu_group,
+        device_group,
+    )
+    graphs, graph_keepalive = capture_candidates(candidates, cpu_group)
+    check_correctness(
+        candidates,
+        inputs,
+        rank,
+        world_size,
+        cpu_group,
+        device_group,
+        graphs,
+    )
+
+    flush = None
+    if args.cold_l2:
+        l2_bytes = torch.cuda.get_device_properties(device).L2_cache_size
+        flush = torch.empty(2 * l2_bytes, dtype=torch.uint8, device=device)
+    times = benchmark_graphs(
+        graphs,
+        args.warmup_replays,
+        args.samples,
+        device_group,
+        device_barrier,
+        flush,
+    )
+    graph_keepalive.clear()
+
+    bus_bytes = message_bytes * (world_size - 1) / world_size
+    fastest = min(times, key=times.__getitem__)
+    return {
+        "message_bytes": message_bytes,
+        "shard_bytes": message_bytes // world_size,
+        "latency_us": times,
+        "logical_bus_gbps": {
+            name: bus_bytes / (latency_us * 1e3) for name, latency_us in times.items()
+        },
+        "fastest": fastest,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    if not args.message_bytes or min(args.message_bytes) <= 0:
+        raise ValueError("--message-bytes values must be positive")
+    if not args.blocks or min(args.blocks) <= 0 or max(args.blocks) > 8:
+        raise ValueError("--blocks values must be in [1, 8]")
+    if min(args.num_workspaces, args.samples) <= 0:
+        raise ValueError("Workspace and sample counts must be positive")
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.accelerator.set_device_index(local_rank)
+    dist.init_process_group("nccl")
+    device_group = dist.group.WORLD
+    cpu_group = dist.new_group(backend="gloo")
+    rank = dist.get_rank(device_group)
+    world_size = dist.get_world_size(device_group)
+    device = torch.device("cuda", local_rank)
+    capability = torch.cuda.get_device_capability(device)
+    if world_size not in (2, 4, 8) or capability not in ((10, 0), (10, 3)):
+        raise RuntimeError(
+            "This benchmark targets TP2/TP4/TP8 SM100/SM103, "
+            f"got TP{world_size} SM{capability}"
+        )
+
+    max_message_bytes = max(args.message_bytes)
+    comm = CustomAllreduce(
+        group=cpu_group,
+        device=device,
+        max_mnnvl_reduce_scatter_size=max_message_bytes,
+        max_mnnvl_multimem_reduce_scatter_size=max_message_bytes,
+    )
+    if comm.disabled:
+        raise RuntimeError("The production custom all-reduce backend is unavailable")
+    comm._init_mnnvl_multimem_reduce_scatter_buffer()
+    if not comm.mnnvl_multimem_rs_multicast_ptr:
+        raise RuntimeError("The production MNNVL multimem backend is unavailable")
+
+    dtype = _DTYPES[args.dtype]
+    warmup = torch.zeros(1, dtype=torch.float32, device=device)
+    dist.all_reduce(warmup, group=device_group)
+    sync = torch.zeros_like(warmup)
+
+    def device_barrier() -> None:
+        dist.all_reduce(sync, group=device_group)
+
+    results = [
+        benchmark_size(
+            args,
+            comm,
+            message_bytes,
+            dtype,
+            device,
+            rank,
+            world_size,
+            device_group,
+            cpu_group,
+            device_barrier,
+        )
+        for message_bytes in args.message_bytes
+    ]
+    if rank == 0:
+        metadata = {
+            "world_size": world_size,
+            "gpu": torch.cuda.get_device_name(device),
+            "sm": f"{capability[0]}{capability[1]}",
+            "num_sms": torch.cuda.get_device_properties(device).multi_processor_count,
+            "dtype": args.dtype,
+            "torch": torch.__version__,
+            "cuda": torch.version.cuda,
+            "blocks": args.blocks,
+            "num_workspaces": args.num_workspaces,
+            "cold_l2": args.cold_l2,
+        }
+        print(json.dumps({"metadata": metadata, "results": results}, indent=2))
+
+    comm.close()
+    dist.barrier(group=cpu_group)
+    dist.destroy_process_group(cpu_group)
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

csrc/custom_all_gather_reduce_scatter.cuh

@@ -7,6 +7,10 @@ namespace vllm {
 constexpr int kMnnvlLamportAgThreads = 128;
 constexpr int kMnnvlLamportRsThreads = 256;
 constexpr int kMnnvlLamportConcurrentPollMaxPacks = 8192;
+constexpr int kMnnvlMultimemRsThreads = 1024;
+constexpr int kMnnvlMultimemRsBlockLimit = 8;
+constexpr int kMnnvlMultimemRsVectorBytes = 16;
+constexpr int kMnnvlMultimemRsUnroll = 8;
 
 using CopyPack = array_t<uint64_t, 2>;
 
@@ -48,6 +52,127 @@ __global__ void __launch_bounds__(512, 1)
   barrier_at_end<ngpus, true>(sg, self_sg, rank);
 }
 
+// Multimem reduce-scatter reduces through an NVLS multicast mapping, which
+// has no AMD equivalent. The whole path is compiled out on ROCm and the host
+// entry point rejects the call instead of launching anything.
+#if !defined(USE_ROCM)
+
+template <typename T>
+DINLINE void multimem_load_reduce_16(uint32_t (&result)[4], const T* address) {
+  #if CUDA_VERSION >= 12020 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  if constexpr (std::is_same<T, nv_bfloat16>::value) {
+    asm volatile(
+        "multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.bf16x2 "
+        "{%0,%1,%2,%3}, [%4];"
+        : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3])
+        : "l"(address)
+        : "memory");
+  } else if constexpr (std::is_same<T, half>::value) {
+    asm volatile(
+        "multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.f16x2 "
+        "{%0,%1,%2,%3}, [%4];"
+        : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3])
+        : "l"(address)
+        : "memory");
+  } else {
+    static_assert(std::is_same<T, float>::value);
+    asm volatile(
+        "multimem.ld_reduce.relaxed.sys.global.add.v4.f32 "
+        "{%0,%1,%2,%3}, [%4];"
+        : "=r"(result[0]), "=r"(result[1]), "=r"(result[2]), "=r"(result[3])
+        : "l"(address)
+        : "memory");
+  }
+  #else
+  asm volatile("trap;");
+  #endif
+}
+
+DINLINE void store_global_16(void* address, const uint32_t (&value)[4]) {
+  asm volatile("st.global.v4.u32 [%0], {%1,%2,%3,%4};"
+               :
+               : "l"(address), "r"(value[0]), "r"(value[1]), "r"(value[2]),
+                 "r"(value[3])
+               : "memory");
+}
+
+DINLINE void mnnvl_multimem_publish_flag(FlagType* flag_addr, FlagType flag) {
+  #if CUDA_VERSION >= 12020 && defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  asm volatile("multimem.st.release.sys.global.u32 [%0], %1;"
+               :
+               : "l"(flag_addr), "r"(flag)
+               : "memory");
+  // The flag is stored through the multicast alias and polled through the
+  // local unicast alias in the same grid.
+  asm volatile("fence.proxy.alias;" ::: "memory");
+  #else
+  asm volatile("trap;");
+  #endif
+}
+
+template <bool ready, int ngpus>
+__global__ void mnnvl_multimem_barrier_kernel(Signal* local_signal,
+                                              Signal* multicast_signal,
+                                              int rank) {
+  FlagType flag = local_signal->_flag[0] + 1;
+  FlagType* multicast_counters =
+      ready ? multicast_signal->start[0] : multicast_signal->end[0];
+  FlagType* local_counters =
+      ready ? local_signal->start[0] : local_signal->end[0];
+  // Replicate this rank's counter into every backing allocation.
+  mnnvl_multimem_publish_flag(&multicast_counters[rank], flag);
+  #pragma unroll
+  for (int peer = 0; peer < ngpus; ++peer) {
+    while (ld_flag_acquire(&local_counters[peer]) != flag);
+  }
+  local_signal->_flag[0] = flag;
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(kMnnvlMultimemRsThreads, 1)
+    mnnvl_multimem_reduce_scatter_kernel(const T* __restrict__ multicast_input,
+                                         T* __restrict__ result, int rank,
+                                         int packs_per_rank) {
+  constexpr int kThreadsPerWarp = 32;
+  constexpr int kWarpsPerCta = kMnnvlMultimemRsThreads / kThreadsPerWarp;
+  constexpr int kPacksPerWarpIteration =
+      kThreadsPerWarp * kMnnvlMultimemRsUnroll;
+
+  int lane;
+  asm volatile("mov.u32 %0, %%laneid;" : "=r"(lane));
+  int warp = blockIdx.x * kWarpsPerCta + threadIdx.x / kThreadsPerWarp;
+  int num_warps = gridDim.x * kWarpsPerCta;
+  int pack_offset = warp * kPacksPerWarpIteration + lane;
+  int pack_stride = num_warps * kPacksPerWarpIteration;
+  auto* rank_input = reinterpret_cast<const char*>(multicast_input) +
+                     rank * packs_per_rank * kMnnvlMultimemRsVectorBytes;
+  auto* rank_output = reinterpret_cast<char*>(result);
+
+  while (pack_offset < packs_per_rank) {
+    uint32_t reduced[kMnnvlMultimemRsUnroll][4];
+  #pragma unroll
+    for (int u = 0; u < kMnnvlMultimemRsUnroll; ++u) {
+      int pack = pack_offset + u * kThreadsPerWarp;
+      if (pack < packs_per_rank) {
+        multimem_load_reduce_16<T>(
+            reduced[u], reinterpret_cast<const T*>(
+                            rank_input + pack * kMnnvlMultimemRsVectorBytes));
+      }
+    }
+  #pragma unroll
+    for (int u = 0; u < kMnnvlMultimemRsUnroll; ++u) {
+      int pack = pack_offset + u * kThreadsPerWarp;
+      if (pack < packs_per_rank) {
+        store_global_16(rank_output + pack * kMnnvlMultimemRsVectorBytes,
+                        reduced[u]);
+      }
+    }
+    pack_offset += pack_stride;
+  }
+}
+
+#endif  // !defined(USE_ROCM)
+
 template <typename P>
 union LamportPack {
   P packed;

csrc/custom_all_reduce.cuh

@@ -339,6 +339,12 @@ class CustomAllreduce {
   void mnnvl_lamport_reduce_scatter(cudaStream_t stream, T* input, T* output,
                                     void* local_buffer, uint32_t* epochs,
                                     int size, int stage_size_bytes);
+  template <typename T>
+  void mnnvl_multimem_reduce_scatter(cudaStream_t stream,
+                                     const T* multicast_input, T* output,
+                                     Signal* local_signal,
+                                     Signal* multicast_signal, int size,
+                                     int block_limit);
 
   ~CustomAllreduce() {
     for (auto [_, ptr] : ipc_handles_) {

csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu

@@ -191,6 +191,55 @@ void CustomAllreduce::mnnvl_lamport_reduce_scatter(cudaStream_t stream,
 #undef MNNVL_LAMPORT_RS_LAUNCH
 }
 
+template <typename T>
+void CustomAllreduce::mnnvl_multimem_reduce_scatter(
+    cudaStream_t stream, const T* multicast_input, T* output,
+    Signal* local_signal, Signal* multicast_signal, int size, int block_limit) {
+#if defined(USE_ROCM)
+  throw std::runtime_error(
+      "MNNVL multimem reduce-scatter requires NVLS multicast and is not "
+      "supported on ROCm");
+#else
+  if (size <= 0)
+    throw std::runtime_error(
+        "MNNVL multimem reduce-scatter requires a non-empty input");
+  int size_bytes = size * sizeof(T);
+  if (size_bytes % (kMnnvlMultimemRsVectorBytes * world_size_) != 0)
+    throw std::runtime_error(
+        "MNNVL multimem reduce-scatter requires each output shard byte size "
+        "to be a multiple of 16");
+  if (block_limit <= 0 || block_limit > kMnnvlMultimemRsBlockLimit)
+    throw std::runtime_error(
+        "MNNVL multimem reduce-scatter block limit is out of range");
+
+  int packs_per_rank = size_bytes / kMnnvlMultimemRsVectorBytes / world_size_;
+  int packs_per_block = kMnnvlMultimemRsThreads * kMnnvlMultimemRsUnroll;
+  int blocks = std::min(
+      block_limit, (packs_per_rank + packs_per_block - 1) / packs_per_block);
+
+  #define MNNVL_MULTIMEM_RS_CASE(ngpus)                                 \
+    case ngpus:                                                         \
+      mnnvl_multimem_barrier_kernel<true, ngpus>                        \
+          <<<1, 1, 0, stream>>>(local_signal, multicast_signal, rank_); \
+      mnnvl_multimem_reduce_scatter_kernel<T, ngpus>                    \
+          <<<blocks, kMnnvlMultimemRsThreads, 0, stream>>>(             \
+              multicast_input, output, rank_, packs_per_rank);          \
+      mnnvl_multimem_barrier_kernel<false, ngpus>                       \
+          <<<1, 1, 0, stream>>>(local_signal, multicast_signal, rank_); \
+      break;
+
+  switch (world_size_) {
+    MNNVL_MULTIMEM_RS_CASE(2)
+    MNNVL_MULTIMEM_RS_CASE(4)
+    MNNVL_MULTIMEM_RS_CASE(8)
+    default:
+      throw std::runtime_error(
+          "MNNVL multimem reduce-scatter only supports num gpus in (2,4,8)");
+  }
+  #undef MNNVL_MULTIMEM_RS_CASE
+#endif
+}
+
 }  // namespace vllm
 
 using fptr_t = int64_t;
@@ -360,3 +409,62 @@ void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
           "bfloat16");
   }
 }
+
+void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                                   torch::stable::Tensor& out,
+                                   fptr_t _local_buffer,
+                                   fptr_t _multicast_buffer,
+                                   int64_t stage_sz_bytes,
+                                   int64_t block_limit) {
+  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      inp.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(inp.get_device_index());
+
+  STD_TORCH_CHECK((inp.scalar_type()) == (out.scalar_type()));
+  STD_TORCH_CHECK((out.numel() * fa->world_size_) == (inp.numel()));
+  STD_TORCH_CHECK(inp.numel() > 0);
+  STD_TORCH_CHECK(_is_weak_contiguous(out));
+  STD_TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  STD_TORCH_CHECK(input_size <= stage_sz_bytes);
+  auto local_buffer = reinterpret_cast<void*>(_local_buffer);
+  auto multicast_buffer = reinterpret_cast<void*>(_multicast_buffer);
+  STD_TORCH_CHECK(local_buffer != nullptr);
+  STD_TORCH_CHECK(multicast_buffer != nullptr);
+  auto* local_signal = reinterpret_cast<vllm::Signal*>(
+      static_cast<char*>(local_buffer) - sizeof(vllm::Signal));
+  auto* multicast_signal = reinterpret_cast<vllm::Signal*>(
+      static_cast<char*>(multicast_buffer) - sizeof(vllm::Signal));
+  STD_CUDA_CHECK(cudaMemcpyAsync(local_buffer, inp.const_data_ptr(), input_size,
+                                 cudaMemcpyDeviceToDevice, stream));
+  switch (out.scalar_type()) {
+    case torch::headeronly::ScalarType::Float: {
+      fa->mnnvl_multimem_reduce_scatter<float>(
+          stream, reinterpret_cast<float*>(multicast_buffer),
+          reinterpret_cast<float*>(out.mutable_data_ptr()), local_signal,
+          multicast_signal, inp.numel(), block_limit);
+      break;
+    }
+    case torch::headeronly::ScalarType::Half: {
+      fa->mnnvl_multimem_reduce_scatter<half>(
+          stream, reinterpret_cast<half*>(multicast_buffer),
+          reinterpret_cast<half*>(out.mutable_data_ptr()), local_signal,
+          multicast_signal, inp.numel(), block_limit);
+      break;
+    }
+#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
+    case torch::headeronly::ScalarType::BFloat16: {
+      fa->mnnvl_multimem_reduce_scatter<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(multicast_buffer),
+          reinterpret_cast<nv_bfloat16*>(out.mutable_data_ptr()), local_signal,
+          multicast_signal, inp.numel(), block_limit);
+      break;
+    }
+#endif
+    default:
+      throw std::runtime_error(
+          "MNNVL multimem reduce-scatter only supports float32, float16 and "
+          "bfloat16");
+  }
+}

csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp

@@ -17,6 +17,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_custom_ar, custom_ag_rs) {
   custom_ag_rs.def(
       "mnnvl_lamport_reduce_scatter(int fa, Tensor inp, Tensor! out, int "
       "local_buffer, int epoch_buffer, int stage_sz_bytes) -> ()");
+  custom_ag_rs.def(
+      "mnnvl_multimem_reduce_scatter(int fa, Tensor inp, Tensor! out, int "
+      "local_buffer, int multicast_buffer, int stage_sz_bytes, int "
+      "block_limit) -> ()");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) {
@@ -26,4 +30,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C_custom_ar, CUDA, custom_ag_rs) {
   custom_ag_rs.impl("custom_reduce_scatter", TORCH_BOX(&custom_reduce_scatter));
   custom_ag_rs.impl("mnnvl_lamport_reduce_scatter",
                     TORCH_BOX(&mnnvl_lamport_reduce_scatter));
+  custom_ag_rs.impl("mnnvl_multimem_reduce_scatter",
+                    TORCH_BOX(&mnnvl_multimem_reduce_scatter));
 }

csrc/libtorch_stable/ops.h

@@ -527,6 +527,10 @@ void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
                                   torch::stable::Tensor& out,
                                   fptr_t local_buffer, fptr_t epoch_buffer,
                                   int64_t stage_sz_bytes);
+void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
+                                   torch::stable::Tensor& out,
+                                   fptr_t local_buffer, fptr_t multicast_buffer,
+                                   int64_t stage_sz_bytes, int64_t block_limit);
 void dispose(fptr_t _fa);
 int64_t meta_size();
 void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);

tests/distributed/test_custom_all_gather_reduce_scatter.py

@@ -23,6 +23,7 @@ import torch
 import torch.distributed as dist
 
 from vllm.distributed.parallel_state import get_tp_group
+from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter
 from vllm.platforms import current_platform
 
 from ..utils import (
@@ -53,6 +54,11 @@ def _supports_multimem():
     )
 
 
+def _supports_sm100_or_sm103():
+    capability = current_platform.get_device_capability()
+    return current_platform.is_cuda() and capability in ((10, 0), (10, 3))
+
+
 def _payload(value, device):
     return torch.full((_ROWS, _HIDDEN), value, dtype=torch.bfloat16, device=device)
 
@@ -324,3 +330,65 @@ def test_mnnvl_lamport_stage_cleanup(monkeypatch: pytest.MonkeyPatch):
     if torch.accelerator.device_count() < 2:
         pytest.skip("Need at least two GPUs to run the test.")
     multi_process_parallel(monkeypatch, 2, 1, _run_stage_cleanup_test)
+
+
+@ray.remote(num_gpus=1, max_calls=1)
+def _run_multimem_reduce_scatter_test(
+    monkeypatch: pytest.MonkeyPatch,
+    tp_size,
+    pp_size,
+    rank,
+    distributed_init_port,
+):
+    with monkeypatch.context() as m:
+        m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+        m.delenv("HIP_VISIBLE_DEVICES", raising=False)
+        device = torch.device(f"cuda:{rank}")
+        torch.accelerator.set_device_index(device)
+        init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
+        ensure_model_parallel_initialized(tp_size, pp_size)
+
+        fa = get_tp_group().device_communicator.ca_comm
+        assert fa is not None and not fa.disabled
+        assert fa.mnnvl_multimem_rs_supported
+        assert not fa.mnnvl_multimem_rs_initialized
+        assert not fa.mnnvl_multimem_rs_multicast_ptr
+
+        message_bytes = {
+            "mnnvl_lamport": 1 * 1024 * 1024,
+            "mnnvl_multimem": 20 * 1024 * 1024,
+            "nccl": 64 * 1024 * 1024 + 128,
+        }
+        for _ in range(2):
+            for backend, size_bytes in message_bytes.items():
+                shard_elements = size_bytes // torch.bfloat16.itemsize // tp_size
+                inp = torch.full(
+                    (tp_size, shard_elements),
+                    rank + 1,
+                    dtype=torch.bfloat16,
+                    device=device,
+                )
+                selected = fa._select_reduce_scatter_backend(inp)
+                if backend == "nccl":
+                    assert selected is None
+                    out = sp_reduce_scatter(inp)
+                else:
+                    assert selected == backend
+                    out = fa.custom_reduce_scatter(inp)
+                    assert out is not None
+                    if backend == "mnnvl_multimem":
+                        assert fa.mnnvl_multimem_rs_initialized
+                        assert fa.mnnvl_multimem_rs_multicast_ptr
+                expected = tp_size * (tp_size + 1) // 2
+                torch.testing.assert_close(out, torch.full_like(out, expected))
+
+
+@pytest.mark.skipif(
+    not _supports_sm100_or_sm103(),
+    reason="The low-SM MNNVL reduce-scatter path requires SM100 or SM103.",
+)
+@pytest.mark.parametrize("tp_size", [2, 4, 8])
+def test_mnnvl_multimem_reduce_scatter(monkeypatch: pytest.MonkeyPatch, tp_size: int):
+    if torch.accelerator.device_count() < tp_size:
+        pytest.skip(f"Need at least {tp_size} GPUs to run the test.")
+    multi_process_parallel(monkeypatch, tp_size, 1, _run_multimem_reduce_scatter_test)

tests/distributed/test_custom_all_reduce.py

@@ -151,9 +151,9 @@ def test_custom_allreduce_filters_dtype(
 ) -> None:
     communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
     communicator.disabled = False
+    communicator._ptr = 0
     communicator.world_size = 2
     communicator.max_size = 1024
-    communicator._ptr = 0
 
     assert communicator.should_custom_ar(torch.empty(16, dtype=dtype)) is expected
 
@@ -219,6 +219,268 @@ def test_local_multicast_support_rejects_non_cuda(monkeypatch):
     assert not car._has_local_multicast_support(torch.device("cuda:0"))
 
 
+@pytest.mark.parametrize(
+    ("world_size", "device_capability", "local_multicast", "expected"),
+    [
+        (2, (10, 0), True, True),
+        (4, (10, 3), True, True),
+        (8, (10, 0), True, True),
+        (8, (10, 3), True, True),
+        (6, (10, 3), True, False),
+        (8, (10, 1), True, False),
+        (8, (9, 0), True, False),
+        (8, (10, 3), False, False),
+    ],
+)
+def test_mnnvl_multimem_reduce_scatter_platform_gate(
+    monkeypatch,
+    world_size,
+    device_capability,
+    local_multicast,
+    expected,
+):
+    def is_device_capability(capability, device_id):
+        assert capability in ((10, 0), (10, 3))
+        assert device_id == 3
+        return device_capability == capability
+
+    monkeypatch.setattr(
+        car.current_platform,
+        "is_device_capability",
+        is_device_capability,
+    )
+    monkeypatch.setattr(
+        car,
+        "_has_local_multicast_support",
+        lambda _device: local_multicast,
+    )
+
+    supported = car._supports_mnnvl_multimem_reduce_scatter(
+        torch.device("cuda:3"), world_size
+    )
+    assert supported is expected
+
+
+@pytest.mark.parametrize(
+    (
+        "message_bytes",
+        "multimem_ptr",
+        "multimem_initialized",
+        "batch_invariant",
+        "expected",
+    ),
+    [
+        (16 * 1024 * 1024, 1, True, False, "mnnvl_lamport"),
+        (16 * 1024 * 1024 + 128, 1, True, False, "mnnvl_multimem"),
+        (64 * 1024 * 1024, 1, True, False, "mnnvl_multimem"),
+        (64 * 1024 * 1024 + 128, 1, True, False, None),
+        (32 * 1024 * 1024, 0, True, False, None),
+        (32 * 1024 * 1024, 0, False, False, "mnnvl_multimem"),
+        (8 * 1024 * 1024, 1, True, True, "mnnvl_lamport"),
+        (32 * 1024 * 1024, 1, True, True, None),
+    ],
+)
+@pytest.mark.parametrize("world_size", [2, 4, 8])
+def test_mnnvl_reduce_scatter_backend_gate(
+    monkeypatch,
+    world_size,
+    message_bytes,
+    multimem_ptr,
+    multimem_initialized,
+    batch_invariant,
+    expected,
+):
+    monkeypatch.setattr(car.current_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(car.envs, "VLLM_BATCH_INVARIANT", batch_invariant)
+    communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
+    communicator.disabled = False
+    communicator._ptr = 0
+    communicator.world_size = world_size
+    communicator.mnnvl_only = False
+    communicator.fully_connected = True
+    communicator.mnnvl_multicast_ptr = 1
+    communicator.mnnvl_multimem_rs_supported = True
+    communicator.mnnvl_multimem_rs_initialized = multimem_initialized
+    communicator.mnnvl_multimem_rs_multicast_ptr = multimem_ptr
+    communicator.max_mnnvl_reduce_scatter_size = 16 * 1024 * 1024
+    communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024
+    communicator.max_reduce_scatter_size = 16 * 1024 * 1024
+    inp = torch.empty(
+        (world_size, message_bytes // torch.bfloat16.itemsize // world_size),
+        dtype=torch.bfloat16,
+    )
+
+    assert inp.nbytes == message_bytes
+    assert communicator._select_reduce_scatter_backend(inp) == expected
+    assert communicator.should_custom_reduce_scatter(inp) is (expected is not None)
+    assert communicator.should_mnnvl_multimem_reduce_scatter(inp) is (
+        expected == "mnnvl_multimem"
+    )
+
+
+def test_mnnvl_multimem_reduce_scatter_skips_rendezvous_after_peer_alloc_failure(
+    monkeypatch,
+):
+    events = []
+
+    class FakeSymmMem:
+        @staticmethod
+        def empty(*_args, **_kwargs):
+            events.append("empty")
+            return torch.empty(1, dtype=torch.uint8)
+
+        @staticmethod
+        def rendezvous(*_args, **_kwargs):
+            events.append("rendezvous")
+            return None
+
+    def report_peer_allocation_failure(group_value, **_kwargs):
+        events.append("all_reduce")
+        assert group_value.item() == 1
+        group_value.zero_()
+
+    monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem)
+    monkeypatch.setattr(car.ops, "meta_size", lambda: 128)
+    monkeypatch.setattr(car.dist, "all_reduce", report_peer_allocation_failure)
+    warnings = []
+    monkeypatch.setattr(
+        car.logger,
+        "warning_once",
+        lambda message, *_args, **_kwargs: warnings.append(message),
+    )
+
+    communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
+    communicator.disabled = True
+    communicator._ptr = 0
+    communicator.group = object()
+    communicator.device = torch.device("cuda:0")
+    communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024
+    communicator.mnnvl_multimem_rs_supported = True
+    communicator.mnnvl_multimem_rs_initialized = False
+    communicator.mnnvl_multimem_rs_buffer = None
+    communicator.mnnvl_multimem_rs_multicast_ptr = 0
+
+    communicator._init_mnnvl_multimem_reduce_scatter_buffer()
+
+    assert events == ["empty", "all_reduce"]
+    assert communicator.mnnvl_multimem_rs_initialized
+    assert communicator.mnnvl_multimem_rs_buffer is None
+    assert communicator.mnnvl_multimem_rs_multicast_ptr == 0
+    assert warnings == [
+        "MNNVL multimem reduce-scatter symmetric-memory allocation "
+        "failed on at least one rank; falling back to NCCL."
+    ]
+
+
+def test_mnnvl_multimem_reduce_scatter_warns_on_rendezvous_failure(monkeypatch):
+    events = []
+
+    class FakeSymmMem:
+        @staticmethod
+        def empty(*_args, **_kwargs):
+            events.append("empty")
+            return torch.empty(1, dtype=torch.uint8)
+
+        @staticmethod
+        def rendezvous(*_args, **_kwargs):
+            events.append("rendezvous")
+            raise RuntimeError("rendezvous failed")
+
+    def preserve_local_result(_group_value, **_kwargs):
+        events.append("all_reduce")
+
+    warnings = []
+    monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem)
+    monkeypatch.setattr(car.ops, "meta_size", lambda: 128)
+    monkeypatch.setattr(car.dist, "all_reduce", preserve_local_result)
+    monkeypatch.setattr(
+        car.logger,
+        "warning_once",
+        lambda message, *_args, **_kwargs: warnings.append(message),
+    )
+
+    communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
+    communicator.disabled = True
+    communicator._ptr = 0
+    communicator.group = type("Group", (), {"group_name": "test"})()
+    communicator.device = torch.device("cuda:0")
+    communicator.max_mnnvl_multimem_reduce_scatter_size = 64 * 1024 * 1024
+    communicator.mnnvl_multimem_rs_supported = True
+    communicator.mnnvl_multimem_rs_initialized = False
+    communicator.mnnvl_multimem_rs_buffer = None
+    communicator.mnnvl_multimem_rs_multicast_ptr = 0
+
+    communicator._init_mnnvl_multimem_reduce_scatter_buffer()
+
+    assert events == ["empty", "all_reduce", "rendezvous", "all_reduce"]
+    assert communicator.mnnvl_multimem_rs_initialized
+    assert communicator.mnnvl_multimem_rs_buffer is None
+    assert communicator.mnnvl_multimem_rs_multicast_ptr == 0
+    assert warnings == [
+        "MNNVL multimem reduce-scatter symmetric-memory rendezvous "
+        "failed on at least one rank; falling back to NCCL."
+    ]
+
+
+def test_mnnvl_multimem_reduce_scatter_initializes_signals(monkeypatch):
+    events = []
+    buffers = []
+
+    class FakeHandle:
+        multicast_ptr = 0x3000
+
+    class FakeSymmMem:
+        @staticmethod
+        def empty(size, **_kwargs):
+            events.append(("empty", size))
+            buffer = torch.ones(size, dtype=torch.uint8)
+            buffers.append(buffer)
+            return buffer
+
+        @staticmethod
+        def rendezvous(*_args, **_kwargs):
+            events.append(("rendezvous", None))
+            return FakeHandle()
+
+    def preserve_local_result(*_args, **_kwargs):
+        events.append(("all_reduce", None))
+
+    monkeypatch.setattr(car, "torch_symm_mem", FakeSymmMem)
+    monkeypatch.setattr(car.ops, "meta_size", lambda: 128)
+    monkeypatch.setattr(
+        car.torch.accelerator,
+        "synchronize",
+        lambda: events.append(("synchronize", None)),
+    )
+    monkeypatch.setattr(car.dist, "all_reduce", preserve_local_result)
+
+    communicator = car.CustomAllreduce.__new__(car.CustomAllreduce)
+    communicator.disabled = True
+    communicator._ptr = 0
+    communicator.group = type("Group", (), {"group_name": "test"})()
+    communicator.device = torch.device("cpu")
+    communicator.max_mnnvl_multimem_reduce_scatter_size = 129
+    communicator.mnnvl_multimem_rs_supported = True
+    communicator.mnnvl_multimem_rs_initialized = False
+    communicator.mnnvl_multimem_rs_buffer = None
+    communicator.mnnvl_multimem_rs_multicast_ptr = 0
+
+    communicator._init_mnnvl_multimem_reduce_scatter_buffer()
+
+    assert events == [
+        ("empty", 257),
+        ("all_reduce", None),
+        ("rendezvous", None),
+        ("synchronize", None),
+        ("all_reduce", None),
+    ]
+    assert torch.all(buffers[0][:128] == 0)
+    assert torch.all(buffers[0][128:] == 1)
+    assert communicator.mnnvl_multimem_rs_buffer_size == 129
+    assert communicator.mnnvl_multimem_rs_local_ptr == buffers[0].data_ptr() + 128
+    assert communicator.mnnvl_multimem_rs_multicast_ptr == 0x3080
+
+
 @ray.remote(num_gpus=1, max_calls=1)
 def graph_allreduce(
     monkeypatch: pytest.MonkeyPatch,

vllm/_custom_ops.py

@@ -3159,6 +3159,26 @@ def mnnvl_lamport_reduce_scatter(
     )
 
 
+def mnnvl_multimem_reduce_scatter(
+    fa: int,
+    inp: torch.Tensor,
+    out: torch.Tensor,
+    local_buffer: int,
+    multicast_buffer: int,
+    stage_sz_bytes: int,
+    block_limit: int,
+) -> None:
+    torch.ops._C_custom_ar.mnnvl_multimem_reduce_scatter(
+        fa,
+        inp,
+        out,
+        local_buffer,
+        multicast_buffer,
+        stage_sz_bytes,
+        block_limit,
+    )
+
+
 def dispose(fa: int) -> None:
     torch.ops._C_custom_ar.dispose(fa)
 

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from contextlib import contextmanager
-from typing import cast
+from typing import Literal, cast
 
 import torch
 import torch.distributed as dist
@@ -32,6 +32,9 @@ except ImportError:
 
 logger = init_logger(__name__)
 
+_ReduceScatterBackend = Literal["legacy", "mnnvl_lamport", "mnnvl_multimem"]
+_MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES = (2, 4, 8)
+
 
 def _has_local_multicast_support(device: torch.device) -> bool:
     """Return whether this CUDA device can allocate multicast memory."""
@@ -53,6 +56,12 @@ def _has_local_multicast_support(device: torch.device) -> bool:
         return False
 
 
+def _all_ranks_true(group: ProcessGroup, local_value: bool) -> bool:
+    group_value = torch.tensor(int(local_value), dtype=torch.int32, device="cpu")
+    dist.all_reduce(group_value, op=dist.ReduceOp.MIN, group=group)
+    return bool(group_value.item())
+
+
 def _group_can_attempt_mnnvl(
     group: ProcessGroup,
     device: torch.device,
@@ -71,17 +80,7 @@ def _group_can_attempt_mnnvl(
         and current_platform.has_device_capability(100, device_index)
         and _has_local_multicast_support(device)
     )
-    group_support = torch.tensor(
-        int(local_support),
-        dtype=torch.int32,
-        device="cpu",
-    )
-    dist.all_reduce(
-        group_support,
-        op=dist.ReduceOp.MIN,
-        group=group,
-    )
-    return bool(group_support.item())
+    return _all_ranks_true(group, local_support)
 
 
 def _can_p2p(rank: int, world_size: int) -> bool:
@@ -101,6 +100,20 @@ def _can_p2p(rank: int, world_size: int) -> bool:
     return True
 
 
+def _supports_mnnvl_multimem_reduce_scatter(
+    device: torch.device, world_size: int
+) -> bool:
+    return (
+        world_size in _MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES
+        and device.index is not None
+        and (
+            current_platform.is_device_capability((10, 0), device.index)
+            or current_platform.is_device_capability((10, 3), device.index)
+        )
+        and _has_local_multicast_support(device)
+    )
+
+
 from vllm.distributed.utils import is_weak_contiguous  # noqa: E402
 
 
@@ -116,6 +129,8 @@ class CustomAllreduce:
     }
     _DEFAULT_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
     _DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE = 16 * 1024 * 1024
+    _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE = 64 * 1024 * 1024
+    _MNNVL_MULTIMEM_REDUCE_SCATTER_BLOCKS = 8
 
     # max_size: max supported allreduce size
     def __init__(
@@ -127,6 +142,9 @@ class CustomAllreduce:
         max_mnnvl_all_gather_size=None,
         max_reduce_scatter_size=_DEFAULT_REDUCE_SCATTER_MAX_SIZE,
         max_mnnvl_reduce_scatter_size=_DEFAULT_MNNVL_REDUCE_SCATTER_MAX_SIZE,
+        max_mnnvl_multimem_reduce_scatter_size=(
+            _DEFAULT_MNNVL_MULTIMEM_REDUCE_SCATTER_MAX_SIZE
+        ),
         symm_mem_enabled=False,
     ) -> None:
         """Args:
@@ -153,6 +171,13 @@ class CustomAllreduce:
         self.mnnvl_lamport_epochs = None
         self.mnnvl_lamport_ag_epoch_ptr = 0
         self.mnnvl_lamport_rs_epoch_ptr = 0
+        self.mnnvl_multimem_rs_supported = False
+        self.mnnvl_multimem_rs_initialized = False
+        self.mnnvl_multimem_rs_buffer = None
+        self.mnnvl_multimem_rs_handle = None
+        self.mnnvl_multimem_rs_buffer_size = 0
+        self.mnnvl_multimem_rs_local_ptr = 0
+        self.mnnvl_multimem_rs_multicast_ptr = 0
         self.mnnvl_only = False
 
         if not custom_ar:
@@ -197,6 +222,9 @@ class CustomAllreduce:
         # now `device` is a `torch.device` object
         assert isinstance(device, torch.device)
         self.device = device
+        mnnvl_multimem_rs_supported = _supports_mnnvl_multimem_reduce_scatter(
+            device, world_size
+        )
         if not same_node and not _group_can_attempt_mnnvl(group, device):
             logger.warning(
                 "Custom collectives are disabled because this multi-node "
@@ -294,6 +322,9 @@ class CustomAllreduce:
         self.max_mnnvl_all_gather_size = max_mnnvl_all_gather_size
         self.max_reduce_scatter_size = max_reduce_scatter_size
         self.max_mnnvl_reduce_scatter_size = max_mnnvl_reduce_scatter_size
+        self.max_mnnvl_multimem_reduce_scatter_size = (
+            max_mnnvl_multimem_reduce_scatter_size
+        )
         self.rank = rank
         self.world_size = world_size
         self.fully_connected = fully_connected
@@ -307,6 +338,11 @@ class CustomAllreduce:
                 max_mnnvl_reduce_scatter_size,
             )
         )
+        if world_size in _MNNVL_MULTIMEM_REDUCE_SCATTER_WORLD_SIZES:
+            self.mnnvl_multimem_rs_supported = _all_ranks_true(
+                self.group,
+                mnnvl_multimem_rs_supported and bool(self.mnnvl_multicast_ptr),
+            )
         if not same_node and not self.mnnvl_multicast_ptr:
             logger.warning(
                 "Custom collectives are disabled because this multi-node "
@@ -368,6 +404,59 @@ class CustomAllreduce:
         except RuntimeError as error:
             logger.debug("MNNVL AG/RS initialization failed: %s", error)
 
+    def _init_mnnvl_multimem_reduce_scatter_buffer(self) -> None:
+        if self.mnnvl_multimem_rs_initialized:
+            return
+        self.mnnvl_multimem_rs_initialized = True
+        if not self.mnnvl_multimem_rs_supported:
+            return
+
+        assert torch_symm_mem is not None
+        buffer = None
+        signal_size = ops.meta_size()
+        try:
+            buffer = torch_symm_mem.empty(
+                signal_size + self.max_mnnvl_multimem_reduce_scatter_size,
+                dtype=torch.uint8,
+                device=self.device,
+            )
+        except RuntimeError as error:
+            logger.debug("MNNVL multimem RS allocation failed: %s", error)
+        if not _all_ranks_true(self.group, buffer is not None):
+            logger.warning_once(
+                "MNNVL multimem reduce-scatter symmetric-memory allocation "
+                "failed on at least one rank; falling back to NCCL.",
+                scope="global",
+            )
+            return
+        assert buffer is not None
+
+        handle = None
+        try:
+            handle = torch_symm_mem.rendezvous(buffer, self.group.group_name)
+            if handle is not None:
+                buffer[:signal_size].zero_()
+                torch.accelerator.synchronize()
+        except RuntimeError as error:
+            logger.debug("MNNVL multimem RS rendezvous failed: %s", error)
+        if not _all_ranks_true(
+            self.group,
+            handle is not None and bool(handle.multicast_ptr),
+        ):
+            logger.warning_once(
+                "MNNVL multimem reduce-scatter symmetric-memory rendezvous "
+                "failed on at least one rank; falling back to NCCL.",
+                scope="global",
+            )
+            return
+        assert handle is not None
+
+        self.mnnvl_multimem_rs_buffer = buffer
+        self.mnnvl_multimem_rs_handle = handle
+        self.mnnvl_multimem_rs_buffer_size = self.max_mnnvl_multimem_reduce_scatter_size
+        self.mnnvl_multimem_rs_local_ptr = buffer.data_ptr() + signal_size
+        self.mnnvl_multimem_rs_multicast_ptr = handle.multicast_ptr + signal_size
+
     @contextmanager
     def capture(self):
         """The main responsibility of this context manager is the
@@ -509,35 +598,77 @@ class CustomAllreduce:
             )
         return out
 
-    def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool:
+    def _select_reduce_scatter_backend(
+        self, inp: torch.Tensor
+    ) -> _ReduceScatterBackend | None:
         if self.disabled or not current_platform.is_cuda():
-            return False
+            return None
         if self.world_size == 16 and not self.mnnvl_only:
-            return False
+            return None
         inp_size = inp.nbytes
         if inp.dtype not in (torch.float32, torch.float16, torch.bfloat16):
-            return False
+            return None
         if inp.shape[0] % self.world_size != 0:
-            return False
+            return None
         output_size = inp_size // self.world_size
-        max_size = (
-            self.max_mnnvl_reduce_scatter_size
-            if self.mnnvl_multicast_ptr
-            else self.max_reduce_scatter_size
-        )
-        return (
-            0 < inp_size <= max_size
-            and output_size % 16 == 0
-            and is_weak_contiguous(inp)
-            and (self.fully_connected or bool(self.mnnvl_multicast_ptr))
-        )
+        if inp_size <= 0 or output_size % 16 != 0 or not is_weak_contiguous(inp):
+            return None
+
+        if self.mnnvl_multicast_ptr:
+            if inp_size <= self.max_mnnvl_reduce_scatter_size:
+                return "mnnvl_lamport"
+            if (
+                (
+                    self.mnnvl_multimem_rs_multicast_ptr
+                    or (
+                        self.mnnvl_multimem_rs_supported
+                        and not self.mnnvl_multimem_rs_initialized
+                    )
+                )
+                and not envs.VLLM_BATCH_INVARIANT
+                and inp_size <= self.max_mnnvl_multimem_reduce_scatter_size
+            ):
+                return "mnnvl_multimem"
+            return None
+
+        if self.fully_connected and inp_size <= self.max_reduce_scatter_size:
+            return "legacy"
+        return None
+
+    def should_custom_reduce_scatter(self, inp: torch.Tensor) -> bool:
+        return self._select_reduce_scatter_backend(inp) is not None
+
+    def should_mnnvl_multimem_reduce_scatter(self, inp: torch.Tensor) -> bool:
+        return self._select_reduce_scatter_backend(inp) == "mnnvl_multimem"
 
     def custom_reduce_scatter(self, inp: torch.Tensor) -> torch.Tensor | None:
-        if not self.should_custom_reduce_scatter(inp):
+        backend = self._select_reduce_scatter_backend(inp)
+        if backend is None:
             return None
+        if backend == "mnnvl_multimem" and not self.mnnvl_multimem_rs_multicast_ptr:
+            if self._IS_CAPTURING:
+                return None
+            self._init_mnnvl_multimem_reduce_scatter_buffer()
+            if not self.mnnvl_multimem_rs_multicast_ptr:
+                return None
+
         out_shape = (inp.shape[0] // self.world_size,) + inp.shape[1:]
         out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
-        if self.mnnvl_multicast_ptr:
+        if backend == "mnnvl_multimem":
+            logger.info_once(
+                "Using the low-SM MNNVL multimem reduce-scatter kernel.",
+                scope="global",
+            )
+            ops.mnnvl_multimem_reduce_scatter(
+                self._ptr,
+                inp,
+                out,
+                self.mnnvl_multimem_rs_local_ptr,
+                self.mnnvl_multimem_rs_multicast_ptr,
+                self.mnnvl_multimem_rs_buffer_size,
+                self._MNNVL_MULTIMEM_REDUCE_SCATTER_BLOCKS,
+            )
+        elif backend == "mnnvl_lamport":
             logger.info_once(
                 "Using the MNNVL Lamport reduce-scatter kernel.",
                 scope="global",
@@ -571,6 +702,8 @@ class CustomAllreduce:
             self.mnnvl_handle = None
             self.mnnvl_buffer = None
             self.mnnvl_lamport_epochs = None
+            self.mnnvl_multimem_rs_handle = None
+            self.mnnvl_multimem_rs_buffer = None
 
     def __del__(self):
         self.close()