| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.utils.benchmark as benchmark |
| 6 | from benchmark_shapes import WEIGHT_SHAPES_MOE |
| 7 | |
| 8 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 9 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 10 | from vllm import _custom_ops as ops |
| 11 | from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config |
| 12 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 13 | maybe_make_prepare_finalize, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config |
| 16 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8 |
| 17 | from vllm.model_executor.layers.fused_moe.fused_moe import ( |
| 18 | fused_experts, |
| 19 | fused_topk, |
| 20 | ) |
| 21 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 22 | from vllm.v1.worker.workspace import init_workspace_manager |
| 23 | |
| 24 | DEFAULT_MODELS = [ |
| 25 | "mistralai/Mixtral-8x7B-Instruct-v0.1", |
| 26 | "deepseek-ai/DeepSeek-V2-Lite", |
| 27 | "ibm-granite/granite-3.0-1b-a400m", |
| 28 | "ibm-granite/granite-3.0-3b-a800m", |
| 29 | ] |
| 30 | DEFAULT_BATCH_SIZES = [1, 4, 8, 16, 32, 64, 128, 256, 512] |
| 31 | DEFAULT_TP_SIZES = [1] |
| 32 | |
| 33 | PER_ACT_TOKEN_OPTS = [False] |
| 34 | PER_OUT_CH_OPTS = [False] |
| 35 | |
| 36 | |
| 37 | def to_fp8(tensor: torch.Tensor): |
| 38 | finfo = torch.finfo(torch.float8_e4m3fn) |
| 39 | return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to( |
| 40 | dtype=torch.float8_e4m3fn |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | def bench_run( |
| 45 | results: list[benchmark.Measurement], |
| 46 | model: str, |
| 47 | num_experts: int, |
| 48 | topk: int, |
| 49 | per_act_token: bool, |
| 50 | per_out_ch: bool, |
| 51 | mkn: tuple[int, int, int], |
| 52 | ): |
| 53 | init_workspace_manager(torch.accelerator.current_device_index()) |
| 54 | label = "Quant Matmul" |
| 55 | |
| 56 | sub_label = ( |
| 57 | "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, MKN=({})".format( |
| 58 | model, num_experts, topk, per_act_token, per_out_ch, mkn |
| 59 | ) |
| 60 | ) |
| 61 | |
| 62 | print(f"Testing: {sub_label}") |
| 63 | |
| 64 | (m, k, n) = mkn |
| 65 | |
| 66 | dtype = torch.half |
| 67 | |
| 68 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 69 | w1 = torch.randn((num_experts, 2 * n, k), device="cuda", dtype=dtype) / 10 |
| 70 | w2 = torch.randn((num_experts, k, n), device="cuda", dtype=dtype) / 10 |
| 71 | |
| 72 | _, a_scale = ops.scaled_fp8_quant(a) |
| 73 | |
| 74 | w1_q = torch.empty( |
| 75 | (num_experts, 2 * n, k), device="cuda", dtype=torch.float8_e4m3fn |
| 76 | ) |
| 77 | w2_q = torch.empty((num_experts, k, n), device="cuda", dtype=torch.float8_e4m3fn) |
| 78 | w1_scale = torch.empty((num_experts, 1, 1), device="cuda", dtype=torch.float32) |
| 79 | w2_scale = torch.empty((num_experts, 1, 1), device="cuda", dtype=torch.float32) |
| 80 | |
| 81 | for expert in range(num_experts): |
| 82 | w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(w1[expert]) |
| 83 | w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(w2[expert]) |
| 84 | |
| 85 | score = torch.randn((m, num_experts), device="cuda", dtype=dtype) |
| 86 | |
| 87 | topk_weights, topk_ids, token_expert_indices = fused_topk( |
| 88 | a, score, topk, renormalize=False |
| 89 | ) |
| 90 | |
| 91 | def run_triton_moe( |
| 92 | a: torch.Tensor, |
| 93 | w1: torch.Tensor, |
| 94 | w2: torch.Tensor, |
| 95 | topk_weights: torch.Tensor, |
| 96 | topk_ids: torch.Tensor, |
| 97 | w1_scale: torch.Tensor, |
| 98 | w2_scale: torch.Tensor, |
| 99 | a_scale: torch.Tensor, |
| 100 | num_repeats: int, |
| 101 | ): |
| 102 | quant_config = fp8_w8a8_moe_quant_config( |
| 103 | w1_scale=w1_scale, |
| 104 | w2_scale=w2_scale, |
| 105 | a1_scale=a_scale, |
| 106 | ) |
| 107 | for _ in range(num_repeats): |
| 108 | fused_experts( |
| 109 | a, |
| 110 | w1, |
| 111 | w2, |
| 112 | topk_weights, |
| 113 | topk_ids, |
| 114 | quant_config=quant_config, |
| 115 | ) |
| 116 | |
| 117 | def run_cutlass_moe( |
| 118 | a: torch.Tensor, |
| 119 | a_scale: torch.Tensor, |
| 120 | w1: torch.Tensor, |
| 121 | w2: torch.Tensor, |
| 122 | w1_scale: torch.Tensor, |
| 123 | w2_scale: torch.Tensor, |
| 124 | topk_weights: torch.Tensor, |
| 125 | topk_ids: torch.Tensor, |
| 126 | per_act_token: bool, |
| 127 | num_repeats: int, |
| 128 | ): |
| 129 | quant_config = fp8_w8a8_moe_quant_config( |
| 130 | w1_scale=w1_scale, |
| 131 | w2_scale=w2_scale, |
| 132 | per_act_token_quant=per_act_token, |
| 133 | ) |
| 134 | moe_config = make_dummy_moe_config( |
| 135 | num_experts=w2.shape[0], |
| 136 | hidden_dim=w2.shape[1], |
| 137 | intermediate_size_per_partition=w2.shape[2], |
| 138 | in_dtype=a.dtype, |
| 139 | ) |
| 140 | |
| 141 | fn = mk.FusedMoEKernel( |
| 142 | maybe_make_prepare_finalize( |
| 143 | moe=moe_config, |
| 144 | quant_config=quant_config, |
| 145 | allow_new_interface=True, |
| 146 | use_monolithic=False, |
| 147 | ), |
| 148 | CutlassExpertsFp8( |
| 149 | moe_config=moe_config, |
| 150 | quant_config=quant_config, |
| 151 | ), |
| 152 | ) |
| 153 | |
| 154 | for _ in range(num_repeats): |
| 155 | fn(a, w1, w2, topk_weights, topk_ids) |
| 156 | |
| 157 | def run_cutlass_from_graph( |
| 158 | a: torch.Tensor, |
| 159 | a_scale: torch.Tensor, |
| 160 | w1: torch.Tensor, |
| 161 | w2: torch.Tensor, |
| 162 | w1_scale: torch.Tensor, |
| 163 | w2_scale: torch.Tensor, |
| 164 | topk_weights: torch.Tensor, |
| 165 | topk_ids: torch.Tensor, |
| 166 | ): |
| 167 | quant_config = fp8_w8a8_moe_quant_config( |
| 168 | w1_scale=w1_scale, |
| 169 | w2_scale=w2_scale, |
| 170 | per_act_token_quant=per_act_token, |
| 171 | ) |
| 172 | moe_config = make_dummy_moe_config( |
| 173 | num_experts=w2.shape[0], |
| 174 | hidden_dim=w2.shape[1], |
| 175 | intermediate_size_per_partition=w2.shape[2], |
| 176 | in_dtype=a.dtype, |
| 177 | ) |
| 178 | |
| 179 | fn = mk.FusedMoEKernel( |
| 180 | maybe_make_prepare_finalize( |
| 181 | moe=moe_config, |
| 182 | quant_config=quant_config, |
| 183 | allow_new_interface=True, |
| 184 | use_monolithic=False, |
| 185 | ), |
| 186 | CutlassExpertsFp8( |
| 187 | moe_config=moe_config, |
| 188 | quant_config=quant_config, |
| 189 | ), |
| 190 | ) |
| 191 | |
| 192 | with set_current_vllm_config( |
| 193 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 194 | ): |
| 195 | return fn(a, w1, w2, topk_weights, topk_ids) |
| 196 | |
| 197 | def run_triton_from_graph( |
| 198 | a: torch.Tensor, |
| 199 | w1: torch.Tensor, |
| 200 | w2: torch.Tensor, |
| 201 | topk_weights: torch.Tensor, |
| 202 | topk_ids: torch.Tensor, |
| 203 | w1_scale: torch.Tensor, |
| 204 | w2_scale: torch.Tensor, |
| 205 | a_scale: torch.Tensor, |
| 206 | ): |
| 207 | quant_config = fp8_w8a8_moe_quant_config( |
| 208 | w1_scale=w1_scale, |
| 209 | w2_scale=w2_scale, |
| 210 | a1_scale=a_scale, |
| 211 | ) |
| 212 | with set_current_vllm_config( |
| 213 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 214 | ): |
| 215 | return fused_experts( |
| 216 | a, |
| 217 | w1, |
| 218 | w2, |
| 219 | topk_weights, |
| 220 | topk_ids, |
| 221 | quant_config=quant_config, |
| 222 | ) |
| 223 | |
| 224 | def replay_graph(graph, num_repeats): |
| 225 | for _ in range(num_repeats): |
| 226 | graph.replay() |
| 227 | torch.accelerator.synchronize() |
| 228 | |
| 229 | cutlass_stream = torch.cuda.Stream() |
| 230 | cutlass_graph = torch.cuda.CUDAGraph() |
| 231 | with torch.cuda.graph(cutlass_graph, stream=cutlass_stream): |
| 232 | run_cutlass_from_graph( |
| 233 | a, |
| 234 | a_scale, |
| 235 | w1_q, |
| 236 | w2_q, |
| 237 | w1_scale, |
| 238 | w2_scale, |
| 239 | topk_weights, |
| 240 | topk_ids, |
| 241 | ) |
| 242 | torch.accelerator.synchronize() |
| 243 | |
| 244 | triton_stream = torch.cuda.Stream() |
| 245 | triton_graph = torch.cuda.CUDAGraph() |
| 246 | with torch.cuda.graph(triton_graph, stream=triton_stream): |
| 247 | run_triton_from_graph( |
| 248 | a, |
| 249 | w1_q, |
| 250 | w2_q, |
| 251 | topk_weights, |
| 252 | topk_ids, |
| 253 | w1_scale, |
| 254 | w2_scale, |
| 255 | a_scale, |
| 256 | ) |
| 257 | torch.accelerator.synchronize() |
| 258 | |
| 259 | min_run_time = 5 |
| 260 | num_warmup = 5 |
| 261 | num_runs = 25 |
| 262 | |
| 263 | globals = { |
| 264 | # Baseline params |
| 265 | "w1": w1, |
| 266 | "w2": w2, |
| 267 | "score": score, |
| 268 | "topk": topk, |
| 269 | # Cutlass params |
| 270 | "a_scale": a_scale, |
| 271 | "w1_q": w1_q, |
| 272 | "w2_q": w2_q, |
| 273 | "w1_scale": w1_scale, |
| 274 | "w2_scale": w2_scale, |
| 275 | "per_act_token": per_act_token, |
| 276 | # cuda graph params |
| 277 | "cutlass_graph": cutlass_graph, |
| 278 | "triton_graph": triton_graph, |
| 279 | # Gen params |
| 280 | "a": a, |
| 281 | "topk_weights": topk_weights, |
| 282 | "topk_ids": topk_ids, |
| 283 | "num_runs": num_runs, |
| 284 | # Kernels |
| 285 | "run_triton_moe": run_triton_moe, |
| 286 | "run_cutlass_moe": run_cutlass_moe, |
| 287 | "replay_graph": replay_graph, |
| 288 | } |
| 289 | |
| 290 | # Warmup |
| 291 | run_triton_moe( |
| 292 | a, |
| 293 | w1_q, |
| 294 | w2_q, |
| 295 | topk_weights, |
| 296 | topk_ids, |
| 297 | w1_scale, |
| 298 | w2_scale, |
| 299 | a_scale, |
| 300 | num_warmup, |
| 301 | ) |
| 302 | |
| 303 | results.append( |
| 304 | benchmark.Timer( |
| 305 | stmt="run_triton_moe(a, w1_q, w2_q, topk_weights, topk_ids, w1_scale, w2_scale, a_scale, num_runs)", # noqa: E501 |
| 306 | globals=globals, |
| 307 | label=label, |
| 308 | sub_label=sub_label, |
| 309 | description="triton_moe", |
| 310 | ).blocked_autorange(min_run_time=min_run_time) |
| 311 | ) |
| 312 | |
| 313 | # Warmup |
| 314 | replay_graph(triton_graph, num_warmup) |
| 315 | |
| 316 | results.append( |
| 317 | benchmark.Timer( |
| 318 | stmt="replay_graph(triton_graph, num_runs)", |
| 319 | globals=globals, |
| 320 | label=label, |
| 321 | sub_label=sub_label, |
| 322 | description="triton_moe_cuda_graphs", |
| 323 | ).blocked_autorange(min_run_time=min_run_time) |
| 324 | ) |
| 325 | |
| 326 | # Warmup |
| 327 | run_cutlass_moe( |
| 328 | a, |
| 329 | a_scale, |
| 330 | w1_q, |
| 331 | w2_q, |
| 332 | w1_scale, |
| 333 | w2_scale, |
| 334 | topk_weights, |
| 335 | topk_ids, |
| 336 | per_act_token, |
| 337 | num_warmup, |
| 338 | ) |
| 339 | |
| 340 | results.append( |
| 341 | benchmark.Timer( |
| 342 | stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501 |
| 343 | globals=globals, |
| 344 | label=label, |
| 345 | sub_label=sub_label, |
| 346 | description="grouped_gemm_moe", |
| 347 | ).blocked_autorange(min_run_time=min_run_time) |
| 348 | ) |
| 349 | |
| 350 | # Warmup |
| 351 | replay_graph(cutlass_graph, num_warmup) |
| 352 | |
| 353 | results.append( |
| 354 | benchmark.Timer( |
| 355 | stmt="replay_graph(cutlass_graph, num_runs)", |
| 356 | globals=globals, |
| 357 | label=label, |
| 358 | sub_label=sub_label, |
| 359 | description="grouped_gemm_moe_cuda_graphs", |
| 360 | ).blocked_autorange(min_run_time=min_run_time) |
| 361 | ) |
| 362 | |
| 363 | |
| 364 | def main(args): |
| 365 | # Initialize workspace manager (required for CUTLASS MoE kernels) |
| 366 | device = torch.device("cuda:0") |
| 367 | init_workspace_manager(device) |
| 368 | |
| 369 | print("Benchmarking models:") |
| 370 | for i, model in enumerate(args.models): |
| 371 | print(f"[{i}] {model}") |
| 372 | |
| 373 | results: list[benchmark.Measurement] = [] |
| 374 | |
| 375 | for model in args.models: |
| 376 | for tp in args.tp_sizes: |
| 377 | for layer in WEIGHT_SHAPES_MOE[model]: |
| 378 | num_experts = layer[0] |
| 379 | topk = layer[1] |
| 380 | size_k = layer[2] |
| 381 | size_n = layer[3] // tp |
| 382 | |
| 383 | if len(args.limit_k) > 0 and size_k not in args.limit_k: |
| 384 | continue |
| 385 | |
| 386 | if len(args.limit_n) > 0 and size_n not in args.limit_n: |
| 387 | continue |
| 388 | |
| 389 | for per_act_token in PER_ACT_TOKEN_OPTS: |
| 390 | for per_out_ch in PER_OUT_CH_OPTS: |
| 391 | for size_m in DEFAULT_BATCH_SIZES: |
| 392 | mkn = (size_m, size_k, size_n) |
| 393 | bench_run( |
| 394 | results, |
| 395 | model, |
| 396 | num_experts, |
| 397 | topk, |
| 398 | per_act_token, |
| 399 | per_out_ch, |
| 400 | mkn, |
| 401 | ) |
| 402 | |
| 403 | compare = benchmark.Compare(results) |
| 404 | compare.print() |
| 405 | |
| 406 | |
| 407 | if __name__ == "__main__": |
| 408 | parser = FlexibleArgumentParser( |
| 409 | description="Benchmark Marlin across specified models/shapes/batches" |
| 410 | ) |
| 411 | parser.add_argument( |
| 412 | "--models", |
| 413 | nargs="+", |
| 414 | type=str, |
| 415 | default=DEFAULT_MODELS, |
| 416 | choices=WEIGHT_SHAPES_MOE.keys(), |
| 417 | ) |
| 418 | parser.add_argument("--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES) |
| 419 | parser.add_argument( |
| 420 | "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES |
| 421 | ) |
| 422 | parser.add_argument("--limit-k", nargs="+", type=int, default=[]) |
| 423 | parser.add_argument("--limit-n", nargs="+", type=int, default=[]) |
| 424 | parser.add_argument("--limit-num-groups", nargs="+", type=int, default=[]) |
| 425 | parser.add_argument("--limit-per-act-token", nargs="+", type=int, default=[]) |
| 426 | parser.add_argument("--limit-per-out-ch", nargs="+", type=int, default=[]) |
| 427 | |
| 428 | args = parser.parse_args() |
| 429 | main(args) |