Skip to content

[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

Merged
merged 6 commits into from
Jul 17, 2025

Conversation

syuoni
Copy link
Collaborator

@syuoni syuoni commented Jul 14, 2025

[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:

  • Launch model forward iteration i (captured by CUDA graph)
  • Launch sampling iteration i
  • Sync sampling state iteration i-1

This PR changes the loop to:

  • Launch model forward iteration i (captured by CUDA graph)
  • Sync sampling state iteration i-1
  • [Optional] Guided decoding iteration i
    • Get token ID (iteration i-1) on CPU and call XGrammar to generate bitmask, asynchronously copy it to device
    • Launch logitsBitmaskKernel
  • Launch sampling iteration i

The key point is, we need to place "sync sampling state i-1" before "launch sampling iteration i", so that guided decoder can have the chance to get the correct token ID on CPU, generate the bitmask, and apply the bitmask before sampling i.

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.

image

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.

@syuoni syuoni requested review from lowsfer, QiJune and nvbrantz July 14, 2025 09:39
@syuoni syuoni requested a review from a team as a code owner July 14, 2025 09:39
@syuoni
Copy link
Collaborator Author

syuoni commented Jul 14, 2025

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11799 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11799 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8743 completed with status: 'SUCCESS'

@syuoni syuoni force-pushed the guided-with-overlap branch from a09f6b3 to ef064c3 Compare July 16, 2025 06:27
@syuoni
Copy link
Collaborator Author

syuoni commented Jul 16, 2025

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12035 [ run ] triggered by Bot

@syuoni syuoni requested a review from Funatiq July 16, 2025 06:42
@tensorrt-cicd
Copy link
Collaborator

PR_Github #12035 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8937 completed with status: 'FAILURE'

@syuoni
Copy link
Collaborator Author

syuoni commented Jul 16, 2025

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12076 [ run ] triggered by Bot

@Funatiq
Copy link
Collaborator

Funatiq commented Jul 16, 2025

The key point is, we need to place "sync sampling state i-1" before "launch sampling iteration i", so that guided decoder can have the chance to get the correct token ID and apply the bitmask before sampling i.

Can you help me understand why this is required? Launch model forward iteration i also gets the token ID (from previous_tensors_device) without requiring the sync. Can we also use the device tensors for the guided decoder?

@syuoni
Copy link
Collaborator Author

syuoni commented Jul 16, 2025

The key point is, we need to place "sync sampling state i-1" before "launch sampling iteration i", so that guided decoder can have the chance to get the correct token ID and apply the bitmask before sampling i.

Can you help me understand why this is required? Launch model forward iteration i also gets the token ID (from previous_tensors_device) without requiring the sync. Can we also use the device tensors for the guided decoder?

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.

syuoni added 6 commits July 16, 2025 12:33
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>
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
@syuoni syuoni force-pushed the guided-with-overlap branch from d7d7a72 to 4cb3d5a Compare July 16, 2025 12:34
@syuoni
Copy link
Collaborator Author

syuoni commented Jul 16, 2025

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12084 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12076 [ run ] completed with state ABORTED

Copy link
Collaborator

@nvbrantz nvbrantz left a 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!

@tensorrt-cicd
Copy link
Collaborator

PR_Github #12084 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8975 completed with status: 'SUCCESS'

Copy link
Collaborator

@QiJune QiJune left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@syuoni syuoni merged commit 21efb50 into NVIDIA:main Jul 17, 2025
3 checks passed
@syuoni syuoni changed the title [TRTLLM-6406] feat: Enable guided decoding with overlap scheduler [TRTLLM-6406, TRTLLM-5172] feat: Enable guided decoding with overlap scheduler Jul 17, 2025
reasonsolo pushed a commit to reasonsolo/TensorRT-LLM that referenced this pull request Jul 21, 2025
…IDIA#6000)

Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com>
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.

5 participants