[lminer]

Description

LLM(batched_logits_processor=...) and SamplingParams(apply_batched_logits_processor=True) are both accepted on the PyTorch backend, but neither does anything there: nothing under _torch/pyexecutor reads the LLM field, and the request flag is dropped when the executor request is built. The batched processor was a TensorRT-backend feature that was never wired into the PyTorch executor, and the docs still present it as available.

This PR makes the PyTorch backend honor both, with the existing public API unchanged:

• DecoderRunner._execute_logit_post_processors already walks the scheduled requests and hands each one's logits rows to its own logits_processor(s). It now also collects the rows, token ids, request ids and client ids of every request that set apply_batched_logits_processor, and hands them to the LLM's batched processor in one call per step, after the walk. Each logits entry is the request's rows viewed as [beam_width, 1, vocab], the shape the per-request processors see, and a view into the step's logits tensor, so in-place edits land. Under beam search the rows are the live beams and the token ids are each beam's own history ([#19902][fix] Give logits processors each beam's hypothesis under beam search NVIDIA/TensorRT-LLM#20915). The row offsets, the context-chunk skip and the vocabulary-sharded spec-worker gather are the ones the per-request path already uses.
• PyExecutor registers llm_args.batched_logits_processor with the model engine, the same way the beam-history and in-graph sampling hooks are registered.
• The request flag is attached to the executor request in BaseWorker, kept on LlmRequest as py_apply_batched_logits_processor, and crosses pipeline stages with the other Python-only request attributes.
• A request with the flag stays on the eager sampling tier, like a request with its own processors: in-graph sampling would otherwise read the logits before the processor runs.
• LLM.generate refuses the flag when the LLM has no batched processor, instead of silently generating without it. PyExecutor raises at construction when a processor is given but the model engine has no post-processor walk to attach it to, rather than serving requests that asked for processing nothing would run. A flagged request that still reaches a walk with nothing registered (only through PyExecutor directly) is left unprocessed with a one-time warning rather than failing the whole step.
• BatchedLogitsProcessor.__call__ types stream_ptr as Optional[int]: on the PyTorch backend it is the CUDA stream the logits live on, or None when they are not on a CUDA device (CPU unit tests). The API-stability reference is updated accordingly.

Two things are deliberately unchanged. Under a vocabulary-sharded spec worker the walk hands processors a gathered copy of the logits, which the per-request path already did; the worker's acceptance ran inside the forward and processors get a read-only view there (SpecWorkerBase.target_logits), so nothing is written back, and the docs now say that such edits do not reach the sampled tokens. And the walk's row stride under speculative decoding is the per-request path's existing one; this PR does not change which rows a request's processors see. The [beam_width, 1, vocab] view of a request's rows is now built once in the walk and handed to both the per-request processors and the batched list.

• docs/source/features/sampling.md gains a short section on the batched processor; the SamplingParams and BatchedLogitsProcessor docstrings say when it is called and what it receives.

Per-request processors run first and the batched one after them, so the batched processor sees their edits. Nothing changes for requests that do not set the flag, or when the LLM has no batched processor.

Test Coverage

• tests/unittest/_torch/executor/test_batched_logits_processor.py (new, CPU only): one call per step in scheduling order with views into the step's logits; per-request processors before the batched one; live beam rows and beam histories; an unfinished context chunk is skipped; nothing is called without opted-in requests; a flagged request without a registered processor is skipped with a one-time warning; beam histories start for batched-only requests; the vocabulary-sharded gather runs for a batched-only request; the flag is refused without an LLM processor.
• tests/unittest/_torch/sampler/test_in_graph_sampling.py: a request with the flag demotes the batch to the FULL tier.
• tests/unittest/_torch/executor/test_request_utils.py: the flag is collected and attached like the other Python-only request attributes.
• tests/unittest/llmapi/test_llm_pytorch.py::test_tinyllama_batched_logits_processor (GPU): two prompts, one opted in; the opted-in one is forced to the biased token, the other is not.

There is no GPU here that runs main, so the GPU checks ran on one L4 with this change ported onto the release:1.4.0rc1 image (together with NVIDIA#20915, which that tag predates); the runner file is laid out differently there, so the port differs in placement, not in logic. On main itself the change is lint- and type-checked, so CI is the first GPU run of this exact code.

Results on that image:

• test_batched_logits_processor.py with the changed test_in_graph_sampling.py and test_request_utils.py cases: 59 passed (the vocabulary-sharded cases are main-only and were not ported). On the stock image the 8 new cases all fail.
• TinyLlama-1.1B end to end, greedy, 6 tokens: the opted-in request is forced to the biased token; a request without the flag is untouched; a request with both its own processor (forcing one token) and the flag ends up with the batched processor's token, so the batched one runs last; the flag is refused with a ValueError on an LLM without a batched processor. On the stock image the opted-in requests come back unforced and the flag is accepted silently.
• Beam search (width 2), the same forcing through the batched processor and through a per-request processor: identical output, beam for beam (with a single allowed token there is only one distinct hypothesis, so the second beam is empty either way).
• Regression on an earlier revision of this branch: test_beam_search.py + test_whisper_suppress_tokens_processor.py: 322 passed, 4 skipped; the 12 errors are the same environment errors as on the stock image.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[lminer]

Opened upstream as NVIDIA#21044; this fork copy served for review.