[Revert] Revert [feature] reduce cases that required DP padding (vllm-project#11772) (vllm-project#13022)

### What this PR does / why we need it?

[Revert] Revert [feature] reduce cases that required DP padding (vllm-project#11772)

### Does this PR introduce _any_ user-facing change?

### How was this patch tested?


- vLLM version: v0.25.1
- vLLM main:
vllm-project/vllm@fe784ff

Signed-off-by: zzzzwwjj <1183291235@qq.com>

Author: zzzzwwjj

.gitleaks.toml

@@ -82,7 +82,3 @@ minEntropy = 3.0
 redact = true
 maxMatchLength = 2048
 scanGitHistory = false
-
-[[allowlists]]
-paths = ["^tests/ut/spec_decode/a2/test_eagle_proposer.py"]
-rules = ["generic-api-key"]

tests/ut/ops/a2/test_token_dispatcher.py

@@ -146,7 +146,11 @@ def setUp(self):
         self.hier_comm_patch = patch(
             "vllm_ascend.ops.fused_moe.token_dispatcher.is_hierarchical_communication_enabled", return_value=False
         )
-        self.mock_hier_comm = self.hier_comm_patch.start()
+        self.hier_comm_patch.start()
+        self.skip_allreduce_patch = patch(
+            "vllm_ascend.ops.fused_moe.token_dispatcher.should_skip_allreduce_across_dp_group", return_value=False
+        )
+        self.mock_skip_allreduce = self.skip_allreduce_patch.start()
 
         kwargs = {"with_quant": False, "top_k": 8, "num_experts": 128}
         self.dispatcher = TokenDispatcherWithMC2(**kwargs)
@@ -161,24 +165,26 @@ def tearDown(self):
         self.ascend_config_patch.stop()
         self.ascend_config_utils_patch.stop()
         self.hier_comm_patch.stop()
+        self.skip_allreduce_patch.stop()
 
     def test_init(self):
         self.assertEqual(self.dispatcher.ep_rank_id, 0)
         self.assertEqual(self.dispatcher.ep_world_size, 8)
         self.assertTrue(self.dispatcher.enable_dispatch_v2)
         self.assertTrue(self.dispatcher.need_extra_args)
-        self.assertEqual(self.dispatcher.global_bs, 1024)
+        self.assertEqual(self.dispatcher.global_bs, 0)
 
     def test_init_uses_mc2_capacity_for_non_uniform_global_bs(self):
         self.mock_get_config.return_value.parallel_config.tensor_parallel_size = 4
+        self.mock_skip_allreduce.return_value = True
 
         dispatcher = TokenDispatcherWithMC2(with_quant=False, top_k=8, num_experts=128)
 
         self.assertEqual(dispatcher.global_bs, 256)
 
-    def test_get_dispatch_mc2_kwargs_without_hier_comm_omits_mc2_mask(self):
+    def test_get_dispatch_mc2_kwargs_with_skip_allreduce_omits_mc2_mask(self):
         self.mock_get_config.return_value.parallel_config.tensor_parallel_size = 4
-        self.mock_hier_comm.return_value = False
+        self.mock_skip_allreduce.return_value = True
         dispatcher = TokenDispatcherWithMC2(with_quant=False, top_k=8, num_experts=128)
 
         hidden_states = torch.randn(10, 128)
@@ -199,9 +205,7 @@ def test_get_dispatch_mc2_kwargs_without_hier_comm_omits_mc2_mask(self):
         self.assertEqual(kwargs["global_bs"], 256)
         self.assertNotIn("x_active_mask", kwargs)
 
-    def test_get_dispatch_mc2_kwargs_with_hier_comm_keeps_mc2_mask(self):
-        self.mock_hier_comm.return_value = True
-        dispatcher = TokenDispatcherWithMC2(with_quant=False, top_k=8, num_experts=128)
+    def test_get_dispatch_mc2_kwargs_without_skip_allreduce_keeps_mc2_mask(self):
         hidden_states = torch.randn(10, 128)
         topk_ids = torch.randint(0, 8, (10, 1))
         topk_weights = torch.randn(10, 1)
@@ -215,7 +219,7 @@ def test_get_dispatch_mc2_kwargs_with_hier_comm_keeps_mc2_mask(self):
             mc2_mask=mc2_mask,
         )
 
-        kwargs = dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
+        kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input)
 
         self.assertEqual(kwargs["global_bs"], 0)
         self.assertIs(kwargs["x_active_mask"], mc2_mask)

