From 51011c8a41ba1cffd17caf4cfd66f55c252739ca Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Wed, 15 Jul 2026 02:42:04 -0700 Subject: [PATCH] [None][chore] Add NVTX ranges to per-iteration ADP sync points in PyExecutor _handle_disagg_cache_errors_synced (tp_allgather error vote at the top of every executor iteration) and _can_queue (post-schedule batch-size tp_allgather under attention DP) are blocking host-side collectives, but unlike the sibling disagg status checks they carry no NVTX range, so they show up as anonymous gaps in nsys traces. Annotate both for profiling visibility. No functional change. Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a792e8f9a634..4a174b3a2ccd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3247,6 +3247,7 @@ def _handle_dynamic_draft_len(self, if spec_config is not None and spec_config.is_linear_tree else self.model_engine.max_total_draft_tokens) + @nvtx_range("_can_queue") def _can_queue(self, scheduled_batch): # can_queue_this_rank is for case that the batch is not empty on this rank, but empty on other ranks @@ -3851,6 +3852,7 @@ def _check_benchmark_disagg_gate(self, scheduled_batch: ScheduledRequests, return can_forward, True return can_forward, False + @nvtx_range("_handle_disagg_cache_errors_synced") def _handle_disagg_cache_errors_synced(self): """Rank-safe disagg cache error and poison handler.