A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 10:00 UTC
[Kernel][MoE] Support GELU on TRT-LLM NvFP4 fused MoE for Gemma4 (#41050)
Signed-off-by: Juhi Mittal <juhim@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Juhi Mittal committed
Apr 30, 2026 at 20:37 UTC
6b6ac6c3c737b69e99264731f010588613dada58
3 files changed
+217
-4
tests/kernels/moe/test_trtllm_nvfp4_moe.py
new
+207
@@ -0,0 +1,207 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for the FlashInfer TRTLLM NvFP4 MoE backend
5
+(`TrtLlmNvFp4ExpertsModular`).
6
+
7
+Covers the activations the wrapper claims to support — SiLU, RELU^2 (non-gated),
8
+and GELU — including a Gemma4-shaped case (128 experts, top-k 8,
9
+intermediate_size 704) that exercises the non-256-aligned padding path.
10
+"""
11
+
12
+import pytest
13
+import torch
14
+
15
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16
+from tests.kernels.moe.utils import make_test_quant_config
17
+from tests.kernels.quantization.nvfp4_utils import (
18
+ FLOAT4_E2M1_MAX,
19
+ FLOAT8_E4M3_MAX,
20
+ dequantize_nvfp4_to_dtype,
21
+)
22
+from tests.kernels.utils import torch_moe
23
+from vllm import _custom_ops as ops
24
+from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
25
+from vllm.model_executor.layers.fused_moe import fused_topk
26
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
27
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
28
+ maybe_make_prepare_finalize,
29
+)
30
+from vllm.model_executor.layers.fused_moe.config import (
31
+ FusedMoEConfig,
32
+ FusedMoEParallelConfig,
33
+ RoutingMethodType,
34
+)
35
+from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
36
+ TrtLlmNvFp4ExpertsModular,
37
+)
38
+from vllm.platforms import current_platform
39
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
40
+from vllm.utils.math_utils import next_power_of_2
41
+from vllm.utils.torch_utils import set_random_seed
42
+
43
+if pytest and (
44
+ not has_flashinfer_trtllm_fused_moe()
45
+ or not current_platform.has_device_capability(100)
46
+):
47
+ pytest.skip(
48
+ "Requires flashinfer TRTLLM fused MoE and NvFP4 (SM100)",
49
+ allow_module_level=True,
50
+ )
51
+
52
+# (m, n, k) = (tokens, intermediate_size_per_partition, hidden_dim).
53
+# The (64, 704, 4096) row matches Gemma4's MoE shape and exercises the
54
+# non-256-aligned intermediate (padded inside the wrapper).
55
+MNK_FACTORS = [
56
+ (2, 1024, 1024),
57
+ (64, 2048, 1536),
58
+ (64, 704, 4096),
59
+]
60
+
61
+
62
+@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
63
+@pytest.mark.parametrize("e", [128])
64
+@pytest.mark.parametrize("topk", [8])
65
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
66
+@pytest.mark.parametrize(
67
+ "activation",
68
+ [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL, MoEActivation.GELU],
69
+)
70
+@torch.inference_mode()
71
+def test_trtllm_fp4_moe_no_graph(
72
+ m: int,
73
+ n: int,
74
+ k: int,
75
+ e: int,
76
+ topk: int,
77
+ dtype: torch.dtype,
78
+ activation: MoEActivation,
79
+ workspace_init,
80
+):
81
+ # FlashInfer's trtllm_batched_gemm_runner has no precompiled tile
82
+ # config for non-gated RELU^2 at non-256-aligned intermediate_size
83
+ # (e.g. Gemma4's 704). Other activations (SiLU/GELU) work at the
84
+ # same shape. Tracked upstream in FlashInfer; unrelated to this
85
+ # PR's GELU enablement (Gemma4 uses GeGLU, not non-gated RELU^2).
86
+ if activation == MoEActivation.RELU2_NO_MUL and (m, n, k) == (64, 704, 4096):
87
+ pytest.skip(
88
+ "FlashInfer trtllm_batched_gemm_runner: no valid tile config "
89
+ "for non-gated RELU^2 at intermediate_size=704 "
90
+ "(getValidConfigIndices throws). Tracked upstream."
91
+ )
92
+
93
+ set_random_seed(7)
94
+ with set_current_vllm_config(
95
+ VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
96
+ ):
97
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
98
+
99
+ quant_blocksize = 16
100
+ is_gated_act = activation.is_gated
101
+
102
+ w1_q, w2_q, quant_config = make_test_quant_config(
103
+ e,
104
+ n,
105
+ k,
106
+ in_dtype=dtype,
107
+ quant_dtype="nvfp4",
108
+ block_shape=None,
109
+ per_act_token_quant=False,
110
+ make_gate=is_gated_act,
111
+ # The TRT-LLM FP4 MoE kernel rejects swizzled (padded) activation
112
+ # scales — its numel-based vec_size check requires numel == M*K/16.
113
+ # Match what oracle/nvfp4.py does for this backend.
114
+ is_nvfp4_scale_swizzled=False,
115
+ )
116
+
117
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
118
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
119
+
120
+ moe_config = FusedMoEConfig(
121
+ num_experts=e,
122
+ experts_per_token=topk,
123
+ hidden_dim=k,
124
+ intermediate_size_per_partition=n,
125
+ num_local_experts=e,
126
+ num_logical_experts=e,
127
+ activation=activation,
128
+ device="cuda",
129
+ moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
130
+ in_dtype=dtype,
131
+ is_act_and_mul=is_gated_act,
132
+ routing_method=RoutingMethodType.TopK,
133
+ max_num_tokens=next_power_of_2(m),
134
+ )
135
+
136
+ trtllm_experts = mk.FusedMoEKernel(
137
+ maybe_make_prepare_finalize(
138
+ moe=moe_config,
139
+ quant_config=quant_config,
140
+ allow_new_interface=True,
141
+ use_monolithic=False,
142
+ ),
143
+ TrtLlmNvFp4ExpertsModular(moe_config=moe_config, quant_config=quant_config),
144
+ inplace=False,
145
+ )
146
+
147
+ trtllm_output = trtllm_experts.apply(
148
+ hidden_states=a,
149
+ w1=w1_q,
150
+ w2=w2_q,
151
+ topk_weights=topk_weights,
152
+ topk_ids=topk_ids,
153
+ activation=activation,
154
+ global_num_experts=e,
155
+ expert_map=None,
156
+ apply_router_weight_on_input=False,
157
+ )
158
+
159
+ # Reference: round-trip activations and weights through FP4
160
+ # quant/dequant so the comparison isolates kernel/activation behavior
161
+ # from quantization error.
162
+ a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / a.abs().max()).to(
163
+ torch.float32
164
+ )
165
+ a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
166
+ a_in_dtype = dequantize_nvfp4_to_dtype(
167
+ a_fp4,
168
+ a_scale_interleaved,
169
+ a_global_scale,
170
+ dtype=a.dtype,
171
+ device=a.device,
172
+ block_size=quant_blocksize,
173
+ )
174
+
175
+ w1_d = torch.empty(
176
+ (e, (2 if is_gated_act else 1) * n, k), device="cuda", dtype=dtype
177
+ )
178
+ w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
179
+ for idx in range(e):
180
+ w1_d[idx] = dequantize_nvfp4_to_dtype(
181
+ w1_q[idx],
182
+ quant_config.w1_scale[idx],
183
+ (1 / quant_config.g1_alphas[idx]),
184
+ dtype=dtype,
185
+ device=w1_q.device,
186
+ block_size=quant_blocksize,
187
+ )
188
+ w2_d[idx] = dequantize_nvfp4_to_dtype(
189
+ w2_q[idx],
190
+ quant_config.w2_scale[idx],
191
+ (1 / quant_config.g2_alphas[idx]),
192
+ dtype=dtype,
193
+ device=w2_q.device,
194
+ block_size=quant_blocksize,
195
+ )
196
+
197
+ torch_output = torch_moe(
198
+ a_in_dtype, w1_d, w2_d, score, topk, activation=activation
199
+ )
200
+
201
+ torch.testing.assert_close(torch_output, trtllm_output, atol=2e-1, rtol=2e-1)
202
+
203
+
204
+if __name__ == "__main__":
205
+ test_trtllm_fp4_moe_no_graph(
206
+ 64, 704, 4096, 128, 8, torch.bfloat16, MoEActivation.GELU, None
207
+ )
tests/kernels/moe/utils.py
+2
@@ -402,6 +402,7 @@ def make_test_quant_config(
402
per_act_token_quant: bool = False,
403
block_shape: list[int] | None = None,
404
make_gate: bool = True,
405
+ is_nvfp4_scale_swizzled: bool = True,
406
) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]:
407
(_, w1, w1_s, w1_gs), (_, w2, w2_s, w2_gs) = make_test_weights(
408
e,
@@ -442,6 +443,7 @@ def make_test_quant_config(
443
# TODO: make sure this is handled properly
444
g1_alphas=(1 / w1_gs) if w1_gs is not None else None,
445
g2_alphas=(1 / w2_gs) if w2_gs is not None else None,
446
+ is_nvfp4_scale_swizzled=is_nvfp4_scale_swizzled,
447
),
448
)
449
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+8
-4
@@ -112,8 +112,12 @@ class TrtLlmNvFp4ExpertsBase:
112
113
@staticmethod
114
def _supports_activation(activation: MoEActivation) -> bool:
115
- """Supports only SiLU and RELU^2 non-gated activation."""
116
- return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
115
+ """Supports only SiLU, RELU^2 non-gated and GELU activation."""
116
+ return activation in [
117
+ MoEActivation.SILU,
118
+ MoEActivation.RELU2_NO_MUL,
119
+ MoEActivation.GELU,
120
+ ]
121
122
@staticmethod
123
def _supports_shape(hidden_dim: int) -> bool:
@@ -190,7 +194,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
194
):
195
import flashinfer
196
193
- assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
197
+ assert self._supports_activation(activation)
198
assert a1q_scale is not None
199
assert self.quant_config.w1_scale is not None
200
assert self.quant_config.w2_scale is not None
@@ -292,7 +296,7 @@ class TrtLlmNvFp4ExpertsMonolithic(
296
) -> torch.Tensor:
297
import flashinfer
298
295
- assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
299
+ assert self._supports_activation(activation)
300
assert a1q_scale is not None
301
assert self.quant_config.w1_scale is not None
302
assert self.quant_config.w2_scale is not None