diff --git a/benchmarks/kernels/benchmark_multimem_reduce_scatter.py b/benchmarks/kernels/benchmark_multimem_reduce_scatter.py new file mode 100644 index 00000000000..0f0966eff44 --- /dev/null +++ b/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() diff --git a/csrc/custom_all_gather_reduce_scatter.cuh b/csrc/custom_all_gather_reduce_scatter.cuh index 0190af728eb..d70ad75d258 100644 --- a/csrc/custom_all_gather_reduce_scatter.cuh +++ b/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; @@ -48,6 +52,127 @@ __global__ void __launch_bounds__(512, 1) barrier_at_end(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 +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::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::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::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 +__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 +__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(multicast_input) + + rank * packs_per_rank * kMnnvlMultimemRsVectorBytes; + auto* rank_output = reinterpret_cast(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( + reduced[u], reinterpret_cast( + 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 union LamportPack { P packed; diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 385a0e0ae79..4bfa3b71dc5 100644 --- a/csrc/custom_all_reduce.cuh +++ b/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 + 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_) { diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu b/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu index 88aed49d181..2e360e90b55 100644 --- a/csrc/libtorch_stable/custom_all_gather_reduce_scatter.cu +++ b/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 +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 \ + <<<1, 1, 0, stream>>>(local_signal, multicast_signal, rank_); \ + mnnvl_multimem_reduce_scatter_kernel \ + <<>>( \ + multicast_input, output, rank_, packs_per_rank); \ + mnnvl_multimem_barrier_kernel \ + <<<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(_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(_local_buffer); + auto multicast_buffer = reinterpret_cast(_multicast_buffer); + STD_TORCH_CHECK(local_buffer != nullptr); + STD_TORCH_CHECK(multicast_buffer != nullptr); + auto* local_signal = reinterpret_cast( + static_cast(local_buffer) - sizeof(vllm::Signal)); + auto* multicast_signal = reinterpret_cast( + static_cast(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( + stream, reinterpret_cast(multicast_buffer), + reinterpret_cast(out.mutable_data_ptr()), local_signal, + multicast_signal, inp.numel(), block_limit); + break; + } + case torch::headeronly::ScalarType::Half: { + fa->mnnvl_multimem_reduce_scatter( + stream, reinterpret_cast(multicast_buffer), + reinterpret_cast(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( + stream, reinterpret_cast(multicast_buffer), + reinterpret_cast(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"); + } +} diff --git a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp b/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp index 198b78d56da..bf7d3f8abdc 100644 --- a/csrc/libtorch_stable/custom_all_gather_reduce_scatter_ops.cpp +++ b/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)); } diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 331cd1fbb87..fba4220d735 100644 --- a/csrc/libtorch_stable/ops.h +++ b/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& fake_ipc_ptrs); diff --git a/tests/distributed/test_custom_all_gather_reduce_scatter.py b/tests/distributed/test_custom_all_gather_reduce_scatter.py index dde5fdf7ff0..61666110e5e 100644 --- a/tests/distributed/test_custom_all_gather_reduce_scatter.py +++ b/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) diff --git a/tests/distributed/test_custom_all_reduce.py b/tests/distributed/test_custom_all_reduce.py index 7119e596a09..1d95801f0bd 100644 --- a/tests/distributed/test_custom_all_reduce.py +++ b/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, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index f08be7b66f5..be9a111a583 100644 --- a/vllm/_custom_ops.py +++ b/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) diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index e18f9194f54..b6fb540c7de 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/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()