| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # Copyright (c) 2025 FlyDSL Project Contributors |
| 4 | |
| 5 | import json |
| 6 | import os |
| 7 | |
| 8 | import torch |
| 9 | from aiter.test_common import run_perftest |
| 10 | |
| 11 | from vllm.model_executor.layers.fused_moe import fused_experts |
| 12 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 13 | from vllm.model_executor.layers.fused_moe.config import ( |
| 14 | int4_w4a16_moe_quant_config, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_moe |
| 17 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 |
| 18 | compressed_tensors_moe_w4a16_flydsl, |
| 19 | ) |
| 20 | from vllm.utils.platform_utils import get_device_name_as_file_name |
| 21 | |
| 22 | RoutingBuffers = tuple[ |
| 23 | torch.Tensor, # sorted_token_ids |
| 24 | torch.Tensor, # sorted_weights |
| 25 | torch.Tensor, # sorted_expert_ids |
| 26 | torch.Tensor, # num_valid_ids (shape [1], i32) |
| 27 | int, # sorted_size |
| 28 | int, # blocks |
| 29 | ] |
| 30 | |
| 31 | MODEL_PARAMS_TO_TUNE = [ |
| 32 | # (num_experts, inter_dim, hidden_size, topk) |
| 33 | (384, 256, 7168, 8), # Kimi K2.5 TP=8 |
| 34 | (384, 512, 7168, 8), # Kimi K2.5 TP=4 |
| 35 | ] |
| 36 | |
| 37 | NUM_TOKENS_TO_TUNE = [ |
| 38 | 1, |
| 39 | 2, |
| 40 | 4, |
| 41 | 8, |
| 42 | 16, |
| 43 | 24, |
| 44 | 32, |
| 45 | 48, |
| 46 | 64, |
| 47 | 128, |
| 48 | 256, |
| 49 | 512, |
| 50 | 1024, |
| 51 | 2048, |
| 52 | 4096, |
| 53 | 8192, |
| 54 | ] |
| 55 | |
| 56 | TILE_M_SEARCH_SPACE = [16, 32, 64, 128, 256] |
| 57 | TILE_N_SEARCH_SPACE = [16, 32, 64, 128, 256] |
| 58 | TILE_K_SEARCH_SPACE = [16, 32, 64, 128, 256, 512] |
| 59 | TILE_N2_SEARCH_SPACE = [16, 32, 64, 128, 256] |
| 60 | TILE_K2_SEARCH_SPACE = [16, 32, 64, 128, 256, 512] |
| 61 | |
| 62 | TILE_CONFIGS = [] |
| 63 | for tile_m in TILE_M_SEARCH_SPACE: |
| 64 | for tile_n in TILE_N_SEARCH_SPACE: |
| 65 | for tile_k in TILE_K_SEARCH_SPACE: |
| 66 | for tile_n2 in TILE_N2_SEARCH_SPACE: |
| 67 | for tile_k2 in TILE_K2_SEARCH_SPACE: |
| 68 | TILE_CONFIGS.append( |
| 69 | { |
| 70 | "tile_m": tile_m, |
| 71 | "tile_n": tile_n, |
| 72 | "tile_k": tile_k, |
| 73 | "tile_n2": tile_n2, |
| 74 | "tile_k2": tile_k2, |
| 75 | } |
| 76 | ) |
| 77 | |
| 78 | |
| 79 | def tune_flydsl_moe_w4a16( |
| 80 | device: str = "cuda", num_iters: int = 100, num_warmup: int = 10 |
| 81 | ): |
| 82 | packed_factor = 8 |
| 83 | w13_num_shards = 2 |
| 84 | params_dtype = torch.bfloat16 |
| 85 | group_size = 32 |
| 86 | scale_factor = 0.01 |
| 87 | |
| 88 | for model_params in MODEL_PARAMS_TO_TUNE: |
| 89 | num_experts = model_params[0] |
| 90 | inter_dim = model_params[1] |
| 91 | hidden_size = model_params[2] |
| 92 | topk = model_params[3] |
| 93 | print( |
| 94 | f"\nTuning: num_experts={num_experts}, inter_dim={inter_dim}, " |
| 95 | f"hidden_size={hidden_size}, topk={topk}...\n" |
| 96 | ) |
| 97 | |
| 98 | w2_scales_size = inter_dim |
| 99 | num_groups_w2 = w2_scales_size // group_size |
| 100 | num_groups_w13 = hidden_size // group_size |
| 101 | |
| 102 | w13_weight = torch.randint( |
| 103 | 0, |
| 104 | 255, |
| 105 | (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim), |
| 106 | dtype=torch.int32, |
| 107 | device=device, |
| 108 | ) |
| 109 | |
| 110 | w2_weight = torch.randint( |
| 111 | 0, |
| 112 | 255, |
| 113 | (num_experts, inter_dim // packed_factor, hidden_size), |
| 114 | dtype=torch.int32, |
| 115 | device=device, |
| 116 | ) |
| 117 | w13_scale = scale_factor * torch.randn( |
| 118 | num_experts, |
| 119 | num_groups_w13, |
| 120 | w13_num_shards * inter_dim, |
| 121 | dtype=params_dtype, |
| 122 | device=device, |
| 123 | ) |
| 124 | w2_scale = scale_factor * torch.randn( |
| 125 | num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device |
| 126 | ) |
| 127 | |
| 128 | w13 = w13_weight |
| 129 | w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13) |
| 130 | w13 = w13.view(-1).contiguous() |
| 131 | |
| 132 | w2 = w2_weight |
| 133 | w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2) |
| 134 | w2 = w2.view(-1).contiguous() |
| 135 | |
| 136 | w13_scale_flydsl = w13_scale |
| 137 | w2_scale_flydsl = w2_scale |
| 138 | |
| 139 | if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1: |
| 140 | E, G, N = w13_scale.shape |
| 141 | w13_scale_flydsl = ( |
| 142 | w13_scale_flydsl.view(E, G // 2, 2, N) |
| 143 | .permute(0, 1, 3, 2) |
| 144 | .contiguous() |
| 145 | .view(-1) |
| 146 | .contiguous() |
| 147 | ) |
| 148 | elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1: |
| 149 | w13_scale_flydsl = w13_scale_flydsl.squeeze(1) |
| 150 | |
| 151 | if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1: |
| 152 | E, G, N = w2_scale.shape |
| 153 | w2_scale_flydsl = ( |
| 154 | w2_scale_flydsl.view(E, G // 2, 2, N) |
| 155 | .permute(0, 1, 3, 2) |
| 156 | .contiguous() |
| 157 | .view(-1) |
| 158 | .contiguous() |
| 159 | ) |
| 160 | elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1: |
| 161 | w2_scale_flydsl = w2_scale_flydsl.squeeze(1) |
| 162 | |
| 163 | w13_scale_flydsl = w13_scale_flydsl.contiguous() |
| 164 | w2_scale_flydsl = w2_scale_flydsl.contiguous() |
| 165 | |
| 166 | w13.is_shuffled = True |
| 167 | w2.is_shuffled = True |
| 168 | |
| 169 | w13_weight_scale = w13_scale.transpose(1, 2).contiguous() |
| 170 | w2_weight_scale = w2_scale.transpose(1, 2).contiguous() |
| 171 | w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8) |
| 172 | w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8) |
| 173 | |
| 174 | moe_quant_config = int4_w4a16_moe_quant_config( |
| 175 | w1_scale=w13_weight_scale, |
| 176 | w2_scale=w2_weight_scale, |
| 177 | w1_zp=None, |
| 178 | w2_zp=None, |
| 179 | block_shape=[0, group_size], |
| 180 | ) |
| 181 | |
| 182 | tuned_config = {} |
| 183 | |
| 184 | for num_tokens in NUM_TOKENS_TO_TUNE: |
| 185 | score = torch.rand( |
| 186 | (num_tokens, num_experts), device=device, dtype=torch.float32 |
| 187 | ) |
| 188 | topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) |
| 189 | topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) |
| 190 | x = torch.randn( |
| 191 | (num_tokens, hidden_size), dtype=torch.bfloat16, device=device |
| 192 | ) |
| 193 | us_best = float("inf") |
| 194 | for tile_config in TILE_CONFIGS: |
| 195 | try: |
| 196 | tile_m = tile_config["tile_m"] |
| 197 | tile_n = tile_config["tile_n"] |
| 198 | tile_k = tile_config["tile_k"] |
| 199 | tile_n2 = tile_config["tile_n2"] |
| 200 | tile_k2 = tile_config["tile_k2"] |
| 201 | |
| 202 | model_dim = x.shape[1] |
| 203 | assert model_dim % 64 == 0 |
| 204 | assert model_dim % tile_k == 0 |
| 205 | assert inter_dim % tile_n == 0 |
| 206 | assert model_dim % tile_n2 == 0 |
| 207 | assert inter_dim % tile_k2 == 0 |
| 208 | assert ((tile_m * tile_k2) % 256) == 0 |
| 209 | bytes_per_thread_x = (tile_m * tile_k2) // 256 |
| 210 | assert (bytes_per_thread_x % 4) == 0 |
| 211 | |
| 212 | out, _us = run_perftest( |
| 213 | fused_flydsl_moe, |
| 214 | x, |
| 215 | w13, |
| 216 | w2, |
| 217 | num_experts, |
| 218 | inter_dim, |
| 219 | topk_weights, |
| 220 | topk_ids, |
| 221 | num_iters=num_iters, |
| 222 | num_warmup=num_warmup, |
| 223 | w1_scale=w13_scale_flydsl, |
| 224 | w2_scale=w2_scale_flydsl, |
| 225 | topk=topk_weights.shape[-1], |
| 226 | group_size=group_size, |
| 227 | doweight_stage1=False, |
| 228 | scale_is_bf16=True, |
| 229 | config=tile_config, |
| 230 | ) |
| 231 | torch.accelerator.synchronize() |
| 232 | except Exception: |
| 233 | torch.accelerator.synchronize() |
| 234 | continue |
| 235 | else: |
| 236 | us = _us.item() |
| 237 | if us < us_best: |
| 238 | out_ref = fused_experts( |
| 239 | x, |
| 240 | w13_weight_packed, |
| 241 | w2_weight_packed, |
| 242 | topk_weights=topk_weights, |
| 243 | topk_ids=topk_ids, |
| 244 | activation=MoEActivation.SILU, |
| 245 | apply_router_weight_on_input=False, |
| 246 | global_num_experts=num_experts, |
| 247 | expert_map=None, |
| 248 | quant_config=moe_quant_config, |
| 249 | ) |
| 250 | try: |
| 251 | assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1) |
| 252 | except Exception: |
| 253 | continue |
| 254 | else: |
| 255 | print( |
| 256 | f"For [num_tokens={num_tokens}, num_experts={num_experts}, " # noqa: E501 |
| 257 | f"inter_dim={inter_dim}] found new best " # noqa: E501 |
| 258 | f"config={tile_config}, us={us:0.3f}" |
| 259 | ) |
| 260 | us_best = us |
| 261 | tuned_config[str(num_tokens)] = tile_config |
| 262 | device_name = get_device_name_as_file_name() |
| 263 | tuned_config_file_name = ( |
| 264 | f"E={num_experts},N={inter_dim},device_name={device_name}," |
| 265 | f"dtype=int4_w4a16,backend=flydsl.json" |
| 266 | ) |
| 267 | tuner_dir_path = os.path.dirname(os.path.realpath(__file__)) |
| 268 | store_path = os.path.join(tuner_dir_path, tuned_config_file_name) |
| 269 | with open(store_path, "w") as f: |
| 270 | json.dump(tuned_config, f, indent=4) |
| 271 | print( |
| 272 | f"\nTuned config for num_tokens={num_tokens} was stored at {store_path}\n" # noqa: E501 |
| 273 | ) |
| 274 | |
| 275 | |
| 276 | if __name__ == "__main__": |
| 277 | tune_flydsl_moe_w4a16(device="cuda") |