| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import gc |
| 6 | import json |
| 7 | import os |
| 8 | import time |
| 9 | from contextlib import nullcontext |
| 10 | from datetime import datetime |
| 11 | from itertools import product |
| 12 | from typing import Any, TypedDict |
| 13 | |
| 14 | import ray |
| 15 | import torch |
| 16 | from ray.experimental.tqdm_ray import tqdm |
| 17 | |
| 18 | from vllm.model_executor.layers.fused_moe import fused_topk |
| 19 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 20 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 21 | maybe_make_prepare_finalize, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.fused_moe.config import ( |
| 24 | FusedMoEConfig, |
| 25 | FusedMoEParallelConfig, |
| 26 | FusedMoEQuantConfig, |
| 27 | RoutingMethodType, |
| 28 | _get_config_dtype_str, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( |
| 31 | TritonOrDeepGemmExperts, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.fused_moe.fused_moe import * |
| 34 | from vllm.transformers_utils.config import get_config |
| 35 | from vllm.triton_utils import triton |
| 36 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 37 | from vllm.utils.torch_utils import set_random_seed |
| 38 | |
| 39 | FP8_DTYPE = current_platform.fp8_dtype() |
| 40 | |
| 41 | # Default interval for clearing Triton JIT cache during tuning |
| 42 | # Set to 0 to disable automatic cache clearing |
| 43 | _CACHE_CLEAR_INTERVAL_ENV = "VLLM_MOE_TUNE_CACHE_CLEAR_INTERVAL" |
| 44 | TRITON_CACHE_CLEAR_INTERVAL = int(os.environ.get(_CACHE_CLEAR_INTERVAL_ENV, "50")) |
| 45 | |
| 46 | |
| 47 | def clear_triton_cache(): |
| 48 | """Clear Triton JIT compilation cache and Python/CUDA memory. |
| 49 | |
| 50 | This helps prevent OOM during tuning with large models (many experts). |
| 51 | """ |
| 52 | # Force Python garbage collection |
| 53 | gc.collect() |
| 54 | |
| 55 | # Clear CUDA memory cache |
| 56 | if torch.cuda.is_available(): |
| 57 | torch.accelerator.empty_cache() |
| 58 | |
| 59 | # Try to clear Triton's runtime cache |
| 60 | try: |
| 61 | if ( |
| 62 | hasattr(triton, "runtime") |
| 63 | and hasattr(triton.runtime, "cache") |
| 64 | and hasattr(triton.runtime.cache, "clear") |
| 65 | ): |
| 66 | triton.runtime.cache.clear() |
| 67 | except ImportError: |
| 68 | # Triton not installed, skip cache clearing |
| 69 | pass |
| 70 | except AttributeError: |
| 71 | # Triton version doesn't have expected cache API |
| 72 | pass |
| 73 | except Exception as e: |
| 74 | print(f"Warning: Failed to clear Triton cache: {e}") |
| 75 | |
| 76 | # Additional garbage collection after clearing caches |
| 77 | gc.collect() |
| 78 | |
| 79 | |
| 80 | def ensure_divisibility(numerator, denominator, text): |
| 81 | """Ensure that numerator is divisible by the denominator.""" |
| 82 | assert numerator % denominator == 0, "{} {} is not divisible by tp {}.".format( |
| 83 | text, numerator, denominator |
| 84 | ) |
| 85 | |
| 86 | |
| 87 | class BenchmarkConfig(TypedDict): |
| 88 | BLOCK_SIZE_M: int |
| 89 | BLOCK_SIZE_N: int |
| 90 | BLOCK_SIZE_K: int |
| 91 | GROUP_SIZE_M: int |
| 92 | num_warps: int |
| 93 | num_stages: int |
| 94 | |
| 95 | |
| 96 | def benchmark_config( |
| 97 | config: BenchmarkConfig, |
| 98 | num_tokens: int, |
| 99 | num_experts: int, |
| 100 | shard_intermediate_size: int, |
| 101 | hidden_size: int, |
| 102 | topk: int, |
| 103 | dtype: torch.dtype, |
| 104 | use_fp8_w8a8: bool, |
| 105 | use_int8_w8a16: bool, |
| 106 | use_int4_w4a16: bool = False, |
| 107 | num_iters: int = 100, |
| 108 | block_quant_shape: list[int] = None, |
| 109 | use_deep_gemm: bool = False, |
| 110 | ) -> float: |
| 111 | init_dtype = torch.float16 if use_fp8_w8a8 else dtype |
| 112 | x = torch.randn(num_tokens, hidden_size, dtype=dtype) |
| 113 | if use_int4_w4a16: |
| 114 | # Int4 packed weights: 2 int4 values per uint8 byte |
| 115 | # K dimension is packed (halved) |
| 116 | intermediate_size = shard_intermediate_size // 2 # after silu_and_mul |
| 117 | w1 = torch.randint( |
| 118 | 0, |
| 119 | 255, |
| 120 | ( |
| 121 | num_experts, |
| 122 | shard_intermediate_size, |
| 123 | hidden_size // 2, # int4 packing |
| 124 | ), |
| 125 | dtype=torch.uint8, |
| 126 | ) |
| 127 | w2 = torch.randint( |
| 128 | 0, |
| 129 | 255, |
| 130 | ( |
| 131 | num_experts, |
| 132 | hidden_size, |
| 133 | intermediate_size // 2, # int4 packing |
| 134 | ), |
| 135 | dtype=torch.uint8, |
| 136 | ) |
| 137 | elif use_int8_w8a16: |
| 138 | w1 = torch.randint( |
| 139 | -127, |
| 140 | 127, |
| 141 | ( |
| 142 | num_experts, |
| 143 | shard_intermediate_size, |
| 144 | hidden_size, |
| 145 | ), |
| 146 | dtype=torch.int8, |
| 147 | ) |
| 148 | w2 = torch.randint( |
| 149 | -127, |
| 150 | 127, |
| 151 | ( |
| 152 | num_experts, |
| 153 | hidden_size, |
| 154 | shard_intermediate_size // 2, |
| 155 | ), |
| 156 | dtype=torch.int8, |
| 157 | ) |
| 158 | else: |
| 159 | w1 = torch.randn( |
| 160 | num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype |
| 161 | ) |
| 162 | w2 = torch.randn( |
| 163 | num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype |
| 164 | ) |
| 165 | gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32) |
| 166 | |
| 167 | w1_scale = None |
| 168 | w2_scale = None |
| 169 | a1_scale = None |
| 170 | a2_scale = None |
| 171 | if use_int4_w4a16: |
| 172 | if block_quant_shape is None: |
| 173 | raise ValueError("block_quant_shape is required for int4_w4a16") |
| 174 | group_size = block_quant_shape[1] |
| 175 | # Scales shape: (E, N, K // group_size) in fp16 |
| 176 | w1_scale = torch.rand( |
| 177 | (num_experts, shard_intermediate_size, hidden_size // group_size), |
| 178 | dtype=dtype, |
| 179 | ) |
| 180 | w2_scale = torch.rand( |
| 181 | (num_experts, hidden_size, intermediate_size // group_size), |
| 182 | dtype=dtype, |
| 183 | ) |
| 184 | elif use_int8_w8a16: |
| 185 | w1_scale = torch.randn( |
| 186 | (num_experts, 2 * shard_intermediate_size), dtype=torch.float32 |
| 187 | ) |
| 188 | w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32) |
| 189 | if use_deep_gemm: |
| 190 | # we use the default block shape for deepgemm |
| 191 | block_quant_shape = [128, 128] |
| 192 | if use_fp8_w8a8: |
| 193 | if block_quant_shape: |
| 194 | block_n, block_k = block_quant_shape[0], block_quant_shape[1] |
| 195 | E = num_experts |
| 196 | N = shard_intermediate_size // 2 |
| 197 | K = hidden_size |
| 198 | factor_for_scale = 1e-2 |
| 199 | n_tiles_w1 = (2 * N + block_n - 1) // block_n |
| 200 | n_tiles_w2 = (K + block_n - 1) // block_n |
| 201 | k_tiles_w1 = (K + block_k - 1) // block_k |
| 202 | k_tiles_w2 = (N + block_k - 1) // block_k |
| 203 | w1_scale = ( |
| 204 | torch.rand((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32) |
| 205 | * factor_for_scale |
| 206 | ) |
| 207 | w2_scale = ( |
| 208 | torch.rand((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32) |
| 209 | * factor_for_scale |
| 210 | ) |
| 211 | else: |
| 212 | w1_scale = torch.randn(num_experts, dtype=torch.float32) |
| 213 | w2_scale = torch.randn(num_experts, dtype=torch.float32) |
| 214 | |
| 215 | a1_scale = torch.randn(1, dtype=torch.float32) |
| 216 | a2_scale = torch.randn(1, dtype=torch.float32) |
| 217 | |
| 218 | w1 = w1.to(FP8_DTYPE) |
| 219 | w2 = w2.to(FP8_DTYPE) |
| 220 | |
| 221 | input_gating = torch.empty(num_tokens, num_experts, dtype=torch.float32) |
| 222 | |
| 223 | def prepare(i: int): |
| 224 | input_gating.copy_(gating_output[i]) |
| 225 | |
| 226 | def run(): |
| 227 | from vllm.model_executor.layers.fused_moe import override_config |
| 228 | |
| 229 | if use_fp8_w8a8: |
| 230 | quant_dtype = torch.float8_e4m3fn |
| 231 | elif use_int8_w8a16: |
| 232 | quant_dtype = torch.int8 |
| 233 | else: |
| 234 | quant_dtype = None |
| 235 | |
| 236 | quant_config = FusedMoEQuantConfig.make( |
| 237 | quant_dtype=quant_dtype, |
| 238 | w1_scale=w1_scale, |
| 239 | w2_scale=w2_scale, |
| 240 | a1_scale=a1_scale, |
| 241 | a2_scale=a2_scale, |
| 242 | block_shape=block_quant_shape, |
| 243 | weight_dtype="int4" if use_int4_w4a16 else None, |
| 244 | ) |
| 245 | |
| 246 | deep_gemm_experts = None |
| 247 | if use_deep_gemm: |
| 248 | moe_config = ( |
| 249 | FusedMoEConfig( |
| 250 | num_experts=num_experts, |
| 251 | experts_per_token=topk, |
| 252 | hidden_dim=hidden_size, |
| 253 | intermediate_size=shard_intermediate_size, |
| 254 | num_local_experts=num_experts, |
| 255 | num_logical_experts=num_experts, |
| 256 | activation=MoEActivation.SILU, |
| 257 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 258 | in_dtype=init_dtype, |
| 259 | routing_method=RoutingMethodType.TopK, |
| 260 | device="cuda", |
| 261 | ), |
| 262 | ) |
| 263 | deep_gemm_experts = mk.FusedMoEKernel( |
| 264 | prepare_finalize=maybe_make_prepare_finalize( |
| 265 | moe=moe_config, |
| 266 | quant_config=quant_config, |
| 267 | allow_new_interface=True, |
| 268 | use_monolithic=False, |
| 269 | ), |
| 270 | fused_experts=TritonOrDeepGemmExperts( |
| 271 | moe_config=moe_config, |
| 272 | quant_config=quant_config, |
| 273 | ), |
| 274 | ) |
| 275 | |
| 276 | with override_config(config): |
| 277 | topk_weights, topk_ids, token_expert_indices = fused_topk( |
| 278 | x, input_gating, topk, renormalize=not use_deep_gemm |
| 279 | ) |
| 280 | |
| 281 | if use_deep_gemm: |
| 282 | return deep_gemm_experts.apply( |
| 283 | x, |
| 284 | w1, |
| 285 | w2, |
| 286 | topk_weights, |
| 287 | topk_ids, |
| 288 | activation=MoEActivation.SILU, |
| 289 | global_num_experts=num_experts, |
| 290 | apply_router_weight_on_input=False, |
| 291 | expert_map=False, |
| 292 | ) |
| 293 | return fused_experts( |
| 294 | x, |
| 295 | w1, |
| 296 | w2, |
| 297 | topk_weights, |
| 298 | topk_ids, |
| 299 | quant_config=quant_config, |
| 300 | ) |
| 301 | |
| 302 | # JIT compilation & warmup |
| 303 | run() |
| 304 | torch.accelerator.synchronize() |
| 305 | |
| 306 | # Capture 10 invocations with CUDA graph |
| 307 | graph = torch.cuda.CUDAGraph() |
| 308 | with torch.cuda.graph(graph): |
| 309 | for _ in range(10): |
| 310 | run() |
| 311 | torch.accelerator.synchronize() |
| 312 | |
| 313 | # Warmup |
| 314 | for _ in range(5): |
| 315 | graph.replay() |
| 316 | torch.accelerator.synchronize() |
| 317 | |
| 318 | start_event = torch.Event(enable_timing=True) |
| 319 | end_event = torch.Event(enable_timing=True) |
| 320 | |
| 321 | latencies: list[float] = [] |
| 322 | for i in range(num_iters): |
| 323 | prepare(i) |
| 324 | torch.accelerator.synchronize() |
| 325 | |
| 326 | start_event.record() |
| 327 | graph.replay() |
| 328 | end_event.record() |
| 329 | end_event.synchronize() |
| 330 | latencies.append(start_event.elapsed_time(end_event)) |
| 331 | avg = sum(latencies) / (num_iters * 10) * 1000 # us |
| 332 | graph.reset() |
| 333 | return avg |
| 334 | |
| 335 | |
| 336 | def get_rocm_tuning_space(use_fp16): |
| 337 | block_mn_range = [16, 32, 64, 128, 256] |
| 338 | block_k_range = [16, 32, 64, 128, 256] |
| 339 | if not use_fp16: |
| 340 | block_k_range.remove(16) # BLOCK_K=16 not supported for fp8 |
| 341 | num_warps_range = [1, 2, 4, 8] |
| 342 | group_m_range = [1, 4, 8, 16, 32] |
| 343 | num_stage_range = [2] |
| 344 | waves_per_eu_range = [0, 1, 2, 4] |
| 345 | matrix_instr_nonkdim_range = [16, 32] if use_fp16 else [] |
| 346 | kpack_range = [1, 2] if use_fp16 else [] |
| 347 | |
| 348 | param_ranges = { |
| 349 | "BLOCK_SIZE_M": block_mn_range, |
| 350 | "BLOCK_SIZE_N": block_mn_range, |
| 351 | "BLOCK_SIZE_K": block_k_range, |
| 352 | "GROUP_SIZE_M": group_m_range, |
| 353 | "num_warps": num_warps_range, |
| 354 | "num_stages": num_stage_range, |
| 355 | "waves_per_eu": waves_per_eu_range, |
| 356 | } |
| 357 | if use_fp16: |
| 358 | param_ranges["matrix_instr_nonkdim"] = matrix_instr_nonkdim_range |
| 359 | param_ranges["kpack"] = kpack_range |
| 360 | |
| 361 | return param_ranges |
| 362 | |
| 363 | |
| 364 | def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int]]: |
| 365 | configs: list[BenchmarkConfig] = [] |
| 366 | |
| 367 | if current_platform.is_rocm(): |
| 368 | param_ranges = get_rocm_tuning_space(use_fp16) |
| 369 | else: |
| 370 | # Reduced search space for faster tuning. |
| 371 | # TODO(woosuk): Increase the search space and use a performance model to |
| 372 | # prune the search space. |
| 373 | block_m_range = [16, 32, 64, 128, 256] |
| 374 | block_n_range = [32, 64, 128, 256] |
| 375 | block_k_range = [64, 128, 256] |
| 376 | num_warps_range = [4, 8] |
| 377 | group_m_range = [1, 16, 32, 64] |
| 378 | num_stage_range = [2, 3, 4, 5] |
| 379 | |
| 380 | param_ranges = { |
| 381 | "BLOCK_SIZE_M": block_m_range, |
| 382 | "BLOCK_SIZE_N": block_n_range, |
| 383 | "BLOCK_SIZE_K": block_k_range, |
| 384 | "GROUP_SIZE_M": group_m_range, |
| 385 | "num_warps": num_warps_range, |
| 386 | "num_stages": num_stage_range, |
| 387 | } |
| 388 | |
| 389 | keys, values = zip(*param_ranges.items()) |
| 390 | for config_values in product(*values): |
| 391 | config = dict(zip(keys, config_values)) |
| 392 | configs.append(config) |
| 393 | |
| 394 | # Drop configs incompatible with fp8 block quantization. A tile must align |
| 395 | # to the quant-block scale grid, i.e. tile and block must divide one |
| 396 | # another. The kernel indexes scales per element (offs_bn // group_n, |
| 397 | # k_start // group_k), so a tile narrower than the block (e.g. N=64 with |
| 398 | # block_n=128) is valid -- and often faster at small batch. An exact |
| 399 | # multiple was required before, which dropped those smaller tiles entirely. |
| 400 | if block_quant_shape is not None and not use_fp16: |
| 401 | block_n, block_k = block_quant_shape[0], block_quant_shape[1] |
| 402 | for config in configs[:]: |
| 403 | bn, bk = config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"] |
| 404 | n_aligned = bn % block_n == 0 or block_n % bn == 0 |
| 405 | k_aligned = bk % block_k == 0 or block_k % bk == 0 |
| 406 | if not (n_aligned and k_aligned): |
| 407 | configs.remove(config) |
| 408 | return configs |
| 409 | |
| 410 | |
| 411 | def prune_rocm_search_space( |
| 412 | num_tokens, shard_intermediate_size, hidden_size, search_space, is_fp16, topk |
| 413 | ): |
| 414 | N1, K1 = shard_intermediate_size, hidden_size |
| 415 | N2, K2 = hidden_size, shard_intermediate_size // 2 |
| 416 | pruned_space_1 = prune_rocm_configs( |
| 417 | num_tokens * topk, N1, K1, search_space, is_fp16 |
| 418 | ) |
| 419 | pruned_space_2 = prune_rocm_configs( |
| 420 | num_tokens * topk, N2, K2, search_space, is_fp16 |
| 421 | ) |
| 422 | search_space = merge_unique_dicts(pruned_space_1, pruned_space_2) |
| 423 | return search_space |
| 424 | |
| 425 | |
| 426 | # The following code is inspired by ROCm/Triton GEMM tuning script: |
| 427 | # https://github.com/ROCm/triton/blob/triton-mlir/scripts/amd/gemm/tune_gemm.py#L89 |
| 428 | def prune_rocm_configs(M, N, K, configs, is_fp16=True): |
| 429 | pruned_configs = [] |
| 430 | elemBytes_a = 2 if is_fp16 else 1 |
| 431 | elemBytes_b = 2 if is_fp16 else 1 |
| 432 | |
| 433 | mfma = 16 if M < 32 or N < 32 else 32 |
| 434 | |
| 435 | # TODO (zhanglx): figure out the boundary between large and small gemms |
| 436 | large_gemm = False |
| 437 | if M >= 2048 and N >= 2048: |
| 438 | large_gemm = True |
| 439 | |
| 440 | for config in configs: |
| 441 | BLOCK_SIZE_M = config.get("BLOCK_SIZE_M") |
| 442 | BLOCK_SIZE_N = config.get("BLOCK_SIZE_N") |
| 443 | BLOCK_SIZE_K = config.get("BLOCK_SIZE_K") |
| 444 | num_warps = config.get("num_warps") |
| 445 | |
| 446 | if is_fp16: |
| 447 | matrix_instr_nonkdim = config.get("matrix_instr_nonkdim") |
| 448 | if matrix_instr_nonkdim > mfma: |
| 449 | continue |
| 450 | if mfma == 4 and BLOCK_SIZE_K < 64: |
| 451 | continue |
| 452 | # some layouts could not work properly in case |
| 453 | # number elements per thread is less 1 |
| 454 | if BLOCK_SIZE_M * BLOCK_SIZE_N < 64: |
| 455 | continue |
| 456 | SPLIT_K = config.get("SPLIT_K", 1) |
| 457 | GROUP_M = config.get("GROUP_SIZE_M") |
| 458 | if is_fp16: |
| 459 | if ( |
| 460 | matrix_instr_nonkdim > BLOCK_SIZE_M |
| 461 | or matrix_instr_nonkdim > BLOCK_SIZE_N |
| 462 | ): |
| 463 | continue |
| 464 | if matrix_instr_nonkdim >= M and matrix_instr_nonkdim != BLOCK_SIZE_M: |
| 465 | continue |
| 466 | if matrix_instr_nonkdim >= N and matrix_instr_nonkdim != BLOCK_SIZE_N: |
| 467 | continue |
| 468 | # Skip BLOCK_SIZE that is too large compare to M/N |
| 469 | # unless BLOCK_SIZE is already small enough |
| 470 | if M * 2 < BLOCK_SIZE_M and BLOCK_SIZE_M != 16: |
| 471 | continue |
| 472 | if N * 2 < BLOCK_SIZE_N and BLOCK_SIZE_N != 16: |
| 473 | continue |
| 474 | # skip large split_k when not necessary |
| 475 | if SPLIT_K != 1 and not need_split_k(M, N, K): |
| 476 | continue |
| 477 | # skip split_k that leads to EVEN_K = false |
| 478 | leap = SPLIT_K * BLOCK_SIZE_K |
| 479 | modv = K % leap |
| 480 | if modv != 0: |
| 481 | continue |
| 482 | # skip large GROUP_M |
| 483 | if GROUP_M * BLOCK_SIZE_M > M and GROUP_M != 1: |
| 484 | continue |
| 485 | # out of shared memory resource |
| 486 | # TODO (zhanglx): This does not consider the LDS usage in the epilogue |
| 487 | LDS = ( |
| 488 | BLOCK_SIZE_K * BLOCK_SIZE_M * elemBytes_a |
| 489 | + BLOCK_SIZE_K * BLOCK_SIZE_N * elemBytes_b |
| 490 | ) |
| 491 | if LDS > 65536: |
| 492 | continue |
| 493 | # Skip small block sizes and num_warps for large gemm |
| 494 | # For fp16 and f8, we want to only use BLOCK_SIZE >= 64 |
| 495 | if large_gemm: |
| 496 | if BLOCK_SIZE_M < 64 or BLOCK_SIZE_N < 64: |
| 497 | continue |
| 498 | if BLOCK_SIZE_K < 64: |
| 499 | continue |
| 500 | if num_warps < 4: |
| 501 | continue |
| 502 | |
| 503 | pruned_configs.append(config) |
| 504 | |
| 505 | return pruned_configs |
| 506 | |
| 507 | |
| 508 | def need_split_k(SIZE_M, SIZE_N, SIZE_K): |
| 509 | return (SIZE_M < 64 or SIZE_N < 64) and SIZE_K > 1024 |
| 510 | |
| 511 | |
| 512 | def merge_unique_dicts(list1, list2): |
| 513 | result = [] |
| 514 | combined_list = list1.copy() |
| 515 | combined_list.extend(list2) |
| 516 | for dictionary in combined_list: |
| 517 | if dictionary not in result: |
| 518 | result.append(dictionary) |
| 519 | return result |
| 520 | |
| 521 | |
| 522 | @ray.remote(num_gpus=1) |
| 523 | class BenchmarkWorker: |
| 524 | def __init__(self, seed: int) -> None: |
| 525 | torch.set_default_device("cuda") |
| 526 | set_random_seed(seed) |
| 527 | self.seed = seed |
| 528 | # Get the device ID to allocate tensors and kernels |
| 529 | # on the respective GPU. This is required for Ray to work |
| 530 | # correctly with multi-GPU tuning on the ROCm platform. |
| 531 | self.device_id = int(ray.get_gpu_ids()[0]) |
| 532 | |
| 533 | def benchmark( |
| 534 | self, |
| 535 | num_tokens: int, |
| 536 | num_experts: int, |
| 537 | shard_intermediate_size: int, |
| 538 | hidden_size: int, |
| 539 | topk: int, |
| 540 | dtype: torch.dtype, |
| 541 | use_fp8_w8a8: bool, |
| 542 | use_int8_w8a16: bool, |
| 543 | use_int4_w4a16: bool = False, |
| 544 | block_quant_shape: list[int] = None, |
| 545 | use_deep_gemm: bool = False, |
| 546 | ) -> tuple[dict[str, int], float]: |
| 547 | # local import to allow serialization by ray |
| 548 | |
| 549 | set_random_seed(self.seed) |
| 550 | dtype_str = _get_config_dtype_str( |
| 551 | dtype, |
| 552 | use_int8_w8a16=use_int8_w8a16, |
| 553 | use_fp8_w8a8=use_fp8_w8a8, |
| 554 | use_int4_w4a16=use_int4_w4a16, |
| 555 | ) |
| 556 | # NOTE(woosuk): The current naming convention uses w2.shape[2], which |
| 557 | # is the intermediate size after silu_and_mul. |
| 558 | block_n = block_quant_shape[0] if block_quant_shape else None |
| 559 | block_k = block_quant_shape[1] if block_quant_shape else None |
| 560 | op_config = get_moe_configs( |
| 561 | num_experts, shard_intermediate_size // 2, dtype_str, block_n, block_k |
| 562 | ) |
| 563 | if op_config is None: |
| 564 | config = get_default_config( |
| 565 | num_tokens, |
| 566 | num_experts, |
| 567 | shard_intermediate_size, |
| 568 | hidden_size, |
| 569 | topk, |
| 570 | dtype_str, |
| 571 | block_quant_shape, |
| 572 | ) |
| 573 | else: |
| 574 | config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))] |
| 575 | kernel_time = benchmark_config( |
| 576 | config, |
| 577 | num_tokens, |
| 578 | num_experts, |
| 579 | shard_intermediate_size, |
| 580 | hidden_size, |
| 581 | topk, |
| 582 | dtype, |
| 583 | use_fp8_w8a8, |
| 584 | use_int8_w8a16, |
| 585 | use_int4_w4a16=use_int4_w4a16, |
| 586 | num_iters=100, |
| 587 | block_quant_shape=block_quant_shape, |
| 588 | use_deep_gemm=use_deep_gemm, |
| 589 | ) |
| 590 | return config, kernel_time |
| 591 | |
| 592 | def tune( |
| 593 | self, |
| 594 | num_tokens: int, |
| 595 | num_experts: int, |
| 596 | shard_intermediate_size: int, |
| 597 | hidden_size: int, |
| 598 | topk: int, |
| 599 | dtype: torch.dtype, |
| 600 | use_fp8_w8a8: bool, |
| 601 | use_int8_w8a16: bool, |
| 602 | use_int4_w4a16: bool, |
| 603 | search_space: list[dict[str, int]], |
| 604 | block_quant_shape: list[int], |
| 605 | use_deep_gemm: bool, |
| 606 | ) -> dict[str, int]: |
| 607 | # local import to allow serialization by ray |
| 608 | from vllm.platforms import current_platform |
| 609 | |
| 610 | best_config = None |
| 611 | best_time = float("inf") |
| 612 | if current_platform.is_rocm(): |
| 613 | is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16) |
| 614 | search_space = prune_rocm_search_space( |
| 615 | num_tokens, |
| 616 | shard_intermediate_size, |
| 617 | hidden_size, |
| 618 | search_space, |
| 619 | is_fp16, |
| 620 | topk, |
| 621 | ) |
| 622 | |
| 623 | need_device_guard = False |
| 624 | if current_platform.is_rocm(): |
| 625 | visible_device = os.environ.get("ROCR_VISIBLE_DEVICES", None) |
| 626 | if visible_device != f"{self.device_id}": |
| 627 | need_device_guard = True |
| 628 | |
| 629 | with ( |
| 630 | # Ray restricts each worker to one GPU; use local index 0 |
| 631 | torch.accelerator.device_index(0) if need_device_guard else nullcontext() |
| 632 | ): |
| 633 | for idx, config in enumerate(tqdm(search_space)): |
| 634 | try: |
| 635 | kernel_time = benchmark_config( |
| 636 | config, |
| 637 | num_tokens, |
| 638 | num_experts, |
| 639 | shard_intermediate_size, |
| 640 | hidden_size, |
| 641 | topk, |
| 642 | dtype, |
| 643 | use_fp8_w8a8, |
| 644 | use_int8_w8a16, |
| 645 | use_int4_w4a16, |
| 646 | num_iters=20, |
| 647 | block_quant_shape=block_quant_shape, |
| 648 | use_deep_gemm=use_deep_gemm, |
| 649 | ) |
| 650 | except triton.runtime.autotuner.OutOfResources: |
| 651 | # Some configurations may be invalid and fail to compile. |
| 652 | continue |
| 653 | |
| 654 | if kernel_time < best_time: |
| 655 | best_time = kernel_time |
| 656 | best_config = config |
| 657 | |
| 658 | # Periodically clear Triton JIT cache to prevent OOM |
| 659 | # This is especially important for large models with many experts |
| 660 | if ( |
| 661 | TRITON_CACHE_CLEAR_INTERVAL > 0 |
| 662 | and idx > 0 |
| 663 | and idx % TRITON_CACHE_CLEAR_INTERVAL == 0 |
| 664 | ): |
| 665 | clear_triton_cache() |
| 666 | |
| 667 | # Final cleanup after tuning completes |
| 668 | clear_triton_cache() |
| 669 | |
| 670 | now = datetime.now() |
| 671 | print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}") |
| 672 | assert best_config is not None |
| 673 | return best_config |
| 674 | |
| 675 | |
| 676 | def sort_config(config: BenchmarkConfig) -> BenchmarkConfig: |
| 677 | return { |
| 678 | "BLOCK_SIZE_M": config["BLOCK_SIZE_M"], |
| 679 | "BLOCK_SIZE_N": config["BLOCK_SIZE_N"], |
| 680 | "BLOCK_SIZE_K": config["BLOCK_SIZE_K"], |
| 681 | "GROUP_SIZE_M": config["GROUP_SIZE_M"], |
| 682 | "num_warps": config["num_warps"], |
| 683 | "num_stages": config["num_stages"], |
| 684 | **( |
| 685 | {"waves_per_eu": config["waves_per_eu"]} if "waves_per_eu" in config else {} |
| 686 | ), |
| 687 | **( |
| 688 | {"matrix_instr_nonkdim": config["matrix_instr_nonkdim"]} |
| 689 | if "matrix_instr_nonkdim" in config |
| 690 | else {} |
| 691 | ), |
| 692 | **({"kpack": config["kpack"]} if "kpack" in config else {}), |
| 693 | **({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}), |
| 694 | } |
| 695 | |
| 696 | |
| 697 | def save_configs( |
| 698 | configs: dict[int, BenchmarkConfig], |
| 699 | num_experts: int, |
| 700 | shard_intermediate_size: int, |
| 701 | hidden_size: int, |
| 702 | topk: int, |
| 703 | dtype: torch.dtype, |
| 704 | use_fp8_w8a8: bool, |
| 705 | use_int8_w8a16: bool, |
| 706 | use_int4_w4a16: bool, |
| 707 | block_quant_shape: list[int], |
| 708 | save_dir: str, |
| 709 | ) -> None: |
| 710 | dtype_str = _get_config_dtype_str( |
| 711 | dtype, |
| 712 | use_int8_w8a16=use_int8_w8a16, |
| 713 | use_fp8_w8a8=use_fp8_w8a8, |
| 714 | use_int4_w4a16=use_int4_w4a16, |
| 715 | ) |
| 716 | |
| 717 | # NOTE(woosuk): The current naming convention uses w2.shape[2], which |
| 718 | # is the intermediate size after silu_and_mul. |
| 719 | filename = get_config_file_name( |
| 720 | num_experts, shard_intermediate_size // 2, dtype_str, block_quant_shape |
| 721 | ) |
| 722 | os.makedirs(save_dir, exist_ok=True) |
| 723 | filename = os.path.join(save_dir, filename) |
| 724 | print(f"Writing best config to {filename}...") |
| 725 | with open(filename, "w") as f: |
| 726 | json.dump({"triton_version": triton.__version__, **configs}, f, indent=4) |
| 727 | f.write("\n") |
| 728 | |
| 729 | |
| 730 | def get_compressed_tensors_block_structure(config, default_value=None): |
| 731 | config_groups = config.get("config_groups", {}) |
| 732 | if len(config_groups) != 1: |
| 733 | return default_value |
| 734 | group = next(iter(config_groups.values())) |
| 735 | weights = group.get("weights", {}) |
| 736 | block_structure = weights.get("block_structure", default_value) |
| 737 | return block_structure |
| 738 | |
| 739 | |
| 740 | def get_weight_block_size_safety(config, default_value=None): |
| 741 | quantization_config = getattr(config, "quantization_config", {}) |
| 742 | if isinstance(quantization_config, dict): |
| 743 | if "weight_block_size" in quantization_config: |
| 744 | return quantization_config["weight_block_size"] |
| 745 | return get_compressed_tensors_block_structure( |
| 746 | quantization_config, default_value |
| 747 | ) |
| 748 | return default_value |
| 749 | |
| 750 | |
| 751 | def get_model_params(config): |
| 752 | architectures = getattr(config, "architectures", None) or [type(config).__name__] |
| 753 | architecture = architectures[0] |
| 754 | |
| 755 | if architecture == "DbrxForCausalLM": |
| 756 | E = config.ffn_config.moe_num_experts |
| 757 | topk = config.ffn_config.moe_top_k |
| 758 | intermediate_size = config.ffn_config.ffn_hidden_size |
| 759 | hidden_size = config.hidden_size |
| 760 | elif architecture == "JambaForCausalLM": |
| 761 | E = config.num_experts |
| 762 | topk = config.num_experts_per_tok |
| 763 | intermediate_size = config.intermediate_size |
| 764 | hidden_size = config.hidden_size |
| 765 | elif architecture in ( |
| 766 | "DeepseekV2ForCausalLM", |
| 767 | "DeepseekV3ForCausalLM", |
| 768 | "DeepseekV32ForCausalLM", |
| 769 | "DeepseekV4ForCausalLM", |
| 770 | "GlmMoeDsaForCausalLM", |
| 771 | "Glm4MoeForCausalLM", |
| 772 | "Glm4MoeLiteForCausalLM", |
| 773 | "NemotronHForCausalLM", |
| 774 | "MistralLarge3ForCausalLM", |
| 775 | ): |
| 776 | E = config.n_routed_experts |
| 777 | topk = config.num_experts_per_tok |
| 778 | intermediate_size = config.moe_intermediate_size |
| 779 | hidden_size = config.hidden_size |
| 780 | elif architecture in ( |
| 781 | "BailingMoeV3ForCausalLM", |
| 782 | "Qwen2MoeForCausalLM", |
| 783 | "Qwen3MoeForCausalLM", |
| 784 | "Qwen3NextForCausalLM", |
| 785 | ): |
| 786 | E = config.num_experts |
| 787 | topk = config.num_experts_per_tok |
| 788 | intermediate_size = config.moe_intermediate_size |
| 789 | hidden_size = config.hidden_size |
| 790 | elif architecture in ( |
| 791 | "Qwen3VLMoeForConditionalGeneration", |
| 792 | "Qwen3_5MoeForConditionalGeneration", |
| 793 | "Qwen3_5MoeTextConfig", |
| 794 | ): |
| 795 | text_config = config.get_text_config() |
| 796 | E = text_config.num_experts |
| 797 | topk = text_config.num_experts_per_tok |
| 798 | intermediate_size = text_config.moe_intermediate_size |
| 799 | hidden_size = text_config.hidden_size |
| 800 | elif architecture == "DiffusionGemmaForBlockDiffusion": |
| 801 | text_config = config.get_text_config() |
| 802 | E = text_config.num_experts |
| 803 | topk = text_config.top_k_experts |
| 804 | intermediate_size = text_config.moe_intermediate_size |
| 805 | hidden_size = text_config.hidden_size |
| 806 | elif architecture == "HunYuanMoEV1ForCausalLM": |
| 807 | E = config.num_experts |
| 808 | topk = config.moe_topk[0] |
| 809 | intermediate_size = config.moe_intermediate_size[0] |
| 810 | hidden_size = config.hidden_size |
| 811 | elif architecture == "Qwen3OmniMoeForConditionalGeneration": |
| 812 | E = config.thinker_config.text_config.num_experts |
| 813 | topk = config.thinker_config.text_config.num_experts_per_tok |
| 814 | intermediate_size = config.thinker_config.text_config.moe_intermediate_size |
| 815 | hidden_size = config.thinker_config.text_config.hidden_size |
| 816 | elif architecture in ( |
| 817 | "KimiK3ForConditionalGeneration", |
| 818 | "KimiLinearForCausalLM", |
| 819 | ): |
| 820 | # Kimi K3 (multimodal) nests its MoE params in a KimiLinearConfig |
| 821 | # text_config and uses ``num_experts_per_token`` rather than the more |
| 822 | # common ``num_experts_per_tok``. get_text_config() returns the config |
| 823 | # itself for the text-only KimiLinearForCausalLM. |
| 824 | text_config = config.get_text_config() |
| 825 | E = text_config.num_experts |
| 826 | topk = text_config.num_experts_per_token |
| 827 | intermediate_size = text_config.moe_intermediate_size |
| 828 | hidden_size = text_config.hidden_size |
| 829 | elif architecture == "PixtralForConditionalGeneration": |
| 830 | # Pixtral can contain different LLM architectures, |
| 831 | # recurse to get their parameters |
| 832 | return get_model_params(config.get_text_config()) |
| 833 | else: |
| 834 | # Support for llama4 |
| 835 | config = config.get_text_config() |
| 836 | # Default: Mixtral. |
| 837 | E = config.num_local_experts |
| 838 | topk = config.num_experts_per_tok |
| 839 | intermediate_size = config.intermediate_size |
| 840 | hidden_size = config.hidden_size |
| 841 | return E, topk, intermediate_size, hidden_size |
| 842 | |
| 843 | |
| 844 | def resolve_dtype(config) -> torch.dtype: |
| 845 | if current_platform.is_rocm(): |
| 846 | return torch.float16 |
| 847 | |
| 848 | dtype = getattr(config, "dtype", None) |
| 849 | if dtype is not None: |
| 850 | return dtype |
| 851 | |
| 852 | if hasattr(config, "get_text_config"): |
| 853 | text_config = config.get_text_config() |
| 854 | dtype = getattr(text_config, "dtype", None) |
| 855 | if dtype is not None: |
| 856 | return dtype |
| 857 | |
| 858 | return torch.bfloat16 |
| 859 | |
| 860 | |
| 861 | def get_quantization_group_size(config) -> int | None: |
| 862 | """Extract the quantization group size from the HF model config. |
| 863 | |
| 864 | This reads directly from the HuggingFace config object (as returned by |
| 865 | ``get_config()``), not from vLLM's quantization config classes. |
| 866 | |
| 867 | Supports AWQ/GPTQ-style configs (direct 'group_size' key) and |
| 868 | compressed-tensors configs (nested inside 'config_groups'). |
| 869 | """ |
| 870 | quantization_config = getattr(config, "quantization_config", {}) |
| 871 | if not isinstance(quantization_config, dict): |
| 872 | return None |
| 873 | # AWQ / GPTQ style: group_size is a top-level key |
| 874 | gs = quantization_config.get("group_size") |
| 875 | if gs is not None: |
| 876 | return gs |
| 877 | # compressed-tensors style: group_size is nested in config_groups |
| 878 | config_groups = quantization_config.get("config_groups", {}) |
| 879 | if not isinstance(config_groups, dict): |
| 880 | return None |
| 881 | for group_cfg in config_groups.values(): |
| 882 | if not isinstance(group_cfg, dict): |
| 883 | continue |
| 884 | weights = group_cfg.get("weights", {}) |
| 885 | if not isinstance(weights, dict): |
| 886 | continue |
| 887 | gs = weights.get("group_size") |
| 888 | if gs is not None: |
| 889 | return gs |
| 890 | return None |
| 891 | |
| 892 | |
| 893 | def main(args: argparse.Namespace): |
| 894 | print(args) |
| 895 | |
| 896 | config = get_config(model=args.model, trust_remote_code=args.trust_remote_code) |
| 897 | if args.model_prefix: |
| 898 | config = getattr(config, args.model_prefix) |
| 899 | E, topk, intermediate_size, hidden_size = get_model_params(config) |
| 900 | enable_ep = bool(args.enable_expert_parallel) |
| 901 | if enable_ep: |
| 902 | ensure_divisibility(E, args.tp_size, "Number of experts") |
| 903 | E = E // args.tp_size |
| 904 | shard_intermediate_size = 2 * intermediate_size |
| 905 | else: |
| 906 | ensure_divisibility(intermediate_size, args.tp_size, "intermediate_size") |
| 907 | shard_intermediate_size = 2 * intermediate_size // args.tp_size |
| 908 | dtype = resolve_dtype(config) |
| 909 | use_fp8_w8a8 = args.dtype == "fp8_w8a8" |
| 910 | use_int8_w8a16 = args.dtype == "int8_w8a16" |
| 911 | use_int4_w4a16 = args.dtype == "int4_w4a16" |
| 912 | block_quant_shape = get_weight_block_size_safety(config) |
| 913 | if use_int4_w4a16: |
| 914 | group_size = get_quantization_group_size(config) |
| 915 | if group_size is None: |
| 916 | raise ValueError( |
| 917 | "Could not determine group_size from model config. " |
| 918 | "The model's quantization_config must contain a 'group_size' " |
| 919 | "field (AWQ/GPTQ) or 'config_groups.*.weights.group_size' " |
| 920 | "(compressed-tensors)." |
| 921 | ) |
| 922 | # For int4_w4a16, block_shape = [0, group_size] |
| 923 | # block_shape[0]=0 means no block quantization on N dimension |
| 924 | block_quant_shape = [0, group_size] |
| 925 | |
| 926 | if args.batch_size is None: |
| 927 | batch_sizes = [ |
| 928 | 1, |
| 929 | 2, |
| 930 | 4, |
| 931 | 8, |
| 932 | 16, |
| 933 | 24, |
| 934 | 32, |
| 935 | 48, |
| 936 | 64, |
| 937 | 96, |
| 938 | 128, |
| 939 | 256, |
| 940 | 512, |
| 941 | 1024, |
| 942 | 1536, |
| 943 | 2048, |
| 944 | 3072, |
| 945 | 4096, |
| 946 | ] |
| 947 | else: |
| 948 | batch_sizes = args.batch_size |
| 949 | |
| 950 | use_deep_gemm = bool(args.use_deep_gemm) |
| 951 | |
| 952 | if current_platform.is_rocm() and "HIP_VISIBLE_DEVICES" in os.environ: |
| 953 | # Ray will set ROCR_VISIBLE_DEVICES for device visibility |
| 954 | logger.warning( |
| 955 | "Ray uses ROCR_VISIBLE_DEVICES to control device accessibility." |
| 956 | "Replacing HIP_VISIBLE_DEVICES with ROCR_VISIBLE_DEVICES." |
| 957 | ) |
| 958 | val = os.environ["HIP_VISIBLE_DEVICES"] |
| 959 | os.environ["ROCR_VISIBLE_DEVICES"] = val |
| 960 | del os.environ["HIP_VISIBLE_DEVICES"] |
| 961 | |
| 962 | ray.init() |
| 963 | num_gpus = int(ray.available_resources()["GPU"]) |
| 964 | workers = [BenchmarkWorker.remote(args.seed) for _ in range(num_gpus)] |
| 965 | |
| 966 | def _distribute(method: str, inputs: list[Any]) -> list[Any]: |
| 967 | outputs = [] |
| 968 | worker_idx = 0 |
| 969 | for input_args in inputs: |
| 970 | worker = workers[worker_idx] |
| 971 | worker_method = getattr(worker, method) |
| 972 | output = worker_method.remote(*input_args) |
| 973 | outputs.append(output) |
| 974 | worker_idx = (worker_idx + 1) % num_gpus |
| 975 | return ray.get(outputs) |
| 976 | |
| 977 | if args.tune: |
| 978 | # int4_w4a16 weights are uint8-packed, not fp16; treat like fp8 for |
| 979 | # search space generation (no matrix_instr_nonkdim/kpack exploration). |
| 980 | is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16) |
| 981 | # For int4_w4a16, the group_size constraint on BLOCK_SIZE_K does not |
| 982 | # apply: the gptq_awq kernel handles arbitrary BLOCK_SIZE_K regardless |
| 983 | # of group_size. Skip block_quant_shape filtering to keep the full |
| 984 | # search space (e.g. BLOCK_SIZE_K=64 with group_size=128). |
| 985 | tune_block_quant_shape = None if use_int4_w4a16 else block_quant_shape |
| 986 | search_space = get_configs_compute_bound(is_fp16, tune_block_quant_shape) |
| 987 | if use_int4_w4a16: |
| 988 | # SPLIT_K is a required kernel constexpr for gptq_awq kernel; |
| 989 | # only SPLIT_K=1 is used at runtime, so fix it during tuning. |
| 990 | for cfg in search_space: |
| 991 | cfg["SPLIT_K"] = 1 |
| 992 | print(f"Start tuning over {len(search_space)} configurations...") |
| 993 | if use_deep_gemm: |
| 994 | raise ValueError( |
| 995 | "Tuning with --use-deep-gemm is not supported as it only tunes Triton " |
| 996 | "kernels. Please remove the flag." |
| 997 | ) |
| 998 | start = time.time() |
| 999 | configs = _distribute( |
| 1000 | "tune", |
| 1001 | [ |
| 1002 | ( |
| 1003 | batch_size, |
| 1004 | E, |
| 1005 | shard_intermediate_size, |
| 1006 | hidden_size, |
| 1007 | topk, |
| 1008 | dtype, |
| 1009 | use_fp8_w8a8, |
| 1010 | use_int8_w8a16, |
| 1011 | use_int4_w4a16, |
| 1012 | search_space, |
| 1013 | block_quant_shape, |
| 1014 | use_deep_gemm, |
| 1015 | ) |
| 1016 | for batch_size in batch_sizes |
| 1017 | ], |
| 1018 | ) |
| 1019 | best_configs = { |
| 1020 | M: sort_config(config) for M, config in zip(batch_sizes, configs) |
| 1021 | } |
| 1022 | save_configs( |
| 1023 | best_configs, |
| 1024 | E, |
| 1025 | shard_intermediate_size, |
| 1026 | hidden_size, |
| 1027 | topk, |
| 1028 | dtype, |
| 1029 | use_fp8_w8a8, |
| 1030 | use_int8_w8a16, |
| 1031 | use_int4_w4a16, |
| 1032 | block_quant_shape, |
| 1033 | args.save_dir, |
| 1034 | ) |
| 1035 | end = time.time() |
| 1036 | print(f"Tuning took {end - start:.2f} seconds") |
| 1037 | else: |
| 1038 | outputs = _distribute( |
| 1039 | "benchmark", |
| 1040 | [ |
| 1041 | ( |
| 1042 | batch_size, |
| 1043 | E, |
| 1044 | shard_intermediate_size, |
| 1045 | hidden_size, |
| 1046 | topk, |
| 1047 | dtype, |
| 1048 | use_fp8_w8a8, |
| 1049 | use_int8_w8a16, |
| 1050 | use_int4_w4a16, |
| 1051 | block_quant_shape, |
| 1052 | use_deep_gemm, |
| 1053 | ) |
| 1054 | for batch_size in batch_sizes |
| 1055 | ], |
| 1056 | ) |
| 1057 | |
| 1058 | for batch_size, (config, kernel_time) in zip(batch_sizes, outputs): |
| 1059 | print(f"Batch size: {batch_size}, config: {config}") |
| 1060 | print(f"Kernel time: {kernel_time:.2f} us") |
| 1061 | |
| 1062 | |
| 1063 | if __name__ == "__main__": |
| 1064 | parser = FlexibleArgumentParser() |
| 1065 | parser.add_argument( |
| 1066 | "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" |
| 1067 | ) |
| 1068 | parser.add_argument( |
| 1069 | "--tp-size", "-tp", "--tensor-parallel-size", type=int, default=2 |
| 1070 | ) |
| 1071 | parser.add_argument("--enable-expert-parallel", "-enable-ep", action="store_true") |
| 1072 | parser.add_argument( |
| 1073 | "--dtype", |
| 1074 | type=str, |
| 1075 | choices=["auto", "fp8_w8a8", "int8_w8a16", "int4_w4a16"], |
| 1076 | default="auto", |
| 1077 | ) |
| 1078 | parser.add_argument("--use-deep-gemm", action="store_true") |
| 1079 | parser.add_argument( |
| 1080 | "--save-dir", type=str, default="./", help="Directory to save tuned results" |
| 1081 | ) |
| 1082 | parser.add_argument("--seed", type=int, default=0) |
| 1083 | parser.add_argument("--batch-size", type=int, nargs="+", required=False) |
| 1084 | parser.add_argument("--tune", action="store_true") |
| 1085 | parser.add_argument("--trust-remote-code", action="store_true") |
| 1086 | parser.add_argument("--model-prefix", type=str, required=False) |
| 1087 | args = parser.parse_args() |
| 1088 | |
| 1089 | main(args) |