[None][feat] Call LLM(batched_logits_processor=...) on the PyTorch backend

docs/source/features/sampling.md

@@ -371,3 +371,34 @@ llm.generate(["Hello, my name is"], sampling_params)
 ```
 
 You can find a more detailed example on logits processors [here](source:examples/llm-api/llm_logits_processor.py).
+
+### Batched logits processor
+
+A per-request logits processor is called once per request on every decoding step. When many requests share the same processing, pass one [`BatchedLogitsProcessor`](source:tensorrt_llm/sampling_params.py) to the `LLM` instead and opt requests in with `apply_batched_logits_processor=True`. It is then called once per step with the logits, token ids, request ids and client ids of every scheduled request that opted in, after each of those requests' own `logits_processor`(s) have run. Each entry of `logits` is that request's rows, shaped `[beam_width, 1, vocab_size]`, and must be modified in place.
+
+```python
+from typing import List, Optional
+
+import torch
+
+from tensorrt_llm import LLM, SamplingParams
+from tensorrt_llm.sampling_params import BatchedLogitsProcessor
+
+class MyBatchedLogitsProcessor(BatchedLogitsProcessor):
+    def __call__(self,
+        req_ids: List[int],
+        logits: List[torch.Tensor],
+        token_ids: List[List[List[int]]],
+        stream_ptr: Optional[int],
+        client_ids: List[Optional[int]]
+    ) -> None:
+        for rows in logits:
+            rows *= rows
+
+llm = LLM(model='nvidia/Llama-3.1-8B-Instruct-FP8',
+          batched_logits_processor=MyBatchedLogitsProcessor())
+sampling_params = SamplingParams(apply_batched_logits_processor=True)
+llm.generate(["Hello, my name is"], sampling_params)
+```
+
+Like a per-request processor, the batched one runs where the model runs, which may be a different process from the caller: it must be picklable, and state it keeps is not shared with the calling process. Also like a per-request processor, under a one-model speculative decoding worker that keeps its logits vocabulary-sharded it is handed a gathered copy after acceptance has already run, so its edits do not reach the sampled tokens there.

tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py

@@ -48,6 +48,7 @@
 )
 from tensorrt_llm._utils import nvtx_range
 from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, PrefillCudaGraphBackend
+from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import Mapping
 
 from ...lora import LoraParamBuilder, make_cuda_graph_lora_manager
@@ -184,6 +185,11 @@ def __init__(
         # post-processors read the result instead of the per-slot token lists
         # (see _start_beam_histories). Left None when the sampler has none.
         self.start_beam_histories_callable = None
+        # LLM(batched_logits_processor=...), registered by PyExecutor. Called
+        # once per step with the logits of every scheduled request that set
+        # SamplingParams.apply_batched_logits_processor, after the request's
+        # own processors (see _execute_logit_post_processors).
+        self.batched_logits_processor = None
         # Stages (or clears) in-graph sampling state once the batch is settled.
         self._stage_in_graph_sampling = None
         # Set by warmup: whether a forward is the TRTLLM-Gen FMHA JIT warmup,
@@ -808,13 +814,9 @@ def _forward_step(
         return outputs
 
     @staticmethod
-    def _apply_logits_processors(
-        request, logits_processors, logits_tensor, beam_width, token_ids, logits_row_offset
-    ) -> None:
-        logits_rows = logits_tensor[logits_row_offset : logits_row_offset + beam_width]
-        # Reshape to align w/ the shape used in the TRT backend,
-        # so the same logit processors can be used across both backends.
-        logits_rows = logits_rows.view(beam_width, 1, -1)
+    def _apply_logits_processors(request, logits_processors, logits_rows, token_ids) -> None:
+        # logits_rows: the request's rows as [beam_width, 1, vocab], a view
+        # into the step's logits tensor (see _execute_logit_post_processors).
         for lp in logits_processors:
             lp_params = inspect.signature(lp).parameters
 
@@ -829,6 +831,49 @@ def _apply_logits_processors(
         # back would be a self-assignment, which torch rejects for the
         # non-contiguous slices a TP-padded vocab produces.
 
+    def _applies_batched_logits_processor(self, request) -> bool:
+        if not getattr(request, "py_apply_batched_logits_processor", False):
+            return False
+        if self.batched_logits_processor is None:
+            # LLM refuses the flag without a processor, so only a direct
+            # caller gets here. Raising would fail the whole step for every
+            # request in it; say so once and leave the request unprocessed.
+            logger.warning_once(
+                "a request sets apply_batched_logits_processor, but no batched "
+                "logits processor is registered with the engine; it generates "
+                "without one",
+                key="batched_logits_processor_unregistered",
+            )
+            return False
+        return True
+
+    def _has_logits_post_processors(self, request) -> bool:
+        """Whether the walk hands this request's logits to any processor."""
+        return bool(
+            getattr(request, "py_logits_post_processors", None)
+        ) or self._applies_batched_logits_processor(request)
+
+    def _apply_batched_logits_processor(
+        self,
+        req_ids: list[int],
+        logits: list[torch.Tensor],
+        token_ids: list[list[list[int]]],
+        client_ids: list[int | None],
+    ) -> None:
+        """One call for every request the walk collected, in scheduling order.
+
+        Each logits entry is the request's rows viewed as [beam_width, 1,
+        vocab], the shape its own processors see, and a view into the step's
+        logits tensor, so the callback's in-place edits land there. The stream
+        is the one the logits live on (None off CUDA, which only tests do),
+        for processors written against the TensorRT backend's contract.
+        """
+        device = logits[0].device
+        stream_ptr = (
+            torch.cuda.current_stream(device).cuda_stream if device.type == "cuda" else None
+        )
+        self.batched_logits_processor(req_ids, logits, token_ids, stream_ptr, client_ids)
+
     def _start_beam_histories(
         self, scheduled_requests: ScheduledRequests
     ) -> tuple[list[LlmRequest], PendingBeamHistories] | None:
@@ -847,7 +892,7 @@ def _start_beam_histories(
         beam_requests = [
             request
             for request in scheduled_requests.generation_requests
-            if getattr(request, "py_logits_post_processors", None)
+            if self._has_logits_post_processors(request)
             and request.get_beam_width_by_iter(for_next_iteration=False) > 1
         ]
         if not beam_requests:
@@ -862,6 +907,11 @@ def _execute_logit_post_processors(
     ) -> None:
         """Apply logit post processors (in-place modify outputs Tensors) if any.
 
+        A request's own processors run as the walk reaches it. The requests
+        that asked for the LLM's batched processor are collected on the way
+        and handed to it in one call once the walk is done, so it sees the
+        per-request edits and every request of the step at once.
+
         pending_beam_histories is what _start_beam_histories returned for this
         step.
         """
@@ -887,15 +937,18 @@ def _execute_logit_post_processors(
                     "not"
                 )
             if not any(
-                getattr(request, "py_logits_post_processors", None)
+                self._has_logits_post_processors(request)
                 for request in scheduled_requests.all_requests()
             ):
                 return
             from tensorrt_llm._torch.distributed import allgather
 
             logits_tensor = allgather(logits_tensor, self._mapping, dim=-1).float()
             # Drop the vocabulary's TP padding, as the LM head's own gather
-            # does.
+            # does. The processors (per-request and batched alike) edit this
+            # gathered copy only: the worker's acceptance already ran inside
+            # the forward (see SpecWorkerBase.target_logits), so the shard
+            # the sampler reads is left as it is.
             vocab_size = getattr(getattr(self._model, "lm_head", None), "num_embeddings", None)
             if vocab_size is not None:
                 logits_tensor = logits_tensor[..., :vocab_size]
@@ -908,6 +961,10 @@ def _execute_logit_post_processors(
             }
 
         logits_row_offset = 0
+        batched_req_ids: list[int] = []
+        batched_logits: list[torch.Tensor] = []
+        batched_token_ids: list[list[list[int]]] = []
+        batched_client_ids: list[int | None] = []
         request_groups = (
             (scheduled_requests.context_requests, True),
             (scheduled_requests.generation_requests, False),
@@ -929,7 +986,8 @@ def _execute_logit_post_processors(
                     row_stride = request.py_beam_width
 
                 logits_processors = getattr(request, "py_logits_post_processors", None)
-                if logits_processors:
+                apply_batched = self._applies_batched_logits_processor(request)
+                if logits_processors or apply_batched:
                     if is_context_request:
                         token_ids = [request.get_tokens(0)]
                     elif (histories := beam_histories.get(id(request))) is not None:
@@ -942,16 +1000,28 @@ def _execute_logit_post_processors(
                         logits_row_offset += row_stride
                         continue
 
-                    self._apply_logits_processors(
-                        request,
-                        logits_processors,
-                        logits_tensor,
-                        beam_width,
-                        token_ids,
-                        logits_row_offset,
-                    )
+                    # The request's rows, reshaped to the TRT backend's
+                    # [beam_width, 1, vocab] so the same processors serve both
+                    # backends; a view, so edits land in the step's logits.
+                    logits_rows = logits_tensor[
+                        logits_row_offset : logits_row_offset + beam_width
+                    ].view(beam_width, 1, -1)
+                    if logits_processors:
+                        self._apply_logits_processors(
+                            request, logits_processors, logits_rows, token_ids
+                        )
+                    if apply_batched:
+                        batched_req_ids.append(request.py_request_id)
+                        batched_logits.append(logits_rows)
+                        batched_token_ids.append(token_ids)
+                        batched_client_ids.append(request.py_client_id)
                 logits_row_offset += row_stride
 
+        if batched_req_ids:
+            self._apply_batched_logits_processor(
+                batched_req_ids, batched_logits, batched_token_ids, batched_client_ids
+            )
+
     def _cuda_graph_metadata_prepare_scope(
         self,
         attn_metadata: AttentionMetadata,

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -719,6 +719,10 @@ def __init__(
             **kwargs):
         self.py_logits_post_processors = kwargs.pop("py_logits_post_processors",
                                                     None)
+        # SamplingParams.apply_batched_logits_processor: the engine hands this
+        # request's logits to LLM(batched_logits_processor=...) each step.
+        self.py_apply_batched_logits_processor: bool = kwargs.pop(
+            "py_apply_batched_logits_processor", False)
         # Owned entirely on the Python side: the TorchSampler is the only
         # consumer, so the C++ request no longer carries a copy.
         self.py_embedding_bias: Optional[torch.Tensor] = kwargs.pop(
@@ -1655,6 +1659,8 @@ def executor_request_to_llm_request(
         guided_decoding_params=executor_request.guided_decoding_params,
         py_logits_post_processors=getattr(executor_request,
                                           "py_logits_post_processors", None),
+        py_apply_batched_logits_processor=getattr(
+            executor_request, "py_apply_batched_logits_processor", False),
         encoder_input_tokens=executor_request.encoder_input_token_ids,
         encoder_input_features=encoder_input_features,
         encoder_output_len=encoder_output_len,

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -812,6 +812,15 @@ def register_beam_histories_callable(self, callable: Optional[Callable]):
         if isinstance(self._runner, DecoderRunner):
             self._runner.start_beam_histories_callable = callable
 
+    def register_batched_logits_processor(self, processor: Optional[Callable]):
+        """Register LLM(batched_logits_processor=...) with the post-processor walk."""
+        if not isinstance(self._runner, DecoderRunner):
+            raise ValueError(
+                "LLM(batched_logits_processor=...) is not supported by "
+                f"{type(self._runner).__name__}: it has no logits "
+                "post-processor walk to call it from.")
+        self._runner.batched_logits_processor = processor
+
     def register_sample_type_resolver(self,
                                       resolver: Optional[Callable],
                                       stage: Optional[Callable] = None):

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -605,6 +605,20 @@ def __init__(
         self.runtime_draft_len = (get_static_draft_len(model_engine.spec_config)
                                   if model_engine.is_spec_decode else 0)
         self.llm_args = self.model_engine.llm_args
+        # LLM(batched_logits_processor=...) is called once per step by the
+        # engine's logits post-processor walk, for the requests that opted in
+        # with SamplingParams.apply_batched_logits_processor.
+        batched_logits_processor = self.llm_args.batched_logits_processor
+        if batched_logits_processor is not None:
+            # Fail at construction rather than serve requests that asked for
+            # processing nothing would run.
+            if not hasattr(model_engine, "register_batched_logits_processor"):
+                raise ValueError(
+                    "LLM(batched_logits_processor=...) is not supported by "
+                    f"{type(model_engine).__name__}: it has no logits "
+                    "post-processor walk to call it from.")
+            model_engine.register_batched_logits_processor(
+                batched_logits_processor)
         self.max_stats_len = self.llm_args.max_stats_len
         self.max_num_tokens = self.llm_args.max_num_tokens
         self.print_log = self.llm_args.print_iter_log

tensorrt_llm/_torch/pyexecutor/request_utils.py

@@ -518,6 +518,9 @@ def _collect_py_objects(self, new_requests: List) -> Tuple:
         py_logits_post_processors = collect_py_objects_from_requests(
             new_requests, "py_logits_post_processors"
         )
+        py_apply_batched_logits_processor = collect_py_objects_from_requests(
+            new_requests, "py_apply_batched_logits_processor"
+        )
         py_multimodal_data = collect_py_objects_from_requests(new_requests, "py_multimodal_data")
         py_scheduling_params = collect_py_objects_from_requests(
             new_requests, "py_scheduling_params"
@@ -537,6 +540,7 @@ def _collect_py_objects(self, new_requests: List) -> Tuple:
                 None,
                 [
                     py_logits_post_processors,
+                    py_apply_batched_logits_processor,
                     py_multimodal_data,
                     py_scheduling_params,
                     py_num_logprobs,

tensorrt_llm/_torch/pyexecutor/sampler/sampler.py

@@ -816,14 +816,17 @@ def _request_bypasses_in_graph_sampling(request: LlmRequest) -> bool:
 
         The hook runs at the tail of the forward, before anything the executor
         does to the logits afterwards: the guided-decoding bitmask applied by
-        ``guided_decoder.execute()``, and the user's own logits processors run
-        by ``_execute_logit_post_processors``. Both happen once
+        ``guided_decoder.execute()``, and the user's own logits processors
+        (per request, or the LLM's batched one) run by
+        ``_execute_logit_post_processors``. Both happen once
         ``_forward_step`` has returned, so sampling in-graph would read unmasked
         logits and emit a token they were meant to forbid, while their output is
         discarded. Any future post-forward logits mutator belongs here too.
         """
-        return request.py_guided_decoding_params is not None or bool(
-            getattr(request, "py_logits_post_processors", None)
+        return (
+            request.py_guided_decoding_params is not None
+            or bool(getattr(request, "py_logits_post_processors", None))
+            or bool(getattr(request, "py_apply_batched_logits_processor", False))
         )
 
     def get_sample_type(self, requests: list[LlmRequest]) -> SampleType:

tensorrt_llm/executor/base_worker.py

@@ -508,6 +508,10 @@ def _deduce_max_tokens(request: GenerationRequest) -> int:
                 lp = request.sampling_params.logits_processor
                 executor_request.py_logits_post_processors = lp if isinstance(
                     lp, list) else [lp]
+            if request.sampling_params.apply_batched_logits_processor:
+                # Only set when asked for, so the PP broadcast of Python-only
+                # request attributes has nothing to carry for other requests.
+                executor_request.py_apply_batched_logits_processor = True
 
             executor_request.py_scheduling_params = None
             if request.scheduling_params is not None:

tensorrt_llm/llmapi/llm.py

@@ -1605,6 +1605,13 @@ def _check_arguments(self, prompt_len: int, sampling_params: SamplingParams,
                     f"max_num_tokens ({max_num_tokens})")
         self._check_one_model_speculative_sampling(sampling_params)
 
+        if (sampling_params.apply_batched_logits_processor
+                and self.args.batched_logits_processor is None):
+            raise ValueError(
+                "`sampling_params.apply_batched_logits_processor=True` requires "
+                "`LLM(batched_logits_processor=...)`: there is no batched logits "
+                "processor to apply.")
+
         if sampling_params.return_routed_experts and not (
                 self.args.backend == "pytorch"
                 and getattr(self.args, "enable_return_routed_experts", False)):

tensorrt_llm/sampling_params.py

@@ -164,7 +164,7 @@ def __call__(
         req_ids: List[int],
         logits: List[torch.Tensor],
         token_ids: List[List[List[int]]],
-        stream_ptr: int,
+        stream_ptr: Optional[int],
         client_ids: List[Optional[int]],
     ) -> None:
         """Batched logits processing callback. The callback is expected to inplace modify the logits.
@@ -173,8 +173,11 @@ def __call__(
             req_ids (List[int]): A batch of request ids.
             logits (List[torch.Tensor]): A batch of the logits tensors.
             token_ids (List[List[List[int]]]): A batch of the token ids produced by the requests so far.
-                The shape is batch * beam_width * sequence_length.
-            stream_ptr (int): The operation stream used by the logits tensors.
+                The shape is batch * beam_width * sequence_length. Each entry is what the
+                request's own ``LogitsProcessor`` would receive (see its ``token_ids``).
+            stream_ptr (int, optional): The operation stream used by the logits tensors.
+                On the PyTorch backend this is the CUDA stream the logits live on,
+                or None when they are not on a CUDA device.
             client_ids (List[Optional[int]]): A batch of optional client ids.
         """
         pass  # noqa
@@ -207,6 +210,8 @@ class SamplingParams:
             If a list, each processor is applied in order during generation (supported in PyTorch backend only).
         apply_batched_logits_processor (bool): Whether to apply batched logits postprocessor callback. Defaults to False.
             The BatchedLogitsProcessor class is recommended for callback creation. The callback must be provided when initializing LLM.
+            The callback is called once per decoding step with every scheduled request that sets this, after each
+            request's own logits_processor(s).
 
         n (int): Number of sequences to generate. Defaults to 1.
         best_of (int, optional): Number of sequences to consider for best output. Defaults to None.