[Performance]: Beam search re-materializes Result.output_token_ids from pybind once per beam

Inizia in tensorrt_llm/executor/result.py: _handle_sequence (intorno alla riga 369) legge response_tensors.output_token_ids nei due punti contrassegnati, e il ciclo sui beam in _handle_response (intorno alla riga 692) lo chiama una volta per ogni beam. Verifica che la proprietà pybind ricostruisca la sua lista a ogni accesso, poi passa in basso una singola copia materializzata, come mostra il diff nella issue. È completato quando gli output della beam search restano invariati e i test esistenti dell'executor della beam search continuano a passare.

Proposal to improve performance

Background

For beam-search requests, GenerationResultBase._handle_response iterates over beams and calls _handle_sequence once per beam:

finish_reasons = response_result.finish_reasons
# output_token_ids = (beams, tokens)
if self.sampling_params.use_beam_search:
    for beam_idx, _ in enumerate(response_result.output_token_ids):   # result.py:692
        self._handle_sequence(finish_reasons, response_result,
                              beam_idx, logprobs_result,
                              req_perf_metrics_dict)

_handle_sequence then reads the same property again, once per call:

if self.sampling_params.use_beam_search:
    output.token_ids = [
        *decoder_output_prefix,
        *response_tensors.output_token_ids[src_idx],      # result.py:402
    ]
else:
    ...
    output.token_ids.extend(response_tensors.output_token_ids[src_idx])   # result.py:407

LlmResult.output_token_ids is not in LlmResult.py_result_properties, so LlmResult.__getattr__ falls through to getattr(self._result, item) on the pybind bindings.executor.Result. That property returns a freshly built nested Python list on every access — it is not cached on either side:

>>> a, b = llm_result.output_token_ids, llm_result.output_token_ids
>>> a is b
False
>>> a[0] is b[0]
False
>>> a == b
True

So for a request with N beams, one response converts the full (N, tokens) structure from C++ N + 1 times when once is sufficient. Because the per-conversion cost is itself O(N · tokens), the redundant work is quadratic in beam width.

This is the dominant host-side cost for the "short decode length, very wide beams" pattern (e.g. SID / generative-recommendation serving with n = best_of = 500, max_tokens = 3), where _handle_sequence does almost no work per beam other than these reads.

Measurement

Offline, no GPU, no model, no executor. A real bindings.executor.Result is constructed and populated (output_token_ids is settable), wrapped in a real LlmResult, and the real GenerationResultBase._handle_sequence is driven once per beam exactly as _handle_response does.

Three arms, interleaved, gc.collect() before every sample, median of 25:

• A stock — real _handle_sequence against the real LlmResult
• C patched — the same real source, textually rewritten so response_tensors.output_token_ids resolves to a binding materialized once per response (i.e. exactly the suggested fix below); every other attribute still reads the real LlmResult per beam, so this arm does not overstate the win
• B snapshot — every attribute pre-materialized; an upper bound

Access counts were obtained by running the real code path against a counting forwarder, not by reading the source by hand. For one 500-beam response:

attribute	reads	served by
time_breakdown_metrics	1000	instance attr (cheap)
output_token_ids	500	pybind
request_perf_metrics	500	pybind
cum_log_probs / log_probs / generation_logits / additional_*_outputs	500 each	py_result (Python-side)
finish_reasons	1	pybind

Total 4501 attribute reads per response, 1001 of them pybind-served.

Results (ms per response, tokens = 3)

beams	one pybind read	A stock	C patched	C saving	B upper bound	B saving
4	0.53 µs	0.122	0.120	−0.002 (1.6%)	0.097	−0.025 (20.5%)
20	1.25 µs	0.195	0.187	−0.008 (4.1%)	0.127	−0.068 (34.9%)
100	6.25 µs	1.016	0.458	−0.558 (54.9%)	0.303	−0.713 (70.2%)
300	16.20 µs	5.848	1.047	−4.801 (82.1%)	0.682	−5.166 (88.3%)
500	25.94 µs	14.428	1.543	−12.885 (89.3%)	0.985	−13.443 (93.2%)

At 500 beams, 89% of the per-response beam loop is redundant pybind re-conversion of one property. 100 → 500 beams is 5× the width and 23× the waste, confirming the quadratic behaviour.

Below ~20 beams the effect is inside measurement noise (arm C's delta at beam=20 is 0.008 ms against a per-sample sd of 0.03–0.04 ms), so this is specifically a wide-beam problem.

Cross-check

Two independent estimates of the same quantity agree:

• micro: 500 reads × 25.94 µs = 12.97 ms
• macro: arm A − arm C = 12.89 ms
• ratio 0.99

The residual between arm C and arm B (0.558 ms at 500 beams) is accounted for by the other 501 pybind reads (request_perf_metrics ×500 at ≈1.1 µs, finish_reasons ×1).

Correctness

Arms A, B and C produce identical token_ids, cumulative_logprob, finish_reason and logprobs for all 500 outputs, verified at every beam width in the sweep.

Version note

Measured against 1.3.0rc29; the full environment is in the last section below. The
structure on main is identical for these reads (_handle_sequence at result.py:369,
the two output_token_ids reads at :402/:407, the beam loop at :692), so the
analysis and the suggested fix carry over unchanged. main does add three further per-beam
getattr calls; quantified under Misc discussion below.

Suggested fix

Materialize once per response and pass it down. Backward compatible: the parameter defaults to None, which preserves the current behaviour for every other call site.

--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -372,10 +372,19 @@
                          sequence_index,
                          logprobs_result=None,
                          req_perf_metrics_dict: Optional[dict[str,
-                                                              float]] = None):
-        """Handle a single sequence in the response."""
+                                                              float]] = None,
+                         output_token_ids=None):
+        """Handle a single sequence in the response.
+
+        ``output_token_ids`` lets the caller pass an already-materialized copy.
+        For the pybind ``Result`` this property rebuilds the whole
+        (beams, tokens) nested list on every access, so the beam-search path --
+        which calls this once per beam -- would otherwise rebuild it N times.
+        """
         seq_idx = sequence_index
         src_idx = sequence_index if self.sampling_params.use_beam_search else 0
