| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 9 | from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig |
| 10 | |
| 11 | from .common import Config |
| 12 | from .mk_objects import ( |
| 13 | MK_ALL_PREPARE_FINALIZE_TYPES, |
| 14 | MK_FUSED_EXPERT_TYPES, |
| 15 | MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, |
| 16 | ) |
| 17 | |
| 18 | |
| 19 | def make_config_arg_parser(description: str): |
| 20 | def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalizeModular: |
| 21 | for pf in MK_ALL_PREPARE_FINALIZE_TYPES: |
| 22 | if pf.__name__ == s: |
| 23 | return pf |
| 24 | raise ValueError(f"Cannot find a PrepareFinalize type that matches {s}") |
| 25 | |
| 26 | def to_experts_class_type(s: str) -> mk.FusedMoEExpertsModular: |
| 27 | for fe in MK_FUSED_EXPERT_TYPES: |
| 28 | if fe.__name__ == s: |
| 29 | return fe |
| 30 | raise ValueError(f"Cannot find a FusedExperts type that matches {s}") |
| 31 | |
| 32 | def to_quant_torch_dtype(s: str) -> torch.dtype: |
| 33 | if s == "torch.float8_e4m3fn": |
| 34 | return torch.float8_e4m3fn |
| 35 | raise ValueError(f"Unsupported quant type {s}") |
| 36 | |
| 37 | parser = argparse.ArgumentParser(description=description) |
| 38 | |
| 39 | parser.add_argument( |
| 40 | "--world-size", |
| 41 | type=int, |
| 42 | default=2, |
| 43 | help="Number of ranks that participate in all2all", |
| 44 | ) |
| 45 | parser.add_argument( |
| 46 | "--pf-type", |
| 47 | type=to_pf_class_type, |
| 48 | required=True, |
| 49 | help=( |
| 50 | "Choose a PrepareFinalize Type : " |
| 51 | f"{[x.__name__ for x in MK_ALL_PREPARE_FINALIZE_TYPES]}" |
| 52 | ), |
| 53 | ) |
| 54 | parser.add_argument( |
| 55 | "--experts-type", |
| 56 | type=to_experts_class_type, |
| 57 | required=True, |
| 58 | help=( |
| 59 | f"Choose a FusedExpert type : {[x.__name__ for x in MK_FUSED_EXPERT_TYPES]}" |
| 60 | ), |
| 61 | ) |
| 62 | parser.add_argument( |
| 63 | "-m", |
| 64 | nargs="+", |
| 65 | type=int, |
| 66 | default=[64], |
| 67 | help="num tokens per rank", |
| 68 | ) |
| 69 | parser.add_argument( |
| 70 | "-k", |
| 71 | type=int, |
| 72 | default=7168, |
| 73 | help="hidden-size", |
| 74 | ) |
| 75 | parser.add_argument( |
| 76 | "-n", |
| 77 | type=int, |
| 78 | default=1024, |
| 79 | help="N dimension of the first fused-moe matmul", |
| 80 | ) |
| 81 | parser.add_argument( |
| 82 | "--num-experts", type=int, default=32, help="Global num experts" |
| 83 | ) |
| 84 | parser.add_argument("--topk", nargs="+", type=int, default=[4, 1], help="num topk") |
| 85 | |
| 86 | # Quant args |
| 87 | parser.add_argument( |
| 88 | "--quant-dtype", type=to_quant_torch_dtype, help="Quant datatype" |
| 89 | ) |
| 90 | parser.add_argument( |
| 91 | "--per-token-quantized-activations", |
| 92 | action="store_true", |
| 93 | help=("The input activations must be per-token quantized"), |
| 94 | ) |
| 95 | parser.add_argument( |
| 96 | "--per-channel-quantized-weights", |
| 97 | action="store_true", |
| 98 | help="The weights must be per-channel quantized.", |
| 99 | ) |
| 100 | parser.add_argument( |
| 101 | "--block-shape", nargs="+", type=int, help="Quantization block shape" |
| 102 | ) |
| 103 | |
| 104 | # Torch trace profile generation args |
| 105 | parser.add_argument( |
| 106 | "--torch-trace-dir-path", |
| 107 | type=str, |
| 108 | default=None, |
| 109 | help="Get torch trace for single execution", |
| 110 | ) |
| 111 | |
| 112 | return parser |
| 113 | |
| 114 | |
| 115 | def _validate_args(args: argparse.Namespace): |
| 116 | if args.quant_dtype is not None: |
| 117 | assert args.quant_dtype == torch.float8_e4m3fn |
| 118 | if args.block_shape is not None: |
| 119 | assert len(args.block_shape) == 2, ( |
| 120 | f"block shape must have 2 elements. got {args.block_shape}" |
| 121 | ) |
| 122 | |
| 123 | if args.experts_type in MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: |
| 124 | assert args.world_size == 1, "Single GPU objects need world size set to 1" |
| 125 | |
| 126 | if args.torch_trace_dir_path is not None: |
| 127 | from pathlib import Path |
| 128 | |
| 129 | assert Path(args.torch_trace_dir_path).is_dir(), ( |
| 130 | f"Please create {args.torch_trace_dir_path}" |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | def make_config(args: argparse.Namespace) -> Config: |
| 135 | _validate_args(args) |
| 136 | |
| 137 | quant_config = None |
| 138 | if args.quant_dtype is not None: |
| 139 | quant_config = FusedMoEQuantConfig.make( |
| 140 | quant_dtype=args.quant_dtype, |
| 141 | per_act_token_quant=args.per_token_quantized_activations, |
| 142 | per_out_ch_quant=args.per_channel_quantized_weights, |
| 143 | block_shape=args.block_shape, |
| 144 | ) |
| 145 | |
| 146 | return Config( |
| 147 | Ms=args.m, |
| 148 | K=args.k, |
| 149 | N=args.n, |
| 150 | E=args.num_experts, |
| 151 | topks=args.topk, |
| 152 | dtype=torch.bfloat16, # hard-code |
| 153 | quant_config=quant_config, |
| 154 | prepare_finalize_type=args.pf_type, |
| 155 | fused_experts_type=args.experts_type, |
| 156 | world_size=args.world_size, |
| 157 | torch_trace_dir_path=args.torch_trace_dir_path, |
| 158 | ) |