| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark the SM100 GEMM-RS/AR kernel against the model's unfused path. |
| 4 | |
| 5 | The baseline is what a model runs without the fusion: the projection's own |
| 6 | linear kernel followed by vLLM's collective (``sp_reduce_scatter`` for RS, |
| 7 | ``tensor_model_parallel_all_reduce`` for AR), with the collective backend |
| 8 | chosen by vLLM's dispatch exactly as in serving. The same two kernels are also |
| 9 | timed bare (the GEMM alone, the collective alone on a precomputed partial); |
| 10 | their max is the floor a perfectly overlapped fused kernel could reach. |
| 11 | |
| 12 | All ranks must belong to one NVLink domain. Run a TP8 sweep over the Kimi-K3 |
| 13 | projections with: |
| 14 | |
| 15 | torchrun --nproc-per-node=8 -- benchmarks/kernels/benchmark_gemm_rs_ar.py |
| 16 | |
| 17 | or the DeepSeek-V4.1 MXFP8 ``wo_b`` projection at TP4 with: |
| 18 | |
| 19 | torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \ |
| 20 | --model deepseek_v41 --backend flashinfer_cutedsl |
| 21 | """ |
| 22 | |
| 23 | import argparse |
| 24 | import json |
| 25 | import os |
| 26 | import statistics |
| 27 | from collections.abc import Callable |
| 28 | |
| 29 | import torch |
| 30 | import torch.distributed as dist |
| 31 | |
| 32 | from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config |
| 33 | from vllm.distributed import ( |
| 34 | cleanup_dist_env_and_memory, |
| 35 | tensor_model_parallel_all_reduce, |
| 36 | ) |
| 37 | from vllm.distributed.parallel_state import ( |
| 38 | get_tp_group, |
| 39 | init_distributed_environment, |
| 40 | initialize_model_parallel, |
| 41 | ) |
| 42 | from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr |
| 43 | from vllm.model_executor.layers.linear import RowParallelLinear |
| 44 | from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter |
| 45 | |
| 46 | # Per model: (hidden size N, global input widths K). Kimi-K3 lists the |
| 47 | # shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the |
| 48 | # ``wo_b`` output projection (o_groups * o_lora_rank). |
| 49 | _MODEL_PROJECTIONS = { |
| 50 | "kimi_k3": (7168, (6144, 12288)), |
| 51 | "deepseek_v41": (5120, (8192,)), |
| 52 | } |
| 53 | |
| 54 | |
| 55 | def parse_args() -> argparse.Namespace: |
| 56 | parser = argparse.ArgumentParser() |
| 57 | parser.add_argument( |
| 58 | "--mode", |
| 59 | choices=("rs", "ar"), |
| 60 | default="rs", |
| 61 | help="Collective mode to benchmark.", |
| 62 | ) |
| 63 | parser.add_argument( |
| 64 | "--m", |
| 65 | type=int, |
| 66 | nargs="+", |
| 67 | default=[128, 512, 2048, 8192, 32768], |
| 68 | help="Global token counts to benchmark.", |
| 69 | ) |
| 70 | parser.add_argument( |
| 71 | "--model", |
| 72 | choices=tuple(_MODEL_PROJECTIONS), |
| 73 | default="kimi_k3", |
| 74 | help="Model whose projection shapes set the default --n and --k.", |
| 75 | ) |
| 76 | parser.add_argument( |
| 77 | "--backend", |
| 78 | choices=("bf16", "flashinfer_cutlass", "flashinfer_cutedsl"), |
| 79 | default="bf16", |
| 80 | help="Weight dtype/layout: BF16 or an online-MXFP8 FlashInfer kernel.", |
| 81 | ) |
| 82 | parser.add_argument( |
| 83 | "--k", |
| 84 | type=int, |
| 85 | nargs="+", |
| 86 | help=( |
| 87 | "Per-rank input dimensions. By default, derive the selected " |
| 88 | "model's projection dimensions from the TP world size." |
| 89 | ), |
| 90 | ) |
| 91 | parser.add_argument( |
| 92 | "--n", |
| 93 | type=int, |
| 94 | help="Output width. Defaults to the selected model's hidden size.", |
| 95 | ) |
| 96 | parser.add_argument( |
| 97 | "--num-workspaces", |
| 98 | type=int, |
| 99 | default=10, |
| 100 | help="Pointer-distinct inputs and CUDA graphs to rotate.", |
| 101 | ) |
| 102 | parser.add_argument("--warmup-replays", type=int, default=5) |
| 103 | parser.add_argument("--samples", type=int, default=20) |
| 104 | parser.add_argument("--json", type=str, help="Also write the results here.") |
| 105 | return parser.parse_args() |
| 106 | |
| 107 | |
| 108 | def make_projection( |
| 109 | gemm_rs_ar: GemmRsAr, |
| 110 | weight: torch.Tensor, |
| 111 | world_size: int, |
| 112 | backend: str, |
| 113 | ) -> RowParallelLinear: |
| 114 | """Wrap a BF16 [N, K] weight the way the model would hand it to GEMM-RS. |
| 115 | |
| 116 | For the MXFP8 backends this quantizes online through the FlashInfer |
| 117 | kernel's post-load step, so the fused kernel and the unfused baseline see |
| 118 | the production weight and scale layouts. |
| 119 | """ |
| 120 | N, K = weight.shape |
| 121 | quant_config = None |
| 122 | if backend != "bf16": |
| 123 | from vllm.config.quantization import QuantizationConfigArgs |
| 124 | from vllm.model_executor.layers.quantization.online.base import ( |
| 125 | OnlineQuantizationConfig, |
| 126 | ) |
| 127 | |
| 128 | quant_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) |
| 129 | with torch.device(weight.device): |
| 130 | linear = RowParallelLinear( |
| 131 | K * world_size, |
| 132 | N, |
| 133 | bias=False, |
| 134 | params_dtype=weight.dtype, |
| 135 | quant_config=quant_config, |
| 136 | # The baseline reduces the local partial itself, like the models. |
| 137 | reduce_results=False, |
| 138 | return_bias=False, |
| 139 | ) |
| 140 | # Select the fused path before online quantization replaces the weights. |
| 141 | assert gemm_rs_ar.can_run(linear) |
| 142 | linear.weight = torch.nn.Parameter(weight, requires_grad=False) |
| 143 | linear.quant_method.process_weights_after_loading(linear) |
| 144 | return linear |
| 145 | |
| 146 | |
| 147 | def capture_graph( |
| 148 | op: Callable[[], torch.Tensor], |
| 149 | stream: torch.cuda.Stream, |
| 150 | cpu_group: dist.ProcessGroup, |
| 151 | ) -> torch.cuda.CUDAGraph: |
| 152 | stream.wait_stream(torch.cuda.current_stream()) |
| 153 | dist.barrier(group=cpu_group) |
| 154 | with torch.cuda.stream(stream): |
| 155 | for _ in range(3): |
| 156 | op() |
| 157 | stream.synchronize() |
| 158 | dist.barrier(group=cpu_group) |
| 159 | |
| 160 | graph = torch.cuda.CUDAGraph() |
| 161 | with torch.cuda.graph(graph, stream=stream): |
| 162 | op() |
| 163 | torch.cuda.current_stream().wait_stream(stream) |
| 164 | dist.barrier(group=cpu_group) |
| 165 | return graph |
| 166 | |
| 167 | |
| 168 | def benchmark_graphs( |
| 169 | candidate_graphs: dict[str, list[torch.cuda.CUDAGraph]], |
| 170 | warmup_replays: int, |
| 171 | samples: int, |
| 172 | device_group: dist.ProcessGroup, |
| 173 | device_barrier: Callable[[], None], |
| 174 | ) -> dict[str, float]: |
| 175 | """Median over samples of the slowest rank's replay time, in us.""" |
| 176 | names = list(candidate_graphs) |
| 177 | for round_index in range(warmup_replays): |
| 178 | for offset in range(len(names)): |
| 179 | name = names[(round_index + offset) % len(names)] |
| 180 | graphs = candidate_graphs[name] |
| 181 | device_barrier() |
| 182 | graphs[round_index % len(graphs)].replay() |
| 183 | torch.accelerator.synchronize() |
| 184 | |
| 185 | timings: dict[str, list[float]] = {name: [] for name in names} |
| 186 | start = torch.cuda.Event(enable_timing=True) |
| 187 | end = torch.cuda.Event(enable_timing=True) |
| 188 | for sample_index in range(samples): |
| 189 | for offset in range(len(names)): |
| 190 | name = names[(sample_index + offset) % len(names)] |
| 191 | graphs = candidate_graphs[name] |
| 192 | device_barrier() |
| 193 | start.record() |
| 194 | graphs[sample_index % len(graphs)].replay() |
| 195 | end.record() |
| 196 | end.synchronize() |
| 197 | |
| 198 | elapsed = torch.tensor( |
| 199 | start.elapsed_time(end) * 1000, |
| 200 | dtype=torch.float64, |
| 201 | device=torch.accelerator.current_device_index(), |
| 202 | ) |
| 203 | dist.all_reduce(elapsed, op=dist.ReduceOp.MAX, group=device_group) |
| 204 | timings[name].append(elapsed.item()) |
| 205 | return {name: statistics.median(values) for name, values in timings.items()} |
| 206 | |
| 207 | |
| 208 | def benchmark_shape( |
| 209 | gemm_rs_ar: GemmRsAr, |
| 210 | mode: str, |
| 211 | backend: str, |
| 212 | M: int, |
| 213 | N: int, |
| 214 | K: int, |
| 215 | num_workspaces: int, |
| 216 | warmup_replays: int, |
| 217 | samples: int, |
| 218 | device_group: dist.ProcessGroup, |
| 219 | cpu_group: dist.ProcessGroup, |
| 220 | device_barrier: Callable[[], None], |
| 221 | ) -> dict[str, float | int | str]: |
| 222 | all_reduce = mode == "ar" |
| 223 | world_size = dist.get_world_size(device_group) |
| 224 | rank = dist.get_rank(device_group) |
| 225 | device = torch.device("cuda", torch.accelerator.current_device_index()) |
| 226 | local_M = (M + world_size - 1) // world_size |
| 227 | |
| 228 | rng = torch.Generator(device=device) |
| 229 | rng.manual_seed(1000 + rank * 10 + M + K) |
| 230 | inputs = [ |
| 231 | torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng) |
| 232 | for _ in range(num_workspaces) |
| 233 | ] |
| 234 | projections = [ |
| 235 | make_projection( |
| 236 | gemm_rs_ar, |
| 237 | torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng), |
| 238 | world_size, |
| 239 | backend, |
| 240 | ) |
| 241 | for _ in range(num_workspaces) |
| 242 | ] |
| 243 | |
| 244 | def collective(partial: torch.Tensor) -> torch.Tensor: |
| 245 | if all_reduce: |
| 246 | return tensor_model_parallel_all_reduce(partial) |
| 247 | return sp_reduce_scatter(partial) |
| 248 | |
| 249 | def gemm(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: |
| 250 | return linear(x) |
| 251 | |
| 252 | def unfused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: |
| 253 | return collective(linear(x)) |
| 254 | |
| 255 | def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: |
| 256 | return gemm_rs_ar.apply(x, linear) |
| 257 | |
| 258 | # The bare collective runs on the GEMM's own output so it sees the same |
| 259 | # message size and layout the unfused path hands to the collective. |
| 260 | partials = [linear(x) for x, linear in zip(inputs, projections)] |
| 261 | |
| 262 | expected = unfused(inputs[0], projections[0]) |
| 263 | actual = fused(inputs[0], projections[0]) |
| 264 | torch.accelerator.synchronize(device) |
| 265 | rows = M if all_reduce else min(max(M - rank * local_M, 0), local_M) |
| 266 | torch.testing.assert_close( |
| 267 | actual[:rows], expected[:rows], rtol=5e-2, atol=4.0, msg=f"M={M}, K={K}" |
| 268 | ) |
| 269 | |
| 270 | candidates: dict[str, list[Callable[[], torch.Tensor]]] = { |
| 271 | "gemm": [lambda x=x, w=w: gemm(x, w) for x, w in zip(inputs, projections)], |
| 272 | "collective": [lambda p=p: collective(p) for p in partials], |
| 273 | "vllm": [lambda x=x, w=w: unfused(x, w) for x, w in zip(inputs, projections)], |
| 274 | "gemm_rs_ar": [ |
| 275 | lambda x=x, w=w: fused(x, w) for x, w in zip(inputs, projections) |
| 276 | ], |
| 277 | } |
| 278 | |
| 279 | candidate_graphs = {} |
| 280 | keepalive: list[object] = [] |
| 281 | for name, ops in candidates.items(): |
| 282 | stream = torch.cuda.Stream() |
| 283 | candidate_graphs[name] = [capture_graph(op, stream, cpu_group) for op in ops] |
| 284 | keepalive.append(stream) |
| 285 | |
| 286 | times = benchmark_graphs( |
| 287 | candidate_graphs, warmup_replays, samples, device_group, device_barrier |
| 288 | ) |
| 289 | floor = max(times["gemm"], times["collective"]) |
| 290 | return { |
| 291 | "mode": mode.upper(), |
| 292 | "backend": backend, |
| 293 | "M": M, |
| 294 | "N": N, |
| 295 | "K": K, |
| 296 | "gemm_us": times["gemm"], |
| 297 | "collective_us": times["collective"], |
| 298 | "floor_us": floor, |
| 299 | "vllm_us": times["vllm"], |
| 300 | "gemm_rs_ar_us": times["gemm_rs_ar"], |
| 301 | "latency_change": (times["gemm_rs_ar"] - times["vllm"]) / times["vllm"], |
| 302 | "floor_gap": (times["gemm_rs_ar"] - floor) / floor, |
| 303 | } |
| 304 | |
| 305 | |
| 306 | def print_results(results: list[dict[str, float | int | str]]) -> None: |
| 307 | collective = results[0]["mode"] |
| 308 | backend = results[0]["backend"] |
| 309 | print(f"### GEMM-{collective} vs vLLM unfused path ({backend} weights)") |
| 310 | print( |
| 311 | f"| M | N | K | GEMM (us) | {collective} (us) | floor (us) | vLLM (us) | " |
| 312 | f"GEMM-{collective} (us) | vs vLLM | vs floor |" |
| 313 | ) |
| 314 | print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|") |
| 315 | for r in results: |
| 316 | print( |
| 317 | f"| {r['M']} | {r['N']} | {r['K']} | {r['gemm_us']:.1f} | " |
| 318 | f"{r['collective_us']:.1f} | {r['floor_us']:.1f} | {r['vllm_us']:.1f} | " |
| 319 | f"{r['gemm_rs_ar_us']:.1f} | {r['latency_change']:+.1%} | " |
| 320 | f"{r['floor_gap']:+.1%} |" |
| 321 | ) |
| 322 | |
| 323 | |
| 324 | def main() -> None: |
| 325 | args = parse_args() |
| 326 | default_n, default_k = _MODEL_PROJECTIONS[args.model] |
| 327 | if args.n is None: |
| 328 | args.n = default_n |
| 329 | assert args.m and min(args.m) > 0 |
| 330 | assert args.n % 128 == 0 |
| 331 | assert args.num_workspaces > 0 |
| 332 | assert args.warmup_replays >= 0 |
| 333 | assert args.samples > 0 |
| 334 | |
| 335 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 336 | torch.accelerator.set_device_index(local_rank) |
| 337 | init_distributed_environment() |
| 338 | world_size = dist.get_world_size() |
| 339 | if args.k is None: |
| 340 | assert all(K % world_size == 0 for K in default_k) |
| 341 | K_values = [K // world_size for K in default_k] |
| 342 | else: |
| 343 | K_values = args.k |
| 344 | k_alignment = 64 if args.backend == "bf16" else 128 |
| 345 | assert all(K % k_alignment == 0 for K in K_values) |
| 346 | |
| 347 | config = VllmConfig() |
| 348 | config.model_config = ModelConfig(dtype="bfloat16") |
| 349 | # Pin the online-MXFP8 kernel the projections quantize through. |
| 350 | config.kernel_config.linear_backend = ( |
| 351 | "auto" if args.backend == "bf16" else args.backend |
| 352 | ) |
| 353 | with set_current_vllm_config(config): |
| 354 | # Builds the TP communicator as in serving, so the baseline's |
| 355 | # collective is whatever vLLM's dispatch selects (FlashInfer all-reduce |
| 356 | # first, then NCCL symmetric memory, custom all-reduce, pynccl). |
| 357 | initialize_model_parallel(tensor_model_parallel_size=world_size) |
| 358 | |
| 359 | tp_group = get_tp_group() |
| 360 | pynccl_comm = tp_group.device_communicator.pynccl_comm |
| 361 | assert pynccl_comm is not None |
| 362 | sync_input = torch.zeros(1, device=torch.accelerator.current_device_index()) |
| 363 | sync_output = torch.empty_like(sync_input) |
| 364 | |
| 365 | def device_barrier() -> None: |
| 366 | # Order the timed launch after a device-side rank rendezvous without |
| 367 | # including the rendezvous itself in the measured event interval. |
| 368 | pynccl_comm.all_reduce(sync_input, sync_output) |
| 369 | |
| 370 | gemm_rs_ar = GemmRsAr( |
| 371 | max_M=max(args.m), |
| 372 | N=args.n, |
| 373 | all_reduce=args.mode == "ar", |
| 374 | ) |
| 375 | with set_current_vllm_config(config): |
| 376 | results = [ |
| 377 | benchmark_shape( |
| 378 | gemm_rs_ar, |
| 379 | args.mode, |
| 380 | args.backend, |
| 381 | M, |
| 382 | args.n, |
| 383 | K, |
| 384 | args.num_workspaces, |
| 385 | args.warmup_replays, |
| 386 | args.samples, |
| 387 | tp_group.device_group, |
| 388 | tp_group.cpu_group, |
| 389 | device_barrier, |
| 390 | ) |
| 391 | for K in K_values |
| 392 | for M in args.m |
| 393 | ] |
| 394 | del gemm_rs_ar |
| 395 | |
| 396 | if tp_group.rank_in_group == 0: |
| 397 | print_results(results) |
| 398 | if args.json: |
| 399 | with open(args.json, "w") as f: |
| 400 | json.dump(results, f, indent=2) |
| 401 | |
| 402 | dist.barrier(group=tp_group.cpu_group) |
| 403 | cleanup_dist_env_and_memory() |
| 404 | |
| 405 | |
| 406 | if __name__ == "__main__": |
| 407 | main() |