diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 8737d9e7372..1eeca142343 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -1573,14 +1573,14 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod): def apply_monolithic( self, - layer: torch.nn.Module, + layer: FusedMoE, x: torch.Tensor, router_logits: torch.Tensor, - expert_map: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: if layer.enable_eplb: raise NotImplementedError( - "EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet." + f"EPLB not supported for {self.__class__.__name__} yet." ) from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501 @@ -1597,7 +1597,7 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod): topk=layer.top_k, renormalize=layer.renormalize, global_num_experts=layer.global_num_experts, - expert_map=expert_map, + expert_map=layer.expert_map, quant_config=self.moe_quant_config, apply_router_weight_on_input=layer.apply_router_weight_on_input, unpadded_N_w1=self.moe.intermediate_size_per_partition_unpadded * 2,