diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py index 3d637258d3..e3a3478e37 100644 --- a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py +++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py @@ -29,14 +29,20 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, FusedMoEParallelConfig, RoutingMethodType, - nvfp4_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( FlashInferCuteDSLExperts, ) +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, + make_nvfp4_moe_quant_config, +) from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( prepare_nvfp4_moe_layer_for_flashinfer_cutedsl, ) +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + quantize_nvfp4_per_token_input, +) from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer_cutedsl_moe_nvfp4 from vllm.utils.math_utils import next_power_of_2 @@ -159,8 +165,12 @@ def _dequantize_nvfp4_linear( return (tensor_f32 * tensor_sf.unsqueeze(-1)).reshape(m, k).to(dtype) -@pytest.mark.parametrize("m,n,k,e,topk", [(16, 128, 512, 4, 2)]) +@pytest.mark.parametrize( + "m,n,k,e,topk", + [(16, 128, 512, 4, 2), (16, 192, 512, 4, 2), (16, 192, 2688, 4, 2)], +) @pytest.mark.parametrize("activation,alpha,beta,limit", _ACT_CASES) +@pytest.mark.parametrize("per_token_activation", [False, True]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @torch.inference_mode() def test_flashinfer_cutedsl_fp4_moe( @@ -173,6 +183,7 @@ def test_flashinfer_cutedsl_fp4_moe( alpha: float | None, beta: float | None, limit: float | None, + per_token_activation: bool, dtype: torch.dtype, workspace_init, ): @@ -181,6 +192,10 @@ def test_flashinfer_cutedsl_fp4_moe( VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1)) ): hidden_states = torch.randn((m, k), device="cuda", dtype=dtype) / 10 + # Keep GEMM1 activation variance stable as K changes, so the fixed + # tolerance against the BF16 GEMM2 reference measures comparable FP4 + # quantization error at each hidden width. + hidden_states *= (512 / k) ** 0.5 w1_rows = 2 * n if activation.is_gated else n w1 = torch.randn((e, w1_rows, k), device="cuda", dtype=dtype) / 15 @@ -193,7 +208,6 @@ def test_flashinfer_cutedsl_fp4_moe( hidden_states, score, topk, renormalize=False ) - activation = MoEActivation.RELU2_NO_MUL moe_config = FusedMoEConfig( num_experts=e, experts_per_token=topk, @@ -231,19 +245,30 @@ def test_flashinfer_cutedsl_fp4_moe( w2_scale_2=(1.0 / w2_global_scale), a2_scale=a2_scale, ) - quant_config = nvfp4_moe_quant_config( - g1_alphas=w1_alpha, - g2_alphas=w2_alpha, - a1_gscale=(1.0 / a1_scale), - a2_gscale=(1.0 / a2_scale), - w1_scale=w1_scale_cutedsl, + # I=192 must remain unpadded for gated experts, but become 256 for + # non-gated experts. The reference below uses the checkpoint shapes. + expected_n = n if activation.is_gated else (n + 127) // 128 * 128 + assert moe_config.intermediate_size_per_partition == expected_n + assert w1_cutedsl.shape[1] == expected_n * (2 if activation.is_gated else 1) + assert w2_cutedsl.shape[2] * 2 == expected_n + expected_k = (k + 255) // 256 * 256 + assert moe_config.hidden_dim == expected_k + assert moe_config.hidden_dim_unpadded == k + assert w1_cutedsl.shape[2] * 2 == expected_k + assert w2_cutedsl.shape[1] == expected_k + quant_config = make_nvfp4_moe_quant_config( + backend=NvFp4MoeBackend.FLASHINFER_CUTEDSL, + w13_scale_2=w1_alpha, + w2_scale_2=w2_alpha, + a13_scale=a1_scale, + a2_scale=a2_scale, + w13_scale=w1_scale_cutedsl, w2_scale=w2_scale_cutedsl, - is_scale_swizzled=False, # Unset params must be omitted rather than forwarded as None into # the kernel's float-typed SwiGLU arguments. - gemm1_alpha=alpha, - gemm1_beta=beta, - gemm1_clamp_limit=limit, + swiglu_alpha=alpha, + swiglu_beta=beta, + swiglu_limit=limit, ) cutedsl_experts = mk.FusedMoEKernel( @@ -256,11 +281,13 @@ def test_flashinfer_cutedsl_fp4_moe( FlashInferCuteDSLExperts( moe_config=moe_config, quant_config=quant_config, + per_token_activation=per_token_activation, ), ) cutedsl_output = cutedsl_experts.apply( - hidden_states=hidden_states, + # The model runner pads inputs and trims outputs around this kernel. + hidden_states=torch.nn.functional.pad(hidden_states, (0, expected_k - k)), w1=w1_cutedsl, w2=w2_cutedsl, topk_weights=topk_weights, @@ -270,19 +297,25 @@ def test_flashinfer_cutedsl_fp4_moe( expert_map=None, apply_router_weight_on_input=False, ) + cutedsl_output = cutedsl_output[:, :k] a_global_scale = torch.ones(1, device="cuda", dtype=torch.float32) - a_q, a_scale = ops.scaled_fp4_quant( - hidden_states, - a_global_scale, - is_sf_swizzled_layout=False, - ) + if per_token_activation: + a_q, a_scale, a_row_scale = quantize_nvfp4_per_token_input(hidden_states) + else: + a_q, a_scale = ops.scaled_fp4_quant( + hidden_states, + a_global_scale, + is_sf_swizzled_layout=False, + ) a_in_dtype = _dequantize_nvfp4_linear( a_q, a_scale, a_global_scale, dtype=dtype, ) + if per_token_activation: + a_in_dtype = (a_in_dtype.float() * a_row_scale.view(-1, 1)).to(dtype) w1_d = torch.empty((e, w1_rows, k), device="cuda", dtype=dtype) w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype) diff --git a/tests/quantization/test_online.py b/tests/quantization/test_online.py index 645ad67062..c7495326e2 100644 --- a/tests/quantization/test_online.py +++ b/tests/quantization/test_online.py @@ -605,6 +605,30 @@ def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None: ModelConfig(model=str(tmp_path), quantization="fp8_per_channel") +def test_nvfp4_per_token_backend_contract() -> None: + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( + FlashInferCuteDSLExperts, + ) + from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( + FlashInferExperts, + ) + from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( + TrtLlmNvFp4ExpertsModular, + TrtLlmNvFp4ExpertsMonolithic, + ) + from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kNvfp4DynamicToken, + kNvfp4Static, + ) + + scheme = (kNvfp4Static, kNvfp4DynamicToken) + assert TrtLlmNvFp4ExpertsMonolithic._supports_quant_scheme(*scheme) + assert TrtLlmNvFp4ExpertsModular._supports_quant_scheme(*scheme) + assert FlashInferCuteDSLExperts._supports_quant_scheme(*scheme) + assert FlashInferCuteDSLExperts._supports_no_act_and_mul() + assert not FlashInferExperts._supports_quant_scheme(*scheme) + + @pytest.mark.parametrize("per_token_activation", [False, True]) @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16]) def test_nvfp4_one_sided_sizes_dispatched_activations( diff --git a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py index be8f984853..9effb4f7fc 100644 --- a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py +++ b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py @@ -11,8 +11,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( prepare_nvfp4_moe_layer_for_fi_or_cutlass, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + align_fp4_moe_hidden_dim_for_fi, align_fp4_moe_weights_for_fi, - align_trtllm_fp4_moe_hidden_dim_for_fi, ) @@ -63,7 +63,7 @@ def test_align_trtllm_fp4_moe_hidden_dim_noop(): w2_scale = torch.arange(2 * 512 * 1, dtype=torch.uint8).reshape(2, 512, 1) out_w13, out_w13_scale, out_w2, out_w2_scale, padded_hidden = ( - align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) + align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) ) assert padded_hidden == 512 @@ -90,7 +90,7 @@ def test_align_trtllm_fp4_moe_hidden_dim_pads_to_256_multiple(): ) out_w13, out_w13_scale, out_w2, out_w2_scale, out_hidden_dim = ( - align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) + align_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) ) assert out_hidden_dim == padded_hidden_dim 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 d525822b91..15fa18b13c 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 @@ -14,11 +14,14 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE, activation_to_flashinfer_int, + quantize_nvfp4_per_token_input, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kNvfp4Dynamic, + kNvfp4DynamicToken, kNvfp4Static, ) from vllm.platforms import current_platform @@ -40,6 +43,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, + per_token_activation: bool = False, ): super().__init__( moe_config=moe_config, @@ -63,6 +67,14 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit self.situ_beta = moe_config.activation_situ_beta self.situ_linear_beta = moe_config.activation_situ_linear_beta + self.per_token_activation = per_token_activation + self.per_token_global_scale = None + if per_token_activation: + assert quant_config.a2_gscale is not None + self.per_token_global_scale = quant_config.a2_gscale.new_full( + (1,), + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE, + ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) @@ -92,6 +104,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): ) -> bool: SUPPORTED_W_A = [ (kNvfp4Static, kNvfp4Dynamic), + (kNvfp4Static, kNvfp4DynamicToken), ] return (weight_key, activation_key) in SUPPORTED_W_A @@ -114,6 +127,10 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() + @property + def expects_unquantized_inputs(self) -> bool: + return self.per_token_activation + def workspace_shapes( self, M: int, @@ -127,8 +144,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: workspace1 = (0,) workspace2 = (0,) - # K is packed (K//2 for uint8), so output uses hidden_dim. - assert self.hidden_dim == K * 2 + expected_hidden_dim = K if self.expects_unquantized_inputs else K * 2 + assert self.hidden_dim == expected_hidden_dim output = (M, self.hidden_dim) return (workspace1, workspace2, output) @@ -151,13 +168,26 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): apply_router_weight_on_input: bool | None, ): assert self.quant_dtype == "nvfp4" - assert a1q_scale is not None assert self.w1_scale is not None assert self.w2_scale is not None + if self.expects_unquantized_inputs: + hidden_states, block_scale, per_token_scale = ( + quantize_nvfp4_per_token_input(hidden_states) + ) + fc2_input_scale = self.per_token_global_scale + else: + assert a1q_scale is not None + block_scale = a1q_scale + per_token_scale = None + fc2_input_scale = self.a2_gscale + + assert block_scale is not None + assert fc2_input_scale is not None + # a1q_scale is (M, K//16) float8_e4m3fn from fp4_quantize. # The functional API expects x_sf with trailing dim: (M, K//16, 1). - x_sf = a1q_scale.unsqueeze(-1) + x_sf = block_scale.unsqueeze(-1) # The kernel defaults swiglu_{alpha,beta,limit} to the plain-SwiGLU # values, so only forward the ones the model actually sets. @@ -188,6 +218,9 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): } swiglu_kwargs = {k: v for k, v in swiglu_params.items() if v is not None} + per_token_kwargs = ( + {"per_token_scale": per_token_scale} if self.per_token_activation else {} + ) flashinfer_cute_dsl_fused_moe_nvfp4( x=hidden_states, x_sf=x_sf, @@ -196,7 +229,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): w1_weight=w1, w1_weight_sf=self.w1_scale, w1_alpha=self.g1_alphas, - fc2_input_scale=self.a2_gscale, + fc2_input_scale=fc2_input_scale, w2_weight=w2, w2_weight_sf=self.w2_scale, w2_alpha=self.g2_alphas, @@ -209,4 +242,5 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular): MoEActivation.SILU if activation == MoEActivation.SITU else activation ), **swiglu_kwargs, + **per_token_kwargs, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index f9cca18a9c..9dcfc6af3f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( activation_to_flashinfer_int, has_flashinfer_situ_activation, + quantize_nvfp4_per_token_input, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -37,10 +38,6 @@ from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe logger = init_logger(__name__) -# Base scale for per-token NVFP4 activation quant; the kernel folds the -# per-token global scale (from the activation amax) on top of it. -_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0) - class TrtLlmNvFp4ExpertsBase: """NvFp4 TRTLLM-Gen MoE kernels. Supports modular and monolithic interface.""" @@ -259,23 +256,6 @@ class TrtLlmNvFp4ExpertsBase: def expects_unquantized_inputs(self) -> bool: return self.per_token_activation - def _quantize_per_token_input( - self, hidden_states: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """NVFP4-quantize activations with a per-token global scale. - - Returns ``(packed_fp4, block_scale, per_token_scale)``. - """ - from flashinfer import SfLayout, nvfp4_quantize - - hs_fp4, hs_block_scale, per_token_scale = nvfp4_quantize( - hidden_states, - _PER_TOKEN_BASE_GLOBAL_SCALE, - sfLayout=SfLayout.layout_linear, - per_token_activation=True, - ) - return hs_fp4, hs_block_scale, per_token_scale - def _get_chunk_size(self) -> int: MAX_GRID_Y = 65535 MAX_TILE_TOKENS_DIM = 128 @@ -361,9 +341,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula # Per-token: input is unquantized, quantize it here. Otherwise it was # already quantized in prepare() with the static global scale. - if self.per_token_activation: + if self.expects_unquantized_inputs: hidden_states, block_scale, per_token_scale = ( - self._quantize_per_token_input(hidden_states) + quantize_nvfp4_per_token_input(hidden_states) ) else: block_scale, per_token_scale = a1q_scale, None @@ -437,7 +417,7 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula ) -> UnfinalizedMoEOutput | None: assert self._supports_activation(activation) # Per-token defers input quant to _invoke_kernel, so a1q_scale is None. - assert a1q_scale is not None or self.per_token_activation + assert a1q_scale is not None or self.expects_unquantized_inputs # DeepEP produces int64 indexes. topk_ids = topk_ids.to(dtype=torch.int32) @@ -539,7 +519,7 @@ class TrtLlmNvFp4ExpertsMonolithic( import flashinfer assert self._supports_activation(activation) - assert a1q_scale is not None or self.per_token_activation + assert a1q_scale is not None or self.expects_unquantized_inputs assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None assert ( @@ -551,9 +531,9 @@ class TrtLlmNvFp4ExpertsMonolithic( ) # Per-token: input is unquantized, quantize it here (see modular apply). - if self.per_token_activation: + if self.expects_unquantized_inputs: hidden_states, block_scale, per_token_scale = ( - self._quantize_per_token_input(hidden_states) + quantize_nvfp4_per_token_input(hidden_states) ) else: block_scale, per_token_scale = a1q_scale, None diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 0fecb65e4b..e1952c0032 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -612,7 +612,14 @@ def make_nvfp4_moe_kernel( logger.info_once("Using %s", prepare_finalize.__class__.__name__) extra_kwargs = {} - if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM and per_token_activation: + if ( + backend + in ( + NvFp4MoeBackend.FLASHINFER_TRTLLM, + NvFp4MoeBackend.FLASHINFER_CUTEDSL, + ) + and per_token_activation + ): extra_kwargs["per_token_activation"] = True # Create Experts. diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 08279d1c89..9cb7cfbb0f 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -720,7 +720,7 @@ class MoERunner(MoERunnerInterface): ) # Record before `_maybe_pad_hidden_states` pads activations to match - # `moe_config.hidden_dim`, e.g. after `align_trtllm_fp4_moe_hidden_dim_for_fi` + # `moe_config.hidden_dim`, e.g. after `align_fp4_moe_hidden_dim_for_fi` # so routed output can be trimmed before # shared+routed add / latent up proj if needed. 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 c5c01e59ca..24b95ef89a 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -9,8 +9,8 @@ import torch from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + align_fp4_moe_hidden_dim_for_fi, align_fp4_moe_weights_for_fi, - align_trtllm_fp4_moe_hidden_dim_for_fi, ) from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( swizzle_blockscale, @@ -128,9 +128,9 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( ]: """Prepare weights for the CuteDSL wrapper-based NvFP4 MoE backend. - Converts weight scale factors to MMA layout expected by CuteDslMoEWrapper, - and interleaves w13 gate/linear rows for gated activations. Non-gated - activations use a single w13 projection and keep its row order unchanged. + Pads the runtime expert tensors to the kernel's GEMM alignment, converts + weight scale factors to the MMA layout expected by CuteDslMoEWrapper, and + interleaves w13 gate/linear rows for gated activations. """ # Global scaling factors (same as other FlashInfer backends). num_experts = w13.shape[0] @@ -140,12 +140,38 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( ) a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(num_experts) - if layer.activation.is_gated: + gated = layer.activation.is_gated + if gated: w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl( layer.activation, w13, w13_scale ) - # Interleave up/gate rows for w13 weights and scales. + # GEMM1's output dimension must be a multiple of 128: 2I for gated + # activations (also required by interleaving), but only I for non-gated. + # Keep the checkpoint tensors unchanged and pad only the kernel's runtime + # representation. Zero rows also make the padded GEMM2 contraction a no-op. + w13, w13_scale, w2, w2_scale, padded_intermediate = align_fp4_moe_weights_for_fi( + w13, + w13_scale, + w2, + w2_scale, + is_act_and_mul=gated, + min_alignment=64 if gated else 128, + ) + layer.moe_config.intermediate_size_per_partition = padded_intermediate + + # GEMM1 gathers full 256-element K tiles of activations and their scales + # without a K-tail predicate. Pad runtime weights to match; the MoE runner + # pads activations and trims the output using the original hidden size. + w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi( + w13, w13_scale, w2, w2_scale + ) + if layer.moe_config.hidden_dim_unpadded is None: + layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim + layer.moe_config.hidden_dim = padded_hidden + + if gated: + # Interleave up/gate rows for the fused gated activation. w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) @@ -371,8 +397,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( # Shuffle weights and scales for FI TRTLLM NVFP4 MoE kernels. if backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: - w13, w13_scale, w2, w2_scale, padded_hidden = ( - align_trtllm_fp4_moe_hidden_dim_for_fi(w13, w13_scale, w2, w2_scale) + w13, w13_scale, w2, w2_scale, padded_hidden = align_fp4_moe_hidden_dim_for_fi( + w13, w13_scale, w2, w2_scale ) if layer.moe_config.hidden_dim_unpadded is None: layer.moe_config.hidden_dim_unpadded = layer.moe_config.hidden_dim diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 1bfd7e61f9..be28de9b66 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -13,6 +13,8 @@ if TYPE_CHECKING: logger = init_logger(__name__) +NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE = 1.0 / (448.0 * 6.0) + def activation_to_flashinfer_int(activation: MoEActivation) -> int: return activation_to_flashinfer_type(activation).value @@ -50,6 +52,20 @@ def activation_to_flashinfer_type(activation: MoEActivation) -> "ActivationType" return ACTIVATION_TO_FI_ACTIVATION[activation] +def quantize_nvfp4_per_token_input( + hidden_states: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize NVFP4 activations with one FP32 decode scale per token.""" + from flashinfer import SfLayout, nvfp4_quantize + + return nvfp4_quantize( + hidden_states, + NVFP4_PER_TOKEN_BASE_GLOBAL_SCALE, + sfLayout=SfLayout.layout_linear, + per_token_activation=True, + ) + + def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor: return ( x.reshape(-1, 2, x.shape[-2] // 2, x.shape[-1]).flip(dims=[1]).reshape(x.shape) @@ -264,7 +280,7 @@ def align_fp4_moe_weights_for_fi( return padded_w13, padded_w13_scale, padded_w2, padded_w2_scale, padded_intermediate -def align_trtllm_fp4_moe_hidden_dim_for_fi( +def align_fp4_moe_hidden_dim_for_fi( w13: torch.Tensor, w13_scale: torch.Tensor, w2: torch.Tensor, @@ -279,7 +295,7 @@ def align_trtllm_fp4_moe_hidden_dim_for_fi( return w13, w13_scale, w2, w2_scale, hidden_size logger.warning_once( - "Padding hidden size from %d to %d for TRTLLM NVFP4 MoE weights. " + "Padding hidden size from %d to %d for FlashInfer NVFP4 MoE weights. " "This requires activation slicing at runtime and may cause " "performance degradation.", hidden_size,