+        if output_token_ids is None:
+            output_token_ids = response_tensors.output_token_ids
 
         output = self._outputs[seq_idx]
         output.disaggregated_params = self.disaggregated_params
@@ -399,12 +408,12 @@
             # Beam search enforces returning all generated tokens
             output.token_ids = [
                 *decoder_output_prefix,
-                *response_tensors.output_token_ids[src_idx],
+                *output_token_ids[src_idx],
             ]
         else:
             if decoder_output_prefix and not output.token_ids:
                 output.token_ids.extend(decoder_output_prefix)
-            output.token_ids.extend(response_tensors.output_token_ids[src_idx])
+            output.token_ids.extend(output_token_ids[src_idx])
 
         if response_tensors.cum_log_probs is not None:
             output.cumulative_logprob = response_tensors.cum_log_probs[src_idx]
@@ -689,10 +698,12 @@
             finish_reasons = response_result.finish_reasons
             # output_token_ids = (beams, tokens)
             if self.sampling_params.use_beam_search:
-                for beam_idx, _ in enumerate(response_result.output_token_ids):
+                output_token_ids = response_result.output_token_ids
+                for beam_idx in range(len(output_token_ids)):
                     self._handle_sequence(finish_reasons, response_result,
-                                          beam_idx, logprobs_result,
-                                          req_perf_metrics_dict)
+                                          beam_idx, logprobs_result,
+                                          req_perf_metrics_dict,
+                                          output_token_ids=output_token_ids)
             else:
                 self._handle_sequence(finish_reasons, response_result,
                                       response_result.sequence_index,

+16 / −4 lines. No behaviour change for non-beam requests, for streaming, or for any caller that does not pass the new argument.

Scope, impact and risk

• Impact: pure host-side latency, incurred once per non-streaming response (and per streamed chunk in streaming mode, where the same loop runs). At 500 beams it is ~13 ms of serial host time per response.
• Risk: low. The change is loop-invariant code motion of a read-only property. Verified output-equivalent above.
• Deliberately out of scope (to keep this PR to a single concern): request_perf_metrics is read per beam through the same pybind path — once at :486 for the is not None test and again at :487 when set. In the measurement above it is unset, so it costs one pybind read per beam, ≈0.56 ms per 500-beam response; when it is set the second read doubles that. It would benefit from the same treatment. Happy to file separately if useful.

Report of performance regression

Not a regression. The behaviour is present in 1.3.0rc29 and unchanged on main; it has
simply not been noticeable at typical beam widths, because the cost is quadratic and only
becomes dominant at wide beams.

Misc discussion on performance

Related, different code path. #20166 vectorizes the final beam assembly inside
_prepare_beam_history_cba (worker/sampler side). This issue is about
GenerationResultBase._handle_response / _handle_sequence (client side), a separate
cost incurred once per response in the API process. The two are complementary.

Two smaller instances of the same pattern, deliberately left out of the suggested fix
to keep it to a single concern:

• request_perf_metrics is read per beam through the same pybind path -- once at
  result.py:486 for the is not None test and again at :487 when set. In the
  measurement above it is unset, so it costs one pybind read per beam, ~0.56 ms per
  500-beam response; when set, the second read doubles that.
• On main, _handle_sequence additionally performs three getattr(response_tensors, ...)
  per beam to build _SpecDecCounters (per_pos_drafted, per_pos_accepted,
  spec_dec_totals). The first two are absent from LlmResult, so each pays an
  AttributeError round trip (~0.67 us measured), adding roughly 0.67 ms per 500-beam
  response on top of the figures above.

Generality. Nothing here is specific to one workload: any beam-search user with wide
beams pays N+1 full materializations of an (N, tokens) structure per response. The cost is
quadratic in beam width, so it grows fastest exactly where beam search is used hardest.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Searched output_token_ids beam (0 results), _handle_sequence, beam search host overhead and pybind beam performance. No existing issue covers this; #20166 is the
closest and is a different code path (linked above).

Your current environment (if you think it is necessary)

System Information:

• OS: Linux x86_64
• Python version: 3.12.7
• CUDA version: 13.2 (V13.2.78)
• GPU model(s): NVIDIA L20
• Driver version: 580.105.08
• TensorRT version: not separately installed (PyTorch backend only)
• PyTorch version: 2.14.0+cu132
• TensorRT-LLM version: 1.3.0rc29

This is a host-side code path: no GPU work is involved, so the measurements do not
depend on the GPU model or driver. Any CPU reproduces them.