diff --git a/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu b/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu index 0b406e103fa8..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. @@ -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..2e6e0822debe 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"); @@ -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,20 +86,22 @@ 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 + && 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(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); } - else + else // fallback to cublas, can be slow { cublas_mm_out(mat_a, mat_b, bias, out); } 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..abafe6a3e271 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 @@ -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)