[Kernel] Add b12x FP4 MoE backend (#52018)
Signed-off-by: Luke Alonso <lalonso@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Luke Alonso committed
Aug 21, 2026 at 08:04 UTC
2740c817ffbe6c1980507c7cbd27b4bc94d2c856
23 files changed
+2417
-60
.buildkite/test_areas/kernels.yaml
+21
-4
@@ -343,9 +343,9 @@ steps:
343
# e2e
344
- pytest -v -s tests/models/quantization/test_nvfp4.py
345
346
-- label: ":nvidia: (DGX) Spark B12X Linear Kernels Nightly"
346
+- label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly"
347
key: b12x-linear-kernels-dgx-spark-nightly
348
- timeout_in_minutes: 30
348
+ timeout_in_minutes: 60
349
device: dgx-spark
350
optional: true
351
num_devices: 1
@@ -353,17 +353,34 @@ steps:
353
- arm64-image-build
354
source_file_dependencies:
355
- setup.py
356
+ - vllm/config/kernel.py
357
+ - vllm/envs.py
358
- vllm/model_executor/kernels/linear/
359
+ - vllm/model_executor/layers/fused_moe/b12x.py
360
+ - vllm/model_executor/layers/fused_moe/config.py
361
+ - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
362
+ - vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
363
+ - vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/
364
+ - vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
365
+ - vllm/model_executor/layers/quantization/modelopt.py
366
+ - vllm/model_executor/layers/quantization/mxfp4.py
367
+ - vllm/model_executor/layers/quantization/utils/b12x_moe.py
368
- vllm/model_executor/warmup/b12x_warmup.py
369
- vllm/utils/b12x.py
370
+ - tests/kernels/moe/test_b12x.py
371
+ - tests/kernels/quantization/nvfp4_utils.py
372
- tests/model_executor/kernels/test_b12x_linear.py
373
- tests/model_executor/test_b12x_warmup.py
374
- tests/kernels/quantization/test_block_fp8.py
375
+ - tests/quantization/reference_mxfp4.py
376
+ - tests/quantization/test_auto_round.py
377
commands:
363
- - uv pip install --system b12x==1.2.4
378
+ - uv pip install --system b12x==1.2.6
379
- pytest -v -s model_executor/kernels/test_b12x_linear.py
380
model_executor/test_b12x_warmup.py
366
- kernels/quantization/test_block_fp8.py -k b12x
381
+ kernels/quantization/test_block_fp8.py
382
+ kernels/moe/test_b12x.py
383
+ quantization/test_auto_round.py -k b12x
384
385
- label: ":nvidia: (H100) Helion Kernels Shard %N"
386
key: kernels-helion-test
docs/features/quantization/b12x.md
+20
-7
@@ -1,24 +1,37 @@
1
-# B12X Linear Backend
1
+# b12x Linear and MoE Backends
2
3
-[B12X](https://pypi.org/project/b12x/) provides optional CUDA kernels for
3
+[b12x](https://pypi.org/project/b12x/) provides optional CUDA kernels for
4
NVIDIA SM120 and SM121 GPUs. Install the dependency with:
5
6
```bash
7
uv pip install "vllm[b12x]"
8
```
9
10
-B12X participates in automatic kernel selection after established optimized
11
-backends and before emulation. Select it explicitly with:
10
+b12x linear kernels participate in automatic selection after established
11
+optimized backends and before emulation. Select the linear and MoE backends
12
+explicitly with:
13
14
```bash
14
-vllm serve <model> --linear-backend b12x
15
+vllm serve <model> \
16
+ --linear-backend b12x \
17
+ --moe-backend b12x
18
```
19
20
+Only pass `--moe-backend b12x` for a compatible NVFP4 or MXFP4 MoE model. The
21
+linear and MoE backends can be selected independently.
22
+
23
+b12x uses MXFP8 activations by default for MXFP4 MoE and the checkpoint's
24
+activation format for NVFP4 MoE. MXFP4 falls back to BF16 when its A8 path does
25
+not support the model configuration. Set `VLLM_B12X_MOE_FP4_FORCE_A16=1` to
26
+force BF16 activations for either FP4 weight format.
27
+
28
## Supported Configurations
29
30
| Backend | Supported configurations |
31
| ------- | ------------------------ |
32
| Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4, and MXFP4 |
33
+| MoE | Tensor-parallel MXFP4 weights with BF16 or MXFP8 activations; NVFP4 weights with BF16, NVFP4, or MXFP8 activations |
34
23
-Dense W4A16 layers are not handled by B12X and continue to use another
24
-compatible backend such as Marlin.
35
+Dense W4A16 layers are not handled by b12x and continue to use another
36
+compatible backend such as Marlin. The b12x MoE backend does not support expert
37
+parallelism, expert maps, EXL3, or NF3.
setup.py
+1
-1
@@ -1530,7 +1530,7 @@ setup(
1530
# only; also needs system GStreamer + libv4l (see docs).
1531
"deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"],
1532
"flashinfer": [], # Kept for backwards compatibility
1533
- "b12x": ["b12x==1.2.4"],
1533
+ "b12x": ["b12x==1.2.6"],
1534
# Optional deps for Helion kernel development
1535
# NOTE: When updating helion version, also update CI files:
1536
# - .buildkite/test_areas/kernels.yaml
tests/kernels/moe/test_b12x.py
new
+1245
@@ -0,0 +1,1245 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the b12x tensor-parallel MoE integration."""
4
+
5
+from dataclasses import dataclass, replace
6
+from types import SimpleNamespace
7
+
8
+import pytest
9
+import torch
10
+
11
+import vllm.model_executor.layers.fused_moe.b12x as b12x
12
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
+import vllm.model_executor.layers.fused_moe.oracle.mxfp4 as mxfp4_oracle
14
+import vllm.model_executor.layers.fused_moe.oracle.nvfp4 as nvfp4_oracle
15
+from tests.kernels.moe.utils import make_dummy_moe_config
16
+from tests.kernels.quantization.nvfp4_utils import (
17
+ dequantize_nvfp4_to_dtype,
18
+ quant_nvfp4_tensor,
19
+)
20
+from tests.kernels.utils import torch_moe
21
+from tests.quantization.reference_mxfp4 import dq_mxfp4_torch
22
+from vllm import _custom_ops as ops
23
+from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
24
+from vllm.model_executor.layers.fused_moe import fused_topk
25
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
26
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
27
+ maybe_make_prepare_finalize,
28
+)
29
+from vllm.model_executor.layers.fused_moe.b12x import B12xExperts
30
+from vllm.model_executor.layers.fused_moe.config import (
31
+ FusedMoEParallelConfig,
32
+ FusedMoEQuantConfig,
33
+ mxfp4_w4a16_moe_quant_config,
34
+ nvfp4_w4a16_moe_quant_config,
35
+)
36
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
37
+ Mxfp4MoeBackend,
38
+ select_deepseek_v4_mxfp4_moe_backend,
39
+ select_mxfp4_moe_backend,
40
+)
41
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
42
+ NvFp4MoeBackend,
43
+ select_nvfp4_moe_backend,
44
+)
45
+from vllm.model_executor.layers.quantization.utils.b12x_moe import (
46
+ prepare_nvfp4_moe_layer_for_b12x,
47
+)
48
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize
49
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
50
+ kMxfp4Static,
51
+ kMxfp8Dynamic,
52
+ kNvfp4Dynamic,
53
+ kNvfp4Static,
54
+)
55
+from vllm.platforms import current_platform
56
+from vllm.utils.torch_utils import set_random_seed
57
+
58
+
59
+def _quantize_nvfp4_linear(
60
+ weight: torch.Tensor,
61
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
62
+ weights_q = []
63
+ scales = []
64
+ global_scales = []
65
+ for expert_weight in weight:
66
+ weight_q, scale, global_scale = quant_nvfp4_tensor(
67
+ expert_weight,
68
+ is_sf_swizzled_layout=False,
69
+ )
70
+ weights_q.append(weight_q)
71
+ scales.append(scale)
72
+ global_scales.append(global_scale)
73
+ return torch.stack(weights_q), torch.stack(scales), torch.stack(global_scales)
74
+
75
+
76
+def _dequantize_nvfp4_linear(
77
+ tensor_fp4: torch.Tensor,
78
+ tensor_sf: torch.Tensor,
79
+ global_scale: torch.Tensor,
80
+ dtype: torch.dtype,
81
+) -> torch.Tensor:
82
+ return dequantize_nvfp4_to_dtype(
83
+ tensor_fp4,
84
+ tensor_sf,
85
+ global_scale,
86
+ dtype=dtype,
87
+ device=tensor_fp4.device,
88
+ is_sf_linear_layout=True,
89
+ )
90
+
91
+
92
+def _nvfp4_activation_reference(
93
+ hidden_states: torch.Tensor,
94
+ w1: torch.Tensor,
95
+ w2: torch.Tensor,
96
+ topk_weights: torch.Tensor,
97
+ topk_ids: torch.Tensor,
98
+ a1_scale: torch.Tensor,
99
+ a2_scale: torch.Tensor,
100
+) -> torch.Tensor:
101
+ tokens, hidden_size = hidden_states.shape
102
+ topk = topk_ids.shape[1]
103
+ routed_input = (
104
+ hidden_states[:, None, :]
105
+ .expand(-1, topk, -1)
106
+ .reshape(tokens * topk, hidden_size)
107
+ )
108
+ routed_output = torch.zeros(
109
+ tokens * topk,
110
+ hidden_size,
111
+ dtype=torch.float32,
112
+ device=hidden_states.device,
113
+ )
114
+ flat_ids = topk_ids.reshape(-1)
115
+
116
+ for expert in range(w1.shape[0]):
117
+ mask = flat_ids == expert
118
+ if not mask.any():
119
+ continue
120
+ a1_q, a1_block_scale = ops.scaled_fp4_quant(
121
+ routed_input[mask],
122
+ a1_scale[expert],
123
+ is_sf_swizzled_layout=False,
124
+ )
125
+ a1 = _dequantize_nvfp4_linear(
126
+ a1_q,
127
+ a1_block_scale,
128
+ a1_scale[expert],
129
+ torch.float32,
130
+ )
131
+ fc1 = a1 @ w1[expert].float().t()
132
+ gate, up = fc1.chunk(2, dim=-1)
133
+ intermediate = (torch.nn.functional.silu(gate) * up).to(torch.bfloat16)
134
+ a2_q, a2_block_scale = ops.scaled_fp4_quant(
135
+ intermediate,
136
+ a2_scale[expert],
137
+ is_sf_swizzled_layout=False,
138
+ )
139
+ a2 = _dequantize_nvfp4_linear(
140
+ a2_q,
141
+ a2_block_scale,
142
+ a2_scale[expert],
143
+ torch.float32,
144
+ )
145
+ routed_output[mask] = a2 @ w2[expert].float().t()
146
+
147
+ return (
148
+ routed_output.view(tokens, topk, hidden_size)
149
+ .mul(topk_weights[..., None])
150
+ .sum(dim=1)
151
+ .to(hidden_states.dtype)
152
+ )
153
+
154
+
155
+def _has_b12x_moe() -> bool:
156
+ return (
157
+ torch.cuda.is_available()
158
+ and current_platform.is_device_capability_family(120)
159
+ and B12xExperts._supports_current_device()
160
+ )
161
+
162
+
163
+def _count_fp4_negative_zeros(packed: torch.Tensor) -> int:
164
+ low = (packed & 0x0F) == 0x08
165
+ high = (packed & 0xF0) == 0x80
166
+ return int(low.sum().item() + high.sum().item())
167
+
168
+
169
+def _make_b12x_moe_kernel(
170
+ hidden_states: torch.Tensor,
171
+ w1: torch.Tensor,
172
+ w2: torch.Tensor,
173
+ topk: int,
174
+ activation: MoEActivation,
175
+ quant_config: FusedMoEQuantConfig,
176
+) -> mk.FusedMoEKernel:
177
+ num_experts = w1.shape[0]
178
+ moe_config = make_dummy_moe_config(
179
+ num_experts=num_experts,
180
+ experts_per_token=topk,
181
+ hidden_dim=hidden_states.shape[1],
182
+ intermediate_size=w2.shape[2] * 2,
183
+ in_dtype=hidden_states.dtype,
184
+ activation=activation,
185
+ )
186
+ experts = B12xExperts(moe_config, quant_config)
187
+ layer = SimpleNamespace(
188
+ activation=activation,
189
+ apply_router_weight_on_input=False,
190
+ w13_weight=w1,
191
+ w2_weight=w2,
192
+ w13_weight_scale=quant_config.w1_scale,
193
+ w2_weight_scale=quant_config.w2_scale,
194
+ )
195
+ if quant_config.weight_quant_dtype == "nvfp4":
196
+ layer.w13_weight_scale_2 = quant_config.g1_alphas
197
+ layer.w2_weight_scale_2 = quant_config.g2_alphas
198
+ if quant_config.quant_dtype is not None:
199
+ assert quant_config.a1_gscale is not None
200
+ assert quant_config.a2_gscale is not None
201
+ layer.w13_input_scale = 1.0 / quant_config.a1_gscale
202
+ layer.w2_input_scale = 1.0 / quant_config.a2_gscale
203
+ experts.process_weights_after_loading(layer)
204
+ return mk.FusedMoEKernel(
205
+ maybe_make_prepare_finalize(
206
+ moe=moe_config,
207
+ quant_config=quant_config,
208
+ allow_new_interface=True,
209
+ use_monolithic=False,
210
+ ),
211
+ experts,
212
+ )
213
+
214
+
215
+def _run_b12x_moe(
216
+ hidden_states: torch.Tensor,
217
+ w1: torch.Tensor,
218
+ w2: torch.Tensor,
219
+ score: torch.Tensor,
220
+ topk: int,
221
+ activation: MoEActivation,
222
+ quant_config: FusedMoEQuantConfig,
223
+) -> torch.Tensor:
224
+ num_experts = w1.shape[0]
225
+ kernel = _make_b12x_moe_kernel(
226
+ hidden_states,
227
+ w1,
228
+ w2,
229
+ topk,
230
+ activation,
231
+ quant_config,
232
+ )
233
+ topk_weights, topk_ids, _ = fused_topk(
234
+ hidden_states, score, topk, renormalize=False
235
+ )
236
+ return kernel.apply(
237
+ hidden_states=hidden_states,
238
+ w1=w1,
239
+ w2=w2,
240
+ topk_weights=topk_weights,
241
+ topk_ids=topk_ids,
242
+ activation=activation,
243
+ global_num_experts=num_experts,
244
+ expert_map=None,
245
+ apply_router_weight_on_input=False,
246
+ )
247
+
248
+
249
+def _quant_config(weight_dtype: str, activation_dtype: str | None):
250
+ scale = torch.ones(1, dtype=torch.float32)
251
+ return FusedMoEQuantConfig.make(
252
+ quant_dtype=activation_dtype,
253
+ weight_dtype=weight_dtype,
254
+ w1_scale=scale,
255
+ w2_scale=scale,
256
+ g1_alphas=scale,
257
+ g2_alphas=scale,
258
+ a1_gscale=scale,
259
+ a2_gscale=scale,
260
+ )
261
+
262
+
263
+def test_b12x_moe_supports_only_tensor_parallel() -> None:
264
+ parallel = FusedMoEParallelConfig.make_no_parallel()
265
+
266
+ assert B12xExperts._supports_parallel_config(parallel)
267
+ assert not B12xExperts._supports_parallel_config(
268
+ replace(parallel, use_ep=True, ep_size=2)
269
+ )
270
+ all2all = replace(parallel, use_ep=True, dp_size=2)
271
+ assert all2all.use_all2all_kernels
272
+ assert not B12xExperts._supports_parallel_config(all2all)
273
+ assert not B12xExperts._supports_parallel_config(
274
+ replace(parallel, enable_eplb=True)
275
+ )
276
+
277
+
278
+_SITU_REASON = "kernel supports only SiTU beta=4 and linear_beta=25"
279
+_UNINTERLEAVED_W4A8_REASON = "kernel does not support swigluoai_uninterleave with W4A8"
280
+
281
+
282
+@pytest.mark.parametrize(
283
+ "config_kwargs,overrides,weight_key,activation_key,expected_reason",
284
+ [
285
+ pytest.param(
286
+ {"hidden_dim": 128, "activation": MoEActivation.SWIGLUOAI},
287
+ {},
288
+ kMxfp4Static,
289
+ None,
290
+ "kernel does not support MoEActivation.SWIGLUOAI activation",
291
+ id="interleaved-swigluoai",
292
+ ),
293
+ pytest.param(
294
+ {
295
+ "hidden_dim": 128,
296
+ "activation": MoEActivation.SWIGLUOAI_UNINTERLEAVE,
297
+ },
298
+ {},
299
+ kMxfp4Static,
300
+ kMxfp8Dynamic,
301
+ _UNINTERLEAVED_W4A8_REASON,
302
+ id="mxfp4-w4a8-uninterleaved-swigluoai",
303
+ ),
304
+ pytest.param(
305
+ {
306
+ "hidden_dim": 128,
307
+ "activation": MoEActivation.SWIGLUOAI_UNINTERLEAVE,
308
+ },
309
+ {},
310
+ kMxfp4Static,
311
+ None,
312
+ None,
313
+ id="w4a16-uninterleaved-swigluoai",
314
+ ),
315
+ pytest.param(
316
+ {"activation": MoEActivation.RELU2_NO_MUL},
317
+ {},
318
+ kMxfp4Static,
319
+ kMxfp8Dynamic,
320
+ "MXFP4 W4A8 supports only SiLU and SiTU",
321
+ id="mxfp4-w4a8-relu2",
322
+ ),
323
+ pytest.param(
324
+ {"hidden_dim": 128},
325
+ {},
326
+ kMxfp4Static,
327
+ kMxfp8Dynamic,
328
+ (
329
+ "MXFP4 W4A8 requires hidden size divisible by 256 and per-rank "
330
+ "intermediate size divisible by 32"
331
+ ),
332
+ id="mxfp4-w4a8-alignment",
333
+ ),
334
+ pytest.param(
335
+ {"hidden_dim": 128, "intermediate_size": 48},
336
+ {"intermediate_size_per_partition": 64},
337
+ kMxfp4Static,
338
+ None,
339
+ "MXFP4 requires the per-rank intermediate size to be divisible by 32",
340
+ id="mxfp4-tp-scale-groups",
341
+ ),
342
+ pytest.param(
343
+ {"activation": MoEActivation.SITU},
344
+ {"activation_situ_beta": 3.0, "activation_situ_linear_beta": 25.0},
345
+ kMxfp4Static,
346
+ None,
347
+ _SITU_REASON,
348
+ id="situ-beta",
349
+ ),
350
+ pytest.param(
351
+ {"activation": MoEActivation.SITU},
352
+ {"activation_situ_beta": 4.0, "activation_situ_linear_beta": 25.0},
353
+ kMxfp4Static,
354
+ None,
355
+ None,
356
+ id="situ-standard-parameters",
357
+ ),
358
+ ],
359
+)
360
+def test_b12x_moe_config_support(
361
+ monkeypatch: pytest.MonkeyPatch,
362
+ config_kwargs,
363
+ overrides,
364
+ weight_key,
365
+ activation_key,
366
+ expected_reason: str | None,
367
+) -> None:
368
+ monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True)
369
+ config = make_dummy_moe_config(
370
+ **{"hidden_dim": 256, "intermediate_size": 64, **config_kwargs}
371
+ )
372
+ for name, value in overrides.items():
373
+ setattr(config, name, value)
374
+
375
+ supported, reason = B12xExperts.is_supported_config(
376
+ B12xExperts,
377
+ config,
378
+ weight_key,
379
+ activation_key,
380
+ mk.FusedMoEActivationFormat.Standard,
381
+ )
382
+
383
+ assert (supported, reason) == (expected_reason is None, expected_reason)
384
+
385
+
386
+@pytest.mark.parametrize(
387
+ "activation_key,force_a16,expected_backend",
388
+ [
389
+ (kMxfp8Dynamic, False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8),
390
+ (None, False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8),
391
+ ],
392
+)
393
+def test_explicit_b12x_mxfp4_selection(
394
+ monkeypatch: pytest.MonkeyPatch,
395
+ activation_key,
396
+ force_a16: bool,
397
+ expected_backend: Mxfp4MoeBackend,
398
+) -> None:
399
+ monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True)
400
+ monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None)
401
+ monkeypatch.setattr(
402
+ mxfp4_oracle.envs,
403
+ "VLLM_B12X_MOE_FP4_FORCE_A16",
404
+ force_a16,
405
+ )
406
+ config = make_dummy_moe_config(hidden_dim=256, intermediate_size=64)
407
+ config.moe_backend = "b12x"
408
+
409
+ backend, experts_cls = select_mxfp4_moe_backend(
410
+ config,
411
+ activation_key=activation_key,
412
+ )
413
+
414
+ assert backend == expected_backend
415
+ assert experts_cls is B12xExperts
416
+
417
+
418
+def test_explicit_b12x_mxfp4_force_a16_uses_a16_contract(
419
+ monkeypatch: pytest.MonkeyPatch,
420
+) -> None:
421
+ monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True)
422
+ monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None)
423
+ monkeypatch.setattr(
424
+ mxfp4_oracle.envs,
425
+ "VLLM_B12X_MOE_FP4_FORCE_A16",
426
+ True,
427
+ )
428
+ config = make_dummy_moe_config(hidden_dim=128, intermediate_size=64)
429
+ config.moe_backend = "b12x"
430
+
431
+ backend, experts_cls = select_mxfp4_moe_backend(
432
+ config,
433
+ activation_key=kMxfp8Dynamic,
434
+ )
435
+
436
+ assert backend == Mxfp4MoeBackend.B12X_MXFP4_BF16
437
+ assert experts_cls is B12xExperts
438
+
439
+
440
+@pytest.mark.parametrize(
441
+ "force_a16,expected_backend",
442
+ [
443
+ (False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8),
444
+ (True, Mxfp4MoeBackend.B12X_MXFP4_BF16),
445
+ ],
446
+)
447
+def test_deepseek_v4_b12x_activation_selection(
448
+ monkeypatch: pytest.MonkeyPatch,
449
+ force_a16: bool,
450
+ expected_backend: Mxfp4MoeBackend,
451
+) -> None:
452
+ monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True)
453
+ monkeypatch.setattr(
454
+ mxfp4_oracle.envs,
455
+ "VLLM_B12X_MOE_FP4_FORCE_A16",
456
+ force_a16,
457
+ )
458
+ config = make_dummy_moe_config(hidden_dim=256, intermediate_size=64)
459
+ config.moe_backend = "b12x"
460
+
461
+ backend, experts_cls = select_deepseek_v4_mxfp4_moe_backend(config)
462
+
463
+ assert backend == expected_backend
464
+ assert experts_cls is B12xExperts
465
+
466
+
467
+def test_compressed_tensors_mxfp4_preserves_checkpoint_packing(
468
+ monkeypatch: pytest.MonkeyPatch,
469
+) -> None:
470
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
471
+ compressed_tensors_moe_w4a4_mxfp4 as ct_mxfp4,
472
+ )
473
+
474
+ monkeypatch.setattr(
475
+ ct_mxfp4.CutlassExpertsMxfp4,
476
+ "_supports_current_device",
477
+ lambda: False,
478
+ )
479
+ monkeypatch.setattr(
480
+ ct_mxfp4,
481
+ "select_mxfp4_moe_backend",
482
+ lambda moe: (Mxfp4MoeBackend.B12X_MXFP4_MXFP8, B12xExperts),
483
+ )
484
+ monkeypatch.setattr(
485
+ ct_mxfp4,
486
+ "prepare_moe_fp4_layer_for_marlin",
487
+ lambda layer: pytest.fail("b12x must not use Marlin packing"),
488
+ )
489
+ moe_config = SimpleNamespace(w13_num_shards=2, moe_backend="b12x")
490
+ method = ct_mxfp4.CompressedTensorsW4A4Mxfp4MoEMethod(moe_config)
491
+ processed_layers: list[torch.nn.Module] = []
492
+ fake_experts = SimpleNamespace(
493
+ process_weights_after_loading=processed_layers.append
494
+ )
495
+ kernel = SimpleNamespace(fused_experts=fake_experts)
496
+ monkeypatch.setattr(method, "get_fused_moe_quant_config", lambda _: object())
497
+ monkeypatch.setattr(ct_mxfp4, "make_mxfp4_moe_kernel", lambda **_: kernel)
498
+ layer = torch.nn.Module()
499
+ layer._expert_routing_tables = lambda: ()
500
+ method.create_weights(
501
+ layer,
502
+ num_experts=2,
503
+ hidden_size=64,
504
+ intermediate_size_per_partition=32,
505
+ params_dtype=torch.bfloat16,
506
+ )
507
+ w13_packed_data = layer.w13_weight_packed.data
508
+ w2_packed_data = layer.w2_weight_packed.data
509
+
510
+ method.process_weights_after_loading(layer)
511
+
512
+ assert layer.w13_weight.data.data_ptr() == w13_packed_data.data_ptr()
513
+ assert layer.w2_weight.data.data_ptr() == w2_packed_data.data_ptr()
514
+ assert processed_layers == [layer]
515
+
516
+
517
+def test_b12x_mxfp4_falls_back_to_a16(
518
+ monkeypatch: pytest.MonkeyPatch,
519
+) -> None:
520
+ monkeypatch.setattr(B12xExperts, "_supports_current_device", lambda: True)
521
+ monkeypatch.setattr(mxfp4_oracle, "_user_moe_activation_override", lambda: None)
522
+ monkeypatch.setattr(
523
+ mxfp4_oracle.envs,
524
+ "VLLM_B12X_MOE_FP4_FORCE_A16",
525
+ False,
526
+ )
527
+ config = make_dummy_moe_config(hidden_dim=128, intermediate_size=64)
528
+ config.moe_backend = "b12x"
529
+
530
+ backend, experts_cls = select_mxfp4_moe_backend(config)
531
+
532
+ assert backend == Mxfp4MoeBackend.B12X_MXFP4_BF16
533
+ assert experts_cls is B12xExperts
534
+
535
+
536
+@pytest.mark.parametrize(
537
+ "activation_key,force_a16,expected_activation_key",
538
+ [
539
+ (kNvfp4Dynamic, False, kNvfp4Dynamic),
540
+ (kMxfp8Dynamic, False, kMxfp8Dynamic),
541
+ (kNvfp4Dynamic, True, None),
542
+ ],
543
+)
544
+def test_explicit_b12x_nvfp4_selection(
545
+ monkeypatch: pytest.MonkeyPatch,
546
+ activation_key,
547
+ force_a16: bool,
548
+ expected_activation_key,
549
+) -> None:
550
+ selected_activation_keys = []
551
+
552
+ def is_supported_config(cls, config, weight_key, activation_key, activation_format):
553
+ selected_activation_keys.append(activation_key)
554
+ return True, None
555
+
556
+ monkeypatch.setattr(B12xExperts, "is_supported_config", is_supported_config)
557
+ monkeypatch.setattr(
558
+ nvfp4_oracle.envs,
559
+ "VLLM_B12X_MOE_FP4_FORCE_A16",
560
+ force_a16,
561
+ )
562
+ config = make_dummy_moe_config(hidden_dim=128, intermediate_size=64)
563
+ config.moe_backend = "b12x"
564
+
565
+ backend, experts_cls = select_nvfp4_moe_backend(
566
+ config,
567
+ weight_key=kNvfp4Static,
568
+ activation_key=activation_key,
569
+ )
570
+
571
+ assert backend == NvFp4MoeBackend.B12X
572
+ assert experts_cls is B12xExperts
573
+ assert selected_activation_keys == [expected_activation_key]
574
+
575
+
576
+@pytest.mark.parametrize(
577
+ "force_a16,expected_quant_dtype", [(False, "nvfp4"), (True, None)]
578
+)
579
+def test_b12x_nvfp4_force_a16_updates_quant_config(
580
+ monkeypatch: pytest.MonkeyPatch,
581
+ force_a16: bool,
582
+ expected_quant_dtype,
583
+) -> None:
584
+ monkeypatch.setattr(
585
+ nvfp4_oracle.envs,
586
+ "VLLM_B12X_MOE_FP4_FORCE_A16",
587
+ force_a16,
588
+ )
589
+ scale = torch.ones(1)
590
+
591
+ quant_config = nvfp4_oracle.make_nvfp4_moe_quant_config(
592
+ backend=NvFp4MoeBackend.B12X,
593
+ w13_scale=scale,
594
+ w2_scale=scale,
595
+ w13_scale_2=scale,
596
+ w2_scale_2=scale,
597
+ a13_scale=scale,
598
+ a2_scale=scale,
599
+ )
600
+
601
+ assert quant_config.quant_dtype == expected_quant_dtype
602
+
603
+
604
+def test_b12x_nvfp4_force_a16_updates_weight_preparation(
605
+ monkeypatch: pytest.MonkeyPatch,
606
+) -> None:
607
+ monkeypatch.setattr(
608
+ nvfp4_oracle.envs,
609
+ "VLLM_B12X_MOE_FP4_FORCE_A16",
610
+ True,
611
+ )
612
+ reorder_w13 = None
613
+
614
+ def prepare_for_b12x(**kwargs):
615
+ nonlocal reorder_w13
616
+ reorder_w13 = kwargs["reorder_w13"]
617
+ return (
618
+ kwargs["w13"],
619
+ kwargs["w13_scale"],
620
+ kwargs["w13_scale_2"],
621
+ kwargs["a13_scale"],
622
+ kwargs["w2"],
623
+ kwargs["w2_scale"],
624
+ kwargs["w2_scale_2"],
625
+ kwargs["a2_scale"],
626
+ )
627
+
628
+ monkeypatch.setattr(
629
+ nvfp4_oracle,
630
+ "prepare_nvfp4_moe_layer_for_b12x",
631
+ prepare_for_b12x,
632
+ )
633
+ tensor = torch.ones(1)
634
+
635
+ nvfp4_oracle.convert_to_nvfp4_moe_kernel_format(
636
+ nvfp4_backend=NvFp4MoeBackend.B12X,
637
+ layer=SimpleNamespace(),
638
+ w13=tensor,
639
+ w13_scale=tensor,
640
+ w13_scale_2=tensor,
641
+ a13_scale=tensor,
642
+ w2=tensor,
643
+ w2_scale=tensor,
644
+ w2_scale_2=tensor,
645
+ a2_scale=tensor,
646
+ is_act_and_mul=True,
647
+ )
648
+
649
+ assert reorder_w13 is True
650
+
651
+
652
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
653
+def test_b12x_nvfp4_preparation_pads_each_gated_half() -> None:
654
+ device = torch.device("cuda")
655
+ num_experts, hidden_size, intermediate_size = 2, 64, 48
656
+ w13 = torch.ones(
657
+ num_experts,
658
+ 2 * intermediate_size,
659
+ hidden_size // 2,
660
+ dtype=torch.uint8,
661
+ device=device,
662
+ )
663
+ w13_scale = torch.ones(
664
+ num_experts,
665
+ 2 * intermediate_size,
666
+ hidden_size // 16,
667
+ dtype=torch.float8_e4m3fn,
668
+ device=device,
669
+ )
670
+ w2 = torch.ones(
671
+ num_experts,
672
+ hidden_size,
673
+ intermediate_size // 2,
674
+ dtype=torch.uint8,
675
+ device=device,
676
+ )
677
+ w2_scale = torch.ones(
678
+ num_experts,
679
+ hidden_size,
680
+ intermediate_size // 16,
681
+ dtype=torch.float8_e4m3fn,
682
+ device=device,
683
+ )
684
+ global_scale = torch.ones(num_experts, device=device)
685
+ input_scale = torch.tensor([[1.0, 2.0], [3.0, 1.0]], device=device)
686
+
687
+ prepared = prepare_nvfp4_moe_layer_for_b12x(
688
+ w13,
689
+ w13_scale,
690
+ global_scale,
691
+ input_scale,
692
+ w2,
693
+ w2_scale,
694
+ global_scale,
695
+ input_scale,
696
+ is_act_and_mul=True,
697
+ )
698
+
699
+ prepared_w13, prepared_w13_scale, _, prepared_a13 = prepared[:4]
700
+ prepared_w2, prepared_w2_scale, _, prepared_a2 = prepared[4:]
701
+ assert prepared_w13.shape == (num_experts, 128, hidden_size // 2)
702
+ assert prepared_w13_scale.shape == (num_experts, 128, hidden_size // 16)
703
+ assert prepared_w2.shape == (num_experts, hidden_size, 32)
704
+ assert prepared_w2_scale.shape == (num_experts, 128, 4)
705
+ torch.testing.assert_close(prepared_a13, torch.tensor([2.0, 3.0], device=device))
706
+ torch.testing.assert_close(prepared_a2, torch.tensor([2.0, 3.0], device=device))
707
+
708
+
709
+def test_b12x_moe_uses_minimax_swiglu_parameters() -> None:
710
+ config = make_dummy_moe_config(
711
+ hidden_dim=128,
712
+ intermediate_size=64,
713
+ activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
714
+ )
715
+ config.swiglu_limit = 7.0
716
+ config.swiglu_alpha = 1.702
717
+ config.swiglu_beta = 1.0
718
+ experts = B12xExperts(config, _quant_config("mxfp4", None))
719
+
720
+ assert experts._swiglu_params(config.activation) == (7.0, 1.702, 1.0)
721
+
722
+
723
+def test_b12x_moe_warmup_runs_each_planner_regime_once(
724
+ monkeypatch: pytest.MonkeyPatch,
725
+) -> None:
726
+ experts = B12xExperts(
727
+ make_dummy_moe_config(
728
+ num_experts=4,
729
+ experts_per_token=2,
730
+ hidden_dim=128,
731
+ intermediate_size=64,
732
+ ),
733
+ _quant_config("mxfp4", None),
734
+ )
735
+ prepared = SimpleNamespace(
736
+ num_experts=4,
737
+ hidden_size=128,
738
+ intermediate_size=64,
739
+ w1_fp4=torch.empty(0),
740
+ )
741
+ layer = SimpleNamespace(
742
+ activation=MoEActivation.SILU,
743
+ apply_router_weight_on_input=False,
744
+ )
745
+ planned_tokens = []
746
+ launched_tokens = []
747
+
748
+ with pytest.raises(RuntimeError, match="process_weights_after_loading"):
749
+ experts.warmup_launches(layer, token_counts=(1,))
750
+ experts._prepared_experts = prepared
751
+
752
+ def fake_execution_plan(**kwargs):
753
+ tokens = kwargs["tokens"]
754
+ if tokens <= 2:
755
+ signature = ("micro", "decode")
756
+ elif tokens <= 4:
757
+ signature = ("dynamic", "small")
758
+ else:
759
+ signature = ("dynamic", "large")
760
+ return SimpleNamespace(
761
+ implementation=signature[0],
762
+ execution=signature[1],
763
+ )
764
+
765
+ def fake_plan(**kwargs):
766
+ planned_tokens.append(kwargs["tokens"])
767
+ return SimpleNamespace(
768
+ scratch_specs=lambda: [SimpleNamespace(dtype=torch.uint8, shape=(64,))]
769
+ )
770
+
771
+ def fake_run(**kwargs):
772
+ launched_tokens.append(kwargs["hidden_states"].shape[0])
773
+
774
+ monkeypatch.setattr(b12x, "_b12x_moe_execution_plan", fake_execution_plan)
775
+ monkeypatch.setattr(b12x, "_run_b12x_moe_plan", fake_run)
776
+ monkeypatch.setattr(experts, "_plan", fake_plan)
777
+
778
+ warmed = experts.warmup_launches(layer, token_counts=(1, 2, 3, 4, 8))
779
+
780
+ assert warmed == 3
781
+ assert planned_tokens == [1, 3, 8]
782
+ assert launched_tokens == planned_tokens
783
+
784
+
785
+def test_b12x_moe_reload_reprepares_current_parameters(
786
+ monkeypatch: pytest.MonkeyPatch,
787
+) -> None:
788
+ plan = SimpleNamespace(
789
+ discards_source_parameters=False,
790
+ quant_modes=("w4a16",),
791
+ io_dtype="bfloat16",
792
+ activation="silu",
793
+ )
794
+ prepared_inputs: list[
795
+ tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
796
+ ] = []
797
+
798
+ def prepare_weights(**kwargs):
799
+ prepared_inputs.append(
800
+ (
801
+ kwargs["w1_fp4"],
802
+ kwargs["w2_fp4"],
803
+ kwargs["w1_blockscale"],
804
+ kwargs["w2_blockscale"],
805
+ )
806
+ )
807
+ return SimpleNamespace(plan=plan)
808
+
809
+ extension = SimpleNamespace(
810
+ plan_weights=lambda **_: plan, prepare_weights=prepare_weights
811
+ )
812
+ monkeypatch.setattr(b12x, "_require_b12x_fused_moe", lambda: extension)
813
+ experts = B12xExperts(
814
+ make_dummy_moe_config(num_experts=2, hidden_dim=4, intermediate_size=8),
815
+ _quant_config("mxfp4", None),
816
+ )
817
+ layer = SimpleNamespace(
818
+ activation=MoEActivation.SILU,
819
+ apply_router_weight_on_input=False,
820
+ w13_weight=torch.full((2, 8, 2), 1, dtype=torch.uint8),
821
+ w2_weight=torch.full((2, 4, 4), 1, dtype=torch.uint8),
822
+ w13_weight_scale=torch.full((2, 8, 1), 1, dtype=torch.uint8),
823
+ w2_weight_scale=torch.full((2, 4, 1), 1, dtype=torch.uint8),
824
+ )
825
+
826
+ experts.process_weights_after_loading(layer)
827
+ layer.w13_weight = torch.full_like(layer.w13_weight, 2)
828
+ layer.w2_weight = torch.full_like(layer.w2_weight, 2)
829
+ layer.w13_weight_scale = torch.full_like(layer.w13_weight_scale, 3)
830
+ layer.w2_weight_scale = torch.full_like(layer.w2_weight_scale, 3)
831
+ experts.process_weights_after_loading(layer)
832
+
833
+ assert len(prepared_inputs) == 2
834
+ assert prepared_inputs[-1][0] is layer.w13_weight
835
+ assert prepared_inputs[-1][1] is layer.w2_weight
836
+ assert prepared_inputs[-1][2] is layer.w13_weight_scale
837
+ assert prepared_inputs[-1][3] is layer.w2_weight_scale
838
+
839
+
840
+def test_b12x_source_release_preserves_prepared_storage_owner() -> None:
841
+ layer = torch.nn.Module()
842
+ for name, shape in (
843
+ ("w13_weight", (4, 32, 16)),
844
+ ("w2_weight", (4, 64, 8)),
845
+ ("w13_weight_scale", (4, 32, 2)),
846
+ ("w2_weight_scale", (4, 64, 1)),
847
+ ):
848
+ layer.register_parameter(
849
+ name,
850
+ torch.nn.Parameter(
851
+ torch.empty(shape, dtype=torch.uint8),
852
+ requires_grad=False,
853
+ ),
854
+ )
855
+ experts = B12xExperts(
856
+ make_dummy_moe_config(hidden_dim=128, intermediate_size=64),
857
+ _quant_config("mxfp4", None),
858
+ )
859
+ owner = SimpleNamespace(
860
+ w1_fp4=layer.w13_weight,
861
+ w2_fp4=layer.w2_weight,
862
+ w1_blockscale=layer.w13_weight_scale,
863
+ w2_blockscale=layer.w2_weight_scale,
864
+ )
865
+ experts._prepared_experts = owner
866
+ owner_tensors = (
867
+ owner.w1_fp4,
868
+ owner.w2_fp4,
869
+ owner.w1_blockscale,
870
+ owner.w2_blockscale,
871
+ )
872
+ owner_ptrs = tuple(tensor.untyped_storage().data_ptr() for tensor in owner_tensors)
873
+
874
+ experts._release_source_parameters(layer)
875
+ experts._release_source_parameters(layer)
876
+
877
+ assert layer.w13_weight.numel() == 0
878
+ assert layer.w2_weight.numel() == 0
879
+ assert layer.w13_weight_scale.numel() == 0
880
+ assert layer.w2_weight_scale.numel() == 0
881
+ assert (
882
+ tuple(tensor.untyped_storage().data_ptr() for tensor in owner_tensors)
883
+ == owner_ptrs
884
+ )
885
+
886
+
887
+def test_b12x_moe_rejects_router_weight_on_input_for_w4a8() -> None:
888
+ experts = B12xExperts(
889
+ make_dummy_moe_config(hidden_dim=256, intermediate_size=64),
890
+ _quant_config("mxfp4", "mxfp8"),
891
+ )
892
+ layer = SimpleNamespace(
893
+ activation=MoEActivation.SILU,
894
+ apply_router_weight_on_input=True,
895
+ )
896
+
897
+ with pytest.raises(
898
+ ValueError,
899
+ match="apply_router_weight_on_input only with W4A16",
900
+ ):
901
+ experts.process_weights_after_loading(layer)
902
+
903
+
904
+def test_b12x_moe_workspace_uses_prepared_router_weight_contract(
905
+ monkeypatch: pytest.MonkeyPatch,
906
+) -> None:
907
+ experts = B12xExperts(
908
+ make_dummy_moe_config(hidden_dim=128, intermediate_size=64),
909
+ _quant_config("mxfp4", None),
910
+ )
911
+ prepared = SimpleNamespace(
912
+ plan=SimpleNamespace(discards_source_parameters=False),
913
+ )
914
+ layer = SimpleNamespace(
915
+ activation=MoEActivation.SILU,
916
+ apply_router_weight_on_input=True,
917
+ w13_weight=torch.empty(0),
918
+ w2_weight=torch.empty(0),
919
+ w13_weight_scale=torch.empty(0),
920
+ w2_weight_scale=torch.empty(0),
921
+ )
922
+ monkeypatch.setattr(experts, "_prepare_experts", lambda **kwargs: prepared)
923
+ planned = []
924
+
925
+ def fake_plan(**kwargs):
926
+ planned.append(kwargs)
927
+ return SimpleNamespace(
928
+ scratch_specs=lambda: [SimpleNamespace(dtype=torch.uint8, shape=(64,))]
929
+ )
930
+
931
+ monkeypatch.setattr(experts, "_plan", fake_plan)
932
+
933
+ experts.process_weights_after_loading(layer)
934
+ assert layer.b12x_warmup_provider is experts
935
+ experts.workspace_shapes(
936
+ 8,
937
+ 128,
938
+ 128,
939
+ 2,
940
+ 4,
941
+ 4,
942
+ None,
943
+ MoEActivation.SILU,
944
+ )
945
+
946
+ assert planned == [
947
+ {
948
+ "tokens": 8,
949
+ "topk": 2,
950
+ "activation": MoEActivation.SILU,
951
+ "apply_router_weight_on_input": True,
952
+ }
953
+ ]
954
+
955
+
956
+@dataclass
957
+class _B12xMoeCase:
958
+ hidden_states: torch.Tensor
959
+ score: torch.Tensor
960
+ w1: torch.Tensor
961
+ w2: torch.Tensor
962
+ w1_ref: torch.Tensor
963
+ w2_ref: torch.Tensor
964
+ quant_config: FusedMoEQuantConfig
965
+ activation: MoEActivation
966
+ activation_dtype: str | None
967
+ topk: int = 2
968
+
969
+
970
+def _make_b12x_moe_case(
971
+ weight_dtype: str,
972
+ activation_dtype: str | None,
973
+ *,
974
+ activation: MoEActivation = MoEActivation.SILU,
975
+ tokens: int = 16,
976
+ seed: int = 19,
977
+) -> _B12xMoeCase:
978
+ set_random_seed(seed)
979
+ num_experts, hidden_size, intermediate_size = 4, 512, 128
980
+ dtype = torch.bfloat16
981
+ hidden_states = torch.randn((tokens, hidden_size), device="cuda", dtype=dtype) / 10
982
+ w1_rows = 2 * intermediate_size if activation.is_gated else intermediate_size
983
+ w1 = (
984
+ torch.randn(
985
+ (num_experts, w1_rows, hidden_size),
986
+ device="cuda",
987
+ dtype=dtype,
988
+ )
989
+ / 15
990
+ )
991
+ w2 = (
992
+ torch.randn(
993
+ (num_experts, hidden_size, intermediate_size),
994
+ device="cuda",
995
+ dtype=dtype,
996
+ )
997
+ / 15
998
+ )
999
+
1000
+ if weight_dtype == "mxfp4":
1001
+ w1_q, w1_scale = mxfp4_quantize(w1)
1002
+ w2_q, w2_scale = mxfp4_quantize(w2)
1003
+ w1_ref = torch.stack(
1004
+ [dq_mxfp4_torch(w1_q[e], w1_scale[e], dtype) for e in range(num_experts)]
1005
+ )
1006
+ w2_ref = torch.stack(
1007
+ [dq_mxfp4_torch(w2_q[e], w2_scale[e], dtype) for e in range(num_experts)]
1008
+ )
1009
+ if activation_dtype is None:
1010
+ quant_config = mxfp4_w4a16_moe_quant_config(
1011
+ w1_scale=w1_scale,
1012
+ w2_scale=w2_scale,
1013
+ )
1014
+ else:
1015
+ quant_config = FusedMoEQuantConfig.make(
1016
+ quant_dtype=activation_dtype,
1017
+ weight_dtype=weight_dtype,
1018
+ w1_scale=w1_scale,
1019
+ w2_scale=w2_scale,
1020
+ )
1021
+ elif weight_dtype == "nvfp4":
1022
+ w1_q, w1_scale, w1_global_scale = _quantize_nvfp4_linear(w1)
1023
+ w2_q, w2_scale, w2_global_scale = _quantize_nvfp4_linear(w2)
1024
+ w1_ref = torch.stack(
1025
+ [
1026
+ _dequantize_nvfp4_linear(
1027
+ w1_q[e], w1_scale[e], w1_global_scale[e], dtype
1028
+ )
1029
+ for e in range(num_experts)
1030
+ ]
1031
+ )
1032
+ w2_ref = torch.stack(
1033
+ [
1034
+ _dequantize_nvfp4_linear(
1035
+ w2_q[e], w2_scale[e], w2_global_scale[e], dtype
1036
+ )
1037
+ for e in range(num_experts)
1038
+ ]
1039
+ )
1040
+ input_scale = torch.full(
1041
+ (num_experts,),
1042
+ 1.0 if activation_dtype is None else 1.0 / 1024.0,
1043
+ device="cuda",
1044
+ dtype=torch.float32,
1045
+ )
1046
+ prepared = prepare_nvfp4_moe_layer_for_b12x(
1047
+ w1_q,
1048
+ w1_scale,
1049
+ 1.0 / w1_global_scale,
1050
+ input_scale,
1051
+ w2_q,
1052
+ w2_scale,
1053
+ 1.0 / w2_global_scale,
1054
+ input_scale,
1055
+ is_act_and_mul=activation.is_gated,
1056
+ reorder_w13=activation_dtype is None and activation.is_gated,
1057
+ )
1058
+ w1_q, w1_scale, w1_alpha, a1_scale = prepared[:4]
1059
+ w2_q, w2_scale, w2_alpha, a2_scale = prepared[4:]
1060
+ if activation_dtype is None:
1061
+ quant_config = nvfp4_w4a16_moe_quant_config(
1062
+ g1_alphas=w1_alpha,
1063
+ g2_alphas=w2_alpha,
1064
+ w1_scale=w1_scale,
1065
+ w2_scale=w2_scale,
1066
+ )
1067
+ else:
1068
+ quant_config = FusedMoEQuantConfig.make(
1069
+ quant_dtype=activation_dtype,
1070
+ weight_dtype=weight_dtype,
1071
+ w1_scale=w1_scale,
1072
+ w2_scale=w2_scale,
1073
+ g1_alphas=w1_alpha,
1074
+ g2_alphas=w2_alpha,
1075
+ a1_gscale=1.0 / a1_scale,
1076
+ a2_gscale=1.0 / a2_scale,
1077
+ )
1078
+ else:
1079
+ raise ValueError(f"unsupported test weight dtype {weight_dtype}")
1080
+
1081
+ return _B12xMoeCase(
1082
+ hidden_states=hidden_states,
1083
+ score=torch.randn((tokens, num_experts), device="cuda", dtype=dtype),
1084
+ w1=w1_q,
1085
+ w2=w2_q,
1086
+ w1_ref=w1_ref,
1087
+ w2_ref=w2_ref,
1088
+ quant_config=quant_config,
1089
+ activation=activation,
1090
+ activation_dtype=activation_dtype,
1091
+ )
1092
+
1093
+
1094
+@pytest.mark.skipif(not _has_b12x_moe(), reason="requires b12x MoE on SM120")
1095
+@pytest.mark.parametrize(
1096
+ "weight_dtype,activation_dtype,activation",
1097
+ [
1098
+ pytest.param("mxfp4", "mxfp8", MoEActivation.SILU, id="mxfp4-mxfp8"),
1099
+ pytest.param("mxfp4", None, MoEActivation.SILU, id="mxfp4-bf16"),
1100
+ pytest.param("nvfp4", "nvfp4", MoEActivation.SILU, id="nvfp4-nvfp4"),
1101
+ pytest.param("nvfp4", "mxfp8", MoEActivation.SILU, id="nvfp4-mxfp8"),
1102
+ pytest.param("nvfp4", None, MoEActivation.SILU, id="nvfp4-bf16-silu"),
1103
+ pytest.param(
1104
+ "nvfp4",
1105
+ None,
1106
+ MoEActivation.RELU2_NO_MUL,
1107
+ id="nvfp4-bf16-relu2",
1108
+ ),
1109
+ ],
1110
+)
1111
+@torch.inference_mode()
1112
+def test_b12x_moe_matches_torch(
1113
+ weight_dtype: str,
1114
+ activation_dtype: str | None,
1115
+ activation: MoEActivation,
1116
+ workspace_init,
1117
+) -> None:
1118
+ with set_current_vllm_config(
1119
+ VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
1120
+ ):
1121
+ case = _make_b12x_moe_case(
1122
+ weight_dtype,
1123
+ activation_dtype,
1124
+ activation=activation,
1125
+ )
1126
+ reference = torch_moe(
1127
+ case.hidden_states,
1128
+ case.w1_ref,
1129
+ case.w2_ref,
1130
+ case.score,
1131
+ case.topk,
1132
+ activation=case.activation,
1133
+ )
1134
+ if activation_dtype == "nvfp4":
1135
+ topk_weights, topk_ids, _ = fused_topk(
1136
+ case.hidden_states,
1137
+ case.score,
1138
+ case.topk,
1139
+ renormalize=False,
1140
+ )
1141
+ reference = _nvfp4_activation_reference(
1142
+ case.hidden_states,
1143
+ case.w1_ref,
1144
+ case.w2_ref,
1145
+ topk_weights,
1146
+ topk_ids,
1147
+ case.quant_config.a1_gscale,
1148
+ case.quant_config.a2_gscale,
1149
+ )
1150
+
1151
+ checks_zero_canonicalization = (
1152
+ weight_dtype == "nvfp4" and activation_dtype is None
1153
+ )
1154
+ if checks_zero_canonicalization:
1155
+ assert _count_fp4_negative_zeros(case.w1) > 0
1156
+ assert _count_fp4_negative_zeros(case.w2) > 0
1157
+ output = _run_b12x_moe(
1158
+ case.hidden_states,
1159
+ case.w1,
1160
+ case.w2,
1161
+ case.score,
1162
+ case.topk,
1163
+ case.activation,
1164
+ case.quant_config,
1165
+ )
1166
+ if checks_zero_canonicalization:
1167
+ assert _count_fp4_negative_zeros(case.w1) == 0
1168
+ assert _count_fp4_negative_zeros(case.w2) == 0
1169
+
1170
+ torch.testing.assert_close(output, reference, atol=2e-1, rtol=2e-1)
1171
+ cosine = torch.nn.functional.cosine_similarity(
1172
+ output.flatten().float(),
1173
+ reference.flatten().float(),
1174
+ dim=0,
1175
+ )
1176
+ assert cosine > 0.99
1177
+
1178
+
1179
+@pytest.mark.skipif(not _has_b12x_moe(), reason="requires b12x MoE on SM120")
1180
+@pytest.mark.parametrize(
1181
+ "weight_dtype,activation_dtype",
1182
+ [
1183
+ pytest.param("mxfp4", "mxfp8", id="w4a8"),
1184
+ pytest.param("nvfp4", None, id="w4a16"),
1185
+ ],
1186
+)
1187
+@torch.inference_mode()
1188
+def test_b12x_moe_cuda_graph_replay(
1189
+ weight_dtype: str,
1190
+ activation_dtype: str | None,
1191
+ workspace_init,
1192
+) -> None:
1193
+ from vllm.v1.worker.workspace import lock_workspace
1194
+
1195
+ with set_current_vllm_config(
1196
+ VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
1197
+ ):
1198
+ case = _make_b12x_moe_case(
1199
+ weight_dtype,
1200
+ activation_dtype,
1201
+ tokens=4,
1202
+ seed=23,
1203
+ )
1204
+ kernel = _make_b12x_moe_kernel(
1205
+ case.hidden_states,
1206
+ case.w1,
1207
+ case.w2,
1208
+ case.topk,
1209
+ case.activation,
1210
+ case.quant_config,
1211
+ )
1212
+ topk_weights, topk_ids, _ = fused_topk(
1213
+ case.hidden_states,
1214
+ case.score,
1215
+ case.topk,
1216
+ renormalize=False,
1217
+ )
1218
+ assert topk_weights.dtype == torch.float32 and topk_weights.is_contiguous()
1219
+ assert topk_ids.dtype == torch.int32 and topk_ids.is_contiguous()
1220
+
1221
+ def apply() -> torch.Tensor:
1222
+ return kernel.apply(
1223
+ hidden_states=case.hidden_states,
1224
+ w1=case.w1,
1225
+ w2=case.w2,
1226
+ topk_weights=topk_weights,
1227
+ topk_ids=topk_ids,
1228
+ activation=case.activation,
1229
+ global_num_experts=case.w1.shape[0],
1230
+ expert_map=None,
1231
+ apply_router_weight_on_input=False,
1232
+ )
1233
+
1234
+ expected = apply().clone()
1235
+ lock_workspace()
1236
+ graph = torch.cuda.CUDAGraph()
1237
+ stream = torch.cuda.Stream()
1238
+ with torch.cuda.graph(graph, stream=stream):
1239
+ actual = apply()
1240
+ graph.replay()
1241
+ torch.accelerator.synchronize()
1242
+
1243
+ assert torch.isfinite(expected).all()
1244
+ assert torch.isfinite(actual).all()
1245
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
tests/kernels/quantization/nvfp4_utils.py
+14
-3
@@ -43,6 +43,7 @@ def dequantize_nvfp4_to_dtype(
43
device,
44
block_size=16,
45
is_sf_128x4_layout=True,
46
+ is_sf_linear_layout=False,
47
):
48
"""Dequantize the fp4 tensor back to high precision."""
49
# Two fp4 values are packed into one uint8.
@@ -52,7 +53,9 @@ def dequantize_nvfp4_to_dtype(
53
tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
54
tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
55
tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
55
- if is_sf_128x4_layout:
56
+ if is_sf_linear_layout:
57
+ tensor_sf = tensor_sf[:m, : k // block_size]
58
+ elif is_sf_128x4_layout:
59
tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
60
else:
61
tensor_sf = convert_swizzled_8x4_layout_to_linear(tensor_sf, m, k, block_size)
@@ -146,7 +149,15 @@ def get_nvfp4_global_scale(a: torch.Tensor):
149
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
150
151
149
-def quant_nvfp4_tensor(a: torch.Tensor):
152
+def quant_nvfp4_tensor(
153
+ a: torch.Tensor,
154
+ *,
155
+ is_sf_swizzled_layout: bool = True,
156
+):
157
a_global_scale = get_nvfp4_global_scale(a)
151
- a_quant, a_block_scale = scaled_fp4_quant(a, a_global_scale)
158
+ a_quant, a_block_scale = scaled_fp4_quant(
159
+ a,
160
+ a_global_scale,
161
+ is_sf_swizzled_layout=is_sf_swizzled_layout,
162
+ )
163
return a_quant, a_block_scale, a_global_scale
tests/model_executor/test_b12x_warmup.py
+9
-3
@@ -181,8 +181,14 @@ def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None:
181
182
worker = SimpleNamespace(
183
get_model=lambda: SimpleNamespace(modules=modules),
184
- scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
184
+ scheduler_config=SimpleNamespace(
185
+ max_num_batched_tokens=8,
186
+ max_num_scheduled_tokens=32,
187
+ ),
188
model_config=SimpleNamespace(dtype=torch.float32),
189
+ vllm_config=SimpleNamespace(
190
+ compilation_config=SimpleNamespace(compile_sizes=[4, 16])
191
+ ),
192
)
193
platform = SimpleNamespace(
194
is_cuda=lambda: True,
@@ -200,7 +206,7 @@ def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None:
206
207
assert scans == 1
208
assert calls == [
203
- ("first", (1, 2, 8), torch.bfloat16),
204
- ("second", (1, 2, 8), torch.bfloat16),
209
+ ("first", (1, 2, 4, 8, 16, 32), torch.bfloat16),
210
+ ("second", (1, 2, 4, 8, 16, 32), torch.bfloat16),
211
]
212
assert synchronized == [True]
tests/quantization/test_auto_round.py
+30
-5
@@ -663,12 +663,28 @@ def test_inc_mxfp4_linear_method_registers_and_processes_weights(
663
assert captured["processed_layer"] is layer
664
665
666
-def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(
666
+@pytest.mark.parametrize(
667
+ ("moe_backend", "is_xpu", "mxfp4_backend"),
668
+ [
669
+ ("auto", True, Mxfp4MoeBackend.XPU),
670
+ ("b12x", False, Mxfp4MoeBackend.B12X_MXFP4_MXFP8),
671
+ ],
672
+)
673
+def test_inc_mxfp4_moe_method_preserves_checkpoint_packing(
674
monkeypatch,
675
+ moe_backend: str,
676
+ is_xpu: bool,
677
+ mxfp4_backend: Mxfp4MoeBackend,
678
) -> None:
679
captured = {}
680
expected_quant_config = object()
671
- expected_kernel = object()
681
+ expected_kernel = SimpleNamespace(
682
+ fused_experts=SimpleNamespace(
683
+ process_weights_after_loading=lambda layer: captured.update(
684
+ {"processed_layer": layer}
685
+ )
686
+ )
687
+ )
688
expected_experts_cls = object()
689
690
monkeypatch.setattr(
@@ -676,11 +692,16 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(
692
"CutlassExpertsMxfp4._supports_current_device",
693
lambda: False,
694
)
679
- monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
695
+ monkeypatch.setattr(current_platform, "is_xpu", lambda: is_xpu)
696
monkeypatch.setattr(
697
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
698
"select_mxfp4_moe_backend",
683
- lambda moe: (Mxfp4MoeBackend.XPU, expected_experts_cls),
699
+ lambda moe: (mxfp4_backend, expected_experts_cls),
700
+ )
701
+ monkeypatch.setattr(
702
+ "vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
703
+ "prepare_moe_fp4_layer_for_marlin",
704
+ lambda layer: pytest.fail("packed backends must not use Marlin packing"),
705
)
706
monkeypatch.setattr(
707
"vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp4_moe."
@@ -698,7 +719,10 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(
719
INCMxfp4MoEMethod,
720
)
721
701
- expected_moe_config = SimpleNamespace(w13_num_shards=2)
722
+ expected_moe_config = SimpleNamespace(
723
+ w13_num_shards=2,
724
+ moe_backend=moe_backend,
725
+ )
726
method = INCMxfp4MoEMethod(moe=cast(Any, expected_moe_config))
727
layer = torch.nn.Module()
728
layer._expert_routing_tables = lambda: "routing-tables"
@@ -731,6 +755,7 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(
755
assert captured["kernel_kwargs"]["moe_config"] is expected_moe_config
756
assert captured["kernel_kwargs"]["experts_cls"] is expected_experts_cls
757
assert captured["kernel_kwargs"]["routing_tables"] == "routing-tables"
758
+ assert captured["processed_layer"] is layer
759
assert method.moe_kernel is expected_kernel
760
761
vllm/config/kernel.py
+2
@@ -130,6 +130,7 @@ MoEBackend = Literal[
130
"flashinfer_cutlass",
131
"flashinfer_cutedsl",
132
"flashinfer_b12x",
133
+ "b12x",
134
"marlin",
135
"humming",
136
"triton_unfused",
@@ -204,6 +205,7 @@ class KernelConfig:
205
- "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
206
- "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x
207
(RTX Pro 6000 / DGX Spark)
208
+ - "b12x": Use b12x FP4 MoE kernels on SM12x
209
- "marlin": Use Marlin kernels (weight-only quantization)
210
- "humming": Use Humming Mixed Precision kernels
211
- "triton_unfused": Use Triton unfused MoE kernels
vllm/envs.py
+5
@@ -190,6 +190,7 @@ if TYPE_CHECKING:
190
VLLM_HUMMING_INPUT_QUANT_CONFIG: dict[str, Any] | None = None
191
VLLM_HUMMING_USE_F16_ACCUM: bool = False
192
VLLM_HUMMING_MOE_GEMM_TYPE: Literal["indexed", "grouped", "auto"] | None = None
193
+ VLLM_B12X_MOE_FP4_FORCE_A16: bool = False
194
VLLM_DEEPEPLL_NVFP4_DISPATCH: bool = False
195
VLLM_V1_USE_OUTLINES_CACHE: bool = False
196
VLLM_TPU_USING_PATHWAYS: bool = False
@@ -1621,6 +1622,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1622
"VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(
1623
int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "1"))
1624
),
1625
+ # Force b12x FP4 MoE to use BF16 activations.
1626
+ "VLLM_B12X_MOE_FP4_FORCE_A16": lambda: bool(
1627
+ int(os.getenv("VLLM_B12X_MOE_FP4_FORCE_A16", "0"))
1628
+ ),
1629
# Allow use of FlashInfer MxInt4 MoE kernels for fused moe ops.
1630
"VLLM_USE_FLASHINFER_MOE_INT4": lambda: bool(
1631
int(os.getenv("VLLM_USE_FLASHINFER_MOE_INT4", "0"))
vllm/model_executor/layers/fused_moe/b12x.py
new
+770
@@ -0,0 +1,770 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""b12x modular tensor-parallel fused MoE backend."""
4
+
5
+from collections.abc import Iterable
6
+from typing import Any
7
+
8
+import torch
9
+
10
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
13
+from vllm.model_executor.layers.fused_moe.config import (
14
+ FusedMoEParallelConfig,
15
+ FusedMoEQuantConfig,
16
+)
17
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
18
+ TopKWeightAndReduceNoOP,
19
+)
20
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
+ QuantKey,
22
+ kMxfp4Static,
23
+ kMxfp8Dynamic,
24
+ kNvfp4Dynamic,
25
+ kNvfp4Static,
26
+)
27
+from vllm.model_executor.utils import replace_parameter
28
+from vllm.platforms import current_platform
29
+from vllm.utils.b12x import (
30
+ B12xWarmupUnit,
31
+ get_b12x_fused_moe,
32
+ reuse_packed_weight_storage,
33
+)
34
+
35
+logger = init_logger(__name__)
36
+
37
+_B12X_MOE_MODES: dict[
38
+ tuple[torch.dtype | str | None, torch.dtype | str | None],
39
+ tuple[str, str, str],
40
+] = {
41
+ ("mxfp4", "mxfp8"): ("w4a8_mx", "fp4_e8m0_k32", "w31"),
42
+ ("mxfp4", None): ("w4a16", "fp4_e8m0_k32", "w31"),
43
+ ("nvfp4", "nvfp4"): ("nvfp4", "modelopt_nvfp4", "w31"),
44
+ ("nvfp4", "mxfp8"): ("w4a8_nvfp4", "modelopt_nvfp4", "w31"),
45
+ ("nvfp4", None): ("w4a16", "modelopt_nvfp4", "w13"),
46
+}
47
+
48
+
49
+def _require_b12x_fused_moe() -> Any:
50
+ fused_moe = get_b12x_fused_moe()
51
+ assert fused_moe is not None
52
+ return fused_moe
53
+
54
+
55
+def _b12x_activation_name(activation: MoEActivation) -> str:
56
+ if activation == MoEActivation.SILU:
57
+ return "silu"
58
+ if activation in (MoEActivation.RELU2, MoEActivation.RELU2_NO_MUL):
59
+ return "relu2"
60
+ return activation.value
61
+
62
+
63
+def _b12x_scratch_nbytes(plan: Any) -> int:
64
+ specs = plan.scratch_specs()
65
+ if len(specs) != 1:
66
+ raise RuntimeError(f"expected one b12x MoE scratch buffer, got {len(specs)}")
67
+ spec = specs[0]
68
+ if spec.dtype != torch.uint8:
69
+ raise TypeError(f"expected b12x MoE scratch dtype uint8, got {spec.dtype}")
70
+ return int(spec.shape[0])
71
+
72
+
73
+def _b12x_moe_execution_plan(
74
+ *,
75
+ tokens: int,
76
+ topk: int,
77
+ prepared: Any,
78
+ quant_mode: str,
79
+ apply_router_weight_on_input: bool,
80
+ swiglu_limit: float | None,
81
+ swiglu_alpha: float | None,
82
+ swiglu_beta: float | None,
83
+) -> Any:
84
+ fused_moe = _require_b12x_fused_moe()
85
+
86
+ return fused_moe.plan_execution(
87
+ num_tokens=max(int(tokens), 1),
88
+ num_topk=int(topk),
89
+ device=prepared.w1_fp4.device,
90
+ weight_plan=prepared.plan,
91
+ quant_mode=quant_mode,
92
+ apply_router_weight_on_input=apply_router_weight_on_input,
93
+ swiglu_limit=swiglu_limit,
94
+ swiglu_alpha=swiglu_alpha,
95
+ swiglu_beta=swiglu_beta,
96
+ )
97
+
98
+
99
+def _run_b12x_moe_plan(
100
+ *,
101
+ plan: Any,
102
+ scratch: torch.Tensor,
103
+ hidden_states: torch.Tensor,
104
+ prepared: Any,
105
+ topk_weights: torch.Tensor,
106
+ topk_ids: torch.Tensor,
107
+ output: torch.Tensor,
108
+ unit_scale_contract: bool,
109
+) -> None:
110
+ fused_moe = _require_b12x_fused_moe()
111
+
112
+ binding = fused_moe.bind(
113
+ plan,
114
+ scratch=scratch,
115
+ a=hidden_states,
116
+ experts=prepared,
117
+ topk_weights=topk_weights,
118
+ topk_ids=topk_ids,
119
+ output=output,
120
+ input_scales_static=True,
121
+ unit_scale_contract=unit_scale_contract,
122
+ )
123
+ fused_moe.run(binding=binding)
124
+
125
+
126
+def _is_current_stream_capturing() -> bool:
127
+ is_capturing = getattr(torch.cuda, "is_current_stream_capturing", None)
128
+ return bool(is_capturing is not None and is_capturing())
129
+
130
+
131
+def _normalize_topk_ids(topk_ids: torch.Tensor) -> torch.Tensor:
132
+ if topk_ids.dtype == torch.int32 and topk_ids.is_contiguous():
133
+ return topk_ids
134
+ if _is_current_stream_capturing():
135
+ raise RuntimeError(
136
+ "b12x MoE topk_ids normalization would allocate during CUDA capture"
137
+ )
138
+ return topk_ids.to(dtype=torch.int32).contiguous()
139
+
140
+
141
+def _normalize_topk_weights(topk_weights: torch.Tensor) -> torch.Tensor:
142
+ if topk_weights.dtype == torch.float32 and topk_weights.is_contiguous():
143
+ return topk_weights
144
+ if _is_current_stream_capturing():
145
+ raise RuntimeError(
146
+ "b12x MoE topk_weights normalization would allocate during CUDA capture"
147
+ )
148
+ return topk_weights.to(dtype=torch.float32).contiguous()
149
+
150
+
151
+def _workspace_as_b12x_scratch(
152
+ workspace: torch.Tensor | None,
153
+ plan: Any,
154
+) -> torch.Tensor:
155
+ if workspace is None:
156
+ raise RuntimeError("b12x MoE requires workspace2 scratch")
157
+ if not workspace.is_contiguous():
158
+ raise ValueError("b12x MoE workspace2 must be contiguous")
159
+ scratch = workspace.view(-1).view(torch.uint8)
160
+ required_nbytes = _b12x_scratch_nbytes(plan)
161
+ if scratch.numel() < required_nbytes:
162
+ raise ValueError(
163
+ "b12x MoE workspace2 is too small: "
164
+ f"have={scratch.numel()} bytes, need={required_nbytes} bytes"
165
+ )
166
+ return scratch
167
+
168
+
169
+def _replace_parameter_with_empty(
170
+ layer: torch.nn.Module,
171
+ name: str,
172
+) -> torch.Tensor | None:
173
+ parameter = getattr(layer, name, None)
174
+ if not isinstance(parameter, torch.Tensor):
175
+ return None
176
+ empty = torch.empty((0,), dtype=parameter.dtype, device=parameter.device)
177
+ replace_parameter(layer, name, empty)
178
+ return getattr(layer, name)
179
+
180
+
181
+def _normalize_expert_scale(scale: torch.Tensor) -> torch.Tensor:
182
+ if scale.ndim == 2:
183
+ if scale.shape[1] not in (1, 2):
184
+ raise ValueError(
185
+ "expected an expert scale with one or two columns, got "
186
+ f"{tuple(scale.shape)}"
187
+ )
188
+ scale = scale[:, 0]
189
+ return scale.to(dtype=torch.float32).contiguous()
190
+
191
+
192
+def _canonicalize_fp4_zero_signs_(packed: torch.Tensor) -> None:
193
+ """Clear sign bits from packed FP4 zero values in place."""
194
+ packed = packed.view(torch.uint8)
195
+ magnitude = packed & 0x77
196
+ nonzero = (magnitude | (magnitude >> 1) | (magnitude >> 2)) & 0x11
197
+ packed.bitwise_and_(0x77 | (nonzero << 3))
198
+
199
+
200
+class B12xExperts(mk.FusedMoEExpertsModular):
201
+ """FP4 MoE experts backed by the b12x SM12x planned API."""
202
+
203
+ def __init__(
204
+ self,
205
+ moe_config: mk.FusedMoEConfig,
206
+ quant_config: FusedMoEQuantConfig,
207
+ ):
208
+ super().__init__(moe_config, quant_config)
209
+ if quant_config.weight_quant_dtype not in ("mxfp4", "nvfp4"):
210
+ raise ValueError(
211
+ "b12x MoE requires MXFP4 or NVFP4 weights, got "
212
+ f"{quant_config.weight_quant_dtype}"
213
+ )
214
+ scheme = (
215
+ quant_config.weight_quant_dtype,
216
+ quant_config.quant_dtype,
217
+ )
218
+ try:
219
+ self._quant_mode, self._source_format, self._w13_layout = _B12X_MOE_MODES[
220
+ scheme
221
+ ]
222
+ except KeyError as exc:
223
+ raise ValueError(
224
+ f"unsupported b12x MoE quantization scheme {scheme}"
225
+ ) from exc
226
+ self._prepared_experts: Any | None = None
227
+ self._source_parameters_released = False
228
+ self._unit_scales: dict[torch.device, torch.Tensor] = {}
229
+ self._plans: dict[tuple[int, int, MoEActivation, bool], Any] = {}
230
+ self._apply_router_weight_on_input = False
231
+
232
+ def _unit_scale(self, device: torch.device, num_experts: int) -> torch.Tensor:
233
+ scale = self._unit_scales.get(device)
234
+ if scale is None or scale.numel() != num_experts:
235
+ scale = torch.ones(num_experts, dtype=torch.float32, device=device)
236
+ self._unit_scales[device] = scale
237
+ return scale
238
+
239
+ def _weight_global_scale(
240
+ self,
241
+ device: torch.device,
242
+ num_experts: int,
243
+ scale: torch.Tensor | None,
244
+ name: str,
245
+ ) -> torch.Tensor:
246
+ if self._source_format != "modelopt_nvfp4":
247
+ return self._unit_scale(device, num_experts)
248
+ if scale is None:
249
+ raise ValueError(f"b12x NVFP4 MoE requires {name}")
250
+ scale = _normalize_expert_scale(scale)
251
+ if scale.numel() != num_experts:
252
+ raise ValueError(
253
+ f"b12x NVFP4 MoE expected {num_experts} {name} values, "
254
+ f"got {scale.numel()}"
255
+ )
256
+ return scale.to(device=device)
257
+
258
+ def _swiglu_params(
259
+ self,
260
+ activation: MoEActivation,
261
+ ) -> tuple[float | None, float | None, float | None]:
262
+ if activation in (
263
+ MoEActivation.SITU,
264
+ MoEActivation.RELU2,
265
+ MoEActivation.RELU2_NO_MUL,
266
+ ):
267
+ return None, None, None
268
+
269
+ limit = self.quant_config.gemm1_clamp_limit
270
+ if limit is None:
271
+ limit = self.moe_config.swiglu_limit
272
+ if activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE:
273
+ return limit, None, None
274
+
275
+ alpha = self.quant_config.gemm1_alpha
276
+ if alpha is None:
277
+ alpha = self.moe_config.swiglu_alpha
278
+ beta = self.quant_config.gemm1_beta
279
+ if beta is None:
280
+ beta = self.moe_config.swiglu_beta
281
+ return limit, alpha, beta
282
+
283
+ def _prepare_experts(
284
+ self,
285
+ *,
286
+ w1: torch.Tensor,
287
+ w2: torch.Tensor,
288
+ activation: MoEActivation,
289
+ params_dtype: torch.dtype,
290
+ ) -> Any:
291
+ quant_mode = self._quant_mode
292
+ if _is_current_stream_capturing():
293
+ raise RuntimeError(
294
+ "b12x MoE weights must be prepared before CUDA graph capture"
295
+ )
296
+ if self.w1_scale is None or self.w2_scale is None:
297
+ raise ValueError("b12x MoE requires w1 and w2 block scales")
298
+
299
+ _canonicalize_fp4_zero_signs_(w1)
300
+ _canonicalize_fp4_zero_signs_(w2)
301
+
302
+ fused_moe = _require_b12x_fused_moe()
303
+
304
+ num_experts = int(w1.shape[0])
305
+ hidden_size = int(w2.shape[1])
306
+ intermediate_size = int(w2.shape[2]) * 2
307
+ unit_scale = self._unit_scale(w1.device, num_experts)
308
+ w1_global_scale = self._weight_global_scale(
309
+ w1.device, num_experts, self.g1_alphas, "w1 global scales"
310
+ )
311
+ w2_global_scale = self._weight_global_scale(
312
+ w2.device, num_experts, self.g2_alphas, "w2 global scales"
313
+ )
314
+
315
+ if quant_mode in ("nvfp4", "w4a8_nvfp4"):
316
+ if self.a1_gscale is None or self.a2_gscale is None:
317
+ raise ValueError("b12x NVFP4 MoE requires activation global scales")
318
+ a1_gscale = _normalize_expert_scale(self.a1_gscale).to(w1.device)
319
+ a2_gscale = _normalize_expert_scale(self.a2_gscale).to(w2.device)
320
+ else:
321
+ a1_gscale = unit_scale
322
+ a2_gscale = unit_scale
323
+
324
+ weight_plan = fused_moe.plan_weights(
325
+ quant_modes=quant_mode,
326
+ source_format=self._source_format,
327
+ activation=_b12x_activation_name(activation),
328
+ params_dtype=params_dtype,
329
+ num_experts=num_experts,
330
+ hidden_size=hidden_size,
331
+ intermediate_size=intermediate_size,
332
+ w13_layout=self._w13_layout,
333
+ )
334
+ return fused_moe.prepare_weights(
335
+ plan=weight_plan,
336
+ w1_fp4=w1,
337
+ w1_blockscale=self.w1_scale,
338
+ w1_global_scale=w1_global_scale,
339
+ a1_gscale=a1_gscale,
340
+ w2_fp4=w2,
341
+ w2_blockscale=self.w2_scale,
342
+ w2_global_scale=w2_global_scale,
343
+ a2_gscale=a2_gscale,
344
+ params_dtype=params_dtype,
345
+ )
346
+
347
+ def _refresh_quant_config(self, layer: torch.nn.Module) -> None:
348
+ self.quant_config._w1.scale = layer.w13_weight_scale
349
+ self.quant_config._w2.scale = layer.w2_weight_scale
350
+ if self._source_format != "modelopt_nvfp4":
351
+ return
352
+
353
+ self.quant_config._w1.alpha_or_gscale = layer.w13_weight_scale_2
354
+ self.quant_config._w2.alpha_or_gscale = layer.w2_weight_scale_2
355
+ if self._quant_mode in ("nvfp4", "w4a8_nvfp4"):
356
+ self.quant_config._a1.alpha_or_gscale = 1.0 / layer.w13_input_scale
357
+ self.quant_config._a2.alpha_or_gscale = 1.0 / layer.w2_input_scale
358
+
359
+ def _release_source_parameters(self, layer: torch.nn.Module) -> None:
360
+ if self._source_parameters_released:
361
+ return
362
+ w1_scale = _replace_parameter_with_empty(layer, "w13_weight_scale")
363
+ w2_scale = _replace_parameter_with_empty(layer, "w2_weight_scale")
364
+ if w1_scale is not None:
365
+ self.quant_config._w1.scale = w1_scale
366
+ if w2_scale is not None:
367
+ self.quant_config._w2.scale = w2_scale
368
+ _replace_parameter_with_empty(layer, "w13_weight")
369
+ _replace_parameter_with_empty(layer, "w2_weight")
370
+ self._source_parameters_released = True
371
+
372
+ def _reuse_prepared_storage(self, layer: torch.nn.Module, prepared: Any) -> Any:
373
+ previous = getattr(layer, "_b12x_prepared_experts", None)
374
+ prepared = reuse_packed_weight_storage(previous, prepared)
375
+ if prepared is not previous:
376
+ self._plans.clear()
377
+ self._prepared_experts = prepared
378
+ layer._b12x_prepared_experts = prepared
379
+ return prepared
380
+
381
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
382
+ self._apply_router_weight_on_input = layer.apply_router_weight_on_input
383
+ if self._apply_router_weight_on_input and self._quant_mode != "w4a16":
384
+ raise ValueError(
385
+ "b12x MoE supports apply_router_weight_on_input only with W4A16"
386
+ )
387
+ self._source_parameters_released = False
388
+ self._refresh_quant_config(layer)
389
+ prepared = self._prepare_experts(
390
+ w1=layer.w13_weight,
391
+ w2=layer.w2_weight,
392
+ activation=layer.activation,
393
+ params_dtype=self.moe_config.in_dtype,
394
+ )
395
+ prepared = self._reuse_prepared_storage(layer, prepared)
396
+ if prepared.plan.discards_source_parameters:
397
+ self._release_source_parameters(layer)
398
+ layer.b12x_warmup_provider = self
399
+
400
+ @staticmethod
401
+ def is_supported_config(
402
+ cls: type[mk.FusedMoEExperts],
403
+ moe_config: mk.FusedMoEConfig,
404
+ weight_key: QuantKey | None,
405
+ activation_key: QuantKey | None,
406
+ activation_format: mk.FusedMoEActivationFormat,
407
+ ) -> tuple[bool, str | None]:
408
+ if moe_config.has_bias:
409
+ return False, "kernel does not support expert biases"
410
+ if moe_config.in_dtype not in (torch.float16, torch.bfloat16):
411
+ return (
412
+ False,
413
+ f"kernel does not support {moe_config.in_dtype} input/output dtype",
414
+ )
415
+ if moe_config.activation == MoEActivation.SITU and (
416
+ moe_config.activation_situ_beta != 4.0
417
+ or moe_config.activation_situ_linear_beta != 25.0
418
+ ):
419
+ return False, "kernel supports only SiTU beta=4 and linear_beta=25"
420
+ if (
421
+ activation_key is not None
422
+ and moe_config.activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE
423
+ ):
424
+ return (
425
+ False,
426
+ "kernel does not support swigluoai_uninterleave with W4A8",
427
+ )
428
+ unpadded_intermediate_size = (
429
+ moe_config.intermediate_size_per_partition_unpadded
430
+ or moe_config.intermediate_size_per_partition
431
+ )
432
+ if weight_key == kMxfp4Static and unpadded_intermediate_size % 32 != 0:
433
+ return (
434
+ False,
435
+ "MXFP4 requires the per-rank intermediate size to be divisible by 32",
436
+ )
437
+ if weight_key == kMxfp4Static and activation_key == kMxfp8Dynamic:
438
+ if moe_config.activation not in (
439
+ MoEActivation.SILU,
440
+ MoEActivation.SITU,
441
+ ):
442
+ return False, "MXFP4 W4A8 supports only SiLU and SiTU"
443
+ if (
444
+ moe_config.hidden_dim % 256 != 0
445
+ or moe_config.intermediate_size_per_partition % 32 != 0
446
+ ):
447
+ return (
448
+ False,
449
+ "MXFP4 W4A8 requires hidden size divisible by 256 and "
450
+ "per-rank intermediate size divisible by 32",
451
+ )
452
+ return mk.FusedMoEExperts.is_supported_config(
453
+ cls, moe_config, weight_key, activation_key, activation_format
454
+ )
455
+
456
+ @staticmethod
457
+ def _supports_current_device() -> bool:
458
+ if not (
459
+ current_platform.is_cuda()
460
+ and current_platform.is_device_capability_family(120)
461
+ ):
462
+ return False
463
+ fused_moe = get_b12x_fused_moe()
464
+ if fused_moe is None:
465
+ return False
466
+ return fused_moe.is_supported()
467
+
468
+ @staticmethod
469
+ def _supports_no_act_and_mul() -> bool:
470
+ return True
471
+
472
+ @staticmethod
473
+ def _supports_quant_scheme(
474
+ weight_key: QuantKey | None,
475
+ activation_key: QuantKey | None,
476
+ ) -> bool:
477
+ return (weight_key, activation_key) in (
478
+ (kMxfp4Static, kMxfp8Dynamic),
479
+ (kMxfp4Static, None),
480
+ (kNvfp4Static, kNvfp4Dynamic),
481
+ (kNvfp4Static, kMxfp8Dynamic),
482
+ (kNvfp4Static, None),
483
+ )
484
+
485
+ @staticmethod
486
+ def _supports_activation(activation: MoEActivation) -> bool:
487
+ return activation in (
488
+ MoEActivation.SILU,
489
+ MoEActivation.SITU,
490
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
491
+ MoEActivation.RELU2_NO_MUL,
492
+ )
493
+
494
+ @staticmethod
495
+ def _supports_parallel_config(
496
+ moe_parallel_config: FusedMoEParallelConfig,
497
+ ) -> bool:
498
+ return (
499
+ not moe_parallel_config.use_ep
500
+ and moe_parallel_config.ep_size == 1
501
+ and not moe_parallel_config.use_all2all_kernels
502
+ and not moe_parallel_config.enable_eplb
503
+ )
504
+
505
+ @staticmethod
506
+ def activation_format() -> mk.FusedMoEActivationFormat:
507
+ return mk.FusedMoEActivationFormat.Standard
508
+
509
+ @property
510
+ def expects_unquantized_inputs(self) -> bool:
511
+ return True
512
+
513
+ def supports_expert_map(self) -> bool:
514
+ return False
515
+
516
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
517
+ return TopKWeightAndReduceNoOP()
518
+
519
+ def _prepared(self) -> Any:
520
+ if self._prepared_experts is None:
521
+ raise RuntimeError(
522
+ "b12x MoE weights must be prepared by process_weights_after_loading"
523
+ )
524
+ return self._prepared_experts
525
+
526
+ def moe_problem_size(
527
+ self,
528
+ a1: torch.Tensor,
529
+ w1: torch.Tensor,
530
+ w2: torch.Tensor,
531
+ topk_ids: torch.Tensor,
532
+ ) -> tuple[int, int, int, int, int]:
533
+ if w1.numel() and w2.numel():
534
+ return super().moe_problem_size(a1, w1, w2, topk_ids)
535
+ prepared = self._prepared()
536
+ tokens = int(a1.shape[0] if a1.ndim == 2 else a1.shape[1])
537
+ return (
538
+ int(prepared.num_experts),
539
+ tokens,
540
+ int(prepared.intermediate_size) * 2,
541
+ int(a1.shape[-1]),
542
+ int(topk_ids.shape[1]),
543
+ )
544
+
545
+ def _plan(
546
+ self,
547
+ *,
548
+ tokens: int,
549
+ topk: int,
550
+ activation: MoEActivation,
551
+ apply_router_weight_on_input: bool = False,
552
+ ) -> Any:
553
+ fused_moe = _require_b12x_fused_moe()
554
+
555
+ key = (
556
+ max(int(tokens), 1),
557
+ int(topk),
558
+ activation,
559
+ bool(apply_router_weight_on_input),
560
+ )
561
+ plan = self._plans.get(key)
562
+ if plan is not None:
563
+ return plan
564
+ if _is_current_stream_capturing():
565
+ raise RuntimeError("b12x MoE plans must be created before CUDA capture")
566
+
567
+ limit, alpha, beta = self._swiglu_params(activation)
568
+ prepared = self._prepared()
569
+ plan = fused_moe.plan(
570
+ fused_moe.Caps(
571
+ max_tokens=key[0],
572
+ num_topk=key[1],
573
+ device=prepared.w1_fp4.device,
574
+ weight_plan=prepared.plan,
575
+ core_token_counts=(key[0],),
576
+ route_num_experts=0,
577
+ quant_mode=self._quant_mode,
578
+ apply_router_weight_on_input=key[3],
579
+ swiglu_limit=limit,
580
+ swiglu_alpha=alpha,
581
+ swiglu_beta=beta,
582
+ frozen=True,
583
+ )
584
+ )
585
+ self._plans[key] = plan
586
+ return plan
587
+
588
+ def get_b12x_warmup_unit(
589
+ self,
590
+ layer: torch.nn.Module,
591
+ token_counts: tuple[int, ...],
592
+ output_dtype: torch.dtype,
593
+ ) -> B12xWarmupUnit:
594
+ assert output_dtype == self.moe_config.in_dtype
595
+ prepared = self._prepared()
596
+ activation = layer.activation
597
+ limit, alpha, beta = self._swiglu_params(activation)
598
+
599
+ def compile() -> None:
600
+ self.warmup_launches(layer, token_counts=token_counts)
601
+
602
+ return B12xWarmupUnit(
603
+ name="MoE",
604
+ key=(
605
+ type(self),
606
+ prepared.w1_fp4.device,
607
+ output_dtype,
608
+ self._quant_mode,
609
+ self._source_format,
610
+ self._w13_layout,
611
+ int(prepared.num_experts),
612
+ int(prepared.hidden_size),
613
+ int(prepared.intermediate_size),
614
+ int(self.moe_config.experts_per_token),
615
+ _b12x_activation_name(activation),
616
+ bool(layer.apply_router_weight_on_input),
617
+ limit,
618
+ alpha,
619
+ beta,
620
+ ),
621
+ compile=compile,
622
+ )
623
+
624
+ @torch.inference_mode()
625
+ def warmup_launches(
626
+ self,
627
+ layer: torch.nn.Module,
628
+ *,
629
+ token_counts: Iterable[int],
630
+ ) -> int:
631
+ """Compile one representative launch for every planned regime."""
632
+ activation = layer.activation
633
+ dtype = self.moe_config.in_dtype
634
+ topk = int(self.moe_config.experts_per_token)
635
+ apply_router_weight_on_input = bool(layer.apply_router_weight_on_input)
636
+ limit, alpha, beta = self._swiglu_params(activation)
637
+ prepared = self._prepared()
638
+ device = prepared.w1_fp4.device
639
+ launch_tokens: dict[tuple[Any, ...], int] = {}
640
+ for tokens in sorted({int(count) for count in token_counts if int(count) > 0}):
641
+ execution_plan = _b12x_moe_execution_plan(
642
+ tokens=tokens,
643
+ topk=topk,
644
+ prepared=prepared,
645
+ quant_mode=self._quant_mode,
646
+ apply_router_weight_on_input=apply_router_weight_on_input,
647
+ swiglu_limit=limit,
648
+ swiglu_alpha=alpha,
649
+ swiglu_beta=beta,
650
+ )
651
+ signature = (execution_plan.implementation, execution_plan.execution)
652
+ launch_tokens.setdefault(signature, tokens)
653
+
654
+ for tokens in launch_tokens.values():
655
+ hidden_states = torch.zeros(
656
+ (tokens, int(prepared.hidden_size)),
657
+ dtype=dtype,
658
+ device=device,
659
+ )
660
+ output = torch.empty_like(hidden_states)
661
+ topk_ids = (
662
+ torch.arange(topk, device=device, dtype=torch.int32)
663
+ .unsqueeze(0)
664
+ .expand(tokens, -1)
665
+ .contiguous()
666
+ )
667
+ topk_ids.remainder_(int(prepared.num_experts))
668
+ topk_weights = torch.full(
669
+ (tokens, topk),
670
+ 1.0 / topk,
671
+ dtype=torch.float32,
672
+ device=device,
673
+ )
674
+ plan = self._plan(
675
+ tokens=tokens,
676
+ topk=topk,
677
+ activation=activation,
678
+ apply_router_weight_on_input=apply_router_weight_on_input,
679
+ )
680
+ scratch = torch.empty(
681
+ (_b12x_scratch_nbytes(plan),),
682
+ dtype=torch.uint8,
683
+ device=device,
684
+ )
685
+ _run_b12x_moe_plan(
686
+ plan=plan,
687
+ scratch=scratch,
688
+ hidden_states=hidden_states,
689
+ prepared=prepared,
690
+ topk_weights=topk_weights,
691
+ topk_ids=topk_ids,
692
+ output=output,
693
+ unit_scale_contract=self._quant_mode == "w4a16",
694
+ )
695
+ return len(launch_tokens)
696
+
697
+ def workspace_shapes(
698
+ self,
699
+ M: int,
700
+ N: int,
701
+ K: int,
702
+ topk: int,
703
+ global_num_experts: int,
704
+ local_num_experts: int,
705
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
706
+ activation: MoEActivation,
707
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
708
+ del N, global_num_experts, local_num_experts, expert_tokens_meta
709
+ plan = self._plan(
710
+ tokens=M,
711
+ topk=topk,
712
+ activation=activation,
713
+ apply_router_weight_on_input=self._apply_router_weight_on_input,
714
+ )
715
+ itemsize = self.moe_config.in_dtype.itemsize
716
+ scratch_elements = max(
717
+ 1, (_b12x_scratch_nbytes(plan) + itemsize - 1) // itemsize
718
+ )
719
+ return (0,), (scratch_elements,), (M, K)
720
+
721
+ def apply(
722
+ self,
723
+ output: torch.Tensor,
724
+ hidden_states: torch.Tensor,
725
+ w1: torch.Tensor,
726
+ w2: torch.Tensor,
727
+ topk_weights: torch.Tensor,
728
+ topk_ids: torch.Tensor,
729
+ activation: MoEActivation,
730
+ global_num_experts: int,
731
+ expert_map: torch.Tensor | None,
732
+ a1q_scale: torch.Tensor | None,
733
+ a2_scale: torch.Tensor | None,
734
+ workspace13: torch.Tensor | None,
735
+ workspace2: torch.Tensor | None,
736
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
737
+ apply_router_weight_on_input: bool | None,
738
+ ) -> None:
739
+ del w1, w2, global_num_experts
740
+ del a1q_scale, a2_scale, workspace13, expert_tokens_meta
741
+ if expert_map is not None:
742
+ raise ValueError("b12x TP MoE does not support expert maps")
743
+ if bool(apply_router_weight_on_input) != self._apply_router_weight_on_input:
744
+ raise ValueError(
745
+ "apply_router_weight_on_input does not match the prepared b12x MoE plan"
746
+ )
747
+ prepared = self._prepared()
748
+ topk_ids = _normalize_topk_ids(topk_ids)
749
+ topk_weights = _normalize_topk_weights(topk_weights)
750
+ plan = self._plan(
751
+ tokens=int(hidden_states.shape[0]),
752
+ topk=int(topk_ids.shape[1]),
753
+ activation=activation,
754
+ apply_router_weight_on_input=bool(apply_router_weight_on_input),
755
+ )
756
+ scratch = _workspace_as_b12x_scratch(workspace2, plan)
757
+
758
+ _run_b12x_moe_plan(
759
+ plan=plan,
760
+ scratch=scratch,
761
+ hidden_states=hidden_states,
762
+ prepared=prepared,
763
+ topk_weights=topk_weights,
764
+ topk_ids=topk_ids,
765
+ output=output,
766
+ unit_scale_contract=self._quant_mode == "w4a16",
767
+ )
768
+
769
+ def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
770
+ raise NotImplementedError("LoRA is not supported for B12xExperts")
vllm/model_executor/layers/fused_moe/config.py
+4
@@ -865,6 +865,8 @@ def nvfp4_w4a16_moe_quant_config(
865
g2_alphas: torch.Tensor,
866
w1_scale: torch.Tensor,
867
w2_scale: torch.Tensor,
868
+ gemm1_alpha: float | None = None,
869
+ gemm1_beta: float | None = None,
870
gemm1_clamp_limit: float | None = None,
871
) -> FusedMoEQuantConfig:
872
"""
@@ -877,6 +879,8 @@ def nvfp4_w4a16_moe_quant_config(
879
g1_alphas=g1_alphas,
880
g2_alphas=g2_alphas,
881
weight_dtype="nvfp4",
882
+ gemm1_alpha=gemm1_alpha,
883
+ gemm1_beta=gemm1_beta,
884
gemm1_clamp_limit=gemm1_clamp_limit,
885
)
886
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+68
-10
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Literal, Union
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+from vllm import envs
10
from vllm.config import get_current_vllm_config
11
from vllm.config.kernel import MoEBackend
12
from vllm.config.quantization import QuantizationConfigArgs
@@ -100,6 +101,9 @@ def _pack_deepgemm_mxfp4_scales(
101
102
class Mxfp4MoeBackend(Enum):
103
NONE = "None"
104
+ # b12x backends
105
+ B12X_MXFP4_MXFP8 = "B12X_MXFP4_MXFP8"
106
+ B12X_MXFP4_BF16 = "B12X_MXFP4_BF16"
107
# DeepGEMM FP8xFP4 backend (SM100+)
108
DEEPGEMM_MXFP4 = "DEEPGEMM_MXFP4"
109
# FlashInfer TRTLLM backends
@@ -141,11 +145,21 @@ TRITON_BACKENDS = (
145
Mxfp4MoeBackend.TRITON_UNFUSED,
146
)
147
148
+B12X_BACKENDS = (
149
+ Mxfp4MoeBackend.B12X_MXFP4_MXFP8,
150
+ Mxfp4MoeBackend.B12X_MXFP4_BF16,
151
+)
152
+
153
154
def backend_to_kernel_cls(
155
backend: Mxfp4MoeBackend,
156
) -> list[type[mk.FusedMoEExperts]]:
148
- if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
157
+ if backend in B12X_BACKENDS:
158
+ from vllm.model_executor.layers.fused_moe.b12x import B12xExperts
159
+
160
+ return [B12xExperts]
161
+
162
+ elif backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
163
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
164
DeepGemmFP4Experts,
165
)
@@ -269,6 +283,7 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:
283
via ``activation_key`` and ``is_supported_config``.
284
"""
285
mapping: dict[str, list[Mxfp4MoeBackend]] = {
286
+ "b12x": list(B12X_BACKENDS),
287
"deep_gemm": [Mxfp4MoeBackend.DEEPGEMM_MXFP4],
288
"flashinfer_trtllm": [
289
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
@@ -355,6 +370,8 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:
370
"""Map backend to its activation key (FP8, MXFP8, or None for BF16)."""
371
if backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
372
return kFp8Dynamic128Sym
373
+ if backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8:
374
+ return kMxfp8Dynamic
375
if backend in (
376
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
377
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
@@ -441,6 +458,22 @@ def _filter_by_activation(
458
return bf16 if bf16 else backends
459
460
461
+def _get_requested_backends(
462
+ runner_backend: MoEBackend,
463
+ requested_activation_key: QuantKey | None,
464
+) -> list[Mxfp4MoeBackend]:
465
+ backends = map_mxfp4_backend(runner_backend)
466
+ if runner_backend == "b12x":
467
+ if envs.VLLM_B12X_MOE_FP4_FORCE_A16:
468
+ return [Mxfp4MoeBackend.B12X_MXFP4_BF16]
469
+ # W4A8 is the high-throughput b12x path and is preferred when the
470
+ # model does not request an activation format. Keep this as backend
471
+ # policy instead of exposing a separate backend name per precision.
472
+ if requested_activation_key is None:
473
+ return backends
474
+ return _filter_by_activation(backends, requested_activation_key)
475
+
476
+
477
def select_mxfp4_moe_backend(
478
config: FusedMoEConfig,
479
activation_key: QuantKey | None = None,
@@ -456,6 +489,7 @@ def select_mxfp4_moe_backend(
489
490
Note: Shape-specific fallbacks may still occur at runtime.
491
"""
492
+ runner_backend = config.moe_backend
493
requested_activation_key = _resolve_activation_key(activation_key)
494
495
activation_format = (
@@ -464,26 +498,25 @@ def select_mxfp4_moe_backend(
498
else mk.FusedMoEActivationFormat.Standard
499
)
500
467
- runner_backend = config.moe_backend
501
if runner_backend != "auto":
469
- requested_backends = map_mxfp4_backend(runner_backend)
502
+ requested_backends = _get_requested_backends(
503
+ runner_backend, requested_activation_key
504
+ )
505
if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
506
requested_backends = [
507
Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b
508
for b in requested_backends
509
]
475
- candidates = _filter_by_activation(requested_backends, requested_activation_key)
476
- if not candidates:
510
+ if not requested_backends:
511
raise ValueError(
512
f"moe_backend={runner_backend!r} does not support "
479
- f"activation={requested_activation_key}; supported variants: "
480
- f"{[b.name for b in requested_backends]}"
513
+ f"activation={requested_activation_key}"
514
)
515
last_error: Exception | None = None
483
- for requested_backend in candidates:
516
+ for requested_backend in requested_backends:
517
act_key = (
518
requested_activation_key
486
- if requested_activation_key is not None
519
+ if requested_backend == Mxfp4MoeBackend.EMULATION
520
else _backend_activation_key(requested_backend)
521
)
522
try:
@@ -577,7 +610,7 @@ def select_deepseek_v4_mxfp4_moe_backend(
610
# falling back to the auto priority list.
611
runner_backend = config.moe_backend
612
if runner_backend != "auto":
580
- requested_backends = map_mxfp4_backend(runner_backend)
613
+ requested_backends = _get_requested_backends(runner_backend, None)
614
if activation_format == mk.FusedMoEActivationFormat.BatchedExperts:
615
requested_backends = [
616
Mxfp4MoeBackend.BATCHED_MARLIN if b == Mxfp4MoeBackend.MARLIN else b
@@ -636,6 +669,10 @@ def mxfp4_round_up_hidden_size_and_intermediate_size(
669
activation: MoEActivation | None = None,
670
) -> tuple[int, int]:
671
"""Round up hidden_size and intermediate_size based on backend requirements."""
672
+ if backend in B12X_BACKENDS:
673
+ # b12x plans for the exact model dimensions. B12xExperts validates the
674
+ # required MXFP4 block alignment before selecting the backend.
675
+ return hidden_size, intermediate_size
676
if backend == Mxfp4MoeBackend.EMULATION:
677
# Emulation has no kernel tile; it only needs OCP MX block alignment so the
678
# per-block scale buffers (`dim // OCP_MX_BLOCK_SIZE`) aren't floor-truncated
@@ -1288,6 +1325,16 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1325
1326
is_gfx1250 = on_gfx1250()
1327
1328
+ if mxfp4_backend in B12X_BACKENDS:
1329
+ return (
1330
+ w13_weight.data,
1331
+ w2_weight.data,
1332
+ w13_weight_scale.data,
1333
+ w2_weight_scale.data,
1334
+ w13_bias,
1335
+ w2_bias,
1336
+ )
1337
+
1338
if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
1339
w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
1340
w13_weight,
@@ -1711,6 +1758,16 @@ def make_mxfp4_moe_quant_config(
1758
layer: "RoutedExperts | None" = None,
1759
) -> FusedMoEQuantConfig | None:
1760
"""Create a FusedMoEQuantConfig for the given MXFP4 backend."""
1761
+ if mxfp4_backend == Mxfp4MoeBackend.B12X_MXFP4_MXFP8:
1762
+ return mxfp4_mxfp8_moe_quant_config(
1763
+ w1_bias=w1_bias,
1764
+ w2_bias=w2_bias,
1765
+ w1_scale=w1_scale,
1766
+ w2_scale=w2_scale,
1767
+ gemm1_alpha=gemm1_alpha,
1768
+ gemm1_beta=gemm1_beta,
1769
+ gemm1_clamp_limit=swiglu_limit,
1770
+ )
1771
if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
1772
from vllm.model_executor.layers.quantization.utils.quant_utils import (
1773
GroupShape,
@@ -1778,6 +1835,7 @@ def make_mxfp4_moe_quant_config(
1835
gemm1_clamp_limit=swiglu_limit,
1836
)
1837
elif mxfp4_backend in (
1838
+ Mxfp4MoeBackend.B12X_MXFP4_BF16,
1839
Mxfp4MoeBackend.MARLIN,
1840
Mxfp4MoeBackend.BATCHED_MARLIN,
1841
Mxfp4MoeBackend.TRITON,
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+60
-4
@@ -5,6 +5,7 @@ from enum import Enum
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm import envs
9
from vllm.config.kernel import MoEBackend
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -17,6 +18,9 @@ from vllm.model_executor.layers.fused_moe.config import (
18
nvfp4_w4a16_moe_quant_config,
19
)
20
from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
21
+from vllm.model_executor.layers.quantization.utils.b12x_moe import (
22
+ prepare_nvfp4_moe_layer_for_b12x,
23
+)
24
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
25
nvfp4_swizzled_scale_to_cutedsl_mma_view,
26
prepare_nvfp4_moe_layer_for_fi_or_cutlass,
@@ -36,6 +40,7 @@ logger = init_logger(__name__)
40
41
42
class NvFp4MoeBackend(Enum):
43
+ B12X = "B12X"
44
FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
45
FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS"
46
FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL"
@@ -67,7 +72,12 @@ def is_global_sf_supported_for_nvfp4_backend(backend: NvFp4MoeBackend) -> bool:
72
def backend_to_kernel_cls(
73
backend: NvFp4MoeBackend,
74
) -> list[type[mk.FusedMoEExperts]]:
70
- if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
75
+ if backend == NvFp4MoeBackend.B12X:
76
+ from vllm.model_executor.layers.fused_moe.b12x import B12xExperts
77
+
78
+ return [B12xExperts]
79
+
80
+ elif backend == NvFp4MoeBackend.FLASHINFER_TRTLLM:
81
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
82
TrtLlmNvFp4ExpertsModular,
83
TrtLlmNvFp4ExpertsMonolithic,
@@ -145,6 +155,7 @@ def backend_to_kernel_cls(
155
def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend:
156
"""Map user's MoEBackend to NvFp4MoeBackend."""
157
mapping = {
158
+ "b12x": NvFp4MoeBackend.B12X,
159
"cutlass": NvFp4MoeBackend.VLLM_CUTLASS,
160
"flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM,
161
"flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS,
@@ -162,6 +173,12 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend:
173
)
174
175
176
+def _use_a16(backend: NvFp4MoeBackend, checkpoint_uses_a16: bool) -> bool:
177
+ return checkpoint_uses_a16 or (
178
+ backend == NvFp4MoeBackend.B12X and envs.VLLM_B12X_MOE_FP4_FORCE_A16
179
+ )
180
+
181
+
182
def select_nvfp4_moe_backend(
183
config: FusedMoEConfig,
184
weight_key: QuantKey | None,
@@ -188,6 +205,7 @@ def select_nvfp4_moe_backend(
205
]
206
207
NVFP4_BACKENDS_WITH_CLAMP = {
208
+ NvFp4MoeBackend.B12X,
209
NvFp4MoeBackend.FLASHINFER_TRTLLM,
210
NvFp4MoeBackend.FLASHINFER_CUTLASS,
211
NvFp4MoeBackend.FLASHINFER_CUTEDSL,
@@ -249,6 +267,8 @@ def select_nvfp4_moe_backend(
267
runner_backend = config.moe_backend
268
if runner_backend != "auto":
269
requested_backend = map_nvfp4_backend(runner_backend)
270
+ if _use_a16(requested_backend, False):
271
+ activation_key = None
272
# For batched activation format, use batched variant if available.
273
if (
274
activation_format == mk.FusedMoEActivationFormat.BatchedExperts
@@ -264,7 +284,7 @@ def select_nvfp4_moe_backend(
284
f"explicitly requested moe_backend={runner_backend!r} does "
285
f"not apply the SwiGLU clamp. Use 'flashinfer_trtllm', "
286
f"'flashinfer_cutlass', 'flashinfer_cutedsl', 'cutlass', "
267
- f"'marlin', or 'humming' instead."
287
+ f"'b12x', 'marlin', or 'humming' instead."
288
)
289
return _return_or_raise(
290
requested_backend, config, weight_key, activation_key, activation_format
@@ -303,6 +323,7 @@ def convert_to_nvfp4_moe_kernel_format(
323
w2_scale_2: torch.Tensor,
324
a2_scale: torch.Tensor | None,
325
is_act_and_mul: bool,
326
+ use_a16: bool = False,
327
) -> tuple[
328
torch.Tensor,
329
torch.Tensor,
@@ -313,7 +334,36 @@ def convert_to_nvfp4_moe_kernel_format(
334
torch.Tensor,
335
torch.Tensor,
336
]:
316
- if nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
337
+ use_a16 = _use_a16(nvfp4_backend, use_a16)
338
+ if nvfp4_backend == NvFp4MoeBackend.B12X:
339
+ if a13_scale is None or a2_scale is None:
340
+ if not use_a16:
341
+ raise ValueError("b12x NVFP4 MoE requires activation scales")
342
+ num_experts = w13.shape[0]
343
+ a13_scale = torch.ones(num_experts, dtype=torch.float32, device=w13.device)
344
+ a2_scale = torch.ones(num_experts, dtype=torch.float32, device=w2.device)
345
+ (
346
+ w13,
347
+ w13_scale,
348
+ w13_scale_2,
349
+ a13_scale,
350
+ w2,
351
+ w2_scale,
352
+ w2_scale_2,
353
+ a2_scale,
354
+ ) = prepare_nvfp4_moe_layer_for_b12x(
355
+ w13=w13,
356
+ w13_scale=w13_scale,
357
+ w13_scale_2=w13_scale_2,
358
+ a13_scale=a13_scale,
359
+ w2=w2,
360
+ w2_scale=w2_scale,
361
+ w2_scale_2=w2_scale_2,
362
+ a2_scale=a2_scale,
363
+ is_act_and_mul=is_act_and_mul,
364
+ reorder_w13=use_a16,
365
+ )
366
+ elif nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
367
(
368
w13,
369
w13_scale,
@@ -470,7 +520,9 @@ def make_nvfp4_moe_quant_config(
520
swiglu_alpha: float | None = None,
521
swiglu_beta: float | None = None,
522
layer: torch.nn.Module | None = None,
523
+ use_a16: bool = False,
524
) -> FusedMoEQuantConfig:
525
+ use_a16 = _use_a16(backend, use_a16)
526
if backend == NvFp4MoeBackend.HUMMING:
527
from vllm.model_executor.layers.fused_moe import RoutedExperts
528
from vllm.model_executor.layers.quantization.utils.humming_utils import (
@@ -484,12 +536,16 @@ def make_nvfp4_moe_quant_config(
536
gemm1_beta=getattr(layer, "swiglu_beta", None),
537
gemm1_clamp_limit=swiglu_limit,
538
)
487
- elif backend == NvFp4MoeBackend.MARLIN:
539
+ elif backend == NvFp4MoeBackend.MARLIN or (
540
+ backend == NvFp4MoeBackend.B12X and use_a16
541
+ ):
542
return nvfp4_w4a16_moe_quant_config(
543
g1_alphas=w13_scale_2,
544
g2_alphas=w2_scale_2,
545
w1_scale=w13_scale,
546
w2_scale=w2_scale,
547
+ gemm1_alpha=swiglu_alpha,
548
+ gemm1_beta=swiglu_beta,
549
gemm1_clamp_limit=swiglu_limit,
550
)
551
elif backend == NvFp4MoeBackend.EMULATION:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+14
-4
@@ -25,9 +25,11 @@ from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
25
XPUExpertsMxFp4,
26
)
27
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
28
+ B12X_BACKENDS,
29
Mxfp4MoeBackend,
30
make_mxfp4_moe_kernel,
31
make_mxfp4_moe_quant_config,
32
+ select_mxfp4_moe_backend,
33
)
34
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501
35
CompressedTensorsMoEMethod,
@@ -46,10 +48,17 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
48
super().__init__(moe)
49
self.group_size = 32
50
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
49
- # use cutlass if supported, otherwise fallback to marlin for weight-only FP4
51
+ # Backend selection must match the weight preparation below: CUTLASS
52
+ # swizzles scales, b12x and XPU consume checkpoint packing, and Marlin
53
+ # repacks weights and scales.
54
self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
55
self.experts_cls: type[mk.FusedMoEExperts]
52
- if self.use_cutlass_mxfp4:
56
+ if moe.moe_backend == "b12x":
57
+ self.mxfp4_backend, experts_cls = select_mxfp4_moe_backend(moe)
58
+ assert experts_cls is not None
59
+ self.experts_cls = experts_cls
60
+ self.use_cutlass_mxfp4 = False
61
+ elif self.use_cutlass_mxfp4:
62
logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE")
63
self.experts_cls = CutlassExpertsMxfp4
64
elif current_platform.is_xpu():
@@ -138,7 +147,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
147
w2_scale=layer.w2_weight_scale,
148
)
149
else:
141
- # W4A16: weight-only via Marlin
150
+ # b12x selects W4A8 or W4A16; XPU uses W4A4; Marlin uses W4A16.
151
return make_mxfp4_moe_quant_config(
152
mxfp4_backend=self.mxfp4_backend,
153
w1_scale=layer.w13_weight_scale,
@@ -188,7 +197,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
197
layer.w2_weight_scale = torch.nn.Parameter(
198
torch.stack(swizzled_w2), requires_grad=False
199
)
191
- elif current_platform.is_xpu():
200
+ elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu():
201
pass
202
else:
203
logger.warning_once(
@@ -208,6 +217,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
217
mxfp4_backend=self.mxfp4_backend,
218
routing_tables=layer._expert_routing_tables(),
219
)
220
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
221
222
def apply(
223
self,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+3
@@ -52,6 +52,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
52
):
53
super().__init__(moe)
54
self.group_size = 16
55
+ self.use_a16 = use_a16
56
57
# Select experts implementation.
58
self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
@@ -231,6 +232,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
232
w2_scale_2=(1.0 / layer.w2_weight_global_scale),
233
a2_scale=(1.0 / layer.w2_input_global_scale),
234
is_act_and_mul=self.moe.is_act_and_mul,
235
+ use_a16=self.use_a16,
236
)
237
238
replace_parameter(layer, "w13_weight", w13)
@@ -267,6 +269,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
269
swiglu_alpha=getattr(layer, "swiglu_alpha", None),
270
swiglu_beta=getattr(layer, "swiglu_beta", None),
271
layer=layer,
272
+ use_a16=self.use_a16,
273
)
274
275
def apply_monolithic(
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py
+18
-15
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
24
FusedMoEMethodBase,
25
)
26
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
27
+ B12X_BACKENDS,
28
Mxfp4MoeBackend,
29
make_mxfp4_moe_kernel,
30
make_mxfp4_moe_quant_config,
@@ -42,26 +43,29 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):
43
"""W4A4 MXFP4 group MoE for AutoRound ``auto_round:llm_compressor`` exports.
44
45
Registers the packed MXFP4 layout (uint8 ``weight_packed`` + uint8 E8M0
45
- ``weight_scale``, ``group_size=32``) and dispatches the fused MoE to the
46
- best backend for the current device: CUTLASS (true W4A4 on supported
47
- GPUs), the native XPU kernel, or Marlin weight-only as a fallback. The
48
- per-expert ``gate_proj`` / ``up_proj`` / ``down_proj`` tensors are folded
49
- into the stacked ``w13`` / ``w2`` parameters by ``make_expert_params_mapping``.
46
+ ``weight_scale``, ``group_size=32``) and dispatches the fused MoE to b12x
47
+ when requested; otherwise it uses CUTLASS W4A4, XPU W4A4, or Marlin W4A16.
48
+ The per-expert ``gate_proj`` / ``up_proj`` / ``down_proj`` tensors are
49
+ folded into the stacked ``w13`` / ``w2`` parameters by
50
+ ``make_expert_params_mapping``.
51
"""
52
53
def __init__(self, moe) -> None:
54
super().__init__(moe)
55
self.group_size = 32
56
# Backend selection must stay consistent with the weight preparation in
56
- # process_weights_after_loading / get_fused_moe_quant_config, which only
57
- # implement three layouts: CUTLASS swizzle (true W4A4), the native XPU
58
- # kernel (packed passthrough), and Marlin weight-only. XPU dispatch is
59
- # deferred to the shared oracle; every other non-CUTLASS device falls
60
- # back to Marlin (mirroring CompressedTensorsW4A4Mxfp4MoEMethod).
57
+ # process_weights_after_loading / get_fused_moe_quant_config, which use
58
+ # three layouts: CUTLASS swizzle (true W4A4), b12x and XPU packed
59
+ # passthrough, and Marlin weight-only. XPU dispatch is deferred to the
60
+ # shared oracle; b12x dispatch is explicit; every remaining non-CUTLASS
61
+ # device falls back to Marlin.
62
self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
63
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
64
self.experts_cls: type[mk.FusedMoEExperts] | None = None
64
- if self.use_cutlass_mxfp4:
65
+ if moe.moe_backend == "b12x":
66
+ self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
67
+ self.use_cutlass_mxfp4 = False
68
+ elif self.use_cutlass_mxfp4:
69
self.experts_cls = CutlassExpertsMxfp4
70
logger.info_once("Using CutlassExpertsMxfp4 for AutoRound MXFP4 MoE")
71
elif current_platform.is_xpu():
@@ -146,7 +150,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):
150
w1_scale=layer.w13_weight_scale,
151
w2_scale=layer.w2_weight_scale,
152
)
149
- # Weight-only (Marlin) or native XPU kernel.
153
+ # b12x selects W4A8 or W4A16; XPU uses W4A4; Marlin uses W4A16.
154
return make_mxfp4_moe_quant_config(
155
mxfp4_backend=self.mxfp4_backend,
156
w1_scale=layer.w13_weight_scale,
@@ -194,9 +198,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):
198
layer.w2_weight_scale = torch.nn.Parameter(
199
torch.stack(swizzled_w2), requires_grad=False
200
)
197
- elif current_platform.is_xpu():
198
- # The XPU fused-MoE kernel consumes the packed layout directly; no
199
- # swizzle / repack / transpose is required.
201
+ elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu():
202
pass
203
else:
204
logger.warning_once(
@@ -216,6 +218,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):
218
mxfp4_backend=self.mxfp4_backend,
219
routing_tables=layer._expert_routing_tables(),
220
)
221
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
222
223
def apply(
224
self,
vllm/model_executor/layers/quantization/modelopt.py
+2
@@ -1574,6 +1574,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1574
w2_scale_2=layer.w2_weight_scale_2,
1575
a2_scale=layer.w2_input_scale,
1576
is_act_and_mul=self.moe.is_act_and_mul,
1577
+ use_a16=self.use_a16,
1578
)
1579
1580
replace_parameter(layer, "w13_weight", w13)
@@ -1610,6 +1611,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1611
swiglu_alpha=getattr(layer, "swiglu_alpha", None),
1612
swiglu_beta=getattr(layer, "swiglu_beta", None),
1613
layer=layer,
1614
+ use_a16=self.use_a16,
1615
)
1616
1617
@property
vllm/model_executor/layers/quantization/mxfp4.py
+2
@@ -376,6 +376,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
376
experts_cls=self.experts_cls,
377
routing_tables=layer._expert_routing_tables(),
378
)
379
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
380
381
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
382
w13 = layer.w13_weight
@@ -769,6 +770,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
770
experts_cls=self.experts_cls,
771
routing_tables=layer._expert_routing_tables(),
772
)
773
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
774
775
def process_weights_after_loading(self, layer):
776
w13 = layer.w13_weight
vllm/model_executor/layers/quantization/online/mxfp4.py
+1
@@ -255,6 +255,7 @@ class Mxfp4OnlineMoEMethod(OnlineMoEMethodBase):
255
experts_cls=self.experts_cls,
256
routing_tables=layer._expert_routing_tables(),
257
)
258
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
259
260
def get_fused_moe_quant_config(
261
self, layer: torch.nn.Module
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
@@ -1334,6 +1334,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1334
experts_cls=self.experts_cls,
1335
routing_tables=layer._expert_routing_tables(),
1336
)
1337
+ self.moe_kernel.fused_experts.process_weights_after_loading(layer)
1338
1339
def get_fused_moe_quant_config(
1340
self, layer: RoutedExperts
vllm/model_executor/layers/quantization/utils/b12x_moe.py
new
+108
@@ -0,0 +1,108 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Weight preparation helpers for b12x MoE kernels."""
4
+
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
9
+ swap_w13_to_w31,
10
+)
11
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
12
+ swizzle_blockscale,
13
+)
14
+from vllm.utils.math_utils import round_up
15
+
16
+
17
+def _pad_gated_weights(
18
+ w13: torch.Tensor,
19
+ w13_scale: torch.Tensor,
20
+ w2: torch.Tensor,
21
+ w2_scale: torch.Tensor,
22
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
23
+ if w13_scale.size(1) != w13.size(1):
24
+ raise ValueError("w13 weight and scale row counts must match")
25
+ half_size = w13.size(1) // 2
26
+ if half_size * 2 != w13.size(1):
27
+ raise ValueError("gated NVFP4 MoE weights must have even row counts")
28
+ if w2.size(2) * 2 != half_size:
29
+ raise ValueError("w2 shape does not match gated w13")
30
+ if w2_scale.size(2) * 16 != half_size:
31
+ raise ValueError("w2 scale shape does not match gated w13")
32
+
33
+ half_pad_size = round_up(half_size, 64) - half_size
34
+ if half_pad_size == 0:
35
+ return w13, w13_scale, w2, w2_scale
36
+
37
+ def pad_rows(tensor: torch.Tensor) -> torch.Tensor:
38
+ shape = tensor.shape
39
+ return F.pad(
40
+ tensor.reshape(shape[0], 2, half_size, *shape[2:]),
41
+ (0, 0, 0, half_pad_size),
42
+ ).flatten(1, 2)
43
+
44
+ return (
45
+ pad_rows(w13),
46
+ pad_rows(w13_scale),
47
+ F.pad(w2, (0, half_pad_size // 2)),
48
+ F.pad(w2_scale, (0, half_pad_size // 16)),
49
+ )
50
+
51
+
52
+def _per_expert_scale(
53
+ scale: torch.Tensor,
54
+ num_experts: int,
55
+ name: str,
56
+) -> torch.Tensor:
57
+ if scale.dim() == 0:
58
+ scale = scale.expand(num_experts)
59
+ elif scale.dim() == 2:
60
+ scale = scale.amax(dim=1)
61
+ if scale.dim() != 1 or scale.numel() != num_experts:
62
+ raise ValueError(f"{name} must contain one value per expert")
63
+ return scale.to(torch.float32).contiguous()
64
+
65
+
66
+def prepare_nvfp4_moe_layer_for_b12x(
67
+ w13: torch.Tensor,
68
+ w13_scale: torch.Tensor,
69
+ w13_scale_2: torch.Tensor,
70
+ a13_scale: torch.Tensor,
71
+ w2: torch.Tensor,
72
+ w2_scale: torch.Tensor,
73
+ w2_scale_2: torch.Tensor,
74
+ a2_scale: torch.Tensor,
75
+ is_act_and_mul: bool,
76
+ reorder_w13: bool = False,
77
+) -> tuple[
78
+ torch.Tensor,
79
+ torch.Tensor,
80
+ torch.Tensor,
81
+ torch.Tensor,
82
+ torch.Tensor,
83
+ torch.Tensor,
84
+ torch.Tensor,
85
+ torch.Tensor,
86
+]:
87
+ """Prepare b12x NVFP4 MoE weights and scales."""
88
+ num_experts = w13.shape[0]
89
+ a13_scale = _per_expert_scale(a13_scale, num_experts, "a13_scale")
90
+ a2_scale = _per_expert_scale(a2_scale, num_experts, "a2_scale")
91
+
92
+ if reorder_w13 and is_act_and_mul:
93
+ w13 = swap_w13_to_w31(w13)
94
+ w13_scale = swap_w13_to_w31(w13_scale)
95
+ if is_act_and_mul:
96
+ w13, w13_scale, w2, w2_scale = _pad_gated_weights(w13, w13_scale, w2, w2_scale)
97
+
98
+ w13_scale = swizzle_blockscale(w13_scale)
99
+ pad_size = w13_scale.size(1) - w13.size(1)
100
+ if pad_size > 0:
101
+ if is_act_and_mul:
102
+ raise RuntimeError("gated NVFP4 MoE padding must precede scale swizzling")
103
+ w13 = F.pad(w13, (0, 0, 0, pad_size))
104
+ w2 = F.pad(w2, (0, pad_size // 2, 0, 0))
105
+ w2_scale = F.pad(w2_scale, (0, pad_size // 16))
106
+
107
+ w2_scale = swizzle_blockscale(w2_scale)
108
+ return w13, w13_scale, w13_scale_2, a13_scale, w2, w2_scale, w2_scale_2, a2_scale
vllm/model_executor/warmup/b12x_warmup.py
+15
-4
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Warm B12X JIT kernels used by a loaded model."""
3
+"""Warm b12x JIT kernels used by a loaded model."""
4
5
from collections import Counter
6
from collections.abc import Iterable
@@ -61,9 +61,20 @@ def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None:
61
)
62
if output_dtype not in (torch.bfloat16, torch.float16):
63
output_dtype = torch.bfloat16
64
+ compile_sizes = worker.vllm_config.compilation_config.compile_sizes or []
65
+ max_num_batched_tokens = worker.scheduler_config.max_num_batched_tokens
66
+ serving_sizes = [
67
+ max_num_batched_tokens,
68
+ *cudagraph_capture_sizes,
69
+ *(size for size in compile_sizes if isinstance(size, int)),
70
+ ]
71
+ max_tokens = max_num_batched_tokens
72
+ max_num_scheduled_tokens = worker.scheduler_config.max_num_scheduled_tokens
73
+ if max_num_scheduled_tokens is not None:
74
+ max_tokens = max(max_tokens, max_num_scheduled_tokens)
75
token_counts = b12x_warmup_token_counts(
65
- max_tokens=worker.scheduler_config.max_num_batched_tokens,
66
- cudagraph_capture_sizes=cudagraph_capture_sizes,
76
+ max_tokens=max_tokens,
77
+ cudagraph_capture_sizes=serving_sizes,
78
)
79
units = _collect_warmup_units(
80
worker.get_model(),
@@ -72,7 +83,7 @@ def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None:
83
)
84
for name, count in _compile_warmup_units(units).items():
85
logger.info_once(
75
- "Warmed up %d B12X %s linear GEMM signatures.",
86
+ "Warmed up %d b12x %s kernel signature(s).",
87
count,
88
name,
89
)
vllm/utils/b12x.py
+4
@@ -52,6 +52,10 @@ def get_b12x_tensor_fp8_linear() -> ModuleType | None:
52
return _get_submodule("b12x.gemm.tensor_fp8_linear")
53
54
55
+def get_b12x_fused_moe() -> ModuleType | None:
56
+ return _get_submodule("b12x.moe.fused_moe")
57
+
58
+
59
def b12x_warmup_token_counts(
60
*,
61
max_tokens: int,