Code refactor

Signed-off-by: Yilin Fan <[email protected]>

Author: nv-yilinf

cpp/include/tensorrt_llm/batch_manager/llmRequest.h

@@ -1261,7 +1261,7 @@ class GenericLlmRequest
     {
         if (mPerfMetrics.timingMetrics.firstScheduledTime == executor::RequestPerfMetrics::TimePoint{})
         {
-            mPerfMetrics.timingMetrics.firstScheduledTime = getCurrentSteadyClock();
+            mPerfMetrics.timingMetrics.firstScheduledTime = getSteadyClockNow();
         }
     }
 
@@ -1677,7 +1677,7 @@ class GenericLlmRequest
         {
             return false;
         }
-        auto const currentTime = getCurrentSteadyClock();
+        auto const currentTime = getSteadyClockNow();
         auto const elapsed = (std::chrono::duration_cast<Duration>(currentTime - mStartTime));
         TLLM_LOG_DEBUG("Checked timeOut for request %ld with allotted Time %ld after time %ld and got %d", mRequestId,
             mAllottedTimeMs->count(), elapsed.count(), (elapsed >= mAllottedTimeMs));
@@ -1794,7 +1794,7 @@ class GenericLlmRequest
         if (finishReason == executor::FinishReason::kTIMED_OUT)
         {
             TLLM_LOG_DEBUG("Request %ld finished by timeout after %f sec", mRequestId,
-                std::chrono::duration<float>(getCurrentSteadyClock() - mStartTime).count());
+                std::chrono::duration<float>(getSteadyClockNow() - mStartTime).count());
         }
         if (finishReason == executor::FinishReason::kCANCELLED)
         {
@@ -1832,10 +1832,9 @@ class GenericLlmRequest
 
     void updatePerfMetrics(executor::IterationType iter)
     {
-        auto const currentTokenTime = getCurrentSteadyClock();
-
         if (!mPerfMetrics.firstIter)
         {
+            auto const currentTokenTime = getSteadyClockNow();
             mPerfMetrics.firstIter = iter;
             mPerfMetrics.timingMetrics.firstTokenTime = currentTokenTime;
         }
@@ -1844,6 +1843,7 @@ class GenericLlmRequest
 
         if (isFinished())
         {
+            auto const currentTokenTime = getSteadyClockNow();
             mPerfMetrics.lastIter = iter;
             mPerfMetrics.timingMetrics.lastTokenTime = currentTokenTime;
         }
@@ -2060,7 +2060,7 @@ class GenericLlmRequest
     // Cache salt id for each request.
     std::optional<CacheSaltIDType> mCacheSaltID{std::nullopt};
 
-    // The offset between local steady clock and glabol steady clock (at rank 0)
+    // The offset between local steady clock and global steady clock (at rank 0)
     std::optional<Duration> mGlobalSteadyClockOffset;
 private:
     void initialize(
@@ -2158,9 +2158,9 @@ class GenericLlmRequest
 
         if (mReturnPerfMetrics)
         {
-            mPerfMetrics.timingMetrics.arrivalTime = arrivalTime.value_or(getCurrentSteadyClock());
+            mPerfMetrics.timingMetrics.arrivalTime = arrivalTime.value_or(getSteadyClockNow());
         }
-        mStartTime = getCurrentSteadyClock();
+        mStartTime = getSteadyClockNow();
     }
 
     TensorPtr createListTensor(std::list<VecTokens> const& wordsList)
@@ -2197,7 +2197,8 @@ class GenericLlmRequest
         }
     }
 
