[None][perf] prepare_inputs: avoid O(seq_len) get_tokens(0) marshalling on the host

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::SizeType32>(&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<GenLlmReq::SizeType32>(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::SizeType32>(&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)

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

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)))