wehub-resource-sync
105 lines
4.1 KiB
C++
105 lines
4.1 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/logit_processor.h
|
|
* \brief The header for logit processor.
|
|
*/
|
|
|
|
#ifndef MLC_LLM_SERVE_LOGIT_PROCESSOR_H_
|
|
#define MLC_LLM_SERVE_LOGIT_PROCESSOR_H_
|
|
|
|
#include <tvm/ffi/extra/module.h>
|
|
#include <tvm/ffi/string.h>
|
|
|
|
#include "../base.h"
|
|
#include "config.h"
|
|
#include "event_trace_recorder.h"
|
|
#include "function_table.h"
|
|
#include "request_state.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
using tvm::Device;
|
|
using namespace tvm::runtime;
|
|
using tvm::ffi::Object;
|
|
using tvm::ffi::ObjectRef;
|
|
|
|
/*!
|
|
* \brief The logit processor class that updates logits with regard
|
|
* presence/frequency penalties, logit bias, etc..
|
|
*/
|
|
class LogitProcessorObj : public Object {
|
|
public:
|
|
/*!
|
|
* \brief In-place update a batch of logits with regard to the given
|
|
* generation config and request states.
|
|
* \param logits The batch of raw logits, in shape (num_total_token, vocab_size),
|
|
* where `num_total_token` may be larger than the number of sequences
|
|
* indicated by `generation_cfg`, in which case some sequences may have
|
|
* more than one token.
|
|
* \param generation_cfg The generation config of each sequence in the batch.
|
|
* \param mstates The request states of each sequence in the batch.
|
|
* \param request_ids The ids of each request.
|
|
* \param cum_num_token The pointer to the cumulative token length of the sequences.
|
|
* If the pointer is nullptr, it means each sequence has only one token.
|
|
* \param draft_mstates Optional. The draft request states of each sequence.
|
|
* \param draft_token_indices Optional. The pointer to the draft token indices of each draft token
|
|
* in the model state (-1 indicates the token is not a draft token) when speculation is enabled.
|
|
* This is used to compute the sequence state with the draft tokens considered (the saved sequence
|
|
* state is not updated with the draft tokens).
|
|
*/
|
|
virtual void InplaceUpdateLogits(
|
|
Tensor logits, const Array<GenerationConfig>& generation_cfg,
|
|
const Array<RequestModelState>& mstates, const Array<String>& request_ids,
|
|
const std::vector<int>* cum_num_token = nullptr,
|
|
const Array<RequestModelState>* draft_mstates = nullptr,
|
|
const std::vector<std::vector<int>>* draft_token_indices = nullptr) = 0;
|
|
|
|
/*!
|
|
* \brief Compute probability distributions for the input batch of logits.
|
|
* \param logits The batch of updated logits.
|
|
* \param generation_cfg The generation config of each sequence in the batch.
|
|
* \param request_ids The ids of each request.
|
|
* \param cum_num_token The pointer to the cumulative token length of the sequences.
|
|
* If the pointer is nullptr, it means each sequence has only one token.
|
|
* \return The batch of computed probability distributions on GPU.
|
|
*/
|
|
virtual Tensor ComputeProbsFromLogits(Tensor logits,
|
|
const Array<GenerationConfig>& generation_cfg,
|
|
const Array<String>& request_ids,
|
|
const std::vector<int>* cum_num_token = nullptr) = 0;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<LogitProcessorObj>();
|
|
}
|
|
|
|
static constexpr const bool _type_has_method_sequal_reduce = false;
|
|
static constexpr const bool _type_has_method_shash_reduce = false;
|
|
static constexpr const bool _type_mutable = true;
|
|
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.LogitProcessor", LogitProcessorObj, Object);
|
|
};
|
|
|
|
class LogitProcessor : public ObjectRef {
|
|
public:
|
|
/*!
|
|
* \brief Constructor.
|
|
* \param max_num_token The max number of tokens in the token processor.
|
|
* \param vocab_size The model's vocabulary size.
|
|
* \param ft The packed function table.
|
|
* \param device The device that the model runs on.
|
|
* \param trace_recorder The event trace recorder.
|
|
*/
|
|
explicit LogitProcessor(int max_num_token, int vocab_size, FunctionTable* ft, DLDevice device,
|
|
Optional<EventTraceRecorder> trace_recorder);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LogitProcessor, ObjectRef, LogitProcessorObj);
|
|
};
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|
|
|
|
#endif // MLC_LLM_SERVE_LOGIT_PROCESSOR_H_
|