| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the FlashInfer TRTLLM NvFP4 MoE backend |
| 4 | (`TrtLlmNvFp4ExpertsModular`). |
| 5 | |
| 6 | Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated), |
| 7 | and GELU — including a Gemma4-shaped case (128 experts, top-k 8, |
| 8 | intermediate_size 704) that exercises the non-256-aligned padding path. |
| 9 | """ |
| 10 | |
| 11 | from dataclasses import replace |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | |
| 16 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 17 | from tests.kernels.moe.utils import check_deferred_moe_finalize, make_test_quant_config |
| 18 | from tests.kernels.quantization.nvfp4_utils import ( |
| 19 | FLOAT4_E2M1_MAX, |
| 20 | FLOAT8_E4M3_MAX, |
| 21 | dequantize_nvfp4_to_dtype, |
| 22 | ) |
| 23 | from tests.kernels.utils import torch_moe |
| 24 | from vllm import _custom_ops as ops |
| 25 | from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config |
| 26 | from vllm.model_executor.custom_op import CustomOp, op_registry |
| 27 | from vllm.model_executor.layers.activation import SiluAndMulWithClamp, SituAndMul |
| 28 | from vllm.model_executor.layers.fused_moe import fused_topk |
| 29 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 30 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 31 | maybe_make_prepare_finalize, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.fused_moe.config import ( |
| 34 | FusedMoEConfig, |
| 35 | FusedMoEParallelConfig, |
| 36 | FusedMoEQuantConfig, |
| 37 | RoutingMethodType, |
| 38 | ) |
| 39 | from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( |
| 40 | TrtLlmNvFp4ExpertsModular, |
| 41 | ) |
| 42 | from vllm.platforms import current_platform |
| 43 | from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe |
| 44 | from vllm.utils.math_utils import next_power_of_2 |
| 45 | from vllm.utils.torch_utils import set_random_seed |
| 46 | |
| 47 | if pytest and ( |
| 48 | not has_flashinfer_trtllm_fused_moe() |
| 49 | or not current_platform.is_device_capability_family(100) |
| 50 | ): |
| 51 | pytest.skip( |
| 52 | "Requires flashinfer TRTLLM fused MoE and NvFP4 (SM100)", |
| 53 | allow_module_level=True, |
| 54 | ) |
| 55 | |
| 56 | # (m, n, k) = (tokens, intermediate_size_per_partition, hidden_dim). |
| 57 | # The (64, 704, 4096) row matches Gemma4's MoE shape and exercises the |
| 58 | # non-256-aligned intermediate (padded inside the wrapper). |
| 59 | MNK_FACTORS = [ |
| 60 | (2, 1024, 1024), |
| 61 | (64, 2048, 1536), |
| 62 | (64, 704, 4096), |
| 63 | ] |
| 64 | |
| 65 | _SWIGLU_LIMIT = 0.1 |
| 66 | _LARGE_OUTPUT1_SCALE = 32768.0 |
| 67 | _CLAMP_OP_NAME = "test_silu_and_mul_with_clamp" |
| 68 | _SITU_OP_NAME = "test_situ_and_mul" |
| 69 | |
| 70 | # Test-only fixed-limit clamp. ``custom_op_name`` makes the class itself |
| 71 | # valid as an ``activation=`` argument to ``torch_moe`` (which only looks |
| 72 | # up ``activation.custom_op_name`` in ``op_registry``), so no |
| 73 | # ``MoEActivation`` enum extension is needed. |
| 74 | if _CLAMP_OP_NAME not in op_registry: |
| 75 | |
| 76 | @CustomOp.register(_CLAMP_OP_NAME) |
| 77 | class _SiluAndMulWithClampTest(SiluAndMulWithClamp): |
| 78 | custom_op_name = _CLAMP_OP_NAME |
| 79 | |
| 80 | def __init__(self, *, compile_native: bool = True) -> None: |
| 81 | super().__init__(_SWIGLU_LIMIT, compile_native=compile_native) |
| 82 | |
| 83 | |
| 84 | if _SITU_OP_NAME not in op_registry: |
| 85 | |
| 86 | @CustomOp.register(_SITU_OP_NAME) |
| 87 | class _SituAndMulTest(SituAndMul): |
| 88 | custom_op_name = _SITU_OP_NAME |
| 89 | |
| 90 | def __init__(self, *, compile_native: bool = True) -> None: |
| 91 | super().__init__(4.0, 25.0, compile_native=compile_native) |
| 92 | |
| 93 | |
| 94 | SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME] |
| 95 | SITU = op_registry[_SITU_OP_NAME] |
| 96 | |
| 97 | |
| 98 | def _make_trtllm_fp4_moe_kernel( |
| 99 | moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig |
| 100 | ) -> mk.FusedMoEKernel: |
| 101 | experts = TrtLlmNvFp4ExpertsModular( |
| 102 | moe_config=moe_config, quant_config=quant_config |
| 103 | ) |
| 104 | # Mimic the production weight-loader path so per-expert tensors that |
| 105 | # are normally precomputed in process_weights_after_loading (g1_scale_c |
| 106 | # and the rescaled gemm1_clamp_limit) get materialized. The test's |
| 107 | # synthetic quant_config has g1_alphas/g2_alphas already at their |
| 108 | # post-fusion values, so we set w13_weight_scale_2 to alias g1_alphas |
| 109 | # (same tensor) and use input_scale=1 to make the in-place |
| 110 | # weight_scale_2 *= input_scale step a no-op. |
| 111 | fake_layer = torch.nn.Module() |
| 112 | fake_layer.w13_weight_scale_2 = quant_config.g1_alphas |
| 113 | fake_layer.w2_weight_scale_2 = quant_config.g2_alphas |
| 114 | fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas) |
| 115 | fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas) |
| 116 | experts.process_weights_after_loading(fake_layer) |
| 117 | return mk.FusedMoEKernel( |
| 118 | maybe_make_prepare_finalize( |
| 119 | moe=moe_config, |
| 120 | quant_config=quant_config, |
| 121 | allow_new_interface=True, |
| 122 | use_monolithic=False, |
| 123 | ), |
| 124 | experts, |
| 125 | ) |
| 126 | |
| 127 | |
| 128 | ACTIVATION_CASES = [ |
| 129 | pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"), |
| 130 | pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"), |
| 131 | pytest.param(MoEActivation.SITU, SITU, None, id="situ"), |
| 132 | pytest.param( |
| 133 | MoEActivation.RELU2_NO_MUL, |
| 134 | MoEActivation.RELU2_NO_MUL, |
| 135 | None, |
| 136 | id="relu2_no_mul", |
| 137 | ), |
| 138 | pytest.param(MoEActivation.GELU, MoEActivation.GELU, None, id="gelu"), |
| 139 | ] |
| 140 | |
| 141 | |
| 142 | @pytest.mark.parametrize("m,n,k", MNK_FACTORS) |
| 143 | @pytest.mark.parametrize("e", [128]) |
| 144 | @pytest.mark.parametrize("topk", [8]) |
| 145 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 146 | @pytest.mark.parametrize("activation,torch_activation,swiglu_limit", ACTIVATION_CASES) |
| 147 | @torch.inference_mode() |
| 148 | def test_trtllm_fp4_moe_no_graph( |
| 149 | m: int, |
| 150 | n: int, |
| 151 | k: int, |
| 152 | e: int, |
| 153 | topk: int, |
| 154 | dtype: torch.dtype, |
| 155 | activation: MoEActivation, |
| 156 | torch_activation: MoEActivation | type[SiluAndMulWithClamp] | type[SituAndMul], |
| 157 | swiglu_limit: float | None, |
| 158 | workspace_init, |
| 159 | ): |
| 160 | # FlashInfer's trtllm_batched_gemm_runner has no precompiled tile |
| 161 | # config for non-gated RELU^2 at non-256-aligned intermediate_size |
| 162 | # (e.g. Gemma4's 704). Other activations (SiLU/GELU) work at the |
| 163 | # same shape. Tracked upstream in FlashInfer; unrelated to this |
| 164 | # PR's GELU enablement (Gemma4 uses GeGLU, not non-gated RELU^2). |
| 165 | if activation == MoEActivation.RELU2_NO_MUL and (m, n, k) == (64, 704, 4096): |
| 166 | pytest.skip( |
| 167 | "FlashInfer trtllm_batched_gemm_runner: no valid tile config " |
| 168 | "for non-gated RELU^2 at intermediate_size=704 " |
| 169 | "(getValidConfigIndices throws). Tracked upstream." |
| 170 | ) |
| 171 | |
| 172 | set_random_seed(7) |
| 173 | with set_current_vllm_config( |
| 174 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 175 | ): |
| 176 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 177 | |
| 178 | quant_blocksize = 16 |
| 179 | is_gated_act = activation.is_gated |
| 180 | |
| 181 | w1_q, w2_q, quant_config = make_test_quant_config( |
| 182 | e, |
| 183 | n, |
| 184 | k, |
| 185 | in_dtype=dtype, |
| 186 | quant_dtype="nvfp4", |
| 187 | block_shape=None, |
| 188 | per_act_token_quant=False, |
| 189 | make_gate=is_gated_act, |
| 190 | # The TRT-LLM FP4 MoE kernel rejects swizzled (padded) activation |
| 191 | # scales — its numel-based vec_size check requires numel == M*K/16. |
| 192 | # Match what oracle/nvfp4.py does for this backend. |
| 193 | is_scale_swizzled=False, |
| 194 | ) |
| 195 | quant_config.gemm1_clamp_limit = swiglu_limit |
| 196 | if swiglu_limit is not None: |
| 197 | assert quant_config.g1_alphas is not None |
| 198 | assert quant_config.a2_gscale is not None |
| 199 | assert torch.all(quant_config.a2_gscale == 1) |
| 200 | # With a2_gscale == 1, g1_alphas is the TRTLLM |
| 201 | # output1_scale_gate_scalar. Make it large enough to catch |
| 202 | # clamp/output-scale coupling in the FlashInfer kernel wrapper. |
| 203 | quant_config.g1_alphas.fill_(_LARGE_OUTPUT1_SCALE) |
| 204 | |
| 205 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 206 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) |
| 207 | |
| 208 | moe_config = FusedMoEConfig( |
| 209 | num_experts=e, |
| 210 | experts_per_token=topk, |
| 211 | hidden_dim=k, |
| 212 | intermediate_size=n, |
| 213 | num_local_experts=e, |
| 214 | num_logical_experts=e, |
| 215 | activation=activation, |
| 216 | device="cuda", |
| 217 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 218 | in_dtype=dtype, |
| 219 | routing_method=RoutingMethodType.TopK, |
| 220 | max_num_tokens=next_power_of_2(m), |
| 221 | activation_situ_beta=4.0 if activation == MoEActivation.SITU else None, |
| 222 | activation_situ_linear_beta=( |
| 223 | 25.0 if activation == MoEActivation.SITU else None |
| 224 | ), |
| 225 | ) |
| 226 | |
| 227 | trtllm_experts = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) |
| 228 | if activation == MoEActivation.SITU: |
| 229 | torch.testing.assert_close( |
| 230 | trtllm_experts.fused_experts.g1_scale_c, quant_config.a2_gscale |
| 231 | ) |
| 232 | |
| 233 | trtllm_output = trtllm_experts.apply( |
| 234 | hidden_states=a, |
| 235 | w1=w1_q, |
| 236 | w2=w2_q, |
| 237 | topk_weights=topk_weights, |
| 238 | topk_ids=topk_ids, |
| 239 | activation=activation, |
| 240 | global_num_experts=e, |
| 241 | expert_map=None, |
| 242 | apply_router_weight_on_input=False, |
| 243 | ) |
| 244 | |
| 245 | # Reference: round-trip activations and weights through FP4 |
| 246 | # quant/dequant so the comparison isolates kernel/activation behavior |
| 247 | # from quantization error. |
| 248 | a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / a.abs().max()).to( |
| 249 | torch.float32 |
| 250 | ) |
| 251 | a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale) |
| 252 | a_in_dtype = dequantize_nvfp4_to_dtype( |
| 253 | a_fp4, |
| 254 | a_scale_interleaved, |
| 255 | a_global_scale, |
| 256 | dtype=a.dtype, |
| 257 | device=a.device, |
| 258 | block_size=quant_blocksize, |
| 259 | ) |
| 260 | |
| 261 | w1_d = torch.empty( |
| 262 | (e, (2 if is_gated_act else 1) * n, k), device="cuda", dtype=dtype |
| 263 | ) |
| 264 | w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) |
| 265 | for idx in range(e): |
| 266 | w1_d[idx] = dequantize_nvfp4_to_dtype( |
| 267 | w1_q[idx], |
| 268 | quant_config.w1_scale[idx], |
| 269 | (1 / quant_config.g1_alphas[idx]), |
| 270 | dtype=dtype, |
| 271 | device=w1_q.device, |
| 272 | block_size=quant_blocksize, |
| 273 | ) |
| 274 | w2_d[idx] = dequantize_nvfp4_to_dtype( |
| 275 | w2_q[idx], |
| 276 | quant_config.w2_scale[idx], |
| 277 | (1 / quant_config.g2_alphas[idx]), |
| 278 | dtype=dtype, |
| 279 | device=w2_q.device, |
| 280 | block_size=quant_blocksize, |
| 281 | ) |
| 282 | |
| 283 | torch_output = torch_moe( |
| 284 | a_in_dtype, w1_d, w2_d, score, topk, activation=torch_activation |
| 285 | ) |
| 286 | |
| 287 | torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1) |
| 288 | |
| 289 | |
| 290 | @torch.inference_mode() |
| 291 | def test_trtllm_fp4_moe_reprocess_does_not_refold_swiglu_params(): |
| 292 | """Weight reloads rerun post-processing on the same experts object; the |
| 293 | g1_alphas fold of clamp/beta must start from the unfolded values.""" |
| 294 | e, n, k = 8, 256, 256 |
| 295 | clamp, beta = 7.0, 1.0 |
| 296 | with set_current_vllm_config( |
| 297 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 298 | ): |
| 299 | _, _, quant_config = make_test_quant_config( |
| 300 | e, |
| 301 | n, |
| 302 | k, |
| 303 | in_dtype=torch.bfloat16, |
| 304 | quant_dtype="nvfp4", |
| 305 | block_shape=None, |
| 306 | per_act_token_quant=False, |
| 307 | is_scale_swizzled=False, |
| 308 | ) |
| 309 | quant_config.gemm1_clamp_limit = clamp |
| 310 | quant_config.gemm1_beta = beta |
| 311 | moe_config = FusedMoEConfig( |
| 312 | num_experts=e, |
| 313 | experts_per_token=2, |
| 314 | hidden_dim=k, |
| 315 | intermediate_size=n, |
| 316 | num_local_experts=e, |
| 317 | num_logical_experts=e, |
| 318 | activation=MoEActivation.SILU, |
| 319 | device="cuda", |
| 320 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 321 | in_dtype=torch.bfloat16, |
| 322 | routing_method=RoutingMethodType.TopK, |
| 323 | max_num_tokens=16, |
| 324 | ) |
| 325 | experts = TrtLlmNvFp4ExpertsModular( |
| 326 | moe_config=moe_config, quant_config=quant_config |
| 327 | ) |
| 328 | g1_alphas = quant_config.g1_alphas |
| 329 | assert g1_alphas is not None |
| 330 | layer = torch.nn.Module() |
| 331 | layer.w13_weight_scale_2 = g1_alphas |
| 332 | layer.w2_weight_scale_2 = quant_config.g2_alphas |
| 333 | layer.w13_input_scale = torch.ones_like(g1_alphas) |
| 334 | layer.w2_input_scale = torch.ones_like(g1_alphas) |
| 335 | |
| 336 | for new_alphas in (0.5, 0.25): |
| 337 | g1_alphas.fill_(new_alphas) |
| 338 | experts.process_weights_after_loading(layer) |
| 339 | expected = torch.full_like(g1_alphas, 1.0 / new_alphas) |
| 340 | torch.testing.assert_close(experts.gemm1_clamp_limit, clamp * expected) |
| 341 | torch.testing.assert_close(experts.gemm1_beta, beta * expected) |
| 342 | |
| 343 | |
| 344 | @pytest.mark.parametrize("m,chunk_size", [(1, None), (16, None), (16, 5)]) |
| 345 | @torch.inference_mode() |
| 346 | def test_trtllm_fp4_moe_deferred_finalize( |
| 347 | m: int, |
| 348 | chunk_size: int | None, |
| 349 | monkeypatch: pytest.MonkeyPatch, |
| 350 | workspace_init, |
| 351 | ): |
| 352 | """TRTLLM-Gen NvFP4 modular experts can leave the top-k finalize to the caller.""" |
| 353 | e, topk, n, k = 128, 8, 1024, 1024 |
| 354 | dtype = torch.bfloat16 |
| 355 | if chunk_size is not None: |
| 356 | monkeypatch.setattr( |
| 357 | TrtLlmNvFp4ExpertsModular, "_get_chunk_size", lambda self: chunk_size |
| 358 | ) |
| 359 | |
| 360 | set_random_seed(7) |
| 361 | with set_current_vllm_config( |
| 362 | VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) |
| 363 | ): |
| 364 | a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 |
| 365 | w1_q, w2_q, quant_config = make_test_quant_config( |
| 366 | e, n, k, in_dtype=dtype, quant_dtype="nvfp4", is_scale_swizzled=False |
| 367 | ) |
| 368 | score = torch.randn((m, e), device="cuda", dtype=dtype) |
| 369 | topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) |
| 370 | |
| 371 | # One rank of a TP group, which deferral needs. |
| 372 | moe_config = FusedMoEConfig( |
| 373 | num_experts=e, |
| 374 | experts_per_token=topk, |
| 375 | hidden_dim=k, |
| 376 | intermediate_size=2 * n, |
| 377 | num_local_experts=e, |
| 378 | num_logical_experts=e, |
| 379 | activation=MoEActivation.SILU, |
| 380 | device="cuda", |
| 381 | moe_parallel_config=replace( |
| 382 | FusedMoEParallelConfig.make_no_parallel(), tp_size=2 |
| 383 | ), |
| 384 | in_dtype=dtype, |
| 385 | routing_method=RoutingMethodType.TopK, |
| 386 | max_num_tokens=next_power_of_2(m), |
| 387 | ) |
| 388 | kernel = _make_trtllm_fp4_moe_kernel(moe_config, quant_config) |
| 389 | |
| 390 | check_deferred_moe_finalize( |
| 391 | moe_config, |
| 392 | lambda: kernel.apply( |
| 393 | hidden_states=a, |
| 394 | w1=w1_q, |
| 395 | w2=w2_q, |
| 396 | topk_weights=topk_weights, |
| 397 | topk_ids=topk_ids, |
| 398 | activation=MoEActivation.SILU, |
| 399 | global_num_experts=e, |
| 400 | expert_map=None, |
| 401 | apply_router_weight_on_input=False, |
| 402 | ), |
| 403 | router_weights=topk_weights, |
| 404 | chunked=chunk_size is not None, |
| 405 | ) |
| 406 | |
| 407 | |
| 408 | if __name__ == "__main__": |
| 409 | test_trtllm_fp4_moe_no_graph( |
| 410 | 64, |
| 411 | 704, |
| 412 | 4096, |
| 413 | 128, |
| 414 | 8, |
| 415 | torch.bfloat16, |
| 416 | MoEActivation.GELU, |
| 417 | MoEActivation.GELU, |
| 418 | None, |
| 419 | None, |
| 420 | ) |