mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
82 lines
3.6 KiB
C++
82 lines
3.6 KiB
C++
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: NVIDIA TensorRT Source Code License Agreement
|
|
*
|
|
* NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
|
* property and proprietary rights in and to this material, related
|
|
* documentation and any modifications thereto. Any use, reproduction,
|
|
* disclosure or distribution of this material and related documentation
|
|
* without an express license agreement from NVIDIA CORPORATION or
|
|
* its affiliates is strictly prohibited.
|
|
*/
|
|
|
|
#include "tensorrt_llm/executor/requestWithId.h"
|
|
#include <gmock/gmock.h>
|
|
#include <gtest/gtest.h>
|
|
|
|
using ::testing::_;
|
|
using ::testing::Invoke;
|
|
|
|
using namespace tensorrt_llm::executor;
|
|
|
|
TEST(RequestWithIdTest, serializeDeserialize)
|
|
{
|
|
std::list<VecTokens> badWords{{1, 2, 4}, {7, 8}};
|
|
std::list<VecTokens> stopWords{{1, 2}, {7, 8, 4}};
|
|
|
|
std::list<VecTokens> badWords2{{1}, {9}};
|
|
std::list<VecTokens> stopWords2{{2}, {11}};
|
|
|
|
auto embeddingBias = Tensor::cpu(DataType::kFP32, Shape({4}));
|
|
auto encoderInputFeatures = Tensor::cpu(DataType::kFP16, Shape({3000, 1280}));
|
|
|
|
float* biasData = reinterpret_cast<float*>(embeddingBias.getData());
|
|
biasData[0] = 16.f;
|
|
biasData[1] = 32.f;
|
|
biasData[2] = 32.f;
|
|
biasData[3] = 48.f;
|
|
|
|
auto request1 = Request({1, 2, 3, 4}, 1000, true, SamplingConfig(1, 4, 0.77), OutputConfig(false, true), 177, 234,
|
|
std::make_optional<std::vector<SizeType32>>({0, 1, 2, 3}), badWords, stopWords, embeddingBias,
|
|
ExternalDraftTokensConfig({11, 22}), std::nullopt, std::nullopt);
|
|
|
|
auto request2 = Request({100, 200, 300, 400}, 77, false, SamplingConfig(1, 1, 0.33), OutputConfig(true, false), 66,
|
|
99, std::make_optional<std::vector<SizeType32>>({0, 1, 1, 2}), badWords2, stopWords2, embeddingBias,
|
|
ExternalDraftTokensConfig({7, 8, 9, 10}), std::nullopt, std::nullopt);
|
|
request2.setEncoderInputFeatures(encoderInputFeatures);
|
|
|
|
auto samplingConfig3 = SamplingConfig(1, 1, 0.9);
|
|
samplingConfig3.setNumReturnSequences(3);
|
|
auto request3 = Request({37, 19, 87, 29}, 4, false, samplingConfig3, OutputConfig(false, false), 66, 99,
|
|
std::make_optional<std::vector<SizeType32>>({0, 1, 1, 2}), badWords2, stopWords2, embeddingBias,
|
|
ExternalDraftTokensConfig({7, 8, 9, 10}), std::nullopt, std::nullopt);
|
|
|
|
std::vector<RequestWithId> reqWithIds;
|
|
reqWithIds.emplace_back(RequestWithId{request1, 1});
|
|
reqWithIds.emplace_back(RequestWithId{request2, 2});
|
|
reqWithIds.emplace_back(RequestWithId{request3, 3, {4, 5}, std::chrono::steady_clock::now()});
|
|
|
|
auto serialized = RequestWithId::serializeReqWithIds(reqWithIds);
|
|
auto reqWithIdsOut = RequestWithId::deserializeReqWithIds(serialized);
|
|
|
|
EXPECT_EQ(reqWithIdsOut.size(), reqWithIds.size());
|
|
|
|
for (int i = 0; i < reqWithIdsOut.size(); ++i)
|
|
{
|
|
auto const& reqWithIdOut = reqWithIdsOut.at(i);
|
|
auto const& reqWithId = reqWithIds.at(i);
|
|
|
|
EXPECT_EQ(reqWithIdOut.id, reqWithId.id);
|
|
EXPECT_EQ(reqWithIdOut.childReqIds, reqWithId.childReqIds);
|
|
EXPECT_EQ(reqWithIdOut.queuedStart, reqWithId.queuedStart);
|
|
auto const& reqOut = reqWithIdOut.req;
|
|
auto const& req = reqWithId.req;
|
|
EXPECT_EQ(reqOut.getInputTokenIds(), req.getInputTokenIds());
|
|
EXPECT_EQ(reqOut.getMaxTokens(), req.getMaxTokens());
|
|
EXPECT_EQ(reqOut.getSamplingConfig(), req.getSamplingConfig());
|
|
EXPECT_EQ(reqOut.getStopWords(), req.getStopWords());
|
|
EXPECT_EQ(reqOut.getExternalDraftTokensConfig().value().getTokens(),
|
|
req.getExternalDraftTokensConfig().value().getTokens());
|
|
}
|
|
}
|