| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fused MoE Triton kernels.""" |
| 4 | |
| 5 | import functools |
| 6 | import json |
| 7 | import os |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.utils.platform_utils import get_device_name_as_file_name |
| 13 | from vllm.utils.torch_utils import direct_register_custom_op |
| 14 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | # Valid tiles for the aiter a16w-mix int4 kernels (aiter#4646): stage1 requires |
| 18 | # tile_n in {64,128} and tile_k in {128,256} (k_wave=1); stage2 requires tile_n2==128 |
| 19 | # and tile_k2 in {128,256}. The pre-#4646 defaults (tile_k=512/64, tile_n2=256) are no |
| 20 | # longer registered and make the port emit an empty-fragment kernel. |
| 21 | _FLYDSL_MOE_DEFAULT_CONFIG = { |
| 22 | 1: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 23 | 2: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 24 | 4: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 25 | 8: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 26 | 16: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 27 | 24: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 28 | 32: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 29 | 48: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 30 | 64: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 31 | 128: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 32 | 256: {"tile_m": 16, "tile_n": 128, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 33 | 512: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 34 | 1024: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 35 | 2048: {"tile_m": 64, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 36 | 4096: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 37 | 8192: {"tile_m": 64, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, |
| 38 | } |
| 39 | |
| 40 | # aiter#4646's a16w-mix int4 kernels only register a fixed tile set. Tiles from before |
| 41 | # the bump (default config or tuned JSONs) carry tile_k=512/64 and tile_n2=256, which |
| 42 | # emit an empty-fragment kernel (IndexError at emit / wrong output). Snap any resolved |
| 43 | # tiles to the nearest registered value so a stale tuned config can't crash the kernel. |
| 44 | _FLYDSL_INT4_TILE_M = (16, 32, 64, 128) |
| 45 | _FLYDSL_INT4_TILE_N1 = (64, 128) |
| 46 | _FLYDSL_INT4_TILE_K = (128, 256) |
| 47 | |
| 48 | |
| 49 | def _snap_flydsl_int4_tiles(tile_m, tile_n, tile_k, tile_n2, tile_k2): |
| 50 | def _nearest(v, allowed): |
| 51 | return min(allowed, key=lambda a: (abs(a - v), a)) |
| 52 | |
| 53 | snapped = ( |
| 54 | _nearest(tile_m, _FLYDSL_INT4_TILE_M), |
| 55 | _nearest(tile_n, _FLYDSL_INT4_TILE_N1), |
| 56 | _nearest(tile_k, _FLYDSL_INT4_TILE_K), |
| 57 | 128, # stage2 N tile is only registered at 128 |
| 58 | _nearest(tile_k2, _FLYDSL_INT4_TILE_K), |
| 59 | ) |
| 60 | if snapped != (tile_m, tile_n, tile_k, tile_n2, tile_k2): |
| 61 | logger.warning_once( |
| 62 | "FlyDSL int4 MoE tiles %s are not registered post-aiter#4646; " |
| 63 | "using %s. Retune the config for the a16w-mix kernel.", |
| 64 | (tile_m, tile_n, tile_k, tile_n2, tile_k2), |
| 65 | snapped, |
| 66 | scope="local", |
| 67 | ) |
| 68 | return snapped |
| 69 | |
| 70 | |
| 71 | def moe_sorting( |
| 72 | topk_ids: torch.Tensor, |
| 73 | topk_weights: torch.Tensor, |
| 74 | *, |
| 75 | num_experts: int, |
| 76 | model_dim: int, |
| 77 | block_m: int, |
| 78 | ): |
| 79 | # Imported lazily so aiter's import-time sort-backend env vars are read |
| 80 | # when this is called, not whenever this module happens to be imported. |
| 81 | from aiter.fused_moe import moe_sorting as aiter_moe_sorting |
| 82 | |
| 83 | topk_ids_i32 = topk_ids.to(torch.int32) |
| 84 | topk_w_f32 = topk_weights.to(torch.float32) |
| 85 | sorted_ids, sorted_w, sorted_expert_ids, num_valid_ids, _moe_buf = ( |
| 86 | aiter_moe_sorting( |
| 87 | topk_ids_i32, |
| 88 | topk_w_f32, |
| 89 | num_experts, |
| 90 | model_dim, |
| 91 | torch.float16, |
| 92 | block_m, |
| 93 | ) |
| 94 | ) |
| 95 | if num_valid_ids.numel() > 1: |
| 96 | num_valid_ids = num_valid_ids[:1].contiguous() |
| 97 | return sorted_ids, sorted_w, sorted_expert_ids, num_valid_ids |
| 98 | |
| 99 | |
| 100 | def build_routing_buffers( |
| 101 | *, |
| 102 | topk_ids: torch.Tensor, |
| 103 | topk_weights: torch.Tensor, |
| 104 | num_experts: int, |
| 105 | model_dim: int, |
| 106 | tile_m: int, |
| 107 | ): |
| 108 | res = moe_sorting( |
| 109 | topk_ids, |
| 110 | topk_weights, |
| 111 | num_experts=num_experts, |
| 112 | model_dim=model_dim, |
| 113 | block_m=tile_m, |
| 114 | ) |
| 115 | if res is None: |
| 116 | raise RuntimeError( |
| 117 | "aiter moe_sorting failed/unavailable; cannot build routing buffers." |
| 118 | ) |
| 119 | sorted_token_ids, sorted_weights, sorted_expert_ids, num_valid_ids = res |
| 120 | |
| 121 | sorted_token_ids = sorted_token_ids.contiguous() |
| 122 | sorted_weights = sorted_weights.contiguous() |
| 123 | sorted_expert_ids = sorted_expert_ids.contiguous() |
| 124 | sorted_size = int(sorted_token_ids.numel()) |
| 125 | blocks = int(sorted_expert_ids.numel()) |
| 126 | return ( |
| 127 | sorted_token_ids, |
| 128 | sorted_weights, |
| 129 | sorted_expert_ids, |
| 130 | num_valid_ids, |
| 131 | sorted_size, |
| 132 | blocks, |
| 133 | ) |
| 134 | |
| 135 | |
| 136 | @functools.lru_cache |
| 137 | def try_get_optimal_config(num_experts, inter_dim): |
| 138 | device_name = get_device_name_as_file_name() |
| 139 | json_file_name = ( |
| 140 | f"E={num_experts},N={inter_dim},device_name={device_name}," |
| 141 | "dtype=int4_w4a16,backend=flydsl.json" |
| 142 | ) |
| 143 | config_file_path = os.path.join( |
| 144 | os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name |
| 145 | ) |
| 146 | if os.path.exists(config_file_path): |
| 147 | with open(config_file_path) as f: |
| 148 | logger.info_once( |
| 149 | "Using tuned FlyDSL MoE config from %s", |
| 150 | config_file_path, |
| 151 | scope="global", |
| 152 | ) |
| 153 | tuned_config = json.load(f) |
| 154 | return {int(key): val for key, val in tuned_config.items()} |
| 155 | |
| 156 | logger.warning_once( |
| 157 | "Using default FlyDSL MoE config. Performance might be sub-optimal! " |
| 158 | "Config file not found at %s", |
| 159 | config_file_path, |
| 160 | scope="local", |
| 161 | ) |
| 162 | return _FLYDSL_MOE_DEFAULT_CONFIG |
| 163 | |
| 164 | |
| 165 | def fused_flydsl_moe_impl( |
| 166 | hidden_states: torch.Tensor, |
| 167 | w1: torch.Tensor, |
| 168 | w2: torch.Tensor, |
| 169 | num_experts: int, |
| 170 | inter_dim: int, |
| 171 | topk_weights: torch.Tensor, |
| 172 | topk_ids: torch.Tensor, |
| 173 | w1_scale: torch.Tensor | None = None, |
| 174 | w2_scale: torch.Tensor | None = None, |
| 175 | topk: int = 8, |
| 176 | group_size: int = 32, |
| 177 | doweight_stage1: bool = False, |
| 178 | in_dtype: str = "int4_bf16", |
| 179 | out_dtype: str = "bf16", |
| 180 | scale_is_bf16: bool = True, |
| 181 | tile_m: int | None = None, |
| 182 | tile_n: int | None = None, |
| 183 | tile_k: int | None = None, |
| 184 | tile_n2: int | None = None, |
| 185 | tile_k2: int | None = None, |
| 186 | ) -> torch.Tensor: |
| 187 | # Imported lazily -- see the comment in moe_sorting() above. |
| 188 | from aiter.ops.flydsl.kernels.moe_2stage_a16wmix import ( |
| 189 | flydsl_a16w4_gemm1, |
| 190 | flydsl_a16w4_gemm2, |
| 191 | ) |
| 192 | |
| 193 | device = hidden_states.device |
| 194 | tokens = hidden_states.shape[0] |
| 195 | model_dim = hidden_states.shape[1] |
| 196 | |
| 197 | tuned_config = {} |
| 198 | if tile_m and tile_n and tile_k and tile_n2 and tile_k2: |
| 199 | tuned_config["tile_m"] = tile_m |
| 200 | tuned_config["tile_n"] = tile_n |
| 201 | tuned_config["tile_k"] = tile_k |
| 202 | tuned_config["tile_n2"] = tile_n2 |
| 203 | tuned_config["tile_k2"] = tile_k2 |
| 204 | else: |
| 205 | tuned_config = try_get_optimal_config(num_experts, inter_dim) |
| 206 | tuned_config = tuned_config[ |
| 207 | min(tuned_config.keys(), key=lambda x: abs(x - tokens)) |
| 208 | ] |
| 209 | out_torch_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 |
| 210 | |
| 211 | tile_m = tuned_config["tile_m"] |
| 212 | tile_n = tuned_config["tile_n"] |
| 213 | tile_k = tuned_config["tile_k"] |
| 214 | tile_n2 = tuned_config["tile_n2"] |
| 215 | tile_k2 = tuned_config["tile_k2"] |
| 216 | tile_m, tile_n, tile_k, tile_n2, tile_k2 = _snap_flydsl_int4_tiles( |
| 217 | tile_m, tile_n, tile_k, tile_n2, tile_k2 |
| 218 | ) |
| 219 | |
| 220 | routing = build_routing_buffers( |
| 221 | topk_ids=topk_ids, |
| 222 | topk_weights=topk_weights, |
| 223 | num_experts=num_experts, |
| 224 | model_dim=model_dim, |
| 225 | tile_m=tile_m, |
| 226 | ) |
| 227 | ( |
| 228 | sorted_token_ids, |
| 229 | sorted_weights, |
| 230 | sorted_expert_ids, |
| 231 | num_valid_ids, |
| 232 | sorted_size, |
| 233 | blocks, |
| 234 | ) = routing |
| 235 | |
| 236 | # aiter#4646 (v0.1.20) removed compile_moe_gemm{1,2} + the flyc.compile pointer ABI |
| 237 | # and replaced them with the a16w-mix wrappers, which compile+launch internally |
| 238 | # (functools-cached) and take torch tensors directly. w_dtype="int4" is the a16wi4 |
| 239 | # (W4A16) path. The stage-1 intermediate now lives in sorted |
| 240 | # [sorted_size, inter_dim] layout, and topk weighting is applied in stage-2 via |
| 241 | # sorted_weights (the port has no stage-1 doweight). |
| 242 | sorted_weights_1d = sorted_weights.view(-1).contiguous() |
| 243 | empty_scale = torch.empty((0,), device=device, dtype=torch.uint8) |
| 244 | cumsum_tensor = num_valid_ids.to(torch.int32).contiguous() |
| 245 | m_indices = sorted_token_ids.to(torch.int32).contiguous() |
| 246 | |
| 247 | # stage1: bf16 A x int4 W1 + SiLU -> bf16 sorted intermediate. |
| 248 | inter_sorted = torch.empty( |
| 249 | (sorted_size, inter_dim), device=device, dtype=torch.bfloat16 |
| 250 | ) |
| 251 | w1_scale_u8 = ( |
| 252 | w1_scale.view(torch.uint8).contiguous().view(-1) |
| 253 | if w1_scale is not None |
| 254 | else empty_scale |
| 255 | ) |
| 256 | flydsl_a16w4_gemm1( |
| 257 | a_bf16=hidden_states.to(torch.bfloat16).contiguous(), |
| 258 | w1_u8=w1.view(torch.uint8).contiguous(), |
| 259 | w1_scale_u8=w1_scale_u8, |
| 260 | sorted_expert_ids=sorted_expert_ids, |
| 261 | cumsum_tensor=cumsum_tensor, |
| 262 | m_indices=m_indices, |
| 263 | inter_sorted_bf16=inter_sorted, |
| 264 | n_tokens=tokens, |
| 265 | NE=num_experts, |
| 266 | D_HIDDEN=model_dim, |
| 267 | D_INTER=inter_dim, |
| 268 | topk=topk, |
| 269 | tile_m=tile_m, |
| 270 | tile_n=tile_n, |
| 271 | tile_k=tile_k, |
| 272 | waves_per_eu=None, |
| 273 | act="silu", |
| 274 | w_dtype="int4", |
| 275 | w_layout="standard", |
| 276 | ) |
| 277 | |
| 278 | # stage2: bf16 inter x int4 W2 -> weighted reduce into [tokens, model_dim]. |
| 279 | out_stage2 = torch.zeros((tokens, model_dim), device=device, dtype=out_torch_dtype) |
| 280 | w2_scale_u8 = ( |
| 281 | w2_scale.view(torch.uint8).contiguous().view(-1) |
| 282 | if w2_scale is not None |
| 283 | else empty_scale |
| 284 | ) |
| 285 | flydsl_a16w4_gemm2( |
| 286 | inter_sorted_bf16=inter_sorted, |
| 287 | w2_u8=w2.view(torch.uint8).contiguous(), |
| 288 | w2_scale_u8=w2_scale_u8, |
| 289 | sorted_expert_ids=sorted_expert_ids, |
| 290 | cumsum_tensor=cumsum_tensor, |
| 291 | sorted_token_ids=sorted_token_ids, |
| 292 | sorted_weights=sorted_weights_1d, |
| 293 | flat_out=out_stage2.view(-1), |
| 294 | M_logical=tokens, |
| 295 | max_sorted=sorted_size, |
| 296 | NE=num_experts, |
| 297 | D_HIDDEN=model_dim, |
| 298 | D_INTER=inter_dim, |
| 299 | topk=topk, |
| 300 | tile_m=tile_m, |
| 301 | tile_n=tile_n2, |
| 302 | tile_k=tile_k2, |
| 303 | waves_per_eu=None, |
| 304 | w_dtype="int4", |
| 305 | ) |
| 306 | return out_stage2 |
| 307 | |
| 308 | |
| 309 | def fused_flydsl_moe_impl_fake( |
| 310 | hidden_states: torch.Tensor, |
| 311 | w1: torch.Tensor, |
| 312 | w2: torch.Tensor, |
| 313 | num_experts: int, |
| 314 | inter_dim: int, |
| 315 | topk_weights: torch.Tensor, |
| 316 | topk_ids: torch.Tensor, |
| 317 | w1_scale: torch.Tensor | None = None, |
| 318 | w2_scale: torch.Tensor | None = None, |
| 319 | topk: int = 8, |
| 320 | group_size: int = 32, |
| 321 | doweight_stage1: bool = False, |
| 322 | in_dtype: str = "int4_bf16", |
| 323 | out_dtype: str = "bf16", |
| 324 | scale_is_bf16: bool = True, |
| 325 | tile_m: int | None = None, |
| 326 | tile_n: int | None = None, |
| 327 | tile_k: int | None = None, |
| 328 | tile_n2: int | None = None, |
| 329 | tile_k2: int | None = None, |
| 330 | ) -> torch.Tensor: |
| 331 | return torch.empty_like(hidden_states) |
| 332 | |
| 333 | |
| 334 | direct_register_custom_op( |
| 335 | op_name="fused_flydsl_moe_impl", |
| 336 | op_func=fused_flydsl_moe_impl, |
| 337 | fake_impl=fused_flydsl_moe_impl_fake, |
| 338 | ) |
| 339 | |
| 340 | |
| 341 | def fused_flydsl_moe( |
| 342 | hidden_states: torch.Tensor, |
| 343 | w1: torch.Tensor, |
| 344 | w2: torch.Tensor, |
| 345 | num_experts: int, |
| 346 | inter_dim: int, |
| 347 | topk_weights: torch.Tensor, |
| 348 | topk_ids: torch.Tensor, |
| 349 | w1_scale: torch.Tensor | None = None, |
| 350 | w2_scale: torch.Tensor | None = None, |
| 351 | topk: int = 8, |
| 352 | group_size: int = 32, |
| 353 | doweight_stage1: bool = False, |
| 354 | in_dtype: str = "int4_bf16", |
| 355 | out_dtype: str = "bf16", |
| 356 | scale_is_bf16: bool = True, |
| 357 | config: dict | None = None, |
| 358 | ) -> torch.Tensor: |
| 359 | tile_m = None |
| 360 | tile_n = None |
| 361 | tile_k = None |
| 362 | tile_n2 = None |
| 363 | tile_k2 = None |
| 364 | if config is not None: |
| 365 | tile_m = config.get("tile_m") |
| 366 | tile_n = config.get("tile_n") |
| 367 | tile_k = config.get("tile_k") |
| 368 | tile_n2 = config.get("tile_n2") |
| 369 | tile_k2 = config.get("tile_k2") |
| 370 | return torch.ops.vllm.fused_flydsl_moe_impl( |
| 371 | hidden_states=hidden_states, |
| 372 | w1=w1, |
| 373 | w2=w2, |
| 374 | num_experts=num_experts, |
| 375 | inter_dim=inter_dim, |
| 376 | topk_weights=topk_weights, |
| 377 | topk_ids=topk_ids, |
| 378 | w1_scale=w1_scale, |
| 379 | w2_scale=w2_scale, |
| 380 | topk=topk, |
| 381 | group_size=group_size, |
| 382 | doweight_stage1=doweight_stage1, |
| 383 | in_dtype=in_dtype, |
| 384 | out_dtype=out_dtype, |
| 385 | scale_is_bf16=scale_is_bf16, |
| 386 | tile_m=tile_m, |
| 387 | tile_n=tile_n, |
| 388 | tile_k=tile_k, |
| 389 | tile_n2=tile_n2, |
| 390 | tile_k2=tile_k2, |
| 391 | ) |