@tduncan / Nvidia-Optimization-Stack-VLLM / commits / f8e0602713

Support kimi k3 nvfp4 checkpoint (#53132)

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

Wei Zhao committed Aug 21, 2026 at 01:59 UTC f8e060271381e352a4eabec97f1839a32e41ec41

6 files changed +151 -35

tests/kernels/moe/test_trtllm_nvfp4_moe.py

+21 -2

@@ -23,7 +23,7 @@ from tests.kernels.utils import torch_moe

23 from vllm import _custom_ops as ops

24 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

25 from vllm.model_executor.custom_op import CustomOp, op_registry

26 -from vllm.model_executor.layers.activation import SiluAndMulWithClamp

26 +from vllm.model_executor.layers.activation import SiluAndMulWithClamp, SituAndMul

27 from vllm.model_executor.layers.fused_moe import fused_topk

28 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

29 from vllm.model_executor.layers.fused_moe.all2all_utils import (

@@ -63,6 +63,7 @@ MNK_FACTORS = [

63 _SWIGLU_LIMIT = 0.1

64 _LARGE_OUTPUT1_SCALE = 32768.0

65 _CLAMP_OP_NAME = "test_silu_and_mul_with_clamp"

66 +_SITU_OP_NAME = "test_situ_and_mul"

68 # Test-only fixed-limit clamp. ``custom_op_name`` makes the class itself

69 # valid as an ``activation=`` argument to ``torch_moe`` (which only looks

@@ -78,12 +79,24 @@ if _CLAMP_OP_NAME not in op_registry:

79 super().__init__(_SWIGLU_LIMIT, compile_native=compile_native)

82 +if _SITU_OP_NAME not in op_registry:

83 +

84 + @CustomOp.register(_SITU_OP_NAME)

85 + class _SituAndMulTest(SituAndMul):

86 + custom_op_name = _SITU_OP_NAME

87 +

88 + def __init__(self, *, compile_native: bool = True) -> None:

89 + super().__init__(4.0, 25.0, compile_native=compile_native)

90 +

91 +

92 SILU_WITH_CLAMP = op_registry[_CLAMP_OP_NAME]

93 +SITU = op_registry[_SITU_OP_NAME]

96 ACTIVATION_CASES = [

97 pytest.param(MoEActivation.SILU, MoEActivation.SILU, None, id="silu"),

98 pytest.param(MoEActivation.SILU, SILU_WITH_CLAMP, _SWIGLU_LIMIT, id="silu_clamp"),

99 + pytest.param(MoEActivation.SITU, SITU, None, id="situ"),

100 pytest.param(

101 MoEActivation.RELU2_NO_MUL,

102 MoEActivation.RELU2_NO_MUL,

@@ -108,7 +121,7 @@ def test_trtllm_fp4_moe_no_graph(

121 topk: int,

122 dtype: torch.dtype,

123 activation: MoEActivation,

111 - torch_activation: MoEActivation | type[SiluAndMulWithClamp],

124 + torch_activation: MoEActivation | type[SiluAndMulWithClamp] | type[SituAndMul],

125 swiglu_limit: float | None,

126 workspace_init,

127 ):

@@ -173,6 +186,10 @@ def test_trtllm_fp4_moe_no_graph(

186 in_dtype=dtype,

187 routing_method=RoutingMethodType.TopK,

188 max_num_tokens=next_power_of_2(m),

189 + activation_situ_beta=4.0 if activation == MoEActivation.SITU else None,

190 + activation_situ_linear_beta=(

191 + 25.0 if activation == MoEActivation.SITU else None

192 + ),

193 )

194

195 trtllm_inner = TrtLlmNvFp4ExpertsModular(

@@ -191,6 +208,8 @@ def test_trtllm_fp4_moe_no_graph(

208 fake_layer.w13_input_scale = torch.ones_like(quant_config.g1_alphas)

209 fake_layer.w2_input_scale = torch.ones_like(quant_config.g2_alphas)

210 trtllm_inner.process_weights_after_loading(fake_layer)

211 + if activation == MoEActivation.SITU:

212 + torch.testing.assert_close(trtllm_inner.g1_scale_c, quant_config.a2_gscale)

213

214 trtllm_experts = mk.FusedMoEKernel(

215 maybe_make_prepare_finalize(

tests/quantization/test_modelopt.py

+59

@@ -23,6 +23,7 @@ from vllm.model_executor.layers.linear import UnquantizedLinearMethod

23 from vllm.model_executor.layers.quantization.modelopt import (

24 ModelOptFp8Config,

25 ModelOptFp8LinearMethod,

26 + ModelOptFp8PbWoLinearMethod,

27 ModelOptMixedPrecisionConfig,

28 ModelOptMxFp8Config,

29 ModelOptNvFp4Config,

@@ -137,6 +138,64 @@ def test_modelopt_fp8_updates_weight_dims_after_transpose():

138 method.fp8_linear.process_weights_after_loading.assert_called_once_with(layer)

139

140

141 +def test_modelopt_fp8_pb_wo_hides_output_padding(default_vllm_config):

142 + config = ModelOptFp8Config(

143 + quant_method="FP8_PB_WO",

144 + is_checkpoint_fp8_serialized=True,

145 + kv_cache_quant_method=None,

146 + exclude_modules=[],

147 + )

148 + default_vllm_config.model_config = Mock(dtype=torch.bfloat16)

149 + kernel = Mock()

150 +

151 + with (

152 + set_current_vllm_config(default_vllm_config),

153 + patch(

154 + "vllm.model_executor.layers.quantization.modelopt.init_fp8_linear_kernel",

155 + return_value=kernel,

156 + ),

157 + patch(

158 + "vllm.model_executor.parameter.get_tensor_model_parallel_rank",

159 + return_value=0,

160 + ),

161 + patch(

162 + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",

163 + return_value=1,

164 + ),

165 + ):

166 + method = ModelOptFp8PbWoLinearMethod(config)

167 + layer = torch.nn.Module()

168 + method.create_weights(

169 + layer,

170 + input_size_per_partition=128,

171 + output_partition_sizes=[128, 65],

172 + input_size=128,

173 + output_size=193,

174 + params_dtype=torch.bfloat16,

175 + )

176 +

177 + assert layer.weight.shape == (193, 128)

178 + assert layer.weight_scale.shape == (2, 1, 1, 1)

179 +

180 + layer.weight.data.fill_(1)

181 + method.process_weights_after_loading(layer)

182 +

183 + assert layer.weight.shape == (256, 128)

184 + assert torch.count_nonzero(layer.weight[193:].float()) == 0

185 + kernel.process_weights_after_loading.assert_called_once_with(layer)

186 +

187 + physical_output = torch.randn(2, 256, dtype=torch.bfloat16)

188 + kernel.apply_weights.return_value = physical_output

189 + bias = torch.randn(193, dtype=torch.bfloat16)

190 + output = method.apply(layer, torch.randn(2, 128), bias)

191 +

192 + torch.testing.assert_close(output, physical_output[:, :193] + bias)

193 + assert output.shape == (2, 193)

194 + assert output.is_contiguous()

195 + kernel.apply_weights.assert_called_once()

196 + assert kernel.apply_weights.call_args.args[2] is None

197 +

198 +

199 def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():

200 config = ModelOptNvFp4Config(

201 is_checkpoint_nvfp4_serialized=True,

vllm/model_executor/layers/attention/mla_attention.py

+11 -2

@@ -371,8 +371,17 @@ def _get_kv_b_proj_input_dtype(

371 return kv_b_proj.params_dtype

372 if weight_dtype == torch.uint8:

373 return None

374 - if weight_dtype == current_platform.fp8_dtype() and not use_fp8_prefill:

375 - return None

374 + if weight_dtype == current_platform.fp8_dtype():

375 + from vllm.model_executor.layers.quantization.modelopt import (

376 + ModelOptFp8PbWoLinearMethod,

377 + )

378 +

379 + quant_method = getattr(kv_b_proj, "quant_method", None)

380 + # FP8_PB_WO dynamically quantizes BF16/FP16 inputs in the linear method.

381 + if isinstance(quant_method, ModelOptFp8PbWoLinearMethod):

382 + return quant_method.input_dtype

383 + if not use_fp8_prefill:

384 + return None

385 return weight_dtype

386

387

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

+18 -16

@@ -70,16 +70,9 @@ class TrtLlmNvFp4ExpertsBase:

70 )

71 self.local_num_experts = moe_config.num_local_experts

72 self.ep_rank = moe_config.moe_parallel_config.ep_rank

73 + self.is_situ = moe_config.activation == MoEActivation.SITU

74 - assert self.quant_config.g1_alphas is not None

75 - assert self.quant_config.a2_gscale is not None

76 - if moe_config.is_act_and_mul:

77 - # g1_alpha_s = a13_scale * w13_scale_2

78 - # a2_gscale = (1 / a2_scale)

79 - # g1_scale_c = a13_scale * w13_scale_2 / a2_scale

80 - self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale

81 - else:

82 - self.g1_scale_c = self.quant_config.a2_gscale.clone()

75 + self.g1_scale_c = self._compute_g1_scale_c()

77 # Fall back to moe_config.swiglu_* when quant_config doesn't carry them

78 # (ModelOpt NVFP4 checkpoints store these on moe_config, not quant_config).

@@ -120,7 +113,6 @@ class TrtLlmNvFp4ExpertsBase:

113 # (gemm1_alpha) and situ linear_beta -> gatedActBeta (gemm1_beta).

114 # These operate on the dequantized gate/up, so they are NOT folded by

115 # g1_alphas in process_weights_after_loading.

123 - self.is_situ = moe_config.activation == MoEActivation.SITU

116 if self.is_situ:

117 situ_beta = moe_config.activation_situ_beta

118 situ_linear_beta = moe_config.activation_situ_linear_beta

@@ -143,18 +135,28 @@ class TrtLlmNvFp4ExpertsBase:

135 clamp,

136 )

137

138 + def _compute_g1_scale_c(self) -> torch.Tensor:

139 + assert self.quant_config.g1_alphas is not None

140 + assert self.quant_config.a2_gscale is not None

141 + if not self.moe_config.is_act_and_mul:

142 + return self.quant_config.a2_gscale.clone()

143 + if self.is_situ:

144 + # SITU applies its nonlinear activation after g1_alphas, so only

145 + # the output quantization factor belongs in g1_scale_c.

146 + return self.quant_config.a2_gscale.clone()

147 +

148 + # g1_alphas = a13_scale * w13_scale_2

149 + # a2_gscale = 1 / a2_scale

150 + # g1_scale_c = a13_scale * w13_scale_2 / a2_scale

151 + return self.quant_config.g1_alphas * self.quant_config.a2_gscale

152 +

153 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

154 layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)

155 layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)

156 # Recompute g1_scale_c since g1_alphas was just fused in-place.

157 # Register as a layer parameter so EPLB rearranges it alongside

158 # other expert weights.

152 - assert self.quant_config.g1_alphas is not None

153 - assert self.quant_config.a2_gscale is not None

154 - if self.moe_config.is_act_and_mul:

155 - g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale

156 - else:

157 - g1_scale_c = self.quant_config.a2_gscale.clone()

159 + g1_scale_c = self._compute_g1_scale_c()

160 layer.register_parameter(

161 "g1_scale_c",

162 torch.nn.Parameter(g1_scale_c, requires_grad=False),

vllm/model_executor/layers/quantization/modelopt.py

+31 -12

@@ -624,6 +624,8 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):

624 where block size is typically 128 for both dims.

625

626 vLLM executes it as FP8 GEMM with *dynamic per-token* activation quant.

627 + Output widths that are not block-aligned are padded and restored

628 + to their logical width before returning to the model.

629 """

630

631 _WEIGHT_BLOCK_SIZE: tuple[int, int] = (128, 128)

@@ -670,6 +672,12 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):

672 # element-space -> block-space for BlockQuantScaleParameter.

673 layer.weight_block_size = self.weight_block_size

674

675 + block_n, block_k = self._WEIGHT_BLOCK_SIZE

676 + remainder = output_size_per_partition % block_n

677 + self.output_padding = 0 if remainder == 0 else block_n - remainder

678 + self.logical_output_size = output_size_per_partition

679 + physical_output_size = output_size_per_partition + self.output_padding

680 +

681 weight = ModelWeightParameter(

682 data=torch.empty(

683 output_size_per_partition,

@@ -682,19 +690,13 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):

690 )

691 layer.register_parameter("weight", weight)

692

685 - block_n, block_k = self._WEIGHT_BLOCK_SIZE

686 - if output_size_per_partition % block_n != 0:

687 - raise ValueError(

688 - "ModelOpt FP8_PB_WO requires out_features divisible by "

689 - f"{block_n}, got {output_size_per_partition}."

690 - )

693 if input_size_per_partition % block_k != 0:

694 raise ValueError(

695 "ModelOpt FP8_PB_WO requires in_features divisible by "

696 f"{block_k}, got {input_size_per_partition}."

697 )

698

697 - out_blks = output_size_per_partition // block_n

699 + out_blks = physical_output_size // block_n

700 in_blks = input_size_per_partition // block_k

701

702 # Match ModelOpt's exported shape so weight loading works without a

@@ -711,7 +713,7 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):

713 self.w8a8_block_fp8_linear = init_fp8_linear_kernel(

714 activation_quant_key=self.activation_quant_key,

715 weight_quant_key=self.weight_quant_key,

714 - weight_shape=layer.weight.shape,

716 + weight_shape=(physical_output_size, input_size_per_partition),

717 input_dtype=self.input_dtype,

718 out_dtype=self.out_dtype,

719 module_name=self.__class__.__name__,

@@ -719,7 +721,15 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):

721

722 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

723 # Keep weight in [out, in] layout for Fp8BlockScaledMMLinearKernel.

722 - layer.weight = Parameter(layer.weight.data, requires_grad=False)

724 + weight = layer.weight.data

725 + if self.output_padding:

726 + padded_weight = weight.new_zeros(

727 + self.logical_output_size + self.output_padding,

728 + weight.shape[1],

729 + )

730 + padded_weight[: self.logical_output_size].copy_(weight)

731 + weight = padded_weight

732 + layer.weight = Parameter(weight, requires_grad=False)

733

734 scale = layer.weight_scale

735 if scale.dim() == 4:

@@ -733,8 +743,7 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):

743

744 layer.weight_scale = Parameter(scale.contiguous(), requires_grad=False)

745

736 - if hasattr(self, "fp8_linear"):

737 - self.fp8_linear.process_weights_after_loading(layer)

746 + self.w8a8_block_fp8_linear.process_weights_after_loading(layer)

747

748 def apply(

749 self,

@@ -742,7 +751,15 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):

751 x: torch.Tensor,

752 bias: torch.Tensor | None = None,

753 ) -> torch.Tensor:

745 - return self.w8a8_block_fp8_linear.apply_weights(layer, x, bias)

754 + kernel_bias = None if self.output_padding else bias

755 + output = self.w8a8_block_fp8_linear.apply_weights(layer, x, kernel_bias)

756 + if not self.output_padding:

757 + return output

758 +

759 + output = output[..., : self.logical_output_size].contiguous()

760 + if bias is not None:

761 + output.add_(bias)

762 + return output

763

764

765 class ModelOptFp8MoEMethod(FusedMoEMethodBase):

@@ -2390,6 +2407,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):

2407 if isinstance(layer, (LinearBase, ParallelLMHead)):

2408 if quant_algo == "FP8":

2409 return ModelOptFp8LinearMethod(self.fp8_config)

2410 + if quant_algo == "FP8_PB_WO":

2411 + return ModelOptFp8PbWoLinearMethod(self.fp8_config)

2412 if quant_algo == "NVFP4":

2413 return ModelOptNvFp4LinearMethod(self.nvfp4_config)

2414 if quant_algo == "W4A16_NVFP4":

vllm/models/kimi_k3/nvidia/kda.py

+11 -3

@@ -36,7 +36,7 @@ from vllm.model_executor.model_loader.weight_utils import (

36 default_weight_loader,

37 sharded_weight_loader,

38 )

39 -from vllm.model_executor.parameter import BasevLLMParameter

39 +from vllm.model_executor.parameter import BasevLLMParameter, BlockQuantScaleParameter

40 from vllm.model_executor.utils import set_weight_attrs

41 from vllm.models.kimi_k3.nvidia.kda_metadata import (

42 KimiK3KDAAttentionBackend,

@@ -102,7 +102,11 @@ class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear):

102 ) -> None:

103 tp_rank = self.tp_rank

104 param_tp_rank = getattr(param, "tp_rank", None)

105 - if loaded_shard_id == self.replicated_shard_id:

105 + replicate_block_scale = (

106 + isinstance(param, BlockQuantScaleParameter)

107 + and loaded_weight.shape[param.output_dim] < self.tp_size

108 + )

109 + if loaded_shard_id == self.replicated_shard_id or replicate_block_scale:

110 self.tp_rank = 0

111 if param_tp_rank is not None:

112 param.tp_rank = 0

@@ -121,7 +125,11 @@ class _KimiGDNMergedColumnParallelLinear(MergedColumnParallelLinear):

125 ) -> None:

126 tp_rank = self.tp_rank

127 param_tp_rank = getattr(param, "tp_rank", None)

124 - if loaded_shard_id == self.replicated_shard_id:

128 + replicate_block_scale = (

129 + isinstance(param, BlockQuantScaleParameter)

130 + and loaded_weight.shape[param.output_dim] < self.tp_size

131 + )

132 + if loaded_shard_id == self.replicated_shard_id or replicate_block_scale:

133 self.tp_rank = 0

134 if param_tp_rank is not None:

135 param.tp_rank = 0