Refactor code

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

Author: Shunkang

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -83,45 +83,13 @@ def _get_from_request_queue(
             pass
         return items
 
-    def _get_from_waiting_queue_attention_tp(
+    def _get_from_waiting_queue(
         self,
         waiting_queue: deque[RequestQueueItem],
         max_req_count: int,
+        enable_attention_dp: bool,
     ) -> List[RequestQueueItem]:
-        """Safely extracts up to max_req_count items from a deque.
 
-        Args:
-            waiting_queue: The queue to pop items from.
-            max_req_count: Maximum items to retrieve. Returns empty list if <=0.
-
-        Returns:
-            List of retrieved items (may be shorter than max_req_count if queue empties first).
-        """
-        # Edge case handling
-        if max_req_count <= 0:  # Handles negative/zero counts
-            return []
-
-        items = []
-        req_count = 0
-        while req_count < max_req_count and waiting_queue:
-            items.append(waiting_queue.popleft())
-            req_count += 1
-        return items
-
-    def _get_from_waiting_queue_attention_dp(
-        self,
-        waiting_queue: deque[RequestQueueItem],
-        max_req_count: int,
-    ) -> List[RequestQueueItem]:
-        """Extract requests from waiting queue with attention DP load balancing.
-
-        Args:
-            waiting_queue: Queue of pending requests
-            max_req_count: Maximum number of requests to extract
-
-        Returns:
-            List of requests that can be processed immediately
-        """
         if max_req_count <= 0:
             return []
 
@@ -130,55 +98,45 @@ def _get_from_waiting_queue_attention_dp(
         pending_requests = []
 
         # Track the request with strict requirements
-        all_ranks_num_active_requests = self.all_ranks_num_active_requests.copy(
-        )
+        scheduling_all_ranks_num_active_requests = self.all_ranks_num_active_requests.copy(
+        ) if enable_attention_dp else None
         while req_count < max_req_count and waiting_queue:
             req_item = waiting_queue.popleft()
-            can_process_now = self._can_process_attention_dp_request(
-                req_item, all_ranks_num_active_requests)
+            can_process = self._can_process_attention_dp_request(
+                req_item, scheduling_all_ranks_num_active_requests
+            ) if enable_attention_dp else True
 
-            if can_process_now:
+            if can_process:
                 items.append(req_item)
                 req_count += 1
             else:
                 pending_requests.append(req_item)
 
         # Put the pending requests back to the waiting queue
         # All ranks should have the same waiting queue
-        self.waiting_queue.extendleft(pending_requests)
+        self.waiting_queue.extendleft(reversed(pending_requests))
 
         return items
 
     def _can_process_attention_dp_request(
             self, req_item: RequestQueueItem,
             all_ranks_num_active_requests: List[int]) -> bool:
-        """Check if a request can be processed immediately.
+        """Return True if the request can be processed immediately, else False."""
 
-        Returns:
-            True if the request can be processed now, False if it should be deferred.
-        """
-        # Handle requests without schedule parameters
-        if req_item.request.py_schedule_params is None:
+        scheduling_params = req_item.request.py_scheduling_params
+        if scheduling_params is None:
             return True
 
-        schedule_params = req_item.request.py_schedule_params
-        target_dp_rank = schedule_params.attention_dp_rank
-        is_relax = schedule_params.attention_dp_relax
-
-        # Handle requests without target rank or in relax mode
-        if target_dp_rank is None or is_relax:
+        target_dp_rank = scheduling_params.attention_dp_rank
+        if target_dp_rank is None or scheduling_params.attention_dp_relax:
             return True
 
-        # Handle strict mode requests - check target rank capacity
-        target_rank_has_capacity = (
-            all_ranks_num_active_requests[target_dp_rank]
-            < self.max_num_active_requests)
-
-        if target_rank_has_capacity:
+        if all_ranks_num_active_requests[
+                target_dp_rank] < self.max_num_active_requests:
             all_ranks_num_active_requests[target_dp_rank] += 1
             return True
-        else:
-            return False
+
+        return False
 
     def enqueue_requests(self, requests: List[ExecutorRequest]):
         req_ids = []
@@ -238,7 +196,9 @@ def can_enqueue_request(self) -> bool:
         return can_enqueue and self.dist.rank == 0
 
     def _fetch_and_process_requests(
-            self, total_num_active_requests: int) -> List[RequestQueueItem]:
+            self, total_num_active_requests: int,
+            total_max_num_active_requests: int,
+            enable_attention_dp: bool) -> List[RequestQueueItem]:
         """Common logic for fetching and processing requests from the queue."""
         # Calculate timeout
         timeout = None if (total_num_active_requests == 0) and len(
@@ -264,6 +224,17 @@ def _fetch_and_process_requests(
 
         self.waiting_queue.extend(new_requests)
 
+        new_requests = self._get_from_waiting_queue(
+            self.waiting_queue,
+            total_max_num_active_requests - total_num_active_requests,
+            enable_attention_dp)
+
+        # Update performance metrics
+        if self.enable_iter_perf_stats and self.dist.rank == 0:
+            self._update_new_active_requests_queue_latency(new_requests)
+
+        return new_requests
+
     @nvtx_range("_fetch_new_requests")
     def fetch_new_requests(self,
                            num_active_requests: int) -> List[RequestQueueItem]:
@@ -280,15 +251,10 @@ def _fetch_new_requests_attention_tp(
         total_max_num_active_requests = self.max_num_active_requests
 
         # fetch and process requests into waiting queue
-        self._fetch_and_process_requests(total_num_active_requests)
-
-        new_requests = self._get_from_waiting_queue_attention_tp(
-            self.waiting_queue,
-            total_max_num_active_requests - total_num_active_requests)
-
-        # Update performance metrics
-        if self.enable_iter_perf_stats and self.dist.rank == 0:
-            self._update_new_active_requests_queue_latency(new_requests)
+        new_requests = self._fetch_and_process_requests(
+            total_num_active_requests,
+            total_max_num_active_requests,
+            enable_attention_dp=False)
 
         # Merge requests and add to active list
         merged_requests = self._merge_requests(new_requests)
@@ -307,16 +273,10 @@ def _fetch_new_requests_attention_dp(
         total_max_num_active_requests = self.dist.tp_size * self.max_num_active_requests
 
         # fetch and process requests into waiting queue
-        self._fetch_and_process_requests(total_num_active_requests)
-
-        new_requests = self._get_from_waiting_queue_attention_dp(
-            self.waiting_queue,
-            total_max_num_active_requests - total_num_active_requests)
-
-        # Update performance metrics
-        # TODO: Check whether we should update the performance metrics for all ranks
-        if self.enable_iter_perf_stats and self.dist.rank == 0:
-            self._update_new_active_requests_queue_latency(new_requests)
+        new_requests = self._fetch_and_process_requests(
+            total_num_active_requests,
+            total_max_num_active_requests,
+            enable_attention_dp=True)
 
         # Schedule attention dp requests
         new_requests_cur_rank = self._schedule_attention_dp_requests(
@@ -342,9 +302,9 @@ def _schedule_attention_dp_requests(
 
         # Prioritize the requests that are not in relax mode
         def get_relax_value(req_item):
-            if req_item.request.py_schedule_params is None:
+            if req_item.request.py_scheduling_params is None:
                 return True
-            return req_item.request.py_schedule_params.attention_dp_relax
+            return req_item.request.py_scheduling_params.attention_dp_relax
 
         new_requests = sorted(new_requests, key=get_relax_value, reverse=True)
 
@@ -353,8 +313,8 @@ def get_relax_value(req_item):
         new_requests_cur_rank = []
         for req_item in new_requests:
             scheduled = False
-            if req_item.request.py_schedule_params is not None:
-                target_dp_rank = req_item.request.py_schedule_params.attention_dp_rank
+            if req_item.request.py_scheduling_params is not None:
+                target_dp_rank = req_item.request.py_scheduling_params.attention_dp_rank
                 if target_dp_rank is not None and self.all_ranks_num_active_requests[
                         target_dp_rank] < self.max_num_active_requests:
                     self.all_ranks_num_active_requests[target_dp_rank] += 1

tensorrt_llm/executor/executor.py

@@ -29,7 +29,7 @@
                             print_colored_debug)
 from ..sampling_params import (BatchedLogitsProcessor, LogprobParams,
                                SamplingParams)
-from ..schedule_params import ScheduleParams
+from ..scheduling_params import SchedulingParams
 from .ipc import FusedIpcQueue
 from .postproc_worker import PostprocParams, PostprocWorkerConfig
 from .request import GenerationRequest, LoRARequest, PromptAdapterRequest
@@ -121,7 +121,7 @@ def generate_async(
         disaggregated_params: Optional[DisaggregatedParams] = None,
         postproc_params: Optional[PostprocParams] = None,
         multimodal_params: Optional[MultimodalParams] = None,
-        schedule_params: Optional[ScheduleParams] = None,
+        scheduling_params: Optional[SchedulingParams] = None,
     ) -> GenerationResult:
         """Generate output for the given prompt token ids in the asynchronous mode.
         Asynchronous generation accepts single prompt only.
@@ -145,7 +145,7 @@ def generate_async(
             kv_cache_retention_config=kv_cache_retention_config,
             disaggregated_params=disaggregated_params,
             multimodal_params=multimodal_params,
-            schedule_params=schedule_params)
+            scheduling_params=scheduling_params)
         result = self.submit(request)
         # release memory in time
         if hasattr(request, "multimodal_params"):

tensorrt_llm/executor/request.py

@@ -10,7 +10,7 @@
 from ..disaggregated_params import DisaggregatedParams
 from ..llmapi.llm_utils import KvCacheRetentionConfig
 from ..sampling_params import SamplingParams
-from ..schedule_params import ScheduleParams
+from ..scheduling_params import SchedulingParams
 from .postproc_worker import PostprocParams
 
 __all__ = [
@@ -96,7 +96,7 @@ def __init__(
         disaggregated_params: Optional[DisaggregatedParams] = None,
         postproc_params: Optional[PostprocParams] = None,
         multimodal_params: Optional[MultimodalParams] = None,
-        schedule_params: Optional[ScheduleParams] = None,
+        scheduling_params: Optional[SchedulingParams] = None,
     ):
         if isinstance(prompt_token_ids, list):
             self.prompt_token_ids = prompt_token_ids
@@ -121,7 +121,7 @@ def __init__(
         self.kv_cache_retention_config = kv_cache_retention_config
         self.id: Optional[int] = None
         self.disaggregated_params = disaggregated_params
-        self.schedule_params = schedule_params
+        self.scheduling_params = scheduling_params
 
     def set_id(self, id):
         assert self.id is None, f"Request ID is already set: {self.id}"

tensorrt_llm/executor/worker.py

@@ -510,8 +510,8 @@ def _deduce_max_tokens(request: GenerationRequest,
                     lp, list) else [lp]
 
             executor_request.py_schedule_params = None
-            if self._is_pytorch_backend and request.schedule_params is not None:
-                executor_request.py_schedule_params = request.schedule_params
+            if self._is_pytorch_backend and request.scheduling_params is not None:
+                executor_request.py_scheduling_params = request.scheduling_params
 
             if request.query_token_ids is not None:
                 # pytorch star attention workflow

tensorrt_llm/llmapi/llm.py

@@ -30,7 +30,7 @@
                       create_input_processor_with_hash, prompt_inputs)
 from ..logger import logger
 from ..sampling_params import SamplingParams
-from ..schedule_params import ScheduleParams
+from ..scheduling_params import SchedulingParams
 from .llm_args import (TORCH_LLMARGS_EXPLICIT_DOCSTRING,
                        TRT_LLMARGS_EXPLICIT_DOCSTRING, PybindMirror,
                        TorchLlmArgs, TrtLlmArgs)
@@ -236,8 +236,8 @@ def generate(
             KvCacheRetentionConfig, Sequence[KvCacheRetentionConfig]]] = None,
         disaggregated_params: Optional[Union[
             DisaggregatedParams, Sequence[DisaggregatedParams]]] = None,
-        schedule_params: Optional[Union[ScheduleParams,
-                                        List[ScheduleParams]]] = None,
+        scheduling_params: Optional[Union[SchedulingParams,
+                                          List[SchedulingParams]]] = None,
     ) -> Union[RequestOutput, List[RequestOutput]]:
         """Generate output for the given prompts in the synchronous mode.
         Synchronous generation accepts either single prompt or batched prompts.
@@ -285,7 +285,7 @@ def _item_at(maybe_batched: Union[Any, Sequence[Any]], pos: int) -> Any:
                 kv_cache_retention_config=_item_at(kv_cache_retention_config,
                                                    i),
                 disaggregated_params=_item_at(disaggregated_params, i),
-                schedule_params=_item_at(schedule_params, i),
+                scheduling_params=_item_at(scheduling_params, i),
                 streaming=False)
             futures.append(future)
 
@@ -311,7 +311,7 @@ def generate_async(
         kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
         disaggregated_params: Optional[DisaggregatedParams] = None,
         _postproc_params: Optional[PostprocParams] = None,
-        schedule_params: Optional[ScheduleParams] = None,
+        scheduling_params: Optional[SchedulingParams] = None,
     ) -> RequestOutput:
         """Generate output for the given prompt in the asynchronous mode.
         Asynchronous generation accepts single prompt only.
@@ -422,7 +422,7 @@ def generate_async(
             disaggregated_params=disaggregated_params,
             postproc_params=_postproc_params,
             multimodal_params=multimodal_params,
-            schedule_params=schedule_params,
+            scheduling_params=scheduling_params,
         )
 
         return RequestOutput._from_generation_result(result, prompt,

tensorrt_llm/schedule_params.py renamed to tensorrt_llm/scheduling_params.py

@@ -3,7 +3,7 @@
 
 
 @dataclass(slots=True, kw_only=True)
-class ScheduleParams:
+class SchedulingParams:
     """Schedule parameters.
 
     Args:

tests/unittest/api_stability/references/llm.yaml

@@ -126,6 +126,9 @@ methods:
       kv_cache_retention_config:
         annotation: Union[tensorrt_llm.bindings.executor.KvCacheRetentionConfig, Sequence[tensorrt_llm.bindings.executor.KvCacheRetentionConfig], NoneType]
         default: null
+      scheduling_params:
+        annotation: Optional[tensorrt_llm.scheduling_params.SchedulingParams]
+        default: null
     return_annotation: Union[tensorrt_llm.llmapi.llm.RequestOutput, List[tensorrt_llm.llmapi.llm.RequestOutput]]
   generate_async:
     parameters:
@@ -135,6 +138,9 @@ methods:
       kv_cache_retention_config:
         annotation: Optional[tensorrt_llm.bindings.executor.KvCacheRetentionConfig]
         default: null
+      scheduling_params:
+        annotation: Optional[tensorrt_llm.scheduling_params.SchedulingParams]
+        default: null
     return_annotation: tensorrt_llm.llmapi.llm.RequestOutput
   get_kv_cache_events:
     parameters: