Skip to content

Respect TensorRT-LLM input limits during generation - #4408

Open
sdivyanshu90 wants to merge 1 commit into
EleutherAI:mainfrom
sdivyanshu90:fix-trtllm-generation-input-limit
Open

sdivyanshu90 wants to merge 1 commit into
EleutherAI:mainfrom
sdivyanshu90:fix-trtllm-generation-input-limit

Conversation

@sdivyanshu90

Copy link
Copy Markdown
Contributor

TRT-LLM generation currently borrows unused output capacity for prompts, so an engine configured with max_input_len=2048 can receive 2304 input tokens. Include the separate engine input cap in the truncation budget while retaining the total sequence cap, requested generation length, and truncation direction.

Add CPU tests using the real constructor and generate_until with a recording engine, covering all three truncation directions, short/boundary/long prompts, and per-request generation lengths. Include the mechanical typing, formatting, and exception-preserving corrections required by pinned lint checks on the touched backend.

Fixes #4406.

Validation:

  • Before the fix: 5 new regression cases failed; 10 controls passed.
  • python -m pytest -q tests/models/test_trtllm_input_limit.py tests/models/test_model_utils.py: 54 passed.
  • pre-commit run --files lm_eval/models/trtllm_causallms.py tests/models/test_trtllm_input_limit.py: all applicable pinned hooks passed.
  • The exact issue reproduction changes from 2304 to 2048 input tokens.
  • Actual TensorRT/GPU inference was not exercised.

@sdivyanshu90

Copy link
Copy Markdown
Contributor Author

Hi @baberabb, could you please review this PR whenever you have some time. Thanks!

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

TRT-LLM generation exceeds the configured max_input_len

1 participant