Skip to content

[None][perf] Use bf16 custom router GEMM kernel for DeepSeek-V4#13646

Merged
hyukn merged 2 commits into
NVIDIA:feat/deepseek_v4from
hyukn:perf/dsv4-router-gemm-opt
May 7, 2026
Merged

[None][perf] Use bf16 custom router GEMM kernel for DeepSeek-V4#13646
hyukn merged 2 commits into
NVIDIA:feat/deepseek_v4from
hyukn:perf/dsv4-router-gemm-opt

Conversation

@hyukn
Copy link
Copy Markdown
Collaborator

@hyukn hyukn commented Apr 30, 2026 •
edited
Loading

Extend the existing dsv3 router GEMM custom kernel to support hidden_dim=4096 (DeepSeek-V4) in addition to 7168 (V3). For batch<=16, the custom kernel avoids cuBLAS overhead (~3us vs ~8us). For larger batches, cuBLAS fallback is preserved.

This eliminates the unnecessary fp32 cast in the router gate forward pass, yielding 3-10% end-to-end throughput improvement on B200.

@coderabbitai summary

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

@hyukn
[None][perf] Use bf16 custom router GEMM kernel for DeepSeek-V4
18a1270 
Extend the existing dsv3 router GEMM custom kernel to support
hidden_dim=4096 (DeepSeek-V4) in addition to 7168 (V3). For batch<=16,
the custom kernel avoids cuBLAS overhead (~3us vs ~8us). For larger
batches, cuBLAS fallback is preserved.

This eliminates the unnecessary fp32 cast in the router gate forward
pass, yielding 3-10% end-to-end throughput improvement on B200.

Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
@hyukn hyukn requested a review from a team as a code owner April 30, 2026 05:21
@hyukn hyukn requested review from 2ez4bz and removed request for a team April 30, 2026 05:21
2ez4bz
2ez4bz approved these changes May 4, 2026
View reviewed changes
Comment thread tensorrt_llm/_torch/models/modeling_deepseekv4.py
Comment thread tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py
@lfr-0531
Copy link
Copy Markdown
Collaborator

lfr-0531 commented May 5, 2026

/bot run

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 5, 2026

PR_Github #46830 [ run ] triggered by Bot. Commit: 18a1270 Link to invocation

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 5, 2026

PR_Github #46830 [ run ] completed with state FAILURE. Commit: 18a1270
/LLM/main/L0_MergeRequest_PR pipeline #36849 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@peihu-nv
Copy link
Copy Markdown
Collaborator

peihu-nv commented May 5, 2026

/bot run

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 5, 2026

PR_Github #46843 [ run ] triggered by Bot. Commit: 18a1270 Link to invocation

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 5, 2026

PR_Github #46843 [ run ] completed with state FAILURE. Commit: 18a1270
/LLM/main/L0_MergeRequest_PR pipeline #36861 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@peihu-nv
Copy link
Copy Markdown
Collaborator

peihu-nv commented May 5, 2026

/bot run

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 5, 2026

PR_Github #46851 [ run ] triggered by Bot. Commit: 591f210 Link to invocation

peihu-nv
peihu-nv approved these changes May 5, 2026
View reviewed changes
Copy link
Copy Markdown
Collaborator

@peihu-nv peihu-nv left a comment

Choose a reason for hiding this comment

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

Tested both the accuracy and perf on B200 and GB200 MTP0/MTP1. All looking good.

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 6, 2026

PR_Github #46851 [ run ] completed with state SUCCESS. Commit: 591f210
/LLM/main/L0_MergeRequest_PR pipeline #36866 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@hyukn
Copy link
Copy Markdown
Collaborator Author

hyukn commented May 7, 2026

/bot run --disable-fail-fast

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 7, 2026

PR_Github #47056 [ run ] triggered by Bot. Commit: 591f210 Link to invocation

@tensorrt-cicd
Copy link
Copy Markdown
Collaborator

tensorrt-cicd commented May 7, 2026

PR_Github #47056 [ run ] completed with state SUCCESS. Commit: 591f210
/LLM/main/L0_MergeRequest_PR pipeline #37029 completed with status: 'SUCCESS'

CI Report

Link to invocation

@hyukn hyukn merged commit a3f3775 into NVIDIA:feat/deepseek_v4 May 7, 2026
6 of 7 checks passed
lfr-0531 pushed a commit that referenced this pull request May 7, 2026
@hyukn @peihu-nv @heyuhhh
[None][perf] Use bf16 custom router GEMM kernel for DeepSeek-V4 (#13646)
efbccca 
Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
Co-authored-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
(cherry picked from commit a3f3775)
Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com>
lfr-0531 pushed a commit that referenced this pull request May 14, 2026
@hyukn @peihu-nv @heyuhhh
[None][perf] Use bf16 custom router GEMM kernel for DeepSeek-V4 (#13646)
aa67e3e 
Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
Co-authored-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
(cherry picked from commit a3f3775)
Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
lfr-0531 pushed a commit to lfr-0531/TensorRT-LLM that referenced this pull request May 29, 2026
@hyukn @peihu-nv @lfr-0531
[None][perf] Use bf16 custom router GEMM kernel for DeepSeek-V4 (NVID…
0f21239 
…IA#13646)

Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
Co-authored-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
(cherry picked from commit a3f3775)
Signed-off-by: Yuhang He <58161490+heyuhhh@users.noreply.github.com>
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
(cherry picked from commit aa67e3e)
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@2ez4bz 2ez4bz 2ez4bz approved these changes

@peihu-nv peihu-nv peihu-nv approved these changes

Assignees

@hyukn hyukn

Labels

deepseek-v4

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@hyukn @lfr-0531 @tensorrt-cicd @peihu-nv @2ez4bz