tests/ut/spec_decode/a2/test_eagle_proposer.py

@@ -1614,7 +1614,7 @@ def check_mock(self):
         assert hasattr(RunnerCls, "_sync_metadata_across_dp")
         sig = inspect.signature(RunnerCls._sync_metadata_across_dp)
         sig_name = self.get_param_names(sig)
-        assert sig_name == ['self', 'num_tokens', 'is_draft_model', 'cudagraph_mode']
+        assert sig_name == ['self', 'num_tokens', 'is_draft_model', 'cudagraph_mode', 'allow_dp_padding']
 
         assert hasattr(RunnerCls, "_pad_query_start_loc_for_fia")
         sig = inspect.signature(RunnerCls._pad_query_start_loc_for_fia)

vllm_ascend/ascend_config.py

@@ -494,15 +494,6 @@ def __init__(self, finegrained_tp_config: dict, vllm_config):
         if any(size > 0 for size in module_tp_sizes) and enabled_configs:
             logger.info("finegrained_tp_config enabled: %s", ", ".join(enabled_configs))
 
-    def get_max_finegrained_tp_size(self) -> int:
-        max_finegrained_tp_size = 1
-        max_finegrained_tp_size = max(max_finegrained_tp_size, self.oproj_tensor_parallel_size)
-        max_finegrained_tp_size = max(max_finegrained_tp_size, self.lmhead_tensor_parallel_size)
-        max_finegrained_tp_size = max(max_finegrained_tp_size, self.embedding_tensor_parallel_size)
-        max_finegrained_tp_size = max(max_finegrained_tp_size, self.mlp_tensor_parallel_size)
-        max_finegrained_tp_size = max(max_finegrained_tp_size, self.olora_tensor_parallel_size)
-        return max_finegrained_tp_size
-
 
 class AscendCompilationConfig:
     """

vllm_ascend/ascend_forward_context.py

@@ -134,7 +134,10 @@ def set_ascend_forward_context(
         from vllm_ascend.ops.fused_moe.moe_comm_method import get_moe_comm_method
 
         max_num_tokens = int(num_tokens_across_dp.max().item()) if num_tokens_across_dp is not None else num_tokens
-        moe_comm_type = select_moe_comm_method(max_num_tokens, vllm_config)
+        moe_comm_type = select_moe_comm_method(
+            max_num_tokens,
+            vllm_config,
+        )
 
         forward_context.moe_comm_type = moe_comm_type
         forward_context.moe_comm_method = get_moe_comm_method(moe_comm_type)
@@ -353,6 +356,7 @@ def select_moe_comm_method(num_tokens: int, vllm_config: VllmConfig) -> MoECommT
     Args:
         num_tokens (int): The number of tokens in the current batch.
         vllm_config (VllmConfig): Runtime configuration for the model.
+        is_draft_model (bool): Whether the model runs in MTP mode.
 
     Raises:
         ValueError: If the soc version is unsupported.

vllm_ascend/ops/fused_moe/token_dispatcher.py

@@ -51,6 +51,7 @@
     AscendDeviceType,
     get_ascend_device_type,
     is_hierarchical_communication_enabled,
+    should_skip_allreduce_across_dp_group,
 )
 
 EXPERT_TOKEN_NUMS_TYPE_CUMSUM = 0
@@ -139,11 +140,11 @@ def __init__(self, **kwargs):
         self.max_num_tokens_per_rank = num_tokens_per_tp_rank
         _max_global_bs = num_tokens_per_tp_rank * self.ep_world_size
 
-        # When hierarchical communication case, tokens are uniform across ranks:
+        # When allreduce across DP is not skipped, tokens are uniform across ranks:
         # use global_bs=0 (uniform mode) and pass mc2_mask.
-        # When it is not hierarchical communication case, we will not do padding across dp,
-        # tokens may differ per rank: use the real global_bs and do NOT pass mc2_mask.
-        self.global_bs = _max_global_bs if not is_hierarchical_communication_enabled() else 0
+        # When allreduce is skipped, tokens may differ per rank:
+        # use the real global_bs and do NOT pass mc2_mask.
+        self.global_bs = _max_global_bs if should_skip_allreduce_across_dp_group(vllm_config) else 0
 
         # NOTE: When enable_mc2_hierarchy_comm is true, we need pass in `comm_alg` to mc2 op.
         self.need_comm_alg = get_ascend_config().enable_mc2_hierarchy_comm

vllm_ascend/spec_decode/extract_hidden_states_proposer.py

@@ -97,6 +97,7 @@ def _determine_batch_execution_and_padding(
                 num_tokens=num_tokens_padded,
                 is_draft_model=True,
                 cudagraph_mode=cudagraph_mode,
+                allow_dp_padding=use_cudagraphs,
             )
 
             if num_tokens_across_dp is not None:

vllm_ascend/worker/model_runner_v1.py

@@ -151,16 +151,17 @@
     AscendDeviceType,
     calc_split_factor,
     check_gdn_layer,
+    embedding_tp_enable,
     enable_sfa_dcp_replicated_indexer,
     enable_sp,
     enable_sp_by_pass,
     get_ascend_device_type,
     get_c_env,
     global_stream,
     is_hidden_state_cache_spec,
-    is_hierarchical_communication_enabled,
     kv_cache_spec_uses_sparse_sfa_c8,
     lmhead_tp_enable,
+    oproj_tp_enable,
     set_potential_max_tokens,
     should_skip_allreduce_across_dp_group,
 )
@@ -635,6 +636,7 @@ def _sync_metadata_across_dp(
         num_tokens: int,
         is_draft_model: bool = False,
         cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+        allow_dp_padding: bool = False,
     ) -> tuple[int, torch.Tensor | None, CUDAGraphMode]:
         # TODO: In vLLM, the only thing that needs to be synced is num_tokens, but in
         # our case, we still need to sync the other two flags as well. So we need to
@@ -660,17 +662,7 @@ def _sync_metadata_across_dp(
         synced_cudagraph_mode = CUDAGraphMode(_post_process_cudagraph_mode(packed_tensor))
 
         # Create a tensor for num_tokens_after_padding
-        comm_method = select_moe_comm_method(max_tokens_across_dp, self.vllm_config)
-        is_mc2_with_hierarchical = (comm_method == MoECommType.MC2 and is_hierarchical_communication_enabled())
-        is_finegrained_tp = self.ascend_config.finegrained_tp_config.get_max_finegrained_tp_size() > 1
-        # There are three cases where padding between DPs is required:
-        # 1. comm_method == ALLGATHER;
-        # 2. comm_method == MC2 and is hierarchical communication, in this case,
-        #    the mc2 operator does not support dynamic batch size.
-        #    TODO(zzzzwwjj): It can be remove after op support this case.
-        # 3. when finegrained_tp is open, we need to ensure num_tokens remains consistent within finegrained_tp_group.
-        #    TODO(zzzzwwjj): We can do dp padding in finegrained_tp_group, instead of world_group.
-        if comm_method == MoECommType.ALLGATHER or is_mc2_with_hierarchical or is_finegrained_tp:
+        if allow_dp_padding or is_draft_model:
             num_tokens_after_padding = torch.tensor(
                 [max_tokens_across_dp] * self.dp_size, device="cpu", dtype=torch.int32
             )
@@ -2745,6 +2737,10 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None):
             _, num_tokens_across_dp, synced_cudagraph_mode = self._sync_metadata_across_dp(
                 num_tokens=num_tokens_padded,
                 cudagraph_mode=cudagraph_mode,
+                allow_dp_padding=((cudagraph_mode != CUDAGraphMode.NONE)
+                                  or enable_sp(self.vllm_config)
+                                  or oproj_tp_enable()
+                                  or embedding_tp_enable()),
             )
 
             # Extract DP padding if there is any