optimize: ADP schedule optimization

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

Author: yunruis

tensorrt_llm/_torch/pyexecutor/config.py

@@ -46,6 +46,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 = 0
+    attention_dp_batching_wait_iters: int = 0
+
     attn_backend: str = 'TRTLLM'
     moe_backend: str = 'CUTLASS'
 

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -176,6 +176,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()
@@ -214,6 +217,9 @@ def __init__(self,
             self.draft_model_engine.warmup(self.resource_manager)
 
         self.is_shutdown = False
+        self.max_batch_size = max_batch_size
+        self.self.adp_ctx_waiting_iters_count = 0
+        self.adp_ctx_batching_wait_iters_count = 0
 
         # request fetcher initialization
         self.executor_request_queue = ExecutorRequestQueue(
@@ -1119,8 +1125,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.self.adp_ctx_waiting_iters_count = 0
+                    # balance number of context requests across ranks
+                    if all_ranks_have_gen_requests:
+                        if self.adp_ctx_batching_wait_iters_count < self.attention_dp_batching_wait_iters:
+                            self.adp_ctx_batching_wait_iters_count += 1
+                            context_requests = []
+                        else:
+                            self.adp_ctx_batching_wait_iters_count = 0
+                else:
+                    self.self.adp_ctx_waiting_iters_count += 1
+                    context_requests = []
+                    if self.self.adp_ctx_waiting_iters_count >= self.attention_dp_time_out_iters or not all_ranks_have_gen_requests:
+                        self.self.adp_ctx_waiting_iters_count = 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

@@ -89,6 +89,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=0,
+        description="The number of iterations to wait for batching.")
+    timeout_iters: int = Field(
+        default=0, description="The number of iterations to timeout.")
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        return cls(**data)
+
+
 class MoeConfig(BaseModel):
     """
     Configuration for MoE.
@@ -1817,6 +1834,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.")
 
@@ -2075,6 +2097,31 @@ def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs':
                 "please update the validator")
         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
@@ -2125,7 +2172,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(
@@ -2142,6 +2196,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: