Repository navigation

Beam search logits processor v2 #14189

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed

kyurious-george wants to merge 204 commits into NVIDIA:main from kyurious-george:beam-search-logits-processor-v2

beam-aware logit_row tracking in _execute_logit_post_processors and b…

…uild coherent per-beam token histories by tracing parentIds

Loading branch information

kyurious-george committed May 15, 2026

commit 71acaf6e7a15681fc96d856735da4cb4fcc6e9ee

tensorrt_llm/_torch/pyexecutor/model_engine.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -3497,7 +3497,17 @@ def load_weights_from_target_model(self, @@
         def _execute_logit_post_processors(self,
                                            scheduled_requests: ScheduledRequests,
                                            outputs: dict):
-            """Apply logit post processors (in-place modify outputs Tensors) if any."""
+            """Apply logit post processors (in-place modify outputs Tensors) if any.
+            For beam search (beam_width > 1), iterates over all beams for each
+            request and passes per-beam token histories to the callback.
+            The logits tensor has ``beam_width`` rows per generation request, laid
+            out as ``[ctx_0, ..., gen_0_beam_0, gen_0_beam_1, ..., gen_1_beam_0, ...]``.
+            On 1.2.x this must run BEFORE make_decoding_batch_input so the
+            decoder sees the modified logits (forward_async reads from
+            decoding_input, not decoder_input_buffers).
+            """
             if not (self.mapping.is_last_pp_rank()):
                 return
@@ Expand All / @@ -3509,29 +3519,66 @@ def _execute_logit_post_processors(self, @@
             num_ctx_req = len(scheduled_requests.context_requests)
             logits_tensor = outputs["logits"]
+            # Track running offset into logits_tensor because during beam search,
+            # generation requests contribute beam_width rows each while context
+            # requests contribute 1 row.
+            logit_row = 0
             for idx, request in enumerate(scheduled_requests.all_requests()):
+                is_ctx = idx < num_ctx_req
+                beam_width = 1 if is_ctx else request.sampling_config.beam_width
                 logits_processors = getattr(request, "py_logits_post_processors",
                                             None)
                 if not logits_processors:
+                    logit_row += beam_width
                     continue
-                token_ids = request.get_tokens(0)
-                if idx < num_ctx_req and request.py_orig_prompt_len < len(
-                        token_ids):
-                    # Skip as we only need to apply logit processor on the last context request
-                    continue
-                logits_row = logits_tensor[idx]
-                # Reshape to align w/ the shape used in the TRT backend,
-                # so the same logit processors can be used across both backends.
-                logits_row = logits_row.view(1, 1, -1)
-                token_ids = [token_ids]
-                for lp in logits_processors:
-                    lp_params = inspect.signature(lp).parameters
-                    assert 4 <= len(lp_params) <= 5, (
-                        "Logit post processor signature must match the `LogitsProcessor` interface "
-                        "defined in `tensorrtllm.sampling_params`.")
-                    lp(request.py_request_id, logits_row, token_ids, None, None)
+                if is_ctx:
+                    token_ids = request.get_tokens(0)
+                    if request.py_orig_prompt_len < len(token_ids):
+                        logit_row += 1
+                        continue
-                logits_tensor[idx] = logits_row.view(-1)
+                if beam_width == 1:
+                    # Single beam: original path.
+                    token_ids = request.get_tokens(0)
+                    logits_row = logits_tensor[logit_row]
+                    logits_row = logits_row.view(1, 1, -1)
+                    token_ids = [token_ids]
+                    for lp in logits_processors:
+                        lp_params = inspect.signature(lp).parameters
+                        assert 4 <= len(lp_params) <= 5, (
+                            "Logit post processor signature must match the `LogitsProcessor` interface "
+                            "defined in `tensorrtllm.sampling_params`.")
+                        lp(request.py_request_id, logits_row, token_ids, None, None)
+                    logits_tensor[logit_row] = logits_row.view(-1)
+                    logit_row += 1
+                else:
+                    # Beam search: process all beams together.
+                    vocab_size = logits_tensor.shape[-1]
+                    logits_block = logits_tensor[logit_row:logit_row + beam_width]
+                    # Shape: [1, beam_width, vocab] to match TRT backend convention.
+                    logits_block_3d = logits_block.view(1, beam_width, vocab_size)
+                    # Use coherent per-beam token histories if available.
+                    # Built by PyExecutor._update_requests using parentIds
+                    # tracing from the previous sampling step.
+                    # Falls back to raw slot histories for the first step
+                    # (before any beam reassignment has occurred).
+                    coherent = getattr(request, 'py_coherent_beam_tokens', None)
+                    if coherent is not None and len(coherent) == beam_width:
+                        all_beam_token_ids = coherent
+                    else:
+                        all_beam_token_ids = [
+                            request.get_tokens(beam) for beam in range(beam_width)
+                        ]
+                    for lp in logits_processors:
+                        lp_params = inspect.signature(lp).parameters
+                        assert 4 <= len(lp_params) <= 5, (
+                            "Logit post processor signature must match the `LogitsProcessor` interface "
+                            "defined in `tensorrtllm.sampling_params`.")
+                        lp(request.py_request_id, logits_block_3d,
+                           all_beam_token_ids, None, None)
+                    logits_tensor[logit_row:logit_row + beam_width] = logits_block_3d.view(beam_width, vocab_size)
+                    logit_row += beam_width

tensorrt_llm/_torch/pyexecutor/py_executor.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -2448,6 +2448,51 @@ def _update_requests(self, @@
                 error_msg = str(e)
                 logger.error(f"Encountered an error in sampling: {error_msg}")
                 self._handle_errors(error_msg)
+                return
+            # Build coherent per-beam token histories for logits processors.
+            # request.get_tokens(beam) returns incoherent slot-accumulated
+            # histories after beam reassignment.  We trace parentIds from
+            # DecoderState to recover the true ancestral token path per beam.
+            # These are stored on the request and used by model_engine's
+            # _execute_logit_post_processors on the NEXT iteration.
+            if isinstance(self.sampler, TRTLLMSampler):
+                try:
+                    decoder_state = self.sampler.store["decoder_state"]
+                    # parentIds shape: [maxNumSequences, maxBeamWidth, maxSequenceLength]
+                    parent_ids_gpu = decoder_state.parent_ids
+                    for req in (sample_state.scheduled_requests.generation_requests
+                                if sample_state.scheduled_requests else []):
+                        bw = req.sampling_config.beam_width
+                        if bw <= 1 or req.py_seq_slot is None:
+                            continue
+                        slot = req.py_seq_slot
+                        prompt_len = req.py_orig_prompt_len
+                        tokens = [req.get_tokens(b) for b in range(bw)]
+                        num_gen = len(tokens[0]) - prompt_len
+                        if num_gen <= 1:
+                            # First step — no beam reassignment yet.
+                            continue
+                        # D2H copy of this slot's parentIds: [beamWidth, maxSeqLen]
+                        pid_host = parent_ids_gpu[slot, :bw, :].cpu()
+                        # Trace backward per beam to build coherent histories.
+                        coherent = []
+                        for b in range(bw):
+                            beam_gen = len(tokens[b]) - prompt_len
+                            slot_at_step = [0] * beam_gen
+                            s = b
+                            for g in range(beam_gen - 1, -1, -1):
+                                slot_at_step[g] = s
+                                if g > 0:
+                                    s = int(pid_host[s, prompt_len + g].item())
+                            path = list(tokens[0][:prompt_len])  # shared prompt
+                            for g in range(beam_gen):
+                                path.append(tokens[slot_at_step[g]][prompt_len + g])
+                            coherent.append(path)
+                        req.py_coherent_beam_tokens = coherent
+                except Exception:
+                    # Non-fatal: fall back to incoherent histories.
+                    pass
         def _handle_errors(self,
                            error_msg: Optional[str] = None,
@@ Expand Down @@

Navigation Menu

Repository navigation

Beam search logits processor v2 #14189

Uh oh!

Diff view

Diff view

Uh oh!

There are no files selected for viewing