-
Notifications
You must be signed in to change notification settings - Fork 1.6k
[TRTLLM-6406, TRTLLM-5172] feat: Enable guided decoding with overlap scheduler #6000
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
Conversation
/bot run |
PR_Github #11799 [ run ] triggered by Bot |
PR_Github #11799 [ run ] completed with state |
a09f6b3
to
ef064c3
Compare
/bot run |
PR_Github #12035 [ run ] triggered by Bot |
PR_Github #12035 [ run ] completed with state |
/bot run --disable-fail-fast |
PR_Github #12076 [ run ] triggered by Bot |
Can you help me understand why this is required? |
Hi @Funatiq , xgrammar (and also other guided decoding backends) is purely a cpu lib, it takes cpu token IDs as input, and generate cpu bitmask. Hence, the gpu token IDs don't work for xgrammar. |
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
d7d7a72
to
4cb3d5a
Compare
/bot run |
PR_Github #12084 [ run ] triggered by Bot |
PR_Github #12076 [ 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.
The execution flow looks correct: enqueue model inference 0 -> build mask 0 -> enqueue apply mask 0 -> enqueue sample 0 -> enqueue model inference 1 -> event sync sample 0 -> build mask 1 -> enqueue apply mask 1 -> enqueue sample 1 -> ...
Thanks!
PR_Github #12084 [ 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.
LGTM
…IDIA#6000) Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
[TRTLLM-6406] feat: Enable guided decoding with overlap scheduler
Description
This PR supports guided decoding with overlap scheduler.
Core changes
For overlap scheduler, the original loop is:
i
(captured by CUDA graph)i
i-1
This PR changes the loop to:
i
(captured by CUDA graph)i-1
i
i-1
) on CPU and call XGrammar to generate bitmask, asynchronously copy it to devicei
The key point is, we need to place "sync sampling state
i-1
" before "launch sampling iterationi
", so that guided decoder can have the chance to get the correct token ID on CPU, generate the bitmask, and apply the bitmask before samplingi
.Normally we can assume model forward is much heavier than sampling, so this change should not affect the effectiveness of overlap scheduler.
Results
Please see the nsys timeline, the overlap scheduler is enabled with guided decoding in the loop.
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.