benchmark_moe.py 31.7 KB
Newer Older
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, Optional
13
14
15
16
17

import ray
import torch
from ray.experimental.tqdm_ray import tqdm

18
19
20
21
from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
    _get_config_dtype_str,
)
22
from vllm.model_executor.layers.fused_moe.fused_moe import *
23
from vllm.transformers_utils.config import get_config
24
from vllm.triton_utils import triton
25
from vllm.utils.argparse_utils import FlexibleArgumentParser
26
from vllm.utils.torch_utils import set_random_seed
27

28
29
# 移除全局的 current_platform 导入,改为在需要时局部导入
# FP8_DTYPE = current_platform.fp8_dtype()
30

31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
# Default interval for clearing Triton JIT cache during tuning
# Set to 0 to disable automatic cache clearing
_CACHE_CLEAR_INTERVAL_ENV = "VLLM_MOE_TUNE_CACHE_CLEAR_INTERVAL"
TRITON_CACHE_CLEAR_INTERVAL = int(os.environ.get(_CACHE_CLEAR_INTERVAL_ENV, "50"))


def clear_triton_cache():
    """Clear Triton JIT compilation cache and Python/CUDA memory.

    This helps prevent OOM during tuning with large models (many experts).
    """
    # Force Python garbage collection
    gc.collect()

    # Clear CUDA memory cache
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Try to clear Triton's runtime cache
    try:
        if (
            hasattr(triton, "runtime")
            and hasattr(triton.runtime, "cache")
            and hasattr(triton.runtime.cache, "clear")
        ):
            triton.runtime.cache.clear()
    except ImportError:
        # Triton not installed, skip cache clearing
        pass
    except AttributeError:
        # Triton version doesn't have expected cache API
        pass
    except Exception as e:
        print(f"Warning: Failed to clear Triton cache: {e}")

    # Additional garbage collection after clearing caches
    gc.collect()

69

70
def ensure_divisibility(numerator, denominator, text):
71
    """Ensure that numerator is divisible by the denominator."""
72
73
    assert numerator % denominator == 0, "{} {} is not divisible by tp {}.".format(
        text, numerator, denominator
74
75
76
    )


77
78
79
80
81
82
83
class BenchmarkConfig(TypedDict):
    BLOCK_SIZE_M: int
    BLOCK_SIZE_N: int
    BLOCK_SIZE_K: int
    GROUP_SIZE_M: int
    num_warps: int
    num_stages: int
zhuwenwen's avatar
zhuwenwen committed
84
    num_ldmatrixes: Optional[int]
85
86


87
88
89
90
91
92
93
94
95
96
97
def benchmark_config(
    config: BenchmarkConfig,
    num_tokens: int,
    num_experts: int,
    shard_intermediate_size: int,
    hidden_size: int,
    topk: int,
    dtype: torch.dtype,
    use_fp8_w8a8: bool,
    use_int8_w8a16: bool,
    num_iters: int = 100,
98
    block_quant_shape: list[int] = None,
99
    use_deep_gemm: bool = False,
zhuwenwen's avatar
zhuwenwen committed
100
    nn_moe: Optional[bool] = False
101
) -> float:
102
103
104
    from vllm.platforms import current_platform
    device = torch.cuda.current_device()

105
    init_dtype = torch.float16 if use_fp8_w8a8 else dtype
106
    x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
107
    if use_int8_w8a16:
zhuwenwen's avatar
zhuwenwen committed
108
        if not nn_moe:
zhuwenwen's avatar
zhuwenwen committed
109
110
111
112
113
114
115
116
117
            w1 = torch.randint(
                -127,
                127,
                (
                    num_experts,
                    shard_intermediate_size,
                    hidden_size,
                ),
                dtype=torch.int8,
118
                device=device,
zhuwenwen's avatar
zhuwenwen committed
119
120
121
122
123
124
125
126
127
128
            )
            w2 = torch.randint(
                -127,
                127,
                (
                    num_experts,
                    hidden_size,
                    shard_intermediate_size // 2,
                ),
                dtype=torch.int8,
129
                device=device,
zhuwenwen's avatar
zhuwenwen committed
130
            )
zhuwenwen's avatar
zhuwenwen committed
131
        else:
zhuwenwen's avatar
zhuwenwen committed
132
133
134
135
136
137
138
139
140
            w1 = torch.randint(
                -127,
                127,
                (
                    num_experts,
                    hidden_size,
                    shard_intermediate_size,
                ),
                dtype=torch.int8,
141
                device=device,
zhuwenwen's avatar
zhuwenwen committed
142
143
144
145
146
147
148
149
150
151
            )
            w2 = torch.randint(
                -127,
                127,
                (
                    num_experts,
                    shard_intermediate_size // 2,
                    hidden_size,
                ),
                dtype=torch.int8,
152
                device=device,
zhuwenwen's avatar
zhuwenwen committed
153
            )
