Enable B12x backend for non-gated MoEs (like Nemotron) (#43328)
Signed-off-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com>
Andrii Skliar committed
Jul 6, 2026 at 21:40 UTC
b1384f5ec6b03140fda8c291aec07d35ce3d5c4c
3 files changed
+225
-40
tests/kernels/moe/test_flashinfer_b12x_moe.py
+157
-22
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from types import SimpleNamespace
5
+
6
import pytest
7
import torch
8
@@ -8,8 +10,7 @@ from vllm.platforms import current_platform
10
11
if not current_platform.is_device_capability_family(120):
12
pytest.skip(
11
- reason="FlashInfer CuteDSL SM12x MoE requires SM120 "
12
- "(RTX Pro 6000 / DGX Spark).",
13
+ reason="FlashInfer B12x MoE requires SM120 (RTX Pro 6000 / DGX Spark).",
14
allow_module_level=True,
15
)
16
@@ -18,8 +19,8 @@ from vllm.utils.flashinfer import has_flashinfer_b12x_moe
19
if not has_flashinfer_b12x_moe():
20
pytest.skip(
21
reason=(
21
- "FlashInfer cute_dsl_fused_moe_nvfp4 / convert_sf_to_mma_layout "
22
- "not available in installed FlashInfer (needs PRs #3051 and #3066)."
22
+ "FlashInfer B12xMoEWrapper not available in installed "
23
+ "FlashInfer (needs PR #3080)."
24
),
25
allow_module_level=True,
26
)
@@ -40,7 +41,6 @@ from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
41
from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import (
42
FlashInferB12xExperts,
43
)
43
-from vllm.utils.flashinfer import flashinfer_convert_sf_to_mma_layout
44
from vllm.utils.torch_utils import set_random_seed
45
46
# Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256,
@@ -59,7 +59,7 @@ def _reorder_gate_up_to_up_gate(
59
) -> tuple[torch.Tensor, torch.Tensor]:
60
"""Swap gate and up-projection halves along dim=1 to [up, gate] order.
61
62
- The SM12x kernel expects weights in [up (w3), gate (w1)] order while the
62
+ The B12x kernel expects weights in [up (w3), gate (w1)] order while the
63
BF16 reference uses [gate (w1), up (w3)]. This replicates the reordering
64
done at model-load time by ``prepare_nvfp4_moe_layer_for_fi_or_cutlass``.
65
"""
@@ -70,6 +70,22 @@ def _reorder_gate_up_to_up_gate(
70
)
71
72
73
+def _process_b12x_weights(
74
+ experts: FlashInferB12xExperts,
75
+ w1_scale: torch.Tensor,
76
+ w2_scale: torch.Tensor,
77
+ w1_scale_2: torch.Tensor,
78
+ w2_scale_2: torch.Tensor,
79
+) -> None:
80
+ layer = SimpleNamespace(
81
+ w13_weight_scale=w1_scale,
82
+ w13_weight_scale_2=w1_scale_2,
83
+ w2_weight_scale=w2_scale,
84
+ w2_weight_scale_2=w2_scale_2,
85
+ )
86
+ experts.process_weights_after_loading(layer)
87
+
88
+
89
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
90
@pytest.mark.parametrize("e", [8, 16])
91
@pytest.mark.parametrize("topk", [1, 2, 4])
@@ -174,22 +190,12 @@ def test_flashinfer_b12x_moe(
190
moe_config=moe_config,
191
quant_config=quant_config,
192
)
177
- # In production, process_weights_after_loading computes these after
178
- # normalizing block scales. In the test the scales are already in final
179
- # form (global_scale=1.0), so we compute the MMA layouts directly.
180
- num_experts_w1, m1, k1_sf = w1_blockscale.shape
181
- experts.w1_sf_mma = flashinfer_convert_sf_to_mma_layout(
182
- w1_blockscale.reshape(num_experts_w1 * m1, k1_sf),
183
- m=m1,
184
- k=k1_sf * 16,
185
- num_groups=num_experts_w1,
186
- )
187
- num_experts_w2, m2, k2_sf = w2_blockscale.shape
188
- experts.w2_sf_mma = flashinfer_convert_sf_to_mma_layout(
189
- w2_blockscale.reshape(num_experts_w2 * m2, k2_sf),
190
- m=m2,
191
- k=k2_sf * 16,
192
- num_groups=num_experts_w2,
193
+ _process_b12x_weights(
194
+ experts,
195
+ w1_blockscale,
196
+ w2_blockscale,
197
+ ones_e,
198
+ ones_e,
199
)
200
201
kernel = mk.FusedMoEKernel(
@@ -224,5 +230,134 @@ def test_flashinfer_b12x_moe(
230
torch.testing.assert_close(sm12x_output, torch_output, atol=2e-1, rtol=2e-1)
231
232
233
+@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
234
+@pytest.mark.parametrize("e", [8, 16])
235
+@pytest.mark.parametrize("topk", [1, 2, 4])
236
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
237
+@torch.inference_mode()
238
+def test_flashinfer_b12x_moe_relu2(
239
+ m: int,
240
+ n: int,
241
+ k: int,
242
+ e: int,
243
+ topk: int,
244
+ dtype: torch.dtype,
245
+ workspace_init,
246
+):
247
+ """Test FlashInferB12xExperts with ReLU2 (non-gated) activation.
248
+
249
+ ReLU2 is used by Nemotron-H style models. Unlike the gated SiLU
250
+ path, w1 has shape [E, N, K] (not [E, 2N, K]) and the activation
251
+ is relu(x)^2 without a gate/up split.
252
+ """
253
+ set_random_seed(7)
254
+ with set_current_vllm_config(
255
+ VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
256
+ ):
257
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
258
+
259
+ # Non-gated: w1 shape is (e, n, k), not (e, 2n, k).
260
+ w1_bf16 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15
261
+ w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15
262
+
263
+ gs = torch.ones(1, device="cuda", dtype=torch.float32)
264
+ sf_vec_size = 16
265
+
266
+ # W1: no gate/up reordering for non-gated.
267
+ w1_flat = w1_bf16.reshape(e * n, k)
268
+ w1_q_flat, w1_sf_flat = fp4_quantize(
269
+ w1_flat,
270
+ global_scale=gs,
271
+ sf_vec_size=sf_vec_size,
272
+ is_sf_swizzled_layout=True,
273
+ )
274
+ w1_q = w1_q_flat.view(e, n, k // 2)
275
+ w1_blockscale = w1_sf_flat.view(e, n, w1_sf_flat.shape[1])
276
+
277
+ w2_flat = w2_bf16.reshape(e * k, n)
278
+ w2_q_flat, w2_sf_flat = fp4_quantize(
279
+ w2_flat,
280
+ global_scale=gs,
281
+ sf_vec_size=sf_vec_size,
282
+ is_sf_swizzled_layout=True,
283
+ )
284
+ w2_q = w2_q_flat.view(e, k, n // 2)
285
+ w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1])
286
+
287
+ ones_e = torch.ones(e, device="cuda", dtype=torch.float32)
288
+
289
+ quant_config = nvfp4_moe_quant_config(
290
+ g1_alphas=ones_e,
291
+ g2_alphas=ones_e,
292
+ a1_gscale=ones_e,
293
+ a2_gscale=ones_e,
294
+ w1_scale=w1_blockscale,
295
+ w2_scale=w2_blockscale,
296
+ )
297
+
298
+ moe_config = make_dummy_moe_config(
299
+ num_experts=e,
300
+ experts_per_token=topk,
301
+ hidden_dim=k,
302
+ intermediate_size=n,
303
+ in_dtype=dtype,
304
+ activation=MoEActivation.RELU2_NO_MUL,
305
+ )
306
+
307
+ experts = FlashInferB12xExperts(
308
+ moe_config=moe_config,
309
+ quant_config=quant_config,
310
+ )
311
+ _process_b12x_weights(
312
+ experts,
313
+ w1_blockscale,
314
+ w2_blockscale,
315
+ ones_e,
316
+ ones_e,
317
+ )
318
+
319
+ kernel = mk.FusedMoEKernel(
320
+ maybe_make_prepare_finalize(
321
+ moe=moe_config,
322
+ quant_config=quant_config,
323
+ allow_new_interface=True,
324
+ use_monolithic=False,
325
+ ),
326
+ experts,
327
+ inplace=False,
328
+ )
329
+
330
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
331
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
332
+
333
+ b12x_output = kernel.apply(
334
+ hidden_states=a,
335
+ w1=w1_q,
336
+ w2=w2_q,
337
+ topk_weights=topk_weights,
338
+ topk_ids=topk_ids,
339
+ global_num_experts=e,
340
+ activation=MoEActivation.RELU2_NO_MUL,
341
+ apply_router_weight_on_input=False,
342
+ expert_map=None,
343
+ )
344
+
345
+ torch_output = torch_moe(
346
+ a,
347
+ w1_bf16,
348
+ w2_bf16,
349
+ score,
350
+ topk,
351
+ activation=MoEActivation.RELU2_NO_MUL,
352
+ )
353
+
354
+ torch.testing.assert_close(
355
+ b12x_output,
356
+ torch_output,
357
+ atol=2e-1,
358
+ rtol=2e-1,
359
+ )
360
+
361
+
362
if __name__ == "__main__":
363
test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16)
tests/kernels/moe/utils.py
+2
-1
@@ -55,6 +55,7 @@ def make_dummy_moe_config(
55
intermediate_size: int = 1,
56
in_dtype: torch.dtype = torch.bfloat16,
57
max_num_tokens: int = 512,
58
+ activation: MoEActivation = MoEActivation.SILU,
59
) -> FusedMoEConfig:
60
"""
61
This is a dummy config for the mk constructor interface
@@ -73,7 +74,7 @@ def make_dummy_moe_config(
74
else num_experts,
75
num_logical_experts=num_experts,
76
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
76
- activation=MoEActivation.SILU,
77
+ activation=activation,
78
in_dtype=in_dtype,
79
device="cuda",
80
routing_method=RoutingMethodType.TopK,
vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py
+66
-17
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from typing import Any
5
+
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -20,7 +22,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
)
23
from vllm.platforms import current_platform
24
from vllm.utils.flashinfer import (
23
- flashinfer_b12x_fused_moe,
25
flashinfer_convert_sf_to_mma_layout,
26
has_flashinfer_b12x_moe,
27
)
@@ -42,6 +43,11 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
43
Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported.
44
"""
45
46
+ _ACTIVATION_MAP: dict[MoEActivation, str] = {
47
+ MoEActivation.SILU: "silu",
48
+ MoEActivation.RELU2_NO_MUL: "relu2",
49
+ }
50
+
51
def __init__(
52
self,
53
moe_config: FusedMoEConfig,
@@ -60,6 +66,30 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
66
# one. Holding it on the instance keeps apply() alloc-free.
67
self._fc2_input_scale: torch.Tensor | None = None
68
69
+ # Shape params for B12xMoEWrapper construction.
70
+ self.global_num_experts = moe_config.num_experts
71
+ self.topk = moe_config.experts_per_token
72
+ self.hidden_dim = moe_config.hidden_dim
73
+ self.intermediate_size_per_partition = (
74
+ moe_config.intermediate_size_per_partition
75
+ )
76
+ self.max_num_tokens = moe_config.max_num_tokens
77
+ self.local_expert_offset = self.ep_rank * self.num_local_experts
78
+
79
+ activation = moe_config.activation
80
+ if activation not in self._ACTIVATION_MAP:
81
+ raise ValueError(
82
+ f"FlashInferB12xExperts does not support "
83
+ f"activation {activation!r}. "
84
+ f"Supported: {list(self._ACTIVATION_MAP.keys())}"
85
+ )
86
+ self._activation_str = self._ACTIVATION_MAP[activation]
87
+
88
+ # Lazily created on first apply() call.
89
+ self._wrapper: Any | None = None
90
+ self.w1_sf_mma: torch.Tensor | None = None
91
+ self.w2_sf_mma: torch.Tensor | None = None
92
+
93
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
94
# Normalise block scales to absorb the per-expert weight global scale
95
# (w_gs). vLLM's NVFP4 convention stores:
@@ -141,7 +171,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
171
172
@staticmethod
173
def _supports_no_act_and_mul() -> bool:
144
- return False
174
+ return True
175
176
@staticmethod
177
def _supports_quant_scheme(
@@ -158,11 +188,13 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
188
189
@staticmethod
190
def _supports_activation(activation: MoEActivation) -> bool:
161
- return activation == MoEActivation.SILU
191
+ return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL)
192
193
@staticmethod
194
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
165
- return True
195
+ # B12xMoEWrapper does not yet support expert parallelism: its local
196
+ # expert count must equal the global expert count.
197
+ return not moe_parallel_config.use_ep
198
199
def supports_expert_map(self) -> bool:
200
return False
@@ -190,13 +222,29 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
222
223
@property
224
def expects_unquantized_inputs(self) -> bool:
193
- # b12x_fused_moe expects BF16 hidden states and performs its own FP4
225
+ # B12xMoEWrapper expects BF16 hidden states and performs its own FP4
226
# quantization internally. Returning True prevents the modular kernel
195
- # from pre-quantizing activations, which would produce an FP4-packed
196
- # tensor with size(-1)=k//2 and break the scale-factor conversion that
197
- # expects size(-1)=k.
227
+ # from pre-quantizing activations.
228
return True
229
230
+ def _ensure_wrapper(self) -> None:
231
+ """Lazily create B12xMoEWrapper on first use."""
232
+ if self._wrapper is not None:
233
+ return
234
+
235
+ from flashinfer.fused_moe import B12xMoEWrapper
236
+
237
+ self._wrapper = B12xMoEWrapper(
238
+ num_experts=self.global_num_experts,
239
+ top_k=self.topk,
240
+ hidden_size=self.hidden_dim,
241
+ intermediate_size=self.intermediate_size_per_partition,
242
+ use_cuda_graph=True,
243
+ max_num_tokens=self.max_num_tokens,
244
+ num_local_experts=self.num_local_experts,
245
+ activation=self._activation_str,
246
+ )
247
+
248
def apply(
249
self,
250
output: torch.Tensor,
@@ -224,13 +272,16 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
272
assert self._fc2_input_scale is not None, (
273
"_fc2_input_scale must be set by process_weights_after_loading"
274
)
275
+ assert self.w1_sf_mma is not None and self.w2_sf_mma is not None, (
276
+ "process_weights_after_loading must run before FlashInferB12xExperts.apply"
277
+ )
278
228
- top_k = topk_ids.shape[1]
279
+ self._ensure_wrapper()
280
+ wrapper = self._wrapper
281
+ assert wrapper is not None
282
230
- flashinfer_b12x_fused_moe(
283
+ wrapper_output = wrapper.run(
284
x=hidden_states,
232
- token_selected_experts=topk_ids.to(torch.int32),
233
- token_final_scales=topk_weights,
285
w1_weight=w1,
286
w1_weight_sf=self.w1_sf_mma,
287
w1_alpha=self.g1_alphas,
@@ -238,9 +289,7 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
289
w2_weight=w2,
290
w2_weight_sf=self.w2_sf_mma,
291
w2_alpha=self.g2_alphas,
241
- num_experts=global_num_experts,
242
- top_k=top_k,
243
- num_local_experts=self.num_local_experts,
244
- output_dtype=self.out_dtype,
245
- output=output,
292
+ token_selected_experts=topk_ids.to(torch.int32),
293
+ token_final_scales=topk_weights,
294
)
295
+ output.copy_(wrapper_output)