feat: Add support for usage in the OpenAI frontend TRT-LLM backend#8326
Conversation
pskiran1
changed the title
usage in the OpenAI frontend TRT-LLM backend (#8264)usage in the OpenAI frontend TRT-LLM backend
There was a problem hiding this comment.
Pull Request Overview
This PR adds support for usage metrics in the OpenAI frontend when using the TRT-LLM backend. Previously, usage reporting (prompt_tokens, completion_tokens, total_tokens) was only available for the vLLM backend.
Key Changes:
- Enabled usage token reporting for TRT-LLM backend by requesting token counts in requests and parsing INT32 data types in responses
- Removed backend-specific checks that limited usage reporting to vLLM only
- Updated tests to expect usage metrics for all backends instead of conditionally based on backend type
Reviewed Changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| qa/L0_openai/test.sh | Added prompt_embedding_table_data_type parameter to TRT-LLM model configurations |
| python/openai/tests/test_openai_client.py | Removed backend parameter from test methods and conditional usage assertions |
| python/openai/tests/test_completions.py | Removed backend-specific usage checks from completion tests |
| python/openai/tests/test_chat_completions.py | Removed backend-specific usage checks from chat completion tests |
| python/openai/openai_frontend/engine/utils/triton.py | Added token count requests and INT32 data type handling for TRT-LLM responses |
| python/openai/openai_frontend/engine/triton_engine.py | Removed vLLM-only restriction for stream_options.include_usage |
| python/openai/README.md | Added placeholder for usage metrics documentation |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull Request Overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
|
||
| ### Helpers ### | ||
|
|
||
| function download_tensorrt_llm_models { |
There was a problem hiding this comment.
This is a temporary workaround for the test script downloading models from the latest TRT-LLM release, which includes a fix for the tensorrt_llm_bls model. I will create a ticket to track this issue and remove the workaround once the next stable TRT-LLM release is available.
What does the PR do?
This PR adds support for usage metrics in the OpenAI frontend when using the TRT-LLM backend. Previously, usage reporting (prompt_tokens, completion_tokens, total_tokens) was only available for the vLLM backend.
Key Changes:
Checklist
<commit_type>: <Title>Commit Type:
Check the conventional commit type
box here and add the label to the github PR.
Related PRs: #8264
Where should the reviewer start?
Test plan:
Caveats:
Background
By default, the OpenAI API supports
usagefor non-streaming requests. For streaming requests, we need to enablestream_options: {"include_usage": true}.https://platform.openai.com/docs/api-reference/chat/create
https://platform.openai.com/docs/api-reference/chat-streaming
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)