Repository navigation
fix(inference): AUT-2539 skip prefill token limit for imported decode pr - #8259
Closed
svcnemo-autobot wants to merge 1 commit into
Closed
svcnemo-autobot wants to merge 1 commit into
svcnemo-autobot wants to merge 1 commit into
Conversation
_finalize_kv_handoff_import() admits an imported decode request through add_request(). Unless chunked prefill was enabled, _add_request() failed any prompt longer than context.max_tokens with TokenOverflowError. That check bounds per-step prefill work, but an imported prompt's KV cache already exists and the request goes straight to decode, so a long disaggregated prompt failed on the decode engine. Add a keyword-only prompt_is_prefilled flag to add_request() and _add_request(), forward it from both add_request() call sites, and skip only the max_tokens check when it is set. The finalizer passes True. The sequence-length and KV-pool checks still apply. The refit test's _add_request double now accepts the new keyword. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>
svcnvidia-nemo-ci
marked this pull request as draft
October 15, 2026 16:04
Contributor
|
This PR has been automatically converted to draft because all PRs must start as drafts. When you are ready for review, click Ready for Review to begin the review process. This will:
See the contribution guide for more details. |
Contributor
|
/ok to test 721672f |
svcnemo-autobot
marked this pull request as ready for review
October 15, 2026 17:32
This branch was previously deployed
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Background
In disaggregated serving,
_finalize_kv_handoff_import()admits an imported decode request throughadd_request(). Unless chunked prefill was enabled,_add_request()failed any prompt longer thancontext.max_tokenswithTokenOverflowError. That limit bounds how much prefill work one step does. An imported prompt's KV cache already exists and the request goes straight to decode, so a long disaggregated prompt failed on the decode engine.What changed
dynamic_engine.py:add_request()and_add_request()take a keyword-onlyprompt_is_prefilled: bool = False.add_request()forwards it from both of its_add_request()call sites (text and VLM). When the flag is set, only themax_tokenscheck is skipped. The max sequence length and KV-pool checks still apply.inference_state_handoff.py: the import finalizer passesprompt_is_prefilled=True.test_add_request_forwards_prompt_is_prefilledintest_dynamic_engine.py.test_disagg_handoff_lifecycle.py._add_requestdouble now accepts the new keyword.Details
The default is
False, so every other caller behaves as before. No subclass overrides_add_request. This was split out of the inference review and is independent of #8258.Tested
3914483f, using main's CI image (sha256:7b4e22a3):tests/unit_tests/inference/**/*.py: 3678 passed, 37 skipped, 0 failed.tests/unit_tests/inference/engines/test_dynamic_engine.py: 255 passed, 2 skipped.