| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Minimal multi-GPU GEMM + reduce-scatter microbenchmark. |
| 5 | |
| 6 | Run on one node with, for example: |
| 7 | |
| 8 | torchrun --standalone --nproc-per-node=8 \ |
| 9 | .agents/skills/kernel-microbenchmark/benchmarks/multi_gpu_gemm_rs.py |
| 10 | """ |
| 11 | |
| 12 | import argparse |
| 13 | import os |
| 14 | import statistics |
| 15 | from collections.abc import Callable |
| 16 | |
| 17 | import pandas as pd |
| 18 | import torch |
| 19 | import torch.distributed as dist |
| 20 | import torch.distributed._symmetric_memory as symm_mem |
| 21 | |
| 22 | from vllm.config import VllmConfig, set_current_vllm_config |
| 23 | from vllm.distributed import cleanup_dist_env_and_memory |
| 24 | from vllm.distributed.parallel_state import ( |
| 25 | get_tp_group, |
| 26 | init_distributed_environment, |
| 27 | initialize_model_parallel, |
| 28 | ) |
| 29 | |
| 30 | |
| 31 | def parse_args() -> argparse.Namespace: |
| 32 | parser = argparse.ArgumentParser() |
| 33 | parser.add_argument("--m", type=int, nargs="+", default=[128, 512, 2048]) |
| 34 | parser.add_argument("--n", type=int, default=4096) |
| 35 | parser.add_argument( |
| 36 | "--k", |
| 37 | type=int, |
| 38 | nargs="+", |
| 39 | default=[4096], |
| 40 | help="Per-rank K values", |
| 41 | ) |
| 42 | parser.add_argument("--num-workspaces", type=int, default=10) |
| 43 | parser.add_argument("--warmup-replays", type=int, default=5) |
| 44 | parser.add_argument("--samples", type=int, default=20) |
| 45 | return parser.parse_args() |
| 46 | |
| 47 | |
| 48 | def make_gemm_rs( |
| 49 | x: torch.Tensor, |
| 50 | weight: torch.Tensor, |
| 51 | partial: torch.Tensor, |
| 52 | output: torch.Tensor, |
| 53 | rows: int, |
| 54 | device_group: dist.ProcessGroup, |
| 55 | ) -> Callable[[], None]: |
| 56 | def run() -> None: |
| 57 | torch.mm(x, weight.T, out=partial[:rows]) |
| 58 | dist.reduce_scatter_single(output, partial, group=device_group) |
| 59 | |
| 60 | return run |
| 61 | |
| 62 | |
| 63 | def check_correctness( |
| 64 | runs: dict[str, Callable[[], None]], |
| 65 | outputs: dict[str, torch.Tensor], |
| 66 | x: torch.Tensor, |
| 67 | weight: torch.Tensor, |
| 68 | padded_rows: int, |
| 69 | rank: int, |
| 70 | device_group: dist.ProcessGroup, |
| 71 | ) -> None: |
| 72 | expected_full = torch.zeros( |
| 73 | (padded_rows, weight.shape[0]), |
| 74 | dtype=x.dtype, |
| 75 | device=x.device, |
| 76 | ) |
| 77 | torch.mm(x, weight.T, out=expected_full[: x.shape[0]]) |
| 78 | dist.all_reduce(expected_full, group=device_group) |
| 79 | expected = expected_full.chunk(dist.get_world_size(device_group))[rank] |
| 80 | for name, run in runs.items(): |
| 81 | run() |
| 82 | torch.accelerator.synchronize() |
| 83 | torch.testing.assert_close( |
| 84 | outputs[name], |
| 85 | expected, |
| 86 | rtol=5e-2, |
| 87 | atol=4.0, |
| 88 | ) |
| 89 | |
| 90 | |
| 91 | def capture_graph( |
| 92 | run: Callable[[], None], |
| 93 | cpu_group: dist.ProcessGroup, |
| 94 | ) -> tuple[torch.cuda.CUDAGraph, torch.cuda.Stream]: |
| 95 | stream = torch.cuda.Stream() |
| 96 | stream.wait_stream(torch.cuda.current_stream()) |
| 97 | dist.barrier(group=cpu_group) |
| 98 | with torch.cuda.stream(stream): |
| 99 | for _ in range(3): |
| 100 | run() |
| 101 | stream.synchronize() |
| 102 | dist.barrier(group=cpu_group) |
| 103 | |
| 104 | graph = torch.cuda.CUDAGraph() |
| 105 | with torch.cuda.graph(graph, stream=stream): |
| 106 | run() |
| 107 | torch.cuda.current_stream().wait_stream(stream) |
| 108 | dist.barrier(group=cpu_group) |
| 109 | return graph, stream |
| 110 | |
| 111 | |
| 112 | def benchmark_graphs( |
| 113 | candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]], |
| 114 | warmup_replays: int, |
| 115 | samples: int, |
| 116 | device_group: dist.ProcessGroup, |
| 117 | device_barrier: Callable[[], None], |
| 118 | ) -> dict[str, float]: |
| 119 | candidate_names = list(candidate_graphs) |
| 120 | for round_index in range(warmup_replays): |
| 121 | for candidate_index in range(len(candidate_names)): |
| 122 | candidate_id = (round_index + candidate_index) % len(candidate_names) |
| 123 | name = candidate_names[candidate_id] |
| 124 | graphs = candidate_graphs[name] |
| 125 | device_barrier() |
| 126 | graphs[round_index % len(graphs)].replay() |
| 127 | torch.accelerator.synchronize() |
| 128 | |
| 129 | timings: dict[str, list[float]] = {name: [] for name in candidate_names} |
| 130 | start = torch.cuda.Event(enable_timing=True) |
| 131 | end = torch.cuda.Event(enable_timing=True) |
| 132 | for sample_index in range(samples): |
| 133 | for candidate_index in range(len(candidate_names)): |
| 134 | candidate_id = (sample_index + candidate_index) % len(candidate_names) |
| 135 | name = candidate_names[candidate_id] |
| 136 | graphs = candidate_graphs[name] |
| 137 | device_barrier() |
| 138 | start.record() |
| 139 | graphs[sample_index % len(graphs)].replay() |
| 140 | end.record() |
| 141 | end.synchronize() |
| 142 | |
| 143 | elapsed_us = torch.tensor( |
| 144 | start.elapsed_time(end) * 1000, |
| 145 | dtype=torch.float64, |
| 146 | device=torch.accelerator.current_device_index(), |
| 147 | ) |
| 148 | dist.all_reduce(elapsed_us, op=dist.ReduceOp.MAX, group=device_group) |
| 149 | timings[name].append(elapsed_us.item()) |
| 150 | return {name: statistics.median(values) for name, values in timings.items()} |
| 151 | |
| 152 | |
| 153 | def benchmark_shape( |
| 154 | m: int, |
| 155 | n: int, |
| 156 | k: int, |
| 157 | num_workspaces: int, |
| 158 | warmup_replays: int, |
| 159 | samples: int, |
| 160 | device: torch.device, |
| 161 | rank: int, |
| 162 | world_size: int, |
| 163 | device_group: dist.ProcessGroup, |
| 164 | cpu_group: dist.ProcessGroup, |
| 165 | device_barrier: Callable[[], None], |
| 166 | ) -> dict[str, float | int]: |
| 167 | padded_m = (m + world_size - 1) // world_size * world_size |
| 168 | local_m = padded_m // world_size |
| 169 | |
| 170 | torch.manual_seed(1000 + rank * 10 + m + k) |
| 171 | inputs = [] |
| 172 | weights = [] |
| 173 | for _ in range(num_workspaces): |
| 174 | inputs.append(torch.randn(m, k, dtype=torch.bfloat16, device=device)) |
| 175 | weights.append(torch.randn(n, k, dtype=torch.bfloat16, device=device)) |
| 176 | |
| 177 | ring_partial = torch.empty(padded_m, n, dtype=torch.bfloat16, device=device) |
| 178 | ldmc_partial = symm_mem.empty( |
| 179 | (padded_m, n), |
| 180 | dtype=torch.bfloat16, |
| 181 | device=device, |
| 182 | ) |
| 183 | ldmc_handle = symm_mem.rendezvous(ldmc_partial, device_group) |
| 184 | ring_output = torch.empty(local_m, n, dtype=torch.bfloat16, device=device) |
| 185 | ldmc_output = torch.empty_like(ring_output) |
| 186 | if padded_m > m: |
| 187 | ring_partial[m:].zero_() |
| 188 | ldmc_partial[m:].zero_() |
| 189 | |
| 190 | candidate_runs = { |
| 191 | "ring_ll_us": [ |
| 192 | make_gemm_rs( |
| 193 | x, |
| 194 | weight, |
| 195 | ring_partial, |
| 196 | ring_output, |
| 197 | m, |
| 198 | device_group, |
| 199 | ) |
| 200 | for x, weight in zip(inputs, weights) |
| 201 | ], |
| 202 | "ldmc_us": [ |
| 203 | make_gemm_rs( |
| 204 | x, |
| 205 | weight, |
| 206 | ldmc_partial, |
| 207 | ldmc_output, |
| 208 | m, |
| 209 | device_group, |
| 210 | ) |
| 211 | for x, weight in zip(inputs, weights) |
| 212 | ], |
| 213 | } |
| 214 | x = inputs[0] |
| 215 | weight = weights[0] |
| 216 | check_correctness( |
| 217 | {name: runs[0] for name, runs in candidate_runs.items()}, |
| 218 | {"ring_ll_us": ring_output, "ldmc_us": ldmc_output}, |
| 219 | x, |
| 220 | weight, |
| 221 | padded_m, |
| 222 | rank, |
| 223 | device_group, |
| 224 | ) |
| 225 | |
| 226 | candidate_graphs = {} |
| 227 | graph_keepalive: list[object] = [ldmc_handle] |
| 228 | for name, runs in candidate_runs.items(): |
| 229 | bundles = [capture_graph(run, cpu_group) for run in runs] |
| 230 | candidate_graphs[name] = [graph for graph, _ in bundles] |
| 231 | graph_keepalive.extend(bundles) |
| 232 | |
| 233 | times = benchmark_graphs( |
| 234 | candidate_graphs, |
| 235 | warmup_replays, |
| 236 | samples, |
| 237 | device_group, |
| 238 | device_barrier, |
| 239 | ) |
| 240 | global_flops = 2 * m * n * k * world_size |
| 241 | return { |
| 242 | "M": m, |
| 243 | "N": n, |
| 244 | "K_per_rank": k, |
| 245 | "K_global": k * world_size, |
| 246 | **times, |
| 247 | "ring_ll_tflops": global_flops / (times["ring_ll_us"] * 1e6), |
| 248 | "ldmc_tflops": global_flops / (times["ldmc_us"] * 1e6), |
| 249 | "ldmc_speedup": times["ring_ll_us"] / times["ldmc_us"], |
| 250 | } |
| 251 | |
| 252 | |
| 253 | def main() -> None: |
| 254 | args = parse_args() |
| 255 | assert args.m and min(args.m) > 0 |
| 256 | assert args.k and min(args.k) > 0 |
| 257 | assert min(args.n, args.num_workspaces, args.samples) > 0 |
| 258 | assert args.warmup_replays >= 0 |
| 259 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 260 | local_world_size = int(os.environ["LOCAL_WORLD_SIZE"]) |
| 261 | torch.accelerator.set_device_index(local_rank) |
| 262 | init_distributed_environment() |
| 263 | world_size = dist.get_world_size() |
| 264 | os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" |
| 265 | symm_mem.set_backend("NCCL") |
| 266 | with set_current_vllm_config(VllmConfig()): |
| 267 | initialize_model_parallel(tensor_model_parallel_size=world_size) |
| 268 | |
| 269 | tp_group = get_tp_group() |
| 270 | device_group = tp_group.device_group |
| 271 | cpu_group = tp_group.cpu_group |
| 272 | rank = tp_group.rank_in_group |
| 273 | device = torch.device("cuda", local_rank) |
| 274 | group_warmup = torch.zeros(1, device=device) |
| 275 | dist.all_reduce(group_warmup, group=device_group) |
| 276 | pynccl_comm = tp_group.device_communicator.pynccl_comm |
| 277 | assert pynccl_comm is not None |
| 278 | sync_input = torch.zeros(1, device=device) |
| 279 | sync_output = torch.empty_like(sync_input) |
| 280 | |
| 281 | def device_barrier() -> None: |
| 282 | # Order the timed launch after a device-side rank rendezvous without |
| 283 | # including the rendezvous itself in the measured event interval. |
| 284 | pynccl_comm.all_reduce(sync_input, sync_output) |
| 285 | |
| 286 | results = [ |
| 287 | benchmark_shape( |
| 288 | m, |
| 289 | args.n, |
| 290 | k, |
| 291 | args.num_workspaces, |
| 292 | args.warmup_replays, |
| 293 | args.samples, |
| 294 | device, |
| 295 | rank, |
| 296 | world_size, |
| 297 | device_group, |
| 298 | cpu_group, |
| 299 | device_barrier, |
| 300 | ) |
| 301 | for k in args.k |
| 302 | for m in args.m |
| 303 | ] |
| 304 | |
| 305 | if rank == 0: |
| 306 | metadata = { |
| 307 | "world_size": world_size, |
| 308 | "local_world_size": local_world_size, |
| 309 | "num_nodes": world_size // local_world_size, |
| 310 | "backend": dist.get_backend(device_group), |
| 311 | "gpu": torch.cuda.get_device_name(local_rank), |
| 312 | "torch": torch.__version__, |
| 313 | "cuda": torch.version.cuda, |
| 314 | } |
| 315 | print(pd.Series(metadata, name="value").to_string()) |
| 316 | df = pd.DataFrame(results) |
| 317 | print(df.to_string(index=False, float_format=lambda x: f"{x:.3f}")) |
| 318 | |
| 319 | dist.barrier(group=cpu_group) |
| 320 | cleanup_dist_env_and_memory() |
| 321 | |
| 322 | |
| 323 | if __name__ == "__main__": |
| 324 | main() |