Support kimi k3 nvfp4 checkpoint (#53132)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>

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(

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,

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
 
 

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),

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":

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