154
    else:
zhuwenwen's avatar
zhuwenwen committed
155
        if not nn_moe:
zhuwenwen's avatar
zhuwenwen committed
156
            w1 = torch.randn(
157
                num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype, device=device
zhuwenwen's avatar
zhuwenwen committed
158
159
            )
            w2 = torch.randn(
160
                num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype, device=device
zhuwenwen's avatar
zhuwenwen committed
161
            )
zhuwenwen's avatar
zhuwenwen committed
162
        else:
zhuwenwen's avatar
zhuwenwen committed
163
            w1 = torch.randn(
164
                num_experts, hidden_size, shard_intermediate_size, dtype=init_dtype, device=device
zhuwenwen's avatar
zhuwenwen committed
165
166
            )
            w2 = torch.randn(
167
                num_experts, shard_intermediate_size // 2, hidden_size, dtype=init_dtype, device=device
zhuwenwen's avatar
zhuwenwen committed
168
            )
169
    gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32, device=device)
170
171
172
173
174

    w1_scale = None
    w2_scale = None
    a1_scale = None
    a2_scale = None
175
    if use_int8_w8a16:
176
        w1_scale = torch.randn(
177
            (num_experts, 2 * shard_intermediate_size), dtype=torch.float32, device=device
178
        )
179
        w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32, device=device)
180
181
182
    if use_deep_gemm:
        # we use the default block shape for deepgemm
        block_quant_shape = [128, 128]
183
    if use_fp8_w8a8:
184
185
186
187
188
189
190
191
192
193
        if block_quant_shape:
            block_n, block_k = block_quant_shape[0], block_quant_shape[1]
            E = num_experts
            N = shard_intermediate_size // 2
            K = hidden_size
            factor_for_scale = 1e-2
            n_tiles_w1 = (2 * N + block_n - 1) // block_n
            n_tiles_w2 = (K + block_n - 1) // block_n
            k_tiles_w1 = (K + block_k - 1) // block_k
            k_tiles_w2 = (N + block_k - 1) // block_k
