[DP-Attn] Clarify MLP sync / idle batch preparation logic (#12843)

python/sglang/bench_one_batch.py

@@ -68,7 +68,7 @@ from sglang.srt.distributed.parallel_state import destroy_distributed_environmen
 from sglang.srt.entrypoints.engine import _set_envs_and_config
 from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
-from sglang.srt.managers.scheduler import Scheduler
+from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.sampling.sampling_params import SamplingParams
@@ -391,15 +391,13 @@ def decode(input_token_ids, batch, model_runner):
 
 def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner):
     if require_mlp_sync(model_runner.server_args):
-        Scheduler.prepare_mlp_sync_batch_raw(
+        prepare_mlp_sync_batch_raw(
             batch,
             dp_size=model_runner.server_args.dp_size,
             attn_tp_size=1,
             tp_group=model_runner.tp_group,
             get_idle_batch=None,
             disable_cuda_graph=model_runner.server_args.disable_cuda_graph,
-            spec_algorithm=SpeculativeAlgorithm.NONE,
-            speculative_num_draft_tokens=None,
             require_mlp_tp_gather=require_mlp_tp_gather(model_runner.server_args),
             disable_overlap_schedule=model_runner.server_args.disable_overlap_schedule,
             offload_tags=set(),

python/sglang/srt/disaggregation/decode.py

@@ -25,7 +25,7 @@ import time
 from collections import deque
 from dataclasses import dataclass
 from http import HTTPStatus
-from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, List, Optional, Type, Union
 
 import torch
 from torch.distributed import ProcessGroup
@@ -63,7 +63,7 @@ from sglang.srt.tracing.trace import (
     trace_slice_batch,
     trace_slice_end,
 )
-from sglang.srt.utils import get_int_env_var, require_mlp_sync
+from sglang.srt.utils import get_int_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
 logger = logging.getLogger(__name__)
@@ -833,8 +833,6 @@ class SchedulerDisaggregationDecodeMixin:
             batch = self.get_next_disagg_decode_batch_to_run()
             self.cur_batch = batch
 
-            prepare_mlp_sync_flag = require_mlp_sync(self.server_args)
-
             if batch:
                 # Generate fake extend output.
                 if batch.forward_mode.is_extend():
@@ -843,15 +841,15 @@ class SchedulerDisaggregationDecodeMixin:
                         batch.reqs, any(req.return_logprob for req in batch.reqs)
                     )
                     trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
-                    if prepare_mlp_sync_flag:
+                    if self.require_mlp_sync:
-                        self._prepare_idle_batch_and_run(None)
+                        self._prepare_idle_batch_and_run()
                 else:
-                    if prepare_mlp_sync_flag:
+                    if self.require_mlp_sync:
                         self.prepare_mlp_sync_batch(batch)
                     result = self.run_batch(batch)
                     self.process_batch_result(batch, result)
-            elif prepare_mlp_sync_flag:
+            elif self.require_mlp_sync:
-                batch, _ = self._prepare_idle_batch_and_run(None)
+                batch, _ = self._prepare_idle_batch_and_run()
 
             queue_size = (
                 len(self.waiting_queue)
@@ -882,8 +880,6 @@ class SchedulerDisaggregationDecodeMixin:
             self.cur_batch = batch
             last_batch_in_queue = False
 
-            prepare_mlp_sync_flag = require_mlp_sync(self.server_args)
-
             batch_result = None
             if batch:
                 # Generate fake extend output.
@@ -893,23 +889,23 @@ class SchedulerDisaggregationDecodeMixin:
                         batch.reqs, any(req.return_logprob for req in batch.reqs)
                     )
                     trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
-                    if prepare_mlp_sync_flag:
+                    if self.require_mlp_sync:
                         batch_, batch_result = self._prepare_idle_batch_and_run(
-                            None, delay_process=True
+                            delay_process=True
                         )
                         if batch_:
                             self.result_queue.append((batch_.copy(), batch_result))
                             last_batch_in_queue = True
                 else:
-                    if prepare_mlp_sync_flag:
+                    if self.require_mlp_sync:
                         self.prepare_mlp_sync_batch(batch)
                     batch_result = self.run_batch(batch)
                     self.result_queue.append((batch.copy(), batch_result))
                     last_batch_in_queue = True
 
-            elif prepare_mlp_sync_flag:
+            elif self.require_mlp_sync:
                 batch, batch_result = self._prepare_idle_batch_and_run(
-                    None, delay_process=True
+                    delay_process=True
                 )
                 if batch:
                     self.result_queue.append((batch.copy(), batch_result))
@@ -936,8 +932,8 @@ class SchedulerDisaggregationDecodeMixin:
             self.last_batch = batch
             self.last_batch_in_queue = last_batch_in_queue
 
-    def _prepare_idle_batch_and_run(self: Scheduler, batch, delay_process=False):
+    def _prepare_idle_batch_and_run(self: Scheduler, delay_process=False):
-        batch = self.prepare_mlp_sync_batch(batch)
+        batch = self.prepare_mlp_sync_batch(None)
         result = None
         if batch:
             result = self.run_batch(batch)
@@ -947,7 +943,7 @@ class SchedulerDisaggregationDecodeMixin:
 
     def get_next_disagg_decode_batch_to_run(
         self: Scheduler,
-    ) -> Optional[Tuple[ScheduleBatch, bool]]:
+    ) -> Optional[ScheduleBatch]:
         """Create fake completed prefill if possible and merge with running batch"""
         # Merge the prefill batch into the running batch
         last_batch = self.last_batch

