| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe |
| 4 | kernel. The cutlass_moe_fp4 kernel takes in fp4 quantized weights and 16-bit |
| 5 | activations. The triton_moe kernel takes in fp8 weights(tensor scaled to fp8) |
| 6 | and 16-bit activations. |
| 7 | """ |
| 8 | |
| 9 | import nvtx |
| 10 | import torch |
| 11 | import torch.utils.benchmark as benchmark |
| 12 | |
| 13 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 14 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 15 | from vllm import _custom_ops as ops |
| 16 | from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config |
| 17 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 18 | maybe_make_prepare_finalize, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.fused_moe.config import ( |
| 21 | fp8_w8a8_moe_quant_config, |
| 22 | nvfp4_moe_quant_config, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( |
| 25 | CutlassExpertsFp4, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk |
| 28 | from vllm.scalar_type import scalar_types |
| 29 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 30 | from vllm.v1.worker.workspace import init_workspace_manager |
| 31 | |
| 32 | WEIGHT_SHAPES_MOE = { |
| 33 | "nvidia/DeepSeek-R1-FP4": [ |
| 34 | [256, 8, 2048, 7168], |
| 35 | ], |
| 36 | } |
| 37 | |
| 38 | DEFAULT_MODELS = [ |
| 39 | "nvidia/DeepSeek-R1-FP4", |
| 40 | ] |
| 41 | |
| 42 | DEFAULT_BATCH_SIZES = [4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048] |
| 43 | DEFAULT_TP_SIZES = [1] |
| 44 | |
| 45 | PER_ACT_TOKEN_OPTS = [False] |
| 46 | PER_OUT_CH_OPTS = [False] |
| 47 | FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max() |
| 48 | FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max |
| 49 | |
| 50 | |
| 51 | def to_fp8(tensor: torch.Tensor): |
| 52 | finfo = torch.finfo(torch.float8_e4m3fn) |
| 53 | return torch.round(tensor.clamp(min=finfo.min, max=finfo.max)).to( |
| 54 | dtype=torch.float8_e4m3fn |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | def bench_run( |
| 59 | results: list[benchmark.Measurement], |
| 60 | model: str, |
| 61 | num_experts: int, |
| 62 | topk: int, |
| 63 | per_act_token: bool, |
| 64 | per_out_ch: bool, |
| 65 | mkn: tuple[int, int, int], |
| 66 | ): |
| 67 | label = "NVFP4 Blockscaled CUTLASS MOE vs FP8 Tensor Scaled Triton" |
| 68 | |
| 69 | sub_label = ( |
| 70 | "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, MKN=({})".format( |
| 71 | model, num_experts, topk, per_act_token, per_out_ch, mkn |
| 72 | ) |
| 73 | ) |
| 74 | |
| 75 | print(f"Testing: {sub_label}") |
| 76 | |
| 77 | (m, k, n) = mkn |
| 78 | |
| 79 | dtype = torch.half |
| 80 | device = "cuda" |
| 81 | a = torch.randn((m, k), device=device, dtype=dtype) / 10 |
| 82 | w1 = torch.randn((num_experts, 2 * n, k), device=device, dtype=dtype) / 10 |
| 83 | w2 = torch.randn((num_experts, k, n), device=device, dtype=dtype) / 10 |
| 84 | |
| 85 | _, a_fp8_scale = ops.scaled_fp8_quant(a) |
| 86 | |
| 87 | w1_fp8q = torch.empty( |
| 88 | (num_experts, 2 * n, k), device=device, dtype=torch.float8_e4m3fn |
| 89 | ) |
| 90 | w2_fp8q = torch.empty((num_experts, k, n), device=device, dtype=torch.float8_e4m3fn) |
| 91 | w1_fp8scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32) |
| 92 | w2_fp8scale = torch.empty((num_experts, 1, 1), device=device, dtype=torch.float32) |
| 93 | |
| 94 | for expert in range(num_experts): |
| 95 | w1_fp8q[expert], w1_fp8scale[expert] = ops.scaled_fp8_quant(w1[expert]) |
| 96 | w2_fp8q[expert], w2_fp8scale[expert] = ops.scaled_fp8_quant(w2[expert]) |
| 97 | |
| 98 | w1_fp8q_notransp = w1_fp8q.clone() |
| 99 | w2_fp8q_notransp = w2_fp8q.clone() |
| 100 | w1_fp8q = w1_fp8q.transpose(1, 2) |
| 101 | w2_fp8q = w2_fp8q.transpose(1, 2) |
| 102 | |
| 103 | score = torch.randn((m, num_experts), device=device, dtype=dtype) |
| 104 | |
| 105 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) |
| 106 | |
| 107 | quant_blocksize = 16 |
| 108 | w1_blockscale = torch.empty( |
| 109 | (num_experts, 2 * n, k // quant_blocksize), |
| 110 | device=device, |
| 111 | dtype=torch.float8_e4m3fn, |
| 112 | ) |
| 113 | w2_blockscale = torch.empty( |
| 114 | (num_experts, k, n // quant_blocksize), device=device, dtype=torch.float8_e4m3fn |
| 115 | ) |
| 116 | |
| 117 | # n_b_scales = 2 * n if per_out_ch else 1 |
| 118 | # k_b_scales = k if per_out_ch else 1 |
| 119 | w1_fp4 = torch.empty((num_experts, 2 * n, k // 2), device=device, dtype=torch.uint8) |
| 120 | w2_fp4 = torch.empty((num_experts, k, n // 2), device=device, dtype=torch.uint8) |
| 121 | |
| 122 | w1_gs = torch.empty((num_experts,), device=device, dtype=torch.float32) |
| 123 | w2_gs = torch.empty((num_experts,), device=device, dtype=torch.float32) |
| 124 | a1_gs = torch.ones((num_experts,), device=device, dtype=torch.float32) |
| 125 | a2_gs = torch.ones((num_experts,), device=device, dtype=torch.float32) |
| 126 | |
| 127 | for expert in range(num_experts): |
| 128 | w1_e = w1[expert] |
| 129 | w2_e = w2[expert] |
| 130 | w1_amax = torch.abs(w1_e).max().to(torch.float32) |
| 131 | w2_amax = torch.abs(w2_e).max().to(torch.float32) |
| 132 | w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax |
| 133 | w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax |
| 134 | |
| 135 | w1_fp4[expert], w1_blockscale[expert] = ops.scaled_fp4_quant( |
| 136 | w1_e, w1_gs[expert] |
| 137 | ) |
| 138 | |
| 139 | w2_fp4[expert], w2_blockscale[expert] = ops.scaled_fp4_quant( |
| 140 | w2_e, w2_gs[expert] |
| 141 | ) |
| 142 | |
| 143 | def run_triton_moe( |
| 144 | a: torch.Tensor, |
| 145 | w1: torch.Tensor, |
| 146 | w2: torch.Tensor, |
| 147 | topk_weights: torch.Tensor, |
| 148 | topk_ids: torch.Tensor, |
| 149 | w1_scale: torch.Tensor, |
| 150 | w2_scale: torch.Tensor, |
| 151 | a_fp8_scale: torch.Tensor, |
| 152 | num_repeats: int, |
| 153 | ): |
| 154 | quant_config = fp8_w8a8_moe_quant_config( |
| 155 | w1_scale=w1_scale, |
| 156 | w2_scale=w2_scale, |
| 157 | a1_scale=a_fp8_scale, |
| 158 | ) |
| 159 | |
| 160 | for _ in range(num_repeats): |
| 161 | fused_experts( |
| 162 | a, |
| 163 | w1, |
| 164 | w2, |
| 165 | topk_weights, |
| 166 | topk_ids, |
| 167 | quant_config=quant_config, |
| 168 | ) |
| 169 | |
| 170 | def run_cutlass_moe_fp4( |
| 171 | a: torch.Tensor, |
| 172 | w1_fp4: torch.Tensor, |
| 173 | w2_fp4: torch.Tensor, |
| 174 | w1_blockscale: torch.Tensor, |
| 175 | w2_blockscale: torch.Tensor, |
| 176 | w1_gs: torch.Tensor, |
| 177 | w2_gs: torch.Tensor, |
| 178 | a1_gs: torch.Tensor, |
| 179 | a2_gs: torch.Tensor, |
| 180 | topk_weights: torch.Tensor, |
| 181 | topk_ids: torch.Tensor, |
| 182 | m: int, |
| 183 | n: int, |
| 184 | k: int, |
| 185 | e: int, |
| 186 | device: torch.device, |
| 187 | num_repeats: int, |
| 188 | ): |
| 189 | quant_config = nvfp4_moe_quant_config( |
| 190 | a1_gscale=a1_gs, |
| 191 | a2_gscale=a2_gs, |
| 192 | w1_scale=w1_blockscale, |
| 193 | w2_scale=w2_blockscale, |
| 194 | g1_alphas=w1_gs, |
| 195 | g2_alphas=w2_gs, |
| 196 | ) |
| 197 | |
| 198 | moe_config = make_dummy_moe_config( |
| 199 | num_experts=num_experts, |
| 200 | hidden_dim=k, |
| 201 | intermediate_size_per_partition=n, |
| 202 | in_dtype=a.dtype, |
| 203 | ) |
| 204 | kernel = mk.FusedMoEKernel( |
| 205 | maybe_make_prepare_finalize( |
| 206 | moe=moe_config, |
| 207 | quant_config=quant_config, |
| 208 | allow_new_interface=True, |
| 209 | use_monolithic=False, |
| 210 | ), |
| 211 | CutlassExpertsFp4( |
| 212 | moe_config=moe_config, |
| 213 | quant_config=quant_config, |
| 214 | ), |
| 215 | ) |
| 216 | |
| 217 | for _ in range(num_repeats): |
| 218 | with nvtx.annotate("cutlass_moe_fp4", color="green"): |
| 219 | kernel( |
| 220 | hidden_states=a, |
| 221 | w1=w1_fp4, |
| 222 | w2=w2_fp4, |
| 223 | topk_weights=topk_weights, |
| 224 | topk_ids=topk_ids, |
| 225 | ) |
| 226 | |
| 227 | def run_cutlass_from_graph( |
| 228 | a: torch.Tensor, |
| 229 | a1_gscale: torch.Tensor, |
| 230 | w1_fp4: torch.Tensor, |
| 231 | w1_blockscale: torch.Tensor, |
| 232 | w1_alphas: torch.Tensor, |
| 233 | a2_gscale: torch.Tensor, |
| 234 | w2_fp4: torch.Tensor, |
| 235 | w2_blockscale: torch.Tensor, |
| 236 | w2_alphas: torch.Tensor, |
| 237 | topk_weights: torch.Tensor, |
| 238 | topk_ids: torch.Tensor, |
| 239 | m: int, |
| 240 | n: int, |
| 241 | k: int, |
| 242 | e: int, |
| 243 | device: torch.device, |
| 244 | ): |
| 245 | quant_config = nvfp4_moe_quant_config( |
| 246 | a1_gscale=a1_gs, |
| 247 | a2_gscale=a2_gs, |
| 248 | w1_scale=w1_blockscale, |
| 249 | w2_scale=w2_blockscale, |
| 250 | g1_alphas=w1_gs, |
| 251 | g2_alphas=w2_gs, |
| 252 | ) |
| 253 | moe_config = make_dummy_moe_config() |
| 254 | |
| 255 | kernel = mk.FusedMoEKernel( |
| 256 | maybe_make_prepare_finalize( |
| 257 | moe=moe_config, |
| 258 | quant_config=quant_config, |
| 259 | allow_new_interface=True, |
| 260 | use_monolithic=False, |
| 261 | ), |
| 262 | CutlassExpertsFp4( |
| 263 | moe_config=moe_config, |
| 264 | quant_config=quant_config, |
| 265 | ), |
| 266 | ) |
| 267 | |
| 268 | with set_current_vllm_config( |
| 269 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 270 | ): |
| 271 | return kernel( |
| 272 | hidden_states=a, |
| 273 | w1=w1_fp4, |
| 274 | w2=w2_fp4, |
| 275 | topk_weights=topk_weights, |
| 276 | topk_ids=topk_ids, |
| 277 | ) |
| 278 | |
| 279 | def run_triton_from_graph( |
| 280 | a: torch.Tensor, |
| 281 | w1: torch.Tensor, |
| 282 | w2: torch.Tensor, |
| 283 | topk_weights: torch.Tensor, |
| 284 | topk_ids: torch.Tensor, |
| 285 | w1_scale: torch.Tensor, |
| 286 | w2_scale: torch.Tensor, |
| 287 | a_fp8_scale: torch.Tensor, |
| 288 | ): |
| 289 | with set_current_vllm_config( |
| 290 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 291 | ): |
| 292 | quant_config = fp8_w8a8_moe_quant_config( |
| 293 | w1_scale=w1_scale, |
| 294 | w2_scale=w2_scale, |
| 295 | a1_scale=a_fp8_scale, |
| 296 | ) |
| 297 | return fused_experts( |
| 298 | a, |
| 299 | w1, |
| 300 | w2, |
| 301 | topk_weights, |
| 302 | topk_ids, |
| 303 | quant_config=quant_config, |
| 304 | ) |
| 305 | |
| 306 | def replay_graph(graph, num_repeats): |
| 307 | for _ in range(num_repeats): |
| 308 | graph.replay() |
| 309 | torch.accelerator.synchronize() |
| 310 | |
| 311 | cutlass_stream = torch.cuda.Stream() |
| 312 | cutlass_graph = torch.cuda.CUDAGraph() |
| 313 | with torch.cuda.graph(cutlass_graph, stream=cutlass_stream): |
| 314 | run_cutlass_from_graph( |
| 315 | a=a, |
| 316 | a1_gscale=a1_gs, |
| 317 | w1_fp4=w1_fp4, |
| 318 | w1_blockscale=w1_blockscale, |
| 319 | w1_alphas=w1_gs, |
| 320 | a2_gscale=a2_gs, |
| 321 | w2_fp4=w2_fp4, |
| 322 | w2_blockscale=w2_blockscale, |
| 323 | w2_alphas=w2_gs, |
| 324 | topk_weights=topk_weights, |
| 325 | topk_ids=topk_ids, |
| 326 | m=m, |
| 327 | n=n, |
| 328 | k=k, |
| 329 | e=num_experts, |
| 330 | device=device, |
| 331 | ) |
| 332 | torch.accelerator.synchronize() |
| 333 | |
| 334 | triton_stream = torch.cuda.Stream() |
| 335 | triton_graph = torch.cuda.CUDAGraph() |
| 336 | with torch.cuda.graph(triton_graph, stream=triton_stream): |
| 337 | run_triton_from_graph( |
| 338 | a, |
| 339 | w1_fp8q_notransp, |
| 340 | w2_fp8q_notransp, |
| 341 | topk_weights, |
| 342 | topk_ids, |
| 343 | w1_fp8scale, |
| 344 | w2_fp8scale, |
| 345 | a_fp8_scale, |
| 346 | ) |
| 347 | torch.accelerator.synchronize() |
| 348 | |
| 349 | min_run_time = 5 |
| 350 | num_warmup = 5 |
| 351 | num_runs = 25 |
| 352 | |
| 353 | globals = { |
| 354 | # Baseline params |
| 355 | "w1": w1, |
| 356 | "w2": w2, |
| 357 | "score": score, |
| 358 | "topk": topk, |
| 359 | "w1_fp8q_notransp": w1_fp8q_notransp, |
| 360 | "w2_fp8q_notransp": w2_fp8q_notransp, |
| 361 | "w1_fp8scale": w1_fp8scale, |
| 362 | "w2_fp8scale": w2_fp8scale, |
| 363 | "a_fp8_scale": a_fp8_scale, |
| 364 | # Cutlass params |
| 365 | "a": a, |
| 366 | "a1_gscale": a1_gs, |
| 367 | "w1_fp4": w1_fp4, |
| 368 | "w1_blockscale": w1_blockscale, |
| 369 | "w1_alphas": w1_gs, |
| 370 | "a2_gscale": a2_gs, |
| 371 | "w2_fp4": w2_fp4, |
| 372 | "w2_blockscale": w2_blockscale, |
| 373 | "w2_alphas": w2_gs, |
| 374 | "topk_weights": topk_weights, |
| 375 | "topk_ids": topk_ids, |
| 376 | "m": m, |
| 377 | "n": n, |
| 378 | "k": k, |
| 379 | "e": num_experts, |
| 380 | "device": device, |
| 381 | # cuda graph params |
| 382 | "cutlass_graph": cutlass_graph, |
| 383 | "triton_graph": triton_graph, |
| 384 | # Gen params |
| 385 | "num_runs": num_runs, |
| 386 | # Kernels |
| 387 | "run_triton_moe": run_triton_moe, |
| 388 | "run_cutlass_moe_fp4": run_cutlass_moe_fp4, |
| 389 | "replay_graph": replay_graph, |
| 390 | } |
| 391 | |
| 392 | # Warmup |
| 393 | run_triton_moe( |
| 394 | a, |
| 395 | w1_fp8q_notransp, |
| 396 | w2_fp8q_notransp, |
| 397 | topk_weights, |
| 398 | topk_ids, |
| 399 | w1_fp8scale, |
| 400 | w2_fp8scale, |
| 401 | a_fp8_scale, |
| 402 | num_warmup, |
| 403 | ) |
| 404 | |
| 405 | results.append( |
| 406 | benchmark.Timer( |
| 407 | stmt="run_triton_moe(a, w1_fp8q_notransp, w2_fp8q_notransp, topk_weights, topk_ids, w1_fp8scale, w2_fp8scale, a_fp8_scale, num_runs)", # noqa: E501 |
| 408 | globals=globals, |
| 409 | label=label, |
| 410 | sub_label=sub_label, |
| 411 | description="triton_moe", |
| 412 | ).blocked_autorange(min_run_time=min_run_time) |
| 413 | ) |
| 414 | |
| 415 | # Warmup |
| 416 | replay_graph(triton_graph, num_warmup) |
| 417 | |
| 418 | results.append( |
| 419 | benchmark.Timer( |
| 420 | stmt="replay_graph(triton_graph, num_runs)", |
| 421 | globals=globals, |
| 422 | label=label, |
| 423 | sub_label=sub_label, |
| 424 | description="triton_moe_cuda_graphs", |
| 425 | ).blocked_autorange(min_run_time=min_run_time) |
| 426 | ) |
| 427 | |
| 428 | # Warmup |
| 429 | |
| 430 | run_cutlass_moe_fp4( |
| 431 | a, |
| 432 | w1_fp4, |
| 433 | w2_fp4, |
| 434 | w1_blockscale, |
| 435 | w2_blockscale, |
| 436 | w1_gs, |
| 437 | w2_gs, |
| 438 | a1_gs, |
| 439 | a2_gs, |
| 440 | topk_weights, |
| 441 | topk_ids, |
| 442 | m, |
| 443 | n, |
| 444 | k, |
| 445 | num_experts, |
| 446 | device, |
| 447 | num_warmup, |
| 448 | ) |
| 449 | |
| 450 | results.append( |
| 451 | benchmark.Timer( |
| 452 | stmt="run_cutlass_moe_fp4(a, w1_fp4, w2_fp4, w1_blockscale, w2_blockscale, w1_alphas, w2_alphas, a1_gscale, a2_gscale, topk_weights, topk_ids, m, n, k, e, device, num_runs)", # noqa: E501 |
| 453 | globals=globals, |
| 454 | label=label, |
| 455 | sub_label=sub_label, |
| 456 | description="cutlass_moe_fp4", |
| 457 | ).blocked_autorange(min_run_time=min_run_time) |
| 458 | ) |
| 459 | |
| 460 | # Warmup |
| 461 | replay_graph(cutlass_graph, num_warmup) |
| 462 | |
| 463 | results.append( |
| 464 | benchmark.Timer( |
| 465 | stmt="replay_graph(cutlass_graph, num_runs)", |
| 466 | globals=globals, |
| 467 | label=label, |
| 468 | sub_label=sub_label, |
| 469 | description="cutlass_moe_fp4_cuda_graphs", |
| 470 | ).blocked_autorange(min_run_time=min_run_time) |
| 471 | ) |
| 472 | |
| 473 | |
| 474 | def main(args): |
| 475 | # Initialize workspace manager (required for CUTLASS MoE kernels) |
| 476 | device = torch.device("cuda:0") |
| 477 | init_workspace_manager(device) |
| 478 | |
| 479 | print("Benchmarking models:") |
| 480 | for i, model in enumerate(args.models): |
| 481 | print(f"[{i}] {model}") |
| 482 | |
| 483 | results: list[benchmark.Measurement] = [] |
| 484 | |
| 485 | for model in args.models: |
| 486 | for tp in args.tp_sizes: |
| 487 | for layer in WEIGHT_SHAPES_MOE[model]: |
| 488 | num_experts = layer[0] |
| 489 | topk = layer[1] |
| 490 | size_k = layer[2] |
| 491 | size_n = layer[3] // tp |
| 492 | |
| 493 | if len(args.limit_k) > 0 and size_k not in args.limit_k: |
| 494 | continue |
| 495 | |
| 496 | if len(args.limit_n) > 0 and size_n not in args.limit_n: |
| 497 | continue |
| 498 | |
| 499 | for per_act_token in PER_ACT_TOKEN_OPTS: |
| 500 | for per_out_ch in PER_OUT_CH_OPTS: |
| 501 | for size_m in args.batch_sizes: |
| 502 | mkn = (size_m, size_k, size_n) |
| 503 | bench_run( |
| 504 | results, |
| 505 | model, |
| 506 | num_experts, |
| 507 | topk, |
| 508 | per_act_token, |
| 509 | per_out_ch, |
| 510 | mkn, |
| 511 | ) |
| 512 | |
| 513 | compare = benchmark.Compare(results) |
| 514 | compare.print() |
| 515 | |
| 516 | |
| 517 | if __name__ == "__main__": |
| 518 | parser = FlexibleArgumentParser( |
| 519 | description="Benchmark NVFP4 CUTLASS MOE across specified models/shapes/batches" |
| 520 | ) |
| 521 | parser.add_argument( |
| 522 | "--models", |
| 523 | nargs="+", |
| 524 | type=str, |
| 525 | default=DEFAULT_MODELS, |
| 526 | choices=WEIGHT_SHAPES_MOE.keys(), |
| 527 | ) |
| 528 | parser.add_argument("--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES) |
| 529 | parser.add_argument( |
| 530 | "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES |
| 531 | ) |
| 532 | parser.add_argument("--limit-k", nargs="+", type=int, default=[]) |
| 533 | parser.add_argument("--limit-n", nargs="+", type=int, default=[]) |
| 534 | parser.add_argument("--limit-num-groups", nargs="+", type=int, default=[]) |
| 535 | parser.add_argument("--limit-per-act-token", nargs="+", type=int, default=[]) |
| 536 | parser.add_argument("--limit-per-out-ch", nargs="+", type=int, default=[]) |
| 537 | |
| 538 | args = parser.parse_args() |
| 539 | main(args) |