Skip to content

[feat] Add TRTLLM MoE nvfp4 cubins for mid-high concurrency; attention_dp for TRTLLM MoE #5723

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.

Sign up for GitHub

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

Jump to bottom
Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

[feat] Add TRTLLM MoE nvfp4 cubins for mid-high concurrency; attention_dp for TRTLLM MoE #5723

wants to merge 6 commits into from
+3,977 −2,056

Conversation

rosenrodt
Copy link
Collaborator

@rosenrodt rosenrodt commented Jul 3, 2025 •
edited
Loading

[feat] Add TRTLLM MoE nvfp4 cubins for mid-high concurrency

Description

  • Update TRTLLM MoE to have better perf or at least on par with CUTLASS MoE at mid-high concurrency (64 tokens per expert).
    • For num_experts=128, top_k=8: 64 tokens/expert means concurrency of 1024 if unrealistically assuming round-robin token distribution
  • Update default heurstics of the kernel selection
  • AttentionDP for DeepSeek and Qwen3 (pytorch backend)

Test Coverage

accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=2-fp8kv=False-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4[moe_backend=TRTLLM-mtp_nextn=2-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=0-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=0-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False]
accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_trtllm]
accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_cutlass]
accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_trtllm]
accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass]
...

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.

@rosenrodt rosenrodt requested a review from a team as a code owner July 3, 2025 16:31
@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch from ef5ea41 to 606ae9d Compare July 3, 2025 16:34
@rosenrodt rosenrodt requested a review from nekorobov July 3, 2025 16:34
@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch from 606ae9d to 8b0eebc Compare July 4, 2025 04:35
@rosenrodt rosenrodt changed the title [feat] Add TRTLLM MoE nvfp4 cubins for mid-high concurrency [feat] [https://nvbugs/5369010] Add TRTLLM MoE nvfp4 cubins for mid-high concurrency Jul 4, 2025
@rosenrodt
Copy link
Collaborator Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10929 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

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

@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch from 105dfad to 3945b94 Compare July 4, 2025 09:10
@rosenrodt
Copy link
Collaborator Author

@nekorobov Let's wait for #5743 to be merged first

@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch from 3945b94 to 9969e33 Compare July 4, 2025 09:54
@rosenrodt
Copy link
Collaborator Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #10989 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

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

@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch 2 times, most recently from 1171165 to 626add3 Compare July 7, 2025 01:35
@rosenrodt rosenrodt changed the title [feat] [https://nvbugs/5369010] Add TRTLLM MoE nvfp4 cubins for mid-high concurrency [feat] Add TRTLLM MoE nvfp4 cubins for mid-high concurrency; attention_dp for TRTLLM MoE Jul 7, 2025
@rosenrodt
Copy link
Collaborator Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11077 [ run ] triggered by Bot

@rosenrodt
Copy link
Collaborator Author

/bot run

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11078 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11077 [ run ] completed with state ABORTED

@tensorrt-cicd
Copy link
Collaborator

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

@rosenrodt
Copy link
Collaborator Author

@nekorobov @byshiue PR is ready for re-review. I pushed updates earlier today to enable TRTLLM MoE + attention DP.

Also, since I cherry-picked #5743 to avoid intermittent failures, I prefer to wait until #5743 is merged.

nekorobov
nekorobov approved these changes Jul 7, 2025
View reviewed changes
@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch from 1b301e9 to 89c2d2c Compare July 7, 2025 10:08
@rosenrodt
Copy link
Collaborator Author

/bot run --stage-list "DGX_B200-4_GPUs-PyTorch-Post-Merge-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11140 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11140 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8236 (Partly Tested) completed with status: 'SUCCESS'

@rosenrodt
Copy link
Collaborator Author

The partial test for newly added cases passed. The previous commit completed full test with success. PR is in good shape.

[2025-07-07T11:12:17.275Z] DGX_B200-4_GPUs-PyTorch-Post-Merge-1/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_cutlass] PASSED [ 64%]
[2025-07-07T11:14:08.816Z] DGX_B200-4_GPUs-PyTorch-Post-Merge-1/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_trtllm] PASSED [ 67%]
[2025-07-07T11:16:15.323Z] DGX_B200-4_GPUs-PyTorch-Post-Merge-1/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_cutlass] PASSED [ 70%]
[2025-07-07T11:18:09.403Z] DGX_B200-4_GPUs-PyTorch-Post-Merge-1/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_trtllm] PASSED [ 72%]

@rosenrodt
Copy link
Collaborator Author

/bot run

@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch 4 times, most recently from 419a0fc to f94a32b Compare July 8, 2025 08:12
@tensorrt-cicd
Copy link
Collaborator

PR_Github #11255 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: Cherry picked from #5816 to avoid build failures on ToT

@rosenrodt
Copy link
Collaborator Author

/bot run --extra-stage "DGX_B200-4_GPUs-PyTorch-Post-Merge-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11256 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11258 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11256 [ run ] completed with state ABORTED

rosenrodt and others added 6 commits July 8, 2025 23:15
@rosenrodt
update moe trtllm nvfp4 cubins
c1f555c 
- add trtllm moe nvfp4 cubins for up to 64 tokens per expert
- update nvfp4 trtllm moe default heuristic
- TLLM_BATCHED_GEMM_PRINT_NAME=1 to print kernel name and problem shape

Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com>
@rosenrodt
enable attention_dp for TRTLLM MoE
09a6e89 
Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com>
@davidclark-nv @rosenrodt
feat: Adds optional module cache for TRT-LLM Gen Gemm interfaces
f462557 
Signed-off-by: David Clark <215764518+davidclark-nv@users.noreply.github.com>
Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com>
@rosenrodt
format
622be6f 
Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com>
@rosenrodt
add more test
e2389b2 
Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com>
@rosenrodt
fix deepseek-v3 TRTLLM MoE + attentionDP
7c07651 
Signed-off-by: Anthony Chang <27950904+rosenrodt@users.noreply.github.com>
@rosenrodt rosenrodt force-pushed the trtllm-moe-nvfp4-cubin-high-concurrency branch from f94a32b to 7c07651 Compare July 8, 2025 15:16
@rosenrodt
Copy link
Collaborator Author

/bot run --extra-stage "DGX_B200-4_GPUs-PyTorch-Post-Merge-1,GB200-4_GPUs-PyTorch-Post-Merge-1"

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11317 [ run ] triggered by Bot

@tensorrt-cicd
Copy link
Collaborator

PR_Github #11258 [ run ] completed with state ABORTED

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@nekorobov nekorobov nekorobov approved these changes

@suyoggupta suyoggupta Awaiting requested review from suyoggupta suyoggupta is a code owner automatically assigned from NVIDIA/trt-llm-torch-devs

@pcastonguay pcastonguay Awaiting requested review from pcastonguay pcastonguay is a code owner automatically assigned from NVIDIA/trt-llm-torch-devs

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@rosenrodt @tensorrt-cicd @nekorobov @davidclark-nv