You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[#22091][fix] Keep streaming logprobs when create_response has no result - #22093
LlmRequest.create_response now checks for an empty serialized result before the streaming logprob handoff. When the result is empty, the method returns None and py_result keeps its pending logprobs. The later empty-result check is removed. Non-empty responses follow the same path as before.
Two points to note:
For empty results, the early return now skips the time-breakdown metrics block that ran before the old check. The caller receives None either way, so no response-visible metric changes. Confirm that no side effect in that block is needed for empty results.
The change depends on the empty-result state being the only signal that a response is not produced. The fix does not change the native create_serialized_result() path.
The test runs LlmRequest directly on a streaming request with logprobs. It checks three steps:
The first response carries token [4] and logprob [[-0.5]].
After reset_for_recompute, create_response() returns None.
The next response carries tokens [5, 6] and logprobs [-0.25, -0.125].
The author reports that the test fails on main and passes with this change. These runs used a Python port and import stubs, not a built TensorRT-LLM. Native CI results are not available for this head.
No integration test-list change applies. The author reports that the file falls under the unittest/_torch/pyexecutor entry of l0_cpu. This was not verified in the test-db lists.
Not covered: the rest of test_request_utils.py, GPU tests, the beam-width greater than 1 path, and end-to-end serving under KV cache pressure.
Coverage verdict: needs follow-up. Confirm the test passes in native CI, and add an end-to-end or pressure test if serving behavior must be verified.
Per-File QA Perspective
tensorrt_llm/_torch/pyexecutor/llm_request.py: Verify that a recompute-paused streaming request with logprobs keeps its pending logprobs. Verify that the next response has one logprob for each token. Verify that non-empty responses behave as before.
tests/unittest/_torch/pyexecutor/test_request_utils.py: Adds a CPU regression test for the recompute-pause logprob loss. Confirm it runs in CI with real bindings.
Description
LlmRequest.create_response hands a streaming request's pending logprobs to the response: it moves py_result._log_probs into the response's copy of py_result and gives the request a fresh LogProbStorage. It did this before checking the serialized native result, and when that result was empty it returned None, so the moved logprobs were dropped.
The native result is empty whenever the request is outside the generation states. With the KV cache manager V2 (the default), KVCacheV2Scheduler recompute-pauses started generation requests under KV cache pressure. In the overlap loop the victim is usually in the previous batch, so a single iteration does three things for it:
_update_requests appends its in-flight token and logprob.
_pause_recompute_paused_requests calls reset_for_recompute, which sets CONTEXT_INIT and py_decoding_iter = 0.
_handle_responses calls create_response() for it.
The token itself was not sent, so after the re-prefill it is streamed together with the next token, but only the next token's logprob is left. GenerationResultBase._handle_sequence then fails assert len(output.logprobs) == output.length, and the request errors out. With stream_interval > 1 every logprob accumulated since the last emitted response is lost the same way, in the non-overlap loop as well.
This change returns None right after create_serialized_result() when the result is empty, so the logprobs stay in the request until a response carries them. Responses that have a result are built exactly as before.
New test_streaming_logprobs_survive_a_response_without_result in tests/unittest/_torch/pyexecutor/test_request_utils.py (cpu_only, covered by the unittest/_torch/pyexecutor entry of l0_cpu). A streaming request with logprobs is recompute-paused after its token and logprob were added. create_response() then returns None, and after the re-prefill the next response must carry both tokens and both logprobs. The test fails on main ([[-0.125]] != [[-0.25, -0.125]]) and passes with this change.
This machine has no TensorRT-LLM build. I ran the new test, plus the four beam-width-1 cases of the existing test_python_logprob_response_snapshots (4 passed), with --noconftest against a Python port of LlmRequest::createResult, createSerializedResult and pause() from llmRequest.cpp/llmRequest.h, with import stubs for the other bindings. In CI the test runs against the real bindings.
With the same port, I drove PyExecutor._handle_responses through the three steps above and passed its responses through pickle into GenerationResult._handle_response. On main the proxy raises AssertionError: logprobs length: 2 != output.length: 3; with this change it receives token_ids [10, 11, 12] with logprobs [-0.1, -0.2, -0.3].
pre-commit run --files on both files: all hooks passed; Pyrefly reported no diagnostics.
Not run: the rest of test_request_utils.py, GPU tests, and an end-to-end serving run under KV cache pressure.
PR Checklist
Please check this after reviewing the above items as appropriate for this PR.
The reason will be displayed to describe this comment to others. Learn more.
LGTM. createResult returns nullopt before setMaxSentTokenLen, so those tokens go out with the next response; returning before the move keeps their logprobs with them. Non-empty responses are built exactly as before, and _handle_responses just skips a None response.
No actionable comments were generated in the recent review.
Walkthrough
create_response now returns before moving streaming logprobs when the serialized result is empty. A regression test checks that logprobs accumulated during a recompute pause are included with the corresponding tokens after generation resumes.
Changes
Streaming logprob preservation
Layer / File(s)
Summary
Guard empty responses before logprob handoff tensorrt_llm/_torch/pyexecutor/llm_request.py, tests/unittest/_torch/pyexecutor/test_request_utils.py
create_response returns None for an empty serialized result before moving streaming logprobs. The regression test checks that a paused step produces no response and that the next response contains the pending tokens with their corresponding logprobs.
This fix keeps streaming logprobs aligned with tokens when a response is skipped during a recompute pause. No concrete merge-blocking risk was identified. Real-binding CI will confirm the test results.
Pre-merge checks: 4
Check name
Status
Explanation
Title check
The title names the main change: streaming logprobs are kept when create_response has no result. It follows the repository template with the issue number and the fix type. The changeset moves the empt…
Description check
The description covers the problem, the cause, the fix, and the test coverage. It explains that the logprobs were dropped when create_response returned None, and that the change keeps them until a res…
Linked Issues check
[#22091] The whole-PR diff against the merge base moves the len(result) == 0 check in LlmRequest.create_response to directly after create_serialized_result(). An empty result now returns None …
Out of Scope Changes check
The whole-PR diff changes tensorrt_llm/_torch/pyexecutor/llm_request.py (a 4-line move of the empty-result check with a explanatory comment) and adds one focused regression test in `tests/unittest/_…
Full details: Title check
Explanation
The title names the main change: streaming logprobs are kept when create_response has no result. It follows the repository template with the issue number and the fix type. The changeset moves the empty-result check ahead of the streaming logprob handoff in LlmRequest.create_response, which matches the title.
Full details: Description check
Explanation
The description covers the problem, the cause, the fix, and the test coverage. It explains that the logprobs were dropped when create_response returned None, and that the change keeps them until a response carries them. It links the issue (#22091) and lists the new test and the checks that were run. It does not use the template's exact title-format guidance, and the PR checklist box is checked without the listed items being marked individually. These gaps are minor, so the description is mostly complete.
Full details: Linked Issues check
Explanation
[#22091] The whole-PR diff against the merge base moves the len(result) == 0 check in LlmRequest.create_response to directly after create_serialized_result(). An empty result now returns None before py_result._log_probs is moved into the response and the request receives a fresh LogProbStorage. The pending logprobs therefore stay in the request until a response carries a result, which matches the issue's expected behavior. The generation-logits branch is unchanged in behavior. The added test test_streaming_logprobs_survive_a_response_without_result covers the recompute-pause sequence and checks that the next response carries both tokens and both logprobs. The test and the code change were not executed in this review; the assessment relies on the diff and the author's reported runs. The change since the previous review is a rebase onto efec60e, and the code diff shows the same fix.
Full details: Out of Scope Changes check
Explanation
The whole-PR diff changes tensorrt_llm/_torch/pyexecutor/llm_request.py (a 4-line move of the empty-result check with a explanatory comment) and adds one focused regression test in tests/unittest/_torch/pyexecutor/test_request_utils.py. Both directly support #22091. One side effect is visible in the diff: the early return now also skips the time-breakdown metrics block that previously ran before the empty check. This affects only empty-result calls and is not described in the PR, but it does not change responses that carry a result. This is a minor note and does not establish unrelated work.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 8 remain after this review.
Could you please rebase this branch onto the latest upstream main to include the CI hang fix in #22101 (6edbe383497a)? The current branch still contains the regression that can cause pipelines to hang for an extended period and worsen CI queueing
…no result
LlmRequest.create_response moved a streaming request's pending logprobs into the response's copy of py_result before checking whether the native result was empty, and dropped them when it was. The KV cache V2 scheduler can recompute-pause a request in the step that adds its token; the response built for it in that step has no result, so the token is streamed after the re-prefill without its logprob and the proxy fails the request on its logprobs length check. Return before touching py_result when there is no result.
AI assistance (Cursor) was used for this change.
Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: YangHaoooooo <233765393+Hang-get@users.noreply.github.com>
Rebased onto current main (efec60ea5), which includes the CI hang fix from #22101 (6edbe383). This branch is now 67e96e290. The patch is unchanged.
This branch has not been deployed
No deployments
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
JunyiXu-nv
left a comment
svc-trtllm-gh-bot
added
the
Community want to contribute
Dev Engineer Review
LlmRequest.create_responsenow checks for an empty serialized result before the streaming logprob handoff. When the result is empty, the method returnsNoneandpy_resultkeeps its pending logprobs. The later empty-result check is removed. Non-empty responses follow the same path as before.Two points to note:
Noneeither way, so no response-visible metric changes. Confirm that no side effect in that block is needed for empty results.create_serialized_result()path.No public API or declaration changes.
QA Engineer Review
Test files changed:
tests/unittest/_torch/pyexecutor/test_request_utils.py: addedtest_streaming_logprobs_survive_a_response_without_result.The test runs
LlmRequestdirectly on a streaming request with logprobs. It checks three steps:[4]and logprob[[-0.5]].reset_for_recompute,create_response()returnsNone.[5, 6]and logprobs[-0.25, -0.125].The author reports that the test fails on
mainand passes with this change. These runs used a Python port and import stubs, not a built TensorRT-LLM. Native CI results are not available for this head.No integration test-list change applies. The author reports that the file falls under the
unittest/_torch/pyexecutorentry ofl0_cpu. This was not verified in the test-db lists.Not covered: the rest of
test_request_utils.py, GPU tests, the beam-width greater than 1 path, and end-to-end serving under KV cache pressure.Coverage verdict: needs follow-up. Confirm the test passes in native CI, and add an end-to-end or pressure test if serving behavior must be verified.
Per-File QA Perspective
tensorrt_llm/_torch/pyexecutor/llm_request.py: Verify that a recompute-paused streaming request with logprobs keeps its pending logprobs. Verify that the next response has one logprob for each token. Verify that non-empty responses behave as before.tests/unittest/_torch/pyexecutor/test_request_utils.py: Adds a CPU regression test for the recompute-pause logprob loss. Confirm it runs in CI with real bindings.Description
LlmRequest.create_responsehands a streaming request's pending logprobs to the response: it movespy_result._log_probsinto the response's copy ofpy_resultand gives the request a freshLogProbStorage. It did this before checking the serialized native result, and when that result was empty it returnedNone, so the moved logprobs were dropped.The native result is empty whenever the request is outside the generation states. With the KV cache manager V2 (the default),
KVCacheV2Schedulerrecompute-pauses started generation requests under KV cache pressure. In the overlap loop the victim is usually in the previous batch, so a single iteration does three things for it:_update_requestsappends its in-flight token and logprob._pause_recompute_paused_requestscallsreset_for_recompute, which setsCONTEXT_INITandpy_decoding_iter = 0._handle_responsescallscreate_response()for it.The token itself was not sent, so after the re-prefill it is streamed together with the next token, but only the next token's logprob is left.
GenerationResultBase._handle_sequencethen failsassert len(output.logprobs) == output.length, and the request errors out. Withstream_interval > 1every logprob accumulated since the last emitted response is lost the same way, in the non-overlap loop as well.This change returns
Noneright aftercreate_serialized_result()when the result is empty, so the logprobs stay in the request until a response carries them. Responses that have a result are built exactly as before.Closes #22091.
Test Coverage
test_streaming_logprobs_survive_a_response_without_resultintests/unittest/_torch/pyexecutor/test_request_utils.py(cpu_only, covered by theunittest/_torch/pyexecutorentry ofl0_cpu). A streaming request with logprobs is recompute-paused after its token and logprob were added.create_response()then returnsNone, and after the re-prefill the next response must carry both tokens and both logprobs. The test fails onmain([[-0.125]] != [[-0.25, -0.125]]) and passes with this change.test_python_logprob_response_snapshots(4 passed), with--noconftestagainst a Python port ofLlmRequest::createResult,createSerializedResultandpause()fromllmRequest.cpp/llmRequest.h, with import stubs for the other bindings. In CI the test runs against the real bindings.PyExecutor._handle_responsesthrough the three steps above and passed its responses through pickle intoGenerationResult._handle_response. Onmainthe proxy raisesAssertionError: logprobs length: 2 != output.length: 3; with this change it receives token_ids[10, 11, 12]with logprobs[-0.1, -0.2, -0.3].pre-commit run --fileson both files: all hooks passed; Pyrefly reported no diagnostics.test_request_utils.py, GPU tests, and an end-to-end serving run under KV cache pressure.PR Checklist
AI assistance (Cursor) was used for this change.