[https://nvbugs/6108841][fix] add hidden_dim=6144 router GEMM instantiation for GLM-5 - #13740
Merged
Merged
Conversation
yijingl-nvidia
changed the title
Contributor
📝 WalkthroughWalkthroughThis PR extends the DeepSeek V3 router GEMM operation to support an additional hidden dimension value of 6144 alongside the existing 7168. The change introduces explicit kernel template instantiations for the new dimension and updates the operator dispatch logic to conditionally select the appropriate kernel variant based on input hidden dimension. ChangesRouter GEMM Hidden Dimension Support
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp (1)
1-3:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUpdate the NVIDIA copyright header year in this modified file.
The header still ends at
2024, but this file is modified in this PR.📝 Suggested header update
- * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.As per coding guidelines: "Include NVIDIA copyright header on all new files; update year on modified files."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp` around lines 1 - 3, Update the file header year range in the SPDX copyright block to include 2025 (or the current year) since this file was modified; locate the top-of-file comment containing "SPDX-FileCopyrightText" and the "SPDX-License-Identifier" lines in dsv3RouterGemmOp.cpp and change the trailing year from 2024 to the current year so the header reflects the modification.cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu (1)
1-3:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUpdate the copyright year in this modified CUDA file.
The header currently ends at
2023, but this file is modified in this PR.📝 Suggested header update
- * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved.As per coding guidelines: "Include NVIDIA copyright header on all new files; update year on modified files."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu` around lines 1 - 3, Update the copyright header line that currently reads "Copyright (c) 2019-2023, NVIDIA CORPORATION." to include the current modification year (e.g., "Copyright (c) 2019-2026, NVIDIA CORPORATION.") so the header reflects this PR's changes; locate the header comment containing that exact string at the top of dsv3RouterGemm.cu and modify the end year accordingly.
🧹 Nitpick comments (1)
tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py (1)
13-13: ⚡ Quick winAdd one unsupported-hidden-size case to validate fallback dispatch.
This matrix now covers both supported custom-kernel shapes, which is great. Please add one unsupported
hidden_size(for example,4096) so this test also exercises the fallback path and catches future dispatch regressions.✅ Minimal test matrix extension
-@pytest.mark.parametrize("hidden_size", [7168, 6144]) +@pytest.mark.parametrize("hidden_size", [7168, 6144, 4096])Also, because this PR only touches unittest scope, QA integration test-list updates are unnecessary here.
As per coding guidelines: "Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py` at line 13, The parametrize matrix in tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py only covers supported hidden sizes and misses an unsupported case to exercise the fallback dispatch; update the pytest.mark.parametrize("hidden_size", ...) invocation (in test_dsv3_router_gemm.py) to include an unsupported hidden size such as 4096 so the test matrix contains [7168, 6144, 4096], thereby exercising the fallback path and catching dispatch regressions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu`:
- Around line 245-292: The PR added new explicit instantiations for
invokeRouterGemm<__nv_bfloat16, N, 256, 6144> (N=1..16) but lacks
perf-regression coverage for the 6144 hidden_dim path; update tests to include a
perf case exercising the new 6144 router-GEMM dispatch: add or modify a perf
test in tests/integration/defs/perf/test_perf_sanity.py to run a
small-but-representative workload that triggers invokeRouterGemm with
hidden_dim=6144 (matching the instantiated templates), then add this test to the
authoritative CI perf list tests/integration/test_lists/test-db/l0_perf.yml and
the corresponding QA perf list(s) under
tests/integration/test_lists/qa/llm_perf_*.yml so CI will catch regressions for
the new 6144 path. Ensure the test references the same configuration (data type
__nv_bfloat16 and block sizes) used by the template instantiations so the 6144
dispatch is exercised during perf runs.
In `@cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp`:
- Around line 89-103: shape_ok currently misses validating the inner dimension
of mat_b, so LoopUnroller::unroll may be dispatched with a compile-time K that
doesn't match mat_b; update the dispatch guard to also verify mat_b.sizes()[0]
== hidden_dim (or equivalent runtime check) before calling LoopUnroller for the
kHiddenDim7168 and kHiddenDim6144 paths (i.e. augment shape_ok or add an extra
condition using hidden_dim and mat_b.sizes()[0] prior to the
LoopUnroller<...>::unroll calls to prevent malformed inputs from reaching the
custom kernel).
---
Outside diff comments:
In `@cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu`:
- Around line 1-3: Update the copyright header line that currently reads
"Copyright (c) 2019-2023, NVIDIA CORPORATION." to include the current
modification year (e.g., "Copyright (c) 2019-2026, NVIDIA CORPORATION.") so the
header reflects this PR's changes; locate the header comment containing that
exact string at the top of dsv3RouterGemm.cu and modify the end year
accordingly.
In `@cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp`:
- Around line 1-3: Update the file header year range in the SPDX copyright block
to include 2025 (or the current year) since this file was modified; locate the
top-of-file comment containing "SPDX-FileCopyrightText" and the
"SPDX-License-Identifier" lines in dsv3RouterGemmOp.cpp and change the trailing
year from 2024 to the current year so the header reflects the modification.
---
Nitpick comments:
In `@tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py`:
- Line 13: The parametrize matrix in
tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py only covers
supported hidden sizes and misses an unsupported case to exercise the fallback
dispatch; update the pytest.mark.parametrize("hidden_size", ...) invocation (in
test_dsv3_router_gemm.py) to include an unsupported hidden size such as 4096 so
the test matrix contains [7168, 6144, 4096], thereby exercising the fallback
path and catching dispatch regressions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8e030565-df06-4a65-853f-7fe7f9a7ae7b
📒 Files selected for processing (3)
cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cucpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpptests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py
Collaborator
Author
|
Coderabbit suggested accuracy tests on new kernel dimension range. Will work on that later |
Collaborator
Author
|
Ran GLM-5 QA tests on nsc-svg, no significant regression |
yijingl-nvidia
force-pushed
the
fix_ampere_fallback
branch
2 times, most recently
from
May 9, 2026 07:00
f226b9a to
301fe88
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #47488 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47488 [ run ] completed with state
|
GLM-5 has hidden_size=6144 with 256 MoE experts, but the
dsv3_router_gemm custom kernel only had compile-time instantiations
for hidden_dim=7168 (DeepSeek-V3). On Blackwell this caused the
routing GEMM to fall through to cublas_mm_out, which dispatched
BF16xBF16->FP32 to cutlass_80_simt_sgemm_64x64_8x5_tn_align1
(Ampere SIMT, no tensor cores) — ~5.4% / ~1 ms per decode iter
on GLM-5 FP8 MTP=3 BS=1 ISL=1K B200 TP=8.
The router_gemm_kernel template is generic over kHiddenDim as long
as kHiddenDim is divisible by VPT*kBlockSize = 1024. 6144/1024 = 6
iterations, so adding the 16 num_token instantiations is sufficient.
dsv3_router_gemm_op now dispatches between K=7168 and K=6144 paths;
unsupported shapes still fall back to cublas_mm_out.
Test coverage extended to hidden_size in {7168, 6144}.
Refs: nvbug 6108841 — TRT-LLM vs TileRT GLM-5 perf gap, target NVIDIA#4
in COMPARISON_TRTLLM_VS_TILERT.md (cutlass_80_simt_sgemm Ampere
fallback).
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com>
yijingl-nvidia
force-pushed
the
fix_ampere_fallback
branch
from
May 10, 2026 03:22
301fe88 to
4990d77
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #47554 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47554 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #47602 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47602 [ run ] completed with state |
yijingl-nvidia
requested review from
a team,
byshiue and
dongxuy04
and removed request for
a team
May 11, 2026 16:54
LarryXFly
approved these changes
May 14, 2026
This was referenced May 19, 2026
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
GLM-5 has hidden_size=6144 with 256 MoE experts, but the dsv3_router_gemm custom kernel only had compile-time instantiations for hidden_dim=7168 (DeepSeek-V3). On Blackwell this caused the routing GEMM to fall through to cublas_mm_out, which dispatched BF16xBF16->FP32 to cutlass_80_simt_sgemm_64x64_8x5_tn_align1 (Ampere SIMT, no tensor cores) — ~5.4% / ~1 ms per decode iter on GLM-5 FP8 MTP=3 BS=1 ISL=1K B200 TP=8.
The router_gemm_kernel template is generic over kHiddenDim as long as kHiddenDim is divisible by VPT*kBlockSize = 1024. 6144/1024 = 6 iterations, so adding the 16 num_token instantiations is sufficient.
dsv3_router_gemm_op now dispatches between K=7168 and K=6144 paths; unsupported shapes still fall back to cublas_mm_out.
Test Coverage
Test coverage extended to hidden_size in {7168, 6144}.
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.Summary by CodeRabbit
New Features
Tests