[Bug]: Streaming requests with logprobs fail after a KV cache V2 recompute pause #22091

System Info

• TensorRT-LLM main at 9ceb00a84, PyTorch backend with the defaults: KV cache manager V2 and the overlap scheduler, aggregated serving.
• Found while reading the executor's response path and reproduced on CPU (details below). No GPU run.

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 streaming request with logprobs loses logprobs when KVCacheV2Scheduler recompute-pauses it in the overlap loop. Under KV cache pressure the scheduler picks a started generation request as the victim. In the overlap loop that request is usually in the previous batch (inflight_req_ids is only maintained by the PP loop), so one iteration does the following for it:

1. _update_requests appends its in-flight token and that token's logprob.
2. _pause_recompute_paused_requests calls reset_for_recompute: the native pause() folds the generated tokens into the prompt and sets CONTEXT_INIT, and py_decoding_iter becomes 0.
3. _handle_responses calls create_response() for it, because 0 % stream_interval == 0.

For a streaming request with logprobs (beam width 1), create_response() moves py_result._log_probs into the response's copy of py_result and gives the request an empty LogProbStorage. Only after that does it check the native result, which is empty in CONTEXT_INIT, and return None, so the moved logprobs are dropped. The native request has not sent the token (mMaxSentTokenLen is unchanged), so after the re-prefill the next response carries that token and the new one, but only the new token's logprob.

The same sequence with LlmRequest directly (it needs no GPU):

from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState
from tensorrt_llm.bindings import executor as trtllm

request = LlmRequest(request_id=42, max_new_tokens=10, input_tokens=[1, 2, 3],
                     sampling_config=trtllm.SamplingConfig(), is_streaming=True,
                     return_log_probs=True, logprobs_simple_format=True)
request.state = LlmRequestState.GENERATION_IN_PROGRESS
request.add_new_tokens([4])
request.py_result.append_log_probs([[-0.5]])
print(request.create_response().result.log_probs)  # [[-0.5]]

request.add_new_tokens([5])                         # _update_requests
request.py_result.append_log_probs([[-0.25]])
request.reset_for_recompute(20)                     # _pause_recompute_paused_requests
print(request.create_response())                    # _handle_responses: None

request.state = LlmRequestState.GENERATION_IN_PROGRESS  # re-prefill finished
request.add_new_tokens([6])
request.py_result.append_log_probs([[-0.125]])
response = request.create_response()
print(response.result.get_result().output_token_ids)  # [[5, 6]]
print(response.result.log_probs)                      # [[-0.125]], expected [[-0.25, -0.125]]

I ran this sequence on CPU against a Python port of LlmRequest::createResult and pause() from llmRequest.cpp/llmRequest.h, because this machine has no TensorRT-LLM build; everything else was the repository code. I also drove PyExecutor._handle_responses through the three steps and passed its responses through pickle into GenerationResult._handle_response. The client received tokens [10, 11, 12] with logprobs [-0.1, -0.3], and the proxy raised AssertionError: logprobs length: 2 != output.length: 3 in _handle_sequence.

Expected behavior

Every streamed token arrives with its logprob. When create_response() produces no response, the pending logprobs stay in the request and go out with the next response.

actual behavior

The logprobs added since the last emitted response are discarded. The next response has more tokens than logprobs, and the proxy's assert len(output.logprobs) == output.length fails, so the streaming request errors out.

additional notes

• Affected: streaming requests with logprobs and beam width 1 on aggregated servers, where enable_recompute_pause is on (not is_disagg). The overlap loop hits it with the default stream_interval=1 whenever a running request is recompute-paused. The non-overlap loop pauses victims before sampling, so it only loses logprobs with stream_interval > 1; in both loops every logprob accumulated since the last emitted response is lost.
• Proposed fix: in LlmRequest.create_response, return None right after create_serialized_result() when the result is empty, before the streaming logprob handoff. The generation-logits branch only copies, so it is unaffected. I will open a PR with a CPU regression test next to test_python_logprob_response_snapshots.
• Searched open and closed issues and PRs (create_response, reset_for_recompute, logprobs with recompute or preemption) and found no report of this.

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.

AI assistance (Cursor) was used to investigate and reproduce this issue.