mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
128 lines
4.5 KiB
C++
128 lines
4.5 KiB
C++
/*
|
|
* Copyright (c) 2022-2024, 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.
|
|
* 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.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "tensorrt_llm/runtime/bufferManager.h"
|
|
#include "tensorrt_llm/runtime/common.h"
|
|
#include "tensorrt_llm/runtime/eagleBuffers.h"
|
|
#include "tensorrt_llm/runtime/explicitDraftTokensBuffers.h"
|
|
#include "tensorrt_llm/runtime/iTensor.h"
|
|
#include "tensorrt_llm/runtime/lookaheadBuffers.h"
|
|
#include <optional>
|
|
#include <utility>
|
|
|
|
namespace tensorrt_llm::batch_manager
|
|
{
|
|
class LookaheadDecodingBuffers;
|
|
} // namespace tensorrt_llm::batch_manager
|
|
|
|
namespace tensorrt_llm::runtime
|
|
{
|
|
class DecodingOutput
|
|
{
|
|
public:
|
|
using TensorPtr = ITensor::SharedPtr;
|
|
|
|
// BS: batch_size, BM: beam_width, MSL: max_seq_length
|
|
// All TensorPtr without special comments are on gpu
|
|
|
|
class BeamHypotheses
|
|
{
|
|
public:
|
|
// Keep same as cpp/tensorrt_llm/kernels/beamSearchKernels.h
|
|
TensorPtr outputIdsCBA; // [BS, BM*2, MSL]
|
|
TensorPtr logProbsCBA; // [BS, BM*2, MSL]
|
|
TensorPtr sequenceLengthsCBA; // [BS, BM*2]
|
|
TensorPtr cumLogProbsCBA; // [BS, BM*2]
|
|
TensorPtr normedScoresCBA; // [BS, BM*2]
|
|
TensorPtr numBeamsCBA; // [BS]
|
|
TensorPtr minNormedScoresCBA; // [BS]
|
|
TensorPtr batchDones; // [BS]
|
|
|
|
void empty(BufferManager const& manager);
|
|
|
|
void reshape(SizeType32 batchSize, SizeType32 beamWidth, SizeType32 maxSequenceLength);
|
|
|
|
void release();
|
|
|
|
void init(BufferManager const& manager, TokenIdType endId);
|
|
|
|
BeamHypotheses slice(SizeType32 batchIndex, SizeType32 size) const;
|
|
};
|
|
|
|
static float constexpr kNegativeInfinity = -1e20f;
|
|
|
|
DecodingOutput() = default;
|
|
|
|
//! Mandatory parameters
|
|
//! Previously generated token ids for all steps before DecodingInput.step, [BS, BM, MSL]
|
|
TensorPtr ids;
|
|
//! The tokens computed during the gatherTree step, [BS, BM, MSL]
|
|
//! Necessary for "Streaming + Beam Search" mode since beam search kernels store ungathered tokens in `ids`.
|
|
TensorPtr gatheredIds;
|
|
//! New tokens at each generated token of maxTokensPerStep, [maxTokensPerStep, BS, BM]
|
|
TensorPtr newTokensSteps;
|
|
//! A view of newTokensSteps for the current token, [BS, BM]
|
|
TensorPtr newTokens;
|
|
//! A Vector of views on newTokensSteps for each token [BS, BM].
|
|
std::vector<TensorPtr> newTokensVec;
|
|
|
|
//! Optional parameters
|
|
//! FinishedState by decoding if any of the stop conditions are met or if DecodingInput.finished is true, [BS, BM]
|
|
TensorPtr finishReasons;
|
|
//! The sum of finished sequences per request, in pinned memory, [BS]
|
|
TensorPtr finishedSum;
|
|
|
|
//! Mandatory parameters for Beam Search
|
|
//! log-probility of generated tokens, [BS, BM, MSL], float
|
|
TensorPtr logProbs;
|
|
//! Sum log-probility of all generated tokens, [BS, BM]
|
|
TensorPtr cumLogProbs;
|
|
//! Index of the beam where the previous token is, [BS, BM, MSL]
|
|
TensorPtr parentIds;
|
|
//! Total sequence lengths including padding, [BS, BM]
|
|
TensorPtr lengths;
|
|
//! K/V indirection for next generation step, [BS, BM, MSL]
|
|
TensorPtr cacheIndirection;
|
|
//! Buffer used to store the transpose of the logProbs, [MSL, BS, BM]
|
|
TensorPtr logProbsTiled;
|
|
|
|
BeamHypotheses beamHypotheses;
|
|
|
|
// Speculative decoding
|
|
class SpeculativeDecodingOutputs
|
|
{
|
|
public:
|
|
TensorPtr nextDraftTokens; // [maxBatchSize, maxDraftTokens]
|
|
TensorPtr nextDraftTokensLen; // [maxBatchSize]
|
|
TensorPtr prevDraftTokensLen; // [maxBatchSize]
|
|
TensorPtr acceptedTokensLen; // [maxBatchSize]
|
|
TensorPtr acceptedLengthsCumSum; // [maxBatchSize + 1]
|
|
TensorPtr pathsOffsets; // [maxBatchSize, maxAcceptedDraftTokensPerStep]
|
|
};
|
|
|
|
std::optional<SpeculativeDecodingOutputs> speculativeDecodingOutputs;
|
|
|
|
std::optional<ExplicitDraftTokensBuffers::Inputs> explicitDraftTokensBuffers;
|
|
|
|
std::optional<LookaheadDecodingBuffers> lookaheadOutputs;
|
|
|
|
std::optional<EagleBuffers::Inputs> eagleBuffers;
|
|
};
|
|
|
|
} // namespace tensorrt_llm::runtime
|