| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Benchmark for FlashInfer fused collective operations vs standard operations. |
| 5 | |
| 6 | This benchmark compares: |
| 7 | 1. FlashInfer's allreduce_fusion with trtllm backend |
| 8 | (fused allreduce + rmsnorm + optional FP8/FP4 quant) |
| 9 | 2. FlashInfer's allreduce_fusion with mnnvl backend |
| 10 | (fused allreduce + rmsnorm only, no quantization support) |
| 11 | 3. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations |
| 12 | |
| 13 | Usage with torchrun: |
| 14 | torchrun --nproc_per_node=2 benchmark_fused_collective.py |
| 15 | |
| 16 | """ |
| 17 | |
| 18 | import argparse |
| 19 | import itertools |
| 20 | import os |
| 21 | import time |
| 22 | |
| 23 | import pandas as pd |
| 24 | import torch # type: ignore |
| 25 | import torch.distributed as dist # type: ignore |
| 26 | |
| 27 | from vllm._custom_ops import create_fp4_output_tensors |
| 28 | from vllm.config.vllm import CompilationConfig, VllmConfig, set_current_vllm_config |
| 29 | from vllm.distributed import ( |
| 30 | tensor_model_parallel_all_reduce, |
| 31 | ) |
| 32 | from vllm.distributed.parallel_state import ( |
| 33 | graph_capture, |
| 34 | init_distributed_environment, |
| 35 | initialize_model_parallel, |
| 36 | ) |
| 37 | from vllm.logger import init_logger |
| 38 | from vllm.model_executor.layers.layernorm import RMSNorm # noqa |
| 39 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 # noqa |
| 40 | from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape # noqa |
| 41 | from vllm.platforms import current_platform # noqa |
| 42 | |
| 43 | RMS_NORM_OP = torch.ops._C.rms_norm |
| 44 | FUSED_ADD_RMS_NORM_OP = torch.ops._C.fused_add_rms_norm |
| 45 | RMS_NORM_STATIC_FP8_QUANT_OP = torch.ops._C.rms_norm_static_fp8_quant |
| 46 | FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT_OP = ( |
| 47 | torch.ops._C.fused_add_rms_norm_static_fp8_quant |
| 48 | ) |
| 49 | SCALED_FP4_QUANT_OUT_OP = torch.ops._C.scaled_fp4_quant.out |
| 50 | |
| 51 | logger = init_logger(__name__) |
| 52 | |
| 53 | # Try to import FlashInfer |
| 54 | TorchDistBackend = None |
| 55 | try: |
| 56 | import flashinfer.comm as flashinfer_comm # type: ignore |
| 57 | from flashinfer.comm.mnnvl import ( # type: ignore |
| 58 | TorchDistBackend, |
| 59 | ) |
| 60 | |
| 61 | if not ( |
| 62 | hasattr(flashinfer_comm, "allreduce_fusion") |
| 63 | and hasattr(flashinfer_comm, "create_allreduce_fusion_workspace") |
| 64 | ): |
| 65 | flashinfer_comm = None |
| 66 | logger.warning("FlashInfer comm module found but missing allreduce_fusion API") |
| 67 | except ImportError: |
| 68 | flashinfer_comm = None |
| 69 | logger.warning("FlashInfer not found, only benchmarking standard operations") |
| 70 | |
| 71 | # Constants |
| 72 | FP8_DTYPE = current_platform.fp8_dtype() |
| 73 | MiB = 1024 * 1024 |
| 74 | |
| 75 | # FlashInfer max sizes per world size |
| 76 | # Enable 64MB for 2, 4, 8 world sizes to verify large input sizes |
| 77 | # use --disable-oneshot to disable oneshot mode for very large input sizes |
| 78 | _FI_MAX_SIZES = { |
| 79 | 2: 64 * MiB, # 64MB |
| 80 | 4: 64 * MiB, # 64MB |
| 81 | 8: 64 * MiB, # 64MB |
| 82 | 16: 64 * MiB, # 64MB (multi-node) |
| 83 | } |
| 84 | |
| 85 | # Global workspace tensors for FlashInfer (keyed by backend name) |
| 86 | _FI_WORKSPACES: dict = {} |
| 87 | |
| 88 | # Backends to benchmark. trtllm is single-node only and can hang cross-node, so |
| 89 | # multi-node sweeps can restrict to mnnvl via FI_BACKENDS=mnnvl. |
| 90 | FLASHINFER_BACKENDS = [ |
| 91 | b for b in os.environ.get("FI_BACKENDS", "trtllm,mnnvl").split(",") if b |
| 92 | ] |
| 93 | |
| 94 | |
| 95 | def setup_flashinfer_workspace( |
| 96 | backend: str, |
| 97 | world_size: int, |
| 98 | rank: int, |
| 99 | hidden_dim: int, |
| 100 | max_token_num: int, |
| 101 | dtype: torch.dtype, |
| 102 | ): |
| 103 | """Setup FlashInfer workspace for fused allreduce operations.""" |
| 104 | global FI_WORKSPACES |
| 105 | |
| 106 | if flashinfer_comm is None: |
| 107 | return None |
| 108 | |
| 109 | if world_size not in _FI_MAX_SIZES: |
| 110 | logger.warning("FlashInfer not supported for world size %s", world_size) |
| 111 | return None |
| 112 | |
| 113 | try: |
| 114 | kwargs = {} |
| 115 | if TorchDistBackend is not None: |
| 116 | kwargs["comm_backend"] = TorchDistBackend(group=dist.group.WORLD) |
| 117 | |
| 118 | workspace = flashinfer_comm.create_allreduce_fusion_workspace( |
| 119 | backend=backend, |
| 120 | world_size=world_size, |
| 121 | rank=rank, |
| 122 | max_token_num=max_token_num, |
| 123 | hidden_dim=hidden_dim, |
| 124 | dtype=dtype, |
| 125 | **kwargs, |
| 126 | ) |
| 127 | |
| 128 | _FI_WORKSPACES[backend] = workspace |
| 129 | return workspace |
| 130 | except Exception as e: |
| 131 | logger.error( |
| 132 | "Failed to setup FlashInfer workspace (backend=%s): %s", backend, e |
| 133 | ) |
| 134 | return None |
| 135 | |
| 136 | |
| 137 | def cleanup_flashinfer_workspaces(): |
| 138 | """Cleanup all FlashInfer workspaces.""" |
| 139 | if flashinfer_comm is None: |
| 140 | return |
| 141 | |
| 142 | for backend, workspace in _FI_WORKSPACES.items(): |
| 143 | try: |
| 144 | workspace.destroy() |
| 145 | except Exception as e: |
| 146 | logger.error( |
| 147 | "Failed to cleanup FlashInfer workspace (backend=%s): %s", |
| 148 | backend, |
| 149 | e, |
| 150 | ) |
| 151 | _FI_WORKSPACES.clear() |
| 152 | |
| 153 | |
| 154 | class FlashInferFusedAllReduceParams: |
| 155 | """Parameters for FlashInfer fused allreduce operations.""" |
| 156 | |
| 157 | def __init__( |
| 158 | self, |
| 159 | max_token_num: int = 1024, |
| 160 | ): |
| 161 | self.launch_with_pdl = True |
| 162 | self.fp32_acc = True |
| 163 | self.max_token_num = max_token_num |
| 164 | |
| 165 | def get_flashinfer_fused_allreduce_kwargs(self): |
| 166 | return { |
| 167 | "launch_with_pdl": self.launch_with_pdl, |
| 168 | "fp32_acc": self.fp32_acc, |
| 169 | } |
| 170 | |
| 171 | |
| 172 | def flashinfer_fused_allreduce_rmsnorm( |
| 173 | input_tensor: torch.Tensor, |
| 174 | residual: torch.Tensor | None, |
| 175 | rms_gamma: torch.Tensor, |
| 176 | rms_eps: float, |
| 177 | allreduce_params: "FlashInferFusedAllReduceParams", |
| 178 | workspace: object, |
| 179 | use_oneshot: bool, |
| 180 | norm_out: torch.Tensor | None = None, |
| 181 | ): |
| 182 | """FlashInfer fused allreduce + rmsnorm operation.""" |
| 183 | if flashinfer_comm is None or workspace is None: |
| 184 | raise RuntimeError("FlashInfer not available or workspace not initialized") |
| 185 | |
| 186 | if norm_out is None: |
| 187 | norm_out = input_tensor |
| 188 | residual_out = residual |
| 189 | else: |
| 190 | residual_out = input_tensor |
| 191 | |
| 192 | layout_code = None |
| 193 | if workspace.backend == "trtllm": |
| 194 | layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4 |
| 195 | |
| 196 | flashinfer_comm.allreduce_fusion( |
| 197 | input=input_tensor, |
| 198 | workspace=workspace, |
| 199 | pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, |
| 200 | residual_in=residual, |
| 201 | residual_out=residual_out, |
| 202 | norm_out=norm_out, |
| 203 | rms_gamma=rms_gamma, |
| 204 | rms_eps=rms_eps, |
| 205 | quant_out=None, |
| 206 | scale_out=None, |
| 207 | layout_code=layout_code, |
| 208 | scale_factor=None, |
| 209 | use_oneshot=use_oneshot, |
| 210 | **allreduce_params.get_flashinfer_fused_allreduce_kwargs(), |
| 211 | ) |
| 212 | |
| 213 | |
| 214 | def flashinfer_fused_allreduce_rmsnorm_fp8_quant( |
| 215 | input_tensor: torch.Tensor, |
| 216 | residual: torch.Tensor | None, |
| 217 | rms_gamma: torch.Tensor, |
| 218 | rms_eps: float, |
| 219 | scale_factor: torch.Tensor, |
| 220 | allreduce_params: FlashInferFusedAllReduceParams, |
| 221 | workspace: object, |
| 222 | use_oneshot: bool = True, |
| 223 | norm_out: torch.Tensor | None = None, |
| 224 | quant_out: torch.Tensor | None = None, |
| 225 | ): |
| 226 | """FlashInfer fused allreduce + rmsnorm + FP8 quantization. |
| 227 | |
| 228 | Note: Only supported by the trtllm backend. |
| 229 | """ |
| 230 | if flashinfer_comm is None or workspace is None: |
| 231 | raise RuntimeError("FlashInfer not available or workspace not initialized") |
| 232 | |
| 233 | if norm_out is None: |
| 234 | norm_out = input_tensor |
| 235 | residual_out = residual |
| 236 | else: |
| 237 | residual_out = input_tensor |
| 238 | |
| 239 | flashinfer_comm.allreduce_fusion( |
| 240 | input=input_tensor, |
| 241 | workspace=workspace, |
| 242 | pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, |
| 243 | residual_in=residual, |
| 244 | residual_out=residual_out, |
| 245 | norm_out=norm_out, |
| 246 | rms_gamma=rms_gamma, |
| 247 | rms_eps=rms_eps, |
| 248 | quant_out=quant_out, |
| 249 | scale_out=None, |
| 250 | layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4, |
| 251 | scale_factor=scale_factor, |
| 252 | use_oneshot=use_oneshot, |
| 253 | **allreduce_params.get_flashinfer_fused_allreduce_kwargs(), |
| 254 | ) |
| 255 | |
| 256 | |
| 257 | def flashinfer_fused_allreduce_rmsnorm_fp4_quant( |
| 258 | input_tensor: torch.Tensor, |
| 259 | residual: torch.Tensor | None, |
| 260 | rms_gamma: torch.Tensor, |
| 261 | rms_eps: float, |
| 262 | input_global_scale: torch.Tensor, |
| 263 | allreduce_params: FlashInferFusedAllReduceParams, |
| 264 | workspace: object, |
| 265 | quant_out: torch.Tensor, |
| 266 | use_oneshot: bool, |
| 267 | output_scale: torch.Tensor, |
| 268 | norm_out: torch.Tensor | None = None, |
| 269 | ): |
| 270 | """FlashInfer fused allreduce + rmsnorm + FP4 quantization. |
| 271 | |
| 272 | Note: Only supported by the trtllm backend. |
| 273 | """ |
| 274 | if flashinfer_comm is None or workspace is None: |
| 275 | raise RuntimeError("FlashInfer not available or workspace not initialized") |
| 276 | |
| 277 | if norm_out is None: |
| 278 | norm_out = input_tensor |
| 279 | residual_out = residual |
| 280 | else: |
| 281 | residual_out = input_tensor |
| 282 | |
| 283 | flashinfer_comm.allreduce_fusion( |
| 284 | input=input_tensor, |
| 285 | workspace=workspace, |
| 286 | pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, |
| 287 | residual_in=residual, |
| 288 | residual_out=residual_out, |
| 289 | norm_out=norm_out, |
| 290 | rms_gamma=rms_gamma, |
| 291 | rms_eps=rms_eps, |
| 292 | quant_out=quant_out, |
| 293 | scale_out=output_scale, |
| 294 | layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4, |
| 295 | scale_factor=input_global_scale, |
| 296 | use_oneshot=use_oneshot, |
| 297 | **allreduce_params.get_flashinfer_fused_allreduce_kwargs(), |
| 298 | ) |
| 299 | |
| 300 | |
| 301 | class VllmFusedAllreduce: |
| 302 | def __init__(self, hidden_dim, dtype): |
| 303 | self.rms_eps = 1e-6 |
| 304 | self.rms_norm = RMSNorm(hidden_dim, eps=self.rms_eps, dtype=dtype) |
| 305 | self.fp8_quant = QuantFP8( |
| 306 | static=True, |
| 307 | group_shape=GroupShape.PER_TENSOR, |
| 308 | ) |
| 309 | |
| 310 | def allreduce_rmsnorm( |
| 311 | self, input_tensor: torch.Tensor, residual: torch.Tensor | None |
| 312 | ): |
| 313 | allreduce_out = tensor_model_parallel_all_reduce(input_tensor) |
| 314 | return self.rms_norm(allreduce_out, residual) |
| 315 | |
| 316 | def allreduce_rmsnorm_fp8_quant( |
| 317 | self, |
| 318 | input_tensor: torch.Tensor, |
| 319 | residual: torch.Tensor | None, |
| 320 | scale_factor: torch.Tensor, |
| 321 | ): |
| 322 | allreduce_out = tensor_model_parallel_all_reduce(input_tensor) |
| 323 | rms_out = self.rms_norm(allreduce_out, residual) |
| 324 | if residual is None: |
| 325 | quant_out = self.fp8_quant(rms_out, scale_factor) |
| 326 | return quant_out |
| 327 | else: |
| 328 | rms_out, residual_out = rms_out |
| 329 | quant_out = self.fp8_quant(rms_out, scale_factor) |
| 330 | return quant_out, residual_out |
| 331 | |
| 332 | def allreduce_rmsnorm_fp4_quant( |
| 333 | self, |
| 334 | input_tensor: torch.Tensor, |
| 335 | residual: torch.Tensor | None, |
| 336 | input_global_scale: torch.Tensor, |
| 337 | quant_out: torch.Tensor, |
| 338 | output_scale: torch.Tensor, |
| 339 | ): |
| 340 | allreduce_out = tensor_model_parallel_all_reduce(input_tensor) |
| 341 | rms_output = self.rms_norm(allreduce_out, residual) |
| 342 | if residual is None: |
| 343 | rms_out = rms_output |
| 344 | else: |
| 345 | rms_out, residual_out = rms_output |
| 346 | |
| 347 | SCALED_FP4_QUANT_OUT_OP( |
| 348 | rms_out, |
| 349 | input_global_scale, |
| 350 | True, |
| 351 | output=quant_out, |
| 352 | output_scale=output_scale, |
| 353 | ) |
| 354 | |
| 355 | if residual is None: |
| 356 | return quant_out, output_scale |
| 357 | else: |
| 358 | return quant_out, residual_out, output_scale |
| 359 | |
| 360 | |
| 361 | def create_test_tensors( |
| 362 | num_tokens: int, hidden_dim: int, dtype: torch.dtype, use_residual: bool = True |
| 363 | ): |
| 364 | """Create test tensors for benchmarking.""" |
| 365 | input_tensor = torch.randn(num_tokens, hidden_dim, dtype=dtype) |
| 366 | residual = ( |
| 367 | torch.randn_like(input_tensor) |
| 368 | if use_residual |
| 369 | else torch.zeros_like(input_tensor) |
| 370 | ) |
| 371 | rms_gamma = torch.ones(hidden_dim, dtype=dtype) |
| 372 | norm_out = None if use_residual else torch.empty_like(input_tensor) |
| 373 | |
| 374 | # Quantization scales |
| 375 | scale_fp8 = torch.tensor(1.0, dtype=torch.float32) |
| 376 | scale_fp4 = torch.tensor(1.0, dtype=torch.float32) |
| 377 | quant_out_fp8 = torch.empty_like(input_tensor, dtype=FP8_DTYPE) |
| 378 | # Pre-allocate FP4 output tensors (to avoid allocation overhead in benchmarks) |
| 379 | fp4_quant_out, fp4_output_scale = create_fp4_output_tensors( |
| 380 | num_tokens, hidden_dim, input_tensor.device, True |
| 381 | ) |
| 382 | |
| 383 | return ( |
| 384 | input_tensor, |
| 385 | norm_out, |
| 386 | residual, |
| 387 | rms_gamma, |
| 388 | scale_fp8, |
| 389 | quant_out_fp8, |
| 390 | scale_fp4, |
| 391 | fp4_quant_out, |
| 392 | fp4_output_scale, |
| 393 | ) |
| 394 | |
| 395 | |
| 396 | def benchmark_operation( |
| 397 | operation_func, *args, warmup: int = 5, trials: int = 20, **kwargs |
| 398 | ): |
| 399 | """Benchmark a single operation using CUDA graphs.""" |
| 400 | # Warmup before graph capture |
| 401 | for _ in range(warmup): |
| 402 | operation_func(*args, **kwargs) |
| 403 | torch.accelerator.synchronize() |
| 404 | |
| 405 | # Create CUDA graph |
| 406 | graph = torch.cuda.CUDAGraph() |
| 407 | num_op_per_cudagraph = 10 |
| 408 | |
| 409 | # Use vLLM's graph_capture to make tensor_model_parallel_all_reduce graph-safe |
| 410 | device = torch.device(f"cuda:{torch.accelerator.current_device_index()}") |
| 411 | with graph_capture(device=device), torch.cuda.graph(graph): |
| 412 | for _ in range(num_op_per_cudagraph): |
| 413 | operation_func(*args, **kwargs) |
| 414 | |
| 415 | # Graph warmup |
| 416 | torch.accelerator.synchronize() |
| 417 | for _ in range(warmup): |
| 418 | graph.replay() |
| 419 | |
| 420 | # Benchmark with CUDA graph |
| 421 | torch.accelerator.synchronize() |
| 422 | start_time = time.perf_counter() |
| 423 | |
| 424 | for _ in range(trials // num_op_per_cudagraph): |
| 425 | # operation_func(*args, **kwargs) |
| 426 | graph.replay() |
| 427 | |
| 428 | torch.accelerator.synchronize() |
| 429 | end_time = time.perf_counter() |
| 430 | |
| 431 | avg_time_ms = ((end_time - start_time) / trials) * 1000 |
| 432 | return avg_time_ms |
| 433 | |
| 434 | |
| 435 | def run_benchmarks( |
| 436 | num_tokens: int, |
| 437 | hidden_dim: int, |
| 438 | dtype: torch.dtype, |
| 439 | use_residual: bool, |
| 440 | allreduce_params: FlashInferFusedAllReduceParams | None, |
| 441 | workspaces: dict, |
| 442 | quant_modes: set[str], |
| 443 | no_oneshot: bool, |
| 444 | ): |
| 445 | """Run all benchmarks for given configuration. |
| 446 | |
| 447 | Args: |
| 448 | num_tokens: Number of tokens in the benchmark input. |
| 449 | hidden_dim: Hidden dimension of the benchmark input. |
| 450 | dtype: dtype of the benchmark input. |
| 451 | use_residual: Whether to include a residual add in the fused op. |
| 452 | no_oneshot: Skip the one-shot allreduce variants. |
| 453 | allreduce_params: Shared parameters for FlashInfer fused allreduce. |
| 454 | workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace. |
| 455 | quant_modes: Set of quantization modes: "none", "fp8", "fp4". |
| 456 | |
| 457 | """ |
| 458 | ( |
| 459 | input_tensor, |
| 460 | norm_out, |
| 461 | residual, |
| 462 | rms_gamma, |
| 463 | scale_fp8, |
| 464 | quant_out_fp8, |
| 465 | scale_fp4, |
| 466 | fp4_quant_out, |
| 467 | fp4_output_scale, |
| 468 | ) = create_test_tensors(num_tokens, hidden_dim, dtype, use_residual) |
| 469 | |
| 470 | rms_eps = 1e-6 |
| 471 | results = {} |
| 472 | use_oneshot_options = [False] if no_oneshot else [True, False] |
| 473 | |
| 474 | if "none" in quant_modes: |
| 475 | # Standard AllReduce + RMSNorm |
| 476 | # Re-create VllmFusedAllreduce per config so CustomOp binds the |
| 477 | # correct forward method (native vs custom kernel). |
| 478 | for custom_op in ["-rms_norm", "+rms_norm"]: |
| 479 | with set_current_vllm_config( |
| 480 | VllmConfig(compilation_config=CompilationConfig(custom_ops=[custom_op])) |
| 481 | ): |
| 482 | try: |
| 483 | vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype) |
| 484 | suffix = ( |
| 485 | "_custom_rms_norm" if "+" in custom_op else "_native_rms_norm" |
| 486 | ) |
| 487 | time_ms = benchmark_operation( |
| 488 | vllm_fused_allreduce.allreduce_rmsnorm, |
| 489 | input_tensor, |
| 490 | residual=residual, |
| 491 | ) |
| 492 | results[f"standard_allreduce_{suffix}"] = time_ms |
| 493 | except Exception as e: |
| 494 | logger.error("Standard AllReduce+RMSNorm failed: %s", e) |
| 495 | results[f"standard_allreduce_{suffix}"] = float("inf") |
| 496 | |
| 497 | # Standard AllReduce + RMSNorm Native Compiled |
| 498 | with set_current_vllm_config( |
| 499 | VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"])) |
| 500 | ): |
| 501 | try: |
| 502 | vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype) |
| 503 | standard_allreduce_rmsnorm_native_compiled = torch.compile( |
| 504 | vllm_fused_allreduce.allreduce_rmsnorm, |
| 505 | fullgraph=True, |
| 506 | dynamic=False, |
| 507 | ) |
| 508 | time_ms = benchmark_operation( |
| 509 | standard_allreduce_rmsnorm_native_compiled, |
| 510 | input_tensor, |
| 511 | residual=residual, |
| 512 | ) |
| 513 | results["standard_allreduce_rmsnorm_native_compiled"] = time_ms |
| 514 | except Exception as e: |
| 515 | logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e) |
| 516 | results["standard_allreduce_rmsnorm_native_compiled"] = float("inf") |
| 517 | |
| 518 | # FlashInfer Fused AllReduce + RMSNorm (all backends) |
| 519 | for backend, workspace in workspaces.items(): |
| 520 | for use_oneshot in use_oneshot_options: |
| 521 | suffix = "_oneshot" if use_oneshot else "_twoshot" |
| 522 | key = f"flashinfer_{backend}_fused_allreduce_rmsnorm{suffix}" |
| 523 | try: |
| 524 | time_ms = benchmark_operation( |
| 525 | flashinfer_fused_allreduce_rmsnorm, |
| 526 | input_tensor, |
| 527 | residual=residual, |
| 528 | norm_out=norm_out, |
| 529 | rms_gamma=rms_gamma, |
| 530 | rms_eps=rms_eps, |
| 531 | allreduce_params=allreduce_params, |
| 532 | workspace=workspace, |
| 533 | use_oneshot=use_oneshot, |
| 534 | ) |
| 535 | results[key] = time_ms |
| 536 | except Exception as e: |
| 537 | logger.error( |
| 538 | "FlashInfer (%s) Fused AllReduce+RMSNorm failed: %s", |
| 539 | backend, |
| 540 | e, |
| 541 | ) |
| 542 | results[key] = float("inf") |
| 543 | |
| 544 | if "fp8" in quant_modes: |
| 545 | # Standard AllReduce + RMSNorm + FP8 Quant |
| 546 | for rms_norm_custom_op in ["-rms_norm", "+rms_norm"]: |
| 547 | suffix = ( |
| 548 | "_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm" |
| 549 | ) |
| 550 | for quant_fp8_custom_op in ["-quant_fp8", "+quant_fp8"]: |
| 551 | op_suffix = suffix + ( |
| 552 | "_custom_quant_fp8" |
| 553 | if "+" in quant_fp8_custom_op |
| 554 | else "_native_quant_fp8" |
| 555 | ) |
| 556 | with set_current_vllm_config( |
| 557 | VllmConfig( |
| 558 | compilation_config=CompilationConfig( |
| 559 | custom_ops=[rms_norm_custom_op, quant_fp8_custom_op] |
| 560 | ) |
| 561 | ) |
| 562 | ): |
| 563 | try: |
| 564 | vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype) |
| 565 | time_ms = benchmark_operation( |
| 566 | vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant, |
| 567 | input_tensor, |
| 568 | residual=residual, |
| 569 | scale_factor=scale_fp8, |
| 570 | ) |
| 571 | results[f"standard_allreduce{op_suffix}"] = time_ms |
| 572 | except Exception as e: |
| 573 | logger.error("Standard AllReduce+RMSNorm+FP8 failed: %s", e) |
| 574 | results[f"standard_allreduce{op_suffix}"] = float("inf") |
| 575 | |
| 576 | # Standard AllReduce + RMSNorm + FP8 Quant Native Compiled |
| 577 | with set_current_vllm_config( |
| 578 | VllmConfig( |
| 579 | compilation_config=CompilationConfig( |
| 580 | custom_ops=["-rms_norm", "-quant_fp8"] |
| 581 | ) |
| 582 | ) |
| 583 | ): |
| 584 | try: |
| 585 | vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype) |
| 586 | standard_allreduce_rmsnorm_fp8_quant_native_compiled = torch.compile( |
| 587 | vllm_fused_allreduce.allreduce_rmsnorm_fp8_quant, |
| 588 | fullgraph=True, |
| 589 | dynamic=False, |
| 590 | ) |
| 591 | time_ms = benchmark_operation( |
| 592 | standard_allreduce_rmsnorm_fp8_quant_native_compiled, |
| 593 | input_tensor, |
| 594 | residual=residual, |
| 595 | scale_factor=scale_fp8, |
| 596 | ) |
| 597 | results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = ( |
| 598 | time_ms |
| 599 | ) |
| 600 | except Exception as e: |
| 601 | logger.error( |
| 602 | "Standard AllReduce+RMSNorm+FP8 Native Compiled failed: %s", e |
| 603 | ) |
| 604 | results["standard_allreduce_rmsnorm_fp8_quant_native_compiled"] = float( |
| 605 | "inf" |
| 606 | ) |
| 607 | |
| 608 | # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant (trtllm only) |
| 609 | if "trtllm" in workspaces: |
| 610 | trtllm_ws = workspaces["trtllm"] |
| 611 | for use_oneshot in use_oneshot_options: |
| 612 | suffix = "_oneshot" if use_oneshot else "_twoshot" |
| 613 | key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp8_quant{suffix}" |
| 614 | try: |
| 615 | time_ms = benchmark_operation( |
| 616 | flashinfer_fused_allreduce_rmsnorm_fp8_quant, |
| 617 | input_tensor, |
| 618 | norm_out=norm_out, |
| 619 | residual=residual, |
| 620 | rms_gamma=rms_gamma, |
| 621 | rms_eps=rms_eps, |
| 622 | scale_factor=scale_fp8, |
| 623 | quant_out=quant_out_fp8, |
| 624 | allreduce_params=allreduce_params, |
| 625 | workspace=trtllm_ws, |
| 626 | use_oneshot=use_oneshot, |
| 627 | ) |
| 628 | results[key] = time_ms |
| 629 | except Exception as e: |
| 630 | logger.error( |
| 631 | "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP8 failed: %s", |
| 632 | e, |
| 633 | ) |
| 634 | results[key] = float("inf") |
| 635 | |
| 636 | if "fp4" in quant_modes and current_platform.has_device_capability(100): |
| 637 | # Standard AllReduce + RMSNorm + FP4 Quant |
| 638 | for rms_norm_custom_op in ["-rms_norm", "+rms_norm"]: |
| 639 | suffix = ( |
| 640 | "_custom_rms_norm" if "+" in rms_norm_custom_op else "_native_rms_norm" |
| 641 | ) |
| 642 | with set_current_vllm_config( |
| 643 | VllmConfig( |
| 644 | compilation_config=CompilationConfig( |
| 645 | custom_ops=[rms_norm_custom_op] |
| 646 | ) |
| 647 | ) |
| 648 | ): |
| 649 | try: |
| 650 | vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype) |
| 651 | time_ms = benchmark_operation( |
| 652 | vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant, |
| 653 | input_tensor, |
| 654 | residual=residual, |
| 655 | input_global_scale=scale_fp4, |
| 656 | quant_out=fp4_quant_out, |
| 657 | output_scale=fp4_output_scale, |
| 658 | ) |
| 659 | results[f"standard_allreduce_{suffix}_fp4_quant"] = time_ms |
| 660 | except Exception as e: |
| 661 | logger.error("Standard AllReduce+RMSNorm+FP4 failed: %s", e) |
| 662 | results[f"standard_allreduce_{suffix}_fp4_quant"] = float("inf") |
| 663 | |
| 664 | # Standard AllReduce + RMSNorm + FP4 Quant Native Compiled |
| 665 | with set_current_vllm_config( |
| 666 | VllmConfig(compilation_config=CompilationConfig(custom_ops=["-rms_norm"])) |
| 667 | ): |
| 668 | try: |
| 669 | vllm_fused_allreduce = VllmFusedAllreduce(hidden_dim, dtype) |
| 670 | standard_allreduce_rmsnorm_fp4_quant_native_compiled = torch.compile( |
| 671 | vllm_fused_allreduce.allreduce_rmsnorm_fp4_quant, |
| 672 | fullgraph=True, |
| 673 | dynamic=False, |
| 674 | ) |
| 675 | time_ms = benchmark_operation( |
| 676 | standard_allreduce_rmsnorm_fp4_quant_native_compiled, |
| 677 | input_tensor, |
| 678 | residual=residual, |
| 679 | quant_out=fp4_quant_out, |
| 680 | input_global_scale=scale_fp4, |
| 681 | output_scale=fp4_output_scale, |
| 682 | ) |
| 683 | results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = ( |
| 684 | time_ms |
| 685 | ) |
| 686 | except Exception as e: |
| 687 | logger.error( |
| 688 | "Standard AllReduce+RMSNorm+FP4 Native Compiled failed: %s", e |
| 689 | ) |
| 690 | results["standard_allreduce_rmsnorm_fp4_quant_native_compiled"] = float( |
| 691 | "inf" |
| 692 | ) |
| 693 | |
| 694 | # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant (trtllm only) |
| 695 | if "trtllm" in workspaces: |
| 696 | trtllm_ws = workspaces["trtllm"] |
| 697 | for use_oneshot in use_oneshot_options: |
| 698 | suffix = "_oneshot" if use_oneshot else "_twoshot" |
| 699 | key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp4_quant{suffix}" |
| 700 | try: |
| 701 | time_ms = benchmark_operation( |
| 702 | flashinfer_fused_allreduce_rmsnorm_fp4_quant, |
| 703 | input_tensor, |
| 704 | residual=residual, |
| 705 | norm_out=norm_out, |
| 706 | rms_gamma=rms_gamma, |
| 707 | rms_eps=rms_eps, |
| 708 | input_global_scale=scale_fp4, |
| 709 | allreduce_params=allreduce_params, |
| 710 | workspace=trtllm_ws, |
| 711 | quant_out=fp4_quant_out, |
| 712 | output_scale=fp4_output_scale, |
| 713 | use_oneshot=use_oneshot, |
| 714 | ) |
| 715 | results[key] = time_ms |
| 716 | except Exception as e: |
| 717 | logger.error( |
| 718 | "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP4 failed: %s", |
| 719 | e, |
| 720 | ) |
| 721 | results[key] = float("inf") |
| 722 | |
| 723 | return results |
| 724 | |
| 725 | |
| 726 | def prepare_results_with_speedups(results_dict): |
| 727 | """Prepare results with speedup calculations based on dynamic baseline selection.""" |
| 728 | prepared_results = [] |
| 729 | |
| 730 | # Determine the fastest baseline for each operation type |
| 731 | def get_fastest_baseline(op_name, results_dict): |
| 732 | """Get the fastest baseline between standard and native_compiled versions.""" |
| 733 | if "fp8_quant" in op_name: |
| 734 | candidates = [ |
| 735 | "standard_allreduce_rmsnorm_fp8_quant", |
| 736 | "standard_allreduce_rmsnorm_fp8_quant_native_compiled", |
| 737 | ] |
| 738 | elif "fp4_quant" in op_name: |
| 739 | candidates = [ |
| 740 | "standard_allreduce_rmsnorm_fp4_quant", |
| 741 | "standard_allreduce_rmsnorm_fp4_quant_native_compiled", |
| 742 | ] |
| 743 | else: |
| 744 | candidates = [ |
| 745 | "standard_allreduce_rmsnorm", |
| 746 | "standard_allreduce_rmsnorm_native_compiled", |
| 747 | ] |
| 748 | |
| 749 | # Find the fastest among available candidates |
| 750 | fastest_time = float("inf") |
| 751 | fastest_baseline = None |
| 752 | |
| 753 | for candidate in candidates: |
| 754 | if ( |
| 755 | candidate in results_dict |
| 756 | and results_dict[candidate] != float("inf") |
| 757 | and results_dict[candidate] < fastest_time |
| 758 | ): |
| 759 | fastest_time = results_dict[candidate] |
| 760 | fastest_baseline = candidate |
| 761 | |
| 762 | return fastest_baseline |
| 763 | |
| 764 | # Create dynamic baseline mapping |
| 765 | dynamic_baseline_mapping = {} |
| 766 | for op_name in results_dict: |
| 767 | if ( |
| 768 | op_name.startswith("flashinfer_") |
| 769 | or op_name.startswith("standard_") |
| 770 | and not op_name.endswith("_native_compiled") |
| 771 | ): |
| 772 | dynamic_baseline_mapping[op_name] = get_fastest_baseline( |
| 773 | op_name, results_dict |
| 774 | ) |
| 775 | |
| 776 | for op_name, time_ms in results_dict.items(): |
| 777 | if time_ms == float("inf"): |
| 778 | speedup_str = "FAILED" |
| 779 | time_str = "FAILED" |
| 780 | else: |
| 781 | time_str = f"{time_ms:.3f}" |
| 782 | # Find the appropriate baseline for this operation |
| 783 | baseline_op = dynamic_baseline_mapping.get(op_name) |
| 784 | if baseline_op and baseline_op in results_dict: |
| 785 | baseline_time = results_dict[baseline_op] |
| 786 | if baseline_time != float("inf") and baseline_time > 0: |
| 787 | speedup = baseline_time / time_ms |
| 788 | speedup_str = f"{speedup:.2f}x" |
| 789 | else: |
| 790 | speedup_str = "N/A" |
| 791 | else: |
| 792 | # For baseline operations, determine if this is the fastest baseline |
| 793 | if op_name.endswith("_native_compiled") or ( |
| 794 | op_name.startswith("standard_") |
| 795 | and not op_name.endswith("_native_compiled") |
| 796 | ): |
| 797 | fastest_baseline = get_fastest_baseline(op_name, results_dict) |
| 798 | if fastest_baseline == op_name: |
| 799 | speedup_str = "baseline" |
| 800 | else: |
| 801 | if fastest_baseline and fastest_baseline in results_dict: |
| 802 | baseline_time = results_dict[fastest_baseline] |
| 803 | if baseline_time != float("inf") and baseline_time > 0: |
| 804 | speedup = baseline_time / time_ms |
| 805 | speedup_str = f"{speedup:.2f}x" |
| 806 | else: |
| 807 | speedup_str = "N/A" |
| 808 | else: |
| 809 | speedup_str = "N/A" |
| 810 | else: |
| 811 | speedup_str = "N/A" |
| 812 | |
| 813 | prepared_results.append( |
| 814 | { |
| 815 | "operation": op_name, |
| 816 | "time_ms": time_ms, |
| 817 | "time_str": time_str, |
| 818 | "speedup_str": speedup_str, |
| 819 | } |
| 820 | ) |
| 821 | |
| 822 | return prepared_results |
| 823 | |
| 824 | |
| 825 | def print_results( |
| 826 | results_dict, |
| 827 | num_tokens, |
| 828 | hidden_dim, |
| 829 | dtype, |
| 830 | use_residual, |
| 831 | quant_modes, |
| 832 | input_size_mb, |
| 833 | ): |
| 834 | """Print benchmark results in a formatted table.""" |
| 835 | print(f"\n{'=' * 80}") |
| 836 | print( |
| 837 | f"Results: num_tokens={num_tokens}, hidden_dim={hidden_dim} " |
| 838 | f"(input size: {input_size_mb:.2f} MB)" |
| 839 | ) |
| 840 | print( |
| 841 | f"dtype={dtype}, residual={'yes' if use_residual else 'no'}, " |
| 842 | f"quant_modes={','.join(sorted(list(quant_modes)))}" |
| 843 | ) |
| 844 | print(f"{'=' * 80}") |
| 845 | print(f"{'Operation':<50} {'Time (ms)':<12} {'Speedup':<10}") |
| 846 | print(f"{'-' * 80}") |
| 847 | |
| 848 | # Prepare results with speedup calculations |
| 849 | prepared_results = prepare_results_with_speedups(results_dict) |
| 850 | |
| 851 | for result in prepared_results: |
| 852 | if result["time_ms"] == float("inf"): |
| 853 | time_display = result["time_str"] |
| 854 | else: |
| 855 | time_display = f"{result['time_ms']:.3f}" |
| 856 | |
| 857 | print( |
| 858 | f"{result['operation']:<50} {time_display:<12} {result['speedup_str']:<10}" |
| 859 | ) |
| 860 | |
| 861 | |
| 862 | def format_results_markdown( |
| 863 | all_results: list[dict], world_size: int, args: argparse.Namespace |
| 864 | ) -> str: |
| 865 | """Format all benchmark results as markdown.""" |
| 866 | lines: list[str] = [] |
| 867 | lines.append("# FlashInfer Fused Collective Operations Benchmark Results") |
| 868 | lines.append("") |
| 869 | lines.append(f"**World Size:** {world_size} ") |
| 870 | lines.append(f"**Hidden Dimension:** {args.hidden_dim} ") |
| 871 | lines.append(f"**Warmup Iterations:** {args.warmup} ") |
| 872 | lines.append(f"**Benchmark Trials:** {args.trials} ") |
| 873 | modes = ",".join(all_results[0]["quant_modes"]) if all_results else "N/A" |
| 874 | lines.append(f"**Quantization Modes:** {modes} ") |
| 875 | lines.append("") |
| 876 | lines.append("---") |
| 877 | lines.append("") |
| 878 | |
| 879 | for entry in all_results: |
| 880 | num_tokens = entry["num_tokens"] |
| 881 | dtype = entry["dtype"] |
| 882 | use_residual = entry["use_residual"] |
| 883 | results_dict = entry["results"] |
| 884 | input_size_mb = entry["input_size_mb"] |
| 885 | residual_str = "with residual" if use_residual else "no residual" |
| 886 | |
| 887 | lines.append( |
| 888 | f"## Configuration: num_tokens={num_tokens}, dtype={dtype}, {residual_str}" |
| 889 | ) |
| 890 | lines.append(f"**Input Size:** {input_size_mb:.2f} MB") |
| 891 | lines.append("") |
| 892 | |
| 893 | prepared = prepare_results_with_speedups(results_dict) |
| 894 | # Build DataFrame for markdown export |
| 895 | rows = [ |
| 896 | { |
| 897 | "Operation": r["operation"].replace("_", " ").title(), |
| 898 | "Time (ms)": r["time_str"], |
| 899 | "Speedup": r["speedup_str"], |
| 900 | } |
| 901 | for r in prepared |
| 902 | ] |
| 903 | df = pd.DataFrame(rows) |
| 904 | if df.empty: |
| 905 | lines.append("No results.") |
| 906 | else: |
| 907 | lines.append(df.to_markdown(index=False)) |
| 908 | lines.append("") |
| 909 | |
| 910 | return "\n".join(lines) |
| 911 | |
| 912 | |
| 913 | def save_results_to_file( |
| 914 | all_results: list[dict], world_size: int, args: argparse.Namespace, rank: int |
| 915 | ): |
| 916 | """Save benchmark results to markdown file (only on rank 0).""" |
| 917 | if rank != 0: |
| 918 | return |
| 919 | |
| 920 | if not all_results: |
| 921 | logger.warning("No results to save") |
| 922 | return |
| 923 | |
| 924 | output_path = args.output_file |
| 925 | |
| 926 | try: |
| 927 | markdown_content = format_results_markdown(all_results, world_size, args) |
| 928 | |
| 929 | with open(output_path, "a") as f: |
| 930 | f.write(markdown_content) |
| 931 | |
| 932 | except Exception as e: |
| 933 | logger.error("Failed to save results to file: %s", e) |
| 934 | |
| 935 | |
| 936 | def main(): |
| 937 | parser = argparse.ArgumentParser( |
| 938 | description="Benchmark fused collective operations" |
| 939 | ) |
| 940 | parser.add_argument( |
| 941 | "--num-tokens", |
| 942 | type=int, |
| 943 | nargs="+", |
| 944 | default=[128, 512, 1024, 2048], |
| 945 | help="Numbers of tokens to test", |
| 946 | ) |
| 947 | parser.add_argument( |
| 948 | "--hidden-dim", type=int, default=8192, help="Hidden dimension size" |
| 949 | ) |
| 950 | parser.add_argument( |
| 951 | "--dtypes", |
| 952 | type=str, |
| 953 | nargs="+", |
| 954 | default=["bfloat16"], |
| 955 | choices=["float16", "bfloat16", "float32"], |
| 956 | help="Data types to test", |
| 957 | ) |
| 958 | parser.add_argument( |
| 959 | "--no-residual", |
| 960 | action="store_true", |
| 961 | help="Skip residual connection tests", |
| 962 | ) |
| 963 | |
| 964 | parser.add_argument( |
| 965 | "--quant-modes", |
| 966 | type=str, |
| 967 | default="none,fp8,fp4", |
| 968 | help=( |
| 969 | "Comma-separated quantization modes to run: none, fp8, fp4. " |
| 970 | "Default: none,fp8,fp4" |
| 971 | ), |
| 972 | ) |
| 973 | |
| 974 | parser.add_argument( |
| 975 | "--warmup", type=int, default=5, help="Number of warmup iterations" |
| 976 | ) |
| 977 | parser.add_argument( |
| 978 | "--trials", type=int, default=20, help="Number of benchmark trials" |
| 979 | ) |
| 980 | parser.add_argument( |
| 981 | "--output-file", |
| 982 | type=str, |
| 983 | help="""Output file path for markdown results |
| 984 | (default: benchmark_results_<timestamp>.md) |
| 985 | """, |
| 986 | ) |
| 987 | |
| 988 | parser.add_argument( |
| 989 | "--no-oneshot", |
| 990 | action="store_true", |
| 991 | help="Skip oneshot benchmarks", |
| 992 | ) |
| 993 | |
| 994 | args = parser.parse_args() |
| 995 | |
| 996 | # Check if running with torchrun (required for collective operations) |
| 997 | if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ: |
| 998 | raise RuntimeError( |
| 999 | "Must run with torchrun for distributed benchmarking. " |
| 1000 | "Example: torchrun --nproc_per_node=2 benchmark_fused_collective.py" |
| 1001 | ) |
| 1002 | |
| 1003 | # Initialize distributed environment |
| 1004 | rank = int(os.environ["RANK"]) |
| 1005 | world_size = int(os.environ["WORLD_SIZE"]) |
| 1006 | |
| 1007 | # Use LOCAL_RANK for the device so multi-node runs (global rank >= GPUs per |
| 1008 | # node) map to a valid local GPU; falls back to global rank single-node. |
| 1009 | local_rank = int(os.environ.get("LOCAL_RANK", rank)) |
| 1010 | device = torch.device(f"cuda:{local_rank}") |
| 1011 | torch.accelerator.set_device_index(device) |
| 1012 | torch.set_default_device(device) |
| 1013 | |
| 1014 | init_distributed_environment() |
| 1015 | initialize_model_parallel(tensor_model_parallel_size=world_size) |
| 1016 | |
| 1017 | # Validate world size (must be > 1 for collective operations) |
| 1018 | if world_size <= 1: |
| 1019 | raise ValueError( |
| 1020 | "World size must be > 1 for collective operations benchmarking. " |
| 1021 | f"Current world size: {world_size}. Use torchrun with --nproc_per_node > 1." |
| 1022 | ) |
| 1023 | |
| 1024 | # Parse quantization modes |
| 1025 | valid_quant_modes = {"none", "fp8", "fp4"} |
| 1026 | raw_modes = [ |
| 1027 | m.strip().lower() for m in (args.quant_modes or "").split(",") if m.strip() |
| 1028 | ] |
| 1029 | quant_modes = set(raw_modes) if raw_modes else {"none", "fp8", "fp4"} |
| 1030 | invalid = sorted(list(quant_modes - valid_quant_modes)) |
| 1031 | if invalid: |
| 1032 | raise ValueError( |
| 1033 | f"Invalid --quant-modes entries: {','.join(invalid)}. " |
| 1034 | f"Valid options are: {','.join(sorted(valid_quant_modes))}." |
| 1035 | ) |
| 1036 | |
| 1037 | if rank == 0: |
| 1038 | logger.info("Running benchmark with world_size=%s, rank=%s", world_size, rank) |
| 1039 | logger.info("Quantization modes: %s", ",".join(sorted(list(quant_modes)))) |
| 1040 | if flashinfer_comm is not None: |
| 1041 | logger.info( |
| 1042 | "FlashInfer available - will benchmark fused operations", |
| 1043 | ) |
| 1044 | else: |
| 1045 | logger.info( |
| 1046 | "FlashInfer not available - only benchmarking standard operations" |
| 1047 | ) |
| 1048 | |
| 1049 | # Convert dtype strings to torch dtypes |
| 1050 | dtype_map = { |
| 1051 | "float16": torch.float16, |
| 1052 | "bfloat16": torch.bfloat16, |
| 1053 | "float32": torch.float32, |
| 1054 | } |
| 1055 | dtypes = [dtype_map[dt] for dt in args.dtypes] |
| 1056 | |
| 1057 | # Test configurations |
| 1058 | residual_options = [True] if not args.no_residual else [False] |
| 1059 | |
| 1060 | configs = list(itertools.product(args.num_tokens, dtypes, residual_options)) |
| 1061 | |
| 1062 | # Setup FlashInfer workspaces for all backends |
| 1063 | allreduce_params = None |
| 1064 | |
| 1065 | if flashinfer_comm is not None: |
| 1066 | # Use the largest hidden dimension for workspace setup |
| 1067 | max_element_size = max(torch.finfo(dt).bits // 8 for dt in dtypes) |
| 1068 | workspace_dtype = ( |
| 1069 | torch.float32 |
| 1070 | if max_element_size == 4 |
| 1071 | else (torch.bfloat16 if torch.bfloat16 in dtypes else torch.float16) |
| 1072 | ) |
| 1073 | max_num_token = _FI_MAX_SIZES.get(world_size) // ( |
| 1074 | args.hidden_dim * max_element_size |
| 1075 | ) |
| 1076 | |
| 1077 | for backend in FLASHINFER_BACKENDS: |
| 1078 | setup_flashinfer_workspace( |
| 1079 | backend=backend, |
| 1080 | world_size=world_size, |
| 1081 | rank=rank, |
| 1082 | hidden_dim=args.hidden_dim, |
| 1083 | max_token_num=max_num_token, |
| 1084 | dtype=workspace_dtype, |
| 1085 | ) |
| 1086 | |
| 1087 | if _FI_WORKSPACES: |
| 1088 | allreduce_params = FlashInferFusedAllReduceParams( |
| 1089 | max_token_num=max_num_token, |
| 1090 | ) |
| 1091 | |
| 1092 | # Collect all results for markdown export |
| 1093 | all_results = [] |
| 1094 | |
| 1095 | try: |
| 1096 | # Run benchmarks |
| 1097 | for num_tokens, dtype, use_residual in configs: |
| 1098 | if rank == 0: |
| 1099 | logger.info( |
| 1100 | "\nTesting: num_tokens=%s, hidden_dim=%s, dtype=%s, residual=%s", |
| 1101 | num_tokens, |
| 1102 | args.hidden_dim, |
| 1103 | dtype, |
| 1104 | use_residual, |
| 1105 | ) |
| 1106 | |
| 1107 | results = run_benchmarks( |
| 1108 | num_tokens, |
| 1109 | args.hidden_dim, |
| 1110 | dtype, |
| 1111 | use_residual, |
| 1112 | allreduce_params, |
| 1113 | workspaces=_FI_WORKSPACES, |
| 1114 | quant_modes=quant_modes, |
| 1115 | no_oneshot=args.no_oneshot, |
| 1116 | ) |
| 1117 | |
| 1118 | # Store results for markdown export |
| 1119 | if rank == 0: |
| 1120 | # Calculate input size in MB |
| 1121 | input_size_mb = ( |
| 1122 | num_tokens * args.hidden_dim * torch.finfo(dtype).bits |
| 1123 | ) / (8 * 1024 * 1024) |
| 1124 | all_results.append( |
| 1125 | { |
| 1126 | "num_tokens": num_tokens, |
| 1127 | "hidden_dim": args.hidden_dim, |
| 1128 | "dtype": str(dtype).replace("torch.", ""), |
| 1129 | "use_residual": use_residual, |
| 1130 | "quant_modes": sorted(list(quant_modes)), |
| 1131 | "input_size_mb": input_size_mb, |
| 1132 | "results": results, |
| 1133 | } |
| 1134 | ) |
| 1135 | |
| 1136 | print_results( |
| 1137 | results, |
| 1138 | num_tokens, |
| 1139 | args.hidden_dim, |
| 1140 | dtype, |
| 1141 | use_residual, |
| 1142 | quant_modes, |
| 1143 | input_size_mb, |
| 1144 | ) |
| 1145 | |
| 1146 | # Save results to markdown file |
| 1147 | if args.output_file and rank == 0: |
| 1148 | save_results_to_file(all_results, world_size, args, rank) |
| 1149 | |
| 1150 | finally: |
| 1151 | # Cleanup |
| 1152 | cleanup_flashinfer_workspaces() |
| 1153 | |
| 1154 | dist.barrier() |
| 1155 | |
| 1156 | |
| 1157 | if __name__ == "__main__": |
| 1158 | from vllm.config import VllmConfig, set_current_vllm_config |
| 1159 | |
| 1160 | with set_current_vllm_config(VllmConfig()): |
| 1161 | main() |