From 269beb47ddc364239f9456df3d310e379320796a Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Fri, 17 Apr 2026 22:59:46 +0000 Subject: [PATCH 01/15] [SM120/121] Add FlashInfer b12x fused MoE backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds FlashInferCuteDSLSM12xExperts targeting SM120/SM121 (RTX Pro 6000 / DGX Spark) using cute_dsl_fused_moe_nvfp4 from FlashInfer PRs #3051 and #3066. The kernel fuses token dispatch, W1 GEMM, SwiGLU, and W2 GEMM into a single call; BF16 hidden states are passed directly as activation quantization is fused internally. - vllm/utils/flashinfer.py: lazy import wrappers for cute_dsl_fused_moe_nvfp4 and convert_sf_to_mma_layout; adds has_flashinfer_cutedsl_sm12x_moe() availability probe - experts/flashinfer_cutedsl_moe.py: FlashInferCuteDSLSM12xExperts with TODO to adopt plan/run() API from PR #3066 - oracle/nvfp4.py: FLASHINFER_CUTEDSL_SM12X backend enum and routing; falls back to FLASHINFER_CUTLASS on SM12x when PRs are absent - flashinfer_fp4_moe.py: SM12X added to FI weight-prep path and w1/w3 → w3/w1 reorder list - tests/kernels/moe/test_cutedsl_sm12x_moe.py: correctness tests vs BF16 torch reference; module-level skip when SM120 hw or FlashInfer PRs are absent Signed-off-by: Meenakshi Venkataraman Co-Authored-By: Claude Sonnet 4.6 --- tests/kernels/moe/test_cutedsl_sm12x_moe.py | 208 ++++++++++ vllm/config/kernel.py | 2 + .../experts/flashinfer_cutedsl_moe.py | 385 ++++++++++++++++++ .../layers/fused_moe/oracle/nvfp4.py | 11 + .../quantization/utils/flashinfer_fp4_moe.py | 3 + vllm/utils/flashinfer.py | 31 +- 6 files changed, 634 insertions(+), 6 deletions(-) create mode 100644 tests/kernels/moe/test_cutedsl_sm12x_moe.py diff --git a/tests/kernels/moe/test_cutedsl_sm12x_moe.py b/tests/kernels/moe/test_cutedsl_sm12x_moe.py new file mode 100644 index 000000000000..59cc8e8a12cb --- /dev/null +++ b/tests/kernels/moe/test_cutedsl_sm12x_moe.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_device_capability_family(120): + pytest.skip( + reason="FlashInfer CuteDSL SM12x MoE requires SM120 (RTX Pro 6000 / DGX Spark).", + allow_module_level=True, + ) + +from vllm.utils.flashinfer import has_flashinfer_cutedsl_sm12x_moe + +if not has_flashinfer_cutedsl_sm12x_moe(): + pytest.skip( + reason=( + "FlashInfer cute_dsl_fused_moe_nvfp4 / convert_sf_to_mma_layout " + "not available in installed FlashInfer (needs PRs #3051 and #3066)." + ), + allow_module_level=True, + ) + +# Import fp4_quantize after the skip guard — FlashInfer must be installed. +from flashinfer.fp4_quantization import fp4_quantize + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from tests.kernels.moe.utils import make_dummy_moe_config +from tests.kernels.utils import torch_moe +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe import fused_topk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) +from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config +from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( + FlashInferCuteDSLSM12xExperts, +) +from vllm.utils.torch_utils import set_random_seed + +# Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256, +# n multiple of 128) while keeping tests fast. +MNK_FACTORS = [ + (2, 128, 256), + (2, 256, 512), + (16, 128, 256), + (64, 256, 512), +] + + +def _reorder_gate_up_to_up_gate( + w: torch.Tensor, + w_s: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Swap gate and up-projection halves along dim=1 to [up, gate] order. + + The SM12x kernel expects weights in [up (w3), gate (w1)] order while the + BF16 reference uses [gate (w1), up (w3)]. This replicates the reordering + done at model-load time by ``prepare_nvfp4_moe_layer_for_fi_or_cutlass``. + """ + n = w.shape[1] // 2 + return ( + torch.cat([w[:, n:, :], w[:, :n, :]], dim=1), + torch.cat([w_s[:, n:, :], w_s[:, :n, :]], dim=1), + ) + + +@pytest.mark.parametrize("m,n,k", MNK_FACTORS) +@pytest.mark.parametrize("e", [8, 16]) +@pytest.mark.parametrize("topk", [1, 2, 4]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.inference_mode() +def test_flashinfer_cutedsl_sm12x_moe( + m: int, + n: int, + k: int, + e: int, + topk: int, + dtype: torch.dtype, + workspace_init, +): + """Test FlashInferCuteDSLSM12xExperts against a BF16 torch reference. + + The SM12x kernel takes BF16 hidden states directly and fuses token + dispatch, W1 GEMM, SwiGLU, and W2 GEMM into one call. We verify + correctness against ``torch_moe`` using generous tolerances to account + for the internal FP4 quantization of activations and weights. + + Scale convention + ---------------- + The SM12x kernel uses ``w1_alpha`` as *both* the activation-quantisation + global scale and the weight dequantisation factor. These two roles are + conflated into a single parameter in ``launch_sm120_moe``, so they must + equal the same value. We use ``global_scale = 1.0`` for + ``fp4_quantize`` so that ``w1_alpha = ones`` satisfies both roles + simultaneously. The alternative — vLLM's convention of baking a large + ``w_gs`` into block-scale values and compensating with + ``g1_alphas = 1/w_gs`` — is incompatible with this kernel. + """ + set_random_seed(7) + with set_current_vllm_config( + VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) + ): + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10 + + # Generate BF16 reference weights in [gate, up] order. + # Shape: w1=(e, 2n, k), w2=(e, k, n). + w1_bf16 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 15 + w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15 + + # ------------------------------------------------------------------ # + # Quantise weights for the SM12x kernel using FlashInfer's convention: + # global_scale = 1.0 → block_scale = max_abs_block / fp4_max + # w1_alpha = 1.0 (no extra global factor to compensate) + # + # The scale factors returned by fp4_quantize(..., is_sf_swizzled_layout=True) + # are already in the swizzled 2D layout expected by convert_sf_to_mma_layout. + # No additional swizzle_blockscale() call is needed. + # ------------------------------------------------------------------ # + gs = torch.ones(1, device="cuda", dtype=torch.float32) + sf_vec_size = 16 + + # W1: reorder BF16 from [gate, up] → [up, gate], then quantise. + w1_reordered = torch.cat( + [w1_bf16[:, n:, :], w1_bf16[:, :n, :]], dim=1 + ) # shape (e, 2n, k), [up, gate] + w1_flat = w1_reordered.reshape(e * 2 * n, k) + w1_q_flat, w1_sf_flat = fp4_quantize( + w1_flat, + global_scale=gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w1_q = w1_q_flat.view(e, 2 * n, k // 2) # uint8, packed FP4 + w1_blockscale = w1_sf_flat.view(e, 2 * n, w1_sf_flat.shape[1]) # float8 + + # W2: no row reordering needed for the down-projection. + w2_flat = w2_bf16.reshape(e * k, n) + w2_q_flat, w2_sf_flat = fp4_quantize( + w2_flat, + global_scale=gs, + sf_vec_size=sf_vec_size, + is_sf_swizzled_layout=True, + ) + w2_q = w2_q_flat.view(e, k, n // 2) # uint8, packed FP4 + w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1]) # float8 + + # All per-expert alphas are 1.0 (global_scale = 1.0, no compensation). + ones_e = torch.ones(e, device="cuda", dtype=torch.float32) + + quant_config = nvfp4_moe_quant_config( + g1_alphas=ones_e, + g2_alphas=ones_e, + a1_gscale=ones_e, + a2_gscale=ones_e, + w1_scale=w1_blockscale, + w2_scale=w2_blockscale, + ) + + moe_config = make_dummy_moe_config( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size_per_partition=n, + in_dtype=dtype, + ) + + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + FlashInferCuteDSLSM12xExperts( + moe_config=moe_config, + quant_config=quant_config, + ), + inplace=False, + ) + + score = torch.randn((m, e), device="cuda", dtype=dtype) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + + sm12x_output = kernel.apply( + hidden_states=a, + w1=w1_q, + w2=w2_q, + topk_weights=topk_weights, + topk_ids=topk_ids, + global_num_experts=e, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + expert_map=None, + ) + + # Reference: BF16 torch MoE using original [gate, up] BF16 weights. + # torch_moe's SiluAndMul expects [gate, up] order, matching w1_bf16. + torch_output = torch_moe(a, w1_bf16, w2_bf16, score, topk) + + torch.testing.assert_close(sm12x_output, torch_output, atol=2e-1, rtol=2e-1) + + +if __name__ == "__main__": + test_flashinfer_cutedsl_sm12x_moe(16, 128, 256, 8, 2, torch.bfloat16) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 66b805383fcb..9531517f0eb4 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -128,6 +128,7 @@ def with_default( "flashinfer_trtllm", "flashinfer_cutlass", "flashinfer_cutedsl", + "flashinfer_cutedsl_sm12x", "marlin", "humming", "triton_unfused", @@ -178,6 +179,7 @@ class KernelConfig: - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) + - "flashinfer_cutedsl_sm12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) - "marlin": Use Marlin kernels (weight-only quantization) - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 2310982792fe..ff68c0bfec98 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -11,6 +11,7 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceDelegate, TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -20,8 +21,14 @@ ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( + flashinfer_convert_sf_to_mma_layout, flashinfer_cute_dsl_fused_moe_nvfp4, + flashinfer_cutedsl_grouped_gemm_nt_masked, + has_flashinfer_cutedsl_grouped_gemm_nt_masked, has_flashinfer_cutedsl_moe_nvfp4, + has_flashinfer_cutedsl_sm12x_moe, + scaled_fp4_grouped_quantize, + silu_and_mul_scaled_nvfp4_experts_quantize, ) @@ -167,3 +174,381 @@ def apply( local_expert_offset=self.local_expert_offset, moe_output=output, ) + + +def get_cute_dtype(input: torch.Tensor) -> str: + if input.dtype == torch.bfloat16: + return "bfloat16" + elif input.dtype == torch.float16: + return "float16" + elif input.dtype == torch.float32: + return "float32" + else: + raise ValueError(f"Unsupported cute dtype {input.dtype}") + + +def flashinfer_cutedsl_moe_masked( + hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + input_global_scale: torch.Tensor, + w1: torch.Tensor, + w1_blockscale: torch.Tensor, + w1_alpha, + w2: torch.Tensor, + a2_global_scale: torch.Tensor, + w2_blockscale: torch.Tensor, + w2_alpha, + masked_m: torch.Tensor, + workspace: torch.Tensor, + out: torch.Tensor, +): + """ + Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL + kernels. + + Args: + hidden_states: Either of the following case + * torch.Tensor: [num_experts, m, k], bf16 + * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2], + uint8, [num_experts, m, k // 16], float8_e4m3fn + input_global_scale (torch.Tensor): (l,) + w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8 + w1_blockscale (torch.Tensor): blockscale factors, e4m3, + w1_alpha (torch.Tensor): (l,) + w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8 + a2_global_scale (torch.Tensor): (l,) + w2_blockscale (torch.Tensor): blockscale factors, e4m3, + w2_alpha (torch.Tensor): (l,) + masked_m (torch.Tensor): Masked dimension indices + workspace (torch.Tensor): For gateup_output + + Notes: + - Assumes max(masked_m) <= m. + """ + + # === Assertions on dtypes === + assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}" + assert w1_blockscale.dtype == torch.float8_e4m3fn, ( + f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}" + ) + assert w1_alpha.dtype == torch.float32, ( + f"w1_alpha must be float32, got {w1_alpha.dtype}" + ) + assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}" + assert a2_global_scale.dtype == torch.float32, ( + f"a2_global_scale must be float32, got {a2_global_scale.dtype}" + ) + assert w2_blockscale.dtype == torch.float8_e4m3fn, ( + f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}" + ) + assert w2_alpha.dtype == torch.float32, ( + f"w2_alpha must be float32, got {w2_alpha.dtype}" + ) + + # === Assertions on shapes === + n = w2.shape[-1] * 2 # intermediate dimension + if isinstance(hidden_states, tuple): + assert input_global_scale is None, ( + "input_global_scale is needed when input needs quant" + ) + + aq = hidden_states[0].view(torch.uint8) + aq_sf = hidden_states[1].view(torch.float8_e4m3fn) + # m, k_by_2, num_experts = aq.shape + num_experts, m, k_by_2 = aq.shape + k = k_by_2 * 2 + aq = aq.permute(1, 2, 0) + else: + num_experts, m, k = hidden_states.shape + + assert input_global_scale.dtype == torch.float32, ( + f"input_global_scale must be float32, got {input_global_scale.dtype}" + ) + assert input_global_scale.shape == (num_experts,), ( + f"input_global_scale must be (l,), got {input_global_scale.shape}" + ) + + aq, aq_sf = scaled_fp4_grouped_quantize( + hidden_states, + masked_m, + input_global_scale, + ) + + assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}" + assert w1.shape[-1] * 2 == k, ( + f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}" + ) + assert w2.shape[-2:] == ( + k, + n // 2, + ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}" + + assert w1_alpha.shape == (num_experts,), ( + f"w1_alpha must be (l,), got {w1_alpha.shape}" + ) + assert a2_global_scale.shape == (num_experts,), ( + f"a2_global_scale must be (l,), got {a2_global_scale.shape}" + ) + assert w2_alpha.shape == (num_experts,), ( + f"w2_alpha must be (l,), got {w2_alpha.shape}" + ) + + workspace = workspace.permute(1, 2, 0) # requirement of kernel + sf_vec_size = 16 + assert aq_sf.dtype == torch.float8_e4m3fn + assert aq.dtype == torch.uint8 + ab_dtype = "float4_e2m1fn" + sf_dtype = "float8_e4m3fn" + + if isinstance(hidden_states, tuple): + c_dtype = "bfloat16" + else: + c_dtype = get_cute_dtype(hidden_states) + + # Gemm1 + flashinfer_cutedsl_grouped_gemm_nt_masked( + (aq, aq_sf), + (w1.permute(1, 2, 0), w1_blockscale), + workspace, + masked_m, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + alpha=w1_alpha.view(1, 1, num_experts), + alpha_dtype=get_cute_dtype(w1_alpha), + ) # in logical [m, n, l] + + # SILU and quantization + diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize( + workspace.permute(2, 0, 1), + masked_m, + a2_global_scale, + ) + + # Gemm2 + out = out.permute(1, 2, 0) # requirement of kernel + flashinfer_cutedsl_grouped_gemm_nt_masked( + (diq, diq_sf), + (w2.permute(1, 2, 0), w2_blockscale), + out, + masked_m, + ab_dtype=ab_dtype, + sf_dtype=sf_dtype, + c_dtype=c_dtype, + sf_vec_size=sf_vec_size, + alpha=w2_alpha.view(1, 1, num_experts), + alpha_dtype=get_cute_dtype(w2_alpha), + ) # in logical [m, k, l] + out = out.permute(2, 0, 1) + + +class FlashInferCuteDSLSM12xExperts(mk.FusedMoEExpertsModular): + """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121, RTX Pro 6000 / DGX Spark). + + Uses ``cute_dsl_fused_moe_nvfp4`` from FlashInfer PR #3066 which fuses + token dispatch, two GEMMs, SwiGLU activation, and topk-weight reduction + into a single kernel call. Input quantization (BF16→FP4) is performed + inside the kernel so BF16 hidden states are passed directly. + + Weight scale factors must be in the MMA layout produced by + ``convert_sf_to_mma_layout``; this conversion is performed on every + forward pass since it is a zero-copy strided view. + + Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config=moe_config, quant_config=quant_config) + assert quant_config.quant_dtype == "nvfp4", ( + "FlashInferCuteDSLSM12xExperts only supports nvfp4 quantization." + ) + self.out_dtype = moe_config.in_dtype + self.num_local_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Normalise block scales to absorb the per-expert weight global scale + # (w_gs). vLLM's NVFP4 convention stores: + # block_scale = max_abs * w_gs / fp4_max, g1_alphas = 1/w_gs + # The SM12x kernel treats w1_alpha (= g1_alphas) as a per-expert weight + # dequant multiplier separate from input_gs (activation scale). We bake + # w_gs into the block scales so that w1_alpha = 1.0 and the kernel sees + # the simpler form: + # block_scale = max_abs / fp4_max, w1_alpha = 1.0 + # The FP4-packed values and dequantised results are identical in both + # representations. We set scale_2 = 1.0 to signal that the bake-in is + # already done. + layer.w13_weight_scale.data = ( + layer.w13_weight_scale.float() + * layer.w13_weight_scale_2.view(-1, 1, 1) + ).to(layer.w13_weight_scale.dtype) + layer.w13_weight_scale_2.data.fill_(1.0) + + layer.w2_weight_scale.data = ( + layer.w2_weight_scale.float() + * layer.w2_weight_scale_2.view(-1, 1, 1) + ).to(layer.w2_weight_scale.dtype) + layer.w2_weight_scale_2.data.fill_(1.0) + + # The SM12x kernel uses dynamic per-block quantization for FC2 input + # activations (the SwiGLU output before the down projection). The + # calibrated a2_gscale from the modelopt checkpoint (~tens to hundreds) + # is intended for static-quantisation backends (TRTLLM/CUTLASS) and + # causes every intermediate activation to saturate at max FP4 when + # multiplied by values that large. Force to 1.0 so the kernel uses + # its own per-block dynamic scale — matching the unit-test convention. + if self.a2_gscale is not None: + self.a2_gscale.fill_(1.0) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return ( + p.is_cuda() + and p.is_device_capability_family(120) + and has_flashinfer_cutedsl_sm12x_moe() + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # cute_dsl_fused_moe_nvfp4 applies topk weights internally. + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # cute_dsl_fused_moe_nvfp4 manages its own internal workspace. + workspace1 = (1,) + workspace2 = (0,) + output_shape = (M, K) + return (workspace1, workspace2, output_shape) + + @property + def expects_unquantized_inputs(self) -> bool: + # cute_dsl_fused_moe_nvfp4 expects BF16 hidden states and performs + # its own FP4 quantization internally. Returning True prevents the + # modular kernel from pre-quantizing activations, which would produce + # an FP4-packed tensor with size(-1)=k//2 and break the scale-factor + # conversion that expects size(-1)=k. + return True + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor | None, + workspace2: torch.Tensor | None, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool | None, + ): + assert self.w1_scale is not None and self.w2_scale is not None, ( + "w1_scale and w2_scale must not be None for FlashInferCuteDSLSM12xExperts" + ) + assert self.g1_alphas is not None and self.g2_alphas is not None, ( + "g1_alphas and g2_alphas must not be None for FlashInferCuteDSLSM12xExperts" + ) + assert self.a2_gscale is not None, ( + "a2_gscale must not be None for FlashInferCuteDSLSM12xExperts" + ) + + top_k = topk_ids.shape[1] + local_expert_offset = self.ep_rank * self.num_local_experts + + # Convert swizzled scale factors to the 6D MMA layout expected by + # cute_dsl_fused_moe_nvfp4. This is a zero-copy view+permute. + num_experts_w1, m1, k1_sf = self.w1_scale.shape + k1 = k1_sf * 16 # sf_vec_size = 16 + w1_sf_mma = flashinfer_convert_sf_to_mma_layout( + self.w1_scale.reshape(num_experts_w1 * m1, k1_sf), + m=m1, + k=k1, + num_groups=num_experts_w1, + ) + + num_experts_w2, m2, k2_sf = self.w2_scale.shape + k2 = k2_sf * 16 + w2_sf_mma = flashinfer_convert_sf_to_mma_layout( + self.w2_scale.reshape(num_experts_w2 * m2, k2_sf), + m=m2, + k=k2, + num_groups=num_experts_w2, + ) + + # x_sf is ignored by the SM12x kernel (quantization is fused + # internally), but the API requires a tensor argument. + x_sf_placeholder = ( + a1q_scale + if a1q_scale is not None + else hidden_states.new_zeros(1) + ) + + # TODO: Use the plan/run() API from FlashInfer PR #3066 instead of + # calling cute_dsl_fused_moe_nvfp4 directly. The plan object can be + # created once in __init__ (shapes are fixed for MoE layers) and + # plan.run() called here, avoiding workspace allocation and kernel + # parameter setup overhead on every forward pass. + flashinfer_cute_dsl_fused_moe_nvfp4( + x=hidden_states, + x_sf=x_sf_placeholder, + token_selected_experts=topk_ids.to(torch.int32), + token_final_scales=topk_weights, + w1_weight=w1, + w1_weight_sf=w1_sf_mma, + w1_alpha=self.g1_alphas, + fc2_input_scale=self.a2_gscale, + w2_weight=w2, + w2_weight_sf=w2_sf_mma, + w2_alpha=self.g2_alphas, + num_experts=global_num_experts, + top_k=top_k, + num_local_experts=self.num_local_experts, + local_expert_offset=local_expert_offset, + output_dtype=self.out_dtype, + moe_output=output, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 902d788da3c8..28bd4d00f82f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -43,6 +43,7 @@ class NvFp4MoeBackend(Enum): FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS" FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL" FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED" + FLASHINFER_CUTEDSL_SM12X = "FLASHINFER_CUTEDSL_SM12X" VLLM_CUTLASS = "VLLM_CUTLASS" MARLIN = "MARLIN" EMULATION = "EMULATION" @@ -53,6 +54,7 @@ class NvFp4MoeBackend(Enum): NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, ] fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = { @@ -106,6 +108,13 @@ def backend_to_kernel_cls( return [FlashInferCuteDSLBatchedExperts] + elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X: + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( # noqa: E501 + FlashInferCuteDSLSM12xExperts, + ) + + return [FlashInferCuteDSLSM12xExperts] + elif backend == NvFp4MoeBackend.VLLM_CUTLASS: from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( CutlassExpertsFp4, @@ -136,6 +145,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS, "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, + "flashinfer_cutedsl_sm12x": NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, "marlin": NvFp4MoeBackend.MARLIN, "emulation": NvFp4MoeBackend.EMULATION, } @@ -162,6 +172,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 5bd08de9a610..84f7ca79cd7f 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -316,7 +316,9 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_TRTLLM, + NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, ] # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels. @@ -328,6 +330,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( in [ NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_TRTLLM, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, ] ): w13, w13_scale = reorder_w1w3_to_w3w1(w13, w13_scale) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 287de921c18a..c32b1fa1ae7d 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -118,6 +118,12 @@ def wrapper(*args, **kwargs): flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper( "flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked" ) +flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper( + "flashinfer.fused_moe", "cute_dsl_fused_moe_nvfp4" +) +flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( + "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout" +) flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize") nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize") silu_and_mul_scaled_nvfp4_experts_quantize = _lazy_import_wrapper( @@ -129,12 +135,6 @@ def wrapper(*args, **kwargs): nvfp4_block_scale_interleave = _lazy_import_wrapper( "flashinfer.fp4_quantization", "block_scale_interleave" ) -flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper( - "flashinfer", "cute_dsl_fused_moe_nvfp4" -) -flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( - "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout" -) trtllm_fp4_block_scale_moe = _lazy_import_wrapper( "flashinfer", "trtllm_fp4_block_scale_moe" ) @@ -267,6 +267,24 @@ def has_flashinfer_cutedsl_moe_nvfp4() -> bool: return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4") +@functools.cache +def has_flashinfer_cutedsl_sm12x_moe() -> bool: + """Return ``True`` if FlashInfer CuteDSL SM12x fused MoE is available.""" + if not has_flashinfer_moe(): + return False + + required_functions = [ + ("flashinfer.fused_moe", "cute_dsl_fused_moe_nvfp4"), + ("flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout"), + ] + + for module_name, attr_name in required_functions: + mod = _get_submodule(module_name) + if not mod or not hasattr(mod, attr_name): + return False + return True + + @functools.cache def has_nvidia_artifactory() -> bool: """Return `True` if NVIDIA's artifactory is accessible. @@ -936,6 +954,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "has_flashinfer_cutlass_fused_moe", "has_flashinfer_cutedsl_grouped_gemm_nt_masked", "has_flashinfer_cutedsl_moe_nvfp4", + "has_flashinfer_cutedsl_sm12x_moe", "has_flashinfer_fp8_blockscale_gemm", "has_nvidia_artifactory", "supports_trtllm_attention", From b16888b1edb6c45e952f6b608120328b676f6d26 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Fri, 17 Apr 2026 23:00:28 +0000 Subject: [PATCH 02/15] [SM120/121] Add FlashInfer b12x FP4 GEMM backend Integrates FlashInfer PR #3051 b12x dense GEMM backend into the NVFP4 linear layer path. b12x uses CuTe DSL warp-level MMA with adaptive tile sizing to improve SM utilization on small-M decode shapes. Changes: - has_flashinfer_b12x_gemm(): availability check via Sm120BlockScaledDenseGemmKernel - FlashInferB12xNvFp4LinearKernel: new NvFp4LinearKernel subclass - Auto-selects b12x on SM120/SM121 (has_device_capability(120)), falls back to FLASHINFER_CUTLASS when unavailable - Adds "flashinfer-b12x" to VLLM_NVFP4_GEMM_BACKEND valid choices - b12x test cases in test_flashinfer_nvfp4_scaled_mm.py Measured on DGX Spark (SM121, Qwen3-30B-A3B-NVFP4, same MoE backend): b12x: 71.81 out tok/s (1P), 229.24 (8P) flashinfer-cutlass: 70.52 out tok/s (1P), 216.28 (8P) (+1.8% 1P, +6.0% 8P) Signed-off-by: Meenakshi Venkataraman Co-Authored-By: Claude Sonnet 4.6 --- .../test_flashinfer_nvfp4_scaled_mm.py | 10 ++- vllm/envs.py | 1 + .../model_executor/kernels/linear/__init__.py | 4 ++ .../kernels/linear/nvfp4/flashinfer.py | 69 ++++++++++++++++++- vllm/utils/flashinfer.py | 10 +++ 5 files changed, 90 insertions(+), 4 deletions(-) diff --git a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py index e414ba7d2cc3..698c679a201c 100644 --- a/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py +++ b/tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py @@ -13,6 +13,7 @@ from vllm.platforms import current_platform from vllm.utils.flashinfer import ( flashinfer_scaled_fp4_mm, + has_flashinfer_b12x_gemm, ) from vllm.utils.torch_utils import set_random_seed @@ -74,7 +75,7 @@ def get_ref_results( @pytest.mark.parametrize("shape", SHAPES) @pytest.mark.parametrize("seed", SEEDS) @pytest.mark.parametrize("device", CUDA_DEVICES) -@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm"]) +@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm", "b12x"]) @pytest.mark.parametrize("autotune", [False, True]) @torch.inference_mode() def test_flashinfer_nvfp4_gemm( @@ -87,6 +88,10 @@ def test_flashinfer_nvfp4_gemm( ) -> None: if "trtllm" in backend and dtype == torch.float16: pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations") + if backend == "b12x" and not current_platform.has_device_capability(120): + pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)") + if backend == "b12x" and not has_flashinfer_b12x_gemm(): + pytest.skip("b12x FP4 GEMM backend not available in installed FlashInfer") set_random_seed(seed) m, n, packed_k = shape @@ -105,8 +110,7 @@ def test_flashinfer_nvfp4_gemm( # ops.scaled_fp4_quant returns swizzled scales, while weights # from checkpoints are in linear scales. - # So instead of needing to swizzle for cutlass as in modelopt.py, - # we need to unswizzle for trtllm here. + # cutlass and b12x use swizzled scales directly; trtllm needs them unswizzled. a_fp4, a_scale_interleaved = ops.scaled_fp4_quant( a_dtype, a_global_scale, is_sf_swizzled_layout=True, backend=backend ) diff --git a/vllm/envs.py b/vllm/envs.py index bdf7298188a1..e89a4bd6b993 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -1535,6 +1535,7 @@ def _get_or_set_default() -> str: "VLLM_NVFP4_GEMM_BACKEND", None, [ + "flashinfer-b12x", "flashinfer-cudnn", "flashinfer-trtllm", "flashinfer-cutlass", diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 277fb28f3c91..1288e183fa66 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -102,6 +102,7 @@ FbgemmNvFp4LinearKernel, ) from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferB12xNvFp4LinearKernel, FlashInferCudnnNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, FlashInferTrtllmNvFp4LinearKernel, @@ -367,6 +368,7 @@ def _filter_kernels_by_backend( _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { PlatformEnum.CUDA: [ + FlashInferB12xNvFp4LinearKernel, FlashInferCutlassNvFp4LinearKernel, CutlassNvFp4LinearKernel, MarlinNvFp4LinearKernel, @@ -812,6 +814,7 @@ def init_wfp8_a16_linear_kernel( # Maps VLLM_NVFP4_GEMM_BACKEND env var values to kernel classes. _NVFP4_BACKEND_TO_KERNEL: dict[str, type[NvFp4LinearKernel]] = { + "flashinfer-b12x": FlashInferB12xNvFp4LinearKernel, "flashinfer-cutlass": FlashInferCutlassNvFp4LinearKernel, "cutlass": CutlassNvFp4LinearKernel, "marlin": MarlinNvFp4LinearKernel, @@ -1041,6 +1044,7 @@ def register_linear_kernel( "CutlassNvFp4LinearKernel", "EmulationNvFp4LinearKernel", "FbgemmNvFp4LinearKernel", + "FlashInferB12xNvFp4LinearKernel", "FlashInferCutlassNvFp4LinearKernel", "FlashInferTrtllmNvFp4LinearKernel", "FlashInferCudnnNvFp4LinearKernel", diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index 399bc3dd2785..25e1c7c91f96 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -11,7 +11,11 @@ swizzle_blockscale, ) from vllm.platforms import current_platform -from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer +from vllm.utils.flashinfer import ( + flashinfer_scaled_fp4_mm, + has_flashinfer, + has_flashinfer_b12x_gemm, +) from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig @@ -216,3 +220,66 @@ def apply_weights( if bias is not None: out = out + bias return out.view(*output_shape) + + +class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel): + """NVFP4 GEMM via FlashInfer's b12x CuTe DSL warp-level MMA kernel (SM120+).""" + + @classmethod + def is_supported( + cls, compute_capability: int | None = None + ) -> tuple[bool, str | None]: + if current_platform.has_device_capability(120) and has_flashinfer_b12x_gemm(): + return True, None + return False, "FlashInfer b12x requires SM120+ and FlashInfer with Sm120BlockScaledDenseGemmKernel" + + @classmethod + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.weight_scale = torch.nn.Parameter( + swizzle_blockscale(layer.weight_scale.data), requires_grad=False + ) + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass( + layer.weight.data + ) + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False) + layer.weights_padding_cols = weights_padding_cols + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + output_size = layer.output_size_per_partition + output_dtype = x.dtype + output_shape = [*x.shape[:-1], output_size] + + x_fp4, x_blockscale = scaled_fp4_quant( + x, + layer.input_global_scale_inv, + is_sf_swizzled_layout=True, + backend="b12x", + ) + + x_fp4 = pad_nvfp4_activation_for_cutlass( + x_fp4, getattr(layer, "weights_padding_cols", 0) + ) + + out = flashinfer_scaled_fp4_mm( + x_fp4, + layer.weight, + x_blockscale, + layer.weight_scale, + layer.alpha, + output_dtype, + backend="b12x", + ) + + out = slice_nvfp4_output(out, output_size) + + if bias is not None: + out = out + bias + return out.view(*output_shape) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index c32b1fa1ae7d..03d308964ffe 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -267,6 +267,15 @@ def has_flashinfer_cutedsl_moe_nvfp4() -> bool: return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4") +@functools.cache +def has_flashinfer_b12x_gemm() -> bool: + """Return True if FlashInfer b12x FP4 GEMM backend is available (SM120+).""" + if not has_flashinfer_cutedsl(): + return False + mod = _get_submodule("flashinfer.gemm") + return mod is not None and hasattr(mod, "Sm120BlockScaledDenseGemmKernel") + + @functools.cache def has_flashinfer_cutedsl_sm12x_moe() -> bool: """Return ``True`` if FlashInfer CuteDSL SM12x fused MoE is available.""" @@ -955,6 +964,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "has_flashinfer_cutedsl_grouped_gemm_nt_masked", "has_flashinfer_cutedsl_moe_nvfp4", "has_flashinfer_cutedsl_sm12x_moe", + "has_flashinfer_b12x_gemm", "has_flashinfer_fp8_blockscale_gemm", "has_nvidia_artifactory", "supports_trtllm_attention", From cdc0382d0209ad11b82742880b3ec68e03c6cb22 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Fri, 17 Apr 2026 23:47:12 +0000 Subject: [PATCH 03/15] [SM120/121] Address review comments on FlashInferCuteDSLSM12xExperts - Preserve a2_gscale; pass torch.ones_like(a2_gscale) to kernel instead of fill_(1.0) which destroyed the calibrated value in-place - Precompute w1_sf_mma/w2_sf_mma in process_weights_after_loading instead of converting on every forward pass - Fix x_sf_placeholder dtype: float8_e4m3fn (was bfloat16) - Pass topk_weights.float() for float32 routing weights as kernel expects Signed-off-by: Meenakshi Venkataraman Co-Authored-By: Claude Sonnet 4.6 --- .../experts/flashinfer_cutedsl_moe.py | 65 ++++++++++--------- 1 file changed, 34 insertions(+), 31 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index ff68c0bfec98..0947b18ef991 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -350,9 +350,9 @@ class FlashInferCuteDSLSM12xExperts(mk.FusedMoEExpertsModular): into a single kernel call. Input quantization (BF16→FP4) is performed inside the kernel so BF16 hidden states are passed directly. - Weight scale factors must be in the MMA layout produced by - ``convert_sf_to_mma_layout``; this conversion is performed on every - forward pass since it is a zero-copy strided view. + Weight scale factors are converted to the MMA layout produced by + ``convert_sf_to_mma_layout`` once during ``process_weights_after_loading`` + and cached as ``w1_sf_mma`` / ``w2_sf_mma``. Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported. """ @@ -399,10 +399,33 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # calibrated a2_gscale from the modelopt checkpoint (~tens to hundreds) # is intended for static-quantisation backends (TRTLLM/CUTLASS) and # causes every intermediate activation to saturate at max FP4 when - # multiplied by values that large. Force to 1.0 so the kernel uses - # its own per-block dynamic scale — matching the unit-test convention. + # multiplied by values that large. Keep the original a2_gscale intact + # and store a separate ones tensor to pass to the SM12x kernel so it + # uses its own per-block dynamic scale. if self.a2_gscale is not None: - self.a2_gscale.fill_(1.0) + self.a2_gscale_ones = torch.ones_like(self.a2_gscale) + else: + self.a2_gscale_ones = None + + # Precompute MMA-layout views of the weight scale factors once here + # rather than recomputing on every forward pass. + num_experts_w1, m1, k1_sf = self.w1_scale.shape + k1 = k1_sf * 16 + self.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( + self.w1_scale.reshape(num_experts_w1 * m1, k1_sf), + m=m1, + k=k1, + num_groups=num_experts_w1, + ) + + num_experts_w2, m2, k2_sf = self.w2_scale.shape + k2 = k2_sf * 16 + self.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( + self.w2_scale.reshape(num_experts_w2 * m2, k2_sf), + m=m2, + k=k2, + num_groups=num_experts_w2, + ) @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: @@ -500,32 +523,12 @@ def apply( top_k = topk_ids.shape[1] local_expert_offset = self.ep_rank * self.num_local_experts - # Convert swizzled scale factors to the 6D MMA layout expected by - # cute_dsl_fused_moe_nvfp4. This is a zero-copy view+permute. - num_experts_w1, m1, k1_sf = self.w1_scale.shape - k1 = k1_sf * 16 # sf_vec_size = 16 - w1_sf_mma = flashinfer_convert_sf_to_mma_layout( - self.w1_scale.reshape(num_experts_w1 * m1, k1_sf), - m=m1, - k=k1, - num_groups=num_experts_w1, - ) - - num_experts_w2, m2, k2_sf = self.w2_scale.shape - k2 = k2_sf * 16 - w2_sf_mma = flashinfer_convert_sf_to_mma_layout( - self.w2_scale.reshape(num_experts_w2 * m2, k2_sf), - m=m2, - k=k2, - num_groups=num_experts_w2, - ) - # x_sf is ignored by the SM12x kernel (quantization is fused # internally), but the API requires a tensor argument. x_sf_placeholder = ( a1q_scale if a1q_scale is not None - else hidden_states.new_zeros(1) + else hidden_states.new_zeros(1, dtype=torch.float8_e4m3fn) ) # TODO: Use the plan/run() API from FlashInfer PR #3066 instead of @@ -537,13 +540,13 @@ def apply( x=hidden_states, x_sf=x_sf_placeholder, token_selected_experts=topk_ids.to(torch.int32), - token_final_scales=topk_weights, + token_final_scales=topk_weights.float(), w1_weight=w1, - w1_weight_sf=w1_sf_mma, + w1_weight_sf=self.w1_sf_mma, w1_alpha=self.g1_alphas, - fc2_input_scale=self.a2_gscale, + fc2_input_scale=self.a2_gscale_ones, w2_weight=w2, - w2_weight_sf=w2_sf_mma, + w2_weight_sf=self.w2_sf_mma, w2_alpha=self.g2_alphas, num_experts=global_num_experts, top_k=top_k, From 75ddb74b433a962233c1162cccb612a966f43b74 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Mon, 20 Apr 2026 20:11:12 +0000 Subject: [PATCH 04/15] rename: flashinfer_cutedsl_sm12x -> flashinfer_b12x Align user-visible backend name with FlashInfer PR #3080 which establishes b12x as the canonical namespace for SM12x MoE kernels. Renames: - MoEBackend Literal: "flashinfer_cutedsl_sm12x" -> "flashinfer_b12x" - NvFp4MoeBackend enum: FLASHINFER_CUTEDSL_SM12X -> FLASHINFER_B12X - vLLM helper: has_flashinfer_cutedsl_sm12x_moe -> has_flashinfer_b12x_moe Signed-off-by: Meenakshi Venkataraman --- tests/kernels/moe/test_cutedsl_sm12x_moe.py | 8 ++++---- vllm/config/kernel.py | 4 ++-- .../layers/fused_moe/experts/flashinfer_cutedsl_moe.py | 4 ++-- vllm/model_executor/layers/fused_moe/oracle/nvfp4.py | 10 +++++----- .../layers/quantization/utils/flashinfer_fp4_moe.py | 4 ++-- vllm/utils/flashinfer.py | 4 ++-- 6 files changed, 17 insertions(+), 17 deletions(-) diff --git a/tests/kernels/moe/test_cutedsl_sm12x_moe.py b/tests/kernels/moe/test_cutedsl_sm12x_moe.py index 59cc8e8a12cb..499c6dddb80a 100644 --- a/tests/kernels/moe/test_cutedsl_sm12x_moe.py +++ b/tests/kernels/moe/test_cutedsl_sm12x_moe.py @@ -12,9 +12,9 @@ allow_module_level=True, ) -from vllm.utils.flashinfer import has_flashinfer_cutedsl_sm12x_moe +from vllm.utils.flashinfer import has_flashinfer_b12x_moe -if not has_flashinfer_cutedsl_sm12x_moe(): +if not has_flashinfer_b12x_moe(): pytest.skip( reason=( "FlashInfer cute_dsl_fused_moe_nvfp4 / convert_sf_to_mma_layout " @@ -73,7 +73,7 @@ def _reorder_gate_up_to_up_gate( @pytest.mark.parametrize("topk", [1, 2, 4]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @torch.inference_mode() -def test_flashinfer_cutedsl_sm12x_moe( +def test_flashinfer_b12x_moe( m: int, n: int, k: int, @@ -205,4 +205,4 @@ def test_flashinfer_cutedsl_sm12x_moe( if __name__ == "__main__": - test_flashinfer_cutedsl_sm12x_moe(16, 128, 256, 8, 2, torch.bfloat16) + test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 9531517f0eb4..bf145b30f0f3 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -128,7 +128,7 @@ def with_default( "flashinfer_trtllm", "flashinfer_cutlass", "flashinfer_cutedsl", - "flashinfer_cutedsl_sm12x", + "flashinfer_b12x", "marlin", "humming", "triton_unfused", @@ -179,7 +179,7 @@ class KernelConfig: - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) - - "flashinfer_cutedsl_sm12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) + - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) - "marlin": Use Marlin kernels (weight-only quantization) - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 0947b18ef991..e46be7a3a9fd 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -26,7 +26,7 @@ flashinfer_cutedsl_grouped_gemm_nt_masked, has_flashinfer_cutedsl_grouped_gemm_nt_masked, has_flashinfer_cutedsl_moe_nvfp4, - has_flashinfer_cutedsl_sm12x_moe, + has_flashinfer_b12x_moe, scaled_fp4_grouped_quantize, silu_and_mul_scaled_nvfp4_experts_quantize, ) @@ -437,7 +437,7 @@ def _supports_current_device() -> bool: return ( p.is_cuda() and p.is_device_capability_family(120) - and has_flashinfer_cutedsl_sm12x_moe() + and has_flashinfer_b12x_moe() ) @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 28bd4d00f82f..bf684aac394c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -43,7 +43,7 @@ class NvFp4MoeBackend(Enum): FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS" FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL" FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED" - FLASHINFER_CUTEDSL_SM12X = "FLASHINFER_CUTEDSL_SM12X" + FLASHINFER_B12X = "FLASHINFER_B12X" VLLM_CUTLASS = "VLLM_CUTLASS" MARLIN = "MARLIN" EMULATION = "EMULATION" @@ -54,7 +54,7 @@ class NvFp4MoeBackend(Enum): NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, - NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, + NvFp4MoeBackend.FLASHINFER_B12X, ] fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = { @@ -108,7 +108,7 @@ def backend_to_kernel_cls( return [FlashInferCuteDSLBatchedExperts] - elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X: + elif backend == NvFp4MoeBackend.FLASHINFER_B12X: from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( # noqa: E501 FlashInferCuteDSLSM12xExperts, ) @@ -145,7 +145,7 @@ def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM, "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS, "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, - "flashinfer_cutedsl_sm12x": NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, + "flashinfer_b12x": NvFp4MoeBackend.FLASHINFER_B12X, "marlin": NvFp4MoeBackend.MARLIN, "emulation": NvFp4MoeBackend.EMULATION, } @@ -172,7 +172,7 @@ def select_nvfp4_moe_backend( NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, - NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, + NvFp4MoeBackend.FLASHINFER_B12X, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 84f7ca79cd7f..6debbb1ba70a 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -318,7 +318,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, - NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, + NvFp4MoeBackend.FLASHINFER_B12X, ] # Reorder [w1, w3] to [w3, w1] for FI NVFP4 MoE kernels. @@ -330,7 +330,7 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( in [ NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_TRTLLM, - NvFp4MoeBackend.FLASHINFER_CUTEDSL_SM12X, + NvFp4MoeBackend.FLASHINFER_B12X, ] ): w13, w13_scale = reorder_w1w3_to_w3w1(w13, w13_scale) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 03d308964ffe..3e15f74de2ca 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -277,7 +277,7 @@ def has_flashinfer_b12x_gemm() -> bool: @functools.cache -def has_flashinfer_cutedsl_sm12x_moe() -> bool: +def has_flashinfer_b12x_moe() -> bool: """Return ``True`` if FlashInfer CuteDSL SM12x fused MoE is available.""" if not has_flashinfer_moe(): return False @@ -963,7 +963,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "has_flashinfer_cutlass_fused_moe", "has_flashinfer_cutedsl_grouped_gemm_nt_masked", "has_flashinfer_cutedsl_moe_nvfp4", - "has_flashinfer_cutedsl_sm12x_moe", + "has_flashinfer_b12x_moe", "has_flashinfer_b12x_gemm", "has_flashinfer_fp8_blockscale_gemm", "has_nvidia_artifactory", From 08af32a2d93b5de2cb825b9a3829ccf243c444b3 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Wed, 22 Apr 2026 22:40:11 +0000 Subject: [PATCH 05/15] [SM120/121] Switch FlashInferCuteDSLSM12xExperts to b12x_fused_moe (FlashInfer PR #3080) FlashInfer PR #3080 split SM100 and SM120/121 MoE APIs: - cute_dsl_fused_moe_nvfp4 is now SM100-only - b12x_fused_moe is the new SM120/121 entry point Update FlashInferCuteDSLSM12xExperts to call b12x_fused_moe instead of cute_dsl_fused_moe_nvfp4. The new API drops x_sf and local_expert_offset, uses output= instead of moe_output=, and takes topk_weights directly (no .float() cast). Also simplify process_weights_after_loading to fill a2_gscale in-place with 1.0 rather than creating a separate ones tensor. Add flashinfer_b12x_fused_moe lazy wrapper and update has_flashinfer_b12x_moe() to check b12x_fused_moe. Signed-off-by: Meenakshi Venkataraman --- .../experts/flashinfer_cutedsl_moe.py | 56 +++++++------------ vllm/utils/flashinfer.py | 6 +- 2 files changed, 24 insertions(+), 38 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index e46be7a3a9fd..6bb876f38d24 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -21,6 +21,7 @@ ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( + flashinfer_b12x_fused_moe, flashinfer_convert_sf_to_mma_layout, flashinfer_cute_dsl_fused_moe_nvfp4, flashinfer_cutedsl_grouped_gemm_nt_masked, @@ -345,10 +346,10 @@ def flashinfer_cutedsl_moe_masked( class FlashInferCuteDSLSM12xExperts(mk.FusedMoEExpertsModular): """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121, RTX Pro 6000 / DGX Spark). - Uses ``cute_dsl_fused_moe_nvfp4`` from FlashInfer PR #3066 which fuses - token dispatch, two GEMMs, SwiGLU activation, and topk-weight reduction - into a single kernel call. Input quantization (BF16→FP4) is performed - inside the kernel so BF16 hidden states are passed directly. + Uses ``b12x_fused_moe`` from FlashInfer PR #3080 which fuses token + dispatch, two GEMMs, SwiGLU activation, and topk-weight reduction into a + single kernel call. Input quantization (BF16→FP4) is performed inside the + kernel so BF16 hidden states are passed directly. Weight scale factors are converted to the MMA layout produced by ``convert_sf_to_mma_layout`` once during ``process_weights_after_loading`` @@ -399,13 +400,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # calibrated a2_gscale from the modelopt checkpoint (~tens to hundreds) # is intended for static-quantisation backends (TRTLLM/CUTLASS) and # causes every intermediate activation to saturate at max FP4 when - # multiplied by values that large. Keep the original a2_gscale intact - # and store a separate ones tensor to pass to the SM12x kernel so it - # uses its own per-block dynamic scale. + # multiplied by values that large. Force to 1.0 so the kernel uses + # its own per-block dynamic scale. if self.a2_gscale is not None: - self.a2_gscale_ones = torch.ones_like(self.a2_gscale) - else: - self.a2_gscale_ones = None + self.a2_gscale.fill_(1.0) # Precompute MMA-layout views of the weight scale factors once here # rather than recomputing on every forward pass. @@ -463,7 +461,7 @@ def supports_expert_map(self) -> bool: return False def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - # cute_dsl_fused_moe_nvfp4 applies topk weights internally. + # b12x_fused_moe applies topk weights internally. return TopKWeightAndReduceNoOP() def workspace_shapes( @@ -477,7 +475,7 @@ def workspace_shapes( expert_tokens_meta: mk.ExpertTokensMetadata | None, activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # cute_dsl_fused_moe_nvfp4 manages its own internal workspace. + # b12x_fused_moe manages its own internal workspace. workspace1 = (1,) workspace2 = (0,) output_shape = (M, K) @@ -485,11 +483,11 @@ def workspace_shapes( @property def expects_unquantized_inputs(self) -> bool: - # cute_dsl_fused_moe_nvfp4 expects BF16 hidden states and performs - # its own FP4 quantization internally. Returning True prevents the - # modular kernel from pre-quantizing activations, which would produce - # an FP4-packed tensor with size(-1)=k//2 and break the scale-factor - # conversion that expects size(-1)=k. + # b12x_fused_moe expects BF16 hidden states and performs its own FP4 + # quantization internally. Returning True prevents the modular kernel + # from pre-quantizing activations, which would produce an FP4-packed + # tensor with size(-1)=k//2 and break the scale-factor conversion that + # expects size(-1)=k. return True def apply( @@ -521,37 +519,21 @@ def apply( ) top_k = topk_ids.shape[1] - local_expert_offset = self.ep_rank * self.num_local_experts - - # x_sf is ignored by the SM12x kernel (quantization is fused - # internally), but the API requires a tensor argument. - x_sf_placeholder = ( - a1q_scale - if a1q_scale is not None - else hidden_states.new_zeros(1, dtype=torch.float8_e4m3fn) - ) - # TODO: Use the plan/run() API from FlashInfer PR #3066 instead of - # calling cute_dsl_fused_moe_nvfp4 directly. The plan object can be - # created once in __init__ (shapes are fixed for MoE layers) and - # plan.run() called here, avoiding workspace allocation and kernel - # parameter setup overhead on every forward pass. - flashinfer_cute_dsl_fused_moe_nvfp4( + flashinfer_b12x_fused_moe( x=hidden_states, - x_sf=x_sf_placeholder, token_selected_experts=topk_ids.to(torch.int32), - token_final_scales=topk_weights.float(), + token_final_scales=topk_weights, w1_weight=w1, w1_weight_sf=self.w1_sf_mma, w1_alpha=self.g1_alphas, - fc2_input_scale=self.a2_gscale_ones, + fc2_input_scale=self.a2_gscale, w2_weight=w2, w2_weight_sf=self.w2_sf_mma, w2_alpha=self.g2_alphas, num_experts=global_num_experts, top_k=top_k, num_local_experts=self.num_local_experts, - local_expert_offset=local_expert_offset, output_dtype=self.out_dtype, - moe_output=output, + output=output, ) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 3e15f74de2ca..63995af6e736 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -121,6 +121,9 @@ def wrapper(*args, **kwargs): flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper( "flashinfer.fused_moe", "cute_dsl_fused_moe_nvfp4" ) +flashinfer_b12x_fused_moe = _lazy_import_wrapper( + "flashinfer.fused_moe", "b12x_fused_moe" +) flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout" ) @@ -283,7 +286,7 @@ def has_flashinfer_b12x_moe() -> bool: return False required_functions = [ - ("flashinfer.fused_moe", "cute_dsl_fused_moe_nvfp4"), + ("flashinfer.fused_moe", "b12x_fused_moe"), ("flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout"), ] @@ -953,6 +956,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: "scaled_fp4_grouped_quantize", "nvfp4_block_scale_interleave", "flashinfer_cute_dsl_fused_moe_nvfp4", + "flashinfer_b12x_fused_moe", "flashinfer_convert_sf_to_mma_layout", "trtllm_fp4_block_scale_moe", "autotune", From e36f109e1bfd467f50f2c6eda75c8b2f9cb1fcbe Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Fri, 24 Apr 2026 20:22:57 +0000 Subject: [PATCH 06/15] pre-commit: fix ruff/formatting issues Signed-off-by: Meenakshi Venkataraman --- tests/kernels/moe/test_cutedsl_sm12x_moe.py | 9 +++++---- .../kernels/linear/nvfp4/flashinfer.py | 6 +++++- .../fused_moe/experts/flashinfer_cutedsl_moe.py | 13 +++++-------- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/tests/kernels/moe/test_cutedsl_sm12x_moe.py b/tests/kernels/moe/test_cutedsl_sm12x_moe.py index 499c6dddb80a..ceddf53ee123 100644 --- a/tests/kernels/moe/test_cutedsl_sm12x_moe.py +++ b/tests/kernels/moe/test_cutedsl_sm12x_moe.py @@ -8,7 +8,8 @@ if not current_platform.is_device_capability_family(120): pytest.skip( - reason="FlashInfer CuteDSL SM12x MoE requires SM120 (RTX Pro 6000 / DGX Spark).", + reason="FlashInfer CuteDSL SM12x MoE requires SM120 " + "(RTX Pro 6000 / DGX Spark).", allow_module_level=True, ) @@ -134,7 +135,7 @@ def test_flashinfer_b12x_moe( sf_vec_size=sf_vec_size, is_sf_swizzled_layout=True, ) - w1_q = w1_q_flat.view(e, 2 * n, k // 2) # uint8, packed FP4 + w1_q = w1_q_flat.view(e, 2 * n, k // 2) # uint8, packed FP4 w1_blockscale = w1_sf_flat.view(e, 2 * n, w1_sf_flat.shape[1]) # float8 # W2: no row reordering needed for the down-projection. @@ -145,8 +146,8 @@ def test_flashinfer_b12x_moe( sf_vec_size=sf_vec_size, is_sf_swizzled_layout=True, ) - w2_q = w2_q_flat.view(e, k, n // 2) # uint8, packed FP4 - w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1]) # float8 + w2_q = w2_q_flat.view(e, k, n // 2) # uint8, packed FP4 + w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1]) # float8 # All per-expert alphas are 1.0 (global_scale = 1.0, no compensation). ones_e = torch.ones(e, device="cuda", dtype=torch.float32) diff --git a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py index 25e1c7c91f96..7ba2127ff98e 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py +++ b/vllm/model_executor/kernels/linear/nvfp4/flashinfer.py @@ -231,7 +231,11 @@ def is_supported( ) -> tuple[bool, str | None]: if current_platform.has_device_capability(120) and has_flashinfer_b12x_gemm(): return True, None - return False, "FlashInfer b12x requires SM120+ and FlashInfer with Sm120BlockScaledDenseGemmKernel" + return ( + False, + "FlashInfer b12x requires SM120+ and FlashInfer " + "with Sm120BlockScaledDenseGemmKernel", + ) @classmethod def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]: diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 6bb876f38d24..14d8f1c030c0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -11,7 +11,6 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( - TopKWeightAndReduceDelegate, TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -25,9 +24,8 @@ flashinfer_convert_sf_to_mma_layout, flashinfer_cute_dsl_fused_moe_nvfp4, flashinfer_cutedsl_grouped_gemm_nt_masked, - has_flashinfer_cutedsl_grouped_gemm_nt_masked, - has_flashinfer_cutedsl_moe_nvfp4, has_flashinfer_b12x_moe, + has_flashinfer_cutedsl_moe_nvfp4, scaled_fp4_grouped_quantize, silu_and_mul_scaled_nvfp4_experts_quantize, ) @@ -344,7 +342,8 @@ def flashinfer_cutedsl_moe_masked( class FlashInferCuteDSLSM12xExperts(mk.FusedMoEExpertsModular): - """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121, RTX Pro 6000 / DGX Spark). + """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121, + RTX Pro 6000 / DGX Spark). Uses ``b12x_fused_moe`` from FlashInfer PR #3080 which fuses token dispatch, two GEMMs, SwiGLU activation, and topk-weight reduction into a @@ -384,14 +383,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # representations. We set scale_2 = 1.0 to signal that the bake-in is # already done. layer.w13_weight_scale.data = ( - layer.w13_weight_scale.float() - * layer.w13_weight_scale_2.view(-1, 1, 1) + layer.w13_weight_scale.float() * layer.w13_weight_scale_2.view(-1, 1, 1) ).to(layer.w13_weight_scale.dtype) layer.w13_weight_scale_2.data.fill_(1.0) layer.w2_weight_scale.data = ( - layer.w2_weight_scale.float() - * layer.w2_weight_scale_2.view(-1, 1, 1) + layer.w2_weight_scale.float() * layer.w2_weight_scale_2.view(-1, 1, 1) ).to(layer.w2_weight_scale.dtype) layer.w2_weight_scale_2.data.fill_(1.0) From 24859bb33a29d4b7c814bb5cf53e134dd904b564 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Fri, 24 Apr 2026 20:38:48 +0000 Subject: [PATCH 07/15] pre-commit: fix ruff E501 and mypy union-attr errors Signed-off-by: Meenakshi Venkataraman --- vllm/config/kernel.py | 3 ++- .../layers/fused_moe/experts/flashinfer_cutedsl_moe.py | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index bf145b30f0f3..c5f44e1563d0 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -179,7 +179,8 @@ class KernelConfig: - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) - - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x (RTX Pro 6000 / DGX Spark) + - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x + (RTX Pro 6000 / DGX Spark) - "marlin": Use Marlin kernels (weight-only quantization) - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 14d8f1c030c0..e51900f8521c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -404,6 +404,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Precompute MMA-layout views of the weight scale factors once here # rather than recomputing on every forward pass. + assert self.w1_scale is not None num_experts_w1, m1, k1_sf = self.w1_scale.shape k1 = k1_sf * 16 self.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( @@ -413,6 +414,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: num_groups=num_experts_w1, ) + assert self.w2_scale is not None num_experts_w2, m2, k2_sf = self.w2_scale.shape k2 = k2_sf * 16 self.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( From ebb4649730c17778c289bcf5a951aca7585ac303 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Mon, 27 Apr 2026 17:52:18 +0000 Subject: [PATCH 08/15] [SM12x] Address review comments: remove dead code, rename FlashInferB12xExperts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove get_cute_dtype, flashinfer_cutedsl_moe_masked, and associated unused imports from flashinfer_cutedsl_moe.py — copy-pasted from flashinfer_cutedsl_batched_moe.py during an earlier draft that used the masked kernel path; the final SM12x implementation calls b12x_fused_moe directly so they were never used. - Rename FlashInferCuteDSLSM12xExperts -> FlashInferB12xExperts to align with the FlashInferB12xNvFp4LinearKernel naming convention. Signed-off-by: Meenakshi Venkataraman --- tests/kernels/moe/test_cutedsl_sm12x_moe.py | 6 +- .../experts/flashinfer_cutedsl_moe.py | 179 +----------------- .../layers/fused_moe/oracle/nvfp4.py | 4 +- 3 files changed, 10 insertions(+), 179 deletions(-) diff --git a/tests/kernels/moe/test_cutedsl_sm12x_moe.py b/tests/kernels/moe/test_cutedsl_sm12x_moe.py index ceddf53ee123..482b5509a5a0 100644 --- a/tests/kernels/moe/test_cutedsl_sm12x_moe.py +++ b/tests/kernels/moe/test_cutedsl_sm12x_moe.py @@ -38,7 +38,7 @@ ) from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( - FlashInferCuteDSLSM12xExperts, + FlashInferB12xExperts, ) from vllm.utils.torch_utils import set_random_seed @@ -83,7 +83,7 @@ def test_flashinfer_b12x_moe( dtype: torch.dtype, workspace_init, ): - """Test FlashInferCuteDSLSM12xExperts against a BF16 torch reference. + """Test FlashInferB12xExperts against a BF16 torch reference. The SM12x kernel takes BF16 hidden states directly and fuses token dispatch, W1 GEMM, SwiGLU, and W2 GEMM into one call. We verify @@ -176,7 +176,7 @@ def test_flashinfer_b12x_moe( allow_new_interface=True, use_monolithic=False, ), - FlashInferCuteDSLSM12xExperts( + FlashInferB12xExperts( moe_config=moe_config, quant_config=quant_config, ), diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index e51900f8521c..756cd4b61554 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -23,11 +23,8 @@ flashinfer_b12x_fused_moe, flashinfer_convert_sf_to_mma_layout, flashinfer_cute_dsl_fused_moe_nvfp4, - flashinfer_cutedsl_grouped_gemm_nt_masked, has_flashinfer_b12x_moe, has_flashinfer_cutedsl_moe_nvfp4, - scaled_fp4_grouped_quantize, - silu_and_mul_scaled_nvfp4_experts_quantize, ) @@ -175,173 +172,7 @@ def apply( ) -def get_cute_dtype(input: torch.Tensor) -> str: - if input.dtype == torch.bfloat16: - return "bfloat16" - elif input.dtype == torch.float16: - return "float16" - elif input.dtype == torch.float32: - return "float32" - else: - raise ValueError(f"Unsupported cute dtype {input.dtype}") - - -def flashinfer_cutedsl_moe_masked( - hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor], - input_global_scale: torch.Tensor, - w1: torch.Tensor, - w1_blockscale: torch.Tensor, - w1_alpha, - w2: torch.Tensor, - a2_global_scale: torch.Tensor, - w2_blockscale: torch.Tensor, - w2_alpha, - masked_m: torch.Tensor, - workspace: torch.Tensor, - out: torch.Tensor, -): - """ - Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL - kernels. - - Args: - hidden_states: Either of the following case - * torch.Tensor: [num_experts, m, k], bf16 - * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2], - uint8, [num_experts, m, k // 16], float8_e4m3fn - input_global_scale (torch.Tensor): (l,) - w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8 - w1_blockscale (torch.Tensor): blockscale factors, e4m3, - w1_alpha (torch.Tensor): (l,) - w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8 - a2_global_scale (torch.Tensor): (l,) - w2_blockscale (torch.Tensor): blockscale factors, e4m3, - w2_alpha (torch.Tensor): (l,) - masked_m (torch.Tensor): Masked dimension indices - workspace (torch.Tensor): For gateup_output - - Notes: - - Assumes max(masked_m) <= m. - """ - - # === Assertions on dtypes === - assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}" - assert w1_blockscale.dtype == torch.float8_e4m3fn, ( - f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}" - ) - assert w1_alpha.dtype == torch.float32, ( - f"w1_alpha must be float32, got {w1_alpha.dtype}" - ) - assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}" - assert a2_global_scale.dtype == torch.float32, ( - f"a2_global_scale must be float32, got {a2_global_scale.dtype}" - ) - assert w2_blockscale.dtype == torch.float8_e4m3fn, ( - f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}" - ) - assert w2_alpha.dtype == torch.float32, ( - f"w2_alpha must be float32, got {w2_alpha.dtype}" - ) - - # === Assertions on shapes === - n = w2.shape[-1] * 2 # intermediate dimension - if isinstance(hidden_states, tuple): - assert input_global_scale is None, ( - "input_global_scale is needed when input needs quant" - ) - - aq = hidden_states[0].view(torch.uint8) - aq_sf = hidden_states[1].view(torch.float8_e4m3fn) - # m, k_by_2, num_experts = aq.shape - num_experts, m, k_by_2 = aq.shape - k = k_by_2 * 2 - aq = aq.permute(1, 2, 0) - else: - num_experts, m, k = hidden_states.shape - - assert input_global_scale.dtype == torch.float32, ( - f"input_global_scale must be float32, got {input_global_scale.dtype}" - ) - assert input_global_scale.shape == (num_experts,), ( - f"input_global_scale must be (l,), got {input_global_scale.shape}" - ) - - aq, aq_sf = scaled_fp4_grouped_quantize( - hidden_states, - masked_m, - input_global_scale, - ) - - assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}" - assert w1.shape[-1] * 2 == k, ( - f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}" - ) - assert w2.shape[-2:] == ( - k, - n // 2, - ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}" - - assert w1_alpha.shape == (num_experts,), ( - f"w1_alpha must be (l,), got {w1_alpha.shape}" - ) - assert a2_global_scale.shape == (num_experts,), ( - f"a2_global_scale must be (l,), got {a2_global_scale.shape}" - ) - assert w2_alpha.shape == (num_experts,), ( - f"w2_alpha must be (l,), got {w2_alpha.shape}" - ) - - workspace = workspace.permute(1, 2, 0) # requirement of kernel - sf_vec_size = 16 - assert aq_sf.dtype == torch.float8_e4m3fn - assert aq.dtype == torch.uint8 - ab_dtype = "float4_e2m1fn" - sf_dtype = "float8_e4m3fn" - - if isinstance(hidden_states, tuple): - c_dtype = "bfloat16" - else: - c_dtype = get_cute_dtype(hidden_states) - - # Gemm1 - flashinfer_cutedsl_grouped_gemm_nt_masked( - (aq, aq_sf), - (w1.permute(1, 2, 0), w1_blockscale), - workspace, - masked_m, - ab_dtype=ab_dtype, - sf_dtype=sf_dtype, - c_dtype=c_dtype, - sf_vec_size=sf_vec_size, - alpha=w1_alpha.view(1, 1, num_experts), - alpha_dtype=get_cute_dtype(w1_alpha), - ) # in logical [m, n, l] - - # SILU and quantization - diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize( - workspace.permute(2, 0, 1), - masked_m, - a2_global_scale, - ) - - # Gemm2 - out = out.permute(1, 2, 0) # requirement of kernel - flashinfer_cutedsl_grouped_gemm_nt_masked( - (diq, diq_sf), - (w2.permute(1, 2, 0), w2_blockscale), - out, - masked_m, - ab_dtype=ab_dtype, - sf_dtype=sf_dtype, - c_dtype=c_dtype, - sf_vec_size=sf_vec_size, - alpha=w2_alpha.view(1, 1, num_experts), - alpha_dtype=get_cute_dtype(w2_alpha), - ) # in logical [m, k, l] - out = out.permute(2, 0, 1) - - -class FlashInferCuteDSLSM12xExperts(mk.FusedMoEExpertsModular): +class FlashInferB12xExperts(mk.FusedMoEExpertsModular): """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121, RTX Pro 6000 / DGX Spark). @@ -364,7 +195,7 @@ def __init__( ): super().__init__(moe_config=moe_config, quant_config=quant_config) assert quant_config.quant_dtype == "nvfp4", ( - "FlashInferCuteDSLSM12xExperts only supports nvfp4 quantization." + "FlashInferB12xExperts only supports nvfp4 quantization." ) self.out_dtype = moe_config.in_dtype self.num_local_experts = moe_config.num_local_experts @@ -508,13 +339,13 @@ def apply( apply_router_weight_on_input: bool | None, ): assert self.w1_scale is not None and self.w2_scale is not None, ( - "w1_scale and w2_scale must not be None for FlashInferCuteDSLSM12xExperts" + "w1_scale and w2_scale must not be None for FlashInferB12xExperts" ) assert self.g1_alphas is not None and self.g2_alphas is not None, ( - "g1_alphas and g2_alphas must not be None for FlashInferCuteDSLSM12xExperts" + "g1_alphas and g2_alphas must not be None for FlashInferB12xExperts" ) assert self.a2_gscale is not None, ( - "a2_gscale must not be None for FlashInferCuteDSLSM12xExperts" + "a2_gscale must not be None for FlashInferB12xExperts" ) top_k = topk_ids.shape[1] diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index bf684aac394c..d51dd509965e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -110,10 +110,10 @@ def backend_to_kernel_cls( elif backend == NvFp4MoeBackend.FLASHINFER_B12X: from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( # noqa: E501 - FlashInferCuteDSLSM12xExperts, + FlashInferB12xExperts, ) - return [FlashInferCuteDSLSM12xExperts] + return [FlashInferB12xExperts] elif backend == NvFp4MoeBackend.VLLM_CUTLASS: from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( From f005eeb2fd8197321fa27503e6a905514b0581c2 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Mon, 27 Apr 2026 17:57:56 +0000 Subject: [PATCH 09/15] [SM12x] Move FlashInferB12xExperts into its own flashinfer_b12x_moe.py Consistent with the one-class-per-file convention in fused_moe/experts/. flashinfer_cutedsl_moe.py now contains only FlashInferCuteDSLExperts (SM100). Also renames the test to test_flashinfer_b12x_moe.py to match. Signed-off-by: Meenakshi Venkataraman --- ...12x_moe.py => test_flashinfer_b12x_moe.py} | 2 +- .../fused_moe/experts/flashinfer_b12x_moe.py | 223 ++++++++++++++++++ .../experts/flashinfer_cutedsl_moe.py | 198 ---------------- .../layers/fused_moe/oracle/nvfp4.py | 2 +- 4 files changed, 225 insertions(+), 200 deletions(-) rename tests/kernels/moe/{test_cutedsl_sm12x_moe.py => test_flashinfer_b12x_moe.py} (98%) create mode 100644 vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py diff --git a/tests/kernels/moe/test_cutedsl_sm12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py similarity index 98% rename from tests/kernels/moe/test_cutedsl_sm12x_moe.py rename to tests/kernels/moe/test_flashinfer_b12x_moe.py index 482b5509a5a0..bb633d7c2071 100644 --- a/tests/kernels/moe/test_cutedsl_sm12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -37,7 +37,7 @@ maybe_make_prepare_finalize, ) from vllm.model_executor.layers.fused_moe.config import nvfp4_moe_quant_config -from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( +from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import ( FlashInferB12xExperts, ) from vllm.utils.torch_utils import set_random_seed diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py new file mode 100644 index 000000000000..6481434f2e78 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kNvfp4Dynamic, + kNvfp4Static, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import ( + flashinfer_b12x_fused_moe, + flashinfer_convert_sf_to_mma_layout, + has_flashinfer_b12x_moe, +) + + +class FlashInferB12xExperts(mk.FusedMoEExpertsModular): + """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121, + RTX Pro 6000 / DGX Spark). + + Uses ``b12x_fused_moe`` from FlashInfer PR #3080 which fuses token + dispatch, two GEMMs, SwiGLU activation, and topk-weight reduction into a + single kernel call. Input quantization (BF16→FP4) is performed inside the + kernel so BF16 hidden states are passed directly. + + Weight scale factors are converted to the MMA layout produced by + ``convert_sf_to_mma_layout`` once during ``process_weights_after_loading`` + and cached as ``w1_sf_mma`` / ``w2_sf_mma``. + + Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported. + """ + + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config=moe_config, quant_config=quant_config) + assert quant_config.quant_dtype == "nvfp4", ( + "FlashInferB12xExperts only supports nvfp4 quantization." + ) + self.out_dtype = moe_config.in_dtype + self.num_local_experts = moe_config.num_local_experts + self.ep_rank = moe_config.moe_parallel_config.ep_rank + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Normalise block scales to absorb the per-expert weight global scale + # (w_gs). vLLM's NVFP4 convention stores: + # block_scale = max_abs * w_gs / fp4_max, g1_alphas = 1/w_gs + # The SM12x kernel treats w1_alpha (= g1_alphas) as a per-expert weight + # dequant multiplier separate from input_gs (activation scale). We bake + # w_gs into the block scales so that w1_alpha = 1.0 and the kernel sees + # the simpler form: + # block_scale = max_abs / fp4_max, w1_alpha = 1.0 + # The FP4-packed values and dequantised results are identical in both + # representations. We set scale_2 = 1.0 to signal that the bake-in is + # already done. + layer.w13_weight_scale.data = ( + layer.w13_weight_scale.float() * layer.w13_weight_scale_2.view(-1, 1, 1) + ).to(layer.w13_weight_scale.dtype) + layer.w13_weight_scale_2.data.fill_(1.0) + + layer.w2_weight_scale.data = ( + layer.w2_weight_scale.float() * layer.w2_weight_scale_2.view(-1, 1, 1) + ).to(layer.w2_weight_scale.dtype) + layer.w2_weight_scale_2.data.fill_(1.0) + + # The SM12x kernel uses dynamic per-block quantization for FC2 input + # activations (the SwiGLU output before the down projection). The + # calibrated a2_gscale from the modelopt checkpoint (~tens to hundreds) + # is intended for static-quantisation backends (TRTLLM/CUTLASS) and + # causes every intermediate activation to saturate at max FP4 when + # multiplied by values that large. Force to 1.0 so the kernel uses + # its own per-block dynamic scale. + if self.a2_gscale is not None: + self.a2_gscale.fill_(1.0) + + # Precompute MMA-layout views of the weight scale factors once here + # rather than recomputing on every forward pass. + assert self.w1_scale is not None + num_experts_w1, m1, k1_sf = self.w1_scale.shape + k1 = k1_sf * 16 + self.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( + self.w1_scale.reshape(num_experts_w1 * m1, k1_sf), + m=m1, + k=k1, + num_groups=num_experts_w1, + ) + + assert self.w2_scale is not None + num_experts_w2, m2, k2_sf = self.w2_scale.shape + k2 = k2_sf * 16 + self.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( + self.w2_scale.reshape(num_experts_w2 * m2, k2_sf), + m=m2, + k=k2, + num_groups=num_experts_w2, + ) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + p = current_platform + return ( + p.is_cuda() + and p.is_device_capability_family(120) + and has_flashinfer_b12x_moe() + ) + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation == MoEActivation.SILU + + @staticmethod + def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + return True + + def supports_expert_map(self) -> bool: + return False + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + # b12x_fused_moe applies topk weights internally. + return TopKWeightAndReduceNoOP() + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # b12x_fused_moe manages its own internal workspace. + workspace1 = (1,) + workspace2 = (0,) + output_shape = (M, K) + return (workspace1, workspace2, output_shape) + + @property + def expects_unquantized_inputs(self) -> bool: + # b12x_fused_moe expects BF16 hidden states and performs its own FP4 + # quantization internally. Returning True prevents the modular kernel + # from pre-quantizing activations, which would produce an FP4-packed + # tensor with size(-1)=k//2 and break the scale-factor conversion that + # expects size(-1)=k. + return True + + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor | None, + workspace2: torch.Tensor | None, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool | None, + ): + assert self.w1_scale is not None and self.w2_scale is not None, ( + "w1_scale and w2_scale must not be None for FlashInferB12xExperts" + ) + assert self.g1_alphas is not None and self.g2_alphas is not None, ( + "g1_alphas and g2_alphas must not be None for FlashInferB12xExperts" + ) + assert self.a2_gscale is not None, ( + "a2_gscale must not be None for FlashInferB12xExperts" + ) + + top_k = topk_ids.shape[1] + + flashinfer_b12x_fused_moe( + x=hidden_states, + token_selected_experts=topk_ids.to(torch.int32), + token_final_scales=topk_weights, + w1_weight=w1, + w1_weight_sf=self.w1_sf_mma, + w1_alpha=self.g1_alphas, + fc2_input_scale=self.a2_gscale, + w2_weight=w2, + w2_weight_sf=self.w2_sf_mma, + w2_alpha=self.g2_alphas, + num_experts=global_num_experts, + top_k=top_k, + num_local_experts=self.num_local_experts, + output_dtype=self.out_dtype, + output=output, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 756cd4b61554..15f4030c1cc5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -20,10 +20,7 @@ ) from vllm.platforms import current_platform from vllm.utils.flashinfer import ( - flashinfer_b12x_fused_moe, - flashinfer_convert_sf_to_mma_layout, flashinfer_cute_dsl_fused_moe_nvfp4, - has_flashinfer_b12x_moe, has_flashinfer_cutedsl_moe_nvfp4, ) @@ -172,198 +169,3 @@ def apply( ) -class FlashInferB12xExperts(mk.FusedMoEExpertsModular): - """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121, - RTX Pro 6000 / DGX Spark). - - Uses ``b12x_fused_moe`` from FlashInfer PR #3080 which fuses token - dispatch, two GEMMs, SwiGLU activation, and topk-weight reduction into a - single kernel call. Input quantization (BF16→FP4) is performed inside the - kernel so BF16 hidden states are passed directly. - - Weight scale factors are converted to the MMA layout produced by - ``convert_sf_to_mma_layout`` once during ``process_weights_after_loading`` - and cached as ``w1_sf_mma`` / ``w2_sf_mma``. - - Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported. - """ - - def __init__( - self, - moe_config: FusedMoEConfig, - quant_config: FusedMoEQuantConfig, - ): - super().__init__(moe_config=moe_config, quant_config=quant_config) - assert quant_config.quant_dtype == "nvfp4", ( - "FlashInferB12xExperts only supports nvfp4 quantization." - ) - self.out_dtype = moe_config.in_dtype - self.num_local_experts = moe_config.num_local_experts - self.ep_rank = moe_config.moe_parallel_config.ep_rank - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Normalise block scales to absorb the per-expert weight global scale - # (w_gs). vLLM's NVFP4 convention stores: - # block_scale = max_abs * w_gs / fp4_max, g1_alphas = 1/w_gs - # The SM12x kernel treats w1_alpha (= g1_alphas) as a per-expert weight - # dequant multiplier separate from input_gs (activation scale). We bake - # w_gs into the block scales so that w1_alpha = 1.0 and the kernel sees - # the simpler form: - # block_scale = max_abs / fp4_max, w1_alpha = 1.0 - # The FP4-packed values and dequantised results are identical in both - # representations. We set scale_2 = 1.0 to signal that the bake-in is - # already done. - layer.w13_weight_scale.data = ( - layer.w13_weight_scale.float() * layer.w13_weight_scale_2.view(-1, 1, 1) - ).to(layer.w13_weight_scale.dtype) - layer.w13_weight_scale_2.data.fill_(1.0) - - layer.w2_weight_scale.data = ( - layer.w2_weight_scale.float() * layer.w2_weight_scale_2.view(-1, 1, 1) - ).to(layer.w2_weight_scale.dtype) - layer.w2_weight_scale_2.data.fill_(1.0) - - # The SM12x kernel uses dynamic per-block quantization for FC2 input - # activations (the SwiGLU output before the down projection). The - # calibrated a2_gscale from the modelopt checkpoint (~tens to hundreds) - # is intended for static-quantisation backends (TRTLLM/CUTLASS) and - # causes every intermediate activation to saturate at max FP4 when - # multiplied by values that large. Force to 1.0 so the kernel uses - # its own per-block dynamic scale. - if self.a2_gscale is not None: - self.a2_gscale.fill_(1.0) - - # Precompute MMA-layout views of the weight scale factors once here - # rather than recomputing on every forward pass. - assert self.w1_scale is not None - num_experts_w1, m1, k1_sf = self.w1_scale.shape - k1 = k1_sf * 16 - self.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( - self.w1_scale.reshape(num_experts_w1 * m1, k1_sf), - m=m1, - k=k1, - num_groups=num_experts_w1, - ) - - assert self.w2_scale is not None - num_experts_w2, m2, k2_sf = self.w2_scale.shape - k2 = k2_sf * 16 - self.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( - self.w2_scale.reshape(num_experts_w2 * m2, k2_sf), - m=m2, - k=k2, - num_groups=num_experts_w2, - ) - - @staticmethod - def activation_format() -> mk.FusedMoEActivationFormat: - return mk.FusedMoEActivationFormat.Standard - - @staticmethod - def _supports_current_device() -> bool: - p = current_platform - return ( - p.is_cuda() - and p.is_device_capability_family(120) - and has_flashinfer_b12x_moe() - ) - - @staticmethod - def _supports_no_act_and_mul() -> bool: - return False - - @staticmethod - def _supports_quant_scheme( - weight_key: QuantKey | None, - activation_key: QuantKey | None, - ) -> bool: - return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) - - @staticmethod - def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SILU - - @staticmethod - def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - return True - - def supports_expert_map(self) -> bool: - return False - - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: - # b12x_fused_moe applies topk weights internally. - return TopKWeightAndReduceNoOP() - - def workspace_shapes( - self, - M: int, - N: int, - K: int, - topk: int, - global_num_experts: int, - local_num_experts: int, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - activation: MoEActivation, - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: - # b12x_fused_moe manages its own internal workspace. - workspace1 = (1,) - workspace2 = (0,) - output_shape = (M, K) - return (workspace1, workspace2, output_shape) - - @property - def expects_unquantized_inputs(self) -> bool: - # b12x_fused_moe expects BF16 hidden states and performs its own FP4 - # quantization internally. Returning True prevents the modular kernel - # from pre-quantizing activations, which would produce an FP4-packed - # tensor with size(-1)=k//2 and break the scale-factor conversion that - # expects size(-1)=k. - return True - - def apply( - self, - output: torch.Tensor, - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: MoEActivation, - global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor | None, - workspace2: torch.Tensor | None, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool | None, - ): - assert self.w1_scale is not None and self.w2_scale is not None, ( - "w1_scale and w2_scale must not be None for FlashInferB12xExperts" - ) - assert self.g1_alphas is not None and self.g2_alphas is not None, ( - "g1_alphas and g2_alphas must not be None for FlashInferB12xExperts" - ) - assert self.a2_gscale is not None, ( - "a2_gscale must not be None for FlashInferB12xExperts" - ) - - top_k = topk_ids.shape[1] - - flashinfer_b12x_fused_moe( - x=hidden_states, - token_selected_experts=topk_ids.to(torch.int32), - token_final_scales=topk_weights, - w1_weight=w1, - w1_weight_sf=self.w1_sf_mma, - w1_alpha=self.g1_alphas, - fc2_input_scale=self.a2_gscale, - w2_weight=w2, - w2_weight_sf=self.w2_sf_mma, - w2_alpha=self.g2_alphas, - num_experts=global_num_experts, - top_k=top_k, - num_local_experts=self.num_local_experts, - output_dtype=self.out_dtype, - output=output, - ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index d51dd509965e..d745de6c5b52 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -109,7 +109,7 @@ def backend_to_kernel_cls( return [FlashInferCuteDSLBatchedExperts] elif backend == NvFp4MoeBackend.FLASHINFER_B12X: - from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( # noqa: E501 + from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import ( # noqa: E501 FlashInferB12xExperts, ) From a229aa2bf450d8b183583b2d323825fd52ecd339 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Mon, 27 Apr 2026 20:50:10 +0000 Subject: [PATCH 10/15] test(b12x): wire w1/w2_sf_mma before kernel call in test_flashinfer_b12x_moe process_weights_after_loading computes MMA-layout block scales after normalizing. The test was constructing FlashInferB12xExperts without setting these, so the kernel received uninitialized scale buffers. Fix: compute w1_sf_mma / w2_sf_mma via flashinfer_convert_sf_to_mma_layout directly from the test's pre-normalized scales before passing the experts object to FusedMoEKernel. Signed-off-by: Meenakshi Venkataraman Co-Authored-By: Claude Sonnet 4.6 --- tests/kernels/moe/test_flashinfer_b12x_moe.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer_b12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py index bb633d7c2071..af5361ef08ed 100644 --- a/tests/kernels/moe/test_flashinfer_b12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -40,6 +40,7 @@ from vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe import ( FlashInferB12xExperts, ) +from vllm.utils.flashinfer import flashinfer_convert_sf_to_mma_layout from vllm.utils.torch_utils import set_random_seed # Dimensions chosen to satisfy FP4 alignment requirements (k multiple of 256, @@ -169,6 +170,24 @@ def test_flashinfer_b12x_moe( in_dtype=dtype, ) + experts = FlashInferB12xExperts( + moe_config=moe_config, + quant_config=quant_config, + ) + # In production, process_weights_after_loading computes these after + # normalizing block scales. In the test the scales are already in final + # form (global_scale=1.0), so we compute the MMA layouts directly. + num_experts_w1, m1, k1_sf = w1_blockscale.shape + experts.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( + w1_blockscale.reshape(num_experts_w1 * m1, k1_sf), + m=m1, k=k1_sf * 16, num_groups=num_experts_w1, + ) + num_experts_w2, m2, k2_sf = w2_blockscale.shape + experts.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( + w2_blockscale.reshape(num_experts_w2 * m2, k2_sf), + m=m2, k=k2_sf * 16, num_groups=num_experts_w2, + ) + kernel = mk.FusedMoEKernel( maybe_make_prepare_finalize( moe=moe_config, @@ -176,10 +195,7 @@ def test_flashinfer_b12x_moe( allow_new_interface=True, use_monolithic=False, ), - FlashInferB12xExperts( - moe_config=moe_config, - quant_config=quant_config, - ), + experts, inplace=False, ) From 1045ba4f57544e5c295aea29ac398400c6c7e6f4 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Mon, 27 Apr 2026 20:58:30 +0000 Subject: [PATCH 11/15] style: remove trailing blank lines in flashinfer_cutedsl_moe.py Signed-off-by: Meenakshi Venkataraman Co-Authored-By: Claude Sonnet 4.6 --- .../layers/fused_moe/experts/flashinfer_cutedsl_moe.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py index 15f4030c1cc5..2310982792fe 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py @@ -167,5 +167,3 @@ def apply( local_expert_offset=self.local_expert_offset, moe_output=output, ) - - From 457112109533a44eab5bb792d1ee2d4552b0c2c4 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Mon, 27 Apr 2026 21:16:30 +0000 Subject: [PATCH 12/15] style: ruff format flashinfer_convert_sf_to_mma_layout calls in test Signed-off-by: Meenakshi Venkataraman Co-Authored-By: Claude Sonnet 4.6 --- tests/kernels/moe/test_flashinfer_b12x_moe.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer_b12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py index af5361ef08ed..ec0a9594fe12 100644 --- a/tests/kernels/moe/test_flashinfer_b12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -180,12 +180,16 @@ def test_flashinfer_b12x_moe( num_experts_w1, m1, k1_sf = w1_blockscale.shape experts.w1_sf_mma = flashinfer_convert_sf_to_mma_layout( w1_blockscale.reshape(num_experts_w1 * m1, k1_sf), - m=m1, k=k1_sf * 16, num_groups=num_experts_w1, + m=m1, + k=k1_sf * 16, + num_groups=num_experts_w1, ) num_experts_w2, m2, k2_sf = w2_blockscale.shape experts.w2_sf_mma = flashinfer_convert_sf_to_mma_layout( w2_blockscale.reshape(num_experts_w2 * m2, k2_sf), - m=m2, k=k2_sf * 16, num_groups=num_experts_w2, + m=m2, + k=k2_sf * 16, + num_groups=num_experts_w2, ) kernel = mk.FusedMoEKernel( From b0c271671ad3c5eafa4d2aed32b49b923d3265cd Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Tue, 28 Apr 2026 17:48:23 +0000 Subject: [PATCH 13/15] refactor(b12x): remove FLASHINFER_CUTEDSL from unreachable assert FLASHINFER_CUTEDSL is routed to prepare_nvfp4_moe_layer_for_flashinfer_cutedsl before the elif that calls prepare_nvfp4_moe_layer_for_fi_or_cutlass, so it can never reach this assert. Added inadvertently in the b12x commit. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Meenakshi Venkataraman --- .../layers/quantization/utils/flashinfer_fp4_moe.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 6debbb1ba70a..082e42f964f4 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -316,7 +316,6 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.FLASHINFER_TRTLLM, - NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, NvFp4MoeBackend.FLASHINFER_B12X, ] From 8df677dfc396866b590fd0f2c9f171ab374d34ef Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Tue, 28 Apr 2026 18:11:19 +0000 Subject: [PATCH 14/15] fix(b12x): exclude b12x MoE and GEMM backends from auto-selection CUTLASS SM121 MMA op guard causes errors loading NVFP4 models on DGX Spark when b12x is auto-selected. Remove FLASHINFER_B12X from AVAILABLE_BACKENDS in the MoE oracle and FlashInferB12xNvFp4LinearKernel from _POSSIBLE_NVFP4_KERNELS. Both remain reachable via explicit opt-in (moe_backend="flashinfer_b12x" and VLLM_NVFP4_GEMM_BACKEND=flashinfer-b12x). Also restore pre-existing lazy imports in utils/flashinfer.py to their original positions and module paths; only flashinfer_b12x_fused_moe is new. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Meenakshi Venkataraman --- vllm/model_executor/kernels/linear/__init__.py | 4 +++- .../layers/fused_moe/oracle/nvfp4.py | 4 +++- vllm/utils/flashinfer.py | 18 +++++++++--------- 3 files changed, 15 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 1288e183fa66..6058a753c71c 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -368,7 +368,9 @@ def _filter_kernels_by_backend( _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = { PlatformEnum.CUDA: [ - FlashInferB12xNvFp4LinearKernel, + # FlashInferB12xNvFp4LinearKernel excluded from auto-selection until + # upstream CUTLASS SM121 MMA op guard is resolved; use + # VLLM_NVFP4_GEMM_BACKEND=flashinfer-b12x to opt in explicitly. FlashInferCutlassNvFp4LinearKernel, CutlassNvFp4LinearKernel, MarlinNvFp4LinearKernel, diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index d745de6c5b52..3e8f38c414ae 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -168,11 +168,13 @@ def select_nvfp4_moe_backend( """ # NOTE: the kernels are selected in the following order. + # FLASHINFER_B12X is intentionally excluded from auto-selection until + # the upstream CUTLASS SM121 MMA op guard is resolved; use + # moe_backend="flashinfer_b12x" to opt in explicitly. AVAILABLE_BACKENDS = [ NvFp4MoeBackend.FLASHINFER_TRTLLM, NvFp4MoeBackend.FLASHINFER_CUTEDSL, NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, - NvFp4MoeBackend.FLASHINFER_B12X, NvFp4MoeBackend.FLASHINFER_CUTLASS, NvFp4MoeBackend.VLLM_CUTLASS, NvFp4MoeBackend.MARLIN, diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 63995af6e736..4b78b9fb3f52 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -118,15 +118,6 @@ def wrapper(*args, **kwargs): flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper( "flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked" ) -flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper( - "flashinfer.fused_moe", "cute_dsl_fused_moe_nvfp4" -) -flashinfer_b12x_fused_moe = _lazy_import_wrapper( - "flashinfer.fused_moe", "b12x_fused_moe" -) -flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( - "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout" -) flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize") nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize") silu_and_mul_scaled_nvfp4_experts_quantize = _lazy_import_wrapper( @@ -138,6 +129,15 @@ def wrapper(*args, **kwargs): nvfp4_block_scale_interleave = _lazy_import_wrapper( "flashinfer.fp4_quantization", "block_scale_interleave" ) +flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper( + "flashinfer", "cute_dsl_fused_moe_nvfp4" +) +flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( + "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout" +) +flashinfer_b12x_fused_moe = _lazy_import_wrapper( + "flashinfer.fused_moe", "b12x_fused_moe" +) trtllm_fp4_block_scale_moe = _lazy_import_wrapper( "flashinfer", "trtllm_fp4_block_scale_moe" ) From 8385c165a91546aea56f056ca4d18fa68828fd75 Mon Sep 17 00:00:00 2001 From: Meenakshi Venkataraman Date: Mon, 11 May 2026 19:26:50 +0000 Subject: [PATCH 15/15] fix(b12x): accept FlashInfer 0.6.11 renamed Sm120 b12x kernel class FlashInfer 0.6.11 (upstream commit 223f2a49) renamed Sm120BlockScaledDenseGemmKernel -> Sm120B12xBlockScaledDenseGemmKernel via git mv with no aliasing. has_flashinfer_b12x_gemm() now checks for either name so the b12x backend works on 0.6.11+ as well as <= 0.6.10. Signed-off-by: Meenakshi Venkataraman --- vllm/utils/flashinfer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 4b78b9fb3f52..f7ed180a7300 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -276,7 +276,13 @@ def has_flashinfer_b12x_gemm() -> bool: if not has_flashinfer_cutedsl(): return False mod = _get_submodule("flashinfer.gemm") - return mod is not None and hasattr(mod, "Sm120BlockScaledDenseGemmKernel") + if mod is None: + return False + # FlashInfer 0.6.11 renamed Sm120BlockScaledDenseGemmKernel -> + # Sm120B12xBlockScaledDenseGemmKernel (commit 223f2a49). Accept either. + return hasattr(mod, "Sm120B12xBlockScaledDenseGemmKernel") or hasattr( + mod, "Sm120BlockScaledDenseGemmKernel" + ) @functools.cache