-    TimePoint getCurrentSteadyClock() const {
+    // If mGlobalSteadyClockOffset is set, return a global steady clock time point, otherwise return local steady clock time point
+    TimePoint getSteadyClockNow() const {
         const TimePoint time_point = std::chrono::steady_clock::now();
 
         return maybeToGlobalSteadyClock(time_point);

requirements.txt

@@ -75,4 +75,3 @@ triton==3.3.1; platform_machine == "x86_64"
 tiktoken
 blobfile
 openai-harmony==0.0.4
-nvidia-cutlass-dsl==4.1.0; python_version >= "3.12"

tensorrt_llm/_torch/pyexecutor/executor_request_queue.py

@@ -44,7 +44,7 @@ class ExecutorRequestQueue:
     def __init__(self, dist: Distributed, enable_attention_dp: bool,
                  max_batch_size: int, max_beam_width: int,
                  max_num_active_requests: int, enable_iter_perf_stats: bool,
-                 batch_wait_timeout_ms: float, is_disaggregated: bool, monotonic_ts_offset: float):
+                 batch_wait_timeout_ms: float, is_disaggregated: bool, global_steady_clock_offset: float):
         self.dist = dist
         self.request_queue: queue.Queue[RequestQueueItem] = queue.Queue()
         self.waiting_queue: deque[RequestQueueItem] = deque()
@@ -60,7 +60,7 @@ def __init__(self, dist: Distributed, enable_attention_dp: bool,
         self.start_times = {}
         self.active = True
         self.batch_wait_timeout_ms = batch_wait_timeout_ms
-        self.monotonic_ts_offset = monotonic_ts_offset
+        self.global_steady_clock_offset = global_steady_clock_offset
 
         # State tracking
         self.num_fetch_requests = 0
@@ -612,7 +612,9 @@ def _merge_requests(
         else:
             req_with_children = []
             for req_item in new_requests:
-                req_item.request.py_global_steady_clock_offset = self.monotonic_ts_offset
+                if self.global_steady_clock_offset:
+                    req_item.request.py_global_steady_clock_offset = self.global_steady_clock_offset
+
                 req = executor_request_to_llm_request(
                     req_item.id, req_item.request, req_item.child_req_ids,
                     self._should_exclude_last_generation_logits())

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -166,7 +166,7 @@ def __init__(self,
         self.device_id = torch.cuda.current_device()
         self.global_rank = global_mpi_rank()
         self.dist = dist
-        self.monotonic_ts_offset = self._get_monotonic_ts_offset()
+        self.global_steady_clock_offset = self._get_global_steady_clock_offset()
 
         self.peft_cache_config = peft_cache_config
 
@@ -262,7 +262,7 @@ def __init__(self,
             enable_iter_perf_stats=self.enable_iter_perf_stats,
             batch_wait_timeout_ms=self.batch_wait_timeout_ms,
             is_disaggregated=kv_cache_transceiver is not None,
-            monotonic_ts_offset = self.monotonic_ts_offset
+            global_steady_clock_offset=self.global_steady_clock_offset,
         )
         self.executor_request_queue.set_exclude_last_generation_logits(
             self.disable_overlap_scheduler, self.dist.pp_size)
@@ -365,14 +365,18 @@ def start_worker(self):
                 self.worker_thread.start()
                 self.worker_started = True
 
-    def _get_monotonic_ts_offset(self):
+    def _get_global_steady_clock_offset(self):
         assert self.global_rank >= 0, "rank should be >= 0"
+
+        # Sync all ranks
         self.dist.barrier()
+        # Immediately take the local steady clock timestamp
         local_timestamp = time.monotonic()
-        timestamps = self.dist.allgather(local_timestamp)
+        all_rank_timestamps = self.dist.allgather(local_timestamp)
         if self.global_rank == 0:
-            logger.info(f"monotonic_ts_offsets for each rank: {[local_timestamp - ts for ts in timestamps]}")
-        return timestamps[0] - local_timestamp
+            logger.info(f"global_steady_clock_offset at each rank: {[local_timestamp - ts for ts in all_rank_timestamps]}")
+        # Compute the steady clock offset between rank 0 and current rank
+        return all_rank_timestamps[0] - local_timestamp
 
     def __enter__(self):
         return self
@@ -1904,6 +1908,7 @@ def _handle_responses(self):
             request.draft_tokens = request.py_draft_tokens
             request.decoding_iter = request.py_decoding_iter
 
+            # Skip active requests that are not scheduled
             if request.return_perf_metrics and request.py_decoding_iter >= 1:
                 request.update_perf_metrics(self.model_engine.iter_counter)
 

tensorrt_llm/executor/result.py

@@ -7,7 +7,6 @@
                     NamedTuple, Optional, TypeAlias, Union)
 from weakref import WeakMethod
 
-from tensorrt_llm.logger import logger
 import torch
 import torch.nn.functional as F
 
@@ -323,11 +322,17 @@ def _handle_response(self,
                 self._outputs[0] = response.res
             else:
                 self._outputs[0]._postprocess_result = response.res
+
             self._outputs[0].request_perf_metrics = response.request_perf_metrics
-            if response.disaggregated_params:
-                self._outputs[0].disaggregated_params = response.disaggregated_params
-            else:
-                self._outputs[0].disaggregated_params = self.disaggregated_params
+            if not self._outputs[0].disaggregated_params:
+                disaggregated_params = response.disaggregated_params
+
+                # Generation only response has no disaggregated_params attached
+                if not disaggregated_params:
+                    disaggregated_params = self.disaggregated_params
+
+                self._outputs[0].disaggregated_params = disaggregated_params
+
             if response.metrics:
                 self.metrics_dict = response.metrics
 

tensorrt_llm/llmapi/llm.py

@@ -354,6 +354,9 @@ def generate_async(
         if self._executor is None or self._executor.is_shutdown():
             raise RuntimeError("LLM is shutting down")
 
+        arrival_time = steady_clock_now(
+        ) if self.args.return_perf_metrics else None
+
         sampling_params = self._prepare_sampling_params(sampling_params)
         cache_salt_id = get_cache_salt_id(
             cache_salt) if cache_salt is not None else None
@@ -464,10 +467,6 @@ def generate_async(
         if _postproc_params:
             _postproc_params.postproc_args.num_prompt_tokens = len(
                 prompt_token_ids)
-
-        arrival_time = steady_clock_now(
-        ) if self.args.return_perf_metrics else None
-
         result = self._executor.generate_async(
             prompt_token_ids,
             query_token_ids=query_token_ids,

tensorrt_llm/serve/openai_disagg_server.py

@@ -57,14 +57,12 @@ def __init__(self,
         self.perf_metrics_max_requests = config.perf_metrics_max_requests
         if self.perf_metrics_max_requests > 0:
             # record corresponding keys of context and generation servers for perf metrics
-            # (ctx_server, gen_server, ctx_request_id, server_start_ts, server_first_token_ts)
+            # (ctx_server, gen_server, ctx_request_id, server_arrival_time, server_first_token_ts)
             self.perf_metrics_keys = deque(maxlen=self.perf_metrics_max_requests)
             self.perf_metrics_keys_lock = asyncio.Lock()
             # server_url -> {ctx_request_id: perf_metrics}
             self.server_perf_metrics: dict[str, dict[int, dict]] = {}
 
-            # server_url -> the perf metric timestamp offset between the disagg server and worker server
-            self.server_perf_ts_offsets: dict[str, float] = {}
         else:
             self.perf_metrics_keys = None
             self.perf_metrics_keys_lock = None
@@ -110,7 +108,7 @@ async def lifespan(app: FastAPI):
             await self.wait_for_servers_ready(server_start_timeout_secs)
 
             if self.perf_metrics_max_requests > 0:
-                await self.query_perf_ts_offsets(self.session)
+                await self.set_steady_clock_offsets(self.session)
 
             if self.metadata_server:
                 logger.info("Starting server monitoring via metadata service")
@@ -143,7 +141,7 @@ async def lifespan(app: FastAPI):
         @self.app.middleware("http")
         async def add_process_time_header(raw_request: Request, call_next):
             start_time = time.monotonic()
-            raw_request.state.server_start_ts = start_time
+            raw_request.state.server_arrival_time = start_time
             response = await call_next(raw_request)
             return response
 
@@ -202,7 +200,7 @@ async def version(self) -> JSONResponse:
 
     async def _add_perf_metrics_keys(self, ctx_server: str, gen_server: str, ctx_request_id: int, raw_request: Request):
         async with self.perf_metrics_keys_lock:
-            self.perf_metrics_keys.append((ctx_server, gen_server, ctx_request_id, raw_request.state.server_start_ts, raw_request.state.server_first_token_ts))
+            self.perf_metrics_keys.append((ctx_server, gen_server, ctx_request_id, raw_request.state.server_arrival_time, raw_request.state.server_first_token_ts))
 
     async def perf_metrics(self) -> JSONResponse:
         if self.perf_metrics_keys is None:
@@ -239,27 +237,23 @@ async def perf_metrics(self) -> JSONResponse:
                 raise exc
 
             remain_keys = []
-            for ctx_server, gen_server, ctx_request_id, server_start_ts, server_first_token_ts in self.perf_metrics_keys:
+            for ctx_server, gen_server, ctx_request_id, server_arrival_time, server_first_token_ts in self.perf_metrics_keys:
                 gen_perf_metrics = self.server_perf_metrics[gen_server].pop(ctx_request_id, None)
                 if gen_perf_metrics is None:
                     # generation not finished
-                    remain_keys.append((ctx_server, gen_server, ctx_request_id, server_start_ts, server_first_token_ts))
+                    remain_keys.append((ctx_server, gen_server, ctx_request_id, server_arrival_time, server_first_token_ts))
                     continue
                 ctx_perf_metrics = self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)
                 return_metrics.append({
                     "ctx_server": ctx_server,
                     "gen_server": gen_server,
-                    "disagg_server_start_ts": server_start_ts,
+                    "disagg_server_arrival_time": server_arrival_time,
                     "disagg_server_first_token_ts": server_first_token_ts,
                     "ctx_perf_metrics": ctx_perf_metrics,
                     "gen_perf_metrics": gen_perf_metrics})
             self.perf_metrics_keys = deque(remain_keys, maxlen=self.perf_metrics_max_requests)
 
-        response = {
-            "server_perf_timestamp_offsets": self.server_perf_ts_offsets,
-            "perf_metrics": return_metrics
-        }
-        return JSONResponse(content=response)
+        return JSONResponse(content=return_metrics)
 
 
     async def openai_completion(self, req: CompletionRequest, raw_request: Request) -> Response:
@@ -514,28 +508,35 @@ async def send_completion_request(self, url: str, request: CompletionRequest) ->
     async def send_chat_request(self, url: str, request: ChatCompletionRequest) -> ChatCompletionResponse:
         return await self.send_request(url, request, "/v1/chat/completions", ChatCompletionResponse, self.create_chat_generator)
 
-    async def query_perf_ts_offsets(self, session: aiohttp.ClientSession):
-        async def query_perf_ts_offset(server_url: str) -> Optional[float]:
+    async def set_steady_clock_offsets(self, session: aiohttp.ClientSession):
+        STEADY_CLOCK_OFFSET_ENDPOINT = "/steady_clock_offset"
+        async def query_steady_clock_offset(server_url: str) -> Optional[float]:
             try:
                 originate_ts = time.monotonic()
-                async with session.get(server_url + '/perf_ts_offset') as response:
+                async with session.get(server_url + STEADY_CLOCK_OFFSET_ENDPOINT) as response:
                     destination_ts = time.monotonic()
                     if response.status == 200:
                         response = await response.json()
+                        # Compute the steady clock timestamp difference using the NTP clock synchronization algorithm. https://en.wikipedia.org/wiki/Network_Time_Protocol#Clock_synchronization_algorithm
                         receive_ts = response['receive_ts']
                         transmit_ts = response['transmit_ts']
                         delay = (destination_ts - originate_ts) - (transmit_ts - receive_ts)
-                        offset = - ((receive_ts - originate_ts) + (transmit_ts - destination_ts)) / 2
+                        offset = ((receive_ts - originate_ts) + (transmit_ts - destination_ts)) / 2
                         return delay, offset
                     else:
                         return None, None
             except Exception:
                 return None
+        async def set_steady_clock_offset(server_url: str, offset: float) -> Optional[float]:
+            payload = {"offset": offset}
+            async with session.post(server_url + STEADY_CLOCK_OFFSET_ENDPOINT, json=payload) as response:
+                if response.status != 200:
+                    logger.warning(f"Cannot set disagg server steady clock offset for server {server_url}, the perf metrics timestamps could be mis-aligned")
         for server_url in self.ctx_servers + self.gen_servers:
-            delay, offset = await query_perf_ts_offset(server_url)
-            self.server_perf_ts_offsets[server_url] = offset
+            delay, offset = await query_steady_clock_offset(server_url)
             logger.info(f'Server: {server_url}, delay: {delay} second, offset: {offset} second')
-        logger.info(f"Server perf metrics timestamp offsets: {self.server_perf_ts_offsets}")
+            # Negate the offset so that worker servers can adjust their steady block by adding the new offset
+            await set_steady_clock_offset(server_url, -offset)
 
     @classmethod
     async def check_server_ready(cls, session: aiohttp.ClientSession, server_url: str) -> bool: