From c45455218037b05cbab177b910912fd19457af5b Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Mon, 26 Oct 2026 09:53:12 -0700 Subject: [PATCH 1/8] [None][feat] Call LLM(batched_logits_processor=...) on the PyTorch backend LLM accepts batched_logits_processor and SamplingParams has apply_batched_logits_processor, but nothing on the PyTorch backend read either: the batched processor was never called, and the flag was dropped when the executor request was built. The engine's logits post-processor walk now collects the logits rows, token ids, request ids and client ids of every scheduled request that set the flag and hands them to the LLM's batched processor in one call per step, after each request's own processors have run. The rows are views into the step's logits tensor, viewed as [beam_width, 1, vocab] like the per-request processors see them, so in-place edits land; under beam search they are the live beams and their histories. The flag travels with the request across pipeline stages like the per-request processors do, keeps the request on the eager sampling tier, and is refused at submission when the LLM has no batched processor. Signed-off-by: Luke Miner --- docs/source/features/sampling.md | 31 +++ .../engine/runners/decoder/runner.py | 115 ++++++-- tensorrt_llm/_torch/pyexecutor/llm_request.py | 5 + .../_torch/pyexecutor/model_engine.py | 10 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 14 + .../_torch/pyexecutor/request_inputs.py | 4 + .../_torch/pyexecutor/sampler/sampler.py | 11 +- tensorrt_llm/executor/base_worker.py | 2 + tensorrt_llm/llmapi/llm.py | 7 + tensorrt_llm/sampling_params.py | 11 +- .../test_batched_logits_processor.py | 248 ++++++++++++++++++ .../_torch/pyexecutor/test_request_utils.py | 15 ++ .../_torch/sampler/test_in_graph_sampling.py | 4 + .../references/batched_logits_processor.yaml | 2 +- tests/unittest/llmapi/test_llm.py | 43 ++- tests/unittest/llmapi/test_llm_pytorch.py | 7 + 16 files changed, 497 insertions(+), 32 deletions(-) create mode 100644 tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index abbc61587709..a0f13029902b 100644 --- a/docs/source/features/sampling.md +++ b/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: 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. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 7419e578f422..1fa077b5b32a 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/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,49 @@ 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 (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: @@ -1043,7 +1087,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 +1102,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,7 +1137,7 @@ 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 @@ -1096,7 +1145,10 @@ def _execute_logit_post_processors( 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 +1161,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 +1186,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)] @@ -1153,16 +1210,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, diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 77a41e8277b3..b3ac9cb74dcc 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/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( @@ -1847,6 +1851,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, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 389f3c8f0242..bbc40cf1eaa1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6f93a575dfd8..e4a5942ca92a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -965,6 +965,20 @@ 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 + # 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 diff --git a/tensorrt_llm/_torch/pyexecutor/request_inputs.py b/tensorrt_llm/_torch/pyexecutor/request_inputs.py index 9cccefe3b8df..843fa1614805 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_inputs.py +++ b/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), diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 027a0803fc07..c0618db38761 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -929,14 +929,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: diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 2fcf106677c3..930a0a44a803 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/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, ), ) diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index bedc8eb0d251..0ffc0fa2bb9d 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1642,6 +1642,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)): diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 3353a0043ee1..ca6b5755b88a 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -166,7 +166,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. @@ -175,8 +175,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 @@ -209,6 +212,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. diff --git a/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py new file mode 100644 index 000000000000..fdd8b5d9c26d --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The engine's logits post-processor walk with ``LLM(batched_logits_processor=...)``. + +Requests that set ``SamplingParams.apply_batched_logits_processor`` are collected +as the walk visits them and handed to the batched processor in one call, after +their own processors, with views into the step's logits tensor. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.distributed as distributed +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm.llmapi.llm import BaseLLM +from tensorrt_llm.sampling_params import SamplingParams + +pytestmark = pytest.mark.cpu_only + +VOCAB = 8 + + +class _Recorder: + """Stands in for a BatchedLogitsProcessor: records each call and adds ``mark`` to every row.""" + + def __init__(self, mark=1.0): + self.calls = [] + self.mark = mark + + def __call__(self, req_ids, logits, token_ids, stream_ptr, client_ids): + self.calls.append((req_ids, logits, token_ids, stream_ptr, client_ids)) + for rows in logits: + rows += self.mark + + +def _request( + req_id, + *, + batched=False, + processors=None, + beam_width=1, + admitted_beam_width=None, + tokens=(1, 2, 3), + orig_prompt_len=None, +): + tokens = list(tokens) + return SimpleNamespace( + py_request_id=req_id, + py_client_id=100 + req_id, + py_logits_post_processors=processors, + py_apply_batched_logits_processor=batched, + py_beam_width=admitted_beam_width or beam_width, + py_prompt_len=len(tokens), + py_orig_prompt_len=len(tokens) if orig_prompt_len is None else orig_prompt_len, + get_beam_width_by_iter=lambda for_next_iteration: beam_width, + get_tokens=lambda beam: list(tokens), + ) + + +def _engine(processor, *, attention_dp=False, vocab_size=None): + engine = object.__new__(DecoderRunner) + engine._mapping = SimpleNamespace( + is_last_pp_rank=lambda: True, enable_attention_dp=attention_dp + ) + engine._model = SimpleNamespace(lm_head=SimpleNamespace(num_embeddings=vocab_size)) + engine.batched_logits_processor = processor + engine.start_beam_histories_callable = None + return engine + + +def _scheduled(context=(), generation=()): + return SimpleNamespace( + context_requests=list(context), + generation_requests=list(generation), + all_requests=lambda: [*context, *generation], + ) + + +def _logits(rows): + # Row r holds the value r everywhere, so edits are attributable to a row. + return torch.arange(rows, dtype=torch.float32)[:, None].repeat(1, VOCAB) + + +def test_flagged_requests_share_one_call_in_scheduling_order(): + processor = _Recorder() + context = _request(1, batched=True, tokens=(4, 5)) + plain = _request(2) + generation = _request(3, batched=True, tokens=(6, 7, 8)) + logits = _logits(3) + + _engine(processor)._execute_logit_post_processors( + _scheduled(context=[context], generation=[plain, generation]), {"logits": logits} + ) + + assert len(processor.calls) == 1 + req_ids, rows, token_ids, stream_ptr, client_ids = processor.calls[0] + assert req_ids == [1, 3] + assert client_ids == [101, 103] + assert token_ids == [[[4, 5]], [[6, 7, 8]]] + assert [tuple(r.shape) for r in rows] == [(1, 1, VOCAB), (1, 1, VOCAB)] + assert stream_ptr is None # the logits are not on a CUDA device + # The callback's in-place edits land in the step's logits; the other request's row is untouched. + assert torch.equal(logits[:, 0], torch.tensor([1.0, 1.0, 3.0])) + + +def test_per_request_processors_run_before_the_batched_one(): + order = [] + + def own(req_id, rows, token_ids, stream_ptr, client_id): + order.append(("own", req_id, rows[0, 0, 0].item())) + rows.fill_(5.0) + + batched = _Recorder() + request = _request(1, batched=True, processors=[own]) + logits = _logits(1) + + _engine(batched)._execute_logit_post_processors( + _scheduled(generation=[request]), {"logits": logits} + ) + + assert order == [("own", 1, 0.0)] + assert batched.calls[0][0] == [1] + assert torch.equal(logits[0], torch.full((VOCAB,), 6.0)) # 5 from `own`, then +1 + + +def test_beam_rows_and_histories_reach_the_batched_processor(): + processor = _Recorder() + first = _request(1, batched=True, beam_width=2, admitted_beam_width=4, tokens=(1, 2)) + second = _request(2, batched=True, beam_width=1, admitted_beam_width=4, tokens=(3,)) + logits = _logits(8) + histories = SimpleNamespace(wait=lambda: [[[7], [8]]]) + + _engine(processor)._execute_logit_post_processors( + _scheduled(generation=[first, second]), + {"logits": logits}, + pending_beam_histories=([first], histories), + ) + + req_ids, rows, token_ids, _, _ = processor.calls[0] + assert req_ids == [1, 2] + # Live beams only, and the prompt plus each beam's own hypothesis. + assert [tuple(r.shape) for r in rows] == [(2, 1, VOCAB), (1, 1, VOCAB)] + assert token_ids == [[[1, 2, 7], [1, 2, 8]], [[3]]] + # The second request's rows start at the admission width, not the live one. + assert torch.equal(logits[:, 0], torch.tensor([1.0, 2.0, 2.0, 3.0, 5.0, 5.0, 6.0, 7.0])) + + +def test_an_unfinished_context_chunk_is_not_handed_over(): + processor = _Recorder() + chunk = _request(1, batched=True, tokens=(1, 2, 3), orig_prompt_len=2) + + _engine(processor)._execute_logit_post_processors( + _scheduled(context=[chunk]), {"logits": _logits(1)} + ) + + assert processor.calls == [] + + +def test_ragged_verification_is_refused_for_batched_only_requests(): + processor = _Recorder() + request = _request(1, batched=True) + request.py_verify_len = 2 + + with pytest.raises(RuntimeError, match="DSpark ragged verification"): + _engine(processor)._execute_logit_post_processors( + _scheduled(generation=[request]), {"logits": _logits(4)} + ) + + assert processor.calls == [] + + +def test_without_opted_in_requests_the_processor_is_not_called(): + processor = _Recorder() + request = _request(1, processors=[lambda req_id, rows, tokens, stream_ptr, client_id: None]) + + _engine(processor)._execute_logit_post_processors( + _scheduled(generation=[request]), {"logits": _logits(1)} + ) + + assert processor.calls == [] + + +def test_the_flag_without_a_registered_processor_is_skipped_with_a_warning(monkeypatch): + # LLM refuses it earlier; a direct caller gets a warning, not a failed step. + import tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner as runner_module + + warnings = [] + monkeypatch.setattr( + runner_module.logger, "warning_once", lambda msg, key=None: warnings.append(key) + ) + logits = _logits(1) + + _engine(None)._execute_logit_post_processors( + _scheduled(generation=[_request(1, batched=True)]), {"logits": logits} + ) + + assert torch.equal(logits, _logits(1)) + assert warnings == ["batched_logits_processor_unregistered"] + + +def test_beam_histories_start_for_batched_only_requests(): + started = [] + engine = _engine(_Recorder()) + engine.start_beam_histories_callable = lambda requests: started.append(requests) or "pending" + beam = _request(1, batched=True, beam_width=2) + single = _request(2, batched=True, beam_width=1) + plain = _request(3, beam_width=2) + + result = engine._start_beam_histories(_scheduled(generation=[beam, single, plain])) + + assert result == ([beam], "pending") + assert started == [[beam]] + + +def test_sharded_logits_are_gathered_for_the_batched_processor(monkeypatch): + calls = [] + + def allgather(tensor, mapping, dim=-1): + calls.append(tensor) + return torch.zeros(tensor.shape[0], 32) + + monkeypatch.setattr(distributed, "allgather", allgather) + processor = _Recorder() + shard = torch.ones(1, 8) + + _engine(processor)._execute_logit_post_processors( + _scheduled(generation=[_request(1, batched=True)]), + {"logits": shard, "logits_vocab_shard": True}, + ) + + assert calls == [shard] + assert [tuple(r.shape) for r in processor.calls[0][1]] == [(1, 1, 32)] + + +def test_the_flag_is_refused_without_a_batched_processor_on_the_llm(): + llm = SimpleNamespace( + args=SimpleNamespace(batched_logits_processor=None, enable_chunked_prefill=False), + _check_one_model_speculative_sampling=lambda params: None, + ) + params = SamplingParams(apply_batched_logits_processor=True) + + with pytest.raises(ValueError, match="batched_logits_processor"): + BaseLLM._check_arguments(llm, 3, params, is_gen_only=True) + + llm.args.batched_logits_processor = _Recorder() + BaseLLM._check_arguments(llm, 3, params, is_gen_only=True) diff --git a/tests/unittest/_torch/pyexecutor/test_request_utils.py b/tests/unittest/_torch/pyexecutor/test_request_utils.py index ebb9da1efb47..00c21fd86294 100644 --- a/tests/unittest/_torch/pyexecutor/test_request_utils.py +++ b/tests/unittest/_torch/pyexecutor/test_request_utils.py @@ -1014,6 +1014,21 @@ def test_single_round_busy_iteration_sends_requests_and_py_objects_once(): assert payload == items +def test_py_objects_carry_the_batched_logits_processor_flag(): + # The flag is only set on requests that opted in, so only those cross the + # pipeline stages with it. + broadcaster = RequestBroadcaster(_SingleRoundDist(rank=0), MagicMock()) + items = [_make_broadcast_item(1), _make_broadcast_item(2)] + items[1].request.py_apply_batched_logits_processor = True + + py_request_objects = broadcaster._collect_py_objects(items) + + assert ("py_apply_batched_logits_processor", {2: True}) in py_request_objects + attach_py_objects_to_requests(items, py_request_objects) + assert not hasattr(items[0].request, "py_apply_batched_logits_processor") + assert items[1].request.py_apply_batched_logits_processor is True + + def test_single_round_non_root_receives_payload_or_nothing(): items = [_make_broadcast_item(7)] dist = _SingleRoundDist(rank=1, received=items) diff --git a/tests/unittest/_torch/sampler/test_in_graph_sampling.py b/tests/unittest/_torch/sampler/test_in_graph_sampling.py index f7afb180398c..811ad27ee4f3 100644 --- a/tests/unittest/_torch/sampler/test_in_graph_sampling.py +++ b/tests/unittest/_torch/sampler/test_in_graph_sampling.py @@ -69,6 +69,7 @@ class _Request: py_bad_words: Any = None py_no_repeat_ngram_size: Any = None py_logits_post_processors: Any = None + py_apply_batched_logits_processor: bool = False py_guided_decoding_params: Any = None py_embedding_bias: Any = None @@ -146,6 +147,9 @@ def test_greedy_is_full(self, sampler): # Consume the logits after the forward returns. pytest.param(dict(py_guided_decoding_params=object()), id="guided_decoding"), pytest.param(dict(py_logits_post_processors=[object()]), id="logits_processor"), + pytest.param( + dict(py_apply_batched_logits_processor=True), id="batched_logits_processor" + ), # Needs outputs the in-graph path does not produce. pytest.param(dict(py_return_log_probs=True), id="logprobs"), ], diff --git a/tests/unittest/api_stability/references/batched_logits_processor.yaml b/tests/unittest/api_stability/references/batched_logits_processor.yaml index 423532e31ca5..bdec874334a4 100644 --- a/tests/unittest/api_stability/references/batched_logits_processor.yaml +++ b/tests/unittest/api_stability/references/batched_logits_processor.yaml @@ -11,7 +11,7 @@ methods: annotation: List[int] default: inspect._empty stream_ptr: - annotation: int + annotation: Optional[int] default: inspect._empty token_ids: annotation: List[List[List[int]]] diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 56b016540fe9..5281d213d100 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -38,7 +38,8 @@ from tensorrt_llm.llmapi.llm_utils import _ParallelConfig from tensorrt_llm.llmapi.tokenizer import (TokenizerBase, TransformersTokenizer, load_hf_tokenizer) -from tensorrt_llm.sampling_params import LogitsProcessor, SamplingParams +from tensorrt_llm.sampling_params import (BatchedLogitsProcessor, + LogitsProcessor, SamplingParams) from tensorrt_llm.serve.openai_protocol import CompletionRequest, StreamOptions from tensorrt_llm.serve.openai_server import OpenAIServer from tensorrt_llm.serve.postprocess_handlers import (ChatPostprocArgs, @@ -707,6 +708,46 @@ def tinyllama_logits_processor_test_harness(backend=None, **llm_kwargs): **llm_kwargs) +class MyBatchedLogitsProcessor(BatchedLogitsProcessor): + + def __init__(self, biased_word_id): + self.biased_word_id = biased_word_id + + def __call__(self, req_ids: List[int], logits: List[torch.Tensor], + ids: List[List[List[int]]], stream_ptr: Optional[int], + client_ids: List[Optional[int]]): + assert len(req_ids) == len(logits) == len(ids) == len(client_ids) + stream = None if stream_ptr is None else torch.cuda.ExternalStream( + stream_ptr) + with torch.cuda.stream(stream): + for rows in logits: + rows[:] = float("-inf") + rows[..., self.biased_word_id] = 0 + + +def tinyllama_batched_logits_processor_test_harness(backend=None, **llm_kwargs): + tokenizer = TransformersTokenizer.from_pretrained(llama_model_path) + biased_word_id = tokenizer.encode("Z", add_special_tokens=False)[-1] + max_tokens = 6 + expected_output = tokenizer.decode([biased_word_id] * max_tokens) + + with assert_resource_freed( + LLM, + llama_model_path, + batched_logits_processor=MyBatchedLogitsProcessor(biased_word_id), + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.4), + backend=backend, + **llm_kwargs) as llm: + # Only the request that opts in is handed to the batched processor. + outputs = llm.generate(["A B C", "A B C"], [ + SamplingParams(max_tokens=max_tokens, + apply_batched_logits_processor=True), + SamplingParams(max_tokens=max_tokens), + ]) + assert outputs[0].outputs[0].text == expected_output + assert outputs[1].outputs[0].text != expected_output + + def test_executor_results_cleanup(): llm = LLM(model=llama_model_path, kv_cache_config=global_kvcache_config) sampling_params = SamplingParams(max_tokens=6) diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index e463c633f69f..4b9e4534f646 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -34,6 +34,7 @@ prompts, run_llm_abort_request, sampling_params_for_aborting_request, run_llm_with_postprocess_parallel_and_result_handler, + tinyllama_batched_logits_processor_test_harness, tinyllama_logits_processor_test_harness) from utils.util import (force_ampere, similar, skip_gpu_memory_less_than_40gb, skip_gpu_memory_less_than_138gb, skip_ray) @@ -73,6 +74,12 @@ def test_tinyllama_logits_processor(enable_chunked_prefill): backend="pytorch", enable_chunked_prefill=enable_chunked_prefill) +@force_ampere +@pytest.mark.part2 +def test_tinyllama_batched_logits_processor(): + tinyllama_batched_logits_processor_test_harness(backend="pytorch") + + @skip_ray @pytest.mark.parametrize( "return_context_logits, use_overlap, enable_chunked_prefill, enable_iter_req_stats", From 7c7d97ab2ef5c56305108a9fab457e569797700c Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Mon, 26 Oct 2026 15:17:27 -0700 Subject: [PATCH 2/8] [None][test] Attach broadcast py objects onto fresh request items Signed-off-by: Luke Miner --- tests/unittest/_torch/pyexecutor/test_request_utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/pyexecutor/test_request_utils.py b/tests/unittest/_torch/pyexecutor/test_request_utils.py index 00c21fd86294..56b14e5d9af0 100644 --- a/tests/unittest/_torch/pyexecutor/test_request_utils.py +++ b/tests/unittest/_torch/pyexecutor/test_request_utils.py @@ -1024,9 +1024,12 @@ def test_py_objects_carry_the_batched_logits_processor_flag(): py_request_objects = broadcaster._collect_py_objects(items) assert ("py_apply_batched_logits_processor", {2: True}) in py_request_objects - attach_py_objects_to_requests(items, py_request_objects) - assert not hasattr(items[0].request, "py_apply_batched_logits_processor") - assert items[1].request.py_apply_batched_logits_processor is True + # Receiving ranks rebuild the requests without the flag; attaching the + # broadcast objects must restore it on request 2 only. + received = [_make_broadcast_item(1), _make_broadcast_item(2)] + attach_py_objects_to_requests(received, py_request_objects) + assert not hasattr(received[0].request, "py_apply_batched_logits_processor") + assert received[1].request.py_apply_batched_logits_processor is True def test_single_round_non_root_receives_payload_or_nothing(): From 5432d1f8a3d858f14eb9c46e31a59a8c366a7e10 Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Thu, 29 Oct 2026 11:55:00 -0700 Subject: [PATCH 3/8] [None][fix] Copy the batched logits processor flag to child requests create_child_request copies each input the child needs by hand, and the flag was not among them, so the extra sequences of a request with n > 1 took the default and skipped the batched processor. Signed-off-by: Luke Miner --- tensorrt_llm/_torch/pyexecutor/llm_request.py | 2 ++ .../_torch/pyexecutor/test_request_utils.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index b3ac9cb74dcc..d2020f95eeca 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1388,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, ) diff --git a/tests/unittest/_torch/pyexecutor/test_request_utils.py b/tests/unittest/_torch/pyexecutor/test_request_utils.py index 56b14e5d9af0..d885c46653eb 100644 --- a/tests/unittest/_torch/pyexecutor/test_request_utils.py +++ b/tests/unittest/_torch/pyexecutor/test_request_utils.py @@ -437,6 +437,25 @@ def __deepcopy__(self, memo): assert len(sibling_processors) == processor_count +@pytest.mark.parametrize("apply_batched", [False, True]) +def test_child_factory_copies_the_batched_logits_processor_flag(apply_batched) -> None: + parent = LlmRequest( + request_id=42, + input_tokens=[1, 2, 3], + max_new_tokens=10, + sampling_config=trtllm.SamplingConfig(num_return_sequences=3), + is_streaming=False, + py_apply_batched_logits_processor=apply_batched, + ) + parent.create_child_request(43) + parent.create_child_request(44) + + assert [child.py_apply_batched_logits_processor for child in parent.child_requests] == [ + apply_batched, + apply_batched, + ] + + @pytest.mark.parametrize("runtime_field", ["py_mm_encoder_state", "py_mm_encoder_event"]) def test_child_factory_rejects_multimodal_encoder_ownership(runtime_field) -> None: parent = LlmRequest( From e7fd46d442e39563e08b6e2c771d5aab958bd268 Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Sun, 1 Nov 2026 07:15:05 -0800 Subject: [PATCH 4/8] [None][fix] Refuse the batched logits processor under speculative decoding; skip middle prefill chunks Review follow-ups for the PyTorch batched logits processor: - LLM.generate refuses apply_batched_logits_processor with speculative decoding on the PyTorch backend (RequestError, so the OpenAI frontend answers 4xx): one-model workers sample inside the forward, and the drafter-based modes sample 1 + draft_len rows per request while the post-processor walk steps by beam width. - The walk skips a context request that is not its last chunk; the prompt-length check alone cannot tell chunks apart. - greedy_tokens_usable lists the flag, next to the per-request processors. - PyExecutor's processor registration moves to a module-level function with its own tests; the l0_cpu constructor stub gains the field it reads. - Docs and docstrings say the processor runs on every TP rank of the last PP stage on the PyTorch backend, so it must be deterministic and free of side effects. Signed-off-by: Luke Miner --- docs/source/features/sampling.md | 2 +- .../engine/runners/decoder/runner.py | 10 ++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 34 ++++---- .../_torch/pyexecutor/sampler/sampler.py | 1 + tensorrt_llm/llmapi/llm.py | 25 ++++-- tensorrt_llm/sampling_params.py | 10 ++- .../test_batched_logits_processor.py | 79 +++++++++++++++++-- .../pyexecutor/test_remote_tail_adp_phase.py | 1 + .../_torch/sampler/test_greedy_tokens.py | 5 ++ 9 files changed, 136 insertions(+), 31 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index a0f13029902b..a4a42d05f916 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -440,4 +440,4 @@ 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. +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. On the PyTorch backend it runs on every tensor-parallel rank of the last pipeline stage, and each rank samples from its own copy of the edits (the TensorRT backend ran it once, on rank 0, and broadcast the tokens): 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. 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. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 1fa077b5b32a..552032c21ace 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -1205,8 +1205,14 @@ 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 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e4a5942ca92a..127ef8e0a6b9 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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,20 +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 - # 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) + _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 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index c0618db38761..e908169c6ae4 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/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 diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 0ffc0fa2bb9d..e631a4268c5c 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1642,12 +1642,25 @@ 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.apply_batched_logits_processor: + if 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 (self.args.backend == "pytorch" + and self.args.speculative_config is not None): + # One-model workers (MTP, Eagle3, ...) sample inside the + # forward, so the processor's edits never reach the tokens; + # the drafter-based modes (NGram, draft-target) sample + # 1 + draft_len rows per request while the post-processor + # walk steps by beam width, so it would hand the processor + # another request's rows. RequestError so the OpenAI frontend + # answers 4xx before the stream opens. + raise RequestError( + "`sampling_params.apply_batched_logits_processor=True` is " + "not supported with speculative decoding on the PyTorch " + "backend. Drop the flag or disable speculative decoding.") if sampling_params.return_routed_experts and not ( self.args.backend == "pytorch" diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index ca6b5755b88a..f13fae0513f2 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -171,6 +171,12 @@ def __call__( ) -> None: """Batched logits processing callback. The callback is expected to inplace modify the logits. + The TensorRT backend calls it once per step, on tensor-parallel rank 0, and broadcasts the + sampled tokens. The PyTorch backend calls it on every tensor-parallel rank of the last + pipeline stage, and each rank samples from its own copy of the edits: a callback with side + effects runs ``tp_size`` times per step, and a nondeterministic one can make the ranks + sample different tokens. + Args: req_ids (List[int]): A batch of request ids. logits (List[torch.Tensor]): A batch of the logits tensors. @@ -213,7 +219,9 @@ class SamplingParams: 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). + request's own logits_processor(s). On the PyTorch backend it runs on every tensor-parallel rank of the + last pipeline stage, each sampling from its own edits, so it must be deterministic and free of side + effects; it is not supported together with speculative decoding there. 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. diff --git a/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py index fdd8b5d9c26d..03b0f49a62cc 100644 --- a/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py +++ b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py @@ -14,6 +14,8 @@ import tensorrt_llm._torch.distributed as distributed from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import DecoderRunner +from tensorrt_llm._torch.pyexecutor.py_executor import _register_batched_logits_processor +from tensorrt_llm.executor.utils import RequestError from tensorrt_llm.llmapi.llm import BaseLLM from tensorrt_llm.sampling_params import SamplingParams @@ -44,9 +46,11 @@ def _request( admitted_beam_width=None, tokens=(1, 2, 3), orig_prompt_len=None, + is_last_context_chunk=True, ): tokens = list(tokens) return SimpleNamespace( + is_last_context_chunk=is_last_context_chunk, py_request_id=req_id, py_client_id=100 + req_id, py_logits_post_processors=processors, @@ -147,15 +151,28 @@ def test_beam_rows_and_histories_reach_the_batched_processor(): assert torch.equal(logits[:, 0], torch.tensor([1.0, 2.0, 2.0, 3.0, 5.0, 5.0, 6.0, 7.0])) -def test_an_unfinished_context_chunk_is_not_handed_over(): +@pytest.mark.parametrize( + "chunk_kwargs", + [ + # get_tokens(0) is the whole prompt during chunked prefill, so only + # the request's own chunk bookkeeping tells a middle chunk apart. + pytest.param({"is_last_context_chunk": False}, id="middle_chunk"), + pytest.param({"orig_prompt_len": 2}, id="prompt_shorter_than_tokens"), + ], +) +def test_an_unfinished_context_chunk_is_not_handed_over(chunk_kwargs): processor = _Recorder() - chunk = _request(1, batched=True, tokens=(1, 2, 3), orig_prompt_len=2) + chunk = _request(1, batched=True, tokens=(1, 2, 3), **chunk_kwargs) + last = _request(2, batched=True, tokens=(4, 5)) + logits = _logits(2) _engine(processor)._execute_logit_post_processors( - _scheduled(context=[chunk]), {"logits": _logits(1)} + _scheduled(context=[chunk, last]), {"logits": logits} ) - assert processor.calls == [] + # The last chunk's row still reaches the processor, at its own offset. + assert [call[0] for call in processor.calls] == [[2]] + assert torch.equal(logits[:, 0], torch.tensor([0.0, 2.0])) def test_ragged_verification_is_refused_for_batched_only_requests(): @@ -234,11 +251,21 @@ def allgather(tensor, mapping, dim=-1): assert [tuple(r.shape) for r in processor.calls[0][1]] == [(1, 1, 32)] -def test_the_flag_is_refused_without_a_batched_processor_on_the_llm(): - llm = SimpleNamespace( - args=SimpleNamespace(batched_logits_processor=None, enable_chunked_prefill=False), +def _llm(**args): + defaults = dict( + backend="pytorch", + batched_logits_processor=None, + enable_chunked_prefill=False, + speculative_config=None, + ) + return SimpleNamespace( + args=SimpleNamespace(**{**defaults, **args}), _check_one_model_speculative_sampling=lambda params: None, ) + + +def test_the_flag_is_refused_without_a_batched_processor_on_the_llm(): + llm = _llm() params = SamplingParams(apply_batched_logits_processor=True) with pytest.raises(ValueError, match="batched_logits_processor"): @@ -246,3 +273,41 @@ def test_the_flag_is_refused_without_a_batched_processor_on_the_llm(): llm.args.batched_logits_processor = _Recorder() BaseLLM._check_arguments(llm, 3, params, is_gen_only=True) + + +def test_the_flag_is_refused_with_speculative_decoding(): + """One-model workers sample before the walk; drafter modes sample rows the walk misaligns.""" + llm = _llm(batched_logits_processor=_Recorder(), speculative_config=object()) + params = SamplingParams(apply_batched_logits_processor=True) + + with pytest.raises(RequestError, match="speculative decoding"): + BaseLLM._check_arguments(llm, 3, params, is_gen_only=True) + + # A request without the flag is unaffected, and so is the flag once + # speculative decoding is off. + BaseLLM._check_arguments(llm, 3, SamplingParams(), is_gen_only=True) + llm.args.speculative_config = None + BaseLLM._check_arguments(llm, 3, params, is_gen_only=True) + + +def test_the_executor_registers_the_processor_with_the_engine(): + processor = _Recorder() + registered = [] + engine = SimpleNamespace(register_batched_logits_processor=registered.append) + + _register_batched_logits_processor(SimpleNamespace(batched_logits_processor=processor), engine) + + assert registered == [processor] + + +def test_the_executor_refuses_an_engine_without_a_post_processor_walk(): + """Fail at construction rather than serve requests nothing would process.""" + engine = SimpleNamespace() # no register_batched_logits_processor + + with pytest.raises(ValueError, match="SimpleNamespace"): + _register_batched_logits_processor( + SimpleNamespace(batched_logits_processor=_Recorder()), engine + ) + + # Without a processor there is nothing to register, whatever the engine. + _register_batched_logits_processor(SimpleNamespace(batched_logits_processor=None), engine) diff --git a/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py b/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py index 5dbd1c39a8b0..73826010da4a 100644 --- a/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py +++ b/tests/unittest/_torch/pyexecutor/test_remote_tail_adp_phase.py @@ -639,6 +639,7 @@ def reserve(resources, draft_len): stream_interval=1, attention_dp_config=None, kv_connector_config=None, + batched_logits_processor=None, batch_wait_timeout_ms=0, batch_wait_timeout_iters=0, batch_wait_max_tokens_ratio=0, diff --git a/tests/unittest/_torch/sampler/test_greedy_tokens.py b/tests/unittest/_torch/sampler/test_greedy_tokens.py index e5100456aadc..8bade494db10 100644 --- a/tests/unittest/_torch/sampler/test_greedy_tokens.py +++ b/tests/unittest/_torch/sampler/test_greedy_tokens.py @@ -63,6 +63,11 @@ def test_plain_greedy_request_is_usable(self): pytest.param( {}, _set("py_logits_post_processors", [print]), id="logits_post_processor" ), + pytest.param( + {}, + _set("py_apply_batched_logits_processor", True), + id="batched_logits_processor", + ), pytest.param({"return_log_probs": True}, None, id="log_probs"), pytest.param({"return_generation_logits": True}, None, id="generation_logits"), pytest.param({"stop_words_list": [[13, 14]]}, None, id="multi_token_stop_word"), From 61533c4735d782c11466fccda64ed701f53cb7a3 Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Sun, 1 Nov 2026 07:24:28 -0800 Subject: [PATCH 5/8] [None][chore] Describe the batched logits processor's ranks without the removed TRT backend The legacy TensorRT backend is gone from main (c77bc6ed6), so the docstring, docs and runner comments describe only where the processor runs now. The child-request flag test keeps only the case that fails without the copy. Signed-off-by: Luke Miner --- docs/source/features/sampling.md | 2 +- .../pyexecutor/engine/runners/decoder/runner.py | 8 ++++---- tensorrt_llm/sampling_params.py | 8 +++----- .../unittest/_torch/pyexecutor/test_request_utils.py | 12 +++++------- 4 files changed, 13 insertions(+), 17 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index a4a42d05f916..40390f649d6e 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -440,4 +440,4 @@ 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. On the PyTorch backend it runs on every tensor-parallel rank of the last pipeline stage, and each rank samples from its own copy of the edits (the TensorRT backend ran it once, on rank 0, and broadcast the tokens): 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. 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. +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. 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. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 552032c21ace..061940d10d2b 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -1061,7 +1061,7 @@ def _apply_batched_logits_processor( 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. + for processors that synchronize on it as the interface documents. """ device = logits[0].device stream_ptr = ( @@ -1216,9 +1216,9 @@ def _execute_logit_post_processors( logits_row_offset += row_stride continue - # 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. + # 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) diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index f13fae0513f2..c0c3a61d2d0c 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -171,11 +171,9 @@ def __call__( ) -> None: """Batched logits processing callback. The callback is expected to inplace modify the logits. - The TensorRT backend calls it once per step, on tensor-parallel rank 0, and broadcasts the - sampled tokens. The PyTorch backend calls it on every tensor-parallel rank of the last - pipeline stage, and each rank samples from its own copy of the edits: a callback with side - effects runs ``tp_size`` times per step, and a nondeterministic one can make the ranks - sample different tokens. + The callback runs on every tensor-parallel rank of the last pipeline stage, and each rank + samples from its own copy of the edits: a callback with side effects runs ``tp_size`` times + per step, and a nondeterministic one can make the ranks sample different tokens. Args: req_ids (List[int]): A batch of request ids. diff --git a/tests/unittest/_torch/pyexecutor/test_request_utils.py b/tests/unittest/_torch/pyexecutor/test_request_utils.py index d885c46653eb..e96dd628aa39 100644 --- a/tests/unittest/_torch/pyexecutor/test_request_utils.py +++ b/tests/unittest/_torch/pyexecutor/test_request_utils.py @@ -437,23 +437,21 @@ def __deepcopy__(self, memo): assert len(sibling_processors) == processor_count -@pytest.mark.parametrize("apply_batched", [False, True]) -def test_child_factory_copies_the_batched_logits_processor_flag(apply_batched) -> None: +def test_child_factory_copies_the_batched_logits_processor_flag() -> None: + # Only the opted-in parent tells a copy apart from the False default. parent = LlmRequest( request_id=42, input_tokens=[1, 2, 3], max_new_tokens=10, sampling_config=trtllm.SamplingConfig(num_return_sequences=3), is_streaming=False, - py_apply_batched_logits_processor=apply_batched, + py_apply_batched_logits_processor=True, ) parent.create_child_request(43) parent.create_child_request(44) - assert [child.py_apply_batched_logits_processor for child in parent.child_requests] == [ - apply_batched, - apply_batched, - ] + assert all(child.py_apply_batched_logits_processor for child in parent.child_requests) + assert len(parent.child_requests) == 2 @pytest.mark.parametrize("runtime_field", ["py_mm_encoder_state", "py_mm_encoder_event"]) From a112d8c572f442be4831b209afff9808c16ef7cc Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Tue, 10 Nov 2026 12:07:27 -0500 Subject: [PATCH 6/8] [None][test] Check the batched logits processor flag survives the request broadcast Main now sends Python-only request inputs inside PythonRequestInputs with the request, so test that path instead of the removed py-object broadcast. Signed-off-by: Luke Miner --- .../_torch/pyexecutor/test_request_utils.py | 31 ++++++++----------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/tests/unittest/_torch/pyexecutor/test_request_utils.py b/tests/unittest/_torch/pyexecutor/test_request_utils.py index e96dd628aa39..3ede9108758c 100644 --- a/tests/unittest/_torch/pyexecutor/test_request_utils.py +++ b/tests/unittest/_torch/pyexecutor/test_request_utils.py @@ -197,6 +197,19 @@ def test_request_broadcast_preserves_python_inputs(enabled, original_end_id) -> assert metrics.timing_metrics.arrival_time == datetime.timedelta(seconds=123) +@pytest.mark.parametrize("enabled", [False, True]) +def test_request_broadcast_carries_the_batched_logits_processor_flag(enabled) -> None: + # Pipeline stages other than the first rebuild requests from the broadcast + # envelope, so the flag must travel inside PythonRequestInputs. + request = trtllm.Request(input_token_ids=[1, 2, 3], max_tokens=10) + if enabled: + request.py_apply_batched_logits_processor = True + restored = _broadcast_round_trip(request) + assert restored.inputs.apply_batched_logits_processor is enabled + translated = executor_request_to_llm_request(42, restored, [], False) + assert translated.py_apply_batched_logits_processor is enabled + + def test_request_broadcast_preserves_absent_python_input_defaults() -> None: restored = _broadcast_round_trip( trtllm.Request(input_token_ids=[1, 2, 3], max_tokens=10, end_id=7) @@ -1031,24 +1044,6 @@ def test_single_round_busy_iteration_sends_requests_and_py_objects_once(): assert payload == items -def test_py_objects_carry_the_batched_logits_processor_flag(): - # The flag is only set on requests that opted in, so only those cross the - # pipeline stages with it. - broadcaster = RequestBroadcaster(_SingleRoundDist(rank=0), MagicMock()) - items = [_make_broadcast_item(1), _make_broadcast_item(2)] - items[1].request.py_apply_batched_logits_processor = True - - py_request_objects = broadcaster._collect_py_objects(items) - - assert ("py_apply_batched_logits_processor", {2: True}) in py_request_objects - # Receiving ranks rebuild the requests without the flag; attaching the - # broadcast objects must restore it on request 2 only. - received = [_make_broadcast_item(1), _make_broadcast_item(2)] - attach_py_objects_to_requests(received, py_request_objects) - assert not hasattr(received[0].request, "py_apply_batched_logits_processor") - assert received[1].request.py_apply_batched_logits_processor is True - - def test_single_round_non_root_receives_payload_or_nothing(): items = [_make_broadcast_item(7)] dist = _SingleRoundDist(rank=1, received=items) From 23b70e04994276a27b91ab4eefe5d20a57c283e5 Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Tue, 10 Nov 2026 13:13:28 -0500 Subject: [PATCH 7/8] [None][chore] Drop the paused-request test case; note attention DP in the batched processor docs The prompt_shorter_than_tokens case asserted that a request whose prompt grew (after a pause folds generated tokens into it) is skipped, which is main's existing behavior rather than something this PR intends. Under attention DP each rank calls the processor with only its own requests. Signed-off-by: Luke Miner --- docs/source/features/sampling.md | 2 +- tensorrt_llm/sampling_params.py | 6 ++++-- .../test_batched_logits_processor.py | 18 +++++------------- 3 files changed, 10 insertions(+), 16 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 40390f649d6e..b41fecb0f1c1 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -440,4 +440,4 @@ 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. 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. +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. diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index c0c3a61d2d0c..8d1dbc41e0bb 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -173,7 +173,8 @@ def __call__( The callback runs on every tensor-parallel rank of the last pipeline stage, and each rank samples from its own copy of the edits: a callback with side effects runs ``tp_size`` times - per step, and a nondeterministic one can make the ranks sample different tokens. + per step, and a nondeterministic one can make the ranks sample different tokens. With + attention data parallelism each rank instead calls it with only the requests it owns. Args: req_ids (List[int]): A batch of request ids. @@ -219,7 +220,8 @@ class SamplingParams: The callback is called once per decoding step with every scheduled request that sets this, after each request's own logits_processor(s). On the PyTorch backend it runs on every tensor-parallel rank of the last pipeline stage, each sampling from its own edits, so it must be deterministic and free of side - effects; it is not supported together with speculative decoding there. + effects; it is not supported together with speculative decoding there. With attention data + parallelism each rank calls it with only the requests that rank owns. 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. diff --git a/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py index 03b0f49a62cc..c0c67c7f6f78 100644 --- a/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py +++ b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py @@ -45,7 +45,6 @@ def _request( beam_width=1, admitted_beam_width=None, tokens=(1, 2, 3), - orig_prompt_len=None, is_last_context_chunk=True, ): tokens = list(tokens) @@ -57,7 +56,7 @@ def _request( py_apply_batched_logits_processor=batched, py_beam_width=admitted_beam_width or beam_width, py_prompt_len=len(tokens), - py_orig_prompt_len=len(tokens) if orig_prompt_len is None else orig_prompt_len, + py_orig_prompt_len=len(tokens), get_beam_width_by_iter=lambda for_next_iteration: beam_width, get_tokens=lambda beam: list(tokens), ) @@ -151,18 +150,11 @@ def test_beam_rows_and_histories_reach_the_batched_processor(): assert torch.equal(logits[:, 0], torch.tensor([1.0, 2.0, 2.0, 3.0, 5.0, 5.0, 6.0, 7.0])) -@pytest.mark.parametrize( - "chunk_kwargs", - [ - # get_tokens(0) is the whole prompt during chunked prefill, so only - # the request's own chunk bookkeeping tells a middle chunk apart. - pytest.param({"is_last_context_chunk": False}, id="middle_chunk"), - pytest.param({"orig_prompt_len": 2}, id="prompt_shorter_than_tokens"), - ], -) -def test_an_unfinished_context_chunk_is_not_handed_over(chunk_kwargs): +def test_a_middle_context_chunk_is_not_handed_over(): + # get_tokens(0) is the whole prompt during chunked prefill, so only the + # request's own chunk bookkeeping tells a middle chunk apart. processor = _Recorder() - chunk = _request(1, batched=True, tokens=(1, 2, 3), **chunk_kwargs) + chunk = _request(1, batched=True, tokens=(1, 2, 3), is_last_context_chunk=False) last = _request(2, batched=True, tokens=(4, 5)) logits = _logits(2) From 2385d8874778a0ef8ee5740681e18433197643c3 Mon Sep 17 00:00:00 2001 From: Luke Miner Date: Tue, 10 Nov 2026 15:43:32 -0500 Subject: [PATCH 8/8] [None][chore] Keep BatchedLogitsProcessor's stream_ptr an int The PyTorch backend's logits are always on a CUDA device, so the engine always passes the current stream and the callback's signature does not change. The CPU unit tests stand in a stream. Signed-off-by: Luke Miner --- docs/source/features/sampling.md | 2 +- .../pyexecutor/engine/runners/decoder/runner.py | 9 +++------ tensorrt_llm/sampling_params.py | 6 ++---- .../pyexecutor/test_batched_logits_processor.py | 14 +++++++++++++- .../references/batched_logits_processor.yaml | 2 +- tests/unittest/llmapi/test_llm.py | 6 ++---- 6 files changed, 22 insertions(+), 17 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index b41fecb0f1c1..b7c4ec51d142 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -428,7 +428,7 @@ class MyBatchedLogitsProcessor(BatchedLogitsProcessor): req_ids: List[int], logits: List[torch.Tensor], token_ids: List[List[List[int]]], - stream_ptr: Optional[int], + stream_ptr: int, client_ids: List[Optional[int]] ) -> None: for rows in logits: diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 061940d10d2b..f399967cda54 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -1060,13 +1060,10 @@ def _apply_batched_logits_processor( 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 that synchronize on it as the interface documents. + is the one the logits live on, for processors that synchronize on it as + the interface documents. """ - device = logits[0].device - stream_ptr = ( - torch.cuda.current_stream(device).cuda_stream if device.type == "cuda" else None - ) + 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( diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 8d1dbc41e0bb..ea66d95d00a0 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -166,7 +166,7 @@ def __call__( req_ids: List[int], logits: List[torch.Tensor], token_ids: List[List[List[int]]], - stream_ptr: Optional[int], + stream_ptr: int, client_ids: List[Optional[int]], ) -> None: """Batched logits processing callback. The callback is expected to inplace modify the logits. @@ -182,9 +182,7 @@ def __call__( 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. 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. + stream_ptr (int): The operation stream used by the logits tensors. client_ids (List[Optional[int]]): A batch of optional client ids. """ pass # noqa diff --git a/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py index c0c67c7f6f78..0cffe33ffa35 100644 --- a/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py +++ b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py @@ -62,6 +62,18 @@ def _request( ) +# The tests' logits are CPU tensors, so stand in for the CUDA stream the +# engine hands the callback. +_STREAM_PTR = 1234 + + +@pytest.fixture(autouse=True) +def _cuda_stream(monkeypatch): + monkeypatch.setattr( + torch.cuda, "current_stream", lambda device=None: SimpleNamespace(cuda_stream=_STREAM_PTR) + ) + + def _engine(processor, *, attention_dp=False, vocab_size=None): engine = object.__new__(DecoderRunner) engine._mapping = SimpleNamespace( @@ -103,7 +115,7 @@ def test_flagged_requests_share_one_call_in_scheduling_order(): assert client_ids == [101, 103] assert token_ids == [[[4, 5]], [[6, 7, 8]]] assert [tuple(r.shape) for r in rows] == [(1, 1, VOCAB), (1, 1, VOCAB)] - assert stream_ptr is None # the logits are not on a CUDA device + assert stream_ptr == _STREAM_PTR # The callback's in-place edits land in the step's logits; the other request's row is untouched. assert torch.equal(logits[:, 0], torch.tensor([1.0, 1.0, 3.0])) diff --git a/tests/unittest/api_stability/references/batched_logits_processor.yaml b/tests/unittest/api_stability/references/batched_logits_processor.yaml index bdec874334a4..423532e31ca5 100644 --- a/tests/unittest/api_stability/references/batched_logits_processor.yaml +++ b/tests/unittest/api_stability/references/batched_logits_processor.yaml @@ -11,7 +11,7 @@ methods: annotation: List[int] default: inspect._empty stream_ptr: - annotation: Optional[int] + annotation: int default: inspect._empty token_ids: annotation: List[List[List[int]]] diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 5281d213d100..a8947d8e0508 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -714,12 +714,10 @@ def __init__(self, biased_word_id): self.biased_word_id = biased_word_id def __call__(self, req_ids: List[int], logits: List[torch.Tensor], - ids: List[List[List[int]]], stream_ptr: Optional[int], + ids: List[List[List[int]]], stream_ptr: int, client_ids: List[Optional[int]]): assert len(req_ids) == len(logits) == len(ids) == len(client_ids) - stream = None if stream_ptr is None else torch.cuda.ExternalStream( - stream_ptr) - with torch.cuda.stream(stream): + with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): for rows in logits: rows[:] = float("-inf") rows[..., self.biased_word_id] = 0