Skip to content

Mtp optimizations round1 - #5689

Merged
Kefeng-Duan merged 27 commits into
NVIDIA:mainfrom
ameynaik-hub:mtp_optimizations_round1
Jul 25, 2025
Merged

Mtp optimizations round1#5689
Kefeng-Duan merged 27 commits into
NVIDIA:mainfrom
ameynaik-hub:mtp_optimizations_round1

Conversation

@ameynaik-hub

@ameynaik-hub ameynaik-hub commented Jul 2, 2025 •
edited by coderabbitai Bot
Loading

Copy link
Copy Markdown
Collaborator

Deepseek MTP optimization.

Optimize mtp head for deepseek

  1. eh_proj gemm sharding
  2. rmsnorm cuda stream parallelization.
  3. lm_head top1 using local argmax --> AG --> argmax
  4. piecewise torch.compile to fuse / opt multiple native kernels.

Test Coverage

BS1 deepseek perf improves from 411-->431 TPS on B200x8 TP8EP1.

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.

Details

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.

Summary by CodeRabbit

  • New Features

    • Improved performance and efficiency for token state extraction and normalization in DeepseekV3 models.
    • Enhanced support for distributed and parallel token selection in speculative decoding with tensor parallelism.
    • Added new helper functions for optimized token sampling and draft token management.

  • Refactor

    • Modularized and compiled key tensor operations for better clarity and speed.
    • Updated constructors and initialization logic to consistently use model configuration context across relevant components.

  • Chores

    • Adjusted internal interfaces to support expanded configuration and parallelism features.

@ameynaik-hub
ameynaik-hub requested a review from a team as a code owner July 2, 2025 23:40
@ameynaik-hub
ameynaik-hub requested review from byshiue and hlu1 July 2, 2025 23:40
@svc-trtllm-gh-bot svc-trtllm-gh-bot added the Community want to contribute PRs initiated from Community label Jul 3, 2025
@ameynaik-hub
ameynaik-hub force-pushed the mtp_optimizations_round1 branch from aba2520 to 2bcb3ac Compare July 7, 2025 18:49
@ameynaik-hub

ameynaik-hub commented Jul 7, 2025

Copy link
Copy Markdown
Collaborator Author

/bot run

@FrankD412
FrankD412 requested a review from Kefeng-Duan July 7, 2025 19:02
@Kefeng-Duan

Kefeng-Duan commented Jul 9, 2025

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 9, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #11430 [ run ] triggered by Bot

@Kefeng-Duan
Kefeng-Duan requested review from lfr-0531 and yweng0828 July 9, 2025 12:35
@Kefeng-Duan

Kefeng-Duan commented Jul 9, 2025

Copy link
Copy Markdown
Collaborator

@lfr-0531 @yweng0828 for viz

@tensorrt-cicd

tensorrt-cicd commented Jul 9, 2025

Copy link
Copy Markdown
Collaborator

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

@ameynaik-hub

ameynaik-hub commented Jul 9, 2025

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 9, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #11459 [ run ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 9, 2025

Copy link
Copy Markdown
Collaborator

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

@ameynaik-hub

ameynaik-hub commented Jul 10, 2025

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@ameynaik-hub
ameynaik-hub force-pushed the mtp_optimizations_round1 branch from 30817d5 to 8d3981d Compare July 11, 2025 23:01
@ameynaik-hub
mtp round1 optimizations without AR+norm fusion in mtp head
f9fc02b 
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
@ameynaik-hub
ameynaik-hub force-pushed the mtp_optimizations_round1 branch from 8d3981d to f9fc02b Compare July 12, 2025 00:09
@Kefeng-Duan
Kefeng-Duan requested a review from hello-11 July 15, 2025 00:36
@ameynaik-hub
:Merge branch 'NVIDIA:main' into mtp_optimizations_round1
bca1a57 
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
@ameynaik-hub
ameynaik-hub force-pushed the mtp_optimizations_round1 branch from ffd8259 to bca1a57 Compare July 15, 2025 06:41
@ameynaik-hub

