From 9206de6cd1f85041268e6ce90e0fd638b35ba301 Mon Sep 17 00:00:00 2001 From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> Date: Mon, 4 May 2026 12:04:11 -0700 Subject: [PATCH 1/4] fix: add hidden_dim=6144 router GEMM instantiation for GLM-5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 #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> --- .../dsv3MinLatencyKernels/dsv3RouterGemm.cu | 49 +++++++++++++++++++ cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp | 19 ++++--- .../thop/parallel/test_dsv3_router_gemm.py | 2 +- 3 files changed, 61 insertions(+), 9 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu index 0b406e103fa8..18c689cdef65 100644 --- a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu +++ b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu @@ -242,6 +242,55 @@ template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__n template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 7168>( float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); +// hidden_dim=6144 instantiations (GLM-5). +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 1, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 2, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 3, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 4, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 5, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 6, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 7, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 8, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 9, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 10, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 11, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 12, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 13, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 14, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 15, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + +template void tensorrt_llm::kernels::dsv3MinLatencyKernels::invokeRouterGemm<__nv_bfloat16, 16, 256, 6144>( + float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); + } // namespace kernels::dsv3MinLatencyKernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp index 6764cbef64cb..2bce9117b5c1 100644 --- a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp @@ -77,7 +77,8 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, auto const out_dtype_ = out_dtype.value_or(mat_a.scalar_type()); auto const data_type = mat_a.scalar_type(); constexpr int kNumExperts = 256; - constexpr int kHiddenDim = 7168; + constexpr int kHiddenDim7168 = 7168; // DeepSeek-V3 / DeepSeek-V3.2 + constexpr int kHiddenDim6144 = 6144; // GLM-5 std::vector output_size = {mat_a.sizes()[0], mat_b.sizes()[1]}; th::Tensor out = th::empty(output_size, mat_a.options().dtype(out_dtype_)); TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2); @@ -85,16 +86,18 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, TORCH_CHECK(mat_b.strides()[0] == 1); // Column-major TORCH_CHECK(!bias.has_value(), "bias is not support yet"); auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device()); - bool use_custom_kernel = false; - if (num_tokens >= 1 && num_tokens <= 16 && num_experts == kNumExperts && hidden_dim == kHiddenDim - && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32) + bool const shape_ok = (num_tokens >= 1 && num_tokens <= 16 && num_experts == kNumExperts + && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32); + + if (shape_ok && hidden_dim == kHiddenDim7168) { - use_custom_kernel = true; + LoopUnroller<1, 16, kNumExperts, kHiddenDim7168>::unroll(num_tokens, + reinterpret_cast(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), + reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); } - - if (use_custom_kernel) + else if (shape_ok && hidden_dim == kHiddenDim6144) { - LoopUnroller<1, 16, kNumExperts, kHiddenDim>::unroll(num_tokens, + LoopUnroller<1, 16, kNumExperts, kHiddenDim6144>::unroll(num_tokens, reinterpret_cast(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); } diff --git a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py index b183ae797593..7872d6f5c8df 100644 --- a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py @@ -10,7 +10,7 @@ def router_gemm_ref(input, weight, bias, dtype): @pytest.mark.parametrize( "num_tokens", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]) @pytest.mark.parametrize("num_experts", [256]) -@pytest.mark.parametrize("hidden_size", [7168]) +@pytest.mark.parametrize("hidden_size", [7168, 6144]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) def test_router_gemm_run(num_tokens, num_experts, hidden_size, dtype): torch.manual_seed(24) From a8cf88465693272f4928afb9ddc20f7cb57a1ea9 Mon Sep 17 00:00:00 2001 From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> Date: Mon, 4 May 2026 15:42:00 -0700 Subject: [PATCH 2/4] update headers Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> --- .../dsv3MinLatencyKernels/dsv3RouterGemm.cu | 2 +- cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp | 2 +- .../_torch/thop/parallel/test_dsv3_router_gemm.py | 15 +++++++++++++++ 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu index 18c689cdef65..e33f46e1e152 100644 --- a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu +++ b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp index 2bce9117b5c1..0d614b889a16 100644 --- a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 * * Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py index 7872d6f5c8df..353f8d4dfae9 100644 --- a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import pytest import torch From edc4352974ccb98df1307654dae5e3fb27299782 Mon Sep 17 00:00:00 2001 From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> Date: Tue, 5 May 2026 14:39:04 -0700 Subject: [PATCH 3/4] minor comment Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp index 0d614b889a16..bfe4258525cc 100644 --- a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp @@ -101,7 +101,7 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, reinterpret_cast(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream); } - else + else // fallback to cublas, can be slow { cublas_mm_out(mat_a, mat_b, bias, out); } From 4990d7704a37a32f6428a3be101a805af44e0e92 Mon Sep 17 00:00:00 2001 From: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> Date: Fri, 8 May 2026 18:21:11 -0700 Subject: [PATCH 4/4] fix input shape guard Signed-off-by: Yijing Li <257409031+yijingl-nvidia@users.noreply.github.com> --- cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp | 2 +- tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp index bfe4258525cc..2e6e0822debe 100644 --- a/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp +++ b/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp @@ -87,7 +87,7 @@ th::Tensor dsv3_router_gemm_op(th::Tensor const& mat_a, th::Tensor const& mat_b, TORCH_CHECK(!bias.has_value(), "bias is not support yet"); auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device()); bool const shape_ok = (num_tokens >= 1 && num_tokens <= 16 && num_experts == kNumExperts - && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32); + && mat_b.sizes()[0] == hidden_dim && data_type == torch::kBFloat16 && out_dtype_ == torch::kFloat32); if (shape_ok && hidden_dim == kHiddenDim7168) { diff --git a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py index 353f8d4dfae9..abafe6a3e271 100644 --- a/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py +++ b/tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py @@ -25,7 +25,8 @@ def router_gemm_ref(input, weight, bias, dtype): @pytest.mark.parametrize( "num_tokens", [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]) @pytest.mark.parametrize("num_experts", [256]) -@pytest.mark.parametrize("hidden_size", [7168, 6144]) +@pytest.mark.parametrize("hidden_size", + [7168, 6144, 4096]) # 4096 will enter fallback kernel @pytest.mark.parametrize("dtype", [torch.bfloat16]) def test_router_gemm_run(num_tokens, num_experts, hidden_size, dtype): torch.manual_seed(24)