diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index abbc61587709..b7c4ec51d142 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: 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. diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 7419e578f422..f399967cda54 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,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,7 +1134,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 +1142,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 +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, diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 77a41e8277b3..d2020f95eeca 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( @@ -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, 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..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,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 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..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 @@ -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: 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..e631a4268c5c 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -1642,6 +1642,26 @@ 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: + 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" 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..ea66d95d00a0 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -171,11 +171,17 @@ def __call__( ) -> None: """Batched logits processing callback. The callback is expected to inplace modify the logits. + 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. 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. 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. + 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): The operation stream used by the logits tensors. client_ids (List[Optional[int]]): A batch of optional client ids. """ @@ -209,6 +215,11 @@ 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). 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. 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 new file mode 100644 index 000000000000..0cffe33ffa35 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_batched_logits_processor.py @@ -0,0 +1,317 @@ +# 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._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 + +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), + 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, + 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), + get_beam_width_by_iter=lambda for_next_iteration: beam_width, + get_tokens=lambda beam: list(tokens), + ) + + +# 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( + 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 == _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])) + + +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_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), is_last_context_chunk=False) + last = _request(2, batched=True, tokens=(4, 5)) + + logits = _logits(2) + _engine(processor)._execute_logit_post_processors( + _scheduled(context=[chunk, last]), {"logits": logits} + ) + + # 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(): + 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 _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"): + 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) + + +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/pyexecutor/test_request_utils.py b/tests/unittest/_torch/pyexecutor/test_request_utils.py index ebb9da1efb47..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) @@ -437,6 +450,23 @@ def __deepcopy__(self, memo): assert len(sibling_processors) == processor_count +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=True, + ) + parent.create_child_request(43) + parent.create_child_request(44) + + 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"]) def test_child_factory_rejects_multimodal_encoder_ownership(runtime_field) -> None: parent = LlmRequest( 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"), 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/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 56b016540fe9..a8947d8e0508 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,44 @@ 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: int, + client_ids: List[Optional[int]]): + assert len(req_ids) == len(logits) == len(ids) == len(client_ids) + with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): + 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",