mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
99 lines
3.8 KiB
C++
99 lines
3.8 KiB
C++
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 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.
|
|
*/
|
|
|
|
#include "tensorrt_llm/common/opUtils.h"
|
|
#include "tensorrt_llm/runtime/torchUtils.h"
|
|
|
|
#include "tensorrt_llm/kernels/noAuxTcKernels.h"
|
|
|
|
// #include <NvInferRuntime.h>
|
|
// #include <c10/cuda/CUDAStream.h>
|
|
// #include <cassert>
|
|
// #include <set>
|
|
// #include <string>
|
|
// #include <torch/extension.h>
|
|
// #include <vector>
|
|
|
|
namespace th = torch;
|
|
namespace tl = tensorrt_llm;
|
|
namespace tk = tensorrt_llm::kernels;
|
|
|
|
namespace torch_ext
|
|
{
|
|
th::Tensor noaux_tc_op(th::Tensor const& scores, th::Tensor const& scores_with_bias, int64_t n_group,
|
|
int64_t topk_group, int64_t topk, double routed_scaling_factor)
|
|
{
|
|
auto data_type = scores_with_bias.scalar_type();
|
|
auto input_size = scores_with_bias.sizes();
|
|
int64_t num_tokens = input_size[0];
|
|
int64_t num_experts = input_size[1];
|
|
TORCH_CHECK(input_size.size() == 2, "scores_with_bias must be a 2D Tensor");
|
|
TORCH_CHECK(num_experts % n_group == 0, "num_experts should be divisible by n_group");
|
|
TORCH_CHECK(
|
|
n_group <= 32, "n_group should be smaller than or equal to 32 for now"); //@todo: remove this restriction later
|
|
TORCH_CHECK(
|
|
topk <= 32, "topk should be smaller than or equal to 32 for now"); //@todo: remove this restriction later
|
|
|
|
th::Tensor group_scores = th::empty({num_tokens, n_group}, th::dtype(data_type).device(torch::kCUDA));
|
|
|
|
auto stream = at::cuda::getCurrentCUDAStream(scores_with_bias.get_device());
|
|
|
|
switch (data_type)
|
|
{
|
|
case torch::kFloat16:
|
|
// Handle Float16
|
|
tk::invokeNoAuxTc<half>(reinterpret_cast<half*>(scores.mutable_data_ptr()),
|
|
reinterpret_cast<half*>(group_scores.mutable_data_ptr()),
|
|
reinterpret_cast<half*>(scores_with_bias.data_ptr()), num_tokens, num_experts, n_group, topk_group, topk,
|
|
routed_scaling_factor, stream);
|
|
break;
|
|
case torch::kFloat32:
|
|
// Handle Float32
|
|
tk::invokeNoAuxTc<float>(reinterpret_cast<float*>(scores.mutable_data_ptr()),
|
|
reinterpret_cast<float*>(group_scores.mutable_data_ptr()),
|
|
reinterpret_cast<float*>(scores_with_bias.data_ptr()), num_tokens, num_experts, n_group, topk_group, topk,
|
|
routed_scaling_factor, stream);
|
|
break;
|
|
case torch::kBFloat16:
|
|
// Handle BFloat16
|
|
tk::invokeNoAuxTc<__nv_bfloat16>(reinterpret_cast<__nv_bfloat16*>(scores.mutable_data_ptr()),
|
|
reinterpret_cast<__nv_bfloat16*>(group_scores.mutable_data_ptr()),
|
|
reinterpret_cast<__nv_bfloat16*>(scores_with_bias.data_ptr()), num_tokens, num_experts, n_group, topk_group,
|
|
topk, routed_scaling_factor, stream);
|
|
break;
|
|
default:
|
|
// Handle other data types
|
|
throw std::invalid_argument("Invalid dtype, only supports float16, float32, and bfloat16");
|
|
break;
|
|
}
|
|
return scores;
|
|
}
|
|
|
|
} // end namespace torch_ext
|
|
|
|
TORCH_LIBRARY_FRAGMENT(trtllm, m)
|
|
{
|
|
m.def(
|
|
"noaux_tc_op(Tensor scores, Tensor scores_with_bias, int n_group, int topk_group, int topk, float "
|
|
"routed_scaling_factor) -> Tensor");
|
|
}
|
|
|
|
TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
|
|
{
|
|
m.impl("noaux_tc_op", &torch_ext::noaux_tc_op);
|
|
}
|