python/sglang/srt/disaggregation/prefill.py

@@ -54,7 +54,7 @@ from sglang.srt.mem_cache.memory_pool import (
     SWAKVPool,
 )
 from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end
-from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj, require_mlp_sync
+from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj
 
 if TYPE_CHECKING:
     from torch.distributed import ProcessGroup
@@ -326,7 +326,7 @@ class SchedulerDisaggregationPrefillMixin:
                 attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
                 trace_event_batch("schedule", batch.reqs, attrs=attrs)
 
-            if require_mlp_sync(self.server_args):
+            if self.require_mlp_sync:
                 batch = self.prepare_mlp_sync_batch(batch)
             self.cur_batch = batch
 
@@ -361,7 +361,7 @@ class SchedulerDisaggregationPrefillMixin:
                 attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
                 trace_event_batch("schedule", batch.reqs, attrs=attrs)
 
-            if require_mlp_sync(self.server_args):
+            if self.require_mlp_sync:
                 batch = self.prepare_mlp_sync_batch(batch)
             self.cur_batch = batch
 

python/sglang/srt/managers/scheduler.py

@@ -128,6 +128,7 @@ from sglang.srt.managers.schedule_policy import (
     PrefillAdder,
     SchedulePolicy,
 )
+from sglang.srt.managers.scheduler_dp_attn_mixin import SchedulerDPAttnMixin
 from sglang.srt.managers.scheduler_input_blocker import SchedulerInputBlocker
 from sglang.srt.managers.scheduler_metrics_mixin import (
     RECORD_STEP_TIME,
@@ -166,7 +167,6 @@ from sglang.srt.tracing.trace import (
     trace_slice_end,
     trace_slice_start,
 )
-from sglang.srt.two_batch_overlap import TboDPAttentionPreparer
 from sglang.srt.utils import (
     DynamicGradMode,
     broadcast_pyobj,
@@ -181,7 +181,6 @@ from sglang.srt.utils import (
     numa_bind_to_node,
     point_to_point_pyobj,
     require_mlp_sync,
-    require_mlp_tp_gather,
     set_gpu_proc_affinity,
     set_random_seed,
     suppress_other_loggers,
@@ -218,6 +217,7 @@ class Scheduler(
     SchedulerMultiplexMixin,
     SchedulerRuntimeCheckerMixin,
     SchedulerPPMixin,
+    SchedulerDPAttnMixin,
 ):
     """A scheduler that manages a tensor parallel GPU worker."""
 
@@ -523,6 +523,9 @@ class Scheduler(
         # Init overlap
         self.init_overlap()
 
+        # Init mlp sync flag
+        self.require_mlp_sync = require_mlp_sync(server_args)
+
         # Init request dispatcher
         self._request_dispatcher = TypeBasedDispatcher(
             [
@@ -1667,13 +1670,14 @@ class Scheduler(
 
         new_batch = self.get_new_batch_prefill()
 
-        need_dp_attn_preparation = require_mlp_sync(self.server_args)
+        need_mlp_sync = self.require_mlp_sync
-
+        if need_mlp_sync and not self.spec_algorithm.is_none():
-        if need_dp_attn_preparation and not self.spec_algorithm.is_none():
+            # NOTE: This branch makes sure prefill and decode batches will not be mixed when spec and dp-attn is enabled.
-            # In speculative decoding, prefill batches and decode batches cannot be processed in the same DP attention group.
+            # Before merging the new batch into running batch:
-            # We prepare idle batches in advance to skip preparing decode batches when there are prefill batches in the group.
+            # 1. All new batches are none -> need_mlp_sync remains true (sync is needed for decode batch).
+            # 2. All new batches are some (prefill / idle) -> we do not need prepare mlp sync one more time.
             new_batch = self.prepare_mlp_sync_batch(new_batch)
-            need_dp_attn_preparation = new_batch is None
+            need_mlp_sync = new_batch is None
 
         if new_batch is not None:
             # Run prefill first if possible
@@ -1687,7 +1691,7 @@ class Scheduler(
                 ret = None
 
         # Handle DP attention
-        if need_dp_attn_preparation:
+        if need_mlp_sync:
             ret = self.prepare_mlp_sync_batch(ret)
 
         if ret:
@@ -2108,142 +2112,6 @@ class Scheduler(
             self.return_health_check_ct -= 1
             self.send_to_tokenizer.send_output(HealthCheckOutput())
 
-    def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch):
-        return self.prepare_mlp_sync_batch_raw(
-            local_batch,
-            dp_size=self.server_args.dp_size,
-            attn_tp_size=self.attn_tp_size,
-            tp_group=self.tp_group,
-            get_idle_batch=self.get_idle_batch,
-            disable_cuda_graph=self.server_args.disable_cuda_graph,
-            spec_algorithm=self.spec_algorithm,
-            speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
-            require_mlp_tp_gather=require_mlp_tp_gather(self.server_args),
-            disable_overlap_schedule=self.server_args.disable_overlap_schedule,
-            offload_tags=self.offload_tags,
-        )
-
-    @staticmethod
-    def prepare_mlp_sync_batch_raw(
-        local_batch: ScheduleBatch,
-        dp_size,
-        attn_tp_size: int,
-        tp_group,
-        get_idle_batch,
-        disable_cuda_graph: bool,
-        spec_algorithm,
-        speculative_num_draft_tokens,
-        require_mlp_tp_gather: bool,
-        disable_overlap_schedule: bool,
-        offload_tags: set[str],
-    ):
-        # Check if other DP workers have running batches
-        if local_batch is None:
-            num_tokens = 0
-            num_tokens_for_logprob = 0
-        elif local_batch.forward_mode.is_decode():
-            num_tokens = local_batch.batch_size()
-            num_tokens_for_logprob = num_tokens
-        else:
-            num_tokens = local_batch.extend_num_tokens
-            if local_batch.return_logprob:
-                num_tokens_for_logprob = sum(
-                    # We should have at least 1 token for sample in every case.
-                    max(extend_len - logprob_start_len, 1)
-                    for logprob_start_len, extend_len in zip(
-                        local_batch.extend_logprob_start_lens,
-                        local_batch.extend_lens,
-                    )
-                )
-            else:
-                # When return_logprob = False, only need last token per request
-                num_tokens_for_logprob = local_batch.batch_size()
-
-        if local_batch is None or local_batch.forward_mode.is_decode_or_idle():
-            can_cuda_graph = 1
-        else:
-            can_cuda_graph = 0
-
-        is_extend_in_batch = (
-            local_batch.forward_mode.is_extend() if local_batch else False
-        )
-
-        tbo_preparer = TboDPAttentionPreparer()
-        if len(offload_tags) == 0 and disable_overlap_schedule:
-            group = tp_group.device_group
-            device = tp_group.device
-        else:
-            group = tp_group.cpu_group
-            device = "cpu"
-
-        local_info = torch.tensor(
-            [
-                num_tokens,
-                can_cuda_graph,
-                num_tokens_for_logprob,
-                is_extend_in_batch,
-                *tbo_preparer.prepare_all_gather(
-                    local_batch,
-                ),
-            ],
-            dtype=torch.int64,
-            device=device,
-        )
-        global_info = torch.empty(
-            (dp_size, attn_tp_size, 6),
-            dtype=torch.int64,
-            device=device,
-        )
-        torch.distributed.all_gather_into_tensor(
-            global_info.flatten(),
-            local_info,
-            group=group,
-        )
-        global_num_tokens = global_info[:, 0, 0].tolist()
-        can_cuda_graph = min(global_info[:, 0, 1].tolist())
-        global_num_tokens_for_logprob = global_info[:, 0, 2].tolist()
-        is_extend_in_batch = global_info[:, 0, 3].tolist()
-
-        tbo_split_seq_index, global_forward_mode = tbo_preparer.compute_output(
-            global_info[:, :, 4:6]
-        )
-
-        if local_batch is None and max(global_num_tokens) > 0:
-            local_batch = get_idle_batch()
-
-        if local_batch is not None:
-            # TODO: handle the case when moe_dense_tp_size != 1
-            if not require_mlp_tp_gather:
-                local_batch.global_num_tokens = [num_tokens]
-                local_batch.global_num_tokens_for_logprob = [num_tokens_for_logprob]
-            else:
-                local_batch.global_num_tokens = global_num_tokens
-                local_batch.global_num_tokens_for_logprob = (
-                    global_num_tokens_for_logprob
-                )
-            local_batch.is_extend_in_batch = any(is_extend_in_batch)
-            local_batch.tbo_split_seq_index = tbo_split_seq_index
-            local_batch.global_forward_mode = global_forward_mode
-
-            # Check forward mode for cuda graph
-            if not disable_cuda_graph:
-                local_batch.can_run_dp_cuda_graph = can_cuda_graph
-
-        return local_batch
-
-    def get_idle_batch(self):
-        idle_batch = ScheduleBatch.init_new(
-            [],
-            self.req_to_token_pool,
-            self.token_to_kv_pool_allocator,
-            self.tree_cache,
-            self.model_config,
-            self.enable_overlap,
-            self.spec_algorithm,
-        )
-        idle_batch.prepare_for_idle()
-        return idle_batch
-
     def move_ready_grammar_requests(self):
         """Move requests whose grammar objects are ready from grammar_queue to waiting_queue."""
 

python/sglang/srt/managers/scheduler_dp_attn_mixin.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.two_batch_overlap import TboDPAttentionPreparer
+from sglang.srt.utils.common import require_mlp_tp_gather
+
+if TYPE_CHECKING:
+    from sglang.srt.distributed.parallel_state import GroupCoordinator
+    from sglang.srt.managers.scheduler import Scheduler
+
+
+@dataclass
+class MLPSyncBatchInfo:
+    dp_size: int
+    tp_size: int
+
+    num_tokens: int
+    num_tokens_for_logprob: int
+    can_cuda_graph: bool
+    is_extend_in_batch: bool
+    local_can_run_tbo: bool
+    local_forward_mode: int
+
+    # some gathered elements
+    tp0_info: torch.Tensor = None
+    global_num_tokens: list[int] = None
+    global_num_tokens_for_logprob: list[int] = None
+
+    def _get_local_tensor(self, device, dtype=torch.int64) -> torch.Tensor:
+        return torch.tensor(
+            [
+                self.num_tokens,
+                self.num_tokens_for_logprob,
+                int(self.can_cuda_graph),
+                int(self.is_extend_in_batch),
+                int(self.local_can_run_tbo),
+                self.local_forward_mode,
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+    def all_gather(self, device, group: torch.distributed.ProcessGroup):
+        local_info_tensor = self._get_local_tensor(device=device)
+        global_info_tensor = torch.empty(
+            (self.dp_size, self.tp_size, 6),
+            dtype=torch.int64,
+            device=device,
+        )
+
+        torch.distributed.all_gather_into_tensor(
+            global_info_tensor.flatten(),
+            local_info_tensor,
+            group=group,
+        )
+
+        tp0_info = global_info_tensor[:, 0, :]
+        self.tp0_info = tp0_info
+        self.global_num_tokens = tp0_info[:, 0].tolist()
+        self.global_num_tokens_for_logprob = tp0_info[:, 1].tolist()
+        self.can_cuda_graph = bool(tp0_info[:, 2].min().item())
+        self.is_extend_in_batch = bool(tp0_info[:, 3].max().item())
+
+
+def prepare_mlp_sync_batch_raw(
+    local_batch: ScheduleBatch,
+    dp_size: int,
+    attn_tp_size: int,
+    tp_group: GroupCoordinator,
+    get_idle_batch: Callable[[], ScheduleBatch],
+    disable_cuda_graph: bool,
+    require_mlp_tp_gather: bool,
+    disable_overlap_schedule: bool,
+    offload_tags: set[str],
+):
+    # Check if other DP workers have running batches
+    if local_batch is None:
+        num_tokens = 0
+        num_tokens_for_logprob = 0
+    elif local_batch.forward_mode.is_decode():
+        num_tokens = local_batch.batch_size()
+        num_tokens_for_logprob = num_tokens
+    else:
+        num_tokens = local_batch.extend_num_tokens
+        if local_batch.return_logprob:
+            num_tokens_for_logprob = sum(
+                # We should have at least 1 token for sample in every case.
+                max(extend_len - logprob_start_len, 1)
+                for logprob_start_len, extend_len in zip(
+                    local_batch.extend_logprob_start_lens,
+                    local_batch.extend_lens,
+                )
+            )
+        else:
+            # When return_logprob = False, only need last token per request
+            num_tokens_for_logprob = local_batch.batch_size()
+
+    if local_batch is None or local_batch.forward_mode.is_decode_or_idle():
+        can_cuda_graph = 1
+    else:
+        can_cuda_graph = 0
+
+    is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
+
+    tbo_preparer = TboDPAttentionPreparer()
+    if len(offload_tags) == 0 and disable_overlap_schedule:
+        group = tp_group.device_group
+        device = tp_group.device
+    else:
+        group = tp_group.cpu_group
+        device = "cpu"
+
+    local_can_run_tbo, local_forward_mode = tbo_preparer.prepare_all_gather(local_batch)
+
+    mlp_sync_info = MLPSyncBatchInfo(
+        dp_size=dp_size,
+        tp_size=attn_tp_size,
+        num_tokens=num_tokens,
+        num_tokens_for_logprob=num_tokens_for_logprob,
+        can_cuda_graph=can_cuda_graph,
+        is_extend_in_batch=is_extend_in_batch,
+        local_can_run_tbo=local_can_run_tbo,
+        local_forward_mode=local_forward_mode,
+    )
+    mlp_sync_info.all_gather(device=device, group=group)
+
+    tbo_split_seq_index, global_forward_mode = tbo_preparer.compute_output(
+        mlp_sync_info.tp0_info[:, 4:6],
+    )
+
+    if local_batch is None and max(mlp_sync_info.global_num_tokens) > 0:
+        local_batch = get_idle_batch()
+
+    if local_batch is not None:
+        # TODO: handle the case when moe_dense_tp_size != 1
+        if not require_mlp_tp_gather:
+            local_batch.global_num_tokens = [num_tokens]
+            local_batch.global_num_tokens_for_logprob = [num_tokens_for_logprob]
+        else:
+            local_batch.global_num_tokens = mlp_sync_info.global_num_tokens
+            local_batch.global_num_tokens_for_logprob = (
+                mlp_sync_info.global_num_tokens_for_logprob
+            )
+        local_batch.is_extend_in_batch = mlp_sync_info.is_extend_in_batch
+        local_batch.tbo_split_seq_index = tbo_split_seq_index
+        local_batch.global_forward_mode = global_forward_mode
+
+        # Check forward mode for cuda graph
+        if not disable_cuda_graph:
+            local_batch.can_run_dp_cuda_graph = mlp_sync_info.can_cuda_graph
+
+    return local_batch
+
+
+class SchedulerDPAttnMixin:
+    def prepare_mlp_sync_batch(self: Scheduler, local_batch: ScheduleBatch):
+        return prepare_mlp_sync_batch_raw(
+            local_batch,
+            dp_size=self.server_args.dp_size,
+            attn_tp_size=self.attn_tp_size,
+            tp_group=self.tp_group,
+            get_idle_batch=self.get_idle_batch,
+            disable_cuda_graph=self.server_args.disable_cuda_graph,
+            require_mlp_tp_gather=require_mlp_tp_gather(self.server_args),
+            disable_overlap_schedule=self.server_args.disable_overlap_schedule,
+            offload_tags=self.offload_tags,
+        )
+
+    def get_idle_batch(self: Scheduler) -> ScheduleBatch:
+        idle_batch = ScheduleBatch.init_new(
+            [],
+            self.req_to_token_pool,
+            self.token_to_kv_pool_allocator,
+            self.tree_cache,
+            self.model_config,
+            self.enable_overlap,
+            self.spec_algorithm,
+        )
+        idle_batch.prepare_for_idle()
+        return idle_batch

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -1,10 +1,15 @@
-from typing import List, Optional
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
 
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
-from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj, require_mlp_sync
+from sglang.srt.utils import DynamicGradMode, point_to_point_pyobj
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.scheduler import Scheduler
 
 
 class SchedulerPPMixin:
@@ -132,7 +137,7 @@ class SchedulerPPMixin:
                 self.self_check_during_idle()
 
     @DynamicGradMode()
-    def event_loop_pp_disagg_prefill(self):
+    def event_loop_pp_disagg_prefill(self: Scheduler):
         """
         An event loop for the prefill server in pipeline parallelism.
 
@@ -238,7 +243,7 @@ class SchedulerPPMixin:
                 self.process_prefill_chunk()
 
                 batch = self.get_new_batch_prefill()
-                if require_mlp_sync(self.server_args):
+                if self.require_mlp_sync:
                     batch = self.prepare_mlp_sync_batch(batch)
                 mbs[mb_id] = batch
 

python/sglang/srt/two_batch_overlap.py

@@ -407,8 +407,8 @@ class TboDPAttentionPreparer:
         return local_can_run_tbo, local_forward_mode
 
     def compute_output(self, partial_global_info):
-        local_can_run_tbo_aggregated = min(partial_global_info[:, 0, 0].tolist())
+        local_can_run_tbo_aggregated = min(partial_global_info[:, 0].tolist())
-        forward_modes = partial_global_info[:, 0, 1].tolist()
+        forward_modes = partial_global_info[:, 1].tolist()
 
         global_forward_mode, forward_mode_agree = self._compute_global_forward_mode(
             forward_modes