ameynaik-hub commented Jul 15, 2025

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 15, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #11905 [ run ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 15, 2025

Copy link
Copy Markdown
Collaborator

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

@ameynaik-hub

ameynaik-hub commented Jul 15, 2025

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@Kefeng-Duan

Kefeng-Duan commented Jul 23, 2025

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #12670 [ run ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 23, 2025

Copy link
Copy Markdown
Collaborator

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

@hello-11

hello-11 commented Jul 24, 2025

Copy link
Copy Markdown
Contributor

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #12788 [ run ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

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

@Kefeng-Duan
Kefeng-Duan enabled auto-merge (squash) July 24, 2025 07:46
@Kefeng-Duan

Kefeng-Duan commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #12824 [ run ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

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

@Kefeng-Duan

Kefeng-Duan commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #12860 [ run ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 24, 2025

Copy link
Copy Markdown
Collaborator

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

@QiJune

QiJune commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator

/bot skip --comment "reuse previous passed CI"

@tensorrt-cicd

tensorrt-cicd commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #12959 [ skip ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #12959 [ skip ] completed with state ABORTED

@Kefeng-Duan

Kefeng-Duan commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #12991 [ run ] triggered by Bot

@ameynaik-hub

ameynaik-hub commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator

PR_Github #13014 [ run ] triggered by Bot

@tensorrt-cicd

tensorrt-cicd commented Jul 25, 2025

Copy link
Copy Markdown
Collaborator

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

@Kefeng-Duan
Kefeng-Duan merged commit 1e5e71a into NVIDIA:main Jul 25, 2025
3 checks passed
NVShreyas pushed a commit to NVShreyas/TensorRT-LLM that referenced this pull request Jul 28, 2025
@ameynaik-hub @Kefeng-Duan @NVShreyas
Mtp optimizations round1 (NVIDIA#5689)
b9e58b7 
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
Co-authored-by: Kefeng-Duan <176893526+Kefeng-Duan@users.noreply.github.com>
Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>
Ransiki pushed a commit to Ransiki/TensorRT-LLM that referenced this pull request Jul 29, 2025
@ameynaik-hub @Kefeng-Duan @Ransiki
Mtp optimizations round1 (NVIDIA#5689)
e8e2a59 
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
Co-authored-by: Kefeng-Duan <176893526+Kefeng-Duan@users.noreply.github.com>
Signed-off-by: Ransiki Zhang <ransikiz@nvidia.com>
lancelly pushed a commit to lancelly/TensorRT-LLM that referenced this pull request Aug 6, 2025
@ameynaik-hub @Kefeng-Duan @lancelly
Mtp optimizations round1 (NVIDIA#5689)
4c788d1 
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
Co-authored-by: Kefeng-Duan <176893526+Kefeng-Duan@users.noreply.github.com>
Signed-off-by: Lanyu Liao <lancelly@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@coderabbitai coderabbitai[bot] coderabbitai[bot] left review comments
@QiJune QiJune QiJune approved these changes
@Kefeng-Duan Kefeng-Duan Kefeng-Duan approved these changes
@byshiue byshiue Awaiting requested review from byshiue byshiue was automatically assigned from NVIDIA/trt-llm-torch-devs
@hlu1 hlu1 Awaiting requested review from hlu1 hlu1 was automatically assigned from NVIDIA/trt-llm-torch-devs
@lfr-0531 lfr-0531 Awaiting requested review from lfr-0531
@yweng0828 yweng0828 Awaiting requested review from yweng0828
@dc3671 dc3671 Awaiting requested review from dc3671
+1 more reviewer
@hello-11 hello-11 hello-11 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

@hello-11 hello-11

@Kefeng-Duan Kefeng-Duan

@ameynaik-hub ameynaik-hub

Labels

Community want to contribute PRs initiated from Community

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

7 participants

@ameynaik-hub @Kefeng-Duan @tensorrt-cicd @hello-11 @QiJune @dc3671 @svc-trtllm-gh-bot