optimize: ADP schedule optimization

Signed-off-by: yunruis <[email protected]>

Author: yunruis

examples/llm-api/quickstart_advanced.py

@@ -50,7 +50,7 @@ def add_llm_args(parser):
     parser.add_argument('--moe_backend',
                         type=str,
                         default='CUTLASS',
-                        choices=['CUTLASS', 'TRTLLM', 'VANILLA'])
+                        choices=['CUTLASS', 'TRTLLM', 'VANILLA', 'WIDEEP'])
     parser.add_argument('--enable_attention_dp',
                         default=False,
                         action='store_true')

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py

@@ -245,7 +245,7 @@ def select_alltoall_method_type(mapping: Mapping, top_k: int,
         if os.environ.get("TRTLLM_MOE_DISABLE_ALLTOALLV", "0") == "1":
             return AlltoallMethodType.NotEnabled
 
-        if mapping.moe_ep_size <= top_k:
+        if mapping.moe_ep_size < top_k:
             return AlltoallMethodType.NotEnabled
 
         if MnnvlMemory.supports_mnnvl():

tensorrt_llm/_torch/pyexecutor/config.py

@@ -44,6 +44,10 @@ class PyTorchConfig:
     moe_max_num_tokens: Optional[int] = None
     moe_load_balancer: Optional[Union[MoeLoadBalancerConfig, dict, str]] = None
 
+    use_attention_dp_config: bool = False
+    attention_dp_time_out_iters: int = 500
+    attention_dp_batching_wait_iters: int = 0
+
     attn_backend: str = 'TRTLLM'
     moe_backend: str = 'CUTLASS'
 

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -241,6 +241,9 @@ def __init__(self,
         self.enable_iter_perf_stats = model_engine.pytorch_backend_config.enable_iter_perf_stats
         self.enable_iter_req_stats = model_engine.pytorch_backend_config.enable_iter_req_stats
         self.stream_interval = model_engine.pytorch_backend_config.stream_interval
+        self.use_attention_dp_config = model_engine.pytorch_backend_config.use_attention_dp_config
+        self.attention_dp_time_out_iters = model_engine.pytorch_backend_config.attention_dp_time_out_iters
+        self.attention_dp_batching_wait_iters = model_engine.pytorch_backend_config.attention_dp_batching_wait_iters
         self.num_fetch_requests_cur_rank = 0
         self.num_fetch_requests = 0
         self.shutdown_event = threading.Event()
@@ -287,6 +290,9 @@ def __init__(self,
             self.draft_model_engine.warmup(self.resource_manager)
 
         self.is_shutdown = False
+        self.max_batch_size = max_batch_size
+        self.adp_ctx_waiting_iters = 0
+        self.adp_ctx_batching_wait_iters = 0
 
         self.stats_lock = threading.Lock()
         self.stats = []
@@ -1228,7 +1234,16 @@ def _broadcast_new_requests(
     def _fetch_new_requests(self) -> List[RequestQueueItem]:
         if self.enable_attention_dp:
             all_ranks_num_active_requests = []
-            responses_list = self.dist.tp_allgather(len(self.active_requests))
+            num_active_requests = len(self.active_requests)
+            responses_list = self.dist.tp_allgather(num_active_requests)
+            # Debug check - remove after verification
+            if not all(isinstance(x, int) for x in responses_list):
+                raise RuntimeError(
+                    f"tp_allgather returned non-integer values: {responses_list} "
+                    +
+                    f"Expected all ranks to return int from {num_active_requests} and {self.active_requests}."
+                )
+
             for num_active_requests in responses_list:
                 all_ranks_num_active_requests.append(num_active_requests)
             total_num_active_requests = sum(all_ranks_num_active_requests)
@@ -1518,8 +1533,66 @@ def _schedule(self):
         scheduler_output = self.scheduler.schedule_request(
             self.active_requests, self.inflight_req_ids)
         scheduled_requests = ScheduledRequests()
+        context_requests = scheduler_output.context_requests
+        if self.enable_attention_dp:
+            num_scheduled_context_requests = len(
+                scheduler_output.context_requests)
+            num_scheduled_generation_requests = len(
+                scheduler_output.generation_requests)
+            num_scheduled_tokens = sum([
+                len(req.get_tokens(0)) for req in context_requests
+            ]) + num_scheduled_generation_requests
+            responses_list = self.dist.tp_allgather([
+                num_scheduled_context_requests,
+                num_scheduled_generation_requests, num_scheduled_tokens
+            ])
+            all_ranks_num_scheduled_context_requests = [
+                response[0] for response in responses_list
+            ]
+            all_ranks_num_scheduled_generation_requests = [
+                response[1] for response in responses_list
+            ]
+            all_ranks_num_scheduled_tokens = [
+                response[2] for response in responses_list
+            ]
+
+            all_ranks_have_free_ctx_slots = all([
+                num_gen < self.max_batch_size
+                for num_gen in all_ranks_num_scheduled_generation_requests
+            ])
+            all_ranks_have_multi_gen = all([
+                num_gen > 1
+                for num_gen in all_ranks_num_scheduled_generation_requests
+            ])
+            all_ranks_have_ctx_requests = all([
+                num_ctx > 0
+                for num_ctx in all_ranks_num_scheduled_context_requests
+            ])
+
+            all_ranks_have_gen_requests = all([
+                num_gen > 0
+                for num_gen in all_ranks_num_scheduled_generation_requests
+            ])
+
+            if self.use_attention_dp_config:
+                # wait for all ranks have context requests
+                if all_ranks_have_free_ctx_slots and all_ranks_have_ctx_requests:
+                    self.adp_ctx_waiting_iters = 0
+                    # balance number of context requests across ranks
+                    if all_ranks_have_gen_requests:
+                        if self.adp_ctx_batching_wait_iters <= self.attention_dp_batching_wait_iters:
+                            self.adp_ctx_batching_wait_iters += 1
+                            context_requests = []
+                        else:
+                            self.adp_ctx_batching_wait_iters = 0
+                else:
+                    self.adp_ctx_waiting_iters += 1
+                    context_requests = []
+                    if self.adp_ctx_waiting_iters >= self.attention_dp_time_out_iters or not all_ranks_have_gen_requests:
+                        self.adp_ctx_waiting_iters = 0
+                        context_requests = scheduler_output.context_requests
 
-        scheduled_requests.context_requests = scheduler_output.context_requests
+        scheduled_requests.context_requests = context_requests
         scheduled_requests.generation_requests = scheduler_output.generation_requests
         scheduled_requests.paused_requests = scheduler_output.paused_requests
         return scheduled_requests, scheduler_output.fitting_disagg_gen_init_requests, scheduler_output.num_fitting_requests

tensorrt_llm/llmapi/llm_args.py

@@ -88,6 +88,23 @@ def validate_cuda_graph_max_batch_size(cls, v):
         return v
 
 
+class AttentionDpConfig(BaseModel):
+    """
+    Configuration for attention DP.
+    """
+    enable_balance: bool = Field(default=False,
+                                 description="Whether to enable balance.")
+    batching_wait_iters: int = Field(
+        default=10,
+        description="The number of iterations to wait for batching.")
+    timeout_iters: int = Field(
+        default=500, description="The number of iterations to timeout.")
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
 class MoeConfig(BaseModel):
     """
     Configuration for MoE.
@@ -1789,6 +1806,11 @@ class TorchLlmArgs(BaseLlmArgs):
         since the input shapes are a function of the sequence lengths).\
          Note that each CUDA graph can use up to 200 MB of extra memory.")
 
+    attention_dp_config: Optional[AttentionDpConfig] = Field(
+        default=None,
+        description=
+        "Attention DP config. If true, use attention DP optimized scheduler.")
+
     disable_overlap_scheduler: bool = Field(
         default=False, description="Disable the overlap scheduler.")
 
@@ -1993,6 +2015,29 @@ def validate_cuda_graph_config(self) -> 'TorchLlmArgs':
 
         return self
 
+    @model_validator(mode='after')
+    def validate_attention_dp_config(self) -> 'TorchLlmArgs':
+        """Validate attention DP configuration.
+
+        Ensures that:
+        1. If attention_dp_config.enable_balance is true, attention_dp_config.batching_wait_iters must be greater than 0
+        2. If attention_dp_config.enable_balance is true, attention_dp_config.timeout_iters must be greater than 0
+        """
+        if self.attention_dp_config is None:
+            return self
+
+        config = self.attention_dp_config
+        if config.enable_balance:
+            if config.batching_wait_iters < 0:
+                raise ValueError(
+                    "attention_dp_config.batching_wait_iters must be greater than 0 when enable_balance is true"
+                )
+            if config.timeout_iters < 0:
+                raise ValueError(
+                    "attention_dp_config.timeout_iters must be greater than 0 when enable_balance is true"
+                )
+        return self
+
     # TODO: Remove this after the PyTorch backend is fully migrated to TorchLlmArgs from ExecutorConfig
     def get_pytorch_backend_config(self) -> "PyTorchConfig":
         from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig
@@ -2040,7 +2085,14 @@ def get_pytorch_backend_config(self) -> "PyTorchConfig":
             enable_min_latency=self.enable_min_latency,
             stream_interval=self.stream_interval,
             force_dynamic_quantization=self.force_dynamic_quantization,
-            allreduce_strategy=self.allreduce_strategy)
+            allreduce_strategy=self.allreduce_strategy,
+            use_attention_dp_config=bool(self.attention_dp_config is not None),
+            attention_dp_time_out_iters=self.attention_dp_config.timeout_iters
+            if self.attention_dp_config is not None else
+            AttentionDpConfig.model_fields['timeout_iters'].default,
+            attention_dp_batching_wait_iters=self.attention_dp_config.
+            batching_wait_iters if self.attention_dp_config is not None else
+            AttentionDpConfig.model_fields['batching_wait_iters'].default)
 
 
 def update_llm_args_with_extra_dict(
@@ -2057,6 +2109,7 @@ def update_llm_args_with_extra_dict(
         "speculative_config": DecodingBaseConfig,
         "lora_config": LoraConfig,
         "moe_config": MoeConfig,
+        "attention_dp_config": AttentionDpConfig,
     }
     for field_name, field_type in field_mapping.items():
         if field_name in llm_args_dict: