From 19164332ee96bc85b05daf4fd96497363b041d5a Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 23 Jul 2026 07:34:50 +0000 Subject: [PATCH] [None][perf] prepare_inputs: avoid O(seq_len) get_tokens(0) marshalling on the host `request.get_tokens(0)` marshals the whole C++ VecTokens (entire sequence) into a fresh Python list of boxed ints -- O(seq_len) host work that grows linearly with ISL. Three call sites in `_prepare_tp_inputs` / `cuda_graph_runner` pay this only to read a length or a small slice, and re-pay it every iteration. In chunked prefill the context loop re-marshals the full prompt for every chunk -> O(L^2/chunk) per request over the prefill. Normal decode is already immune (get_last_tokens(0), O(1)); the waste is in the prefill/context and MTP first-draft paths. Changes: - New nanobind binding `LlmRequest.get_tokens_range(beam, begin, end)` that copies only [begin, end) (O(chunk)) instead of the whole VecTokens. - context loop and first_draft loop: use get_tokens_range for the chunk and get_num_tokens(0) (O(1)) for lengths, instead of get_tokens(0) + slice. - cuda_graph_runner first-draft branch: len(get_tokens(0)) -> get_num_tokens(0). Output is bit-identical (get_tokens_range returns the same subrange). Measured on a DSv4-Pro disagg CTX worker (c3120, non-overlap, CTX nsys): _prepare_inputs p50 27.53 ms -> 15.10 ms (-45%), p90 43.29 -> 20.43 ms; 100% request success. Complements #16734 (which removes the overlap device-scalar sync, 268 -> 19.96 ms); together they target the two independent costs in _prepare_inputs. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 4.8 --- .../nanobind/batch_manager/bindings.cpp | 28 +++++++++++++++++++ .../_torch/pyexecutor/cuda_graph_runner.py | 4 ++- .../_torch/pyexecutor/model_engine.py | 23 ++++++++++----- 3 files changed, 47 insertions(+), 8 deletions(-) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index 202228b394d8..776a3bdd22a4 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -101,6 +101,34 @@ void initBindings(nb::module_& m) .def("get_token", &GenLlmReq::getToken, nb::arg("beam"), nb::arg("pos")) .def("get_tokens", nb::overload_cast(&GenLlmReq::getTokens, nb::const_), nb::arg("beam")) .def("get_tokens", nb::overload_cast<>(&GenLlmReq::getTokens, nb::const_)) + // Copies only [begin, end) -> O(end-begin), vs get_tokens(beam) which + // marshals the whole O(seq_len) VecTokens into a Python list. + .def( + "get_tokens_range", + [](GenLlmReq const& self, GenLlmReq::SizeType32 beam, GenLlmReq::SizeType32 begin, + GenLlmReq::SizeType32 end) + { + auto const& tokens = self.getTokens(beam); + auto const n = static_cast(tokens.size()); + if (begin < 0) + { + begin = 0; + } + if (begin > n) + { + begin = n; + } + if (end < begin) + { + end = begin; + } + if (end > n) + { + end = n; + } + return GenLlmReq::VecTokens(tokens.begin() + begin, tokens.begin() + end); + }, + nb::arg("beam"), nb::arg("begin"), nb::arg("end")) .def("get_last_tokens", nb::overload_cast(&GenLlmReq::getLastTokens), nb::arg("beam")) .def("get_last_tokens", nb::overload_cast<>(&GenLlmReq::getLastTokens)) .def("get_beam_width_by_iter", &GenLlmReq::getBeamWidthByIter, nb::arg("for_next_iteration") = false) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index b8b67ed5ded8..c2a6a5bb4de1 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -204,7 +204,9 @@ def _get_seq_len_mode( num_draft_tokens = self.spec_config.max_draft_len if is_spec_request else 0 # First draft if request.py_is_first_draft: - total_seq_len = len(request.get_tokens(0)) + # get_num_tokens is O(1); len(get_tokens(0)) marshals the + # whole O(seq_len) VecTokens into a Python list just for len. + total_seq_len = request.get_num_tokens(0) # With overlap scheduler disabled or dummy request or not assigned to a batch, elif not overlap_scheduler_enabled or request.is_dummy or request.py_batch_idx is None: total_seq_len = request.max_beam_num_tokens + num_draft_tokens diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 50aa317e9e43..7e37c35cbacd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3938,11 +3938,16 @@ def append_cross_attention_state(request: LlmRequest, for request in scheduled_requests.context_requests: request_ids.append(request.py_request_id) - all_prompt_tokens = request.get_tokens(0) draft_lens.append(0) begin_compute = request.context_current_position end_compute = begin_compute + request.context_chunk_size - prompt_tokens = all_prompt_tokens[begin_compute:end_compute] + # Fetch only the current chunk. get_tokens(0) marshals the whole + # O(seq_len) VecTokens into a Python list of boxed ints; chunked + # prefill re-enters this loop for every chunk of the same prompt, so + # that is O(L) per chunk = O(L^2/chunk) over the prefill. + # get_tokens_range copies only [begin, end) -> O(chunk). + prompt_tokens = request.get_tokens_range(0, begin_compute, + end_compute) position_ids.extend( range(begin_compute, begin_compute + len(prompt_tokens))) @@ -3987,7 +3992,7 @@ def append_cross_attention_state(request: LlmRequest, request.py_multimodal_data, begin_compute=past_seen_token_num, end_compute=end_compute, - prompt_len=len(all_prompt_tokens), + prompt_len=request.get_num_tokens(0), ) mm_data = request.py_multimodal_data or {} cumsum = mm_data.get('multimodal_embed_mask_cumsum') @@ -4209,12 +4214,16 @@ def append_cross_attention_state(request: LlmRequest, for request in first_draft_requests: request_ids.append(request.py_request_id) - all_prompt_tokens = request.get_tokens(0) draft_lens.append(0) - begin_compute = len( - all_prompt_tokens) - self.original_max_draft_len - 1 + # Only the length and the last (original_max_draft_len+1) tokens are + # needed here; get_num_tokens is O(1) and get_tokens_range copies only + # the requested window, whereas get_tokens(0) marshals the whole + # O(seq_len) VecTokens into a Python list. + _num_tokens = request.get_num_tokens(0) + begin_compute = _num_tokens - self.original_max_draft_len - 1 end_compute = begin_compute + self.original_max_draft_len + 1 - prompt_tokens = all_prompt_tokens[begin_compute:end_compute] + prompt_tokens = request.get_tokens_range(0, begin_compute, + end_compute) position_ids.extend( range(begin_compute, begin_compute + len(prompt_tokens)))