format code

Signed-off-by: Kurumi5210 <[email protected]>

Author: Kurumi5210

vllm_ascend/distributed/mooncake_connector.py

@@ -40,7 +40,6 @@
 from vllm_ascend.distributed.mooncake.transfer_engine import get_global_te
 from vllm_ascend.distributed.utils import get_transfer_timeout_value
 from vllm_ascend.utils import prefill_context_parallel_enable, vllm_version_is
-from typing import List
 
 # isort: off
 if prefill_context_parallel_enable():
@@ -1289,21 +1288,22 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata):
                         self.kv_recv_thread.add_request(
                             request_id=req_id,
                             local_block_ids=local_block_ids_list[pcp_dcp_rank],
-                            remote_block_ids=remote_block_ids_list[pcp_dcp_rank],
+                            remote_block_ids=remote_block_ids_list[
+                                pcp_dcp_rank],
                             remote_engine_id=meta.remote_engine_id,
                             remote_host=meta.remote_host,
                             remote_handshake_port=remote_handshake_port_list[
                                 pcp_dcp_rank][i],
                             offset=i,
                             tp_num_need_pulls=self.tp_num_need_pulls,
-                            all_task_done=(pcp_dcp_rank
-                                        == len(remote_handshake_port_list) - 1
-                                        and i == self.tp_num_need_pulls - 1 ))
-            else:   #TODO: support prefill context parallel and pipeline parallel open at the same time
+                            all_task_done=(
+                                pcp_dcp_rank
+                                == len(remote_handshake_port_list) - 1
+                                and i == self.tp_num_need_pulls - 1))
+            else:  #TODO: support prefill context parallel and pipeline parallel open at the same time
                 choosen_rank_list = self._get_remote_tp_rank(req_id)
-                remote_handshake_port_list = [
-                    [x + meta.remote_port] for x in choosen_rank_list
-                ]
+                remote_handshake_port_list = [[x + meta.remote_port]
+                                              for x in choosen_rank_list]
                 for i in range(self.tp_num_need_pulls * self._prefill_pp_size):
                     assert self.kv_recv_thread is not None
                     self.kv_recv_thread.add_request(
@@ -1315,7 +1315,8 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata):
                         remote_handshake_port=remote_handshake_port_list[i][0],
                         offset=i,
                         tp_num_need_pulls=self.tp_num_need_pulls,
-                        all_task_done=(i == self.tp_num_need_pulls * self._prefill_pp_size - 1 ))
+                        all_task_done=(i == self.tp_num_need_pulls *
+                                       self._prefill_pp_size - 1))
 
         if self.kv_send_thread is not None:
             for req_id, delay_start_time in metadata.requests_to_send.items():

vllm_ascend/worker/model_runner_v1.py

@@ -46,7 +46,8 @@
 from vllm.compilation.monitor import set_cudagraph_capturing_enabled
 from vllm.config import (CompilationMode, CUDAGraphMode, VllmConfig,
                          get_layers_from_vllm_config)
-from vllm.distributed import tensor_model_parallel_all_gather, get_tensor_model_parallel_world_size
+from vllm.distributed import (get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_gather)
 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
                                           has_kv_transfer_group)
 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1

vllm_ascend/worker/worker_v1.py

@@ -52,7 +52,7 @@
 from vllm_ascend.utils import (enable_sp, init_ascend_soc_version,
                                is_enable_nz, prefill_context_parallel_enable,
                                register_ascend_customop, sleep_mode_enabled,
-                               try_register_lib, vllm_version_is)
+                               try_register_lib)
 from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
 
 torch._dynamo.trace_rules.clear_lru_cache()  # noqa: E402