| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark production MNNVL reduce-scatter backends on SM100/SM103. |
| 4 | |
| 5 | The multimem candidate includes the device-to-device copy into its persistent |
| 6 | symmetric input buffer. Timings use pointer-distinct CUDA graphs and report the |
| 7 | median of the per-sample maximum across ranks. |
| 8 | |
| 9 | Run the default BF16 message-size sweep with: |
| 10 | |
| 11 | .venv/bin/python -m torch.distributed.run --standalone --nproc-per-node=8 \ |
| 12 | benchmarks/kernels/benchmark_multimem_reduce_scatter.py |
| 13 | """ |
| 14 | |
| 15 | import argparse |
| 16 | import json |
| 17 | import os |
| 18 | import statistics |
| 19 | from collections.abc import Callable |
| 20 | from dataclasses import dataclass |
| 21 | |
| 22 | import torch |
| 23 | import torch.distributed as dist |
| 24 | |
| 25 | import vllm._custom_ops as ops |
| 26 | from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce |
| 27 | |
| 28 | _DTYPES = { |
| 29 | "bf16": torch.bfloat16, |
| 30 | "f16": torch.float16, |
| 31 | "f32": torch.float32, |
| 32 | } |
| 33 | |
| 34 | |
| 35 | @dataclass |
| 36 | class Candidate: |
| 37 | name: str |
| 38 | runs: list[Callable[[], None]] |
| 39 | outputs: list[torch.Tensor] |
| 40 | |
| 41 | |
| 42 | def parse_size(value: str) -> int: |
| 43 | suffixes = {"k": 1 << 10, "m": 1 << 20, "g": 1 << 30} |
| 44 | value = value.strip().lower() |
| 45 | multiplier = suffixes.get(value[-1], 1) |
| 46 | number = value[:-1] if multiplier != 1 else value |
| 47 | return int(float(number) * multiplier) |
| 48 | |
| 49 | |
| 50 | def parse_args() -> argparse.Namespace: |
| 51 | parser = argparse.ArgumentParser() |
| 52 | parser.add_argument( |
| 53 | "--message-bytes", |
| 54 | type=parse_size, |
| 55 | nargs="+", |
| 56 | default=[ |
| 57 | 16 << 20, |
| 58 | 32 << 20, |
| 59 | 64 << 20, |
| 60 | 68 << 20, |
| 61 | 72 << 20, |
| 62 | 76 << 20, |
| 63 | 96 << 20, |
| 64 | ], |
| 65 | help="Full input bytes on each rank. K/M/G suffixes are accepted.", |
| 66 | ) |
| 67 | parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16") |
| 68 | parser.add_argument( |
| 69 | "--blocks", |
| 70 | type=int, |
| 71 | nargs="+", |
| 72 | default=[8], |
| 73 | help="Maximum CTAs used by each multimem candidate.", |
| 74 | ) |
| 75 | parser.add_argument("--num-workspaces", type=int, default=4) |
| 76 | parser.add_argument("--warmup-replays", type=int, default=8) |
| 77 | parser.add_argument("--samples", type=int, default=31) |
| 78 | parser.add_argument( |
| 79 | "--cold-l2", |
| 80 | action="store_true", |
| 81 | help="Flush approximately twice L2 capacity before each replay.", |
| 82 | ) |
| 83 | return parser.parse_args() |
| 84 | |
| 85 | |
| 86 | def make_lamport_run( |
| 87 | comm: CustomAllreduce, |
| 88 | input_tensor: torch.Tensor, |
| 89 | output: torch.Tensor, |
| 90 | ) -> Callable[[], None]: |
| 91 | def run() -> None: |
| 92 | ops.mnnvl_lamport_reduce_scatter( |
| 93 | comm._ptr, |
| 94 | input_tensor, |
| 95 | output, |
| 96 | comm.mnnvl_lamport_rs_local_ptr, |
| 97 | comm.mnnvl_lamport_rs_epoch_ptr, |
| 98 | comm.mnnvl_buffer_size, |
| 99 | ) |
| 100 | |
| 101 | return run |
| 102 | |
| 103 | |
| 104 | def make_multimem_run( |
| 105 | comm: CustomAllreduce, |
| 106 | input_tensor: torch.Tensor, |
| 107 | output: torch.Tensor, |
| 108 | blocks: int, |
| 109 | ) -> Callable[[], None]: |
| 110 | def run() -> None: |
| 111 | ops.mnnvl_multimem_reduce_scatter( |
| 112 | comm._ptr, |
| 113 | input_tensor, |
| 114 | output, |
| 115 | comm.mnnvl_multimem_rs_local_ptr, |
| 116 | comm.mnnvl_multimem_rs_multicast_ptr, |
| 117 | comm.mnnvl_multimem_rs_buffer_size, |
| 118 | blocks, |
| 119 | ) |
| 120 | |
| 121 | return run |
| 122 | |
| 123 | |
| 124 | def make_nccl_run( |
| 125 | input_tensor: torch.Tensor, |
| 126 | output: torch.Tensor, |
| 127 | device_group: dist.ProcessGroup, |
| 128 | ) -> Callable[[], None]: |
| 129 | def run() -> None: |
| 130 | dist.reduce_scatter_single(output, input_tensor, group=device_group) |
| 131 | |
| 132 | return run |
| 133 | |
| 134 | |
| 135 | def build_candidates( |
| 136 | comm: CustomAllreduce, |
| 137 | inputs: list[torch.Tensor], |
| 138 | local_elements: int, |
| 139 | blocks_values: list[int], |
| 140 | device_group: dist.ProcessGroup, |
| 141 | ) -> list[Candidate]: |
| 142 | candidates: list[Candidate] = [] |
| 143 | |
| 144 | lamport_outputs = [ |
| 145 | input_tensor.new_empty(local_elements) for input_tensor in inputs |
| 146 | ] |
| 147 | candidates.append( |
| 148 | Candidate( |
| 149 | "lamport", |
| 150 | [ |
| 151 | make_lamport_run(comm, input_tensor, output) |
| 152 | for input_tensor, output in zip(inputs, lamport_outputs) |
| 153 | ], |
| 154 | lamport_outputs, |
| 155 | ) |
| 156 | ) |
| 157 | |
| 158 | for blocks in blocks_values: |
| 159 | outputs = [input_tensor.new_empty(local_elements) for input_tensor in inputs] |
| 160 | candidates.append( |
| 161 | Candidate( |
| 162 | f"multimem_b{blocks}", |
| 163 | [ |
| 164 | make_multimem_run(comm, input_tensor, output, blocks) |
| 165 | for input_tensor, output in zip(inputs, outputs) |
| 166 | ], |
| 167 | outputs, |
| 168 | ) |
| 169 | ) |
| 170 | |
| 171 | nccl_outputs = [input_tensor.new_empty(local_elements) for input_tensor in inputs] |
| 172 | candidates.append( |
| 173 | Candidate( |
| 174 | "nccl", |
| 175 | [ |
| 176 | make_nccl_run(input_tensor, output, device_group) |
| 177 | for input_tensor, output in zip(inputs, nccl_outputs) |
| 178 | ], |
| 179 | nccl_outputs, |
| 180 | ) |
| 181 | ) |
| 182 | return candidates |
| 183 | |
| 184 | |
| 185 | def check_correctness( |
| 186 | candidates: list[Candidate], |
| 187 | inputs: list[torch.Tensor], |
| 188 | rank: int, |
| 189 | world_size: int, |
| 190 | cpu_group: dist.ProcessGroup, |
| 191 | device_group: dist.ProcessGroup, |
| 192 | graphs: dict[str, list[torch.cuda.CUDAGraph]] | None = None, |
| 193 | ) -> None: |
| 194 | expected_outputs = [] |
| 195 | for input_tensor in inputs: |
| 196 | expected = input_tensor.float().clone() |
| 197 | dist.all_reduce(expected, group=device_group) |
| 198 | expected_outputs.append(expected.to(input_tensor.dtype).chunk(world_size)[rank]) |
| 199 | |
| 200 | if inputs[0].dtype == torch.float32: |
| 201 | rtol, atol = 1e-5, 1e-5 |
| 202 | elif inputs[0].dtype == torch.float16: |
| 203 | rtol, atol = 1e-2, 1e-2 |
| 204 | else: |
| 205 | rtol, atol = 5e-2, 1.25e-1 |
| 206 | |
| 207 | for candidate in candidates: |
| 208 | for workspace_index, expected in enumerate(expected_outputs): |
| 209 | dist.barrier(group=cpu_group) |
| 210 | if graphs is None: |
| 211 | candidate.runs[workspace_index]() |
| 212 | else: |
| 213 | graphs[candidate.name][workspace_index].replay() |
| 214 | torch.accelerator.synchronize() |
| 215 | error = None |
| 216 | try: |
| 217 | torch.testing.assert_close( |
| 218 | candidate.outputs[workspace_index], |
| 219 | expected, |
| 220 | rtol=rtol, |
| 221 | atol=atol, |
| 222 | ) |
| 223 | except AssertionError as assertion_error: |
| 224 | error = assertion_error |
| 225 | correct = torch.tensor( |
| 226 | error is None, |
| 227 | dtype=torch.int32, |
| 228 | device=candidate.outputs[workspace_index].device, |
| 229 | ) |
| 230 | dist.all_reduce(correct, op=dist.ReduceOp.MIN, group=device_group) |
| 231 | if not correct.item(): |
| 232 | location = f"{candidate.name} workspace {workspace_index}" |
| 233 | detail = str(error) if error is not None else "failed on another rank" |
| 234 | raise AssertionError( |
| 235 | f"{location} produced an incorrect shard: {detail}" |
| 236 | ) |
| 237 | dist.barrier(group=cpu_group) |
| 238 | |
| 239 | |
| 240 | def capture_graph( |
| 241 | run: Callable[[], None], |
| 242 | cpu_group: dist.ProcessGroup, |
| 243 | ) -> tuple[torch.cuda.CUDAGraph, torch.cuda.Stream]: |
| 244 | stream = torch.cuda.Stream() |
| 245 | stream.wait_stream(torch.cuda.current_stream()) |
| 246 | dist.barrier(group=cpu_group) |
| 247 | with torch.cuda.stream(stream): |
| 248 | for _ in range(3): |
| 249 | run() |
| 250 | stream.synchronize() |
| 251 | dist.barrier(group=cpu_group) |
| 252 | |
| 253 | graph = torch.cuda.CUDAGraph() |
| 254 | with torch.cuda.graph(graph, stream=stream): |
| 255 | run() |
| 256 | torch.cuda.current_stream().wait_stream(stream) |
| 257 | dist.barrier(group=cpu_group) |
| 258 | return graph, stream |
| 259 | |
| 260 | |
| 261 | def capture_candidates( |
| 262 | candidates: list[Candidate], |
| 263 | cpu_group: dist.ProcessGroup, |
| 264 | ) -> tuple[dict[str, list[torch.cuda.CUDAGraph]], list[object]]: |
| 265 | graphs = {} |
| 266 | keepalive: list[object] = [] |
| 267 | for candidate in candidates: |
| 268 | bundles = [capture_graph(run, cpu_group) for run in candidate.runs] |
| 269 | graphs[candidate.name] = [graph for graph, _ in bundles] |
| 270 | keepalive.extend(bundles) |
| 271 | return graphs, keepalive |
| 272 | |
| 273 | |
| 274 | def benchmark_graphs( |
| 275 | graphs: dict[str, list[torch.cuda.CUDAGraph]], |
| 276 | warmup_replays: int, |
| 277 | samples: int, |
| 278 | device_group: dist.ProcessGroup, |
| 279 | device_barrier: Callable[[], None], |
| 280 | flush: torch.Tensor | None, |
| 281 | ) -> dict[str, float]: |
| 282 | names = list(graphs) |
| 283 | for round_index in range(warmup_replays): |
| 284 | for candidate_index in range(len(names)): |
| 285 | name = names[(round_index + candidate_index) % len(names)] |
| 286 | device_barrier() |
| 287 | graphs[name][round_index % len(graphs[name])].replay() |
| 288 | torch.accelerator.synchronize() |
| 289 | |
| 290 | timings = {name: [] for name in names} |
| 291 | start = torch.cuda.Event(enable_timing=True) |
| 292 | end = torch.cuda.Event(enable_timing=True) |
| 293 | for sample_index in range(samples): |
| 294 | for candidate_index in range(len(names)): |
| 295 | name = names[(sample_index + candidate_index) % len(names)] |
| 296 | if flush is not None: |
| 297 | flush.zero_() |
| 298 | device_barrier() |
| 299 | start.record() |
| 300 | graphs[name][sample_index % len(graphs[name])].replay() |
| 301 | end.record() |
| 302 | end.synchronize() |
| 303 | |
| 304 | elapsed = torch.tensor( |
| 305 | start.elapsed_time(end) * 1000, |
| 306 | dtype=torch.float64, |
| 307 | device=torch.accelerator.current_device_index(), |
| 308 | ) |
| 309 | dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group) |
| 310 | timings[name].append(elapsed.item()) |
| 311 | |
| 312 | return {name: statistics.median(values) for name, values in timings.items()} |
| 313 | |
| 314 | |
| 315 | def benchmark_size( |
| 316 | args: argparse.Namespace, |
| 317 | comm: CustomAllreduce, |
| 318 | message_bytes: int, |
| 319 | dtype: torch.dtype, |
| 320 | device: torch.device, |
| 321 | rank: int, |
| 322 | world_size: int, |
| 323 | device_group: dist.ProcessGroup, |
| 324 | cpu_group: dist.ProcessGroup, |
| 325 | device_barrier: Callable[[], None], |
| 326 | ) -> dict[str, object]: |
| 327 | element_size = dtype.itemsize |
| 328 | if message_bytes % (element_size * world_size) != 0: |
| 329 | raise ValueError("Message bytes must divide evenly into TP shards") |
| 330 | num_elements = message_bytes // element_size |
| 331 | local_elements = num_elements // world_size |
| 332 | if local_elements * element_size % 16 != 0: |
| 333 | raise ValueError("Each output shard must be a multiple of 16 bytes") |
| 334 | |
| 335 | inputs = [] |
| 336 | for workspace_index in range(args.num_workspaces): |
| 337 | generator = torch.Generator(device=device) |
| 338 | generator.manual_seed(1000 + rank * 101 + workspace_index) |
| 339 | input_tensor = torch.empty(num_elements, dtype=dtype, device=device) |
| 340 | input_tensor.uniform_(-1.0, 1.0, generator=generator) |
| 341 | inputs.append(input_tensor) |
| 342 | |
| 343 | candidates = build_candidates( |
| 344 | comm, |
| 345 | inputs, |
| 346 | local_elements, |
| 347 | args.blocks, |
| 348 | device_group, |
| 349 | ) |
| 350 | torch.accelerator.synchronize() |
| 351 | dist.barrier(group=cpu_group) |
| 352 | check_correctness( |
| 353 | candidates, |
| 354 | inputs, |
| 355 | rank, |
| 356 | world_size, |
| 357 | cpu_group, |
| 358 | device_group, |
| 359 | ) |
| 360 | graphs, graph_keepalive = capture_candidates(candidates, cpu_group) |
| 361 | check_correctness( |
| 362 | candidates, |
| 363 | inputs, |
| 364 | rank, |
| 365 | world_size, |
| 366 | cpu_group, |
| 367 | device_group, |
| 368 | graphs, |
| 369 | ) |
| 370 | |
| 371 | flush = None |
| 372 | if args.cold_l2: |
| 373 | l2_bytes = torch.cuda.get_device_properties(device).L2_cache_size |
| 374 | flush = torch.empty(2 * l2_bytes, dtype=torch.uint8, device=device) |
| 375 | times = benchmark_graphs( |
| 376 | graphs, |
| 377 | args.warmup_replays, |
| 378 | args.samples, |
| 379 | device_group, |
| 380 | device_barrier, |
| 381 | flush, |
| 382 | ) |
| 383 | graph_keepalive.clear() |
| 384 | |
| 385 | bus_bytes = message_bytes * (world_size - 1) / world_size |
| 386 | fastest = min(times, key=times.__getitem__) |
| 387 | return { |
| 388 | "message_bytes": message_bytes, |
| 389 | "shard_bytes": message_bytes // world_size, |
| 390 | "latency_us": times, |
| 391 | "logical_bus_gbps": { |
| 392 | name: bus_bytes / (latency_us * 1e3) for name, latency_us in times.items() |
| 393 | }, |
| 394 | "fastest": fastest, |
| 395 | } |
| 396 | |
| 397 | |
| 398 | def main() -> None: |
| 399 | args = parse_args() |
| 400 | if not args.message_bytes or min(args.message_bytes) <= 0: |
| 401 | raise ValueError("--message-bytes values must be positive") |
| 402 | if not args.blocks or min(args.blocks) <= 0 or max(args.blocks) > 8: |
| 403 | raise ValueError("--blocks values must be in [1, 8]") |
| 404 | if min(args.num_workspaces, args.samples) <= 0: |
| 405 | raise ValueError("Workspace and sample counts must be positive") |
| 406 | |
| 407 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 408 | torch.accelerator.set_device_index(local_rank) |
| 409 | dist.init_process_group("nccl") |
| 410 | device_group = dist.group.WORLD |
| 411 | cpu_group = dist.new_group(backend="gloo") |
| 412 | rank = dist.get_rank(device_group) |
| 413 | world_size = dist.get_world_size(device_group) |
| 414 | device = torch.device("cuda", local_rank) |
| 415 | capability = torch.cuda.get_device_capability(device) |
| 416 | if world_size not in (2, 4, 8) or capability not in ((10, 0), (10, 3)): |
| 417 | raise RuntimeError( |
| 418 | "This benchmark targets TP2/TP4/TP8 SM100/SM103, " |
| 419 | f"got TP{world_size} SM{capability}" |
| 420 | ) |
| 421 | |
| 422 | max_message_bytes = max(args.message_bytes) |
| 423 | comm = CustomAllreduce( |
| 424 | group=cpu_group, |
| 425 | device=device, |
| 426 | max_mnnvl_reduce_scatter_size=max_message_bytes, |
| 427 | max_mnnvl_multimem_reduce_scatter_size=max_message_bytes, |
| 428 | ) |
| 429 | if comm.disabled: |
| 430 | raise RuntimeError("The production custom all-reduce backend is unavailable") |
| 431 | comm._init_mnnvl_multimem_reduce_scatter_buffer() |
| 432 | if not comm.mnnvl_multimem_rs_multicast_ptr: |
| 433 | raise RuntimeError("The production MNNVL multimem backend is unavailable") |
| 434 | |
| 435 | dtype = _DTYPES[args.dtype] |
| 436 | warmup = torch.zeros(1, dtype=torch.float32, device=device) |
| 437 | dist.all_reduce(warmup, group=device_group) |
| 438 | sync = torch.zeros_like(warmup) |
| 439 | |
| 440 | def device_barrier() -> None: |
| 441 | dist.all_reduce(sync, group=device_group) |
| 442 | |
| 443 | results = [ |
| 444 | benchmark_size( |
| 445 | args, |
| 446 | comm, |
| 447 | message_bytes, |
| 448 | dtype, |
| 449 | device, |
| 450 | rank, |
| 451 | world_size, |
| 452 | device_group, |
| 453 | cpu_group, |
| 454 | device_barrier, |
| 455 | ) |
| 456 | for message_bytes in args.message_bytes |
| 457 | ] |
| 458 | if rank == 0: |
| 459 | metadata = { |
| 460 | "world_size": world_size, |
| 461 | "gpu": torch.cuda.get_device_name(device), |
| 462 | "sm": f"{capability[0]}{capability[1]}", |
| 463 | "num_sms": torch.cuda.get_device_properties(device).multi_processor_count, |
| 464 | "dtype": args.dtype, |
| 465 | "torch": torch.__version__, |
| 466 | "cuda": torch.version.cuda, |
| 467 | "blocks": args.blocks, |
| 468 | "num_workspaces": args.num_workspaces, |
| 469 | "cold_l2": args.cold_l2, |
| 470 | } |
| 471 | print(json.dumps({"metadata": metadata, "results": results}, indent=2)) |
| 472 | |
| 473 | comm.close() |
| 474 | dist.barrier(group=cpu_group) |
| 475 | dist.destroy_process_group(cpu_group) |
| 476 | dist.destroy_process_group() |
| 477 | |
| 478 | |
| 479 | if __name__ == "__main__": |
| 480 | main() |