[TRTLLM-4983] feat: enable overlap scheduler between draft forwards

[lfr-0531]

Description

• Enable the overlap scheduler between different draft forwards.
• Add a new Eagle3ResourceManager to manage the hidden states, and remove the extra model input.
• Move the eagle3 fc into the model forward.
• Move the h2d copy to the end of the _prepare_tp_inputs to hide the CPU time.
• Disable CUDA graph for the 1st draft forward.

I'll collect more performance and accuracy data and update the results here.

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #7897 [ run ] triggered by Bot

[tensorrt-cicd]

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

[mikeiovine]

Shall we measure the performance gain? Since our initial goal is to hit parity with vLLM, I think using Llama 3.3 70B on Hopper makes sense. I got these numbers on 8xH200 today.

Framework	Max Batch Size	OSL	Eagle (draft len = 3)?	Output tok/sec
TRTLLM	8	256	No	876
TRTLLM	8	256	Yes	1154
vLLM	8	256	No	715
vLLM	8	256	Yes	1430

A note about the dataset: I used gsm8k from here. For the llama 3.3 eagle drafters we have from the paper author's, you have to apply the tokenizer's chat template or the AR will drop significantly and performance will be regressed. I will send you a preprocessed dataset that is compatible with trtllm-bench.

[lfr-0531]

/bot run

[tensorrt-cicd]

PR_Github #8288 [ run ] triggered by Bot

[tensorrt-cicd]

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

[lfr-0531]

/bot run

[tensorrt-cicd]

PR_Github #8319 [ run ] triggered by Bot

[lfr-0531]

/bot kill

[tensorrt-cicd]

PR_Github #8377 [ kill ] triggered by Bot

[tensorrt-cicd]

PR_Github #8802 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #8802 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #6390 completed with status: 'FAILURE'

[mikeiovine]

@lfr-0531 That plan sounds good to me. It can still be useful for other use case (@IzzyPutterman's draft/target speculative decoding should land soon, it'll be useful there since the draft models are a bit bigger than EAGLE).

[IzzyPutterman]

Do we already have logic to make sure that extra pass of overlap doesnt run (and infringe on the verify pass)? I think I have an internal MR that does this from a while ago.

[IzzyPutterman]

Something like this: #5211

[lfr-0531]

/bot run

[tensorrt-cicd]

PR_Github #8907 [ run ] triggered by Bot

[tensorrt-cicd]

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

[lfr-0531]

/bot run

[tensorrt-cicd]

PR_Github #8925 [ run ] triggered by Bot

[lfr-0531]

Do we already have logic to make sure that extra pass of overlap doesnt run (and infringe on the verify pass)? I think I have an internal MR that does this from a while ago.

We didn't add it in this PR. This PR is only used to overlap the different forwards in the same iteration. After we have PR-5211, the last iteration, including the _prepare_draft_tokens, can be skipped.

[tensorrt-cicd]

PR_Github #8925 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #6509 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

[jhaotingc]

Hi @lfr-0531,
So, is overlap scheduler supported in Eagle3 Two Model yet?
support_overlap_scheduler
cc @mikeiovine

Thanks!