[Bug]: PyTorch backend: per-request logits_processor plus return_context_logits gives wrong tokens #22044

System Info

• TensorRT-LLM main at 1f8bb80 (PyTorch backend); also before [TRTLLM-17516][fix] Bound context logits memory with CPU fallback #21315
• Tested on one NVIDIA L4 with nvcr.io/nvidia/tensorrt-llm/release:1.4.0rc3 and main's Python code, model Qwen/Qwen3-0.6B

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

A per-request logits_processor turns off the context-logits plan from #21315. If a request in the same step sets return_context_logits=True, the model then returns a row of logits for every prompt token of every prompt request, but the sampler and the logits processors expect one row for each request that did not ask for context logits.

import asyncio, torch
from tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.llmapi import KvCacheConfig
from tensorrt_llm.sampling_params import LogitsProcessor

class Force(LogitsProcessor):
    def __init__(self, token_id): self.token_id = token_id
    def __call__(self, req_id, logits, token_ids, stream_ptr, client_id):
        stream = None if stream_ptr is None else torch.cuda.ExternalStream(stream_ptr)
        with torch.cuda.stream(stream):
            logits[:] = float("-inf")
            logits[..., self.token_id] = 0

async def main():
    llm = LLM("Qwen/Qwen3-0.6B", kv_cache_config=KvCacheConfig(enable_block_reuse=False))
    z = llm.tokenizer.encode("Z", add_special_tokens=False)[-1]
    g = dict(max_tokens=3, temperature=0.0)
    a, b = "The capital of France is", "One, two, three, four, five, six, seven,"

    # 1. One request with context logits and a processor.
    out = await llm.generate_async(a, SamplingParams(**g, return_context_logits=True, logits_processor=Force(z)))
    print(repr(out.outputs[0].text))

    # 2. That request sharing a prompt step with plain requests (the first request
    #    is scheduled alone; the next three share a step).
    futures = [llm.generate_async(p, sp) for p, sp in [
        (a, SamplingParams(**g)),
        (a, SamplingParams(**g, return_context_logits=True, logits_processor=Force(z))),
        (b, SamplingParams(**g)),
        (b, SamplingParams(**g)),
    ]]
    print([repr((await f).outputs[0].text) for f in futures])

asyncio.run(main())

Expected behavior

1. 'ZZZ'
2. ' Paris. The', 'ZZZ', ' eight, nine', ' eight, nine' (what each prompt gives on its own)

actual behavior

1. ' ParisZZ': the processor edits the first prompt row, so the first generated token ignores it.
2. ' Paris, and', ' ParisZZ', ' Answer: ', ' two, three': every request after the one with context logits, including one already generating, samples from another request's rows. Its output is wrong and nothing reports an error.

additional notes

• The context-logits plan stores context logits during the forward and keeps one sampling row per prompt request. The fallback taken when the plan is off still returns every row. Besides a per-request processor anywhere in the step, the plan is off for quantized, LoRA or row-parallel LM heads, Cohere2 and Gemma4, attention DP and context parallelism, so mixed steps there can misread rows the same way without any processor.
• Separately, _execute_logit_post_processors decides a prompt request's final chunk with py_orig_prompt_len < len(tokens); after a pause folds generated tokens into the prompt, that skips the resumed request's final chunk, so one token misses its processor.
• test_tinyllama_logits_processor compares with an 80% text match, which passes with one wrong token in six.
• Fix proposed in [#22044][fix] Keep logits rows aligned with the sampler when context logits are gathered without the plan #22038.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.