| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import multiprocessing |
| 4 | from dataclasses import dataclass, replace |
| 5 | from types import SimpleNamespace |
| 6 | from unittest.mock import MagicMock |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 12 | from tests.kernels.moe.utils import ( |
| 13 | check_deferred_moe_finalize, |
| 14 | make_dummy_moe_config, |
| 15 | make_test_weights, |
| 16 | ) |
| 17 | from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config |
| 18 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 19 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 20 | maybe_make_prepare_finalize, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.config import ( |
| 23 | FusedMoEConfig, |
| 24 | FusedMoEParallelConfig, |
| 25 | FusedMoEQuantConfig, |
| 26 | RoutingMethodType, |
| 27 | fp8_w8a8_moe_quant_config, |
| 28 | ) |
| 29 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( |
| 30 | FlashInferExperts, |
| 31 | ) |
| 32 | from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( |
| 33 | view_as_block_major_k, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import ( |
| 36 | TrtLlmFp8ExpertsModular, |
| 37 | TrtLlmFp8ExpertsMonolithic, |
| 38 | ) |
| 39 | from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts |
| 40 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 41 | Fp8MoeBackend, |
| 42 | convert_to_fp8_moe_kernel_format, |
| 43 | make_fp8_moe_quant_config, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 46 | fused_topk_bias, |
| 47 | ) |
| 48 | from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk |
| 49 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 50 | rotate_weights_for_fi_trtllm_fp8_per_tensor_moe, |
| 51 | swap_w13_to_w31, |
| 52 | ) |
| 53 | from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8 |
| 54 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 55 | mxfp8_e4m3_quantize, |
| 56 | ) |
| 57 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 58 | QuantKey, |
| 59 | kFp8Dynamic128Sym, |
| 60 | kFp8DynamicTensorSym, |
| 61 | kFp8Static128BlockSym, |
| 62 | kFp8StaticTensorSym, |
| 63 | kMxfp8Dynamic, |
| 64 | kMxfp8Static, |
| 65 | ) |
| 66 | from vllm.model_executor.models.llama4 import Llama4MoE |
| 67 | from vllm.platforms import current_platform |
| 68 | from vllm.utils.math_utils import next_power_of_2 |
| 69 | from vllm.utils.torch_utils import set_random_seed |
| 70 | |
| 71 | try: |
| 72 | from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe |
| 73 | except ImportError: |
| 74 | if current_platform.is_rocm(): |
| 75 | pytest.skip( |
| 76 | "flashinfer not supported for vLLM on ROCm", allow_module_level=True |
| 77 | ) |
| 78 | |
| 79 | if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability( |
| 80 | 90 |
| 81 | ): |
| 82 | pytest.skip( |
| 83 | "Supported for sm >= 90", |
| 84 | allow_module_level=True, |
| 85 | ) |
| 86 | |
| 87 | NUM_EXPERTS = [16] |
| 88 | TOP_KS = [1] |
| 89 | |
| 90 | MNK_FACTORS = [ |
| 91 | (256, 8192, 5120), |
| 92 | (127, 4096, 5120), |
| 93 | (10, 8192, 5120), |
| 94 | (10, 4096, 5120), |
| 95 | (1, 8192, 5120), |
| 96 | (1, 4096, 5120), |
| 97 | ] |
| 98 | |
| 99 | vllm_config = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 100 | |
| 101 | |
| 102 | def quant_fp8_per_tensor_batches(a): |
| 103 | num_batches = a.size(0) |
| 104 | a_quant = [] |
| 105 | a_scales = [] |
| 106 | |
| 107 | for i in range(num_batches): |
| 108 | a_fp8, a_global_sf = input_to_float8(a[i]) |
| 109 | if a_global_sf.numel() == 1: |
| 110 | a_global_sf = a_global_sf.view(1, 1) |
| 111 | a_quant.append(a_fp8) |
| 112 | a_scales.append(a_global_sf) |
| 113 | |
| 114 | result_a_quant = torch.stack(a_quant) |
| 115 | result_a_scales = torch.stack(a_scales) |
| 116 | |
| 117 | return result_a_quant, result_a_scales |
| 118 | |
| 119 | |
| 120 | def check_accuracy(ref_output, actual_output, atol=0.1, rtol=0.85, percent=0.925): |
| 121 | close = torch.isclose(ref_output, actual_output, atol=atol, rtol=rtol) |
| 122 | match_ratio = close.float().mean() |
| 123 | assert match_ratio >= percent, ( |
| 124 | f"Match ratio {match_ratio:.4f} is below the threshold {percent:.4f}" |
| 125 | ) |
| 126 | |
| 127 | mismatch_percent = 1.0 - match_ratio.item() |
| 128 | assert mismatch_percent <= 1 - percent, ( |
| 129 | f"Mismatch percentage {mismatch_percent:.4f} is above the threshold " |
| 130 | f"{1 - percent:.4f}" |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | @dataclass |
| 135 | class TestData: |
| 136 | hidden_states: torch.Tensor |
| 137 | w13_quantized: torch.Tensor |
| 138 | w2_quantized: torch.Tensor |
| 139 | a1_scale: torch.Tensor |
| 140 | a2_scale: torch.Tensor |
| 141 | w13_weight_scale: torch.Tensor |
| 142 | w2_weight_scale: torch.Tensor |
| 143 | layer: torch.nn.Module |
| 144 | |
| 145 | @staticmethod |
| 146 | def make_moe_tensors_8bit( |
| 147 | m: int, |
| 148 | k: int, |
| 149 | n: int, |
| 150 | e: int, |
| 151 | is_trtllm: bool, |
| 152 | activation: MoEActivation = MoEActivation.SILU, |
| 153 | topk: int = 1, |
| 154 | ) -> "TestData": |
| 155 | is_gated = activation.is_gated |
| 156 | |
| 157 | hidden_states = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 |
| 158 | w13 = ( |
| 159 | torch.randn( |
| 160 | (e, (2 * n) if is_gated else n, k), device="cuda", dtype=torch.bfloat16 |
| 161 | ) |
| 162 | / 10 |
| 163 | ) |
| 164 | w2 = torch.randn((e, k, n), device="cuda", dtype=torch.bfloat16) / 10 |
| 165 | |
| 166 | # Scale to fp8 |
| 167 | _, a1_scale = input_to_float8(hidden_states) |
| 168 | a2_scale = torch.scalar_tensor(1.0).to(device="cuda").to(dtype=torch.float32) |
| 169 | w13_quantized, w13_weight_scale = quant_fp8_per_tensor_batches(w13) |
| 170 | w2_quantized, w2_weight_scale = quant_fp8_per_tensor_batches(w2) |
| 171 | |
| 172 | layer = torch.nn.Module() |
| 173 | layer.orig_dtype = torch.bfloat16 |
| 174 | layer.w13_weight = w13_quantized.clone() |
| 175 | layer.w2_weight = w2_quantized.clone() |
| 176 | layer.w13_input_scale = a1_scale |
| 177 | layer.w2_input_scale = a2_scale |
| 178 | layer.w13_weight_scale = w13_weight_scale |
| 179 | layer.w2_weight_scale = w2_weight_scale |
| 180 | layer.activation = activation |
| 181 | # Setup dummy config. |
| 182 | layer.moe_parallel_config = mk.FusedMoEParallelConfig.make_no_parallel() |
| 183 | |
| 184 | # flashinfer expects swapped rows for w13 |
| 185 | if is_gated: |
| 186 | layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data) |
| 187 | if is_trtllm: |
| 188 | rotate_weights_for_fi_trtllm_fp8_per_tensor_moe( |
| 189 | layer.w13_weight, layer.w2_weight, is_gated |
| 190 | ) |
| 191 | layer.custom_routing_function = Llama4MoE.custom_routing_function |
| 192 | layer.routing_method_type = RoutingMethodType.Llama4 |
| 193 | layer.renormalize = False |
| 194 | layer.intermediate_size_per_partition = n |
| 195 | layer.ep_rank = 0 |
| 196 | layer.local_num_experts = e |
| 197 | |
| 198 | layer.moe = FusedMoEConfig( |
| 199 | num_experts=e, |
| 200 | experts_per_token=topk, |
| 201 | hidden_dim=k, |
| 202 | intermediate_size=n, |
| 203 | num_local_experts=e, |
| 204 | num_logical_experts=e, |
| 205 | moe_parallel_config=layer.moe_parallel_config, |
| 206 | in_dtype=hidden_states.dtype, |
| 207 | routing_method=layer.routing_method_type, |
| 208 | activation=activation, |
| 209 | device=w13_quantized.device, |
| 210 | max_num_tokens=next_power_of_2(m), |
| 211 | ) |
| 212 | |
| 213 | return TestData( |
| 214 | hidden_states=hidden_states, |
| 215 | w13_quantized=w13_quantized, |
| 216 | w2_quantized=w2_quantized, |
| 217 | a1_scale=a1_scale, |
| 218 | a2_scale=a2_scale, |
| 219 | w13_weight_scale=w13_weight_scale, |
| 220 | w2_weight_scale=w2_weight_scale, |
| 221 | layer=layer, |
| 222 | ) |
| 223 | |
| 224 | |
| 225 | @pytest.mark.parametrize("m,n,k", MNK_FACTORS) |
| 226 | @pytest.mark.parametrize("e", NUM_EXPERTS) |
| 227 | @pytest.mark.parametrize("topk", TOP_KS) |
| 228 | @pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]) |
| 229 | def test_flashinfer_per_tensor_moe_fp8_no_graph( |
| 230 | m: int, |
| 231 | n: int, |
| 232 | k: int, |
| 233 | e: int, |
| 234 | topk: int, |
| 235 | activation: MoEActivation, |
| 236 | monkeypatch, |
| 237 | ): |
| 238 | if not current_platform.has_device_capability(100): |
| 239 | pytest.skip("Test is only supported for sm >= 100") |
| 240 | set_random_seed(7) |
| 241 | with set_current_vllm_config(vllm_config): |
| 242 | td = TestData.make_moe_tensors_8bit( |
| 243 | m, k, n, e, is_trtllm=True, activation=activation |
| 244 | ) |
| 245 | |
| 246 | score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) |
| 247 | topk_weights, topk_ids = Llama4MoE.custom_routing_function( |
| 248 | hidden_states=td.hidden_states, |
| 249 | gating_output=score, |
| 250 | topk=topk, |
| 251 | renormalize=False, |
| 252 | ) |
| 253 | |
| 254 | quant_config = fp8_w8a8_moe_quant_config( |
| 255 | w1_scale=td.w13_weight_scale, |
| 256 | w2_scale=td.w2_weight_scale, |
| 257 | a1_scale=td.a1_scale, |
| 258 | a2_scale=td.a2_scale, |
| 259 | per_act_token_quant=False, |
| 260 | ) |
| 261 | |
| 262 | output = fused_experts( |
| 263 | td.hidden_states, |
| 264 | td.w13_quantized, |
| 265 | td.w2_quantized, |
| 266 | topk_weights=topk_weights, |
| 267 | topk_ids=topk_ids, |
| 268 | activation=activation, |
| 269 | global_num_experts=e, |
| 270 | expert_map=None, |
| 271 | apply_router_weight_on_input=True, |
| 272 | quant_config=quant_config, |
| 273 | ) |
| 274 | |
| 275 | kernel = mk.FusedMoEKernel( |
| 276 | maybe_make_prepare_finalize( |
| 277 | moe=td.layer.moe, |
| 278 | quant_config=quant_config, |
| 279 | allow_new_interface=True, |
| 280 | use_monolithic=True, |
| 281 | ), |
| 282 | TrtLlmFp8ExpertsMonolithic( |
| 283 | moe_config=td.layer.moe, |
| 284 | quant_config=quant_config, |
| 285 | ), |
| 286 | ) |
| 287 | |
| 288 | flashinfer_output = kernel.apply_monolithic( |
| 289 | hidden_states=td.hidden_states, |
| 290 | w1=td.layer.w13_weight, |
| 291 | w2=td.layer.w2_weight, |
| 292 | router_logits=score, |
| 293 | activation=activation, |
| 294 | global_num_experts=e, |
| 295 | expert_map=None, |
| 296 | apply_router_weight_on_input=True, |
| 297 | routed_scaling_factor=1.0, |
| 298 | ) |
| 299 | |
| 300 | check_accuracy( |
| 301 | ref_output=output, |
| 302 | actual_output=flashinfer_output, |
| 303 | atol=0.1, |
| 304 | rtol=0.85, |
| 305 | percent=0.925, |
| 306 | ) |
| 307 | |
| 308 | |
| 309 | @pytest.mark.parametrize("m,n,k", MNK_FACTORS) |
| 310 | @pytest.mark.parametrize("e", NUM_EXPERTS) |
| 311 | @pytest.mark.parametrize("topk", TOP_KS) |
| 312 | @pytest.mark.parametrize("activation", [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]) |
| 313 | def test_flashinfer_cutlass_moe_fp8_no_graph( |
| 314 | m: int, |
| 315 | n: int, |
| 316 | k: int, |
| 317 | e: int, |
| 318 | topk: int, |
| 319 | activation: MoEActivation, |
| 320 | monkeypatch, |
| 321 | workspace_init, |
| 322 | ): |
| 323 | set_random_seed(7) |
| 324 | with set_current_vllm_config(vllm_config): |
| 325 | td = TestData.make_moe_tensors_8bit( |
| 326 | m, k, n, e, is_trtllm=False, activation=activation |
| 327 | ) |
| 328 | |
| 329 | score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) |
| 330 | topk_weights, topk_ids = Llama4MoE.custom_routing_function( |
| 331 | hidden_states=td.hidden_states, |
| 332 | gating_output=score, |
| 333 | topk=topk, |
| 334 | renormalize=False, |
| 335 | ) |
| 336 | |
| 337 | quant_config = fp8_w8a8_moe_quant_config( |
| 338 | w1_scale=td.w13_weight_scale, |
| 339 | g1_alphas=(td.w13_weight_scale * td.a1_scale).squeeze(), |
| 340 | w2_scale=td.w2_weight_scale, |
| 341 | g2_alphas=(td.w2_weight_scale * td.a2_scale).squeeze(), |
| 342 | a1_scale=td.a1_scale, |
| 343 | a1_gscale=td.a1_scale, |
| 344 | a2_scale=td.a2_scale, |
| 345 | a2_gscale=1.0 / td.a2_scale, |
| 346 | per_act_token_quant=False, |
| 347 | ) |
| 348 | |
| 349 | output = fused_experts( |
| 350 | td.hidden_states, |
| 351 | td.w13_quantized, |
| 352 | td.w2_quantized, |
| 353 | topk_weights=topk_weights, |
| 354 | topk_ids=topk_ids, |
| 355 | activation=activation, |
| 356 | global_num_experts=e, |
| 357 | expert_map=None, |
| 358 | apply_router_weight_on_input=True, |
| 359 | quant_config=quant_config, |
| 360 | ) |
| 361 | |
| 362 | td.layer.dp_size = 1 |
| 363 | |
| 364 | def get_fused_moe_quant_config(n: torch.nn.Module) -> FusedMoEQuantConfig: |
| 365 | return quant_config |
| 366 | |
| 367 | td.layer.get_fused_moe_quant_config = get_fused_moe_quant_config |
| 368 | td.layer.quant_method = td.layer |
| 369 | |
| 370 | moe_config = FusedMoEConfig( |
| 371 | num_experts=e, |
| 372 | experts_per_token=topk, |
| 373 | hidden_dim=k, |
| 374 | intermediate_size=n, |
| 375 | num_local_experts=e, |
| 376 | num_logical_experts=e, |
| 377 | activation=activation, |
| 378 | device="cuda", |
| 379 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 380 | in_dtype=torch.bfloat16, |
| 381 | routing_method=RoutingMethodType.TopK, |
| 382 | max_num_tokens=next_power_of_2(m), |
| 383 | ) |
| 384 | |
| 385 | kernel = mk.FusedMoEKernel( |
| 386 | maybe_make_prepare_finalize( |
| 387 | moe=moe_config, |
| 388 | quant_config=quant_config, |
| 389 | allow_new_interface=True, |
| 390 | use_monolithic=False, |
| 391 | ), |
| 392 | FlashInferExperts( |
| 393 | moe_config=moe_config, |
| 394 | quant_config=quant_config, |
| 395 | ), |
| 396 | ) |
| 397 | |
| 398 | flashinfer_cutlass_output = kernel.apply( |
| 399 | td.hidden_states, |
| 400 | td.layer.w13_weight, |
| 401 | td.layer.w2_weight, |
| 402 | topk_weights, |
| 403 | topk_ids, |
| 404 | activation=activation, |
| 405 | global_num_experts=e, |
| 406 | expert_map=None, |
| 407 | apply_router_weight_on_input=True, |
| 408 | ) |
| 409 | |
| 410 | check_accuracy( |
| 411 | ref_output=output, |
| 412 | actual_output=flashinfer_cutlass_output, |
| 413 | atol=0.1, |
| 414 | rtol=0.85, |
| 415 | percent=0.925, |
| 416 | ) |
| 417 | |
| 418 | |
| 419 | @pytest.mark.parametrize( |
| 420 | "num_experts,intermediate,hidden", |
| 421 | [ |
| 422 | (8, 2048, 1536), |
| 423 | (64, 4096, 4096), |
| 424 | ], |
| 425 | ) |
| 426 | def test_convert_moe_weights_to_flashinfer_trtllm_block_layout( |
| 427 | num_experts, intermediate, hidden |
| 428 | ): |
| 429 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 430 | convert_moe_weights_to_flashinfer_trtllm_block_layout, |
| 431 | ) |
| 432 | |
| 433 | w13 = torch.randn( |
| 434 | (num_experts, 2 * intermediate, hidden), dtype=torch.bfloat16, device="cuda" |
| 435 | ) |
| 436 | w2 = torch.randn( |
| 437 | (num_experts, hidden, intermediate), dtype=torch.bfloat16, device="cuda" |
| 438 | ) |
| 439 | |
| 440 | cache: dict[torch.Size, torch.Tensor] = {} |
| 441 | w13_converted, w2_converted = convert_moe_weights_to_flashinfer_trtllm_block_layout( |
| 442 | cache, w13, w2 |
| 443 | ) |
| 444 | |
| 445 | assert w13_converted.ndim == 4, ( |
| 446 | f"Expected 4D tensor, got shape {w13_converted.shape}" |
| 447 | ) |
| 448 | assert w2_converted.ndim == 4, f"Expected 4D tensor, got shape {w2_converted.shape}" |
| 449 | |
| 450 | assert w13_converted.numel() == w13.numel(), "W13 element count should be preserved" |
| 451 | assert w2_converted.numel() == w2.numel(), "W2 element count should be preserved" |
| 452 | |
| 453 | assert w13_converted.dtype == torch.bfloat16 |
| 454 | assert w2_converted.dtype == torch.bfloat16 |
| 455 | |
| 456 | assert w13_converted.shape[0] == num_experts |
| 457 | assert w2_converted.shape[0] == num_experts |
| 458 | assert w13_converted.data_ptr() == w13.data_ptr() |
| 459 | assert w2_converted.data_ptr() == w2.data_ptr() |
| 460 | |
| 461 | |
| 462 | @pytest.mark.parametrize("is_gated_act_gemm", [True, False]) |
| 463 | def test_convert_moe_weights_to_flashinfer_trtllm_block_layout_values( |
| 464 | is_gated_act_gemm, |
| 465 | ): |
| 466 | from flashinfer.fused_moe.core import ( |
| 467 | _maybe_get_cached_w3_w1_permute_indices, |
| 468 | get_w2_permute_indices_with_cache, |
| 469 | ) |
| 470 | |
| 471 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 472 | convert_moe_weights_to_flashinfer_trtllm_block_layout, |
| 473 | ) |
| 474 | |
| 475 | num_experts, intermediate, hidden = 2, 256, 256 |
| 476 | w13_multiplier = 2 if is_gated_act_gemm else 1 |
| 477 | w13 = torch.randn( |
| 478 | (num_experts, w13_multiplier * intermediate, hidden), |
| 479 | dtype=torch.bfloat16, |
| 480 | device="cuda", |
| 481 | ) |
| 482 | w2 = torch.randn( |
| 483 | (num_experts, hidden, intermediate), |
| 484 | dtype=torch.bfloat16, |
| 485 | device="cuda", |
| 486 | ) |
| 487 | |
| 488 | def _reference_block_layout( |
| 489 | weight: torch.Tensor, |
| 490 | is_w13: bool, |
| 491 | cache: dict[torch.Size, torch.Tensor], |
| 492 | ) -> torch.Tensor: |
| 493 | outputs = [] |
| 494 | for expert in weight: |
| 495 | expert_uint8 = expert.view(torch.uint8) |
| 496 | if is_w13: |
| 497 | indices = _maybe_get_cached_w3_w1_permute_indices( |
| 498 | cache, |
| 499 | expert_uint8, |
| 500 | 128, |
| 501 | is_gated_act_gemm=is_gated_act_gemm, |
| 502 | ) |
| 503 | if is_gated_act_gemm: |
| 504 | indices = (indices + expert_uint8.shape[0] // 2) % ( |
| 505 | expert_uint8.shape[0] |
| 506 | ) |
| 507 | else: |
| 508 | indices = get_w2_permute_indices_with_cache( |
| 509 | cache, |
| 510 | expert_uint8, |
| 511 | 128, |
| 512 | ) |
| 513 | rows, cols = expert_uint8.shape |
| 514 | blocks = expert_uint8.view(rows, cols // 128, 128).permute(1, 0, 2) |
| 515 | outputs.append(torch.index_select(blocks, 1, indices.to(weight.device))) |
| 516 | return torch.stack(outputs).view(torch.bfloat16) |
| 517 | |
| 518 | reference_cache: dict[torch.Size, torch.Tensor] = {} |
| 519 | expected_w13 = _reference_block_layout(w13, is_w13=True, cache=reference_cache) |
| 520 | expected_w2 = _reference_block_layout(w2, is_w13=False, cache=reference_cache) |
| 521 | w13_ptr = w13.data_ptr() |
| 522 | w2_ptr = w2.data_ptr() |
| 523 | |
| 524 | actual_w13, actual_w2 = convert_moe_weights_to_flashinfer_trtllm_block_layout( |
| 525 | {}, |
| 526 | w13, |
| 527 | w2, |
| 528 | is_gated_act_gemm=is_gated_act_gemm, |
| 529 | ) |
| 530 | |
| 531 | assert actual_w13.data_ptr() == w13_ptr |
| 532 | assert actual_w2.data_ptr() == w2_ptr |
| 533 | assert torch.equal(actual_w13, expected_w13) |
| 534 | assert torch.equal(actual_w2, expected_w2) |
| 535 | |
| 536 | |
| 537 | def _make_unquantized_flashinfer_test_layer( |
| 538 | monkeypatch, intermediate, is_gated, *, device="cuda" |
| 539 | ): |
| 540 | from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( |
| 541 | UnquantizedMoeBackend, |
| 542 | ) |
| 543 | from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( |
| 544 | UnquantizedFusedMoEMethod, |
| 545 | ) |
| 546 | |
| 547 | moe_config = make_dummy_moe_config( |
| 548 | num_experts=2, |
| 549 | hidden_dim=256, |
| 550 | intermediate_size=intermediate, |
| 551 | activation=MoEActivation.SILU if is_gated else MoEActivation.RELU2_NO_MUL, |
| 552 | ) |
| 553 | moe_config.intermediate_size_per_partition = (intermediate + 127) // 128 * 128 |
| 554 | method = object.__new__(UnquantizedFusedMoEMethod) |
| 555 | method.moe = moe_config |
| 556 | method.unquantized_backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM |
| 557 | method.moe_kernel = None |
| 558 | mock_kernel = MagicMock() |
| 559 | monkeypatch.setattr( |
| 560 | method, |
| 561 | "_init_moe_kernel", |
| 562 | lambda _: setattr(method, "moe_kernel", mock_kernel), |
| 563 | ) |
| 564 | |
| 565 | layer = torch.nn.Module() |
| 566 | layer.moe_config = moe_config |
| 567 | with torch.device(device): |
| 568 | method.create_weights( |
| 569 | layer, |
| 570 | num_experts=2, |
| 571 | hidden_size=256, |
| 572 | intermediate_size_per_partition=moe_config.intermediate_size_per_partition, |
| 573 | params_dtype=torch.bfloat16, |
| 574 | ) |
| 575 | return method, layer |
| 576 | |
| 577 | |
| 578 | @pytest.mark.parametrize("intermediate", [192, 256]) |
| 579 | @pytest.mark.parametrize("is_gated", [True, False]) |
| 580 | def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed( |
| 581 | monkeypatch, intermediate, is_gated |
| 582 | ): |
| 583 | """Each raw reload clears padding left dirty by prior in-place conversion.""" |
| 584 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 585 | convert_moe_weights_to_flashinfer_trtllm_block_layout, |
| 586 | ) |
| 587 | |
| 588 | method, layer = _make_unquantized_flashinfer_test_layer( |
| 589 | monkeypatch, intermediate, is_gated |
| 590 | ) |
| 591 | w13_shape, w2_shape = layer.w13_weight.shape, layer.w2_weight.shape |
| 592 | padded = w2_shape[-1] |
| 593 | w13_ptr = layer.w13_weight.data_ptr() |
| 594 | w2_ptr = layer.w2_weight.data_ptr() |
| 595 | load_numel = ( |
| 596 | layer.w13_weight.weight_loader_numel, |
| 597 | layer.w2_weight.weight_loader_numel, |
| 598 | ) |
| 599 | |
| 600 | for _ in range(2): |
| 601 | reloaded_w13 = torch.randn_like(layer.w13_weight) |
| 602 | reloaded_w2 = torch.randn_like(layer.w2_weight) |
| 603 | reloaded_w13[:, intermediate:padded].zero_() |
| 604 | if is_gated: |
| 605 | reloaded_w13[:, padded + intermediate :].zero_() |
| 606 | reloaded_w2[:, :, intermediate:].zero_() |
| 607 | expected_w13, expected_w2 = ( |
| 608 | convert_moe_weights_to_flashinfer_trtllm_block_layout( |
| 609 | {}, |
| 610 | reloaded_w13.clone(), |
| 611 | reloaded_w2.clone(), |
| 612 | is_gated_act_gemm=is_gated, |
| 613 | ) |
| 614 | ) |
| 615 | |
| 616 | # The loader writes only logical checkpoint slices, not padding. |
| 617 | layer.w13_weight.fill_(float("nan")) |
| 618 | layer.w2_weight.fill_(float("nan")) |
| 619 | layer.w13_weight[:, :intermediate].copy_(reloaded_w13[:, :intermediate]) |
| 620 | if is_gated: |
| 621 | layer.w13_weight[:, padded : padded + intermediate].copy_( |
| 622 | reloaded_w13[:, padded : padded + intermediate] |
| 623 | ) |
| 624 | layer.w2_weight[:, :, :intermediate].copy_(reloaded_w2[:, :, :intermediate]) |
| 625 | method.process_weights_after_loading(layer) |
| 626 | |
| 627 | assert layer.w13_weight.shape == w13_shape |
| 628 | assert layer.w2_weight.shape == w2_shape |
| 629 | assert layer.w13_weight.data_ptr() == w13_ptr |
| 630 | assert layer.w2_weight.data_ptr() == w2_ptr |
| 631 | assert ( |
| 632 | layer.w13_weight.weight_loader_numel, |
| 633 | layer.w2_weight.weight_loader_numel, |
| 634 | ) == load_numel |
| 635 | kernel_w13 = view_as_block_major_k(layer.w13_weight) |
| 636 | kernel_w2 = view_as_block_major_k(layer.w2_weight) |
| 637 | assert torch.equal(kernel_w13, expected_w13) |
| 638 | assert torch.equal(kernel_w2, expected_w2) |
| 639 | |
| 640 | |
| 641 | def _check_flashinfer_ipc_weights(entries, expected, mode, device_index): |
| 642 | from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader |
| 643 | from vllm.model_executor.model_loader.weight_cache.protocol import ( |
| 644 | WeightCacheState, |
| 645 | ) |
| 646 | from vllm.model_executor.utils import weights_already_processed |
| 647 | |
| 648 | torch.accelerator.set_device_index(device_index) |
| 649 | with pytest.MonkeyPatch.context() as monkeypatch: |
| 650 | method, layer = _make_unquantized_flashinfer_test_layer( |
| 651 | monkeypatch, 192, is_gated=True, device="meta" |
| 652 | ) |
| 653 | loader = object.__new__(IpcModelLoader) |
| 654 | loader.mode = mode |
| 655 | loader._apply_entries(layer, WeightCacheState(entries, {}, {}), device_index) |
| 656 | pointers = (layer.w13_weight.data_ptr(), layer.w2_weight.data_ptr()) |
| 657 | |
| 658 | # Like ipc_cache: a fresh method, only tensor metadata, no _setup_kernel. |
| 659 | with weights_already_processed(): |
| 660 | method.process_weights_after_loading(layer) |
| 661 | assert method.moe_kernel is not None |
| 662 | actual = ( |
| 663 | view_as_block_major_k(layer.w13_weight), |
| 664 | view_as_block_major_k(layer.w2_weight), |
| 665 | ) |
| 666 | for weight, reference, pointer in zip(actual, expected, pointers): |
| 667 | assert weight.data_ptr() == pointer |
| 668 | assert torch.equal(weight.cpu(), reference) |
| 669 | # Cross-process addresses need not match; writes prove sharing/isolation. |
| 670 | with torch.no_grad(): |
| 671 | for weight in actual: |
| 672 | weight.fill_(1.0) |
| 673 | torch.accelerator.synchronize() |
| 674 | |
| 675 | |
| 676 | @pytest.mark.parametrize("cache_ndim", [3, 4]) |
| 677 | @pytest.mark.parametrize("mode", ["copy", "zero_copy"]) |
| 678 | def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state( |
| 679 | monkeypatch, cache_ndim, mode |
| 680 | ): |
| 681 | """A spawned IPC consumer restores views without transient method state.""" |
| 682 | from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( |
| 683 | convert_moe_weights_to_flashinfer_trtllm_block_layout, |
| 684 | ) |
| 685 | from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry |
| 686 | |
| 687 | method, layer = _make_unquantized_flashinfer_test_layer( |
| 688 | monkeypatch, 192, is_gated=True, device="meta" |
| 689 | ) |
| 690 | packed = convert_moe_weights_to_flashinfer_trtllm_block_layout( |
| 691 | {}, |
| 692 | torch.randn_like(layer.w13_weight, device="cuda"), |
| 693 | torch.randn_like(layer.w2_weight, device="cuda"), |
| 694 | ) |
| 695 | expected = [weight.cpu() for weight in packed] |
| 696 | entries = {} |
| 697 | for name, weight in zip(("w13_weight", "w2_weight"), packed): |
| 698 | cached = weight.view_as(getattr(layer, name)) if cache_ndim == 3 else weight |
| 699 | entries[name] = TensorEntry.from_tensor(cached, kind="param") |
| 700 | consumer = multiprocessing.get_context("spawn").Process( |
| 701 | target=_check_flashinfer_ipc_weights, |
| 702 | args=( |
| 703 | entries, |
| 704 | expected, |
| 705 | mode, |
| 706 | torch.accelerator.current_device_index(), |
| 707 | ), |
| 708 | ) |
| 709 | consumer.start() |
| 710 | try: |
| 711 | # Keep producer allocations alive until the consumer releases its views. |
| 712 | consumer.join(timeout=180) |
| 713 | assert not consumer.is_alive(), "IPC consumer timed out" |
| 714 | assert consumer.exitcode == 0, "IPC consumer failed; see child traceback" |
| 715 | for source, reference in zip(packed, expected): |
| 716 | expected_source = ( |
| 717 | torch.ones_like(reference) if mode == "zero_copy" else reference |
| 718 | ) |
| 719 | assert torch.equal(source.cpu(), expected_source) |
| 720 | finally: |
| 721 | if consumer.is_alive(): |
| 722 | consumer.terminate() |
| 723 | consumer.join(timeout=10) |
| 724 | if consumer.is_alive(): |
| 725 | consumer.kill() |
| 726 | consumer.join(timeout=10) |
| 727 | consumer.close() |
| 728 | |
| 729 | |
| 730 | @pytest.mark.parametrize( |
| 731 | ("weight_key", "activation_key", "activation", "expected"), |
| 732 | [ |
| 733 | (kMxfp8Static, kMxfp8Dynamic, MoEActivation.SILU, True), |
| 734 | (kFp8Static128BlockSym, kFp8Dynamic128Sym, MoEActivation.SILU, True), |
| 735 | ( |
| 736 | kFp8Static128BlockSym, |
| 737 | kFp8Dynamic128Sym, |
| 738 | MoEActivation.SWIGLUOAI_UNINTERLEAVE, |
| 739 | True, |
| 740 | ), |
| 741 | (kFp8StaticTensorSym, kFp8DynamicTensorSym, MoEActivation.SILU, False), |
| 742 | # FlashInfer takes the clamp only with a SwiGLU activation. |
| 743 | ( |
| 744 | kFp8Static128BlockSym, |
| 745 | kFp8Dynamic128Sym, |
| 746 | MoEActivation.RELU2_NO_MUL, |
| 747 | False, |
| 748 | ), |
| 749 | ], |
| 750 | ) |
| 751 | def test_trtllm_fp8_swiglu_clamp_support( |
| 752 | weight_key: QuantKey, |
| 753 | activation_key: QuantKey, |
| 754 | activation: MoEActivation, |
| 755 | expected: bool, |
| 756 | ): |
| 757 | """FlashInfer >= 0.6.18 applies the SwiGLU clamp for both block-scaled |
| 758 | kernels with a SwiGLU activation (DeepSeek-V4 sets swiglu_limit); the |
| 759 | per-tensor kernel has no clamp, and Relu2 rejects the parameters.""" |
| 760 | |
| 761 | class _Experts(TrtLlmFp8ExpertsModular): |
| 762 | @staticmethod |
| 763 | def _supports_current_device() -> bool: |
| 764 | return True |
| 765 | |
| 766 | @staticmethod |
| 767 | def _supports_quant_scheme( |
| 768 | weight_key: QuantKey | None, activation_key: QuantKey | None |
| 769 | ) -> bool: |
| 770 | return True |
| 771 | |
| 772 | moe_config = make_dummy_moe_config() |
| 773 | moe_config.swiglu_limit = 7.0 |
| 774 | moe_config.activation = activation |
| 775 | |
| 776 | supported, reason = _Experts.is_supported_config( |
| 777 | _Experts, |
| 778 | moe_config, |
| 779 | weight_key, |
| 780 | activation_key, |
| 781 | mk.FusedMoEActivationFormat.Standard, |
| 782 | ) |
| 783 | |
| 784 | assert supported == expected, reason |
| 785 | if not expected: |
| 786 | assert "SwiGLU" in reason |
| 787 | |
| 788 | |
| 789 | def _make_mxfp8_moe_weights( |
| 790 | e: int, n: int, k: int |
| 791 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 792 | def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 793 | q, s = zip( |
| 794 | *(mxfp8_e4m3_quantize(w[i], is_sf_swizzled_layout=False) for i in range(e)) |
| 795 | ) |
| 796 | return torch.stack(q), torch.stack(s) |
| 797 | |
| 798 | w1, w1_scale = quantize( |
| 799 | torch.randn(e, 2 * n, k, device="cuda", dtype=torch.bfloat16) / 10 |
| 800 | ) |
| 801 | w2, w2_scale = quantize( |
| 802 | torch.randn(e, k, n, device="cuda", dtype=torch.bfloat16) / 10 |
| 803 | ) |
| 804 | return w1, w2, w1_scale, w2_scale |
| 805 | |
| 806 | |
| 807 | @pytest.mark.parametrize("block_shape", [[128, 128], [1, 32]], ids=["block", "mxfp8"]) |
| 808 | @pytest.mark.parametrize("m", [1, 16]) |
| 809 | @pytest.mark.skipif( |
| 810 | not current_platform.is_device_capability_family(100), |
| 811 | reason="Requires TRTLLM-Gen FP8 MoE (SM100)", |
| 812 | ) |
| 813 | def test_trtllm_fp8_block_moe_deferred_finalize( |
| 814 | m: int, block_shape: list[int], workspace_init |
| 815 | ): |
| 816 | """TRTLLM-Gen block-FP8 and MXFP8 modular experts can leave the top-k |
| 817 | finalize to the caller.""" |
| 818 | e, topk, n, k = 32, 4, 1024, 1024 |
| 819 | set_random_seed(7) |
| 820 | with set_current_vllm_config(vllm_config): |
| 821 | if block_shape == [1, 32]: |
| 822 | w1, w2, w1_scale, w2_scale = _make_mxfp8_moe_weights(e, n, k) |
| 823 | else: |
| 824 | (_, w1, w1_scale, _), (_, w2, w2_scale, _) = make_test_weights( |
| 825 | e, n, k, quant_dtype=torch.float8_e4m3fn, block_shape=block_shape |
| 826 | ) |
| 827 | w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( |
| 828 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 829 | SimpleNamespace( |
| 830 | weight_block_size=block_shape, |
| 831 | moe_config=SimpleNamespace( |
| 832 | is_act_and_mul=True, intermediate_size_per_partition=n |
| 833 | ), |
| 834 | activation=MoEActivation.SILU, |
| 835 | ), |
| 836 | w1, |
| 837 | w2, |
| 838 | w1_scale, |
| 839 | w2_scale, |
| 840 | w13_input_scale=None, |
| 841 | w2_input_scale=None, |
| 842 | ) |
| 843 | quant_config = make_fp8_moe_quant_config( |
| 844 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 845 | w1_scale, |
| 846 | w2_scale, |
| 847 | a1_scale=None, |
| 848 | a2_scale=None, |
| 849 | block_shape=block_shape, |
| 850 | ) |
| 851 | # One rank of a TP group, which deferral needs. |
| 852 | moe_config = FusedMoEConfig( |
| 853 | num_experts=e, |
| 854 | experts_per_token=topk, |
| 855 | hidden_dim=k, |
| 856 | intermediate_size=2 * n, |
| 857 | num_local_experts=e, |
| 858 | num_logical_experts=e, |
| 859 | activation=MoEActivation.SILU, |
| 860 | device="cuda", |
| 861 | moe_parallel_config=replace( |
| 862 | FusedMoEParallelConfig.make_no_parallel(), tp_size=2 |
| 863 | ), |
| 864 | in_dtype=torch.bfloat16, |
| 865 | routing_method=RoutingMethodType.TopK, |
| 866 | max_num_tokens=next_power_of_2(m), |
| 867 | ) |
| 868 | kernel = mk.FusedMoEKernel( |
| 869 | maybe_make_prepare_finalize( |
| 870 | moe=moe_config, quant_config=quant_config, allow_new_interface=True |
| 871 | ), |
| 872 | TrtLlmFp8ExpertsModular(moe_config=moe_config, quant_config=quant_config), |
| 873 | ) |
| 874 | |
| 875 | a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 |
| 876 | score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16) |
| 877 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) |
| 878 | check_deferred_moe_finalize( |
| 879 | moe_config, |
| 880 | lambda: kernel.apply( |
| 881 | hidden_states=a, |
| 882 | w1=w1, |
| 883 | w2=w2, |
| 884 | topk_weights=topk_weights, |
| 885 | topk_ids=topk_ids, |
| 886 | activation=MoEActivation.SILU, |
| 887 | global_num_experts=e, |
| 888 | expert_map=None, |
| 889 | apply_router_weight_on_input=False, |
| 890 | ), |
| 891 | router_weights=topk_weights, |
| 892 | ) |
| 893 | |
| 894 | |
| 895 | @pytest.mark.parametrize("m", [1, 16]) |
| 896 | @pytest.mark.skipif( |
| 897 | not current_platform.is_device_capability_family(100), |
| 898 | reason="Requires TRTLLM-Gen FP8 MoE (SM100)", |
| 899 | ) |
| 900 | def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_init): |
| 901 | """FlashInfer's fused MiniMax2 routing must apply routed_scaling_factor, |
| 902 | matching vLLM's own sigmoid+bias top-k routing (MiniMax-M3 uses 2.0).""" |
| 903 | e, topk, n, k, scale = 32, 4, 1024, 1024, 2.0 |
| 904 | block_shape = [1, 32] |
| 905 | set_random_seed(7) |
| 906 | with set_current_vllm_config(vllm_config): |
| 907 | w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( |
| 908 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 909 | SimpleNamespace( |
| 910 | weight_block_size=block_shape, |
| 911 | moe_config=SimpleNamespace( |
| 912 | is_act_and_mul=True, intermediate_size_per_partition=n |
| 913 | ), |
| 914 | activation=MoEActivation.SILU, |
| 915 | ), |
| 916 | *_make_mxfp8_moe_weights(e, n, k), |
| 917 | w13_input_scale=None, |
| 918 | w2_input_scale=None, |
| 919 | ) |
| 920 | quant_config = make_fp8_moe_quant_config( |
| 921 | Fp8MoeBackend.FLASHINFER_TRTLLM, |
| 922 | w1_scale, |
| 923 | w2_scale, |
| 924 | a1_scale=None, |
| 925 | a2_scale=None, |
| 926 | block_shape=block_shape, |
| 927 | ) |
| 928 | moe_config = FusedMoEConfig( |
| 929 | num_experts=e, |
| 930 | experts_per_token=topk, |
| 931 | hidden_dim=k, |
| 932 | intermediate_size=n, |
| 933 | num_local_experts=e, |
| 934 | num_logical_experts=e, |
| 935 | activation=MoEActivation.SILU, |
| 936 | device="cuda", |
| 937 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 938 | in_dtype=torch.bfloat16, |
| 939 | routing_method=RoutingMethodType.MiniMax2, |
| 940 | max_num_tokens=next_power_of_2(m), |
| 941 | ) |
| 942 | |
| 943 | def make_kernel(experts_cls) -> mk.FusedMoEKernel: |
| 944 | return mk.FusedMoEKernel( |
| 945 | maybe_make_prepare_finalize( |
| 946 | moe=moe_config, |
| 947 | quant_config=quant_config, |
| 948 | allow_new_interface=True, |
| 949 | use_monolithic=experts_cls is TrtLlmFp8ExpertsMonolithic, |
| 950 | ), |
| 951 | experts_cls(moe_config=moe_config, quant_config=quant_config), |
| 952 | ) |
| 953 | |
| 954 | a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10 |
| 955 | logits = torch.randn((m, e), device="cuda", dtype=torch.float32) |
| 956 | bias = torch.randn(e, device="cuda", dtype=torch.float32) / 10 |
| 957 | common = dict( |
| 958 | w1=w1, |
| 959 | w2=w2, |
| 960 | activation=MoEActivation.SILU, |
| 961 | global_num_experts=e, |
| 962 | expert_map=None, |
| 963 | apply_router_weight_on_input=False, |
| 964 | ) |
| 965 | |
| 966 | fused = make_kernel(TrtLlmFp8ExpertsMonolithic).apply_monolithic( |
| 967 | hidden_states=a, |
| 968 | router_logits=logits, |
| 969 | e_score_correction_bias=bias, |
| 970 | routed_scaling_factor=scale, |
| 971 | **common, |
| 972 | ) |
| 973 | topk_weights, topk_ids = fused_topk_bias( |
| 974 | a, |
| 975 | logits, |
| 976 | scoring_func="sigmoid", |
| 977 | e_score_correction_bias=bias, |
| 978 | topk=topk, |
| 979 | renormalize=True, |
| 980 | routed_scaling_factor=scale, |
| 981 | ) |
| 982 | reference = make_kernel(TrtLlmFp8ExpertsModular).apply( |
| 983 | hidden_states=a, topk_weights=topk_weights, topk_ids=topk_ids, **common |
| 984 | ) |
| 985 | |
| 986 | torch.testing.assert_close(fused, reference, atol=2e-2, rtol=2e-2) |