remove cpu all_reduce in disaggregated-prefill scenario

Signed-off-by: whx-sjtu <2952154980@qq.com>

Author: whx-sjtu

vllm_ascend/worker/model_runner_v1.py

@@ -386,8 +386,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
 
         # kv role
         self.is_kv_producer = False
+        self.is_kv_consumer = False
         if vllm_config.kv_transfer_config is not None:
             self.is_kv_producer = vllm_config.kv_transfer_config.is_kv_producer
+            self.is_kv_consumer = vllm_config.kv_transfer_config.is_kv_consumer
 
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         """Update the cached states and the persistent batch with the scheduler
@@ -605,6 +607,15 @@ def _get_forward_metadata_across_dp(
     ) -> tuple[int, Optional[torch.Tensor], bool, bool]:
         if self.dp_size == 1:
             return num_tokens, None, with_prefill, enable_dbo
+        if self.is_kv_producer:
+            return num_tokens, None, True, enable_dbo
+        if self.is_kv_consumer:
+            max_num_decode_tokens = self.max_num_reqs * self.decode_token_per_req
+            num_tokens_across_dp = torch.tensor([max_num_decode_tokens] *
+                                                self.dp_size,
+                                                device="cpu",
+                                                dtype=torch.int32)
+            return max_num_decode_tokens, num_tokens_across_dp, False, enable_dbo
 
         forward_metadata = torch.tensor(
             [num_tokens, with_prefill, not enable_dbo],