194
            w1_scale = (
195
                torch.rand((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32, device=device)
196
197
198
                * factor_for_scale
            )
            w2_scale = (
199
                torch.rand((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32, device=device)
200
201
                * factor_for_scale
            )
202
        else:
203
204
            w1_scale = torch.randn(num_experts, dtype=torch.float32, device=device)
            w2_scale = torch.randn(num_experts, dtype=torch.float32, device=device)
205

206
207
        a1_scale = torch.randn(1, dtype=torch.float32, device=device)
        a2_scale = torch.randn(1, dtype=torch.float32, device=device)
208

209
210
        # 获取 FP8_DTYPE
        FP8_DTYPE = current_platform.fp8_dtype()
211
212
        w1 = w1.to(FP8_DTYPE)
        w2 = w2.to(FP8_DTYPE)
213

214
    input_gating = torch.empty(num_tokens, num_experts, dtype=torch.float32, device=device)
215
216
217
218
219

    def prepare(i: int):
        input_gating.copy_(gating_output[i])

    def run():
220
        from vllm.model_executor.layers.fused_moe import override_config
221

222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
        if use_fp8_w8a8:
            quant_dtype = torch.float8_e4m3fn
        elif use_int8_w8a16:
            quant_dtype = torch.int8
        else:
            quant_dtype = None

        quant_config = FusedMoEQuantConfig.make(
            quant_dtype=quant_dtype,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            block_shape=block_quant_shape,
        )

238
        with override_config(config):
239
240
241
242
243
244
245
246
247
248
249
250
            topk_weights, topk_ids, token_expert_indices = fused_topk(
                x, input_gating, topk, renormalize=not use_deep_gemm
            )
            return fused_experts(
                x,
                w1,
                w2,
                topk_weights,
                topk_ids,
                inplace=True,
                quant_config=quant_config,
                allow_deep_gemm=use_deep_gemm,
zhuwenwen's avatar
zhuwenwen committed
251
                use_nn_moe=nn_moe,
252
            )
253
254
255
256
257
258

    # JIT compilation & warmup
    run()
    torch.cuda.synchronize()

    # Capture 10 invocations with CUDA graph
259
260
261
262
263
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        for _ in range(10):
            run()
    torch.cuda.synchronize()
264
265
266

    # Warmup
    for _ in range(5):
267
268
        graph.replay()
        # run()
269
270
    torch.cuda.synchronize()

271
272
    start_event = torch.Event(enable_timing=True)
    end_event = torch.Event(enable_timing=True)
273

274
    latencies: list[float] = []
275
276
277
278
279
    for i in range(num_iters):
        prepare(i)
        torch.cuda.synchronize()

        start_event.record()
280
281
        graph.replay()
        # run()
282
283
284
285
        end_event.record()
        end_event.synchronize()
        latencies.append(start_event.elapsed_time(end_event))
    avg = sum(latencies) / (num_iters * 10) * 1000  # us
286
    graph.reset()
287
288
289
    return avg


zhuwenwen's avatar
zhuwenwen committed
290
def get_rocm_tuning_space(use_fp16, nn_moe: Optional[bool] = False):
291
292
    block_m_range = [16, 32, 64, 128, 256]
    block_n_range = [32, 64, 128, 256]
293
    block_k_range = [32, 64, 128, 256]
294
295
296
297
298
    if not use_fp16:
        block_k_range.remove(16)  # BLOCK_K=16 not supported for fp8
    num_warps_range = [1, 2, 4, 8]
    group_m_range = [1, 4, 8, 16, 32]
    num_stage_range = [2]
299
    # waves_per_eu_range = [0, 1, 2, 4]
300
301
    # matrix_instr_nonkdim_range = [16, 32] if use_fp16 else []
    # kpack_range = [1, 2] if use_fp16 else []
302
303

    param_ranges = {
304
305
        "BLOCK_SIZE_M": block_m_range,
        "BLOCK_SIZE_N": block_n_range,
306
307
308
309
        "BLOCK_SIZE_K": block_k_range,
        "GROUP_SIZE_M": group_m_range,
        "num_warps": num_warps_range,
        "num_stages": num_stage_range,
310
        # "waves_per_eu": waves_per_eu_range,
311
    }
zhuwenwen's avatar
zhuwenwen committed
312
    if nn_moe:
313
314
315
316
317
318
        param_ranges["num_ldmatrixes"] = [1]
    
    # DCU currently does not support the following parameters
    # if use_fp16:
    #     param_ranges["matrix_instr_nonkdim"] = matrix_instr_nonkdim_range
    #     param_ranges["kpack"] = kpack_range
319
320
321
322

    return param_ranges


zhuwenwen's avatar
zhuwenwen committed
323
def get_configs_compute_bound(use_fp16, block_quant_shape, nn_moe: Optional[bool] = False) -> list[dict[str, int]]:
324
    configs: list[BenchmarkConfig] = []
325
326
327
    
    # 局部导入 current_platform
    from vllm.platforms import current_platform
328
329

    if current_platform.is_rocm():
zhuwenwen's avatar
zhuwenwen committed
330
        param_ranges = get_rocm_tuning_space(use_fp16, nn_moe)
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
    else:
        # Reduced search space for faster tuning.
        # TODO(woosuk): Increase the search space and use a performance model to
        # prune the search space.
        block_m_range = [16, 32, 64, 128, 256]
        block_n_range = [32, 64, 128, 256]
        block_k_range = [64, 128, 256]
        num_warps_range = [4, 8]
        group_m_range = [1, 16, 32, 64]
        num_stage_range = [2, 3, 4, 5]

        param_ranges = {
            "BLOCK_SIZE_M": block_m_range,
            "BLOCK_SIZE_N": block_n_range,
            "BLOCK_SIZE_K": block_k_range,
            "GROUP_SIZE_M": group_m_range,
            "num_warps": num_warps_range,
            "num_stages": num_stage_range,
        }

    keys, values = zip(*param_ranges.items())
    for config_values in product(*values):
        config = dict(zip(keys, config_values))
        configs.append(config)
355
356
357
358
359
360
361

    # Remove configs that are not compatible with fp8 block quantization
    # BLOCK_SIZE_K must be a multiple of block_k
    # BLOCK_SIZE_N must be a multiple of block_n
    if block_quant_shape is not None and not use_fp16:
        block_n, block_k = block_quant_shape[0], block_quant_shape[1]
        for config in configs[:]:
362
363
364
365
            if (
                config["BLOCK_SIZE_K"] % block_k != 0
                or config["BLOCK_SIZE_N"] % block_n != 0
            ):
366
                configs.remove(config)
367
368
369
    return configs


370
371
372
def prune_rocm_search_space(
    num_tokens, shard_intermediate_size, hidden_size, search_space, is_fp16, topk
):
373
374
    N1, K1 = shard_intermediate_size, hidden_size
    N2, K2 = hidden_size, shard_intermediate_size // 2
375
376
377
378
379
380
    pruned_space_1 = prune_rocm_configs(
        num_tokens * topk, N1, K1, search_space, is_fp16
    )
    pruned_space_2 = prune_rocm_configs(
        num_tokens * topk, N2, K2, search_space, is_fp16
    )
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
    search_space = merge_unique_dicts(pruned_space_1, pruned_space_2)
    return search_space


# The following code is inspired by ROCm/Triton GEMM tuning script:
# https://github.com/ROCm/triton/blob/triton-mlir/scripts/amd/gemm/tune_gemm.py#L89
def prune_rocm_configs(M, N, K, configs, is_fp16=True):
    pruned_configs = []
    elemBytes_a = 2 if is_fp16 else 1
    elemBytes_b = 2 if is_fp16 else 1

    mfma = 16 if M < 32 or N < 32 else 32

    # TODO (zhanglx): figure out the boundary between large and small gemms
    large_gemm = False
    if M >= 2048 and N >= 2048:
        large_gemm = True

    for config in configs:
        BLOCK_SIZE_M = config.get("BLOCK_SIZE_M")
        BLOCK_SIZE_N = config.get("BLOCK_SIZE_N")
        BLOCK_SIZE_K = config.get("BLOCK_SIZE_K")
        num_warps = config.get("num_warps")

405
406
407
408
409
        # DCU currently does not support matrix_instr_nonkdim param
        # if is_fp16:
        #     matrix_instr_nonkdim = config.get("matrix_instr_nonkdim")
        #     if matrix_instr_nonkdim > mfma:
        #         continue
410
411
412
413
414
415
416
417
        if mfma == 4 and BLOCK_SIZE_K < 64:
            continue
        # some layouts could not work properly in case
        # number elements per thread is less 1
        if BLOCK_SIZE_M * BLOCK_SIZE_N < 64:
            continue
        SPLIT_K = config.get("SPLIT_K", 1)
        GROUP_M = config.get("GROUP_SIZE_M")
418
419
420

        # DCU currently does not support matrix_instr_nonkdim param
        # if is_fp16:
zhuwenwen's avatar
zhuwenwen committed
421
422
423
424
        #     if (
        #         matrix_instr_nonkdim > BLOCK_SIZE_M
        #         or matrix_instr_nonkdim > BLOCK_SIZE_N
        #     ):
425
        #         continue
zhuwenwen's avatar
zhuwenwen committed
426
        #     if matrix_instr_nonkdim >= M and matrix_instr_nonkdim != BLOCK_SIZE_M:
427
        #         continue
zhuwenwen's avatar
zhuwenwen committed
428
        #     if matrix_instr_nonkdim >= N and matrix_instr_nonkdim != BLOCK_SIZE_N:
429
        #         continue
zhuwenwen's avatar
zhuwenwen committed
430
        
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
        # Skip BLOCK_SIZE that is too large compare to M/N
        # unless BLOCK_SIZE is already small enough
        if M * 2 < BLOCK_SIZE_M and BLOCK_SIZE_M != 16:
            continue
        if N * 2 < BLOCK_SIZE_N and BLOCK_SIZE_N != 16:
            continue
        # skip large split_k when not necessary
        if SPLIT_K != 1 and not need_split_k(M, N, K):
            continue
        # skip split_k that leads to EVEN_K = false
        leap = SPLIT_K * BLOCK_SIZE_K
        modv = K % leap
        if modv != 0:
            continue
        # skip large GROUP_M
        if GROUP_M * BLOCK_SIZE_M > M and GROUP_M != 1:
            continue
        # out of shared memory resource
        # TODO (zhanglx): This does not consider the LDS usage in the epilogue
450
451
452
453
        LDS = (
            BLOCK_SIZE_K * BLOCK_SIZE_M * elemBytes_a
            + BLOCK_SIZE_K * BLOCK_SIZE_N * elemBytes_b
        )
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
        if LDS > 65536:
            continue
        # Skip small block sizes and num_warps for large gemm
        # For fp16 and f8, we want to only use BLOCK_SIZE >= 64
        if large_gemm:
            if BLOCK_SIZE_M < 64 or BLOCK_SIZE_N < 64:
                continue
            if BLOCK_SIZE_K < 64:
                continue
            if num_warps < 4:
                continue

        pruned_configs.append(config)

    return pruned_configs


def need_split_k(SIZE_M, SIZE_N, SIZE_K):
    return (SIZE_M < 64 or SIZE_N < 64) and SIZE_K > 1024


def merge_unique_dicts(list1, list2):
    result = []
    combined_list = list1.copy()
    combined_list.extend(list2)
    for dictionary in combined_list:
        if dictionary not in result:
            result.append(dictionary)
    return result


485
486
@ray.remote(num_gpus=1)
class BenchmarkWorker:
王敏's avatar
王敏 committed
487
    def __init__(self, seed: int, device_id: int) -> None:
488
489
490
491
492
493
494
495
496
        from vllm.platforms import current_platform
        import os
        
        if current_platform.is_rocm():
            # In ROCm environment with Ray, let Ray handle device assignment
            # Don't manually set default device as it may conflict with Ray's device mapping
            pass
        else:
            torch.set_default_device("cuda:"+ str(device_id))
497

498
        set_random_seed(seed)
499
        self.seed = seed
500
        # Store the logical device ID for Ray
王敏's avatar
王敏 committed
501
        self.device_id = device_id
502
503
504
505
506
507
508
509
510

    def benchmark(
        self,
        num_tokens: int,
        num_experts: int,
        shard_intermediate_size: int,
        hidden_size: int,
        topk: int,
        dtype: torch.dtype,
511
512
        use_fp8_w8a8: bool,
        use_int8_w8a16: bool,
513
        block_quant_shape: list[int] = None,
514
        use_deep_gemm: bool = False,
515
        nn_moe: Optional[bool] = False,
516
    ) -> tuple[dict[str, int], float]:
517
518
        # 局部导入 current_platform
        from vllm.platforms import current_platform
519
520
521
        
        from vllm.model_executor.layers.fused_moe.fused_moe import  get_moe_configs, get_default_config

522
        set_random_seed(self.seed)
523
        dtype_str = _get_config_dtype_str(
524
525
            dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
        )
526
527
        # NOTE(woosuk): The current naming convention uses w2.shape[2], which
        # is the intermediate size after silu_and_mul.
528
529
        block_n = block_quant_shape[0] if block_quant_shape else None
        block_k = block_quant_shape[1] if block_quant_shape else None
530
        op_config = get_moe_configs(
531
            num_experts, shard_intermediate_size // 2, dtype_str, block_n, block_k, use_nn_moe=nn_moe
532
        )
533
        if op_config is None:
534
535
536
537
538
539
540
            config = get_default_config(
                num_tokens,
                num_experts,
                shard_intermediate_size,
                hidden_size,
                topk,
                dtype_str,
541
                block_quant_shape,
542
                use_nn_moe=nn_moe,
543
            )
544
        else:
545
546
547
548
549
550
551
552
553
554
555
556
557
558
            config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))]
        kernel_time = benchmark_config(
            config,
            num_tokens,
            num_experts,
            shard_intermediate_size,
            hidden_size,
            topk,
            dtype,
            use_fp8_w8a8,
            use_int8_w8a16,
            num_iters=100,
            block_quant_shape=block_quant_shape,
            use_deep_gemm=use_deep_gemm,
zhuwenwen's avatar
zhuwenwen committed
559
            use_nn_moe=nn_moe
560
        )
