Skip to content

[https://nvbugs/6108841][fix] add hidden_dim=6144 router GEMM instantiation for GLM-5 #13740

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

Merged
yijingl-nvidia merged 4 commits into from
May 14, 2026
Merged

[https://nvbugs/6108841][fix] add hidden_dim=6144 router GEMM instantiation for GLM-5 #13740

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9206de6
fix: add hidden_dim=6144 router GEMM instantiation for GLM-5
yijingl-nvidia May 4, 2026
a8cf884
update headers
yijingl-nvidia May 4, 2026
edc4352
minor comment
yijingl-nvidia May 5, 2026
4990d77
fix input shape guard
yijingl-nvidia May 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 50 additions & 1 deletion cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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);
Comment thread
coderabbitai[bot] marked this conversation as resolved.

} // namespace kernels::dsv3MinLatencyKernels

TRTLLM_NAMESPACE_END
23 changes: 13 additions & 10 deletions cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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");
Expand Down Expand Up @@ -77,28 +77,31 @@ 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<int64_t> 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);
TORCH_CHECK(mat_a.strides()[1] == 1 && out.strides()[1] == 1); // Row-major
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
&& mat_b.sizes()[0] == hidden_dim && 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<float*>(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<float*>(out.mutable_data_ptr()), reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
}
Comment thread
yijingl-nvidia marked this conversation as resolved.
else
else // fallback to cublas, can be slow
{
cublas_mm_out(mat_a, mat_b, bias, out);
}
Expand Down
18 changes: 17 additions & 1 deletion tests/unittest/_torch/thop/parallel/test_dsv3_router_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -10,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])
@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)
Expand Down
Loading