@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 363fc84407
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 05:30 UTC

Integrate flashinfer b12x MoE and FP4 GEMM kernels for SM120/121 (#40082)

Signed-off-by: Meenakshi Venkataraman <meenakshiv@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

meena-at-work committed May 20, 2026 at 10:21 UTC 363fc84407f8c966c1cee6786e45e9e6ab289684
10 files changed +595 -4
tests/kernels/moe/test_flashinfer_b12x_moe.py new
+229
@@ -0,0 +1,229 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.platforms import current_platform
8 +
9 +if not current_platform.is_device_capability_family(120):
10 + pytest.skip(
11 + reason="FlashInfer CuteDSL SM12x MoE requires SM120 "
12 + "(RTX Pro 6000 / DGX Spark).",
13 + allow_module_level=True,
14 + )
15 +
16 +from vllm.utils.flashinfer import has_flashinfer_b12x_moe
17 +
18 +if not has_flashinfer_b12x_moe():
19 + pytest.skip(
20 + reason=(
21 + "FlashInfer cute_dsl_fused_moe_nvfp4 / convert_sf_to_mma_layout "
22 + "not available in installed FlashInfer (needs PRs #3051 and #3066)."
23 + ),
24 + allow_module_level=True,
25 + )
26 +
27 +# Import fp4_quantize after the skip guard — FlashInfer must be installed.
28 +from flashinfer.fp4_quantization import fp4_quantize
29 +
30 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
31 +from tests.kernels.moe.utils import make_dummy_moe_config
32 +from tests.kernels.utils import torch_moe
33 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
34 +from vllm.model_executor.layers.fused_moe import fused_topk
35 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
36 +from vllm.model_executor.layers.fused_moe.all2all_utils import (
37 + maybe_make_prepare_finalize,
38 +)
39 +from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config
40 +from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import (
41 + FlashInferB12xExperts,
42 +)
43 +from vllm.utils.flashinfer import flashinfer_convert_sf_to_mma_layout
44 +from vllm.utils.torch_utils import set_random_seed
45 +
46 +# Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256,
47 +# n multiple of 128) while keeping tests fast.
48 +MNK_FACTORS = [
49 + (2, 128, 256),
50 + (2, 256, 512),
51 + (16, 128, 256),
52 + (64, 256, 512),
53 +]
54 +
55 +
56 +def _reorder_gate_up_to_up_gate(
57 + w: torch.Tensor,
58 + w_s: torch.Tensor,
59 +) -> tuple[torch.Tensor, torch.Tensor]:
60 + """Swap gate and up-projection halves along dim=1 to [up, gate] order.
61 +
62 + The SM12x kernel expects weights in [up (w3), gate (w1)] order while the
63 + BF16 reference uses [gate (w1), up (w3)]. This replicates the reordering
64 + done at model-load time by ``prepare_nvfp4_moe_layer_for_fi_or_cutlass``.
65 + """
66 + n = w.shape[1] // 2
67 + return (
68 + torch.cat([w[:, n:, :], w[:, :n, :]], dim=1),
69 + torch.cat([w_s[:, n:, :], w_s[:, :n, :]], dim=1),
70 + )
71 +
72 +
73 +@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
74 +@pytest.mark.parametrize("e", [8, 16])
75 +@pytest.mark.parametrize("topk", [1, 2, 4])
76 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
77 +@torch.inference_mode()
78 +def test_flashinfer_b12x_moe(
79 + m: int,
80 + n: int,
81 + k: int,
82 + e: int,
83 + topk: int,
84 + dtype: torch.dtype,
85 + workspace_init,
86 +):
87 + """Test FlashInferB12xExperts against a BF16 torch reference.
88 +
89 + The SM12x kernel takes BF16 hidden states directly and fuses token
90 + dispatch, W1 GEMM, SwiGLU, and W2 GEMM into one call. We verify
91 + correctness against ``torch_moe`` using generous tolerances to account
92 + for the internal FP4 quantization of activations and weights.
93 +
94 + Scale convention
95 + ----------------
96 + The SM12x kernel uses ``w1_alpha`` as *both* the activation-quantisation
97 + global scale and the weight dequantisation factor. These two roles are
98 + conflated into a single parameter in ``launch_sm120_moe``, so they must
99 + equal the same value. We use ``global_scale = 1.0`` for
100 + ``fp4_quantize`` so that ``w1_alpha = ones`` satisfies both roles
101 + simultaneously. The alternative — vLLM's convention of baking a large
102 + ``w_gs`` into block-scale values and compensating with
103 + ``g1_alphas = 1/w_gs`` — is incompatible with this kernel.
104 + """
105 + set_random_seed(7)
106 + with set_current_vllm_config(
107 + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
108 + ):
109 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
110 +
111 + # Generate BF16 reference weights in [gate, up] order.
112 + # Shape: w1=(e, 2n, k), w2=(e, k, n).
113 + w1_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 15
114 + w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15
115 +
116 + # ------------------------------------------------------------------ #
117 + # Quantise weights for the SM12x kernel using FlashInfer's convention:
118 + # global_scale = 1.0 → block_scale = max_abs_block / fp4_max
119 + # w1_alpha = 1.0 (no extra global factor to compensate)
120 + #
121 + # The scale factors returned by fp4_quantize(..., is_sf_swizzled_layout=True)
122 + # are already in the swizzled 2D layout expected by convert_sf_to_mma_layout.
123 + # No additional swizzle_blockscale() call is needed.
124 + # ------------------------------------------------------------------ #
125 + gs = torch.ones(1, device="cuda", dtype=torch.float32)
126 + sf_vec_size = 16
127 +
128 + # W1: reorder BF16 from [gate, up] → [up, gate], then quantise.
129 + w1_reordered = torch.cat(
130 + [w1_bf16[:, n:, :], w1_bf16[:, :n, :]], dim=1
131 + ) # shape (e, 2n, k), [up, gate]
132 + w1_flat = w1_reordered.reshape(e * 2 * n, k)
133 + w1_q_flat, w1_sf_flat = fp4_quantize(
134 + w1_flat,
135 + global_scale=gs,
136 + sf_vec_size=sf_vec_size,
137 + is_sf_swizzled_layout=True,
138 + )
139 + w1_q = w1_q_flat.view(e, 2 * n, k // 2) # uint8, packed FP4
140 + w1_blockscale = w1_sf_flat.view(e, 2 * n, w1_sf_flat.shape[1]) # float8
141 +
142 + # W2: no row reordering needed for the down-projection.
143 + w2_flat = w2_bf16.reshape(e * k, n)
144 + w2_q_flat, w2_sf_flat = fp4_quantize(
145 + w2_flat,
146 + global_scale=gs,
147 + sf_vec_size=sf_vec_size,
148 + is_sf_swizzled_layout=True,
149 + )
150 + w2_q = w2_q_flat.view(e, k, n // 2) # uint8, packed FP4
151 + w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1]) # float8
152 +
153 + # All per-expert alphas are 1.0 (global_scale = 1.0, no compensation).
154 + ones_e = torch.ones(e, device="cuda", dtype=torch.float32)
155 +
156 + quant_config = nvfp4_moe_quant_config(
157 + g1_alphas=ones_e,
158 + g2_alphas=ones_e,
159 + a1_gscale=ones_e,
160 + a2_gscale=ones_e,
161 + w1_scale=w1_blockscale,
162 + w2_scale=w2_blockscale,
163 + )
164 +
165 + moe_config = make_dummy_moe_config(
166 + num_experts=e,
167 + experts_per_token=topk,
168 + hidden_dim=k,
169 + intermediate_size_per_partition=n,
170 + in_dtype=dtype,
171 + )
172 +
173 + experts = FlashInferB12xExperts(
174 + moe_config=moe_config,
175 + quant_config=quant_config,
176 + )
177 + # In production, process_weights_after_loading computes these after
178 + # normalizing block scales. In the test the scales are already in final
179 + # form (global_scale=1.0), so we compute the MMA layouts directly.
180 + num_experts_w1, m1, k1_sf = w1_blockscale.shape
181 + experts.w1_sf_mma = flashinfer_convert_sf_to_mma_layout(
182 + w1_blockscale.reshape(num_experts_w1 * m1, k1_sf),
183 + m=m1,
184 + k=k1_sf * 16,
185 + num_groups=num_experts_w1,
186 + )
187 + num_experts_w2, m2, k2_sf = w2_blockscale.shape
188 + experts.w2_sf_mma = flashinfer_convert_sf_to_mma_layout(
189 + w2_blockscale.reshape(num_experts_w2 * m2, k2_sf),
190 + m=m2,
191 + k=k2_sf * 16,
192 + num_groups=num_experts_w2,
193 + )
194 +
195 + kernel = mk.FusedMoEKernel(
196 + maybe_make_prepare_finalize(
197 + moe=moe_config,
198 + quant_config=quant_config,
199 + allow_new_interface=True,
200 + use_monolithic=False,
201 + ),
202 + experts,
203 + inplace=False,
204 + )
205 +
206 + score = torch.randn((m, e), device="cuda", dtype=dtype)
207 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
208 +
209 + sm12x_output = kernel.apply(
210 + hidden_states=a,
211 + w1=w1_q,
212 + w2=w2_q,
213 + topk_weights=topk_weights,
214 + topk_ids=topk_ids,
215 + global_num_experts=e,
216 + activation=MoEActivation.SILU,
217 + apply_router_weight_on_input=False,
218 + expert_map=None,
219 + )
220 +
221 + # Reference: BF16 torch MoE using original [gate, up] BF16 weights.
222 + # torch_moe's SiluAndMul expects [gate, up] order, matching w1_bf16.
223 + torch_output = torch_moe(a, w1_bf16, w2_bf16, score, topk)
224 +
225 + torch.testing.assert_close(sm12x_output, torch_output, atol=2e-1, rtol=2e-1)
226 +
227 +
228 +if __name__ == "__main__":
229 + test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16)
tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+7 -3
@@ -13,6 +13,7 @@ from vllm import _custom_ops as ops
13 from vllm.platforms import current_platform
14 from vllm.utils.flashinfer import (
15 flashinfer_scaled_fp4_mm,
16 + has_flashinfer_b12x_gemm,
17 )
18 from vllm.utils.torch_utils import set_random_seed
19
@@ -74,7 +75,7 @@ def get_ref_results(
75 @pytest.mark.parametrize("shape", SHAPES)
76 @pytest.mark.parametrize("seed", SEEDS)
77 @pytest.mark.parametrize("device", CUDA_DEVICES)
77 -@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm"])
78 +@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm", "b12x"])
79 @pytest.mark.parametrize("autotune", [False, True])
80 @torch.inference_mode()
81 def test_flashinfer_nvfp4_gemm(
@@ -87,6 +88,10 @@ def test_flashinfer_nvfp4_gemm(
88 ) -> None:
89 if "trtllm" in backend and dtype == torch.float16:
90 pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")
91 + if backend == "b12x" and not current_platform.has_device_capability(120):
92 + pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)")
93 + if backend == "b12x" and not has_flashinfer_b12x_gemm():
94 + pytest.skip("b12x FP4 GEMM backend not available in installed FlashInfer")
95
96 set_random_seed(seed)
97 m, n, packed_k = shape
@@ -105,8 +110,7 @@ def test_flashinfer_nvfp4_gemm(
110
111 # ops.scaled_fp4_quant returns swizzled scales, while weights
112 # from checkpoints are in linear scales.
108 - # So instead of needing to swizzle for cutlass as in modelopt.py,
109 - # we need to unswizzle for trtllm here.
113 + # cutlass and b12x use swizzled scales directly; trtllm needs them unswizzled.
114 a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
115 a_dtype, a_global_scale, is_sf_swizzled_layout=True, backend=backend
116 )
vllm/config/kernel.py
+3
@@ -128,6 +128,7 @@ MoEBackend = Literal[
128 "flashinfer_trtllm",
129 "flashinfer_cutlass",
130 "flashinfer_cutedsl",
131 + "flashinfer_b12x",
132 "marlin",
133 "humming",
134 "triton_unfused",
@@ -178,6 +179,8 @@ class KernelConfig:
179 - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels
180 - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
181 - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)
182 + - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x
183 + (RTX Pro 6000 / DGX Spark)
184 - "marlin": Use Marlin kernels (weight-only quantization)
185 - "humming": Use Humming Mixed Precision kernels
186 - "triton_unfused": Use Triton unfused MoE kernels
vllm/envs.py
+1
@@ -1540,6 +1540,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
1540 "VLLM_NVFP4_GEMM_BACKEND",
1541 None,
1542 [
1543 + "flashinfer-b12x",
1544 "flashinfer-cudnn",
1545 "flashinfer-trtllm",
1546 "flashinfer-cutlass",
vllm/model_executor/kernels/linear/__init__.py
+6
@@ -102,6 +102,7 @@ from vllm.model_executor.kernels.linear.nvfp4.fbgemm import (
102 FbgemmNvFp4LinearKernel,
103 )
104 from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
105 + FlashInferB12xNvFp4LinearKernel,
106 FlashInferCudnnNvFp4LinearKernel,
107 FlashInferCutlassNvFp4LinearKernel,
108 FlashInferTrtllmNvFp4LinearKernel,
@@ -367,6 +368,9 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
368
369 _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
370 PlatformEnum.CUDA: [
371 + # FlashInferB12xNvFp4LinearKernel excluded from auto-selection until
372 + # upstream CUTLASS SM121 MMA op guard is resolved; use
373 + # VLLM_NVFP4_GEMM_BACKEND=flashinfer-b12x to opt in explicitly.
374 FlashInferCutlassNvFp4LinearKernel,
375 CutlassNvFp4LinearKernel,
376 MarlinNvFp4LinearKernel,
@@ -812,6 +816,7 @@ def init_wfp8_a16_linear_kernel(
816
817 # Maps VLLM_NVFP4_GEMM_BACKEND env var values to kernel classes.
818 _NVFP4_BACKEND_TO_KERNEL: dict[str, type[NvFp4LinearKernel]] = {
819 + "flashinfer-b12x": FlashInferB12xNvFp4LinearKernel,
820 "flashinfer-cutlass": FlashInferCutlassNvFp4LinearKernel,
821 "cutlass": CutlassNvFp4LinearKernel,
822 "marlin": MarlinNvFp4LinearKernel,
@@ -1041,6 +1046,7 @@ __all__ = [
1046 "CutlassNvFp4LinearKernel",
1047 "EmulationNvFp4LinearKernel",
1048 "FbgemmNvFp4LinearKernel",
1049 + "FlashInferB12xNvFp4LinearKernel",
1050 "FlashInferCutlassNvFp4LinearKernel",
1051 "FlashInferTrtllmNvFp4LinearKernel",
1052 "FlashInferCudnnNvFp4LinearKernel",
vllm/model_executor/kernels/linear/nvfp4/flashinfer.py
+72 -1
@@ -10,7 +10,11 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
10 swizzle_blockscale,
11 )
12 from vllm.platforms import current_platform
13 -from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer
13 +from vllm.utils.flashinfer import (
14 + flashinfer_scaled_fp4_mm,
15 + has_flashinfer,
16 + has_flashinfer_b12x_gemm,
17 +)
18
19 from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
20
@@ -211,3 +215,70 @@ class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel):
215 if bias is not None:
216 out = out + bias
217 return out.view(*output_shape)
218 +
219 +
220 +class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel):
221 + """NVFP4 GEMM via FlashInfer's b12x CuTe DSL warp-level MMA kernel (SM120+)."""
222 +
223 + @classmethod
224 + def is_supported(
225 + cls, compute_capability: int | None = None
226 + ) -> tuple[bool, str | None]:
227 + if current_platform.has_device_capability(120) and has_flashinfer_b12x_gemm():
228 + return True, None
229 + return (
230 + False,
231 + "FlashInfer b12x requires SM120+ and FlashInfer "
232 + "with Sm120BlockScaledDenseGemmKernel",
233 + )
234 +
235 + @classmethod
236 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
237 + return True, None
238 +
239 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
240 + layer.weight_scale = torch.nn.Parameter(
241 + swizzle_blockscale(layer.weight_scale.data), requires_grad=False
242 + )
243 + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
244 + layer.weight.data
245 + )
246 + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
247 + layer.weights_padding_cols = weights_padding_cols
248 +
249 + def apply_weights(
250 + self,
251 + layer: torch.nn.Module,
252 + x: torch.Tensor,
253 + bias: torch.Tensor | None = None,
254 + ) -> torch.Tensor:
255 + output_size = layer.output_size_per_partition
256 + output_dtype = x.dtype
257 + output_shape = [*x.shape[:-1], output_size]
258 +
259 + x_fp4, x_blockscale = scaled_fp4_quant(
260 + x,
261 + layer.input_global_scale_inv,
262 + is_sf_swizzled_layout=True,
263 + backend="b12x",
264 + )
265 +
266 + x_fp4 = pad_nvfp4_activation_for_cutlass(
267 + x_fp4, getattr(layer, "weights_padding_cols", 0)
268 + )
269 +
270 + out = flashinfer_scaled_fp4_mm(
271 + x_fp4,
272 + layer.weight,
273 + x_blockscale,
274 + layer.weight_scale,
275 + layer.alpha,
276 + output_dtype,
277 + backend="b12x",
278 + )
279 +
280 + out = slice_nvfp4_output(out, output_size)
281 +
282 + if bias is not None:
283 + out = out + bias
284 + return out.view(*output_shape)
vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py new
+223
@@ -0,0 +1,223 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8 +from vllm.model_executor.layers.fused_moe.config import (
9 + FusedMoEConfig,
10 + FusedMoEParallelConfig,
11 + FusedMoEQuantConfig,
12 +)
13 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14 + TopKWeightAndReduceNoOP,
15 +)
16 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 + QuantKey,
18 + kNvfp4Dynamic,
19 + kNvfp4Static,
20 +)
21 +from vllm.platforms import current_platform
22 +from vllm.utils.flashinfer import (
23 + flashinfer_b12x_fused_moe,
24 + flashinfer_convert_sf_to_mma_layout,
25 + has_flashinfer_b12x_moe,
26 +)
27 +
28 +
29 +class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
30 + """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121,
31 + RTX Pro 6000 / DGX Spark).
32 +
33 + Uses ``b12x_fused_moe`` from FlashInfer PR #3080 which fuses token
34 + dispatch, two GEMMs, SwiGLU activation, and topk-weight reduction into a
35 + single kernel call. Input quantization (BF16→FP4) is performed inside the
36 + kernel so BF16 hidden states are passed directly.
37 +
38 + Weight scale factors are converted to the MMA layout produced by
39 + ``convert_sf_to_mma_layout`` once during ``process_weights_after_loading``
40 + and cached as ``w1_sf_mma`` / ``w2_sf_mma``.
41 +
42 + Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported.
43 + """
44 +
45 + def __init__(
46 + self,
47 + moe_config: FusedMoEConfig,
48 + quant_config: FusedMoEQuantConfig,
49 + ):
50 + super().__init__(moe_config=moe_config, quant_config=quant_config)
51 + assert quant_config.quant_dtype == "nvfp4", (
52 + "FlashInferB12xExperts only supports nvfp4 quantization."
53 + )
54 + self.out_dtype = moe_config.in_dtype
55 + self.num_local_experts = moe_config.num_local_experts
56 + self.ep_rank = moe_config.moe_parallel_config.ep_rank
57 +
58 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
59 + # Normalise block scales to absorb the per-expert weight global scale
60 + # (w_gs). vLLM's NVFP4 convention stores:
61 + # block_scale = max_abs * w_gs / fp4_max, g1_alphas = 1/w_gs
62 + # The SM12x kernel treats w1_alpha (= g1_alphas) as a per-expert weight
63 + # dequant multiplier separate from input_gs (activation scale). We bake
64 + # w_gs into the block scales so that w1_alpha = 1.0 and the kernel sees
65 + # the simpler form:
66 + # block_scale = max_abs / fp4_max, w1_alpha = 1.0
67 + # The FP4-packed values and dequantised results are identical in both
68 + # representations. We set scale_2 = 1.0 to signal that the bake-in is
69 + # already done.
70 + layer.w13_weight_scale.data = (
71 + layer.w13_weight_scale.float() * layer.w13_weight_scale_2.view(-1, 1, 1)
72 + ).to(layer.w13_weight_scale.dtype)
73 + layer.w13_weight_scale_2.data.fill_(1.0)
74 +
75 + layer.w2_weight_scale.data = (
76 + layer.w2_weight_scale.float() * layer.w2_weight_scale_2.view(-1, 1, 1)
77 + ).to(layer.w2_weight_scale.dtype)
78 + layer.w2_weight_scale_2.data.fill_(1.0)
79 +
80 + # The SM12x kernel uses dynamic per-block quantization for FC2 input
81 + # activations (the SwiGLU output before the down projection). The
82 + # calibrated a2_gscale from the modelopt checkpoint (~tens to hundreds)
83 + # is intended for static-quantisation backends (TRTLLM/CUTLASS) and
84 + # causes every intermediate activation to saturate at max FP4 when
85 + # multiplied by values that large. Force to 1.0 so the kernel uses
86 + # its own per-block dynamic scale.
87 + if self.a2_gscale is not None:
88 + self.a2_gscale.fill_(1.0)
89 +
90 + # Precompute MMA-layout views of the weight scale factors once here
91 + # rather than recomputing on every forward pass.
92 + assert self.w1_scale is not None
93 + num_experts_w1, m1, k1_sf = self.w1_scale.shape
94 + k1 = k1_sf * 16
95 + self.w1_sf_mma = flashinfer_convert_sf_to_mma_layout(
96 + self.w1_scale.reshape(num_experts_w1 * m1, k1_sf),
97 + m=m1,
98 + k=k1,
99 + num_groups=num_experts_w1,
100 + )
101 +
102 + assert self.w2_scale is not None
103 + num_experts_w2, m2, k2_sf = self.w2_scale.shape
104 + k2 = k2_sf * 16
105 + self.w2_sf_mma = flashinfer_convert_sf_to_mma_layout(
106 + self.w2_scale.reshape(num_experts_w2 * m2, k2_sf),
107 + m=m2,
108 + k=k2,
109 + num_groups=num_experts_w2,
110 + )
111 +
112 + @staticmethod
113 + def activation_format() -> mk.FusedMoEActivationFormat:
114 + return mk.FusedMoEActivationFormat.Standard
115 +
116 + @staticmethod
117 + def _supports_current_device() -> bool:
118 + p = current_platform
119 + return (
120 + p.is_cuda()
121 + and p.is_device_capability_family(120)
122 + and has_flashinfer_b12x_moe()
123 + )
124 +
125 + @staticmethod
126 + def _supports_no_act_and_mul() -> bool:
127 + return False
128 +
129 + @staticmethod
130 + def _supports_quant_scheme(
131 + weight_key: QuantKey | None,
132 + activation_key: QuantKey | None,
133 + ) -> bool:
134 + return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)
135 +
136 + @staticmethod
137 + def _supports_activation(activation: MoEActivation) -> bool:
138 + return activation == MoEActivation.SILU
139 +
140 + @staticmethod
141 + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
142 + return True
143 +
144 + def supports_expert_map(self) -> bool:
145 + return False
146 +
147 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
148 + # b12x_fused_moe applies topk weights internally.
149 + return TopKWeightAndReduceNoOP()
150 +
151 + def workspace_shapes(
152 + self,
153 + M: int,
154 + N: int,
155 + K: int,
156 + topk: int,
157 + global_num_experts: int,
158 + local_num_experts: int,
159 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
160 + activation: MoEActivation,
161 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
162 + # b12x_fused_moe manages its own internal workspace.
163 + workspace1 = (1,)
164 + workspace2 = (0,)
165 + output_shape = (M, K)
166 + return (workspace1, workspace2, output_shape)
167 +
168 + @property
169 + def expects_unquantized_inputs(self) -> bool:
170 + # b12x_fused_moe expects BF16 hidden states and performs its own FP4
171 + # quantization internally. Returning True prevents the modular kernel
172 + # from pre-quantizing activations, which would produce an FP4-packed
173 + # tensor with size(-1)=k//2 and break the scale-factor conversion that
174 + # expects size(-1)=k.
175 + return True
176 +
177 + def apply(
178 + self,
179 + output: torch.Tensor,
180 + hidden_states: torch.Tensor,
181 + w1: torch.Tensor,
182 + w2: torch.Tensor,
183 + topk_weights: torch.Tensor,
184 + topk_ids: torch.Tensor,
185 + activation: MoEActivation,
186 + global_num_experts: int,
187 + expert_map: torch.Tensor | None,
188 + a1q_scale: torch.Tensor | None,
189 + a2_scale: torch.Tensor | None,
190 + workspace13: torch.Tensor | None,
191 + workspace2: torch.Tensor | None,
192 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
193 + apply_router_weight_on_input: bool | None,
194 + ):
195 + assert self.w1_scale is not None and self.w2_scale is not None, (
196 + "w1_scale and w2_scale must not be None for FlashInferB12xExperts"
197 + )
198 + assert self.g1_alphas is not None and self.g2_alphas is not None, (
199 + "g1_alphas and g2_alphas must not be None for FlashInferB12xExperts"
200 + )
201 + assert self.a2_gscale is not None, (
202 + "a2_gscale must not be None for FlashInferB12xExperts"
203 + )
204 +
205 + top_k = topk_ids.shape[1]
206 +
207 + flashinfer_b12x_fused_moe(
208 + x=hidden_states,
209 + token_selected_experts=topk_ids.to(torch.int32),
210 + token_final_scales=topk_weights,
211 + w1_weight=w1,
212 + w1_weight_sf=self.w1_sf_mma,
213 + w1_alpha=self.g1_alphas,
214 + fc2_input_scale=self.a2_gscale,
215 + w2_weight=w2,
216 + w2_weight_sf=self.w2_sf_mma,
217 + w2_alpha=self.g2_alphas,
218 + num_experts=global_num_experts,
219 + top_k=top_k,
220 + num_local_experts=self.num_local_experts,
221 + output_dtype=self.out_dtype,
222 + output=output,
223 + )
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+13
@@ -43,6 +43,7 @@ class NvFp4MoeBackend(Enum):
43 FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS"
44 FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL"
45 FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED"
46 + FLASHINFER_B12X = "FLASHINFER_B12X"
47 VLLM_CUTLASS = "VLLM_CUTLASS"
48 MARLIN = "MARLIN"
49 EMULATION = "EMULATION"
@@ -53,6 +54,7 @@ FLASHINFER_NVFP4_MOE_BACKENDS = [
54 NvFp4MoeBackend.FLASHINFER_CUTLASS,
55 NvFp4MoeBackend.FLASHINFER_CUTEDSL,
56 NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
57 + NvFp4MoeBackend.FLASHINFER_B12X,
58 ]
59
60 fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = {
@@ -106,6 +108,13 @@ def backend_to_kernel_cls(
108
109 return [FlashInferCuteDSLBatchedExperts]
110
111 + elif backend == NvFp4MoeBackend.FLASHINFER_B12X:
112 + from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import ( # noqa: E501
113 + FlashInferB12xExperts,
114 + )
115 +
116 + return [FlashInferB12xExperts]
117 +
118 elif backend == NvFp4MoeBackend.VLLM_CUTLASS:
119 from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
120 CutlassExpertsFp4,
@@ -136,6 +145,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend:
145 "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM,
146 "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS,
147 "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL,
148 + "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X,
149 "marlin": NvFp4MoeBackend.MARLIN,
150 "emulation": NvFp4MoeBackend.EMULATION,
151 }
@@ -158,6 +168,9 @@ def select_nvfp4_moe_backend(
168 """
169
170 # NOTE: the kernels are selected in the following order.
171 + # FLASHINFER_B12X is intentionally excluded from auto-selection until
172 + # the upstream CUTLASS SM121 MMA op guard is resolved; use
173 + # moe_backend="flashinfer_b12x" to opt in explicitly.
174 AVAILABLE_BACKENDS = [
175 NvFp4MoeBackend.FLASHINFER_TRTLLM,
176 NvFp4MoeBackend.FLASHINFER_CUTEDSL,
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+2
@@ -317,6 +317,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
317 NvFp4MoeBackend.FLASHINFER_CUTLASS,
318 NvFp4MoeBackend.FLASHINFER_TRTLLM,
319 NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
320 + NvFp4MoeBackend.FLASHINFER_B12X,
321 ]
322
323 # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels.
@@ -328,6 +329,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
329 in [
330 NvFp4MoeBackend.FLASHINFER_CUTLASS,
331 NvFp4MoeBackend.FLASHINFER_TRTLLM,
332 + NvFp4MoeBackend.FLASHINFER_B12X,
333 ]
334 ):
335 w13, w13_scale = reorder_w1w3_to_w3w1(w13, w13_scale)
vllm/utils/flashinfer.py
+39
@@ -135,6 +135,9 @@ flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper(
135 flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper(
136 "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout"
137 )
138 +flashinfer_b12x_fused_moe = _lazy_import_wrapper(
139 + "flashinfer.fused_moe", "b12x_fused_moe"
140 +)
141 trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
142 "flashinfer", "trtllm_fp4_block_scale_moe"
143 )
@@ -267,6 +270,39 @@ def has_flashinfer_cutedsl_moe_nvfp4() -> bool:
270 return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4")
271
272
273 +@functools.cache
274 +def has_flashinfer_b12x_gemm() -> bool:
275 + """Return True if FlashInfer b12x FP4 GEMM backend is available (SM120+)."""
276 + if not has_flashinfer_cutedsl():
277 + return False
278 + mod = _get_submodule("flashinfer.gemm")
279 + if mod is None:
280 + return False
281 + # FlashInfer 0.6.11 renamed Sm120BlockScaledDenseGemmKernel ->
282 + # Sm120B12xBlockScaledDenseGemmKernel (commit 223f2a49). Accept either.
283 + return hasattr(mod, "Sm120B12xBlockScaledDenseGemmKernel") or hasattr(
284 + mod, "Sm120BlockScaledDenseGemmKernel"
285 + )
286 +
287 +
288 +@functools.cache
289 +def has_flashinfer_b12x_moe() -> bool:
290 + """Return ``True`` if FlashInfer CuteDSL SM12x fused MoE is available."""
291 + if not has_flashinfer_moe():
292 + return False
293 +
294 + required_functions = [
295 + ("flashinfer.fused_moe", "b12x_fused_moe"),
296 + ("flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout"),
297 + ]
298 +
299 + for module_name, attr_name in required_functions:
300 + mod = _get_submodule(module_name)
301 + if not mod or not hasattr(mod, attr_name):
302 + return False
303 + return True
304 +
305 +
306 @functools.cache
307 def has_nvidia_artifactory() -> bool:
308 """Return `True` if NVIDIA's artifactory is accessible.
@@ -926,6 +962,7 @@ __all__ = [
962 "scaled_fp4_grouped_quantize",
963 "nvfp4_block_scale_interleave",
964 "flashinfer_cute_dsl_fused_moe_nvfp4",
965 + "flashinfer_b12x_fused_moe",
966 "flashinfer_convert_sf_to_mma_layout",
967 "trtllm_fp4_block_scale_moe",
968 "autotune",
@@ -936,6 +973,8 @@ __all__ = [
973 "has_flashinfer_cutlass_fused_moe",
974 "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
975 "has_flashinfer_cutedsl_moe_nvfp4",
976 + "has_flashinfer_b12x_moe",
977 + "has_flashinfer_b12x_gemm",
978 "has_flashinfer_fp8_blockscale_gemm",
979 "has_nvidia_artifactory",
980 "supports_trtllm_attention",