561
562
563
564
565
566
567
568
569
570
        return config, kernel_time

    def tune(
        self,
        num_tokens: int,
        num_experts: int,
        shard_intermediate_size: int,
        hidden_size: int,
        topk: int,
        dtype: torch.dtype,
571
572
        use_fp8_w8a8: bool,
        use_int8_w8a16: bool,
573
        search_space: list[dict[str, int]],
574
        block_quant_shape: list[int],
575
        use_deep_gemm: bool,
王敏's avatar
王敏 committed
576
        nn_moe: Optional[bool] = False,
577
    ) -> dict[str, int]:
578
579
580
        from vllm.platforms import current_platform
        import os

581
582
        best_config = None
        best_time = float("inf")
583
584
        if current_platform.is_rocm():
            is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
585
586
587
588
589
590
591
592
            search_space = prune_rocm_search_space(
                num_tokens,
                shard_intermediate_size,
                hidden_size,
                search_space,
                is_fp16,
                topk,
            )
593

594
595
        # In ROCm environments with Ray, device context is already handled by Ray
        # Using torch.cuda.device() may cause device ordinal conflicts
596
597
        need_device_guard = False
        if current_platform.is_rocm():
598
599
600
601
602
603
            # For ROCm with Ray, skip additional device context management
            need_device_guard = False
        else:
            # For other platforms, use device guard if needed
            visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
            if visible_devices is not None and len(visible_devices.split(',')) > 1:
