[TRTLLM-6406, TRTLLM-5172] feat: Enable guided decoding with overlap scheduler

[syuoni]

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

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]

/bot run

[tensorrt-cicd]

PR_Github #11799 [ run ] triggered by Bot

[tensorrt-cicd]

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

[syuoni]

/bot run

[tensorrt-cicd]

PR_Github #12035 [ run ] triggered by Bot

[tensorrt-cicd]

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

[syuoni]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #12076 [ run ] triggered by Bot

[Funatiq]

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]

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]

/bot run

[tensorrt-cicd]

PR_Github #12084 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #12076 [ run ] completed with state ABORTED

[nvbrantz]

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]

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

[QiJune]

LGTM