a couple mypy edits

vllm/worker/cpu_worker.py

@@ -333,9 +333,8 @@ def execute_worker(
     def prepare_worker_input(
             self, execute_model_req: ExecuteModelRequest) -> WorkerInput:
         assert execute_model_req is not None
-        virtual_engine = execute_model_req.virtual_engine
+        virtual_engine: int = execute_model_req.virtual_engine
         num_seq_groups: int = len(execute_model_req.seq_group_metadata_list)
-        blocks_to_copy = execute_model_req.blocks_to_copy
         blocks_to_copy = torch.tensor(execute_model_req.blocks_to_copy,
                                       device="cpu",
                                       dtype=torch.int64).view(-1, 2)

vllm/worker/multi_step_model_runner.py

@@ -407,7 +407,7 @@ def _async_process_outputs(self, model_input: StatefulModelInput,
                 break
 
     def _final_process_outputs(self, model_input: StatefulModelInput,
-                               output_proc_callback: Optional[Callable]):
+            output_proc_callback: Optional[Callable]) -> List[SamplerOutput]:
         assert model_input.frozen_model_input is not None
 
         has_async_callback = output_proc_callback is not None
@@ -594,8 +594,9 @@ def execute_model(
         # should be [SamplerOutput]
         return output
 
-    def _update_sampling_metadata(self, sampling_metadata, num_seqs,
-                                  num_queries):
+    def _update_sampling_metadata(self, sampling_metadata: SamplingMetadata, 
+                                  num_seqs: Optional[int],
+                                  num_queries: int):
 
         assert sampling_metadata.num_prompts == 0
         assert len(sampling_metadata.seq_groups) == num_queries
@@ -858,7 +859,7 @@ def _pythonize_sampler_output(
             seq_outputs: List[
                 SequenceOutput] = completion_seq_group_output.samples
         else:
-            seq_outputs = []
+            seq_outputs: List[SequenceOutput] = []
 
         for tdx, (parent_id,
                   next_token_id) in enumerate(zip(parent_ids, next_token_ids)):

vllm/worker/worker_base.py

@@ -452,7 +452,7 @@ def init_worker(self, *args, **kwargs):
         self.worker = worker_class(*args, **kwargs)
         assert self.worker is not None
 
-    def execute_method(self, method, *args, **kwargs):
+    def execute_method(self, method: str, *args, **kwargs):
         try:
             target = self if self.worker is None else self.worker
             executor = getattr(target, method)