604
                need_device_guard = True
605

606
        with torch.cuda.device(self.device_id) if need_device_guard else nullcontext():
607
            for idx, config in enumerate(tqdm(search_space)):
608
                try:
609
610
611
612
613
614
615
616
617
618
619
                    kernel_time = benchmark_config(
                        config,
                        num_tokens,
                        num_experts,
                        shard_intermediate_size,
                        hidden_size,
                        topk,
                        dtype,
                        use_fp8_w8a8,
                        use_int8_w8a16,
                        num_iters=20,
zhuwenwen's avatar
zhuwenwen committed
620
                        block_quant_shape=block_quant_shape,
zhuwenwen's avatar
zhuwenwen committed
621
622
                        use_deep_gemm=use_deep_gemm,
                        nn_moe=nn_moe)
623
624
625
626
627
628
629
                except triton.runtime.autotuner.OutOfResources:
                    # Some configurations may be invalid and fail to compile.
                    continue

                if kernel_time < best_time:
                    best_time = kernel_time
                    best_config = config
630
631
632
633
634
635
636
637
638
639
640
641
642

                # Periodically clear Triton JIT cache to prevent OOM
                # This is especially important for large models with many experts
                if (
                    TRITON_CACHE_CLEAR_INTERVAL > 0
                    and idx > 0
                    and idx % TRITON_CACHE_CLEAR_INTERVAL == 0
                ):
                    clear_triton_cache()

        # Final cleanup after tuning completes
        clear_triton_cache()

