-
Notifications
You must be signed in to change notification settings - Fork 1.8k
refactor: Enhanced handling of decoder requests and logits within the batch manager #6055
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
refactor: Enhanced handling of decoder requests and logits within the batch manager #6055
Conversation
Signed-off-by: Robin Kobus <[email protected]>
…ssor Signed-off-by: Robin Kobus <[email protected]>
… batch manager - Filter requests and store logits in HandleContextLogits and HandleGenerationLogits. - Refactored DecoderInputBuffers to hold decoder requests and logits in request order. - Updated GuidedDecoder and LogitsPostProcessor to utilize DecoderInputBuffers directly. - Adjusted related tests and bindings to reflect changes in the DecoderInputBuffers interface. This change streamlines the interaction between components and improves clarity in managing decoder inputs and outputs. Signed-off-by: Robin Kobus <[email protected]>
/bot run |
PR_Github #11929 [ run ] triggered by Bot |
PR_Github #11929 [ run ] completed with state |
/bot run --disable-fail-fast |
PR_Github #11953 [ run ] triggered by Bot |
PR_Github #11953 [ run ] completed with state |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
This PR refactors the decoder request handling and logits management within the batch manager to improve component interaction clarity. The changes centralize logits storage and request filtering within DecoderInputBuffers
, moving away from a slot-based approach to a direct request-order approach.
- Filter requests and store logits in
HandleContextLogits
andHandleGenerationLogits
functions - Refactor
DecoderInputBuffers
to hold decoder requests and logits in request order rather than by sequence slots - Update
GuidedDecoder
andLogitsPostProcessor
to utilizeDecoderInputBuffers
directly instead of separate request vectors
Reviewed Changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.
Show a summary per file
File | Description |
---|---|
tensorrt_llm/_torch/pyexecutor/sampler.py | Remove maxNumSequences parameter from DecoderInputBuffers constructor |
cpp/tests/runtime/gptDecoderBatchedTest.cpp | Update DecoderInputBuffers constructor calls in tests |
cpp/tests/batch_manager/guidedDecoderTest.cpp | Update test to use new DecoderInputBuffers interface and fix batch index references |
cpp/tensorrt_llm/pybind/batch_manager/bindings.cpp | Update Python bindings for DecoderInputBuffers constructor and add decoderRequests field |
cpp/tensorrt_llm/pybind/batch_manager/algorithms.cpp | Update algorithm function signatures to use DecoderInputBuffers instead of separate request vectors |
cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp | Update to use new DecoderInputBuffers interface and pass buffers directly to components |
cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp | Refactor to process requests from DecoderInputBuffers and improve logits handling logic |
cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp | Update to work with DecoderInputBuffers and use stream instead of runtime |
cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp | Add logic to populate decoderRequests and logits vectors in DecoderInputBuffers |
cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp | Add logic to populate decoderRequests and logits vectors for context requests |
cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp | Update to use DecoderInputBuffers interface and process requests in order |
cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp | Remove maxNumSequences parameter and logits vector initialization |
cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h | Update function signature to use DecoderInputBuffers |
cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h | Update interface to use DecoderInputBuffers and stream pointer |
cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h | Update execute method to use DecoderInputBuffers |
cpp/include/tensorrt_llm/batch_manager/decoderBuffers.h | Add decoderRequests field and update constructor signature |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM but you probably want to talk to @Linda-Stadter about the bindings so we can coordinate with her PR #6149 and #5961
You can merge it now and I can include the changes |
… batch manager (NVIDIA#6055) Signed-off-by: Robin Kobus <[email protected]>
… batch manager (NVIDIA#6055) Signed-off-by: Robin Kobus <[email protected]>
… batch manager (NVIDIA#6055) Signed-off-by: Robin Kobus <[email protected]> Signed-off-by: Shreyas Misra <[email protected]>
Description
This change streamlines the interaction between components and improves clarity in managing decoder inputs and outputs.
Test Coverage
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...
Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]
to print this help message.See details below for each supported subcommand.
run [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]
Launch build/test pipelines. All previously running jobs will be killed.
--disable-fail-fast
(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test
(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-1, xxx"
(OPTIONAL) : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"
(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--only-multi-gpu-test
(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test
(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test
(OPTIONAL) : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.--post-merge
(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-[Post-Merge]-1, xxx"
(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.md
.kill
kill
Kill all running builds associated with pull request.
skip
skip --comment COMMENT
Skip testing for latest commit on pull request.
--comment "Reason for skipping build/test"
is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.