diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index 2653b711d9..bbe6277220 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -23,7 +23,7 @@ from tests.kernels.utils import torch_moe from vllm import _custom_ops as ops from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.custom_op import CustomOp, op_registry -from vllm.model_executor.layers.activation import SiluAndMulWithClamp +from vllm.model_executor.layers.activation import SiluAndMulWithClamp, SituAndMul 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 ( @@ -63,6 +63,7 @@ MNK_FACTORS = [ _SWIGLU_LIMIT = 0.1 _LARGE_OUTPUT1_SCALE = 32768.0 _CLAMP_OP_NAME = "test_silu_and_mul_with_clamp" +_SITU_OP_NAME = "test_situ_and_mul" # Test-only fixed-limit clamp. ``custom_op_name`` makes the class itself # valid as an ``activation=`` argument to ``torch_moe`` (which only looks @@ -78,12 +79,24 @@ if _CLAMP_OP_NAME not in op_registry: super().__init__(_SWIGLU_LIMIT, compile_native=compile_native) +if _SITU_OP_NAME not in op_registry: + + @CustomOp.register(_SITU_OP_NAME) + class _SituAndMulTest(SituAndMul): + custom_op_name = _SITU_OP_NAME + + def __init__(self, *, compile_native: bool = True) -> None: + super().__init__(4.0, 25.0, compile_native=compile_native) + + SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME] +SITU = op_registry[_SITU_OP_NAME] ACTIVATION_CASES = [ pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"), pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"), + pytest.param(MoEActivation.SITU, SITU, None, id="situ"), pytest.param( MoEActivation.RELU2_NO_MUL, MoEActivation.RELU2_NO_MUL, @@ -108,7 +121,7 @@ def test_trtllm_fp4_moe_no_graph( topk: int, dtype: torch.dtype, activation: MoEActivation, - torch_activation: MoEActivation | type[SiluAndMulWithClamp], + torch_activation: MoEActivation | type[SiluAndMulWithClamp] | type[SituAndMul], swiglu_limit: float | None, workspace_init, ): @@ -173,6 +186,10 @@ def test_trtllm_fp4_moe_no_graph( in_dtype=dtype, routing_method=RoutingMethodType.TopK, max_num_tokens=next_power_of_2(m), + activation_situ_beta=4.0 if activation == MoEActivation.SITU else None, + activation_situ_linear_beta=( + 25.0 if activation == MoEActivation.SITU else None + ), ) trtllm_inner = TrtLlmNvFp4ExpertsModular( @@ -191,6 +208,8 @@ def test_trtllm_fp4_moe_no_graph( fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas) fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas) trtllm_inner.process_weights_after_loading(fake_layer) + if activation == MoEActivation.SITU: + torch.testing.assert_close(trtllm_inner.g1_scale_c, quant_config.a2_gscale) trtllm_experts = mk.FusedMoEKernel( maybe_make_prepare_finalize( diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 41cedca9d3..7225642259 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -23,6 +23,7 @@ from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( ModelOptFp8Config, ModelOptFp8LinearMethod, + ModelOptFp8PbWoLinearMethod, ModelOptMixedPrecisionConfig, ModelOptMxFp8Config, ModelOptNvFp4Config, @@ -137,6 +138,64 @@ def test_modelopt_fp8_updates_weight_dims_after_transpose(): method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer) +def test_modelopt_fp8_pb_wo_hides_output_padding(default_vllm_config): + config = ModelOptFp8Config( + quant_method="FP8_PB_WO", + is_checkpoint_fp8_serialized=True, + kv_cache_quant_method=None, + exclude_modules=[], + ) + default_vllm_config.model_config = Mock(dtype=torch.bfloat16) + kernel = Mock() + + with ( + set_current_vllm_config(default_vllm_config), + patch( + "vllm.model_executor.layers.quantization.modelopt.init_fp8_linear_kernel", + return_value=kernel, + ), + patch( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", + return_value=0, + ), + patch( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", + return_value=1, + ), + ): + method = ModelOptFp8PbWoLinearMethod(config) + layer = torch.nn.Module() + method.create_weights( + layer, + input_size_per_partition=128, + output_partition_sizes=[128, 65], + input_size=128, + output_size=193, + params_dtype=torch.bfloat16, + ) + + assert layer.weight.shape == (193, 128) + assert layer.weight_scale.shape == (2, 1, 1, 1) + + layer.weight.data.fill_(1) + method.process_weights_after_loading(layer) + + assert layer.weight.shape == (256, 128) + assert torch.count_nonzero(layer.weight[193:].float()) == 0 + kernel.process_weights_after_loading.assert_called_once_with(layer) + + physical_output = torch.randn(2, 256, dtype=torch.bfloat16) + kernel.apply_weights.return_value = physical_output + bias = torch.randn(193, dtype=torch.bfloat16) + output = method.apply(layer, torch.randn(2, 128), bias) + + torch.testing.assert_close(output, physical_output[:, :193] + bias) + assert output.shape == (2, 193) + assert output.is_contiguous() + kernel.apply_weights.assert_called_once() + assert kernel.apply_weights.call_args.args[2] is None + + def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized(): config = ModelOptNvFp4Config( is_checkpoint_nvfp4_serialized=True, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index b41c810d76..fe732e2850 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -371,8 +371,17 @@ def _get_kv_b_proj_input_dtype( return kv_b_proj.params_dtype if weight_dtype == torch.uint8: return None - if weight_dtype == current_platform.fp8_dtype() and not use_fp8_prefill: - return None + if weight_dtype == current_platform.fp8_dtype(): + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptFp8PbWoLinearMethod, + ) + + quant_method = getattr(kv_b_proj, "quant_method", None) + # FP8_PB_WO dynamically quantizes BF16/FP16 inputs in the linear method. + if isinstance(quant_method, ModelOptFp8PbWoLinearMethod): + return quant_method.input_dtype + if not use_fp8_prefill: + return None return weight_dtype 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 29ccd5770e..0bff8264f7 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 @@ -70,16 +70,9 @@ class TrtLlmNvFp4ExpertsBase: ) self.local_num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank + self.is_situ = moe_config.activation == MoEActivation.SITU - assert self.quant_config.g1_alphas is not None - assert self.quant_config.a2_gscale is not None - if moe_config.is_act_and_mul: - # g1_alpha_s = a13_scale * w13_scale_2 - # a2_gscale = (1 / a2_scale) - # g1_scale_c = a13_scale * w13_scale_2 / a2_scale - self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale - else: - self.g1_scale_c = self.quant_config.a2_gscale.clone() + self.g1_scale_c = self._compute_g1_scale_c() # Fall back to moe_config.swiglu_* when quant_config doesn't carry them # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config). @@ -120,7 +113,6 @@ class TrtLlmNvFp4ExpertsBase: # (gemm1_alpha) and situ linear_beta -> gatedActBeta (gemm1_beta). # These operate on the dequantized gate/up, so they are NOT folded by # g1_alphas in process_weights_after_loading. - self.is_situ = moe_config.activation == MoEActivation.SITU if self.is_situ: situ_beta = moe_config.activation_situ_beta situ_linear_beta = moe_config.activation_situ_linear_beta @@ -143,18 +135,28 @@ class TrtLlmNvFp4ExpertsBase: clamp, ) + def _compute_g1_scale_c(self) -> torch.Tensor: + assert self.quant_config.g1_alphas is not None + assert self.quant_config.a2_gscale is not None + if not self.moe_config.is_act_and_mul: + return self.quant_config.a2_gscale.clone() + if self.is_situ: + # SITU applies its nonlinear activation after g1_alphas, so only + # the output quantization factor belongs in g1_scale_c. + return self.quant_config.a2_gscale.clone() + + # g1_alphas = a13_scale * w13_scale_2 + # a2_gscale = 1 / a2_scale + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale + return self.quant_config.g1_alphas * self.quant_config.a2_gscale + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) # Recompute g1_scale_c since g1_alphas was just fused in-place. # Register as a layer parameter so EPLB rearranges it alongside # other expert weights. - assert self.quant_config.g1_alphas is not None - assert self.quant_config.a2_gscale is not None - if self.moe_config.is_act_and_mul: - g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale - else: - g1_scale_c = self.quant_config.a2_gscale.clone() + g1_scale_c = self._compute_g1_scale_c() layer.register_parameter( "g1_scale_c", torch.nn.Parameter(g1_scale_c, requires_grad=False), diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e24c5656ad..185a55d00b 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -624,6 +624,8 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase): where block size is typically 128 for both dims. vLLM executes it as FP8 GEMM with *dynamic per-token* activation quant. + Output widths that are not block-aligned are padded and restored + to their logical width before returning to the model. """ _WEIGHT_BLOCK_SIZE: tuple[int, int] = (128, 128) @@ -670,6 +672,12 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase): # element-space -> block-space for BlockQuantScaleParameter. layer.weight_block_size = self.weight_block_size + block_n, block_k = self._WEIGHT_BLOCK_SIZE + remainder = output_size_per_partition % block_n + self.output_padding = 0 if remainder == 0 else block_n - remainder + self.logical_output_size = output_size_per_partition + physical_output_size = output_size_per_partition + self.output_padding + weight = ModelWeightParameter( data=torch.empty( output_size_per_partition, @@ -682,19 +690,13 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase): ) layer.register_parameter("weight", weight) - block_n, block_k = self._WEIGHT_BLOCK_SIZE - if output_size_per_partition % block_n != 0: - raise ValueError( - "ModelOpt FP8_PB_WO requires out_features divisible by " - f"{block_n}, got {output_size_per_partition}." - ) if input_size_per_partition % block_k != 0: raise ValueError( "ModelOpt FP8_PB_WO requires in_features divisible by " f"{block_k}, got {input_size_per_partition}." ) - out_blks = output_size_per_partition // block_n + out_blks = physical_output_size // block_n in_blks = input_size_per_partition // block_k # Match ModelOpt's exported shape so weight loading works without a @@ -711,7 +713,7 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase): self.w8a8_block_fp8_linear = init_fp8_linear_kernel( activation_quant_key=self.activation_quant_key, weight_quant_key=self.weight_quant_key, - weight_shape=layer.weight.shape, + weight_shape=(physical_output_size, input_size_per_partition), input_dtype=self.input_dtype, out_dtype=self.out_dtype, module_name=self.__class__.__name__, @@ -719,7 +721,15 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Keep weight in [out, in] layout for Fp8BlockScaledMMLinearKernel. - layer.weight = Parameter(layer.weight.data, requires_grad=False) + weight = layer.weight.data + if self.output_padding: + padded_weight = weight.new_zeros( + self.logical_output_size + self.output_padding, + weight.shape[1], + ) + padded_weight[: self.logical_output_size].copy_(weight) + weight = padded_weight + layer.weight = Parameter(weight, requires_grad=False) scale = layer.weight_scale if scale.dim() == 4: @@ -733,8 +743,7 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase): layer.weight_scale = Parameter(scale.contiguous(), requires_grad=False) - if hasattr(self, "fp8_linear"): - self.fp8_linear.process_weights_after_loading(layer) + self.w8a8_block_fp8_linear.process_weights_after_loading(layer) def apply( self, @@ -742,7 +751,15 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: - return self.w8a8_block_fp8_linear.apply_weights(layer, x, bias) + kernel_bias = None if self.output_padding else bias + output = self.w8a8_block_fp8_linear.apply_weights(layer, x, kernel_bias) + if not self.output_padding: + return output + + output = output[..., : self.logical_output_size].contiguous() + if bias is not None: + output.add_(bias) + return output class ModelOptFp8MoEMethod(FusedMoEMethodBase): @@ -2390,6 +2407,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): if isinstance(layer, (LinearBase, ParallelLMHead)): if quant_algo == "FP8": return ModelOptFp8LinearMethod(self.fp8_config) + if quant_algo == "FP8_PB_WO": + return ModelOptFp8PbWoLinearMethod(self.fp8_config) if quant_algo == "NVFP4": return ModelOptNvFp4LinearMethod(self.nvfp4_config) if quant_algo == "W4A16_NVFP4": diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index fe49610094..ec6b514f88 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -36,7 +36,7 @@ from vllm.model_executor.model_loader.weight_utils import ( default_weight_loader, sharded_weight_loader, ) -from vllm.model_executor.parameter import BasevLLMParameter +from vllm.model_executor.parameter import BasevLLMParameter, BlockQuantScaleParameter from vllm.model_executor.utils import set_weight_attrs from vllm.models.kimi_k3.nvidia.kda_metadata import ( KimiK3KDAAttentionBackend, @@ -102,7 +102,11 @@ class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): ) -> None: tp_rank = self.tp_rank param_tp_rank = getattr(param, "tp_rank", None) - if loaded_shard_id == self.replicated_shard_id: + replicate_block_scale = ( + isinstance(param, BlockQuantScaleParameter) + and loaded_weight.shape[param.output_dim] < self.tp_size + ) + if loaded_shard_id == self.replicated_shard_id or replicate_block_scale: self.tp_rank = 0 if param_tp_rank is not None: param.tp_rank = 0 @@ -121,7 +125,11 @@ class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear): ) -> None: tp_rank = self.tp_rank param_tp_rank = getattr(param, "tp_rank", None) - if loaded_shard_id == self.replicated_shard_id: + replicate_block_scale = ( + isinstance(param, BlockQuantScaleParameter) + and loaded_weight.shape[param.output_dim] < self.tp_size + ) + if loaded_shard_id == self.replicated_shard_id or replicate_block_scale: self.tp_rank = 0 if param_tp_rank is not None: param.tp_rank = 0