643
644
        now = datetime.now()
        print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}")
645
        assert best_config is not None
646
647
648
        return best_config


649
def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:
650
651

    return {
652
653
654
655
656
657
        "BLOCK_SIZE_M": config["BLOCK_SIZE_M"],
        "BLOCK_SIZE_N": config["BLOCK_SIZE_N"],
        "BLOCK_SIZE_K": config["BLOCK_SIZE_K"],
        "GROUP_SIZE_M": config["GROUP_SIZE_M"],
        "num_warps": config["num_warps"],
        "num_stages": config["num_stages"],
zhuwenwen's avatar
zhuwenwen committed
658
659
660
        **(
            {"num_ldmatrixes": config["num_ldmatrixes"]} if "num_ldmatrixes" in config else {}
        ),
661
662
663
664
665
666
667
668
669
        **(
            {"waves_per_eu": config["waves_per_eu"]} if "waves_per_eu" in config else {}
        ),
        **(
            {"matrix_instr_nonkdim": config["matrix_instr_nonkdim"]}
            if "matrix_instr_nonkdim" in config
            else {}
        ),
        **({"kpack": config["kpack"]} if "kpack" in config else {}),
670
671
672
    }


673
674
675
676
677
678
679
680
681
def save_configs(
    configs: dict[int, BenchmarkConfig],
    num_experts: int,
    shard_intermediate_size: int,
    hidden_size: int,
    topk: int,
    dtype: torch.dtype,
    use_fp8_w8a8: bool,
    use_int8_w8a16: bool,
682
    block_quant_shape: list[int],
683
    save_dir: str,
zhuwenwen's avatar
zhuwenwen committed
684
    use_nn_moe: Optional[bool] = False,
685
) -> None:
686
    dtype_str = _get_config_dtype_str(
687
688
        dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
    )
689

690
691
    # NOTE(woosuk): The current naming convention uses w2.shape[2], which
    # is the intermediate size after silu_and_mul.
692
    filename = get_config_file_name(
zhuwenwen's avatar
zhuwenwen committed
693
        num_experts, shard_intermediate_size // 2, dtype_str, block_quant_shape, use_nn_moe=use_nn_moe
694
    )
695
696
    os.makedirs(save_dir, exist_ok=True)
    filename = os.path.join(save_dir, filename)
697
698
    print(f"Writing best config to {filename}...")
    with open(filename, "w") as f:
699
        json.dump({"triton_version": triton.__version__, **configs}, f, indent=4)
700
701
702
        f.write("\n")


703
def get_weight_block_size_safety(config, default_value=None):
704
    quantization_config = getattr(config, "quantization_config", {})
705
    if isinstance(quantization_config, dict):
706
        return quantization_config.get("weight_block_size", default_value)
707
708
709
    return default_value


710
def main(args: argparse.Namespace):
711
712
    import os
    import logging
713

714
715
716
717
    from vllm.platforms import current_platform
    
    logger = logging.getLogger(__name__)

718
    print(args)
zhuwenwen's avatar
zhuwenwen committed
719
    
王敏's avatar
王敏 committed
720
    tp_size = args.tp_size
721
    config = get_config(model=args.model, trust_remote_code=args.trust_remote_code)
722
723
    if args.model_prefix:
        config = getattr(config, args.model_prefix)
王敏's avatar
王敏 committed
724

725
726
727
728
    if config.architectures[0] == "DbrxForCausalLM":
        E = config.ffn_config.moe_num_experts
        topk = config.ffn_config.moe_top_k
        intermediate_size = config.ffn_config.ffn_hidden_size
729
        hidden_size = config.hidden_size
