[main][refrator] Refractor forward metadata retrieval across DP nodes to reduce redundant padding.

Signed-off-by: yx0716 <jinyx1007@foxmail.com>

Author: yx0716

vllm_ascend/worker/model_runner_v1.py

@@ -569,45 +569,55 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
             self.input_batch.refresh_sampling_metadata()
 
     def _get_forward_metadata_across_dp(
-        self,
-        maybe_padded_num_tokens: int,
-        num_tokens: int,
-        with_prefill: bool,
-        enable_dbo: bool = False,
+            self, num_tokens: int,
+            with_prefill: bool, enable_dbo: bool
     ) -> tuple[int, Optional[torch.Tensor], bool, bool]:
-        if self.dp_size == 1:
-            return maybe_padded_num_tokens, None, with_prefill, enable_dbo
-
-        num_tokens_across_dp = [0] * self.dp_size * 2
-        num_tokens_across_dp[self.dp_rank] = maybe_padded_num_tokens
-        num_tokens_across_dp[self.dp_size + self.dp_rank] = num_tokens
-        forward_metadata = torch.tensor(num_tokens_across_dp +
-                                        [with_prefill, not enable_dbo],
-                                        device="cpu",
-                                        dtype=torch.int32)
-        dist.all_reduce(forward_metadata, group=get_dp_group().cpu_group)
-        with_prefill = bool(forward_metadata[-2])
-
-        # NOTE: when with_prefill is false before all_reduce and true after all_reduce, we need to revert pad.
-        if with_prefill:
-            num_tokens_across_dp = forward_metadata[self.dp_size:self.dp_size *
-                                                    2]
-            maybe_padded_num_tokens = num_tokens
-        else:
-            num_tokens_across_dp = forward_metadata[:self.dp_size]
 
-        # NOTE: when in torchair_graph_mode, we need to pad local_num_tokens to
-        # `max_tokens_across_dp`, in other situation it is not necessary.
-        if self.torchair_graph_enabled and not with_prefill:
-            maybe_padded_num_tokens = torch.max(num_tokens_across_dp).item()
-            num_tokens_across_dp = torch.tensor([maybe_padded_num_tokens] *
+        # Compose: all_reduce metadata (num_tokens of each rank, with_prefill, enable_dbo)
+        num_tokens_across_dp = torch.zeros(self.dp_size + 2, dtype=torch.int32, device="cpu")
+        num_tokens_across_dp[self.dp_rank] = num_tokens
+        num_tokens_across_dp[-2]   = int(with_prefill)
+        num_tokens_across_dp[-1]   = int(not enable_dbo)
+        dist.all_reduce(num_tokens_across_dp, group=get_dp_group().cpu_group) 
+        with_prefill  = bool(num_tokens_across_dp[-2])
+        enable_dbo    = not bool(num_tokens_across_dp[-1])
+        num_tokens_across_dp = num_tokens_across_dp[:-2]
+        return num_tokens_across_dp, with_prefill, enable_dbo
+
+    def _get_forward_metadata_across_dp_and_pad(
+            self, num_tokens: int,
+            with_prefill: bool, enable_dbo: bool
+    ) -> 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 and not envs_ascend.VLLM_ASCEND_ENABLE_CHUNK_MC2:
+            num_tokens_across_dp = torch.tensor([num_tokens] * self.dp_size,
+                                                device="cpu",
+                                                dtype=torch.int32)
+            return num_tokens, num_tokens_across_dp, True, enable_dbo
+        
+        if self.is_kv_consumer and self.torchair_graph_enabled and len(
+                self.torchair_graph_batch_sizes
+        ) == 1 and not self.in_profile_run:
+            max_num_decode_tokens = self.torchair_graph_batch_sizes[0]
+            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
+        
+        maybe_padded_num_tokens = num_tokens
+        num_tokens_across_dp, with_prefill, enable_dbo = self._get_forward_metadata_across_dp(num_tokens, with_prefill,enable_dbo)
 
-        return maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, not bool(
-            forward_metadata[-1])
-
+        if self.torchair_graph_enabled and not with_prefill:
+            max_num_token = num_tokens_across_dp.max().item()
+            maybe_padded_num_tokens = self.select_torchair_padded_batch_size(
+                max_num_token )
+            num_tokens_across_dp = torch.full((self.dp_size,), maybe_padded_num_tokens, dtype=torch.int32, device="cpu")
+            
+        return maybe_padded_num_tokens, num_tokens_across_dp, with_prefill, enable_dbo
+    
     def get_eagle_atten_dict(
         self,
         scheduler_output: "SchedulerOutput",
@@ -1083,13 +1093,12 @@ def _process_reqs(
             AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding
         ]
 
-        maybe_padded_num_tokens = total_num_scheduled_tokens
-        if self.torchair_graph_enabled and not with_prefill:
-            maybe_padded_num_tokens = self.select_torchair_padded_batch_size(
-                total_num_scheduled_tokens)
+        enable_dbo = self._check_dbo_is_valid(self.query_lens.tolist(),
+                                        attn_state,
+                                        total_num_scheduled_tokens)
         (padded_num_tokens_across_dp, num_tokens_across_dp, with_prefill,
-         enable_dbo) = self._get_forward_metadata_across_dp(
-             maybe_padded_num_tokens, total_num_scheduled_tokens, with_prefill)
+         enable_dbo) = self._get_forward_metadata_across_dp_and_pad(
+             total_num_scheduled_tokens, with_prefill, enable_dbo)
 
         if self.torchair_graph_enabled and not with_prefill:
             graph_pad_size = padded_num_tokens_across_dp - total_num_scheduled_tokens
@@ -1722,15 +1731,10 @@ def _dummy_run(
         with_prefill: bool = False,
         is_torchair_compile: bool = False,
     ) -> torch.Tensor:
-        maybe_padded_num_tokens = num_tokens
-        if self.torchair_graph_enabled and not with_prefill:
-            maybe_padded_num_tokens = self.select_torchair_padded_batch_size(
-                num_tokens)
-
         # Padding for DP
         (num_tokens, num_tokens_across_dp, with_prefill,
-         enable_dbo) = self._get_forward_metadata_across_dp(
-             maybe_padded_num_tokens, num_tokens, with_prefill, False)
+         enable_dbo) = self._get_forward_metadata_across_dp_and_pad(
+             num_tokens, with_prefill, False)
 
         # Set num_scheduled_tokens based on num_tokens and max_num_seqs
         # for dummy run with LoRA so that the num_reqs collectively
@@ -2590,7 +2594,7 @@ def select_torchair_padded_batch_size(self, batch_size: int):
         return selected_batch_size
 
     def get_supported_pooling_tasks(self):
-        model = self.get_model()
+        model=self.get_model()
         if not is_pooling_model(model):
             return []