diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index 621b33dd8f3..376ef3cee54 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -413,6 +413,7 @@ def batched_fused_marlin_moe( is_k_full: bool = True, output: torch.Tensor | None = None, inplace: bool = False, + clamp_limit: float | None = None, ) -> torch.Tensor: """ This function massages the inputs so the batched hidden_states can be @@ -535,6 +536,7 @@ def batched_fused_marlin_moe( intermediate_cache2=intermediate_cache2, output=output.view(-1, K) if output is not None else output, is_k_full=is_k_full, + clamp_limit=clamp_limit, ) output = output.view(B, BATCH_TOKENS_MAX, K) @@ -768,6 +770,7 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase): sort_indices2=self.w2_g_idx_sort_indices, is_k_full=self.is_k_full, input_dtype=self.input_dtype, + clamp_limit=self.gemm1_clamp_limit, ) return @@ -970,4 +973,5 @@ class BatchedMarlinExperts(MarlinExpertsBase): sort_indices1=self.w13_g_idx_sort_indices, sort_indices2=self.w2_g_idx_sort_indices, is_k_full=self.is_k_full, + clamp_limit=self.gemm1_clamp_limit, )