730
731
732
733
    elif config.architectures[0] == "JambaForCausalLM":
        E = config.num_experts
        topk = config.num_experts_per_tok
        intermediate_size = config.intermediate_size
734
        hidden_size = config.hidden_size
Yuxuan Zhang's avatar
Yuxuan Zhang committed
735
736
    elif config.architectures[0] in (
        "DeepseekV2ForCausalLM",
737
738
        "DeepseekV3ForCausalLM",
        "DeepseekV32ForCausalLM",
Yuxuan Zhang's avatar
Yuxuan Zhang committed
739
        "Glm4MoeForCausalLM",
740
        "NemotronHForCausalLM",
Yuxuan Zhang's avatar
Yuxuan Zhang committed
741
    ):
742
        E = config.n_routed_experts
王敏's avatar
王敏 committed
743
744
        topk = config.num_experts_per_tok
        intermediate_size = config.moe_intermediate_size
745
        hidden_size = config.hidden_size
746
747
748
749
750
    elif config.architectures[0] in (
        "Qwen2MoeForCausalLM",
        "Qwen3MoeForCausalLM",
        "Qwen3NextForCausalLM",
    ):
王敏's avatar
王敏 committed
751
        E = config.num_experts
752
753
        topk = config.num_experts_per_tok
        intermediate_size = config.moe_intermediate_size
754
755
756
757
758
759
760
        hidden_size = config.hidden_size
    elif config.architectures[0] == "Qwen3VLMoeForConditionalGeneration":
        text_config = config.get_text_config()
        E = text_config.num_experts
        topk = text_config.num_experts_per_tok
        intermediate_size = text_config.moe_intermediate_size
        hidden_size = text_config.hidden_size
761
762
763
764
    elif config.architectures[0] in ("HunYuanMoEV1ForCausalLM"):
        E = config.num_experts
        topk = config.moe_topk[0]
        intermediate_size = config.moe_intermediate_size[0]
765
        hidden_size = config.hidden_size
zhuwenwen's avatar
zhuwenwen committed
766
767
768
769
    elif config.architectures[0] in ("Step3VLForConditionalGeneration"):
        E = config.text_config.moe_num_experts
        topk = config.text_config.moe_top_k
        intermediate_size = config.text_config.moe_intermediate_size
770
771
772
773
774
    elif config.architectures[0] in ["Qwen3OmniMoeForConditionalGeneration"]:
        E = config.thinker_config.text_config.num_experts
        topk = config.thinker_config.text_config.num_experts_per_tok
        intermediate_size = config.thinker_config.text_config.moe_intermediate_size
        hidden_size = config.thinker_config.text_config.hidden_size
775
    else:
776
777
        # Support for llama4
        config = config.get_text_config()
778
779
780
781
        # Default: Mixtral.
        E = config.num_local_experts
        topk = config.num_experts_per_tok
        intermediate_size = config.intermediate_size
782
        hidden_size = config.hidden_size
783
784
    enable_ep = bool(args.enable_expert_parallel)
    if enable_ep:
zhuwenwen's avatar
zhuwenwen committed
785
786
        ensure_divisibility(E, tp_size, "Number of experts")
        E = E // tp_size
787
788
789
        shard_intermediate_size = 2 * intermediate_size
    else:
        ensure_divisibility(intermediate_size, args.tp_size, "intermediate_size")
zhuwenwen's avatar
zhuwenwen committed
790
        shard_intermediate_size = 2 * intermediate_size // tp_size
791
    dtype = torch.float16 if current_platform.is_rocm() else config.dtype
792
793
    use_fp8_w8a8 = args.dtype == "fp8_w8a8"
    use_int8_w8a16 = args.dtype == "int8_w8a16"
794
    block_quant_shape = get_weight_block_size_safety(config)
795
796

    if args.batch_size is None:
797
        batch_sizes = [
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
            1,
            2,
            4,
            8,
            16,
            24,
            32,
            48,
            64,
            96,
            128,
            256,
            512,
            1024,
            1536,
            2048,
            3072,
            4096,
816
        ]
817
    else:
818
        batch_sizes = args.batch_size
819

820
821
    use_deep_gemm = bool(args.use_deep_gemm)

822
823
824
825
    if current_platform.is_rocm() and "HIP_VISIBLE_DEVICES" in os.environ:
        # Ray will set ROCR_VISIBLE_DEVICES for device visibility
        logger.warning(
            "Ray uses ROCR_VISIBLE_DEVICES to control device accessibility."
826
827
            "Replacing HIP_VISIBLE_DEVICES with ROCR_VISIBLE_DEVICES."
        )
