Repository navigation
[https://nvbugs/6858505][fix] Reject an encoder input longer than one encoder iteration at admission - #21199
[https://nvbugs/6858505][fix] Reject an encoder input longer than one encoder iteration at admission#21199trtllm-agent wants to merge 1 commit into
Conversation
… encoder iteration at admission An encoder input is never chunked: it runs whole in one iteration, which is bounded by max_num_tokens (the micro-batch schedulers), encoder_max_num_tokens (the encoder step) and the per-sequence capacity of the cross-attention KV cache that holds the encoder output. Admission did not compare the encoder length with these bounds, so a longer request was admitted and then failed inside the executor loop. The micro-batch schedulers raised for it (the Python assertion, the C++ check, the KV cache V2 scheduler's deadlock error), the KV cache V2 manager raised while resizing the cross KV cache, or the forward or encoder step raised for the whole batch and the error counted against the executor's error budget. PyExecutor._validate_request_budget now rejects such a request with a ValueError that names the encoder length and each bound. The request fails on its own, without counting against the error budget, and the executor keeps serving the other requests. Requests within the bounds and decoder-only requests are admitted as before. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
|
No actionable comments were generated in the recent review.
|
| Layer / File(s) | Summary |
|---|---|
Encoder output length validation tensorrt_llm/_torch/pyexecutor/py_executor.py |
_validate_request_budget checks encoder output length before KV-cache block budgets. It rejects lengths above the smallest applicable nonzero limit among max_num_tokens, encoder_max_num_tokens, and cross-KV cache max_seq_len. |
Unit coverage for admission and limits tests/unittest/_torch/executor/test_resource_manager.py, tests/unittest/_torch/executor/test_py_executor.py |
Tests cover configured limits, cross-KV cache sequence length, unknown encoder output length, boundary lengths, and rejected requests without budget charging. |
T5 integration coverage tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py, tests/integration/test_lists/test-db/l0_h100.yml |
T5 tests cover oversized encoder inputs across three KV manager and scheduler configurations, accepted boundary inputs, and reference output. The H100 test list includes the cases. |
Priority:
Estimated code review effort: 2 (Simple) | ~12 minutes
Change: Bug fix
Suggested reviewers: litaotju, lori-ren, brnguyen2
Merge Risk: 
Minimal · up to 5f41c
-
Oversized encoder requests are checked before scheduling. No merge-blocking risk was identified.
Pre-merge checks: 
4 · 
1
-
Failed checks: 1 inconclusive Passed checks: 4 passed-
Full details: Description check
-
Explanation
The description clearly explains the problem, implementation, expected behavior, and extensive unit and integration test coverage. It omits the PR Checklist section, but the required technical information is otherwise complete.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 1 functions across 3 files. (2 skipped: 1 unsupported, 1 too large.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The check sits at the right point: an encoder input that can never fit now fails only its own request at admission, and doesn't reach the scheduler or forward step. Decoder-only and VLM requests return None from try_get_encoder_output_len(), so they skip it. The tests cover each bound plus the "error budget not charged" path.
One gap for a follow-up. The check bounds one request, but encoder_max_num_tokens also caps the packed encoder batch (encoder_decoder.py:173 / :205 compare the summed length). When a user sets encoder_max_num_tokens < max_num_tokens, the micro-batch scheduler budgets only against max_num_tokens. It can then co-schedule several encoder requests that each pass this check but whose sum is above encoder_max_num_tokens, and the encoder step fails the whole batch. I couldn't find an enc-dec scheduler bound on this (_apply_mm_encoder_admission covers only MM item scheduling). Please either add an encoder token budget to the enc-dec scheduling path, or say in the description that only the single-request case is covered.
Design note, non-blocking: this method keeps its own copy of the limits that the scheduler, encoder runner, and cross KV manager each enforce. If one of them changes its bound, admission will fall out of sync again. A small accessor on each component (or one helper that the scheduler assert and this check both call) would keep them in sync.
| with pytest.raises( | ||
| RequestError, | ||
| match=( | ||
| rf"encoder tokens \({len(too_long_ids)}\) exceeds the limit value " |
There was a problem hiding this comment.
max_input_len, max_num_tokens, and max_seq_len are all _MAX_SEQUENCE_LENGTH here, so the cross KV cache max_seq_len is also 64. Any of the three bounds could produce this error, so the test doesn't confirm that the test name's "above max_num_tokens" case is the one that fired. Optional: add max_num_tokens=64 to the regex (the details list prints each bound), or raise max_seq_len so max_num_tokens is the only limit at 64.
|
/bot run |
|
PR_Github #81669 [ run ] triggered by Bot. Commit: |
|
PR_Github #81669 [ run ] completed with state
|
Description
An encoder-decoder request (T5, BART, mBART, Whisper) runs its whole encoder input in one iteration: the encoder input is never chunked. That iteration is bounded by
max_num_tokens(the micro-batch schedulers), byencoder_max_num_tokens(the encoder step), and by the per-sequence capacity of the cross-attention KV cache that holds the encoder output (max_seq_lenof the cross KV cache manager).Before this change, request admission did not compare the encoder length with these bounds. A request with a longer encoder input was admitted and then failed later:
max_num_tokens, the Python micro-batch scheduler raised anAssertionError, the C++ micro-batch scheduler raised a check failure, and the KV cache V2 scheduler left the request unscheduled until it raised its scheduler-deadlock error. Each of these errors ended the executor loop.max_seq_lenbut withinmax_num_tokens, the attention metadata check raised in the forward step, and the error counted against the executor's error budget. With the KV cache V2 manager, a longer input instead made the manager raise while resizing the cross KV cache, which ended the executor loop.encoder_max_num_tokens, when it is set belowmax_num_tokens, the encoder step raised for the whole encoder batch, and the error counted against the error budget.With this change,
PyExecutor._validate_request_budgetrejects, at admission, a request whose encoder output length exceeds the smallest of these bounds. TheValueErrornames the length and each bound. The request fails on its own, like a request rejected by the existing KV block budget check, without counting against the error budget, and the executor keeps serving the other requests. Requests within the bounds, and decoder-only requests, are admitted as before.Test Coverage
tests/unittest/_torch/executor/test_resource_manager.py: the encoder length is checked againstmax_num_tokens, against a smallerencoder_max_num_tokens, and against the cross KV cachemax_seq_len; decoder-only prompt length is not checked here.tests/unittest/_torch/executor/test_py_executor.py::test_encoder_input_above_max_num_tokens_fails_only_that_request: through_fetch_and_activate_new_requests, only the request over the limit fails, and the error budget is not charged.tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py::test_t5_encoder_input_above_max_num_tokens_fails_alone(KV cache V2, V1 with the Python scheduler, V1 with the C++ scheduler; added tol0_h100.yml): an encoder input one token over the limit fails with the new error, then the longest encoder input that fits and the reference prompt are served.Ran on H100: the new unit tests and all three integration cases fail without the source change and pass with it. With the change, these also pass:
test_resource_manager.py,test_py_executor.py,test_py_scheduler.py,test_encoder_executor.py,test_kv_connector_initialization.pyandtest_request_utils.pyintests/unittest/_torch/executor/, plusengine/test_encoder_decoder.py,kv_cache/test_dual_pool_kv_cache.py,kv_cache/test_kv_cache_v2_scheduler.py,tests/unittest/disaggregated/test_chunked_transfer.pyandtests/unittest/executor/test_forward_error_survival.py;test_t5_pytorch_generate_encoder_decoder_end_to_end(bf16 t5-small greedy, beam search, greedy with the overlap scheduler),test_t5_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch(4 cases),test_t5_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs,test_t5_pytorch_block_reuse_is_disabled_for_v2andtest_t5_v2_decode_beyond_encoder_length(2 cases);test_bart_pytorch_generate_encoder_decoder_end_to_end(7 single-GPU bart-large-cnn cases),test_bart_pytorch_generate_encoder_decoder_mixed_encoder_lengths_batch,test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs[tp1]andtest_mbart_pytorch_generate_encoder_decoder_end_to_end;test_whisper_pytorch_transcribe_end_to_end,test_whisper_pytorch_block_reuse_requested,test_whisper_pytorch_beam_search(2 cases) andtest_whisper_pytorch_feature_combinations(5 single-GPU cases).Dev Engineer Review
PyExecutor._validate_request_budgetnow rejects encoder-decoder requests when the known encoder output length exceeds an applicable token limit or the cross-KV cache’s per-sequence limit. This prevents requests that cannot fit in one encoder iteration from proceeding to scheduling. Decoder-only requests are not subject to this encoder-length check.QA Engineer Review
The unit tests cover each configured encoder-length limit, boundary-sized inputs, decoder-only requests, and rejection without charging the executor error budget. The T5 integration test covers KV cache V2 and KV cache V1 with Python and C++ schedulers. Its H100 CI selectors are listed in
test-db/l0_h100.yml. Coverage verdict: sufficient for the stated admission behavior. Test execution results were not independently verified.Per-File QA Perspective
tensorrt_llm/_torch/pyexecutor/py_executor.py— Verify that oversized encoder outputs are rejected against the applicable token and cross-KV sequence-length limits. Also verify that decoder-only requests and requests without a known encoder output length retain their existing behavior.tests/integration/defs/llmapi/test_llm_api_pytorch_t5.py— The added integration test checks the rejection error and subsequent service of a boundary-sized request. Its H100 CI selectors are listed intests/integration/test_lists/test-db/l0_h100.yml.tests/integration/test_lists/test-db/l0_h100.yml— Adds the T5 oversized-encoder-input cases for KV cache V2 and KV cache V1 with Python and C++ schedulers to the H100 CI list.tests/unittest/_torch/executor/test_py_executor.py— Covers rejection of only the oversized encoder request and verifies that the rejection does not charge the executor error budget. This unit test is not an integration test-list entry.tests/unittest/_torch/executor/test_resource_manager.py— Covers the encoder token limits, cross-KVmax_seq_len, inputs at the limits, and acceptance of a long decoder prompt without encoder output. This unit test is not an integration test-list entry.