[0.9.1][Bugfix] fix dp error in dbo (#1291)

Fix running error in dbo when dp_size>1. Add conditional logic in
`_get_forward_metadata_across_dp` to enable dbo.

Signed-off-by: shikang-hangzhou <459956190@qq.com>

Author: shikang-hangzhou

vllm_ascend/attention/attention_v1.py

@@ -134,6 +134,7 @@ class AscendMetadata:
     num_input_tokens: int = 0  # Number of tokens including padding.
 
     with_prefill_across_dp: bool = False
+    enable_dbo_across_dp: bool = False
 
 
 class AscendAttentionMetadataBuilder:
@@ -150,7 +151,8 @@ def build(self,
               num_actual_tokens,
               max_query_len,
               common_prefix_len,
-              with_prefill_across_dp: bool = False):
+              with_prefill_across_dp: bool = False,
+              enable_dbo_across_dp: bool = False):
 
         block_table = self.runner.input_batch.block_table[0].get_device_tensor(
         )
@@ -177,7 +179,8 @@ def build(self,
             slot_mapping=slot_mapping,
             attn_mask=attn_mask,
             attn_state=attn_state,
-            with_prefill_across_dp=with_prefill_across_dp)
+            with_prefill_across_dp=with_prefill_across_dp,
+            enable_dbo_across_dp=enable_dbo_across_dp)
         return attn_metadata
 
 

vllm_ascend/attention/mla_v1.py

@@ -138,6 +138,7 @@ class AscendMLAMetadata:
 
     max_num_tokens_across_dp: int = 0
     with_prefill_across_dp: bool = False
+    enable_dbo_across_dp: bool = False
 
     query_lens: Optional[list[int]] = None
     # The dimension of the attention heads
@@ -367,6 +368,7 @@ def build(
         graph_pad_size: int = -1,
         max_num_tokens_across_dp: int = 0,
         with_prefill_across_dp: bool = False,
+        enable_dbo_across_dp: bool = False,
     ) -> AscendMLAMetadata:
         assert self._num_decodes + self._num_prefills == num_reqs
 
@@ -513,7 +515,7 @@ def build(
             seq_lens=seq_lens,
             max_num_tokens_across_dp=max_num_tokens_across_dp,
             with_prefill_across_dp=with_prefill_across_dp,
-        )
+            enable_dbo_across_dp=enable_dbo_across_dp)
 
 
 class AscendMLAImpl(MLAAttentionImpl):

vllm_ascend/models/deepseek_dbo.py

@@ -74,7 +74,6 @@
 from vllm_ascend.multistream.metadata import (MultiStreamConfig,
                                               MultiStreamStepMetadata,
                                               make_multistream_metadata_ds)
-from vllm_ascend.multistream.ms_split import compute_split_seq_index
 from vllm_ascend.ops.fused_moe import AscendFusedMoE
 from vllm_ascend.utils import dispose_tensor
 
@@ -881,22 +880,8 @@ def forward(
 
     def can_run_ms(self):
         attn_metadata = get_forward_context().attn_metadata
-        # support mla attention and V1 engine at present
-        if not self.use_mla or not envs.VLLM_USE_V1:
-            return False
         # enable prefill overlap
-        if attn_metadata is None or attn_metadata.num_prefills == 0:
-            return False
-        else:
-            [token_index, seq_index
-             ] = compute_split_seq_index(attn_metadata.query_lens,
-                                         attn_metadata.attn_state,
-                                         attn_metadata.num_decode_tokens)
-            if token_index == 0 or seq_index == 0 or seq_index == len(
-                    attn_metadata.query_lens):
-                return False
-        # check whether the total tokens exceed the threshold
-        if self.multistream_config is None or attn_metadata.num_actual_tokens < self.multistream_config.min_total_tokens_to_split:
+        if attn_metadata is None or attn_metadata.num_prefills == 0 or not attn_metadata.enable_dbo_across_dp:
             return False
         return True
 

vllm_ascend/multistream/ms_split.py

@@ -96,10 +96,12 @@ def model_input_split_v1_mla_attn(
     seq_lens = attn_metadata.prefill.seq_lens if attn_metadata.num_prefills > 0 else attn_metadata.decode.seq_lens
     [seq_lens_pre, seq_lens_post] = split_attn_tensor_type(seq_lens, seq_index)
 
-    query_start_loc_pre = attn_metadata.query_start_loc[:seq_index + 1]
-    query_start_loc_post = deepcopy(
-        attn_metadata.query_start_loc[seq_index:]
-    ) - attn_metadata.query_start_loc[seq_index]
+    query_start_loc_pre = query_start_loc_post = None
+    if attn_metadata.query_start_loc is not None:
+        query_start_loc_pre = attn_metadata.query_start_loc[:seq_index + 1]
+        query_start_loc_post = deepcopy(
+            attn_metadata.query_start_loc[seq_index:]
+        ) - attn_metadata.query_start_loc[seq_index]
     [block_table_pre,
      block_table_post] = split_attn_tensor_type(attn_metadata.block_tables,
                                                 seq_index)
@@ -224,6 +226,7 @@ def model_input_split_v1_mla_attn(
         prefill=prefill_pre,
         decode=decode_pre,
         with_prefill_across_dp=attn_metadata.with_prefill_across_dp,
+        enable_dbo_across_dp=attn_metadata.enable_dbo_across_dp,
     )
     attention_metadata_post = _metadata_cls(
         num_actual_tokens=attn_metadata.num_actual_tokens - token_index,
@@ -241,5 +244,6 @@ def model_input_split_v1_mla_attn(
         prefill=prefill_post,
         decode=decode_post,
         with_prefill_across_dp=attn_metadata.with_prefill_across_dp,
+        enable_dbo_across_dp=attn_metadata.enable_dbo_across_dp,
     )
     return [attention_metadata_pre, attention_metadata_post]

vllm_ascend/worker/model_runner_v1.py

@@ -77,6 +77,7 @@
 from vllm_ascend.attention.attention import AttentionMaskBuilder
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
 from vllm_ascend.attention.mla_v1 import CommonAttentionMetadata
+from vllm_ascend.multistream.ms_split import compute_split_seq_index
 from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.sample.rejection_sampler import AscendRejectionSampler
 from vllm_ascend.utils import ProfileExecuteDuration
@@ -569,16 +570,38 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
             self.input_batch.refresh_sampling_metadata()
 
     def _get_forward_metadata_across_dp(
-            self, total_num_scheduled_tokens: int,
-            with_prefill: bool) -> tuple[int, bool]:
+            self, total_num_scheduled_tokens: int, with_prefill: bool,
+            enable_dbo: bool) -> tuple[int, bool, bool]:
         forward_metadata = torch.tensor(
-            [total_num_scheduled_tokens, with_prefill],
+            [total_num_scheduled_tokens, with_prefill, not enable_dbo],
             device="cpu",
             dtype=torch.int32)
         dist.all_reduce(forward_metadata,
                         op=ReduceOp.MAX,
                         group=get_dp_group().cpu_group)
-        return int(forward_metadata[0]), bool(forward_metadata[1] > 0)
+        return int(forward_metadata[0]), bool(
+            forward_metadata[1] > 0), not bool(forward_metadata[2] > 0)
+
+    def _check_dbo_is_valid(self, query_lens: torch.Tensor,
+                            attn_state: AscendAttentionState,
+                            num_tokens: int) -> bool:
+        # do the checks for dp + dbo
+        if attn_state in [
+                AscendAttentionState.DecodeOnly,
+                AscendAttentionState.SpecDecoding
+        ]:
+            return False
+        # considering the case that one dp rank may enable dbo while others may not
+        if not self.vllm_config.model_config.use_mla or not envs_ascend.VLLM_ASCEND_ENABLE_DBO:
+            return False
+        # TODO: remove it if token-level microbatch is enabled
+        [token_index,
+         seq_index] = compute_split_seq_index(query_lens, attn_state,
+                                              num_tokens)
+        if token_index == 0 or seq_index == 0 or seq_index == len(
+                query_lens) or num_tokens < 256:
+            return False
+        return True
 
     def get_model(self) -> nn.Module:
         return self.model
@@ -900,12 +923,16 @@ def _process_reqs(
         with_prefill = attn_state not in [
             AscendAttentionState.DecodeOnly, AscendAttentionState.SpecDecoding
         ]
+        enable_dbo = self._check_dbo_is_valid(self.query_lens.tolist(),
+                                              attn_state,
+                                              total_num_scheduled_tokens)
 
         if self.dp_size > 1:
-            max_num_tokens, with_prefill = self._get_forward_metadata_across_dp(
-                total_num_scheduled_tokens, with_prefill)
+            max_num_tokens, with_prefill, enable_dbo = self._get_forward_metadata_across_dp(
+                total_num_scheduled_tokens, with_prefill, enable_dbo)
             extra_builder_kwargs['max_num_tokens_across_dp'] = max_num_tokens
             extra_builder_kwargs['with_prefill_across_dp'] = with_prefill
+        extra_builder_kwargs['enable_dbo_across_dp'] = enable_dbo
 
         # Add graph_pad_size here
         if self.torchair_graph_enabled and not with_prefill:

vllm_ascend/worker/worker_v1.py

@@ -251,9 +251,10 @@ def execute_dummy_batch(self) -> None:
         runner = self.model_runner
         max_num_tokens = 1
         with_prefill = False
+        enable_dbo = False
         if runner.dp_size > 1:
-            max_num_tokens, with_prefill = runner._get_forward_metadata_across_dp(
-                max_num_tokens, with_prefill)
+            max_num_tokens, with_prefill, _ = runner._get_forward_metadata_across_dp(
+                max_num_tokens, with_prefill, enable_dbo)
         if runner.torchair_graph_enabled and not with_prefill:
             max_num_tokens = runner.select_torchair_padded_batch_size(
                 max_num_tokens)