From c0409ab86f9137a5d89892c1fdf0b41aa873edeb Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 21 May 2026 19:33:36 +0200 Subject: [PATCH 1/2] [B12x][NVFP4] Add W4A16 support to FlashInfer SM12x MoE wrapper Extends FlashInferB12xExperts to handle both: - W4A4 NVFP4 from modelopt checkpoints (existing path) - W4A16 NVFP4 from compressed-tensors `nvfp4-pack-quantized` Key changes - FlashInferB12xExperts: * activation_precision plumbing ("fp4" / "bf16") auto-detected from quant_config.a1_gscale; no env-var flag. * _supports_quant_scheme accepts the compressed-tensors W4A16 key shape (uint8 weights + kNvfp4StaticGroupScale + kStaticTensorScale, no activation key). * source_format ("modelopt" / "compressed_tensors") read from FusedMoEQuantConfig and forwarded to B12xMoEWrapper. * Precomputed MMA-layout views (w1_sf_mma / w2_sf_mma) cached in process_weights_after_loading. - FusedMoEQuantConfig: new source_format field and plumbing through nvfp4_moe_quant_config / nvfp4_w4a16_moe_quant_config / make_nvfp4_moe_quant_config; ModelOptNvFp4FusedMoE and CompressedTensorsW4A4Nvfp4MoEMethod set it explicitly. - prepare_nvfp4_moe_layer_for_fi_or_cutlass: bump intermediate_size_per_partition when pad is applied. - prepare_fp4_layer_for_marlin: fall back to torch.get_default_dtype() when layer.params_dtype is absent (ParallelLMHead / VocabParallelEmbedding). - VocabParallelEmbedding.weight_loader: reshape (not assert) when AutoQuantize emits scalar FP4 scales (shape () vs (1,)). - ModelOptMixedPrecisionConfig.get_quant_method: route ParallelLMHead through the FP8/NVFP4 linear methods (previously fell through to unquantized). - FlashInfer bumped to 0.6.11.post3. Test plan - pytest tests/kernels/moe/test_flashinfer_b12x_moe.py -v on SM120/SM121. AI assistance was used for refactoring and commit-message drafting; every changed line was reviewed locally by the submitter. Co-authored-by: Claude Signed-off-by: Andrii Skliar --- docker/versions.json | 2 +- requirements/cuda.txt | 4 +- tests/kernels/moe/test_flashinfer_b12x_moe.py | 182 +++++++++++++++--- tests/kernels/moe/utils.py | 5 +- .../model_executor/layers/fused_moe/config.py | 12 ++ .../fused_moe/experts/flashinfer_b12x_moe.py | 139 +++++++++---- .../layers/fused_moe/oracle/nvfp4.py | 4 + .../compressed_tensors_moe_w4a4_nvfp4.py | 1 + .../layers/quantization/modelopt.py | 9 + .../quantization/utils/flashinfer_fp4_moe.py | 3 + .../quantization/utils/marlin_utils_fp4.py | 6 +- .../layers/vocab_parallel_embedding.py | 10 + 12 files changed, 315 insertions(+), 62 deletions(-) diff --git a/docker/versions.json b/docker/versions.json index ee23b5baf04f..6301c6bab81a 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -68,7 +68,7 @@ "default": "true" }, "FLASHINFER_VERSION": { - "default": "0.6.11.post2" + "default": "0.6.11.post3" }, "GDRCOPY_CUDA_VERSION": { "default": "12.8" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 20fec60f27d3..bc31330eee43 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -9,8 +9,8 @@ torchaudio==2.11.0 # These must be updated alongside torch torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version # FlashInfer should be updated together with the Dockerfile -flashinfer-python==0.6.11.post2 -flashinfer-cubin==0.6.11.post2 +flashinfer-python==0.6.11.post3 +flashinfer-cubin==0.6.11.post3 apache-tvm-ffi==0.1.9 tilelang==0.1.9 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to diff --git a/tests/kernels/moe/test_flashinfer_b12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py index ec0a9594fe12..733f7d2f0589 100644 --- a/tests/kernels/moe/test_flashinfer_b12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch @@ -8,8 +10,7 @@ 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 B12x MoE requires SM120 (RTX Pro 6000 / DGX Spark).", allow_module_level=True, ) @@ -18,8 +19,8 @@ if not has_flashinfer_b12x_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)." + "FlashInfer B12xMoEWrapper not available in installed " + "FlashInfer (needs PR #3080)." ), allow_module_level=True, ) @@ -40,7 +41,6 @@ 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, @@ -59,9 +59,9 @@ def _reorder_gate_up_to_up_gate( ) -> 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 + The B12x 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``. + done at model-load time by the FP4 layer-prep helper. """ n = w.shape[1] // 2 return ( @@ -70,6 +70,22 @@ def _reorder_gate_up_to_up_gate( ) +def _process_b12x_weights( + experts: FlashInferB12xExperts, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w1_scale_2: torch.Tensor, + w2_scale_2: torch.Tensor, +) -> None: + layer = SimpleNamespace( + w13_weight_scale=w1_scale, + w13_weight_scale_2=w1_scale_2, + w2_weight_scale=w2_scale, + w2_weight_scale_2=w2_scale_2, + ) + experts.process_weights_after_loading(layer) + + @pytest.mark.parametrize("m,n,k", MNK_FACTORS) @pytest.mark.parametrize("e", [8, 16]) @pytest.mark.parametrize("topk", [1, 2, 4]) @@ -174,22 +190,12 @@ def test_flashinfer_b12x_moe( 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, + _process_b12x_weights( + experts, + w1_blockscale, + w2_blockscale, + ones_e, + ones_e, ) kernel = mk.FusedMoEKernel( @@ -225,5 +231,135 @@ def test_flashinfer_b12x_moe( torch.testing.assert_close(sm12x_output, torch_output, atol=2e-1, rtol=2e-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_b12x_moe_relu2( + m: int, + n: int, + k: int, + e: int, + topk: int, + dtype: torch.dtype, + workspace_init, +): + """Test FlashInferB12xExperts with ReLU2 (non-gated) activation. + + ReLU2 is used by Nemotron-H style models. Unlike the gated SiLU + path, w1 has shape [E, N, K] (not [E, 2N, K]) and the activation + is relu(x)^2 without a gate/up split. + """ + 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 + + # Non-gated: w1 shape is (e, n, k), not (e, 2n, k). + w1_bf16 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 15 + w2_bf16 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 15 + + gs = torch.ones(1, device="cuda", dtype=torch.float32) + sf_vec_size = 16 + + # W1: no gate/up reordering for non-gated. + w1_flat = w1_bf16.reshape(e * 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, n, k // 2) + w1_blockscale = w1_sf_flat.view(e, n, w1_sf_flat.shape[1]) + + 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) + w2_blockscale = w2_sf_flat.view(e, k, w2_sf_flat.shape[1]) + + 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, + activation=MoEActivation.RELU2_NO_MUL, + is_act_and_mul=False, + ) + + experts = FlashInferB12xExperts( + moe_config=moe_config, + quant_config=quant_config, + ) + _process_b12x_weights( + experts, + w1_blockscale, + w2_blockscale, + ones_e, + ones_e, + ) + + kernel = mk.FusedMoEKernel( + maybe_make_prepare_finalize( + moe=moe_config, + quant_config=quant_config, + allow_new_interface=True, + use_monolithic=False, + ), + experts, + inplace=False, + ) + + score = torch.randn((m, e), device="cuda", dtype=dtype) + topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False) + + b12x_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.RELU2_NO_MUL, + apply_router_weight_on_input=False, + expert_map=None, + ) + + torch_output = torch_moe( + a, + w1_bf16, + w2_bf16, + score, + topk, + activation=MoEActivation.RELU2_NO_MUL, + ) + + torch.testing.assert_close( + b12x_output, + torch_output, + atol=2e-1, + rtol=2e-1, + ) + + if __name__ == "__main__": test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16) diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index acb2c21b3896..9141586c0d08 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -53,6 +53,8 @@ def make_dummy_moe_config( hidden_dim: int = 1, intermediate_size_per_partition: int = 1, in_dtype: torch.dtype = torch.bfloat16, + activation: MoEActivation = MoEActivation.SILU, + is_act_and_mul: bool = True, ) -> FusedMoEConfig: """ This is a dummy config for the mk constructor interface @@ -69,7 +71,8 @@ def make_dummy_moe_config( num_local_experts=num_experts, num_logical_experts=num_experts, moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), - activation=MoEActivation.SILU, + activation=activation, + is_act_and_mul=is_act_and_mul, in_dtype=in_dtype, device="cuda", routing_method=RoutingMethodType.TopK, diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index d6f915fa3444..858f3ad85cb3 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -256,6 +256,12 @@ class FusedMoEQuantConfig: mx_alignment: int = 0 + # Identifies the upstream quantization format the weights were exported + # from (e.g. "modelopt", "compressed_tensors"). Backends that need to + # interpret a shared weight payload differently per source (FlashInfer + # B12x NVFP4) read this; most MoE methods can ignore it. + source_format: str | None = None + def __post_init__(self): assert not self.per_act_token_quant or self.block_shape is None, ( "illegal quantization" @@ -506,6 +512,7 @@ def make( gemm1_alpha: float | None = None, gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, + source_format: str | None = None, ) -> "FusedMoEQuantConfig": """ General builder function for a FusedMoEQuantConfig. @@ -577,6 +584,7 @@ def make( gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, + source_format=source_format, ) assert quant_config.per_act_token_quant == per_act_token_quant assert quant_config.per_out_ch_quant == per_out_ch_quant @@ -805,6 +813,7 @@ def nvfp4_moe_quant_config( w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, is_scale_swizzled: bool = True, + source_format: str | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for mxfp4 activations and nvp4 weights. @@ -823,6 +832,7 @@ def nvfp4_moe_quant_config( per_out_ch_quant=False, block_shape=None, is_scale_swizzled=is_scale_swizzled, + source_format=source_format, ) @@ -850,6 +860,7 @@ def nvfp4_w4a16_moe_quant_config( g2_alphas: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor, + source_format: str | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for 16-but activations and nvp4 weights. @@ -861,6 +872,7 @@ def nvfp4_w4a16_moe_quant_config( g1_alphas=g1_alphas, g2_alphas=g2_alphas, weight_dtype="nvfp4", + source_format=source_format, ) 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 index 6481434f2e78..146e8892ee64 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py @@ -17,10 +17,11 @@ QuantKey, kNvfp4Dynamic, kNvfp4Static, + kNvfp4StaticGroupScale, + kStaticTensorScale, ) 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, ) @@ -42,6 +43,11 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular): Only NVFP4 (kNvfp4Static/kNvfp4Dynamic) quantization is supported. """ + _ACTIVATION_MAP: dict[MoEActivation, str] = { + MoEActivation.SILU: "silu", + MoEActivation.RELU2_NO_MUL: "relu2", + } + def __init__( self, moe_config: FusedMoEConfig, @@ -55,6 +61,41 @@ def __init__( self.num_local_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank + # Shape params for B12xMoEWrapper construction. + self.global_num_experts = moe_config.num_experts + self.topk = moe_config.experts_per_token + self.hidden_dim = moe_config.hidden_dim + self.intermediate_size_per_partition = ( + moe_config.intermediate_size_per_partition + ) + self.max_num_tokens = moe_config.max_num_tokens + self.local_expert_offset = self.ep_rank * self.num_local_experts + + activation = moe_config.activation + if activation not in self._ACTIVATION_MAP: + raise ValueError( + f"FlashInferB12xExperts does not support " + f"activation {activation!r}. " + f"Supported: {list(self._ACTIVATION_MAP.keys())}" + ) + self._activation_str = self._ACTIVATION_MAP[activation] + + self.activation_precision = ( + "fp4" if quant_config.a1_gscale is not None else "bf16" + ) + + # source_format selects how the FlashInfer kernel interprets the + # FP4 byte payload. Set by the parent quant method on + # FusedMoEQuantConfig; fall back to "modelopt" if upstream hasn't + # been wired through yet. + self.source_format = quant_config.source_format or "modelopt" + + # Lazily created on first apply() call. + self._wrapper: object | None = None + # Populated in process_weights_after_loading. + self.w1_sf_mma: torch.Tensor | None = None + self.w2_sf_mma: torch.Tensor | None = None + 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: @@ -87,26 +128,27 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: 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 + # Precompute MMA-layout views of the (now-rewritten) weight scale + # factors once here rather than recomputing on every forward pass. + # Converts swizzled 3D scale factors [E, M, K_sf] to the 6D MMA + # layout expected by the SM12x kernel's _get_weight_views(). + assert self.w1_scale is not None and self.w2_scale is not None + sf_vec_size = 16 + E_w1, M_w1, K_sf_w1 = self.w1_scale.shape 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, + self.w1_scale.reshape(E_w1 * M_w1, K_sf_w1), + m=M_w1, + k=K_sf_w1 * sf_vec_size, + num_groups=E_w1, + sf_vec_size=sf_vec_size, ) - - assert self.w2_scale is not None - num_experts_w2, m2, k2_sf = self.w2_scale.shape - k2 = k2_sf * 16 + E_w2, M_w2, K_sf_w2 = self.w2_scale.shape 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, + self.w2_scale.reshape(E_w2 * M_w2, K_sf_w2), + m=M_w2, + k=K_sf_w2 * sf_vec_size, + num_groups=E_w2, + sf_vec_size=sf_vec_size, ) @staticmethod @@ -124,18 +166,30 @@ def _supports_current_device() -> bool: @staticmethod def _supports_no_act_and_mul() -> bool: - return False + return True @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) + # Original W4A4 NVFP4 (modelopt format). + if (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic): + return True + + # W4A16 NVFP4 compressed-tensors `nvfp4-pack-quantized` + return ( + weight_key is not None + and weight_key.dtype == torch.uint8 + and weight_key.scale == kNvfp4StaticGroupScale + and weight_key.scale2 == kStaticTensorScale + and weight_key.symmetric + and activation_key is None + ) @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation == MoEActivation.SILU + return activation in (MoEActivation.SILU, MoEActivation.RELU2_NO_MUL) @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -167,13 +221,31 @@ def workspace_shapes( @property def expects_unquantized_inputs(self) -> bool: - # b12x_fused_moe expects BF16 hidden states and performs its own FP4 + # B12xMoEWrapper 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. + # from pre-quantizing activations. return True + def _ensure_wrapper(self) -> None: + """Lazily create B12xMoEWrapper on first use.""" + if self._wrapper is not None: + return + + from flashinfer.fused_moe import B12xMoEWrapper + + self._wrapper = B12xMoEWrapper( + num_experts=self.global_num_experts, + top_k=self.topk, + hidden_size=self.hidden_dim, + intermediate_size=self.intermediate_size_per_partition, + use_cuda_graph=True, + max_num_tokens=self.max_num_tokens, + num_local_experts=self.num_local_experts, + activation=self._activation_str, + activation_precision=self.activation_precision, + source_format=self.source_format, + ) + def apply( self, output: torch.Tensor, @@ -201,13 +273,14 @@ def apply( assert self.a2_gscale is not None, ( "a2_gscale must not be None for FlashInferB12xExperts" ) + assert self.w1_sf_mma is not None and self.w2_sf_mma is not None, ( + "process_weights_after_loading must run before FlashInferB12xExperts.apply" + ) - top_k = topk_ids.shape[1] + self._ensure_wrapper() - flashinfer_b12x_fused_moe( + result = self._wrapper.run( 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, @@ -215,9 +288,7 @@ def apply( 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, + token_selected_experts=topk_ids.to(torch.int32), + token_final_scales=topk_weights, ) + output.copy_(result) diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 3e8f38c414ae..b83f74c496be 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -439,6 +439,7 @@ def make_nvfp4_moe_quant_config( w2_scale_2: torch.Tensor, a13_scale: torch.Tensor, a2_scale: torch.Tensor, + source_format: str | None = None, ) -> FusedMoEQuantConfig: if backend == NvFp4MoeBackend.MARLIN: return nvfp4_w4a16_moe_quant_config( @@ -446,6 +447,7 @@ def make_nvfp4_moe_quant_config( g2_alphas=w2_scale_2, w1_scale=w13_scale, w2_scale=w2_scale, + source_format=source_format, ) elif backend == NvFp4MoeBackend.EMULATION: return nvfp4_moe_quant_config( @@ -455,6 +457,7 @@ def make_nvfp4_moe_quant_config( a2_gscale=a2_scale, w1_scale=w13_scale, w2_scale=w2_scale, + source_format=source_format, ) # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas. @@ -478,6 +481,7 @@ def make_nvfp4_moe_quant_config( NvFp4MoeBackend.FLASHINFER_CUTEDSL, ) ), + source_format=source_format, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index e1df4484f17e..199cb4538934 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -258,6 +258,7 @@ def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantCon w2_scale_2=layer.w2_weight_scale_2, a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, + source_format="compressed_tensors", ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 4a3b7619b94c..7d670bfbaede 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -85,6 +85,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( requantize_with_max_scale, ) +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.parameter import ( BlockQuantScaleParameter, ChannelQuantScaleParameter, @@ -1601,6 +1602,7 @@ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfi w2_scale_2=layer.w2_weight_scale_2, a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, + source_format="modelopt", ) @property @@ -2357,6 +2359,13 @@ def get_quant_method( # Layer not in quantized_layers — leave unquantized return UnquantizedLinearMethod() + if isinstance(layer, ParallelLMHead): + if quant_algo == "FP8": + return ModelOptFp8LinearMethod(self.fp8_config) + if quant_algo == "NVFP4": + return ModelOptNvFp4LinearMethod(self.nvfp4_config) + return UnquantizedLinearMethod() + if isinstance(layer, RoutedExperts): if quant_algo == "FP8": return ModelOptFp8MoEMethod( 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 082e42f964f4..df509adbfca8 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -388,5 +388,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( w2_scale = torch.nn.functional.pad(w2_scale, (0, pad_size // 16)) w2_scale = swizzle_blockscale(w2_scale) + layer.moe_config.intermediate_size_per_partition = ( + layer.moe_config.intermediate_size_per_partition + pad_size + ) return w13, w13_scale, w13_scale_2, a13_scale, w2, w2_scale, w2_scale_2, a2_scale diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py index c02d39c17a02..69b26c05a0be 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py @@ -216,7 +216,11 @@ def prepare_fp4_layer_for_marlin( part_size_n = layer.output_size_per_partition part_size_k = layer.input_size_per_partition - param_dtype = layer.params_dtype + # VocabParallelEmbedding / ParallelLMHead does not store params_dtype as an + # attribute + param_dtype = getattr(layer, "params_dtype", None) + if param_dtype is None: + param_dtype = torch.get_default_dtype() assert layer.weight.shape == (part_size_n, part_size_k // 2) diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index ddae01856da0..31dd59caa10e 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -438,6 +438,16 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor): # If parameter does not have output dim, then it should # be copied onto all gpus (e.g. g_idx for act_order gptq). if output_dim is None: + # AutoQuantize-quantized lm_head emits FP4 scalar scales + # (input_scale, weight_scale_2). Their on-disk shape is () while + # PerTensorScaleParameter materializes them as (1,) -- same numel, + # different rank. Reshape rather than asserting. + if ( + param.data.shape != loaded_weight.shape + and param.data.numel() == loaded_weight.numel() + ): + param.data.copy_(loaded_weight.reshape(param.data.shape)) + return assert param.data.shape == loaded_weight.shape param.data.copy_(loaded_weight) return From 957a9ec2394f9294705be0610ef292a1b6d03f9a Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 21 May 2026 19:54:40 +0200 Subject: [PATCH 2/2] [B12x][Test] Parametrize SiLU and ReLU2 tests over W4A4 / W4A16 Adds activation_precision in {"fp4", "bf16"} parametrization to both existing FlashInferB12xExperts tests: - test_flashinfer_b12x_moe (SiLU + Mul) - test_flashinfer_b12x_moe_relu2 (ReLU2 + non-gated, Nemotron-H 3.5) For the W4A16 branch, a1_gscale and a2_gscale are passed as None and source_format is set to "compressed_tensors", matching the production path used by compressed-tensors `nvfp4-pack-quantized` checkpoints. Each test now also asserts that FlashInferB12xExperts auto-detects activation_precision from a1_gscale and reads source_format off the quant config, locking in the contract added in the parent PR. Co-authored-by: Claude Signed-off-by: Andrii Skliar --- tests/kernels/moe/test_flashinfer_b12x_moe.py | 55 +++++++++++++++++-- 1 file changed, 50 insertions(+), 5 deletions(-) diff --git a/tests/kernels/moe/test_flashinfer_b12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py index 733f7d2f0589..240f7c4bbe96 100644 --- a/tests/kernels/moe/test_flashinfer_b12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -90,6 +90,7 @@ def _process_b12x_weights( @pytest.mark.parametrize("e", [8, 16]) @pytest.mark.parametrize("topk", [1, 2, 4]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("activation_precision", ["fp4", "bf16"]) @torch.inference_mode() def test_flashinfer_b12x_moe( m: int, @@ -98,6 +99,7 @@ def test_flashinfer_b12x_moe( e: int, topk: int, dtype: torch.dtype, + activation_precision: str, workspace_init, ): """Test FlashInferB12xExperts against a BF16 torch reference. @@ -107,6 +109,14 @@ def test_flashinfer_b12x_moe( correctness against ``torch_moe`` using generous tolerances to account for the internal FP4 quantization of activations and weights. + Two activation precisions are exercised: + * ``fp4`` (W4A4, modelopt path): activations re-quantized to FP4 + before each GEMM. Set ``a1_gscale`` / ``a2_gscale`` on the quant + config (any non-None tensor enables this branch). + * ``bf16`` (W4A16, compressed-tensors `nvfp4-pack-quantized` path): + activations kept in BF16; the kernel dequantizes weights instead. + ``a1_gscale`` / ``a2_gscale`` are left as ``None``. + Scale convention ---------------- The SM12x kernel uses ``w1_alpha`` as *both* the activation-quantisation @@ -169,13 +179,26 @@ def test_flashinfer_b12x_moe( # All per-expert alphas are 1.0 (global_scale = 1.0, no compensation). ones_e = torch.ones(e, device="cuda", dtype=torch.float32) + # W4A4: pass identity activation global scales; B12x re-quantizes + # activations to FP4. W4A16: leave them as None; activations stay + # in BF16 and the kernel dequantizes weights instead. + if activation_precision == "fp4": + a1_gscale = ones_e + a2_gscale = ones_e + source_format = "modelopt" + else: + a1_gscale = None + a2_gscale = None + source_format = "compressed_tensors" + quant_config = nvfp4_moe_quant_config( g1_alphas=ones_e, g2_alphas=ones_e, - a1_gscale=ones_e, - a2_gscale=ones_e, + a1_gscale=a1_gscale, + a2_gscale=a2_gscale, w1_scale=w1_blockscale, w2_scale=w2_blockscale, + source_format=source_format, ) moe_config = make_dummy_moe_config( @@ -190,6 +213,9 @@ def test_flashinfer_b12x_moe( moe_config=moe_config, quant_config=quant_config, ) + # Cross-check the precision/source-format auto-detection. + assert experts.activation_precision == activation_precision + assert experts.source_format == source_format _process_b12x_weights( experts, w1_blockscale, @@ -235,6 +261,7 @@ def test_flashinfer_b12x_moe( @pytest.mark.parametrize("e", [8, 16]) @pytest.mark.parametrize("topk", [1, 2, 4]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("activation_precision", ["fp4", "bf16"]) @torch.inference_mode() def test_flashinfer_b12x_moe_relu2( m: int, @@ -243,6 +270,7 @@ def test_flashinfer_b12x_moe_relu2( e: int, topk: int, dtype: torch.dtype, + activation_precision: str, workspace_init, ): """Test FlashInferB12xExperts with ReLU2 (non-gated) activation. @@ -250,6 +278,10 @@ def test_flashinfer_b12x_moe_relu2( ReLU2 is used by Nemotron-H style models. Unlike the gated SiLU path, w1 has shape [E, N, K] (not [E, 2N, K]) and the activation is relu(x)^2 without a gate/up split. + + Parametrized over ``activation_precision`` (``"fp4"`` for W4A4 / + modelopt and ``"bf16"`` for W4A16 / compressed-tensors); Nemotron-H 3.5 + is the production user of the ReLU2 + W4A16 combination. """ set_random_seed(7) with set_current_vllm_config( @@ -287,13 +319,23 @@ def test_flashinfer_b12x_moe_relu2( ones_e = torch.ones(e, device="cuda", dtype=torch.float32) + if activation_precision == "fp4": + a1_gscale = ones_e + a2_gscale = ones_e + source_format = "modelopt" + else: + a1_gscale = None + a2_gscale = None + source_format = "compressed_tensors" + quant_config = nvfp4_moe_quant_config( g1_alphas=ones_e, g2_alphas=ones_e, - a1_gscale=ones_e, - a2_gscale=ones_e, + a1_gscale=a1_gscale, + a2_gscale=a2_gscale, w1_scale=w1_blockscale, w2_scale=w2_blockscale, + source_format=source_format, ) moe_config = make_dummy_moe_config( @@ -310,6 +352,8 @@ def test_flashinfer_b12x_moe_relu2( moe_config=moe_config, quant_config=quant_config, ) + assert experts.activation_precision == activation_precision + assert experts.source_format == source_format _process_b12x_weights( experts, w1_blockscale, @@ -362,4 +406,5 @@ def test_flashinfer_b12x_moe_relu2( if __name__ == "__main__": - test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16) + test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16, "fp4") + test_flashinfer_b12x_moe(16, 128, 256, 8, 2, torch.bfloat16, "bf16")