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

docs/source/features/sampling.md

@@ -410,3 +410,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: 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. It runs on every tensor-parallel rank of the last pipeline stage, and each rank samples from its own copy of the edits: a processor with side effects runs `tp_size` times per step, and a nondeterministic one can make the ranks sample different tokens, so keep it deterministic and side-effect free. With attention data parallelism, each rank owns its own requests, so each call sees only that rank's share of the step's requests. The flag is refused with speculative decoding on the PyTorch backend: one-model workers sample inside the forward, before the processor could run, and the drafter-based modes sample more than one row per request, which the processor would not see.

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

@@ -256,6 +256,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,
@@ -971,13 +976,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
 
@@ -1025,6 +1026,46 @@ def _logits_for_greedy_tokens(
             self.sample_in_graph_callable(outputs)
         return outputs
 
+    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, for processors that synchronize on it as
+        the interface documents.
+        """
+        stream_ptr = torch.cuda.current_stream(logits[0].device).cuda_stream
+        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:
@@ -1043,7 +1084,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:
@@ -1058,6 +1099,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.
         """
@@ -1088,15 +1134,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]
@@ -1109,6 +1158,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),
@@ -1130,16 +1183,17 @@ 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 (
                         not is_context_request
                         and getattr(request, "py_verify_len", None) is not None
                     ):
                         raise RuntimeError(
-                            "per-request logits post-processors are not supported with DSpark "
-                            "ragged verification: the processor interface assumes one uniform "
-                            "row stride and would modify another request's packed logits after "
-                            "a short verify window"
+                            "logits post-processors (per-request or batched) are not supported "
+                            "with DSpark ragged verification: the processor interface assumes one "
+                            "uniform row stride and would modify another request's packed logits "
+                            "after a short verify window"
                         )
                     if is_context_request:
                         token_ids = [request.get_tokens(0)]
@@ -1148,21 +1202,39 @@ def _execute_logit_post_processors(
                         token_ids = [prompt + history for history in histories]
                     else:
                         token_ids = [request.get_tokens(beam_idx) for beam_idx in range(beam_width)]
-                    if is_context_request and request.py_orig_prompt_len < len(token_ids[0]):
-                        # Skip as we only need to apply logit processor on the last context request
+                    if is_context_request and (
+                        not request.is_last_context_chunk
+                        or request.py_orig_prompt_len < len(token_ids[0])
+                    ):
+                        # Only the last context chunk's row is sampled from; the
+                        # sampler drops the row of an earlier chunk. get_tokens(0)
+                        # is already the whole prompt during chunked prefill, so
+                        # the length check alone does not tell the chunks apart.
                         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 as [beam_width, 1, vocab], the shape
+                    # the LogitsProcessor interface documents; 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

@@ -835,6 +835,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(
@@ -1384,6 +1388,8 @@ def create_child_request(self, child_id: int) -> None:
             py_logits_post_processors=(self.py_logits_post_processors.copy()
                                        if self.py_logits_post_processors
                                        is not None else None),
+            py_apply_batched_logits_processor=self.
+            py_apply_batched_logits_processor,
             py_multimodal_data=deepcopy(self.py_multimodal_data),
             encoder_input_features=self.py_encoder_input_features,
         )
@@ -1847,6 +1853,7 @@ def executor_request_to_llm_request(
         exclude_input_from_output,
         guided_decoding_params=executor_request.guided_decoding_params,
         py_logits_post_processors=inputs.logits_post_processors,
+        py_apply_batched_logits_processor=inputs.apply_batched_logits_processor,
         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

@@ -723,6 +723,16 @@ def register_beam_histories_callable(self, callable: Callable | None) -> None:
         if isinstance(self._runner, DecoderRunner):
             self._runner.start_beam_histories_callable = callable
 
+    def register_batched_logits_processor(self, processor: Callable | None) -> None:
+        """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: Callable | None, stage: Callable | None = None
     ) -> None:

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -672,6 +672,25 @@ class _ExecutorExecutionConfig:
     transfer_budget_is_explicit: bool = False
 
 
+def _register_batched_logits_processor(llm_args, model_engine) -> None:
+    """Hand LLM(batched_logits_processor=...) to the model engine.
+
+    The engine's logits post-processor walk calls it once per step, for the
+    requests that opted in with SamplingParams.apply_batched_logits_processor.
+    Raises at construction when the engine has no such walk, rather than serve
+    requests that asked for processing nothing would run.
+    """
+    batched_logits_processor = llm_args.batched_logits_processor
+    if batched_logits_processor is None:
+        return
+    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)
+
+
 class PyExecutor:
     # True while the overlap loop runs with the execution stream current
     # (_overlap_loop_stream_scope).
@@ -965,6 +984,7 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None:
         self._decode_only = is_decode_only(
             getattr(self.llm_args, "sparse_attention_config", None))
         self._kv_connector_config = self.llm_args.kv_connector_config
+        _register_batched_logits_processor(self.llm_args, model_engine)
         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_inputs.py

@@ -24,6 +24,7 @@ class PythonRequestInputs:
     """
 
     logits_post_processors: list[LogitsProcessor] | None = None
+    apply_batched_logits_processor: bool = False
     multimodal_data: dict[str, Any] | None = None
     scheduling_params: SchedulingParams | None = None
     num_logprobs: int | None = 0
@@ -52,6 +53,9 @@ def from_native_request(cls, request: Request) -> "PythonRequestInputs":
         """
         return cls(
             logits_post_processors=getattr(request, "py_logits_post_processors", None),
+            apply_batched_logits_processor=getattr(
+                request, "py_apply_batched_logits_processor", False
+            ),
             multimodal_data=getattr(request, "py_multimodal_data", None),
             scheduling_params=getattr(request, "py_scheduling_params", None),
             num_logprobs=getattr(request, "py_num_logprobs", 0),

tensorrt_llm/_torch/pyexecutor/sampler/sampler.py

@@ -147,6 +147,7 @@ def greedy_tokens_usable(request: LlmRequest) -> bool:
         and not request.py_return_log_probs
         and not request.py_return_generation_logits
         and not request.py_logits_post_processors
+        and not request.py_apply_batched_logits_processor
         and request.py_guided_decoding_params is None
         and not check_stop_words_length(request)
         and request_row_params(request).mode is SamplingMode.GREEDY
@@ -929,14 +930,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

@@ -583,6 +583,8 @@ def _deduce_max_tokens(request: GenerationRequest) -> int:
                     kv_hints=request.kv_hints,
                     arrival_time=request.arrival_time,
                     watermarking=request.sampling_params.watermarking,
+                    apply_batched_logits_processor=request.sampling_params.
+                    apply_batched_logits_processor,
                 ),
             )