828
829
830
831
        val = os.environ["HIP_VISIBLE_DEVICES"]
        os.environ["ROCR_VISIBLE_DEVICES"] = val
        del os.environ["HIP_VISIBLE_DEVICES"]

zhuwenwen's avatar
zhuwenwen committed
832
    ray.init(address=None, ignore_reinit_error=True, num_gpus=args.num_gpus)
833
    num_gpus = int(ray.available_resources()["GPU"])
王敏's avatar
王敏 committed
834
    workers = [BenchmarkWorker.remote(args.seed, i) for i in range(num_gpus)]
835

836
    def _distribute(method: str, inputs: list[Any]) -> list[Any]:
837
838
839
840
841
842
843
844
845
846
847
        outputs = []
        worker_idx = 0
        for input_args in inputs:
            worker = workers[worker_idx]
            worker_method = getattr(worker, method)
            output = worker_method.remote(*input_args)
            outputs.append(output)
            worker_idx = (worker_idx + 1) % num_gpus
        return ray.get(outputs)

    if args.tune:
848
        is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
zhuwenwen's avatar
zhuwenwen committed
849
        search_space = get_configs_compute_bound(is_fp16, block_quant_shape, args.nn_moe)
850
        print(f"Start tuning over {len(search_space)} configurations...")
851
852
853
854
855
        if use_deep_gemm:
            raise ValueError(
                "Tuning with --use-deep-gemm is not supported as it only tunes Triton "
                "kernels. Please remove the flag."
            )
856
857
        start = time.time()
        configs = _distribute(
858
859
860
861
862
863
864
865
866
867
868
869
870
871
            "tune",
            [
                (
                    batch_size,
                    E,
                    shard_intermediate_size,
                    hidden_size,
                    topk,
                    dtype,
                    use_fp8_w8a8,
                    use_int8_w8a16,
                    search_space,
                    block_quant_shape,
                    use_deep_gemm,
zhuwenwen's avatar
zhuwenwen committed
872
                    args.nn_moe,
873
874
875
876
                )
                for batch_size in batch_sizes
            ],
        )
877
        best_configs = {
878
            M: sort_config(config) for M, config in zip(batch_sizes, configs)
879
        }
880
881
882
883
884
885
886
887
888
889
        save_configs(
            best_configs,
            E,
            shard_intermediate_size,
            hidden_size,
            topk,
            dtype,
            use_fp8_w8a8,
            use_int8_w8a16,
            block_quant_shape,
890
            args.save_dir,
zhuwenwen's avatar
zhuwenwen committed
891
            use_nn_moe=args.nn_moe,
892
        )
893
894
895
        end = time.time()
        print(f"Tuning took {end - start:.2f} seconds")
    else:
896
        outputs = _distribute(
897
            "benchmark",
898
899
900
901
902
903
904
905
906
907
908
909
            [
                (
                    batch_size,
                    E,
                    shard_intermediate_size,
                    hidden_size,
                    topk,
                    dtype,
                    use_fp8_w8a8,
                    use_int8_w8a16,
                    block_quant_shape,
                    use_deep_gemm,
zhuwenwen's avatar
zhuwenwen committed
910
                    args.nn_moe,
911
912
913
914
                )
                for batch_size in batch_sizes
            ],
        )
915
916
917
918
919
920
921

        for batch_size, (config, kernel_time) in zip(batch_sizes, outputs):
            print(f"Batch size: {batch_size}, config: {config}")
            print(f"Kernel time: {kernel_time:.2f} us")


if __name__ == "__main__":
922
    parser = FlexibleArgumentParser()
923
924
925
926
927
928
    parser.add_argument(
        "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1"
    )
    parser.add_argument(
        "--tp-size", "-tp", "--tensor-parallel-size", type=int, default=2
    )
929
    parser.add_argument("--enable-expert-parallel", "-enable-ep", action="store_true")
930
931
932
    parser.add_argument(
        "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
    )
933
    parser.add_argument("--use-deep-gemm", action="store_true")
934
935
936
    parser.add_argument(
        "--save-dir", type=str, default="./", help="Directory to save tuned results"
    )
937
    parser.add_argument("--seed", type=int, default=0)
938
    parser.add_argument("--batch-size", type=int, nargs="+", required=False)
939
    parser.add_argument("--tune", action="store_true")
王敏's avatar
王敏 committed
940
    parser.add_argument("--nn-moe", action='store_true', default=False)
941
    parser.add_argument("--trust-remote-code", action="store_true")
942
    parser.add_argument("--model-prefix", type=str, required=False)
王敏's avatar
王敏 committed
943
    parser.add_argument("--num-gpus", type=int, default=1)
944
945
946
    args = parser.parse_args()

    main(args)