| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | import vllm._custom_ops as ops |
| 8 | from tests.kernels.quant_utils import per_block_cast_to_int8 |
| 9 | from tests.kernels.quantization.nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX |
| 10 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 11 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 12 | maybe_make_prepare_finalize, |
| 13 | ) |
| 14 | from vllm.model_executor.layers.fused_moe.config import ( |
| 15 | FusedMoEConfig, |
| 16 | FusedMoEParallelConfig, |
| 17 | FusedMoEQuantConfig, |
| 18 | RoutingMethodType, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import ( |
| 21 | BatchedTritonExperts, |
| 22 | NaiveBatchedExperts, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( |
| 25 | TritonExperts, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.fused_moe.fused_moe import ( |
| 28 | fused_experts, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel |
| 31 | from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput |
| 32 | from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import moe_unpermute |
| 33 | from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import ( |
| 34 | BatchedPrepareAndFinalize, |
| 35 | ) |
| 36 | from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk |
| 37 | from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input |
| 38 | from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 |
| 39 | from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( |
| 40 | ref_nvfp4_quant, |
| 41 | ) |
| 42 | from vllm.platforms import current_platform |
| 43 | from vllm.utils.deep_gemm import per_block_cast_to_fp8 |
| 44 | from vllm.utils.math_utils import round_up |
| 45 | |
| 46 | DEVICE = current_platform.device_type |
| 47 | |
| 48 | |
| 49 | def shuffle_weight(w: torch.Tensor) -> torch.Tensor: |
| 50 | """Fold weights to adjacent locations for Triton MoE / SwiGLU kernel layout.""" |
| 51 | shape = w.shape |
| 52 | n = shape[-1] |
| 53 | first = w[..., : n // 2] |
| 54 | second = w[..., n // 2 :] |
| 55 | stacked = torch.stack((first, second), dim=-1) |
| 56 | return stacked.reshape(shape) |
| 57 | |
| 58 | |
| 59 | def make_dummy_moe_config( |
| 60 | num_experts: int = 1, |
| 61 | num_local_experts: int | None = None, |
| 62 | experts_per_token: int = 1, |
| 63 | hidden_dim: int = 1, |
| 64 | intermediate_size: int = 1, |
| 65 | in_dtype: torch.dtype = torch.bfloat16, |
| 66 | max_num_tokens: int = 512, |
| 67 | activation: MoEActivation = MoEActivation.SILU, |
| 68 | ) -> FusedMoEConfig: |
| 69 | """This is a dummy config for the mk constructor interface |
| 70 | as most kernels like DeepGEMM, CUTLASSFp4, Triton, MARLIN |
| 71 | do not actually use this config. |
| 72 | |
| 73 | CUTLASSFp8 needs to set some params for workshapes. |
| 74 | """ |
| 75 | return FusedMoEConfig( |
| 76 | num_experts=num_experts, |
| 77 | experts_per_token=experts_per_token, |
| 78 | hidden_dim=hidden_dim, |
| 79 | intermediate_size=intermediate_size, |
| 80 | num_local_experts=num_local_experts |
| 81 | if num_local_experts is not None |
| 82 | else num_experts, |
| 83 | num_logical_experts=num_experts, |
| 84 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 85 | activation=activation, |
| 86 | in_dtype=in_dtype, |
| 87 | device=DEVICE, |
| 88 | routing_method=RoutingMethodType.TopK, |
| 89 | max_num_tokens=max_num_tokens, |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | def triton_moe( |
| 94 | a: torch.Tensor, |
| 95 | w1: torch.Tensor, |
| 96 | w2: torch.Tensor, |
| 97 | topk_weight: torch.Tensor, |
| 98 | topk_ids: torch.Tensor, |
| 99 | w1_scale: torch.Tensor | None = None, |
| 100 | w2_scale: torch.Tensor | None = None, |
| 101 | a1_scale: torch.Tensor | None = None, |
| 102 | a2_scale: torch.Tensor | None = None, |
| 103 | quant_dtype: torch.dtype | None = None, |
| 104 | per_act_token_quant=False, |
| 105 | block_shape: list[int] | None = None, |
| 106 | ) -> torch.Tensor: |
| 107 | quant_config = FusedMoEQuantConfig.make( |
| 108 | quant_dtype, |
| 109 | per_act_token_quant=per_act_token_quant, |
| 110 | block_shape=block_shape, |
| 111 | w1_scale=w1_scale, |
| 112 | w2_scale=w2_scale, |
| 113 | a1_scale=a1_scale, |
| 114 | a2_scale=a2_scale, |
| 115 | ) |
| 116 | |
| 117 | return fused_experts(a, w1, w2, topk_weight, topk_ids, quant_config=quant_config) |
| 118 | |
| 119 | |
| 120 | def batched_moe( |
| 121 | a: torch.Tensor, |
| 122 | w1: torch.Tensor, |
| 123 | w2: torch.Tensor, |
| 124 | topk_weight: torch.Tensor, |
| 125 | topk_ids: torch.Tensor, |
| 126 | w1_scale: torch.Tensor | None = None, |
| 127 | w2_scale: torch.Tensor | None = None, |
| 128 | a1_scale: torch.Tensor | None = None, |
| 129 | a2_scale: torch.Tensor | None = None, |
| 130 | quant_dtype: torch.dtype | None = None, |
| 131 | per_act_token_quant: bool = False, |
| 132 | block_shape: list[int] | None = None, |
| 133 | moe_config: FusedMoEConfig | None = None, |
| 134 | ) -> torch.Tensor: |
| 135 | max_num_tokens = round_up(a.shape[0], 64) |
| 136 | |
| 137 | quant_config = FusedMoEQuantConfig.make( |
| 138 | quant_dtype, |
| 139 | per_act_token_quant=per_act_token_quant, |
| 140 | block_shape=block_shape, |
| 141 | w1_scale=w1_scale, |
| 142 | w2_scale=w2_scale, |
| 143 | a1_scale=a1_scale, |
| 144 | a2_scale=a2_scale, |
| 145 | ) |
| 146 | |
| 147 | if moe_config is None: |
| 148 | moe_config = make_dummy_moe_config() |
| 149 | |
| 150 | fused_experts = FusedMoEKernel( |
| 151 | BatchedPrepareAndFinalize( |
| 152 | max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 |
| 153 | ), |
| 154 | BatchedTritonExperts( |
| 155 | max_num_tokens=max_num_tokens, |
| 156 | num_dispatchers=1, |
| 157 | quant_config=quant_config, |
| 158 | moe_config=moe_config, |
| 159 | ), |
| 160 | ) |
| 161 | |
| 162 | return fused_experts.apply( |
| 163 | a, |
| 164 | w1, |
| 165 | w2, |
| 166 | topk_weight, |
| 167 | topk_ids, |
| 168 | global_num_experts=w1.shape[0], |
| 169 | activation=moe_config.activation, |
| 170 | apply_router_weight_on_input=False, |
| 171 | expert_map=None, |
| 172 | ) |
| 173 | |
| 174 | |
| 175 | def naive_batched_moe( |
| 176 | a: torch.Tensor, |
| 177 | w1: torch.Tensor, |
| 178 | w2: torch.Tensor, |
| 179 | topk_weight: torch.Tensor, |
| 180 | topk_ids: torch.Tensor, |
| 181 | w1_scale: torch.Tensor | None = None, |
| 182 | w2_scale: torch.Tensor | None = None, |
| 183 | a1_scale: torch.Tensor | None = None, |
| 184 | a2_scale: torch.Tensor | None = None, |
| 185 | quant_dtype: torch.dtype | None = None, |
| 186 | per_act_token_quant: bool = False, |
| 187 | block_shape: list[int] | None = None, |
| 188 | ) -> torch.Tensor: |
| 189 | max_num_tokens = round_up(a.shape[0], 64) |
| 190 | |
| 191 | quant_config = FusedMoEQuantConfig.make( |
| 192 | quant_dtype, |
| 193 | per_act_token_quant=per_act_token_quant, |
| 194 | block_shape=block_shape, |
| 195 | w1_scale=w1_scale, |
| 196 | w2_scale=w2_scale, |
| 197 | a1_scale=a1_scale, |
| 198 | a2_scale=a2_scale, |
| 199 | ) |
| 200 | moe_config = make_dummy_moe_config() |
| 201 | |
| 202 | fused_experts = FusedMoEKernel( |
| 203 | BatchedPrepareAndFinalize( |
| 204 | max_num_tokens, num_dispatchers=1, num_local_experts=w1.shape[0], rank=0 |
| 205 | ), |
| 206 | NaiveBatchedExperts( |
| 207 | max_num_tokens=max_num_tokens, |
| 208 | num_dispatchers=1, |
| 209 | quant_config=quant_config, |
| 210 | moe_config=moe_config, |
| 211 | ), |
| 212 | ) |
| 213 | |
| 214 | return fused_experts.apply( |
| 215 | a, |
| 216 | w1, |
| 217 | w2, |
| 218 | topk_weight, |
| 219 | topk_ids, |
| 220 | global_num_experts=w1.shape[0], |
| 221 | activation=moe_config.activation, |
| 222 | apply_router_weight_on_input=False, |
| 223 | expert_map=None, |
| 224 | ) |
| 225 | |
| 226 | |
| 227 | def chunk_scales( |
| 228 | scales: torch.Tensor | None, start: int, end: int |
| 229 | ) -> torch.Tensor | None: |
| 230 | if scales is not None: |
| 231 | if scales.numel() == 1: |
| 232 | return scales |
| 233 | else: |
| 234 | return scales[start:end] |
| 235 | return None |
| 236 | |
| 237 | |
| 238 | def make_quantized_test_activations( |
| 239 | E: int, |
| 240 | m: int, |
| 241 | k: int, |
| 242 | in_dtype: torch.dtype, |
| 243 | quant_dtype: torch.dtype | None = None, |
| 244 | block_shape: list[int] | None = None, |
| 245 | per_act_token_quant: bool = False, |
| 246 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: |
| 247 | a = torch.randn((E, m, k), device=DEVICE, dtype=in_dtype) / 10 |
| 248 | a_q = a |
| 249 | a_scale = None |
| 250 | |
| 251 | if quant_dtype is not None: |
| 252 | assert quant_dtype == torch.float8_e4m3fn or quant_dtype == torch.int8, ( |
| 253 | "only fp8/int8 supported" |
| 254 | ) |
| 255 | a_q = torch.zeros_like(a, dtype=quant_dtype) |
| 256 | a_scale_l = [None] * E |
| 257 | for e in range(E): |
| 258 | a_q[e], a_scale_l[e] = moe_kernel_quantize_input( |
| 259 | a[e], None, quant_dtype, per_act_token_quant, block_shape |
| 260 | ) |
| 261 | a_scale = torch.stack(a_scale_l) |
| 262 | |
| 263 | if not per_act_token_quant and block_shape is None: |
| 264 | a_scale = a_scale.view(E, 1, 1) |
| 265 | |
| 266 | return a, a_q, a_scale |
| 267 | |
| 268 | |
| 269 | def moe_quantize_weights_2d( |
| 270 | w: torch.Tensor, |
| 271 | w_s: torch.Tensor | None, |
| 272 | quant_dtype: torch.dtype | str | None, |
| 273 | per_token_quant: bool, |
| 274 | block_shape: list[int] | None, |
| 275 | ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: |
| 276 | is_fp8_dtype = is_fp8(quant_dtype) |
| 277 | assert is_fp8_dtype or quant_dtype == torch.int8 or quant_dtype == "nvfp4", ( |
| 278 | "only fp8/int8/nvfp4 supported" |
| 279 | ) |
| 280 | |
| 281 | w_gs = None |
| 282 | |
| 283 | if block_shape is not None: |
| 284 | assert not per_token_quant |
| 285 | if quant_dtype == torch.int8: |
| 286 | w, w_s = per_block_cast_to_int8(w, block_shape) |
| 287 | elif is_fp8_dtype: |
| 288 | w, w_s = per_block_cast_to_fp8(w, block_shape) |
| 289 | elif quant_dtype == "nvfp4": |
| 290 | raise RuntimeError("blocked quantization not supported for nvfp4") |
| 291 | else: |
| 292 | raise RuntimeError(f"Unsupported quant type {quant_dtype}") |
| 293 | else: |
| 294 | if quant_dtype == torch.int8: |
| 295 | w, w_s = ops.scaled_int8_quant( |
| 296 | w, w_s, use_per_token_if_dynamic=per_token_quant |
| 297 | ) |
| 298 | elif is_fp8_dtype: |
| 299 | w, w_s = ops.scaled_fp8_quant( |
| 300 | w, w_s, use_per_token_if_dynamic=per_token_quant |
| 301 | ) |
| 302 | elif quant_dtype == "nvfp4": |
| 303 | assert not per_token_quant |
| 304 | w_amax = torch.abs(w).max().to(torch.float32) |
| 305 | w_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w_amax |
| 306 | if current_platform.is_rocm(): |
| 307 | w, w_s = _scaled_fp4_quant_emulated(w, w_gs) |
| 308 | else: |
| 309 | w, w_s = ops.scaled_fp4_quant(w, w_gs) |
| 310 | else: |
| 311 | raise RuntimeError(f"Unsupported quant type {quant_dtype}") |
| 312 | |
| 313 | return w, w_s, w_gs |
| 314 | |
| 315 | |
| 316 | def _pack_e2m1_fp4(fp4_values: torch.Tensor) -> torch.Tensor: |
| 317 | assert fp4_values.shape[-1] % 2 == 0 |
| 318 | |
| 319 | abs_values = fp4_values.abs() |
| 320 | codes = torch.empty_like(abs_values, dtype=torch.uint8) |
| 321 | for code, value in enumerate((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)): |
| 322 | codes[abs_values == value] = code |
| 323 | codes = codes | ((fp4_values < 0).to(torch.uint8) << 3) |
| 324 | return codes[..., 0::2] | (codes[..., 1::2] << 4) |
| 325 | |
| 326 | |
| 327 | def _scaled_fp4_quant_emulated( |
| 328 | w: torch.Tensor, w_gs: torch.Tensor |
| 329 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 330 | fp4_values, w_s = ref_nvfp4_quant(w, w_gs, block_size=16) |
| 331 | return _pack_e2m1_fp4(fp4_values), w_s.to(torch.float8_e4m3fn) |
| 332 | |
| 333 | |
| 334 | def moe_quantize_weights( |
| 335 | w: torch.Tensor, |
| 336 | w_s: torch.Tensor | None, |
| 337 | quant_dtype: torch.dtype | str | None, |
| 338 | per_token_quant: bool, |
| 339 | block_shape: list[int] | None, |
| 340 | ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: |
| 341 | assert w.dim() == 3 |
| 342 | e, rows, cols = w.shape |
| 343 | w_l = [None] * e |
| 344 | w_s_l = [None] * e |
| 345 | w_gs_l = [None] * e |
| 346 | for idx in range(e): |
| 347 | w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights_2d( |
| 348 | w[idx], None, quant_dtype, per_token_quant, block_shape |
| 349 | ) |
| 350 | |
| 351 | w = torch.stack(w_l) |
| 352 | w_s = torch.stack(w_s_l) |
| 353 | w_gs = torch.stack(w_gs_l) if e > 0 and w_gs_l[0] is not None else None |
| 354 | |
| 355 | if w_s.ndim == 2: |
| 356 | assert w_s.shape[-1] == 1 |
| 357 | w_s = w_s.view(-1, 1, 1) |
| 358 | |
| 359 | if block_shape is not None: |
| 360 | block_n, block_k = block_shape |
| 361 | n_tiles = (rows + block_n - 1) // block_n |
| 362 | k_tiles = (cols + block_k - 1) // block_k |
| 363 | assert w_s.shape == (e, n_tiles, k_tiles) |
| 364 | |
| 365 | return w, w_s, w_gs |
| 366 | |
| 367 | |
| 368 | def make_test_weight( |
| 369 | e: int, |
| 370 | rows: int, |
| 371 | cols: int, |
| 372 | in_dtype: torch.dtype = torch.bfloat16, |
| 373 | quant_dtype: torch.dtype | str | None = None, |
| 374 | block_shape: list[int] | None = None, |
| 375 | per_out_ch_quant: bool = False, |
| 376 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: |
| 377 | w_16 = torch.randn((e, rows, cols), device=DEVICE, dtype=in_dtype) / 15 |
| 378 | |
| 379 | if quant_dtype is not None: |
| 380 | w, w_s, w_gs = moe_quantize_weights( |
| 381 | w_16, None, quant_dtype, per_out_ch_quant, block_shape |
| 382 | ) |
| 383 | else: |
| 384 | w = w_16 |
| 385 | w_s = None |
| 386 | w_gs = None |
| 387 | |
| 388 | return w_16, w, w_s, w_gs |
| 389 | |
| 390 | |
| 391 | def make_test_weights( |
| 392 | e: int, |
| 393 | n: int, |
| 394 | k: int, |
| 395 | in_dtype: torch.dtype = torch.bfloat16, |
| 396 | quant_dtype: torch.dtype | str | None = None, |
| 397 | block_shape: list[int] | None = None, |
| 398 | per_out_ch_quant: bool = False, |
| 399 | make_gate: bool = True, |
| 400 | ) -> tuple[ |
| 401 | tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None], |
| 402 | tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None], |
| 403 | ]: |
| 404 | return ( |
| 405 | make_test_weight( |
| 406 | e, |
| 407 | (2 if make_gate else 1) * n, |
| 408 | k, |
| 409 | in_dtype, |
| 410 | quant_dtype, |
| 411 | block_shape, |
| 412 | per_out_ch_quant, |
| 413 | ), |
| 414 | make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape, per_out_ch_quant), |
| 415 | ) |
| 416 | |
| 417 | |
| 418 | def per_token_cast_to_fp8( |
| 419 | x: torch.Tensor, block_size: int = 128 |
| 420 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 421 | assert x.dim() == 2 |
| 422 | m, n = x.shape |
| 423 | pad_size = (block_size - (n % block_size)) % block_size |
| 424 | x = torch.nn.functional.pad(x, (0, pad_size), value=0) if pad_size > 0 else x |
| 425 | x_view = x.view(m, -1, block_size) |
| 426 | x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) |
| 427 | fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn) |
| 428 | return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1) |
| 429 | |
| 430 | |
| 431 | def make_test_quant_config( |
| 432 | e: int, |
| 433 | n: int, |
| 434 | k: int, |
| 435 | in_dtype: torch.dtype, |
| 436 | quant_dtype: torch.dtype | str | None = None, |
| 437 | per_act_token_quant: bool = False, |
| 438 | block_shape: list[int] | None = None, |
| 439 | make_gate: bool = True, |
| 440 | is_scale_swizzled: bool = True, |
| 441 | ) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]: |
| 442 | (_, w1, w1_s, w1_gs), (_, w2, w2_s, w2_gs) = make_test_weights( |
| 443 | e, |
| 444 | n, |
| 445 | k, |
| 446 | in_dtype, |
| 447 | quant_dtype, |
| 448 | per_out_ch_quant=per_act_token_quant, |
| 449 | block_shape=block_shape, |
| 450 | make_gate=make_gate, |
| 451 | ) |
| 452 | |
| 453 | # Hacky/trivial scales for nvfp4. |
| 454 | a1_gscale: torch.Tensor | None = None |
| 455 | a2_gscale: torch.Tensor | None = None |
| 456 | if quant_dtype == "nvfp4": |
| 457 | a1_gscale = torch.ones((e,), device=DEVICE, dtype=torch.float32) |
| 458 | a2_gscale = torch.ones((e,), device=DEVICE, dtype=torch.float32) |
| 459 | a1_scale = a1_gscale |
| 460 | a2_scale = a2_gscale |
| 461 | else: |
| 462 | a1_scale = None |
| 463 | a2_scale = None |
| 464 | |
| 465 | return ( |
| 466 | w1, |
| 467 | w2, |
| 468 | FusedMoEQuantConfig.make( |
| 469 | quant_dtype, |
| 470 | per_act_token_quant=per_act_token_quant, |
| 471 | block_shape=block_shape, |
| 472 | w1_scale=w1_s, |
| 473 | w2_scale=w2_s, |
| 474 | a1_gscale=a1_gscale, |
| 475 | a2_gscale=a2_gscale, |
| 476 | a1_scale=a1_scale, |
| 477 | a2_scale=a2_scale, |
| 478 | # TODO: make sure this is handled properly |
| 479 | g1_alphas=(1 / w1_gs) if w1_gs is not None else None, |
| 480 | g2_alphas=(1 / w2_gs) if w2_gs is not None else None, |
| 481 | is_scale_swizzled=is_scale_swizzled, |
| 482 | ), |
| 483 | ) |
| 484 | |
| 485 | |
| 486 | def fused_moe( |
| 487 | hidden_states: torch.Tensor, |
| 488 | w1: torch.Tensor, |
| 489 | w2: torch.Tensor, |
| 490 | score: torch.Tensor, |
| 491 | topk: int, |
| 492 | renormalize: bool = False, |
| 493 | quant_config: FusedMoEQuantConfig | None = None, |
| 494 | global_num_experts: int = -1, |
| 495 | expert_map: torch.Tensor | None = None, |
| 496 | ) -> torch.Tensor: |
| 497 | topk_weights, topk_ids, _ = fused_topk( |
| 498 | hidden_states, score.float(), topk, renormalize |
| 499 | ) |
| 500 | return fused_experts( |
| 501 | hidden_states, |
| 502 | w1, |
| 503 | w2, |
| 504 | topk_weights, |
| 505 | topk_ids, |
| 506 | global_num_experts=global_num_experts, |
| 507 | expert_map=expert_map, |
| 508 | quant_config=quant_config, |
| 509 | ) |
| 510 | |
| 511 | |
| 512 | class BaselineMM(torch.nn.Module): |
| 513 | def __init__( |
| 514 | self, |
| 515 | b: torch.Tensor, |
| 516 | out_dtype: torch.dtype, |
| 517 | ): |
| 518 | super().__init__() |
| 519 | self.b = torch.nn.Parameter(b.to(dtype=torch.float32)) |
| 520 | self.out_dtype = out_dtype |
| 521 | |
| 522 | def forward(self, a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 523 | return torch.mm(a.to(dtype=torch.float32), self.b).to(self.out_dtype), None |
| 524 | |
| 525 | |
| 526 | class BaselineSiluAndMul(torch.nn.Module): |
| 527 | def __init__(self): |
| 528 | super().__init__() |
| 529 | |
| 530 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 531 | d = x.shape[-1] // 2 |
| 532 | return torch.nn.functional.silu(x[..., :d]) * x[..., d:] |
| 533 | |
| 534 | |
| 535 | class TestMLP(torch.nn.Module): |
| 536 | def __init__( |
| 537 | self, |
| 538 | w1: torch.Tensor, |
| 539 | w2: torch.Tensor, |
| 540 | out_dtype: torch.dtype, |
| 541 | ): |
| 542 | super().__init__() |
| 543 | self.gate_up_proj = BaselineMM(w1, out_dtype) |
| 544 | self.down_proj = BaselineMM(w2, out_dtype) |
| 545 | self.act_fn = BaselineSiluAndMul() |
| 546 | |
| 547 | def forward(self, x): |
| 548 | x, _ = self.gate_up_proj(x) |
| 549 | x = self.act_fn(x) |
| 550 | x, _ = self.down_proj(x) |
| 551 | return x |
| 552 | |
| 553 | |
| 554 | def modular_triton_fused_moe( |
| 555 | moe_config: FusedMoEConfig, |
| 556 | quant_config: FusedMoEQuantConfig, |
| 557 | ) -> FusedMoEKernel: |
| 558 | return FusedMoEKernel( |
| 559 | maybe_make_prepare_finalize( |
| 560 | moe=moe_config, |
| 561 | quant_config=quant_config, |
| 562 | allow_new_interface=True, |
| 563 | use_monolithic=False, |
| 564 | ), |
| 565 | TritonExperts(moe_config, quant_config), |
| 566 | ) |
| 567 | |
| 568 | |
| 569 | def check_accuracy(a, b, atol, rtol, percent): |
| 570 | """Allow a mismatch percentage of 1 - percent.""" |
| 571 | if torch.any(torch.isnan(a)): |
| 572 | raise Exception("NaN in reference output") |
| 573 | if torch.any(torch.isnan(b)): |
| 574 | raise Exception("NaN in actual output") |
| 575 | if torch.any(torch.isinf(a)): |
| 576 | raise Exception("Inf in reference output") |
| 577 | if torch.any(torch.isinf(b)): |
| 578 | raise Exception("Inf in actual output") |
| 579 | assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}" |
| 580 | |
| 581 | left = torch.abs(a - b) |
| 582 | right = atol + rtol * torch.abs(b) |
| 583 | count = torch.sum(left > right) |
| 584 | mismatch_percent = count / a.numel() |
| 585 | if mismatch_percent > 1 - percent: |
| 586 | raise Exception( |
| 587 | f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " |
| 588 | f"(threshold: {1 - percent:.4f})" |
| 589 | ) |
| 590 | |
| 591 | |
| 592 | def check_deferred_moe_finalize( |
| 593 | moe_config: FusedMoEConfig, |
| 594 | run: Callable[[], torch.Tensor | UnfinalizedMoEOutput], |
| 595 | router_weights: torch.Tensor | None = None, |
| 596 | chunked: bool = False, |
| 597 | ) -> None: |
| 598 | """Check a kernel that defers its finalize against the finalize it skips. |
| 599 | |
| 600 | ``run`` calls the kernel on fixed inputs, first as built and then with |
| 601 | ``moe_config`` asking to defer, which ``should_defer_moe_finalize`` must |
| 602 | report truthfully. A deferred output reduced by ``moe_unpermute``, the |
| 603 | TRT-LLM finalize kernel, must give the kernel's own finalized output bit for |
| 604 | bit, and modular experts hand the router's weights back as-is. A call the |
| 605 | experts split across kernel launches finalizes instead. |
| 606 | """ |
| 607 | # The deferred output views the router's buffer, so compare with a copy. |
| 608 | expected_weights = None if router_weights is None else router_weights.clone() |
| 609 | finalized = run() |
| 610 | assert isinstance(finalized, torch.Tensor) |
| 611 | moe_config.defer_moe_finalize() |
| 612 | output = run() |
| 613 | assert moe_config.should_defer_moe_finalize(finalized.shape[0]) != chunked |
| 614 | if chunked: |
| 615 | torch.testing.assert_close(output, finalized, atol=0, rtol=0) |
| 616 | return |
| 617 | |
| 618 | assert isinstance(output, UnfinalizedMoEOutput) |
| 619 | if expected_weights is not None: |
| 620 | torch.testing.assert_close( |
| 621 | output.expert_weights, expected_weights, atol=0, rtol=0 |
| 622 | ) |
| 623 | reference = torch.empty_like(finalized) |
| 624 | moe_unpermute( |
| 625 | reference, |
| 626 | output.gemm2_permuted, |
| 627 | output.expert_weights.float(), |
| 628 | output.expanded_idx_to_permuted_idx, |
| 629 | # The kernel reads its valid-row count through this pointer either way. |
| 630 | expert_first_token_offset=output.gemm2_permuted.new_full( |
| 631 | (1,), output.gemm2_permuted.shape[0], dtype=torch.int64 |
| 632 | ), |
| 633 | ) |
| 634 | torch.testing.assert_close(reference, finalized, atol=0, rtol=0) |
| 635 | |
| 636 | |
| 637 | def mxfp4_w_layouts(mx_axis: int, num_warps: int = 8): |
| 638 | """Weight/scale layouts for mxfp4 MoE, as (layout, opts) pairs. |
| 639 | |
| 640 | triton_kernels 3.8 returns layout instances; earlier versions return a |
| 641 | (layout, opts) tuple. |
| 642 | """ |
| 643 | from triton_kernels.tensor_details import layout |
| 644 | |
| 645 | from vllm.utils.import_utils import get_triton_kernels_version |
| 646 | |
| 647 | if get_triton_kernels_version() == "3.8": |
| 648 | w = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis) |
| 649 | s = layout.make_default_matmul_mxfp4_w_scale_layout( |
| 650 | mx_axis=mx_axis, num_warps=num_warps |
| 651 | ) |
| 652 | return w, {}, s, {} |
| 653 | |
| 654 | w, w_opts = layout.make_default_matmul_mxfp4_w_layout(mx_axis=mx_axis) |
| 655 | s, s_opts = layout.make_default_matmul_mxfp4_w_scale_layout( |
| 656 | mx_axis=mx_axis, num_warps=num_warps |
| 657 | ) |
| 658 | return w, w_opts, s, s_opts |