Skip to content

fix(inference): AUT-2539 skip prefill token limit for imported decode pr - #8259

Closed
svcnemo-autobot wants to merge 1 commit into
NVIDIA:mainfrom
svcnemo-autobot:fix/inference-prefilled-prompt-token-limit
Closed

svcnemo-autobot wants to merge 1 commit into
NVIDIA:mainfrom
svcnemo-autobot:fix/inference-prefilled-prompt-token-limit

Conversation

@svcnemo-autobot

@svcnemo-autobot svcnemo-autobot commented Oct 15, 2026 •

Copy link
Copy Markdown
Collaborator

Background

In disaggregated serving, _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 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-only prompt_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 the max_tokens check is skipped. The max sequence length and KV-pool checks still apply.
  • inference_state_handoff.py: the import finalizer passes prompt_is_prefilled=True.
  • Tests:
    • test_add_request_forwards_prompt_is_prefilled in test_dynamic_engine.py.
    • Parametrized finalizer and limit tests in test_disagg_handoff_lifecycle.py.
    • The refit test's _add_request double 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

  • black, isort, pylint (10.00/10) and ruff are clean on the 4 changed files.
  • Run on 8×H100 at this head merged with main 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.
  • No functional test case exercises disaggregation.

_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>
@svcnemo-autobot
svcnemo-autobot requested review from a team as code owners October 15, 2026 16:04
@copy-pr-bot

copy-pr-bot Bot commented Oct 15, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@svcnvidia-nemo-ci
svcnvidia-nemo-ci marked this pull request as draft October 15, 2026 16:04
@github-actions

Copy link
Copy Markdown
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:

  1. Add the oncall reviewer (optional reviewer)
  2. Add required review teams based on your changes

See the contribution guide for more details.

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

/ok to test 721672f

@svcnemo-autobot svcnemo-autobot changed the title fix(inference): AUT-2539 skip prefill token limit for imported decode prompts fix(inference): AUT-2539 skip prefill token limit for imported decode pr Oct 15, 2026
@svcnemo-autobot
svcnemo-autobot marked this pull request as ready for review October 15, 2026 17:32
@ko3n1g ko3n1g closed this Oct 17, 2026

This branch was previously deployed

2 inactive deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants