| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import types |
| 5 | from dataclasses import dataclass, replace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from tests.kernels.moe.utils import check_accuracy, check_deferred_moe_finalize |
| 11 | from vllm._aiter_ops import ( |
| 12 | is_aiter_found, |
| 13 | is_aiter_found_and_supported, |
| 14 | rocm_aiter_ops, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import ( |
| 17 | activation_quant_dtype, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input |
| 20 | from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.utils.flashinfer import has_flashinfer |
| 23 | |
| 24 | TRTLLM_GEN_MXFP4_AVAILABLE = ( |
| 25 | current_platform.is_cuda() and current_platform.is_device_capability_family(100) |
| 26 | ) |
| 27 | |
| 28 | TRTLLM_GEN_MXFP8_AVAILABLE = TRTLLM_GEN_MXFP4_AVAILABLE |
| 29 | |
| 30 | HOPPER_MXFP4_BF16_AVAILABLE = ( |
| 31 | current_platform.is_cuda() |
| 32 | and current_platform.is_device_capability(90) |
| 33 | and has_flashinfer() |
| 34 | ) |
| 35 | |
| 36 | # ROCm platform and dependencies |
| 37 | ROCM_AVAILABLE = current_platform.is_rocm() |
| 38 | ROCM_TRITON_KERNELS_AVAILABLE = False |
| 39 | ROCM_AITER_AVAILABLE = is_aiter_found() |
| 40 | ROCM_GFX950 = False |
| 41 | |
| 42 | if ROCM_AVAILABLE: |
| 43 | from vllm.platforms.rocm import on_gfx950 |
| 44 | from vllm.utils.import_utils import has_triton_kernels |
| 45 | |
| 46 | ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels() |
| 47 | ROCM_GFX950 = on_gfx950() |
| 48 | |
| 49 | if ROCM_AITER_AVAILABLE: |
| 50 | from aiter.ops.triton.quant import dynamic_mxfp4_quant |
| 51 | |
| 52 | if TRTLLM_GEN_MXFP4_AVAILABLE: |
| 53 | from flashinfer import ( |
| 54 | fp4_quantize, |
| 55 | mxfp8_quantize, |
| 56 | reorder_rows_for_gated_act_gemm, |
| 57 | shuffle_matrix_a, |
| 58 | shuffle_matrix_sf_a, |
| 59 | trtllm_fp4_block_scale_moe, |
| 60 | trtllm_fp8_block_scale_moe, |
| 61 | ) |
| 62 | from flashinfer.fp4_quantization import nvfp4_block_scale_interleave |
| 63 | |
| 64 | if TRTLLM_GEN_MXFP8_AVAILABLE: |
| 65 | from flashinfer.fused_moe.core import ( |
| 66 | Fp8QuantizationType, |
| 67 | get_w2_permute_indices_with_cache, |
| 68 | ) |
| 69 | |
| 70 | |
| 71 | @dataclass |
| 72 | class ModelCase: |
| 73 | model_id: str |
| 74 | tp: int |
| 75 | |
| 76 | |
| 77 | @pytest.fixture(scope="function", autouse=True) |
| 78 | def enable_pickle(monkeypatch): |
| 79 | """`LLM.apply_model` requires pickling a function.""" |
| 80 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 81 | |
| 82 | |
| 83 | @pytest.mark.parametrize( |
| 84 | "model_case", |
| 85 | [ |
| 86 | ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=2), |
| 87 | ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8), |
| 88 | ModelCase("mawong-amd/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1), |
| 89 | ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=1), |
| 90 | ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=4), |
| 91 | ], |
| 92 | ) |
| 93 | @pytest.mark.skipif( |
| 94 | not ROCM_AVAILABLE, |
| 95 | reason="This test is only enabled on platforms that support amd-quark", |
| 96 | ) |
| 97 | def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase): |
| 98 | if torch.accelerator.device_count() < model_case.tp: |
| 99 | pytest.skip( |
| 100 | f"This test requires >={model_case.tp} gpus, got only " |
| 101 | f"{torch.accelerator.device_count()}" |
| 102 | ) |
| 103 | |
| 104 | # `cudagraph_capture_sizes=[16]` to reduce load time. |
| 105 | with vllm_runner( |
| 106 | model_case.model_id, |
| 107 | tensor_parallel_size=model_case.tp, |
| 108 | load_format="dummy", |
| 109 | compilation_config={"cudagraph_capture_sizes": [16]}, |
| 110 | gpu_memory_utilization=0.8, # mxfp6 models use more scratch space |
| 111 | ) as llm: |
| 112 | # Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562 |
| 113 | # def check_model(model): |
| 114 | # from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501 |
| 115 | # QuarkLinearMethod) |
| 116 | # from vllm.model_executor.layers.quantization.quark.schemes.quark_ocp_mx import QuarkOCP_MX # noqa: E501 |
| 117 | # from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501 |
| 118 | # QuarkOCP_MX_MoEMethod) |
| 119 | |
| 120 | # layer = model.model.layers[0] |
| 121 | |
| 122 | # qkv_proj = layer.self_attn.qkv_proj |
| 123 | |
| 124 | # assert isinstance(qkv_proj.quant_method, QuarkLinearMethod) |
| 125 | # assert isinstance(qkv_proj.scheme, QuarkOCP_MX) |
| 126 | |
| 127 | # assert isinstance(layer.mlp.experts.quant_method, |
| 128 | # QuarkOCP_MX_MoEMethod) |
| 129 | |
| 130 | # if model_case.model_id == "fxmarty/qwen_1.5-moe-a2.7b-mxfp4": |
| 131 | # llm.apply_model(check_model) |
| 132 | output = llm.generate_greedy("Today I am in the French Alps and", max_tokens=20) |
| 133 | assert output |
| 134 | |
| 135 | |
| 136 | def swiglu( |
| 137 | x, |
| 138 | alpha: float = 1.702, |
| 139 | beta: float = 1.0, |
| 140 | limit: float | None = None, |
| 141 | interleaved: bool = False, |
| 142 | ): |
| 143 | if interleaved: |
| 144 | x_glu, x_linear = x[..., ::2], x[..., 1::2] |
| 145 | else: |
| 146 | x_glu, x_linear = torch.chunk(x, 2, dim=-1) |
| 147 | if limit is not None: |
| 148 | x_glu = x_glu.clamp(max=limit) |
| 149 | x_linear = x_linear.clamp(min=-limit, max=limit) |
| 150 | out_glu = x_glu * torch.sigmoid(alpha * x_glu) |
| 151 | return out_glu * (x_linear + beta) |
| 152 | |
| 153 | |
| 154 | fp4_lookup_table = [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0, -0.5, -1, -1.5, -2, -3, -4, -6] |
| 155 | |
| 156 | |
| 157 | def mxfp4_dequantize(x, scale): |
| 158 | assert x.dtype == torch.uint8 |
| 159 | x = x.view(torch.uint8).to(torch.int32) |
| 160 | x_unpacked = torch.zeros( |
| 161 | *x.shape[:-1], x.shape[-1] * 2, dtype=torch.int32, device=x.device |
| 162 | ) |
| 163 | x_unpacked[..., 0::2].copy_(x & 0xF) |
| 164 | x_unpacked[..., 1::2].copy_((x >> 4) & 0xF) |
| 165 | |
| 166 | x_float = torch.zeros(x_unpacked.shape, dtype=torch.float32, device=x.device) |
| 167 | for i, val in enumerate(fp4_lookup_table): |
| 168 | x_float[x_unpacked == i] = val |
| 169 | |
| 170 | scale = scale.view(torch.uint8).to(torch.int32) |
| 171 | scale = (scale << 23).view(torch.float32) |
| 172 | scale = scale.reshape(*x.shape[:-1], -1) |
| 173 | scale = torch.stack([scale] * 32, dim=-1).reshape(*x_float.shape) |
| 174 | |
| 175 | return x_float * scale |
| 176 | |
| 177 | |
| 178 | def mxfp8_dequantize(x, scale): |
| 179 | assert x.dtype == torch.float8_e4m3fn |
| 180 | x_float = x.to(torch.float32) |
| 181 | |
| 182 | scale = scale.view(torch.uint8).to(torch.int32) |
| 183 | scale = (scale << 23).view(torch.float32) |
| 184 | scale = scale.reshape(*x.shape[:-1], -1) |
| 185 | scale = torch.stack([scale] * 32, dim=-1).reshape(*x_float.shape) |
| 186 | |
| 187 | return x_float * scale |
| 188 | |
| 189 | |
| 190 | def mxfp4_quant_dequant(x: torch.Tensor) -> torch.Tensor: |
| 191 | shape = x.shape |
| 192 | quantized, scale = dynamic_mxfp4_quant(x.to(torch.bfloat16).flatten(0, -2)) |
| 193 | return mxfp4_dequantize(quantized.view(torch.uint8), scale).reshape(shape) |
| 194 | |
| 195 | |
| 196 | def reference_moe( |
| 197 | roouting_logits, |
| 198 | topk, |
| 199 | num_experts, |
| 200 | hidden_states, |
| 201 | w13, |
| 202 | bias13, |
| 203 | w2, |
| 204 | bias2, |
| 205 | alpha, |
| 206 | beta, |
| 207 | limit, |
| 208 | act_type, |
| 209 | activation: str = "swiglu", |
| 210 | use_interleaved_layout: bool = False, |
| 211 | ): |
| 212 | """Reference MoE implementation for accuracy testing. |
| 213 | |
| 214 | Args: |
| 215 | activation: One of "swiglu", "silu", "relu2". Controls the activation |
| 216 | function used after the first MLP. |
| 217 | use_interleaved_layout: If True, uses interleaved gate/up layout |
| 218 | (gate=x[..., ::2], up=x[..., 1::2]) as used by SWIGLUOAI. |
| 219 | If False, uses chunked layout (gate, up = chunk(x, 2)) as used |
| 220 | by standard swiglu/silu. |
| 221 | |
| 222 | """ |
| 223 | # renormalize routing |
| 224 | experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True) |
| 225 | expert_weights = torch.nn.functional.softmax(experts.values, dim=1) |
| 226 | expert_indices = experts.indices |
| 227 | t = hidden_states.clone() |
| 228 | if act_type == "mxfp4": |
| 229 | t = mxfp4_quant_dequant(t) |
| 230 | # MLP #1 |
| 231 | mlp1_weight = w13[expert_indices, ...] |
| 232 | mlp1_bias = bias13[expert_indices, ...] |
| 233 | t = torch.einsum("beck,bk->bec", mlp1_weight, t) + mlp1_bias |
| 234 | |
| 235 | # Apply activation |
| 236 | if activation in ("swiglu", "silu"): |
| 237 | t = swiglu( |
| 238 | t, alpha=alpha, beta=beta, limit=limit, interleaved=use_interleaved_layout |
| 239 | ) |
| 240 | elif activation == "relu2": |
| 241 | # RELU2_NO_MUL: relu(x)^2 |
| 242 | t = torch.relu(t) |
| 243 | t = t * t |
| 244 | else: |
| 245 | raise ValueError(f"Unknown activation: {activation}") |
| 246 | |
| 247 | if act_type == "mxfp8": |
| 248 | t_quantized, t_scale = mxfp8_quantize( |
| 249 | t.to(torch.bfloat16), is_sf_swizzled_layout=False |
| 250 | ) |
| 251 | t = mxfp8_dequantize(t_quantized, t_scale) |
| 252 | elif act_type == "mxfp4": |
| 253 | t = mxfp4_quant_dequant(t) |
| 254 | # MLP #2 |
| 255 | mlp2_weight = w2[expert_indices, ...] |
| 256 | mlp2_bias = bias2[expert_indices, ...] |
| 257 | t = torch.einsum("beck,bek->bec", mlp2_weight, t) + mlp2_bias |
| 258 | # Weighted sum of experts |
| 259 | t = torch.einsum("bec,be->bc", t, expert_weights) |
| 260 | assert t.shape == hidden_states.shape |
| 261 | return t.to(torch.bfloat16) |
| 262 | |
| 263 | |
| 264 | def tg_mxfp4_moe( |
| 265 | router_logits, |
| 266 | topk, |
| 267 | num_experts, |
| 268 | intermediate_size, |
| 269 | hidden_size, |
| 270 | hidden_states, |
| 271 | hidden_states_scale, |
| 272 | w13_weight, |
| 273 | w13_weight_scale, |
| 274 | w13_bias, |
| 275 | w2_weight, |
| 276 | w2_weight_scale, |
| 277 | w2_bias, |
| 278 | act_type, |
| 279 | alpha, |
| 280 | beta, |
| 281 | limit, |
| 282 | transpose_optimized: bool = False, |
| 283 | ) -> torch.Tensor: |
| 284 | sf_block_size = 32 |
| 285 | assert ( |
| 286 | w13_weight.dim() == 3 |
| 287 | and w13_weight.shape[0] == num_experts |
| 288 | and w13_weight.shape[1] == intermediate_size * 2 |
| 289 | and w13_weight.shape[2] == hidden_size // 2 |
| 290 | ) |
| 291 | assert ( |
| 292 | w13_weight_scale.dim() == 3 |
| 293 | and w13_weight_scale.shape[0] == num_experts |
| 294 | and w13_weight_scale.shape[1] == intermediate_size * 2 |
| 295 | and w13_weight_scale.shape[2] == hidden_size // sf_block_size |
| 296 | ) |
| 297 | assert ( |
| 298 | w2_weight.dim() == 3 |
| 299 | and w2_weight.shape[0] == num_experts |
| 300 | and w2_weight.shape[1] == hidden_size |
| 301 | and w2_weight.shape[2] == intermediate_size // 2 |
| 302 | ) |
| 303 | assert ( |
| 304 | w2_weight_scale.dim() == 3 |
| 305 | and w2_weight_scale.shape[1] == hidden_size |
| 306 | and w2_weight_scale.shape[2] == intermediate_size // sf_block_size |
| 307 | ) |
| 308 | assert ( |
| 309 | w13_bias.dim() == 2 |
| 310 | and w13_bias.shape[0] == num_experts |
| 311 | and w13_bias.shape[1] == intermediate_size * 2 |
| 312 | ) |
| 313 | assert ( |
| 314 | w2_bias.dim() == 2 |
| 315 | and w2_bias.shape[0] == num_experts |
| 316 | and w2_bias.shape[1] == hidden_size |
| 317 | ) |
| 318 | |
| 319 | # Swap w1 and w3 as the definition of |
| 320 | # swiglu is different in the trtllm-gen |
| 321 | w13_weight_scale_ = w13_weight_scale.clone() |
| 322 | w13_weight_ = w13_weight.clone() |
| 323 | w13_bias_ = w13_bias.clone() |
| 324 | w13_weight[:, :intermediate_size, :].copy_(w13_weight_[:, intermediate_size:, :]) |
| 325 | w13_weight[:, intermediate_size:, :].copy_(w13_weight_[:, :intermediate_size, :]) |
| 326 | w13_weight_scale[:, :intermediate_size, :].copy_( |
| 327 | w13_weight_scale_[:, intermediate_size:, :] |
| 328 | ) |
| 329 | w13_weight_scale[:, intermediate_size:, :].copy_( |
| 330 | w13_weight_scale_[:, :intermediate_size, :] |
| 331 | ) |
| 332 | w13_bias[:, :intermediate_size].copy_(w13_bias_[:, intermediate_size:]) |
| 333 | w13_bias[:, intermediate_size:].copy_(w13_bias_[:, :intermediate_size]) |
| 334 | |
| 335 | # Interleave the weights and scaling factors for activation |
| 336 | w13_weight_interleaved = [] |
| 337 | w13_weight_scale_interleaved = [] |
| 338 | w13_bias_interleaved = [] |
| 339 | for i in range(num_experts): |
| 340 | w13_weight_interleaved.append( |
| 341 | reorder_rows_for_gated_act_gemm(w13_weight[i].clone()) |
| 342 | ) |
| 343 | w13_weight_scale_interleaved.append( |
| 344 | reorder_rows_for_gated_act_gemm(w13_weight_scale[i].clone()) |
| 345 | ) |
| 346 | w13_bias_interleaved.append( |
| 347 | reorder_rows_for_gated_act_gemm(w13_bias[i].clone().reshape(-1, 1)) |
| 348 | ) |
| 349 | w13_weight = torch.stack(w13_weight_interleaved).reshape( |
| 350 | num_experts, 2 * intermediate_size, hidden_size // 2 |
| 351 | ) |
| 352 | w13_weight_scale = torch.stack(w13_weight_scale_interleaved).reshape( |
| 353 | num_experts, 2 * intermediate_size, hidden_size // 32 |
| 354 | ) |
| 355 | w13_bias = torch.stack(w13_bias_interleaved).reshape( |
| 356 | num_experts, 2 * intermediate_size |
| 357 | ) |
| 358 | |
| 359 | # Shuffle weights and scaling factors for transposed mma output |
| 360 | gemm1_weights_shuffled = [] |
| 361 | gemm1_scales_shuffled = [] |
| 362 | gemm2_weights_shuffled = [] |
| 363 | gemm2_scales_shuffled = [] |
| 364 | gemm1_bias_shuffled = [] |
| 365 | gemm2_bias_shuffled = [] |
| 366 | epilogue_tile_m = 128 # FIXME: this depends on the kernel internals |
| 367 | _cache_permute_indices: dict[torch.Size, torch.Tensor] = {} |
| 368 | if transpose_optimized: |
| 369 | for i in range(num_experts): |
| 370 | # w13 weight shuffling |
| 371 | permute_indices = get_w2_permute_indices_with_cache( |
| 372 | _cache_permute_indices, |
| 373 | w13_weight[i].view(torch.uint8), |
| 374 | epilogue_tile_m, |
| 375 | ) |
| 376 | gemm1_weights_shuffled.append( |
| 377 | w13_weight[i] |
| 378 | .view(torch.uint8)[permute_indices.to(w13_weight.device)] |
| 379 | .contiguous() |
| 380 | ) |
| 381 | # w13 scale shuffling |
| 382 | permute_sf_indices = get_w2_permute_indices_with_cache( |
| 383 | _cache_permute_indices, |
| 384 | w13_weight_scale[i].view(torch.uint8), |
| 385 | epilogue_tile_m, |
| 386 | num_elts_per_sf=16, |
| 387 | ) |
| 388 | gemm1_scales_shuffled.append( |
| 389 | nvfp4_block_scale_interleave( |
| 390 | w13_weight_scale[i] |
| 391 | .view(torch.uint8)[permute_sf_indices.to(w13_weight_scale.device)] |
| 392 | .contiguous() |
| 393 | ) |
| 394 | ) |
| 395 | # w13 bias shuffling |
| 396 | permute_bias_indices = get_w2_permute_indices_with_cache( |
| 397 | _cache_permute_indices, |
| 398 | w13_bias[i].clone().reshape(-1, 1), |
| 399 | epilogue_tile_m, |
| 400 | ) |
| 401 | gemm1_bias_shuffled.append( |
| 402 | w13_bias[i] |
| 403 | .clone() |
| 404 | .reshape(-1, 1)[permute_bias_indices.to(w13_bias.device)] |
| 405 | .contiguous() |
| 406 | ) |
| 407 | # w2 weight shuffling |
| 408 | permute_indices = get_w2_permute_indices_with_cache( |
| 409 | _cache_permute_indices, |
| 410 | w2_weight[i].view(torch.uint8), |
| 411 | epilogue_tile_m, |
| 412 | ) |
| 413 | gemm2_weights_shuffled.append( |
| 414 | w2_weight[i] |
| 415 | .view(torch.uint8)[permute_indices.to(w2_weight.device)] |
| 416 | .contiguous() |
| 417 | ) |
| 418 | # w2 scale shuffling |
| 419 | permute_sf_indices = get_w2_permute_indices_with_cache( |
| 420 | _cache_permute_indices, |
| 421 | w2_weight_scale[i].view(torch.uint8), |
| 422 | epilogue_tile_m, |
| 423 | num_elts_per_sf=16, |
| 424 | ) |
| 425 | gemm2_scales_shuffled.append( |
| 426 | nvfp4_block_scale_interleave( |
| 427 | w2_weight_scale[i] |
| 428 | .view(torch.uint8)[permute_sf_indices.to(w2_weight_scale.device)] |
| 429 | .contiguous() |
| 430 | ) |
| 431 | ) |
| 432 | # w2 bias shuffling |
| 433 | permute_indices = get_w2_permute_indices_with_cache( |
| 434 | _cache_permute_indices, |
| 435 | w2_bias[i].clone().reshape(-1, 1), |
| 436 | epilogue_tile_m, |
| 437 | ) |
| 438 | gemm2_bias_shuffled.append( |
| 439 | w2_bias[i] |
| 440 | .clone() |
| 441 | .reshape(-1, 1)[permute_indices.to(w2_bias.device)] |
| 442 | .contiguous() |
| 443 | ) |
| 444 | |
| 445 | else: |
| 446 | for i in range(num_experts): |
| 447 | gemm1_weights_shuffled.append( |
| 448 | shuffle_matrix_a(w13_weight[i].view(torch.uint8), epilogue_tile_m) |
| 449 | ) |
| 450 | gemm1_scales_shuffled.append( |
| 451 | shuffle_matrix_sf_a( |
| 452 | w13_weight_scale[i].view(torch.uint8), epilogue_tile_m |
| 453 | ) |
| 454 | ) |
| 455 | |
| 456 | gemm2_weights_shuffled.append( |
| 457 | shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m) |
| 458 | ) |
| 459 | gemm2_scales_shuffled.append( |
| 460 | shuffle_matrix_sf_a( |
| 461 | w2_weight_scale[i].view(torch.uint8), epilogue_tile_m |
| 462 | ) |
| 463 | ) |
| 464 | gemm1_bias_shuffled.append( |
| 465 | shuffle_matrix_a(w13_bias[i].reshape(-1, 1), epilogue_tile_m) |
| 466 | ) |
| 467 | gemm2_bias_shuffled.append( |
| 468 | shuffle_matrix_a(w2_bias[i].reshape(-1, 1), epilogue_tile_m) |
| 469 | ) |
| 470 | |
| 471 | w13_weight = torch.stack(gemm1_weights_shuffled) |
| 472 | w13_weight_scale = ( |
| 473 | torch.stack(gemm1_scales_shuffled) |
| 474 | .reshape(num_experts, 2 * intermediate_size, hidden_size // sf_block_size) |
| 475 | .view(torch.float8_e4m3fn) |
| 476 | ) |
| 477 | w13_bias = torch.stack(gemm1_bias_shuffled).reshape(num_experts, -1) |
| 478 | |
| 479 | w2_weight = torch.stack(gemm2_weights_shuffled) |
| 480 | w2_weight_scale = ( |
| 481 | torch.stack(gemm2_scales_shuffled) |
| 482 | .reshape(num_experts, hidden_size, intermediate_size // sf_block_size) |
| 483 | .view(torch.float8_e4m3fn) |
| 484 | ) |
| 485 | w2_bias = torch.stack(gemm2_bias_shuffled).reshape(num_experts, -1) |
| 486 | |
| 487 | tg_result = trtllm_fp4_block_scale_moe( |
| 488 | routing_logits=router_logits.to(torch.bfloat16), |
| 489 | routing_bias=None, |
| 490 | hidden_states=hidden_states, |
| 491 | hidden_states_scale=hidden_states_scale, |
| 492 | gemm1_weights=w13_weight, |
| 493 | gemm1_weights_scale=w13_weight_scale, |
| 494 | gemm1_bias=w13_bias, |
| 495 | gemm1_alpha=alpha, |
| 496 | gemm1_beta=beta, |
| 497 | gemm1_clamp_limit=limit, |
| 498 | gemm2_weights=w2_weight, |
| 499 | gemm2_weights_scale=w2_weight_scale, |
| 500 | gemm2_bias=w2_bias, |
| 501 | output1_scale_scalar=None, |
| 502 | output1_scale_gate_scalar=None, |
| 503 | output2_scale_scalar=None, |
| 504 | num_experts=num_experts, |
| 505 | top_k=topk, |
| 506 | n_group=None, |
| 507 | topk_group=None, |
| 508 | intermediate_size=intermediate_size, |
| 509 | local_expert_offset=0, |
| 510 | local_num_experts=num_experts, |
| 511 | routed_scaling_factor=None, |
| 512 | routing_method_type=1, # renormalize |
| 513 | do_finalize=True, |
| 514 | )[0] |
| 515 | return tg_result |
| 516 | |
| 517 | |
| 518 | @pytest.mark.parametrize("topk", [1, 4]) |
| 519 | @pytest.mark.parametrize("num_experts", [32, 128]) |
| 520 | @pytest.mark.parametrize("num_tokens", [1, 128, 1024]) |
| 521 | @pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)]) |
| 522 | @pytest.mark.parametrize("alpha,beta,limit", [(1.0, 1.0, None), (1.702, 1.0, 7.0)]) |
| 523 | @pytest.mark.parametrize("act_type", ["mxfp8", "bf16"]) |
| 524 | @pytest.mark.parametrize("transpose_optimized", [False, True]) |
| 525 | @pytest.mark.skipif( |
| 526 | not TRTLLM_GEN_MXFP4_AVAILABLE, |
| 527 | reason="nvidia gpu and compute capability sm100 is required for this test", |
| 528 | ) |
| 529 | def test_trtllm_gen_mxfp4_fused_moe( |
| 530 | topk: int, |
| 531 | num_experts: int, |
| 532 | num_tokens: int, |
| 533 | intermediate_size: int, |
| 534 | hidden_size: int, |
| 535 | alpha: float, |
| 536 | beta: float, |
| 537 | limit: float | None, |
| 538 | act_type: str, |
| 539 | transpose_optimized: bool, |
| 540 | ): |
| 541 | seed = 42 |
| 542 | torch.manual_seed(seed) |
| 543 | hidden_states = torch.randn( |
| 544 | num_tokens, hidden_size, device="cuda:0", dtype=torch.bfloat16 |
| 545 | ) |
| 546 | w13 = torch.randn( |
| 547 | num_experts, |
| 548 | intermediate_size * 2, |
| 549 | hidden_size, |
| 550 | device="cuda:0", |
| 551 | dtype=torch.bfloat16, |
| 552 | ) |
| 553 | w2 = torch.randn( |
| 554 | num_experts, |
| 555 | hidden_size, |
| 556 | intermediate_size, |
| 557 | device="cuda:0", |
| 558 | dtype=torch.bfloat16, |
| 559 | ) |
| 560 | bias13 = torch.randn(num_experts, intermediate_size * 2, device="cuda:0") * 10 |
| 561 | bias2 = torch.randn(num_experts, hidden_size, device="cuda:0") * 10 |
| 562 | router_logits = torch.rand(num_tokens, num_experts, dtype=torch.float32).cuda() |
| 563 | |
| 564 | w13, w13_scale = fp4_quantize( |
| 565 | w13, |
| 566 | torch.tensor(1.0, device="cuda:0"), |
| 567 | 32, |
| 568 | sf_use_ue8m0=True, |
| 569 | is_sf_swizzled_layout=False, |
| 570 | ) |
| 571 | w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape( |
| 572 | num_experts, intermediate_size * 2, hidden_size // 32 |
| 573 | ) |
| 574 | w2, w2_scale = fp4_quantize( |
| 575 | w2, |
| 576 | torch.tensor(1.0, device="cuda:0"), |
| 577 | 32, |
| 578 | sf_use_ue8m0=True, |
| 579 | is_sf_swizzled_layout=False, |
| 580 | ) |
| 581 | w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape( |
| 582 | num_experts, hidden_size, intermediate_size // 32 |
| 583 | ) |
| 584 | if act_type == "mxfp8": |
| 585 | hidden_states, hidden_states_scale = mxfp8_quantize( |
| 586 | hidden_states, is_sf_swizzled_layout=False |
| 587 | ) |
| 588 | hidden_states_scale = hidden_states_scale.view(torch.float8_e4m3fn).reshape( |
| 589 | *hidden_states.shape[:-1], -1 |
| 590 | ) |
| 591 | else: |
| 592 | hidden_states_scale = None |
| 593 | |
| 594 | # reference result |
| 595 | ref_result = torch.empty_like(hidden_states, dtype=torch.bfloat16) |
| 596 | w13_ref = mxfp4_dequantize(w13.clone(), w13_scale.clone()) |
| 597 | w2_ref = mxfp4_dequantize(w2.clone(), w2_scale.clone()) |
| 598 | bias13_ref = bias13 |
| 599 | bias2_ref = bias2 |
| 600 | if act_type == "mxfp8": |
| 601 | hidden_states_ref = mxfp8_dequantize(hidden_states, hidden_states_scale).to( |
| 602 | torch.float32 |
| 603 | ) |
| 604 | else: |
| 605 | hidden_states_ref = hidden_states.to(torch.float32) |
| 606 | # Process tokens in chunks of 32 to reduce memory usage |
| 607 | chunk_size = 32 |
| 608 | num_chunks = (num_tokens + chunk_size - 1) // chunk_size |
| 609 | for i in range(num_chunks): |
| 610 | start_idx = i * chunk_size |
| 611 | end_idx = min(start_idx + chunk_size, num_tokens) |
| 612 | chunk_result = reference_moe( |
| 613 | router_logits[start_idx:end_idx].to(torch.float32), |
| 614 | topk, |
| 615 | num_experts, |
| 616 | hidden_states_ref[start_idx:end_idx], |
| 617 | w13_ref, |
| 618 | bias13_ref, |
| 619 | w2_ref, |
| 620 | bias2_ref, |
| 621 | alpha, |
| 622 | beta, |
| 623 | limit, |
| 624 | act_type, |
| 625 | activation="swiglu", |
| 626 | use_interleaved_layout=False, |
| 627 | ) |
| 628 | ref_result[start_idx:end_idx].copy_(chunk_result) |
| 629 | |
| 630 | # trtllm-gen result |
| 631 | if alpha is not None: |
| 632 | alpha = torch.full((num_experts,), alpha, device=hidden_states.device) |
| 633 | if limit is not None: |
| 634 | limit = torch.full((num_experts,), limit, device=hidden_states.device) |
| 635 | if beta is not None: |
| 636 | beta = torch.full((num_experts,), beta, device=hidden_states.device) |
| 637 | tg_result = tg_mxfp4_moe( |
| 638 | router_logits=router_logits, |
| 639 | topk=topk, |
| 640 | num_experts=num_experts, |
| 641 | intermediate_size=intermediate_size, |
| 642 | hidden_size=hidden_size, |
| 643 | hidden_states=hidden_states, |
| 644 | hidden_states_scale=hidden_states_scale, |
| 645 | w13_weight=w13, |
| 646 | w13_weight_scale=w13_scale, |
| 647 | w13_bias=bias13, |
| 648 | w2_weight=w2, |
| 649 | w2_weight_scale=w2_scale, |
| 650 | w2_bias=bias2, |
| 651 | act_type=act_type, |
| 652 | alpha=alpha, |
| 653 | beta=beta, |
| 654 | limit=limit, |
| 655 | transpose_optimized=transpose_optimized, |
| 656 | ) |
| 657 | # relatively loose check since the mxfp4 quantization is less accurate |
| 658 | check_accuracy(ref_result, tg_result, atol=0, rtol=0.3, percent=0.8) |
| 659 | |
| 660 | |
| 661 | @pytest.mark.parametrize( |
| 662 | "monolithic,num_tokens,chunk_size", |
| 663 | [(False, 1, None), (False, 16, None), (False, 16, 5), (True, 16, None)], |
| 664 | ) |
| 665 | @pytest.mark.skipif( |
| 666 | not TRTLLM_GEN_MXFP4_AVAILABLE, |
| 667 | reason="nvidia gpu and compute capability sm100 is required for this test", |
| 668 | ) |
| 669 | def test_trtllm_mxfp4_deferred_finalize( |
| 670 | monolithic: bool, |
| 671 | num_tokens: int, |
| 672 | chunk_size: int | None, |
| 673 | monkeypatch: pytest.MonkeyPatch, |
| 674 | workspace_init, |
| 675 | ): |
| 676 | """TRTLLM-Gen MXFP4 experts can leave the top-k finalize to the caller.""" |
| 677 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 678 | from vllm.model_executor.layers.fused_moe.config import ( |
| 679 | FusedMoEConfig, |
| 680 | FusedMoEParallelConfig, |
| 681 | RoutingMethodType, |
| 682 | ) |
| 683 | from vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe import ( |
| 684 | TrtLlmMxfp4ExpertsModular, |
| 685 | TrtLlmMxfp4ExpertsMonolithic, |
| 686 | ) |
| 687 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 688 | Mxfp4MoeBackend, |
| 689 | convert_weight_to_mxfp4_moe_kernel_format, |
| 690 | make_mxfp4_moe_kernel, |
| 691 | make_mxfp4_moe_quant_config, |
| 692 | ) |
| 693 | |
| 694 | num_experts, topk, hidden_size, intermediate_size = 32, 4, 1024, 512 |
| 695 | device = torch.device("cuda:0") |
| 696 | torch.manual_seed(0) |
| 697 | if chunk_size is not None: |
| 698 | monkeypatch.setattr( |
| 699 | TrtLlmMxfp4ExpertsModular, |
| 700 | "_max_supported_tokens", |
| 701 | lambda self, top_k, global_num_experts: chunk_size, |
| 702 | ) |
| 703 | |
| 704 | def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 705 | w_q, w_scale = fp4_quantize( |
| 706 | w, |
| 707 | torch.tensor(1.0, device=device), |
| 708 | 32, |
| 709 | sf_use_ue8m0=True, |
| 710 | is_sf_swizzled_layout=False, |
| 711 | ) |
| 712 | return ( |
| 713 | w_q.reshape(*w.shape[:-1], -1), |
| 714 | w_scale.view(torch.uint8).reshape(*w.shape[:-1], -1), |
| 715 | ) |
| 716 | |
| 717 | w13, w13_scale = quantize( |
| 718 | torch.randn( |
| 719 | num_experts, |
| 720 | 2 * intermediate_size, |
| 721 | hidden_size, |
| 722 | device=device, |
| 723 | dtype=torch.bfloat16, |
| 724 | ) |
| 725 | ) |
| 726 | w2, w2_scale = quantize( |
| 727 | torch.randn( |
| 728 | num_experts, |
| 729 | hidden_size, |
| 730 | intermediate_size, |
| 731 | device=device, |
| 732 | dtype=torch.bfloat16, |
| 733 | ) |
| 734 | ) |
| 735 | backend = Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8 |
| 736 | w13, w2, w13_scale, w2_scale, _, _ = convert_weight_to_mxfp4_moe_kernel_format( |
| 737 | backend, |
| 738 | torch.nn.Module(), |
| 739 | w13, |
| 740 | w2, |
| 741 | w13_scale, |
| 742 | w2_scale, |
| 743 | _cache_permute_indices={}, |
| 744 | ) |
| 745 | # One rank of a TP group, which deferral needs. |
| 746 | moe_config = FusedMoEConfig( |
| 747 | num_experts=num_experts, |
| 748 | experts_per_token=topk, |
| 749 | hidden_dim=hidden_size, |
| 750 | intermediate_size=2 * intermediate_size, |
| 751 | num_local_experts=num_experts, |
| 752 | num_logical_experts=num_experts, |
| 753 | activation=MoEActivation.SILU, |
| 754 | device=device, |
| 755 | routing_method=RoutingMethodType.Renormalize, |
| 756 | moe_parallel_config=replace( |
| 757 | FusedMoEParallelConfig.make_no_parallel(), tp_size=2 |
| 758 | ), |
| 759 | in_dtype=torch.bfloat16, |
| 760 | ) |
| 761 | kernel = make_mxfp4_moe_kernel( |
| 762 | make_mxfp4_moe_quant_config(backend, w1_scale=w13_scale, w2_scale=w2_scale), |
| 763 | moe_config, |
| 764 | TrtLlmMxfp4ExpertsMonolithic if monolithic else TrtLlmMxfp4ExpertsModular, |
| 765 | backend, |
| 766 | ) |
| 767 | |
| 768 | hidden_states = torch.randn( |
| 769 | num_tokens, hidden_size, device=device, dtype=torch.bfloat16 |
| 770 | ) |
| 771 | if monolithic: |
| 772 | router_logits = torch.randn(num_tokens, num_experts, device=device) |
| 773 | check_deferred_moe_finalize( |
| 774 | moe_config, |
| 775 | lambda: kernel.apply_monolithic( |
| 776 | hidden_states, |
| 777 | w13, |
| 778 | w2, |
| 779 | router_logits, |
| 780 | activation=MoEActivation.SILU, |
| 781 | global_num_experts=num_experts, |
| 782 | expert_map=None, |
| 783 | apply_router_weight_on_input=False, |
| 784 | ), |
| 785 | ) |
| 786 | return |
| 787 | |
| 788 | topk_ids = torch.rand(num_tokens, num_experts, device=device).argsort(dim=-1) |
| 789 | topk_ids = topk_ids[:, :topk].to(torch.int32) |
| 790 | topk_weights = torch.rand(num_tokens, topk, device=device).softmax(dim=-1) |
| 791 | check_deferred_moe_finalize( |
| 792 | moe_config, |
| 793 | lambda: kernel.apply( |
| 794 | hidden_states=hidden_states, |
| 795 | w1=w13, |
| 796 | w2=w2, |
| 797 | topk_weights=topk_weights, |
| 798 | topk_ids=topk_ids, |
| 799 | activation=MoEActivation.SILU, |
| 800 | global_num_experts=num_experts, |
| 801 | expert_map=None, |
| 802 | apply_router_weight_on_input=False, |
| 803 | ), |
| 804 | router_weights=topk_weights, |
| 805 | chunked=chunk_size is not None, |
| 806 | ) |
| 807 | |
| 808 | |
| 809 | @pytest.mark.parametrize("topk", [1, 4]) |
| 810 | @pytest.mark.parametrize("num_experts", [32]) |
| 811 | @pytest.mark.parametrize("num_tokens", [1, 128]) |
| 812 | @pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)]) |
| 813 | @pytest.mark.parametrize("alpha,beta,limit", [(1.0, 1.0, None), (1.702, 1.0, 7.0)]) |
| 814 | @pytest.mark.skipif( |
| 815 | not HOPPER_MXFP4_BF16_AVAILABLE, |
| 816 | reason="nvidia gpu sm90 and flashinfer are required for this test", |
| 817 | ) |
| 818 | def test_flashinfer_cutlass_mxfp4_fused_moe( |
| 819 | topk: int, |
| 820 | num_experts: int, |
| 821 | num_tokens: int, |
| 822 | intermediate_size: int, |
| 823 | hidden_size: int, |
| 824 | alpha: float, |
| 825 | beta: float, |
| 826 | limit: float | None, |
| 827 | ): |
| 828 | torch.manual_seed(42) |
| 829 | device = "cuda:0" |
| 830 | |
| 831 | # Inputs |
| 832 | hidden_states = torch.randn( |
| 833 | num_tokens, hidden_size, device=device, dtype=torch.bfloat16 |
| 834 | ) |
| 835 | # Random MXFP4 weights and scales (uint8), contiguous [w1; w3] |
| 836 | w13_q = torch.randint( |
| 837 | 0, |
| 838 | 256, |
| 839 | (num_experts, 2 * intermediate_size, hidden_size // 2), |
| 840 | device=device, |
| 841 | dtype=torch.uint8, |
| 842 | ) |
| 843 | w13_scale = torch.randint( |
| 844 | 118, |
| 845 | 123, |
| 846 | (num_experts, 2 * intermediate_size, hidden_size // 32), |
| 847 | device=device, |
| 848 | dtype=torch.uint8, |
| 849 | ) |
| 850 | |
| 851 | w2_q = torch.randint( |
| 852 | 0, |
| 853 | 256, |
| 854 | (num_experts, hidden_size, intermediate_size // 2), |
| 855 | device=device, |
| 856 | dtype=torch.uint8, |
| 857 | ) |
| 858 | w2_scale = torch.randint( |
| 859 | 118, |
| 860 | 123, |
| 861 | (num_experts, hidden_size, intermediate_size // 32), |
| 862 | device=device, |
| 863 | dtype=torch.uint8, |
| 864 | ) |
| 865 | # Bias contiguous [b1; b3] |
| 866 | bias13 = ( |
| 867 | torch.randn( |
| 868 | num_experts, 2 * intermediate_size, device=device, dtype=torch.bfloat16 |
| 869 | ) |
| 870 | * 10 |
| 871 | ) |
| 872 | bias2 = ( |
| 873 | torch.randn(num_experts, hidden_size, device=device, dtype=torch.bfloat16) * 10 |
| 874 | ) |
| 875 | router_logits = torch.rand( |
| 876 | num_tokens, num_experts, dtype=torch.float32, device=device |
| 877 | ) |
| 878 | |
| 879 | w13_ref = mxfp4_dequantize(w13_q.clone(), w13_scale.clone()).reshape( |
| 880 | num_experts, 2 * intermediate_size, hidden_size |
| 881 | ) |
| 882 | w2_ref = mxfp4_dequantize(w2_q.clone(), w2_scale.clone()).reshape( |
| 883 | num_experts, hidden_size, intermediate_size |
| 884 | ) |
| 885 | ref = reference_moe( |
| 886 | router_logits.to(torch.float32), |
| 887 | topk, |
| 888 | num_experts, |
| 889 | hidden_states.to(torch.float32), |
| 890 | w13_ref, |
| 891 | bias13.to(torch.float32), |
| 892 | w2_ref, |
| 893 | bias2.to(torch.float32), |
| 894 | alpha, |
| 895 | beta, |
| 896 | limit, |
| 897 | "bf16", |
| 898 | activation="swiglu", |
| 899 | use_interleaved_layout=False, |
| 900 | ) |
| 901 | |
| 902 | from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe |
| 903 | |
| 904 | # Swap halves to arrange as [w3; w1] (kernel expectation) |
| 905 | w1_w, w3_w = torch.chunk(w13_q, 2, dim=1) |
| 906 | w13_q_swapped = torch.cat([w3_w, w1_w], dim=1) |
| 907 | |
| 908 | # SM90 mixed-input GEMM expects weights/scales in an interleaved layout; |
| 909 | # without it the FP4->BF16 LUT reads bytes from wrong positions for K>128. |
| 910 | from flashinfer.fused_moe import ( |
| 911 | interleave_moe_scales_for_sm90_mixed_gemm, |
| 912 | interleave_moe_weights_for_sm90_mixed_gemm, |
| 913 | ) |
| 914 | |
| 915 | w13_q_swapped = interleave_moe_weights_for_sm90_mixed_gemm( |
| 916 | w13_q_swapped, quant_type="fp4" |
| 917 | ) |
| 918 | w2_q = interleave_moe_weights_for_sm90_mixed_gemm(w2_q, quant_type="fp4") |
| 919 | |
| 920 | b1, b3 = torch.chunk(bias13.to(torch.float32), 2, dim=-1) |
| 921 | w13_b = torch.cat([b3, b1], dim=-1).to(torch.bfloat16) |
| 922 | |
| 923 | w1_s, w3_s = torch.chunk(w13_scale, 2, dim=1) |
| 924 | w13_s = torch.cat([w3_s, w1_s], dim=1) |
| 925 | w13_s_inter = interleave_moe_scales_for_sm90_mixed_gemm(w13_s) |
| 926 | w2_s_inter = interleave_moe_scales_for_sm90_mixed_gemm(w2_scale) |
| 927 | |
| 928 | routing_weights = torch.nn.functional.softmax( |
| 929 | router_logits, dim=1, dtype=torch.float32 |
| 930 | ) |
| 931 | token_final_scales, token_selected_experts = torch.topk( |
| 932 | routing_weights, topk, dim=-1 |
| 933 | ) |
| 934 | token_final_scales = token_final_scales / token_final_scales.sum( |
| 935 | dim=-1, keepdim=True |
| 936 | ) |
| 937 | token_selected_experts = token_selected_experts.to(torch.int).contiguous() |
| 938 | |
| 939 | out = torch.empty_like(hidden_states, dtype=torch.bfloat16) |
| 940 | if alpha is not None: |
| 941 | alpha = torch.full((num_experts,), alpha, device=hidden_states.device) |
| 942 | if beta is not None: |
| 943 | beta = torch.full((num_experts,), beta, device=hidden_states.device) |
| 944 | if limit is not None: |
| 945 | limit = torch.full((num_experts,), limit, device=hidden_states.device) |
| 946 | |
| 947 | _ = flashinfer_cutlass_fused_moe( |
| 948 | input=hidden_states, |
| 949 | token_selected_experts=token_selected_experts, |
| 950 | token_final_scales=token_final_scales, |
| 951 | fc1_expert_weights=w13_q_swapped, |
| 952 | fc2_expert_weights=w2_q, |
| 953 | output_dtype=torch.bfloat16, |
| 954 | output=out, |
| 955 | quant_scales=[w13_s_inter.to(torch.uint8), w2_s_inter.to(torch.uint8)], |
| 956 | fc1_expert_biases=w13_b, |
| 957 | fc2_expert_biases=bias2.to(torch.bfloat16), |
| 958 | swiglu_alpha=alpha, |
| 959 | swiglu_beta=beta, |
| 960 | swiglu_limit=limit, |
| 961 | tp_size=1, |
| 962 | tp_rank=0, |
| 963 | ep_size=1, |
| 964 | ep_rank=0, |
| 965 | use_w4_group_scaling=True, |
| 966 | ) |
| 967 | |
| 968 | # Allow some mismatch due to MXFP4 quantization |
| 969 | check_accuracy(ref, out, atol=0, rtol=0.3, percent=0.8) |
| 970 | |
| 971 | |
| 972 | @pytest.mark.parametrize("topk", [1, 4]) |
| 973 | @pytest.mark.parametrize("num_experts", [32]) |
| 974 | @pytest.mark.parametrize("num_tokens", [1, 128]) |
| 975 | @pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)]) |
| 976 | @pytest.mark.parametrize("alpha,beta,limit", [(1.0, 1.0, None), (1.702, 1.0, 7.0)]) |
| 977 | @pytest.mark.skipif( |
| 978 | not ( |
| 979 | current_platform.is_cuda() |
| 980 | and current_platform.is_device_capability_family(100) |
| 981 | and has_flashinfer() |
| 982 | ), |
| 983 | reason="NVIDIA GPU sm100 and flashinfer are required for this test", |
| 984 | ) |
| 985 | def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe( |
| 986 | topk: int, |
| 987 | num_experts: int, |
| 988 | num_tokens: int, |
| 989 | intermediate_size: int, |
| 990 | hidden_size: int, |
| 991 | alpha: float | None, |
| 992 | beta: float | None, |
| 993 | limit: float | None, |
| 994 | ): |
| 995 | torch.manual_seed(42) |
| 996 | device = "cuda:0" |
| 997 | |
| 998 | # Inputs |
| 999 | hidden_states = torch.randn( |
| 1000 | num_tokens, hidden_size, device=device, dtype=torch.bfloat16 |
| 1001 | ) |
| 1002 | # Float weights in w13 format [w1; w3] |
| 1003 | w13 = ( |
| 1004 | torch.randn( |
| 1005 | num_experts, |
| 1006 | 2 * intermediate_size, |
| 1007 | hidden_size, |
| 1008 | device=device, |
| 1009 | dtype=torch.bfloat16, |
| 1010 | ) |
| 1011 | / 10 |
| 1012 | ) |
| 1013 | w2 = ( |
| 1014 | torch.randn( |
| 1015 | num_experts, |
| 1016 | hidden_size, |
| 1017 | intermediate_size, |
| 1018 | device=device, |
| 1019 | dtype=torch.bfloat16, |
| 1020 | ) |
| 1021 | / 10 |
| 1022 | ) |
| 1023 | # Bias contiguous [b1; b3] |
| 1024 | bias13 = ( |
| 1025 | torch.randn( |
| 1026 | num_experts, 2 * intermediate_size, device=device, dtype=torch.bfloat16 |
| 1027 | ) |
| 1028 | * 10 |
| 1029 | ) |
| 1030 | bias2 = ( |
| 1031 | torch.randn(num_experts, hidden_size, device=device, dtype=torch.bfloat16) * 10 |
| 1032 | ) |
| 1033 | router_logits = torch.rand( |
| 1034 | num_tokens, num_experts, dtype=torch.float32, device=device |
| 1035 | ) |
| 1036 | |
| 1037 | # Quantize weights to MXFP4 per expert (SM100 path) |
| 1038 | from flashinfer import mxfp4_quantize |
| 1039 | |
| 1040 | def quant_mxfp4_batches(a: torch.Tensor, e: int): |
| 1041 | qs, sfs = [], [] |
| 1042 | for i in range(e): |
| 1043 | q, sf = mxfp4_quantize(a[i].cuda()) |
| 1044 | qs.append(q) |
| 1045 | sfs.append(sf) |
| 1046 | return torch.stack(qs), torch.stack(sfs) |
| 1047 | |
| 1048 | def dequant_mxfp4_batches(mat_fp4: torch.Tensor, scale_tensor: torch.Tensor): |
| 1049 | num_batches = mat_fp4.size(0) |
| 1050 | scale_tensor = scale_tensor.view(num_batches, -1) |
| 1051 | from flashinfer import mxfp4_dequantize |
| 1052 | |
| 1053 | return torch.stack( |
| 1054 | [ |
| 1055 | mxfp4_dequantize(mat_fp4[b, :, :], scale_tensor[b, :]) |
| 1056 | for b in range(num_batches) |
| 1057 | ] |
| 1058 | ) |
| 1059 | |
| 1060 | w13_q, w13_scale = quant_mxfp4_batches(w13, num_experts) |
| 1061 | w2_q, w2_scale = quant_mxfp4_batches(w2, num_experts) |
| 1062 | |
| 1063 | # Reference result using dequantized tensors and reference_moe |
| 1064 | w13_ref = ( |
| 1065 | dequant_mxfp4_batches( |
| 1066 | w13_q.view(torch.uint8), w13_scale.view(torch.uint8).reshape(-1) |
| 1067 | ) |
| 1068 | .to(torch.float32) |
| 1069 | .reshape(num_experts, 2 * intermediate_size, hidden_size) |
| 1070 | .to(device) |
| 1071 | ) |
| 1072 | w2_ref = ( |
| 1073 | dequant_mxfp4_batches( |
| 1074 | w2_q.view(torch.uint8), w2_scale.view(torch.uint8).reshape(-1) |
| 1075 | ) |
| 1076 | .to(torch.float32) |
| 1077 | .reshape(num_experts, hidden_size, intermediate_size) |
| 1078 | .to(device) |
| 1079 | ) |
| 1080 | |
| 1081 | # Quantize activations for SM100 path and dequantize for reference |
| 1082 | hidden_states_q, hidden_states_sf = mxfp8_quantize(hidden_states, True, 32) |
| 1083 | # Reference uses BF16 input but quantizes intermediate activation to MXFP8 |
| 1084 | ref = reference_moe( |
| 1085 | router_logits.to(torch.float32), |
| 1086 | topk, |
| 1087 | num_experts, |
| 1088 | hidden_states.to(torch.float32), |
| 1089 | w13_ref, |
| 1090 | bias13.to(torch.float32), |
| 1091 | w2_ref, |
| 1092 | bias2.to(torch.float32), |
| 1093 | alpha, |
| 1094 | beta, |
| 1095 | limit, |
| 1096 | "mxfp8", |
| 1097 | activation="swiglu", |
| 1098 | use_interleaved_layout=False, |
| 1099 | ) |
| 1100 | |
| 1101 | # Prepare inputs for FlashInfer CUTLASS fused MoE |
| 1102 | from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe |
| 1103 | |
| 1104 | # Swap halves to arrange as [w3; w1] (kernel expectation) |
| 1105 | w1_w, w3_w = torch.chunk(w13_q, 2, dim=1) |
| 1106 | w13_q_swapped = torch.cat([w3_w, w1_w], dim=1) |
| 1107 | |
| 1108 | # Swap scales halves to match swapped weights |
| 1109 | s1, s3 = torch.chunk(w13_scale, 2, dim=1) |
| 1110 | w13_scale_swapped = torch.cat([s3, s1], dim=1) |
| 1111 | |
| 1112 | b1, b3 = torch.chunk(bias13.to(torch.float32), 2, dim=-1) |
| 1113 | w13_b = torch.cat([b3, b1], dim=-1).to(torch.bfloat16) |
| 1114 | |
| 1115 | # Build routing for kernel |
| 1116 | routing_weights = torch.nn.functional.softmax( |
| 1117 | router_logits, dim=1, dtype=torch.float32 |
| 1118 | ) |
| 1119 | token_final_scales, token_selected_experts = torch.topk( |
| 1120 | routing_weights, topk, dim=-1 |
| 1121 | ) |
| 1122 | token_final_scales = token_final_scales / token_final_scales.sum( |
| 1123 | dim=-1, keepdim=True |
| 1124 | ) |
| 1125 | token_selected_experts = token_selected_experts.to(torch.int).contiguous() |
| 1126 | |
| 1127 | out = torch.empty_like(hidden_states, dtype=torch.bfloat16) |
| 1128 | if alpha is not None: |
| 1129 | alpha_t = torch.full((num_experts,), alpha, device=hidden_states.device) |
| 1130 | else: |
| 1131 | alpha_t = None |
| 1132 | if beta is not None: |
| 1133 | beta_t = torch.full((num_experts,), beta, device=hidden_states.device) |
| 1134 | else: |
| 1135 | beta_t = None |
| 1136 | if limit is not None: |
| 1137 | limit_t = torch.full((num_experts,), limit, device=hidden_states.device) |
| 1138 | else: |
| 1139 | limit_t = None |
| 1140 | |
| 1141 | # Quant scales for SM100 MXFP8+MXFP4 path |
| 1142 | fake_input_scale = torch.ones(num_experts, device=device) |
| 1143 | quant_scales = [ |
| 1144 | w13_scale_swapped.view(torch.int32), |
| 1145 | fake_input_scale, |
| 1146 | w2_scale.view(torch.int32), |
| 1147 | fake_input_scale, |
| 1148 | ] |
| 1149 | |
| 1150 | _ = flashinfer_cutlass_fused_moe( |
| 1151 | input=hidden_states_q, |
| 1152 | token_selected_experts=token_selected_experts, |
| 1153 | token_final_scales=token_final_scales, |
| 1154 | fc1_expert_weights=w13_q_swapped.contiguous().view(torch.long), |
| 1155 | fc2_expert_weights=w2_q.contiguous().view(torch.long), |
| 1156 | output_dtype=torch.bfloat16, |
| 1157 | output=out, |
| 1158 | quant_scales=quant_scales, |
| 1159 | fc1_expert_biases=w13_b, |
| 1160 | fc2_expert_biases=bias2.to(torch.bfloat16), |
| 1161 | swiglu_alpha=alpha_t, |
| 1162 | swiglu_beta=beta_t, |
| 1163 | swiglu_limit=limit_t, |
| 1164 | tp_size=1, |
| 1165 | tp_rank=0, |
| 1166 | ep_size=1, |
| 1167 | ep_rank=0, |
| 1168 | use_mxfp8_act_scaling=True, |
| 1169 | input_sf=hidden_states_sf, |
| 1170 | ) |
| 1171 | |
| 1172 | # Allow some mismatch due to MXFP4 quantization |
| 1173 | check_accuracy(ref, out, atol=0, rtol=0.3, percent=0.8) |
| 1174 | |
| 1175 | |
| 1176 | @pytest.mark.parametrize( |
| 1177 | ("backend_name", "interleaves_weights"), |
| 1178 | [ |
| 1179 | ("FLASHINFER_CUTLASS_MXFP4_MXFP8", False), |
| 1180 | ("FLASHINFER_CUTLASS_MXFP4_BF16", True), |
| 1181 | ], |
| 1182 | ) |
| 1183 | @pytest.mark.skipif(not has_flashinfer(), reason="flashinfer is required") |
| 1184 | def test_convert_standard_mxfp4_weights_for_flashinfer_cutlass( |
| 1185 | monkeypatch, backend_name, interleaves_weights |
| 1186 | ): |
| 1187 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 1188 | Mxfp4MoeBackend, |
| 1189 | convert_weight_to_mxfp4_moe_kernel_format, |
| 1190 | ) |
| 1191 | |
| 1192 | w13 = torch.arange(16, dtype=torch.uint8).reshape(1, 4, 4) |
| 1193 | w2 = torch.arange(8, dtype=torch.uint8).reshape(1, 4, 2) |
| 1194 | w13_scale = torch.arange(8, dtype=torch.uint8).reshape(1, 4, 2) |
| 1195 | w2_scale = torch.arange(4, dtype=torch.uint8).reshape(1, 4, 1) |
| 1196 | w13_bias = torch.arange(4, dtype=torch.float32).reshape(1, 4) |
| 1197 | w2_bias = torch.arange(4, dtype=torch.float32).reshape(1, 4) |
| 1198 | interleaved_weights = [] |
| 1199 | interleaved_scales = [] |
| 1200 | |
| 1201 | def record_weight_interleave(weight, quant_type): |
| 1202 | assert quant_type == "fp4" |
| 1203 | interleaved_weights.append(weight.clone()) |
| 1204 | return weight |
| 1205 | |
| 1206 | def record_scale_interleave(scale): |
| 1207 | interleaved_scales.append(scale.clone()) |
| 1208 | return scale |
| 1209 | |
| 1210 | monkeypatch.setattr("flashinfer.block_scale_interleave", record_scale_interleave) |
| 1211 | monkeypatch.setattr( |
| 1212 | "flashinfer.fused_moe.interleave_moe_weights_for_sm90_mixed_gemm", |
| 1213 | record_weight_interleave, |
| 1214 | ) |
| 1215 | monkeypatch.setattr( |
| 1216 | "flashinfer.fused_moe.interleave_moe_scales_for_sm90_mixed_gemm", |
| 1217 | record_scale_interleave, |
| 1218 | ) |
| 1219 | converted = convert_weight_to_mxfp4_moe_kernel_format( |
| 1220 | getattr(Mxfp4MoeBackend, backend_name), |
| 1221 | types.SimpleNamespace(), |
| 1222 | w13, |
| 1223 | w2, |
| 1224 | w13_scale, |
| 1225 | w2_scale, |
| 1226 | w13_bias, |
| 1227 | w2_bias, |
| 1228 | ) |
| 1229 | |
| 1230 | expected_w13 = torch.cat([w13[:, 2:], w13[:, :2]], dim=1) |
| 1231 | expected_w13_scale = torch.cat([w13_scale[:, 2:], w13_scale[:, :2]], dim=1) |
| 1232 | torch.testing.assert_close(converted[0], expected_w13) |
| 1233 | torch.testing.assert_close(converted[1], w2) |
| 1234 | torch.testing.assert_close(converted[2], expected_w13_scale) |
| 1235 | torch.testing.assert_close(converted[3], w2_scale) |
| 1236 | expected_w13_bias = torch.cat([w13_bias[:, 2:], w13_bias[:, :2]], dim=1).to( |
| 1237 | torch.bfloat16 |
| 1238 | ) |
| 1239 | torch.testing.assert_close(converted[4], expected_w13_bias) |
| 1240 | torch.testing.assert_close(converted[5], w2_bias.to(torch.bfloat16)) |
| 1241 | |
| 1242 | if interleaves_weights: |
| 1243 | assert len(interleaved_weights) == 2 |
| 1244 | torch.testing.assert_close(interleaved_weights[0], expected_w13) |
| 1245 | torch.testing.assert_close(interleaved_weights[1], w2) |
| 1246 | else: |
| 1247 | assert not interleaved_weights |
| 1248 | assert len(interleaved_scales) == 2 |
| 1249 | torch.testing.assert_close(interleaved_scales[0], expected_w13_scale) |
| 1250 | torch.testing.assert_close(interleaved_scales[1], w2_scale) |
| 1251 | |
| 1252 | |
| 1253 | def test_gpt_oss_quant_config_supplies_clamped_swiglu_params(): |
| 1254 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend |
| 1255 | from vllm.model_executor.layers.quantization.mxfp4 import GptOssMxfp4MoEMethod |
| 1256 | |
| 1257 | fake_method = types.SimpleNamespace( |
| 1258 | mxfp4_backend=Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, |
| 1259 | w13_precision_config=None, |
| 1260 | w2_precision_config=None, |
| 1261 | ) |
| 1262 | fake_layer = types.SimpleNamespace( |
| 1263 | w13_weight_scale=torch.zeros(2, 4, 2, dtype=torch.uint8), |
| 1264 | w2_weight_scale=torch.zeros(2, 4, 2, dtype=torch.uint8), |
| 1265 | ) |
| 1266 | |
| 1267 | quant_config = GptOssMxfp4MoEMethod.get_fused_moe_quant_config( |
| 1268 | fake_method, fake_layer |
| 1269 | ) |
| 1270 | |
| 1271 | assert quant_config is not None |
| 1272 | assert quant_config.gemm1_alpha == 1.702 |
| 1273 | assert quant_config.gemm1_beta == 1.0 |
| 1274 | assert quant_config.gemm1_clamp_limit == 7.0 |
| 1275 | |
| 1276 | |
| 1277 | @pytest.mark.skipif( |
| 1278 | not current_platform.is_cuda(), |
| 1279 | reason="CUDA is required for FlashInferExperts parameter storage", |
| 1280 | ) |
| 1281 | @pytest.mark.parametrize("supplied", [False, True]) |
| 1282 | def test_flashinfer_experts_swiglu_params_follow_quant_config(supplied: bool): |
| 1283 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 1284 | from vllm.model_executor.layers.fused_moe.config import ( |
| 1285 | FusedMoEConfig, |
| 1286 | FusedMoEParallelConfig, |
| 1287 | RoutingMethodType, |
| 1288 | mxfp4_w4a16_moe_quant_config, |
| 1289 | ) |
| 1290 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( |
| 1291 | FlashInferExperts, |
| 1292 | ) |
| 1293 | |
| 1294 | num_experts, intermediate_size, hidden_size = 2, 128, 256 |
| 1295 | w1_scale = torch.zeros( |
| 1296 | (num_experts, 2 * intermediate_size, hidden_size // 32), |
| 1297 | dtype=torch.uint8, |
| 1298 | device="cuda", |
| 1299 | ) |
| 1300 | w2_scale = torch.zeros( |
| 1301 | (num_experts, hidden_size, intermediate_size // 32), |
| 1302 | dtype=torch.uint8, |
| 1303 | device="cuda", |
| 1304 | ) |
| 1305 | extra = ( |
| 1306 | {"gemm1_alpha": 0.5, "gemm1_beta": 0.25, "gemm1_clamp_limit": 3.0} |
| 1307 | if supplied |
| 1308 | else {} |
| 1309 | ) |
| 1310 | quant_config = mxfp4_w4a16_moe_quant_config( |
| 1311 | w1_scale=w1_scale, w2_scale=w2_scale, **extra |
| 1312 | ) |
| 1313 | moe_config = FusedMoEConfig( |
| 1314 | num_experts=num_experts, |
| 1315 | experts_per_token=1, |
| 1316 | hidden_dim=hidden_size, |
| 1317 | intermediate_size=intermediate_size, |
| 1318 | num_local_experts=num_experts, |
| 1319 | num_logical_experts=num_experts, |
| 1320 | activation=MoEActivation.SILU, |
| 1321 | device="cuda", |
| 1322 | routing_method=RoutingMethodType.TopK, |
| 1323 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 1324 | in_dtype=torch.bfloat16, |
| 1325 | ) |
| 1326 | experts = FlashInferExperts(moe_config=moe_config, quant_config=quant_config) |
| 1327 | |
| 1328 | if supplied: |
| 1329 | for name, value in ( |
| 1330 | ("gemm1_alpha", 0.5), |
| 1331 | ("gemm1_beta", 0.25), |
| 1332 | ("gemm1_clamp_limit", 3.0), |
| 1333 | ): |
| 1334 | parameter = getattr(experts, name) |
| 1335 | assert isinstance(parameter, torch.Tensor) |
| 1336 | assert parameter.dtype == torch.float32 |
| 1337 | assert parameter.shape == (num_experts,) |
| 1338 | torch.testing.assert_close( |
| 1339 | parameter.cpu(), |
| 1340 | torch.full((num_experts,), value, dtype=torch.float32), |
| 1341 | ) |
| 1342 | else: |
| 1343 | assert experts.gemm1_alpha is None |
| 1344 | assert experts.gemm1_beta is None |
| 1345 | assert experts.gemm1_clamp_limit is None |
| 1346 | |
| 1347 | |
| 1348 | @pytest.mark.parametrize("topk", [1, 4]) |
| 1349 | @pytest.mark.parametrize("num_experts", [32]) |
| 1350 | @pytest.mark.parametrize("num_tokens", [1, 128]) |
| 1351 | @pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)]) |
| 1352 | @pytest.mark.parametrize("is_gated", [True], ids=["gated"]) |
| 1353 | @pytest.mark.skipif( |
| 1354 | not TRTLLM_GEN_MXFP8_AVAILABLE, |
| 1355 | reason="nvidia gpu and compute capability sm100 is required for this test", |
| 1356 | ) |
| 1357 | def test_trtllm_gen_mxfp8_block_scale_moe( |
| 1358 | topk: int, |
| 1359 | num_experts: int, |
| 1360 | num_tokens: int, |
| 1361 | intermediate_size: int, |
| 1362 | hidden_size: int, |
| 1363 | is_gated: bool, |
| 1364 | ): |
| 1365 | torch.manual_seed(42) |
| 1366 | device = "cuda:0" |
| 1367 | |
| 1368 | inter_size = intermediate_size * (2 if is_gated else 1) |
| 1369 | |
| 1370 | hidden_states = ( |
| 1371 | torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) / 20 |
| 1372 | ) |
| 1373 | w13 = ( |
| 1374 | torch.randn( |
| 1375 | num_experts, |
| 1376 | inter_size, |
| 1377 | hidden_size, |
| 1378 | device=device, |
| 1379 | dtype=torch.bfloat16, |
| 1380 | ) |
| 1381 | / 20 |
| 1382 | ) |
| 1383 | w2 = ( |
| 1384 | torch.randn( |
| 1385 | num_experts, |
| 1386 | hidden_size, |
| 1387 | intermediate_size, |
| 1388 | device=device, |
| 1389 | dtype=torch.bfloat16, |
| 1390 | ) |
| 1391 | / 20 |
| 1392 | ) |
| 1393 | router_logits = torch.rand( |
| 1394 | num_tokens, num_experts, dtype=torch.float32, device=device |
| 1395 | ) |
| 1396 | router_logits_kernel = router_logits.to(torch.bfloat16) |
| 1397 | |
| 1398 | # Quantize weights to MXFP8 and normalize scales to [E, M, K//32]. |
| 1399 | w13_q, w13_scale = mxfp8_quantize(w13, is_sf_swizzled_layout=False) |
| 1400 | w2_q, w2_scale = mxfp8_quantize(w2, is_sf_swizzled_layout=False) |
| 1401 | if w13_scale.ndim == 1: |
| 1402 | w13_scale = w13_scale.view( |
| 1403 | num_experts, |
| 1404 | inter_size, |
| 1405 | hidden_size // 32, |
| 1406 | ) |
| 1407 | if w2_scale.ndim == 1: |
| 1408 | w2_scale = w2_scale.view(num_experts, hidden_size, intermediate_size // 32) |
| 1409 | |
| 1410 | # Quantize activations to MXFP8. |
| 1411 | hidden_states_q, hidden_states_scale = mxfp8_quantize( |
| 1412 | hidden_states, is_sf_swizzled_layout=False |
| 1413 | ) |
| 1414 | if hidden_states_scale.ndim == 1: |
| 1415 | hidden_states_scale = hidden_states_scale.view(num_tokens, hidden_size // 32) |
| 1416 | |
| 1417 | # Reference output using dequantized tensors + MXFP8 intermediate quantization. |
| 1418 | w13_ref = mxfp8_dequantize(w13_q, w13_scale).to(torch.float32) |
| 1419 | w2_ref = mxfp8_dequantize(w2_q, w2_scale).to(torch.float32) |
| 1420 | hidden_states_ref = mxfp8_dequantize(hidden_states_q, hidden_states_scale).to( |
| 1421 | torch.float32 |
| 1422 | ) |
| 1423 | bias13 = torch.zeros( |
| 1424 | num_experts, |
| 1425 | intermediate_size * (2 if is_gated else 1), |
| 1426 | device=device, |
| 1427 | ) |
| 1428 | bias2 = torch.zeros(num_experts, hidden_size, device=device) |
| 1429 | ref = reference_moe( |
| 1430 | router_logits_kernel.to(torch.float32), |
| 1431 | topk, |
| 1432 | num_experts, |
| 1433 | hidden_states_ref, |
| 1434 | w13_ref, |
| 1435 | bias13, |
| 1436 | w2_ref, |
| 1437 | bias2, |
| 1438 | alpha=1.0, |
| 1439 | beta=0.0, |
| 1440 | limit=None, |
| 1441 | act_type="mxfp8", |
| 1442 | activation="swiglu" if is_gated else "relu2", |
| 1443 | use_interleaved_layout=False, |
| 1444 | ) |
| 1445 | |
| 1446 | # Shuffle weights/scales with the same indexed layout used by TRTLLM kernels. |
| 1447 | epilogue_tile_m = 128 |
| 1448 | gemm1_weights_shuffled = [] |
| 1449 | gemm1_scales_shuffled = [] |
| 1450 | gemm2_weights_shuffled = [] |
| 1451 | gemm2_scales_shuffled = [] |
| 1452 | for i in range(num_experts): |
| 1453 | w13_rows = intermediate_size * (2 if is_gated else 1) |
| 1454 | w13_interleaved = w13_q[i].clone().reshape(w13_rows, -1) |
| 1455 | w13_scale_interleaved = w13_scale[i].clone().reshape(w13_rows, -1) |
| 1456 | if is_gated: |
| 1457 | w13_interleaved = reorder_rows_for_gated_act_gemm(w13_interleaved) |
| 1458 | w13_scale_interleaved = reorder_rows_for_gated_act_gemm( |
| 1459 | w13_scale_interleaved |
| 1460 | ) |
| 1461 | gemm1_weights_shuffled.append( |
| 1462 | shuffle_matrix_a(w13_interleaved.view(torch.uint8), epilogue_tile_m) |
| 1463 | .contiguous() |
| 1464 | .view(w13_q.dtype) |
| 1465 | ) |
| 1466 | gemm2_weights_shuffled.append( |
| 1467 | shuffle_matrix_a(w2_q[i].view(torch.uint8), epilogue_tile_m) |
| 1468 | .contiguous() |
| 1469 | .view(w2_q.dtype) |
| 1470 | ) |
| 1471 | |
| 1472 | gemm1_scales_shuffled.append( |
| 1473 | shuffle_matrix_sf_a( |
| 1474 | w13_scale_interleaved.view(torch.uint8).reshape(w13_rows, -1), |
| 1475 | epilogue_tile_m, |
| 1476 | ) |
| 1477 | .contiguous() |
| 1478 | .view(w13_scale.dtype) |
| 1479 | ) |
| 1480 | gemm2_scales_shuffled.append( |
| 1481 | shuffle_matrix_sf_a( |
| 1482 | w2_scale[i].view(torch.uint8).reshape(hidden_size, -1), epilogue_tile_m |
| 1483 | ) |
| 1484 | .contiguous() |
| 1485 | .view(w2_scale.dtype) |
| 1486 | ) |
| 1487 | |
| 1488 | out = trtllm_fp8_block_scale_moe( |
| 1489 | routing_logits=router_logits_kernel, |
| 1490 | routing_bias=None, |
| 1491 | hidden_states=hidden_states_q, |
| 1492 | hidden_states_scale=hidden_states_scale, |
| 1493 | gemm1_weights=torch.stack(gemm1_weights_shuffled), |
| 1494 | gemm1_weights_scale=torch.stack(gemm1_scales_shuffled), |
| 1495 | gemm2_weights=torch.stack(gemm2_weights_shuffled), |
| 1496 | gemm2_weights_scale=torch.stack(gemm2_scales_shuffled), |
| 1497 | num_experts=num_experts, |
| 1498 | top_k=topk, |
| 1499 | n_group=None, |
| 1500 | topk_group=None, |
| 1501 | intermediate_size=intermediate_size, |
| 1502 | local_expert_offset=0, |
| 1503 | local_num_experts=num_experts, |
| 1504 | routed_scaling_factor=None, |
| 1505 | routing_method_type=1, # renormalize routing |
| 1506 | use_shuffled_weight=True, |
| 1507 | weight_layout=0, # MajorK |
| 1508 | fp8_quantization_type=Fp8QuantizationType.MxFp8, |
| 1509 | ) |
| 1510 | |
| 1511 | # Block-scale MXFP8 kernels are approximate; require majority close. |
| 1512 | check_accuracy(ref, out, atol=0.1, rtol=0.85, percent=0.8) |
| 1513 | |
| 1514 | |
| 1515 | # ----------------------------------------------------------------------------- |
| 1516 | # ROCm Oracle-based kernel execution tests |
| 1517 | # ----------------------------------------------------------------------------- |
| 1518 | # TODO: Further tighten the accuracy threshold. |
| 1519 | # - More accurate ref moe to include activation quantization |
| 1520 | # - Check aiter kernel accuracy. E.g., quant / dequant details. |
| 1521 | ROCM_BACKEND_CONFIGS = { |
| 1522 | "TRITON": { |
| 1523 | "activation": "SWIGLUOAI", |
| 1524 | "rtol": 0.3, |
| 1525 | "percent": 0.95, |
| 1526 | "requires_aiter": False, |
| 1527 | "requires_gfx950": False, |
| 1528 | }, |
| 1529 | "TRITON_UNFUSED": { |
| 1530 | "activation": "SWIGLUOAI", |
| 1531 | "rtol": 0.3, |
| 1532 | "percent": 0.95, |
| 1533 | "requires_aiter": False, |
| 1534 | "requires_gfx950": False, |
| 1535 | }, |
| 1536 | "AITER_MXFP4_BF16": { |
| 1537 | "activation": "SWIGLUOAI", |
| 1538 | "rtol": 0.1, |
| 1539 | "percent": 0.7, |
| 1540 | "requires_aiter": True, |
| 1541 | "requires_gfx950": True, |
| 1542 | }, |
| 1543 | "AITER_TRITON_MXFP4_BF16": { |
| 1544 | "activation": "SILU", |
| 1545 | "rtol": 0.3, |
| 1546 | "percent": 0.95, |
| 1547 | "requires_aiter": True, |
| 1548 | "requires_gfx950": False, |
| 1549 | "interleaved_layout": True, |
| 1550 | }, |
| 1551 | "AITER_MXFP4_FP8": { |
| 1552 | "activation": "SWIGLUOAI", |
| 1553 | "rtol": 0.5, |
| 1554 | "percent": 0.9, |
| 1555 | "requires_aiter": True, |
| 1556 | "requires_gfx950": True, |
| 1557 | }, |
| 1558 | "AITER_MXFP4_MXFP4": { |
| 1559 | "activation": "SILU", |
| 1560 | "act_type": "mxfp4", |
| 1561 | "rtol": 1.0, |
| 1562 | "percent": 0.8, |
| 1563 | "requires_aiter": True, |
| 1564 | "requires_gfx950": True, |
| 1565 | }, |
| 1566 | } |
| 1567 | |
| 1568 | |
| 1569 | @pytest.mark.parametrize("backend_name", list(ROCM_BACKEND_CONFIGS.keys())) |
| 1570 | @pytest.mark.parametrize("topk", [4]) |
| 1571 | @pytest.mark.parametrize("num_experts", [8]) |
| 1572 | @pytest.mark.parametrize("num_tokens,hidden_size,intermediate_size", [(16, 256, 256)]) |
| 1573 | @pytest.mark.skipif( |
| 1574 | not ROCM_AVAILABLE, |
| 1575 | reason="ROCm is required for this test", |
| 1576 | ) |
| 1577 | @torch.inference_mode() |
| 1578 | def test_rocm_mxfp4_moe_oracle( |
| 1579 | backend_name: str, |
| 1580 | topk: int, |
| 1581 | num_experts: int, |
| 1582 | num_tokens: int, |
| 1583 | hidden_size: int, |
| 1584 | intermediate_size: int, |
| 1585 | monkeypatch: pytest.MonkeyPatch, |
| 1586 | ): |
| 1587 | """Test ROCm MXFP4 MoE using oracle functions. |
| 1588 | |
| 1589 | This test validates that the oracle functions work end-to-end: |
| 1590 | - select_mxfp4_moe_backend() selects a valid backend |
| 1591 | - convert_gpt_oss_weight_to_mxfp4_moe_kernel_format() converts weights without error |
| 1592 | - make_mxfp4_moe_quant_config() builds a valid quant config |
| 1593 | - make_mxfp4_moe_kernel() creates a kernel that runs without error |
| 1594 | - The kernel output is within accuracy tolerance of reference |
| 1595 | """ |
| 1596 | config = ROCM_BACKEND_CONFIGS[backend_name] |
| 1597 | |
| 1598 | # Check platform requirements |
| 1599 | if not ROCM_TRITON_KERNELS_AVAILABLE: |
| 1600 | pytest.skip("triton_kernels required for quantization") |
| 1601 | if config["requires_aiter"] and not ROCM_AITER_AVAILABLE: |
| 1602 | pytest.skip(f"Backend {backend_name} requires AITER") |
| 1603 | if config["requires_gfx950"] and not ROCM_GFX950: |
| 1604 | pytest.skip(f"Backend {backend_name} requires GFX950") |
| 1605 | |
| 1606 | import vllm.distributed.parallel_state as ps |
| 1607 | from vllm.config import VllmConfig, set_current_vllm_config |
| 1608 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 1609 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 1610 | Mxfp4MoeBackend, |
| 1611 | backend_to_kernel_cls, |
| 1612 | convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, |
| 1613 | make_mxfp4_moe_kernel, |
| 1614 | make_mxfp4_moe_quant_config, |
| 1615 | ) |
| 1616 | from vllm.v1.worker.workspace import init_workspace_manager |
| 1617 | |
| 1618 | # Initialize workspace manager (needed for modular kernels) |
| 1619 | init_workspace_manager(torch.accelerator.current_device_index()) |
| 1620 | |
| 1621 | # Set up the TP Group to prevent failure on should_use_cdna4_mx_scale_swizzle check |
| 1622 | monkeypatch.setattr(ps, "_TP", types.SimpleNamespace(world_size=1)) |
| 1623 | |
| 1624 | # AITER must be enabled or aiter_mxfp4_w4a8_moe asserts before dispatch. |
| 1625 | monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True) |
| 1626 | |
| 1627 | # Map string to enum |
| 1628 | backend = Mxfp4MoeBackend[backend_name] |
| 1629 | |
| 1630 | # Get experts class from oracle |
| 1631 | experts_cls_list = backend_to_kernel_cls(backend) |
| 1632 | if experts_cls_list is None or len(experts_cls_list) == 0: |
| 1633 | pytest.skip(f"Backend {backend_name} not available") |
| 1634 | |
| 1635 | # Use first experts class |
| 1636 | experts_cls = experts_cls_list[0] |
| 1637 | |
| 1638 | torch.manual_seed(42) |
| 1639 | dtype = torch.bfloat16 |
| 1640 | device = "cuda:0" |
| 1641 | |
| 1642 | # Create MoE config with Renormalize routing (required by monolithic kernels) |
| 1643 | from vllm.model_executor.layers.fused_moe import FusedMoEConfig |
| 1644 | from vllm.model_executor.layers.fused_moe.config import ( |
| 1645 | FusedMoEParallelConfig, |
| 1646 | RoutingMethodType, |
| 1647 | ) |
| 1648 | |
| 1649 | moe_config = FusedMoEConfig( |
| 1650 | num_experts=num_experts, |
| 1651 | experts_per_token=topk, |
| 1652 | hidden_dim=hidden_size, |
| 1653 | intermediate_size=intermediate_size, |
| 1654 | num_local_experts=num_experts, |
| 1655 | num_logical_experts=num_experts, |
| 1656 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 1657 | activation=MoEActivation[config["activation"]], |
| 1658 | in_dtype=dtype, |
| 1659 | device="cuda", |
| 1660 | routing_method=RoutingMethodType.Renormalize, |
| 1661 | ) |
| 1662 | |
| 1663 | # Create float weights in checkpoint format: |
| 1664 | # w13: [num_experts, 2*intermediate_size, hidden_size] |
| 1665 | # w2: [num_experts, hidden_size, intermediate_size] |
| 1666 | w13_float = torch.randn( |
| 1667 | num_experts, 2 * intermediate_size, hidden_size, dtype=dtype, device=device |
| 1668 | ) |
| 1669 | w2_float = torch.randn( |
| 1670 | num_experts, hidden_size, intermediate_size, dtype=dtype, device=device |
| 1671 | ) |
| 1672 | |
| 1673 | # dynamic_mxfp4_quant expects 2D input, so reshape 3D weights |
| 1674 | # w13: [E, 2*I, H] -> [E*2*I, H] -> quantize -> [E, 2*I, H//2] |
| 1675 | # w2: [E, H, I] -> [E*H, I] -> quantize -> [E, H, I//2] |
| 1676 | w13_2d = w13_float.reshape(-1, hidden_size) |
| 1677 | w13_quant_2d, w13_scale_2d = dynamic_mxfp4_quant(w13_2d) |
| 1678 | w13_quant = w13_quant_2d.reshape(num_experts, 2 * intermediate_size, -1) |
| 1679 | w13_scale = w13_scale_2d.reshape(num_experts, 2 * intermediate_size, -1) |
| 1680 | |
| 1681 | w2_2d = w2_float.reshape(-1, intermediate_size) |
| 1682 | w2_quant_2d, w2_scale_2d = dynamic_mxfp4_quant(w2_2d) |
| 1683 | w2_quant = w2_quant_2d.reshape(num_experts, hidden_size, -1) |
| 1684 | w2_scale = w2_scale_2d.reshape(num_experts, hidden_size, -1) |
| 1685 | |
| 1686 | # AITER conversion mutates checkpoint-format weights and scales in place. |
| 1687 | # Preserve their original layout for the reference calculation below. |
| 1688 | w13_quant_ref = w13_quant.clone() |
| 1689 | w13_scale_ref = w13_scale.clone() |
| 1690 | w2_quant_ref = w2_quant.clone() |
| 1691 | w2_scale_ref = w2_scale.clone() |
| 1692 | |
| 1693 | w13_bias = torch.randn( |
| 1694 | num_experts, 2 * intermediate_size, dtype=dtype, device=device |
| 1695 | ) |
| 1696 | w2_bias = torch.randn(num_experts, hidden_size, dtype=dtype, device=device) |
| 1697 | |
| 1698 | # Create static input scales for W4A8 backend (AITER_MXFP4_FP8) |
| 1699 | w13_input_scale: torch.Tensor | None = None |
| 1700 | w2_input_scale: torch.Tensor | None = None |
| 1701 | if backend_name == "AITER_MXFP4_FP8": |
| 1702 | # Static FP8 scales: one scale per expert |
| 1703 | w13_input_scale = torch.ones(num_experts, dtype=torch.float32, device=device) |
| 1704 | w2_input_scale = torch.ones(num_experts, dtype=torch.float32, device=device) |
| 1705 | |
| 1706 | # Create mock layer for oracle functions |
| 1707 | class MockLayer: |
| 1708 | w13_weight: torch.Tensor |
| 1709 | w2_weight: torch.Tensor |
| 1710 | w13_weight_scale: torch.Tensor |
| 1711 | w2_weight_scale: torch.Tensor |
| 1712 | w13_input_scale: torch.Tensor | None |
| 1713 | w2_input_scale: torch.Tensor | None |
| 1714 | |
| 1715 | layer = MockLayer() |
| 1716 | layer.w13_weight = w13_quant |
| 1717 | layer.w2_weight = w2_quant |
| 1718 | layer.w13_weight_scale = w13_scale |
| 1719 | layer.w2_weight_scale = w2_scale |
| 1720 | layer.w13_input_scale = w13_input_scale |
| 1721 | layer.w2_input_scale = w2_input_scale |
| 1722 | |
| 1723 | # Convert weights using oracle |
| 1724 | w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = ( |
| 1725 | convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( |
| 1726 | mxfp4_backend=backend, |
| 1727 | layer=layer, |
| 1728 | w13_weight=w13_quant, |
| 1729 | w2_weight=w2_quant, |
| 1730 | w13_weight_scale=w13_scale, |
| 1731 | w2_weight_scale=w2_scale, |
| 1732 | w13_bias=w13_bias, |
| 1733 | w2_bias=w2_bias, |
| 1734 | ) |
| 1735 | ) |
| 1736 | |
| 1737 | # Build quant config using oracle |
| 1738 | quant_config = make_mxfp4_moe_quant_config( |
| 1739 | mxfp4_backend=backend, |
| 1740 | w1_scale=w13_scale_conv, |
| 1741 | w2_scale=w2_scale_conv, |
| 1742 | w1_bias=w13_bias_conv, |
| 1743 | w2_bias=w2_bias_conv, |
| 1744 | a1_scale=w13_input_scale, |
| 1745 | a2_scale=w2_input_scale, |
| 1746 | ) |
| 1747 | |
| 1748 | # Select activation based on backend |
| 1749 | activation_name = str(config["activation"]) |
| 1750 | activation = MoEActivation[activation_name] |
| 1751 | |
| 1752 | # Build kernel using oracle |
| 1753 | assert quant_config is not None, "Failed to create quant config" |
| 1754 | with set_current_vllm_config(VllmConfig()): |
| 1755 | kernel = make_mxfp4_moe_kernel( |
| 1756 | moe_quant_config=quant_config, |
| 1757 | moe_config=moe_config, |
| 1758 | mxfp4_backend=backend, |
| 1759 | experts_cls=experts_cls, |
| 1760 | routing_tables=None, |
| 1761 | ) |
| 1762 | |
| 1763 | # Create inputs |
| 1764 | x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) |
| 1765 | router_logits = torch.randn( |
| 1766 | num_tokens, num_experts, dtype=torch.float32, device=device |
| 1767 | ) |
| 1768 | topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) |
| 1769 | topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) |
| 1770 | |
| 1771 | # Run kernel - use appropriate method based on impl type |
| 1772 | if kernel.is_monolithic: |
| 1773 | # Monolithic impl uses router_logits |
| 1774 | out = kernel.apply_monolithic( |
| 1775 | hidden_states=x, |
| 1776 | w1=w13_conv, |
| 1777 | w2=w2_conv, |
| 1778 | router_logits=router_logits, |
| 1779 | activation=activation, |
| 1780 | global_num_experts=num_experts, |
| 1781 | expert_map=None, |
| 1782 | apply_router_weight_on_input=False, |
| 1783 | ) |
| 1784 | else: |
| 1785 | # Modular impl uses topk_weights and topk_ids |
| 1786 | out = kernel.apply( |
| 1787 | hidden_states=x, |
| 1788 | w1=w13_conv, |
| 1789 | w2=w2_conv, |
| 1790 | topk_weights=topk_weights, |
| 1791 | topk_ids=topk_ids, |
| 1792 | activation=activation, |
| 1793 | global_num_experts=num_experts, |
| 1794 | expert_map=None, |
| 1795 | apply_router_weight_on_input=False, |
| 1796 | ) |
| 1797 | |
| 1798 | # Verify output is valid (no NaN/Inf) and has expected shape |
| 1799 | assert out.shape == (num_tokens, hidden_size), f"Unexpected shape: {out.shape}" |
| 1800 | assert not torch.any(torch.isnan(out)), "Output contains NaN" |
| 1801 | assert not torch.any(torch.isinf(out)), "Output contains Inf" |
| 1802 | |
| 1803 | # Verify output has reasonable magnitude (not all zeros) |
| 1804 | assert out.abs().max() > 0.01, "Output is effectively zero" |
| 1805 | |
| 1806 | # Dequantize weights for reference computation |
| 1807 | w13_dq = mxfp4_dequantize(w13_quant_ref.view(torch.uint8), w13_scale_ref) |
| 1808 | w2_dq = mxfp4_dequantize(w2_quant_ref.view(torch.uint8), w2_scale_ref) |
| 1809 | |
| 1810 | # Determine activation type and layout |
| 1811 | # SWIGLUOAI uses interleaved layout (gate/up alternating) |
| 1812 | # SILU uses chunked layout (first half gate, second half up) |
| 1813 | use_interleaved = bool( |
| 1814 | config.get("interleaved_layout", activation == MoEActivation.SWIGLUOAI) |
| 1815 | ) |
| 1816 | if activation in [MoEActivation.SWIGLUOAI, MoEActivation.SILU]: |
| 1817 | act_name = "swiglu" |
| 1818 | else: |
| 1819 | act_name = "relu2" |
| 1820 | |
| 1821 | ref = reference_moe( |
| 1822 | router_logits, |
| 1823 | topk, |
| 1824 | num_experts, |
| 1825 | x.to(torch.float32), |
| 1826 | w13_dq.to(torch.float32), |
| 1827 | w13_bias.to(torch.float32), |
| 1828 | w2_dq.to(torch.float32), |
| 1829 | w2_bias.to(torch.float32), |
| 1830 | alpha=1.702 if activation == MoEActivation.SWIGLUOAI else 1.0, |
| 1831 | beta=1.0 if activation == MoEActivation.SWIGLUOAI else 0.0, |
| 1832 | limit=7.0 if activation == MoEActivation.SWIGLUOAI else None, |
| 1833 | act_type=str(config.get("act_type", "bf16")), |
| 1834 | activation=act_name, |
| 1835 | use_interleaved_layout=use_interleaved, |
| 1836 | ) |
| 1837 | |
| 1838 | # Compute and print accuracy statistics |
| 1839 | diff = (ref.float() - out.float()).abs() |
| 1840 | rel_diff = diff / (ref.float().abs() + 1e-6) |
| 1841 | |
| 1842 | print(f"\n[{backend_name}] Accuracy statistics:") |
| 1843 | print( |
| 1844 | f" Reference: min={ref.min():.4f}, max={ref.max():.4f}, mean={ref.mean():.4f}" |
| 1845 | ) |
| 1846 | print( |
| 1847 | f" Output: min={out.min():.4f}, max={out.max():.4f}, mean={out.mean():.4f}" |
| 1848 | ) |
| 1849 | print( |
| 1850 | f" Abs diff: min={diff.min():.4f}, max={diff.max():.4f}, " |
| 1851 | f"mean={diff.mean():.4f}" |
| 1852 | ) |
| 1853 | print( |
| 1854 | f" Rel diff: min={rel_diff.min():.4f}, max={rel_diff.max():.4f}, " |
| 1855 | f"mean={rel_diff.mean():.4f}" |
| 1856 | ) |
| 1857 | |
| 1858 | # Check what percentage of values are within various tolerances |
| 1859 | for rtol in [0.1, 0.5, 1.0, 2.0]: |
| 1860 | within_tol = (diff <= rtol * out.float().abs()).float().mean() |
| 1861 | print(f" Within rtol={rtol}: {within_tol * 100:.1f}%") |
| 1862 | |
| 1863 | # Check accuracy using per-backend thresholds |
| 1864 | check_accuracy(ref, out, atol=0.1, rtol=config["rtol"], percent=config["percent"]) |
| 1865 | |
| 1866 | |
| 1867 | # ----------------------------------------------------------------------------- |
| 1868 | # MXFP4 size-rounding tests |
| 1869 | # ----------------------------------------------------------------------------- |
| 1870 | |
| 1871 | |
| 1872 | @pytest.mark.skipif(not ROCM_AVAILABLE, reason="ROCm-specific test") |
| 1873 | @pytest.mark.skipif( |
| 1874 | not is_aiter_found_and_supported(), reason="AITER is not installed or supported" |
| 1875 | ) |
| 1876 | def test_aiter_mxfp4_bf16_silu_preserves_native_tp_shard(): |
| 1877 | """Keep 128-aligned SiLU shards at their native widths.""" |
| 1878 | from vllm.model_executor.layers.fused_moe import FusedMoEConfig |
| 1879 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 1880 | from vllm.model_executor.layers.fused_moe.config import ( |
| 1881 | FusedMoEParallelConfig, |
| 1882 | RoutingMethodType, |
| 1883 | ) |
| 1884 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 1885 | Mxfp4MoeBackend, |
| 1886 | ) |
| 1887 | from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod |
| 1888 | |
| 1889 | moe_parallel_config = replace( |
| 1890 | FusedMoEParallelConfig.make_no_parallel(), |
| 1891 | tp_size=8, |
| 1892 | ) |
| 1893 | moe_config = FusedMoEConfig( |
| 1894 | num_experts=384, |
| 1895 | experts_per_token=6, |
| 1896 | hidden_dim=7168, |
| 1897 | intermediate_size=3072, |
| 1898 | num_local_experts=384, |
| 1899 | num_logical_experts=384, |
| 1900 | activation=MoEActivation.SILU, |
| 1901 | device="cpu", |
| 1902 | routing_method=RoutingMethodType.Renormalize, |
| 1903 | moe_parallel_config=moe_parallel_config, |
| 1904 | in_dtype=torch.bfloat16, |
| 1905 | ) |
| 1906 | method = object.__new__(Mxfp4MoEMethod) |
| 1907 | method.moe = moe_config |
| 1908 | method.mxfp4_backend = Mxfp4MoeBackend.AITER_MXFP4_BF16 |
| 1909 | |
| 1910 | rounded_shape = method.maybe_roundup_sizes( |
| 1911 | hidden_size=7168, |
| 1912 | intermediate_size_per_partition=moe_config.intermediate_size_per_partition, |
| 1913 | act_dtype=torch.bfloat16, |
| 1914 | moe_parallel_config=moe_config.moe_parallel_config, |
| 1915 | ) |
| 1916 | |
| 1917 | assert rounded_shape == (7168, 384) |
| 1918 | |
| 1919 | |
| 1920 | # Emulation needs each per-partition dim rounded up to OCP_MX_BLOCK_SIZE (32); |
| 1921 | # a non-block-aligned shard (e.g. GPT-OSS 2880 // 4 = 720) otherwise truncates |
| 1922 | # the scale buffer and fails weight loading. |
| 1923 | # NOTE: gated to ROCm since it is the emulation backend's current target; |
| 1924 | # remove this skip if the backend is enabled on non-ROCm platforms. |
| 1925 | @pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm") |
| 1926 | @pytest.mark.parametrize( |
| 1927 | "hidden_size,intermediate_size,expected_hidden,expected_intermediate", |
| 1928 | [ |
| 1929 | (2880, 720, 2880, 736), # GPT-OSS TP=4 shard: 720 -> round_up(720, 32) |
| 1930 | (2880, 360, 2880, 384), # GPT-OSS TP=8 shard: 360 -> 384 |
| 1931 | (2880, 2880, 2880, 2880), # already block-aligned: unchanged |
| 1932 | (90, 90, 96, 96), # both dims unaligned |
| 1933 | ], |
| 1934 | ) |
| 1935 | def test_mxfp4_emulation_rounds_up_to_block_size( |
| 1936 | hidden_size: int, |
| 1937 | intermediate_size: int, |
| 1938 | expected_hidden: int, |
| 1939 | expected_intermediate: int, |
| 1940 | ): |
| 1941 | """Emulation must block-align per-partition dims to OCP_MX_BLOCK_SIZE.""" |
| 1942 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 1943 | Mxfp4MoeBackend, |
| 1944 | mxfp4_round_up_hidden_size_and_intermediate_size, |
| 1945 | ) |
| 1946 | from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( |
| 1947 | OCP_MX_BLOCK_SIZE, |
| 1948 | ) |
| 1949 | |
| 1950 | rounded_hidden, rounded_intermediate = ( |
| 1951 | mxfp4_round_up_hidden_size_and_intermediate_size( |
| 1952 | Mxfp4MoeBackend.EMULATION, hidden_size, intermediate_size |
| 1953 | ) |
| 1954 | ) |
| 1955 | |
| 1956 | assert rounded_hidden == expected_hidden |
| 1957 | assert rounded_intermediate == expected_intermediate |
| 1958 | # The block-scale buffer (dim // OCP_MX_BLOCK_SIZE) must not floor-truncate. |
| 1959 | assert rounded_hidden % OCP_MX_BLOCK_SIZE == 0 |
| 1960 | assert rounded_intermediate % OCP_MX_BLOCK_SIZE == 0 |
| 1961 | |
| 1962 | |
| 1963 | def test_select_mxfp4_moe_backend_raises_with_unsupported_reasons( |
| 1964 | monkeypatch: pytest.MonkeyPatch, |
| 1965 | ): |
| 1966 | """select_mxfp4_moe_backend() must raise NotImplementedError, with the |
| 1967 | collected per-backend unsupported reasons in the message, when no |
| 1968 | backend supports the requested deployment configuration. |
| 1969 | """ |
| 1970 | import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle |
| 1971 | from vllm.model_executor.layers.fused_moe import FusedMoEConfig |
| 1972 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 1973 | from vllm.model_executor.layers.fused_moe.config import ( |
| 1974 | FusedMoEParallelConfig, |
| 1975 | RoutingMethodType, |
| 1976 | ) |
| 1977 | |
| 1978 | class UnsupportedExperts: |
| 1979 | @staticmethod |
| 1980 | def is_supported_config( |
| 1981 | cls, moe_config, weight_key, activation_key, activation_format |
| 1982 | ): |
| 1983 | return False, f"unsupported reason for {cls.__name__}" |
| 1984 | |
| 1985 | monkeypatch.setattr( |
| 1986 | mxfp4_oracle, "backend_to_kernel_cls", lambda backend: [UnsupportedExperts] |
| 1987 | ) |
| 1988 | monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None) |
| 1989 | monkeypatch.setattr(current_platform, "is_xpu", lambda: False) |
| 1990 | monkeypatch.setattr(current_platform, "is_cpu", lambda: False) |
| 1991 | |
| 1992 | moe_config = FusedMoEConfig( |
| 1993 | num_experts=8, |
| 1994 | experts_per_token=2, |
| 1995 | hidden_dim=256, |
| 1996 | intermediate_size=256, |
| 1997 | num_local_experts=8, |
| 1998 | num_logical_experts=8, |
| 1999 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 2000 | activation=MoEActivation.SILU, |
| 2001 | in_dtype=torch.bfloat16, |
| 2002 | device="cpu", |
| 2003 | routing_method=RoutingMethodType.Renormalize, |
| 2004 | ) |
| 2005 | |
| 2006 | with pytest.raises(NotImplementedError, match="Unsupported reasons"): |
| 2007 | mxfp4_oracle.select_mxfp4_moe_backend(moe_config) |
| 2008 | |
| 2009 | |
| 2010 | @pytest.mark.parametrize("backend_name", ["AITER_TRITON_MXFP4_BF16", "AITER_MXFP4_FP8"]) |
| 2011 | @pytest.mark.parametrize("use_ep", [False, True]) |
| 2012 | def test_aiter_mxfp4_monolithic_rejects_expert_parallel( |
| 2013 | monkeypatch: pytest.MonkeyPatch, backend_name: str, use_ep: bool |
| 2014 | ): |
| 2015 | """Global routing cannot index local expert weights, even without all-to-all.""" |
| 2016 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 2017 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 2018 | from vllm.model_executor.layers.fused_moe.config import RoutingMethodType |
| 2019 | from vllm.model_executor.layers.fused_moe.modular_kernel import ( |
| 2020 | FusedMoEActivationFormat, |
| 2021 | ) |
| 2022 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 2023 | Mxfp4MoeBackend, |
| 2024 | backend_to_kernel_cls, |
| 2025 | ) |
| 2026 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 2027 | kFp8StaticTensorSym, |
| 2028 | kMxfp4Static, |
| 2029 | ) |
| 2030 | |
| 2031 | experts_cls = backend_to_kernel_cls(Mxfp4MoeBackend[backend_name])[0] |
| 2032 | monkeypatch.setattr(experts_cls, "_supports_current_device", lambda: True) |
| 2033 | config = make_dummy_moe_config( |
| 2034 | num_experts=256, |
| 2035 | num_local_experts=64 if use_ep else 256, |
| 2036 | experts_per_token=6, |
| 2037 | hidden_dim=4096, |
| 2038 | intermediate_size=2048, |
| 2039 | activation=MoEActivation.SWIGLUOAI, |
| 2040 | ) |
| 2041 | config = replace( |
| 2042 | config, |
| 2043 | routing_method=RoutingMethodType.Renormalize, |
| 2044 | moe_parallel_config=replace( |
| 2045 | config.moe_parallel_config, |
| 2046 | tp_size=1 if use_ep else 4, |
| 2047 | ep_size=4 if use_ep else 1, |
| 2048 | use_ep=use_ep, |
| 2049 | ), |
| 2050 | ) |
| 2051 | assert not config.moe_parallel_config.use_all2all_kernels |
| 2052 | supported, reason = experts_cls.is_supported_config( |
| 2053 | experts_cls, |
| 2054 | config, |
| 2055 | kMxfp4Static, |
| 2056 | kFp8StaticTensorSym if backend_name == "AITER_MXFP4_FP8" else None, |
| 2057 | FusedMoEActivationFormat.Standard, |
| 2058 | ) |
| 2059 | |
| 2060 | assert supported is not use_ep |
| 2061 | if use_ep: |
| 2062 | assert "parallel config" in reason |
| 2063 | |
| 2064 | |
| 2065 | # Every activation-quantizing OCP MX scheme must map to a `quant_dtype` that |
| 2066 | # `moe_kernel_quantize_input` actually dispatches on. Its final `else` returns |
| 2067 | # the activation untouched, so a name it does not know (e.g. "mxfp6" instead of |
| 2068 | # "mxfp6_e3m2") silently skips the fake-quantization the emulation exists for. |
| 2069 | @pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm") |
| 2070 | @pytest.mark.parametrize("ocp_mx_scheme", list(OCP_MX_Scheme)) |
| 2071 | def test_emulation_activation_quant_dtype_is_dispatchable(ocp_mx_scheme): |
| 2072 | quant_dtype = activation_quant_dtype(ocp_mx_scheme) |
| 2073 | |
| 2074 | if "_a_" not in ocp_mx_scheme.value: |
| 2075 | assert quant_dtype is None, "weight-only schemes must not quantize activations" |
| 2076 | return |
| 2077 | |
| 2078 | a = torch.randn(64, 128, dtype=torch.bfloat16, device="cuda") |
| 2079 | a_scale = torch.ones(1, dtype=torch.float32, device="cuda") |
| 2080 | out, _ = moe_kernel_quantize_input( |
| 2081 | a, a_scale, quant_dtype, False, None, quantization_emulation=True |
| 2082 | ) |
| 2083 | assert not torch.equal(out, a), ( |
| 2084 | f"{ocp_mx_scheme.value} -> quant_dtype={quant_dtype!r} left the activation" |
| 2085 | " unquantized; moe_kernel_quantize_input does not dispatch on it" |
| 2086 | ) |
| 2087 | |
| 2088 | |
| 2089 | @pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm") |
| 2090 | @torch.inference_mode() |
| 2091 | def test_emulation_a_mxfp6_moe_forward_quantizes_activations(): |
| 2092 | """The same property observed through a full MoE forward. |
| 2093 | |
| 2094 | `w_mxfp4_a_mxfp6_e3m2` and the weight-only `w_mxfp4` differ only in whether |
| 2095 | activations are fake-quantized, so with identical weights, activations and |
| 2096 | routing their layer outputs must differ. When the emulation selects a |
| 2097 | `quant_dtype` `moe_kernel_quantize_input` does not dispatch on, the QDQ is |
| 2098 | skipped and the two outputs come out bit-identical. |
| 2099 | """ |
| 2100 | from vllm.config import VllmConfig, set_current_vllm_config |
| 2101 | from vllm.model_executor.layers.fused_moe import FusedMoEConfig |
| 2102 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 2103 | from vllm.model_executor.layers.fused_moe.config import ( |
| 2104 | FusedMoEParallelConfig, |
| 2105 | RoutingMethodType, |
| 2106 | mxfp4_w4a16_moe_quant_config, |
| 2107 | ocp_mx_moe_quant_config, |
| 2108 | ) |
| 2109 | from vllm.model_executor.layers.fused_moe.experts.ocp_mx_emulation_moe import ( |
| 2110 | OCP_MXQuantizationEmulationTritonExperts, |
| 2111 | ) |
| 2112 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 2113 | Mxfp4MoeBackend, |
| 2114 | make_mxfp4_moe_kernel, |
| 2115 | ) |
| 2116 | from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( |
| 2117 | mxfp4_quantize, |
| 2118 | ) |
| 2119 | from vllm.v1.worker.workspace import init_workspace_manager |
| 2120 | |
| 2121 | init_workspace_manager(torch.accelerator.current_device_index()) |
| 2122 | |
| 2123 | num_experts, topk = 8, 2 |
| 2124 | hidden_size, intermediate_size, num_tokens = 256, 256, 64 |
| 2125 | dtype, device = torch.bfloat16, "cuda:0" |
| 2126 | |
| 2127 | torch.manual_seed(0) |
| 2128 | w13, w13_scale = mxfp4_quantize( |
| 2129 | torch.randn( |
| 2130 | num_experts, 2 * intermediate_size, hidden_size, dtype=dtype, device=device |
| 2131 | ) |
| 2132 | / 8 |
| 2133 | ) |
| 2134 | w2, w2_scale = mxfp4_quantize( |
| 2135 | torch.randn( |
| 2136 | num_experts, hidden_size, intermediate_size, dtype=dtype, device=device |
| 2137 | ) |
| 2138 | / 8 |
| 2139 | ) |
| 2140 | w13, w2 = w13.contiguous(), w2.contiguous() |
| 2141 | w13_scale, w2_scale = w13_scale.contiguous(), w2_scale.contiguous() |
| 2142 | |
| 2143 | torch.manual_seed(1) |
| 2144 | hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) |
| 2145 | topk_weights, topk_ids = torch.topk( |
| 2146 | torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device), |
| 2147 | k=topk, |
| 2148 | dim=-1, |
| 2149 | ) |
| 2150 | topk_weights = torch.softmax(topk_weights, dim=-1).to(dtype) |
| 2151 | topk_ids = topk_ids.to(torch.int32) |
| 2152 | |
| 2153 | def run(quant_config): |
| 2154 | moe_config = FusedMoEConfig( |
| 2155 | num_experts=num_experts, |
| 2156 | experts_per_token=topk, |
| 2157 | hidden_dim=hidden_size, |
| 2158 | intermediate_size=intermediate_size, |
| 2159 | num_local_experts=num_experts, |
| 2160 | num_logical_experts=num_experts, |
| 2161 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 2162 | activation=MoEActivation.SILU, |
| 2163 | in_dtype=dtype, |
| 2164 | device=device, |
| 2165 | routing_method=RoutingMethodType.Renormalize, |
| 2166 | ) |
| 2167 | with set_current_vllm_config(VllmConfig()): |
| 2168 | kernel = make_mxfp4_moe_kernel( |
| 2169 | moe_quant_config=quant_config, |
| 2170 | moe_config=moe_config, |
| 2171 | mxfp4_backend=Mxfp4MoeBackend.EMULATION, |
| 2172 | experts_cls=OCP_MXQuantizationEmulationTritonExperts, |
| 2173 | routing_tables=None, |
| 2174 | ) |
| 2175 | return kernel.apply( |
| 2176 | hidden_states=hidden_states, |
| 2177 | w1=w13, |
| 2178 | w2=w2, |
| 2179 | topk_weights=topk_weights, |
| 2180 | topk_ids=topk_ids, |
| 2181 | activation=MoEActivation.SILU, |
| 2182 | global_num_experts=num_experts, |
| 2183 | expert_map=None, |
| 2184 | apply_router_weight_on_input=False, |
| 2185 | ) |
| 2186 | |
| 2187 | # A fresh quant config per run: the experts nulls the weight scales on it. |
| 2188 | a_mxfp6 = ocp_mx_moe_quant_config( |
| 2189 | quant_dtype="mxfp6_e3m2", |
| 2190 | weight_dtype="mxfp4", |
| 2191 | w1_scale=w13_scale, |
| 2192 | w2_scale=w2_scale, |
| 2193 | ) |
| 2194 | assert a_mxfp6.ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2 |
| 2195 | out_a_mxfp6 = run(a_mxfp6) |
| 2196 | |
| 2197 | weight_only = mxfp4_w4a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale) |
| 2198 | assert weight_only.ocp_mx_scheme == OCP_MX_Scheme.w_mxfp4 |
| 2199 | out_weight_only = run(weight_only) |
| 2200 | |
| 2201 | max_diff = (out_a_mxfp6.float() - out_weight_only.float()).abs().max().item() |
| 2202 | assert max_diff > 0.0, ( |
| 2203 | "w_mxfp4_a_mxfp6_e3m2 output is bit-identical to weight-only w_mxfp4:" |
| 2204 | " the emulation never fake-quantized the activations" |
| 2205 | ) |