wehub-resource-sync
258 lines
8.2 KiB
C++
258 lines
8.2 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/data.h
|
|
*/
|
|
#ifndef MLC_LLM_SERVE_DATA_H_
|
|
#define MLC_LLM_SERVE_DATA_H_
|
|
|
|
#include <tvm/ffi/container/array.h>
|
|
#include <tvm/ffi/container/shape.h>
|
|
#include <tvm/ffi/object.h>
|
|
#include <tvm/ffi/optional.h>
|
|
#include <tvm/ffi/reflection/registry.h>
|
|
#include <tvm/ffi/string.h>
|
|
#include <tvm/runtime/tensor.h>
|
|
|
|
#include <atomic>
|
|
#include <optional>
|
|
|
|
#include "../tokenizers/tokenizers.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
using namespace tvm::runtime;
|
|
using tvm::ffi::Object;
|
|
using tvm::ffi::ObjectPtr;
|
|
using tvm::ffi::ObjectRef;
|
|
using tvm::ffi::Optional;
|
|
using tvm::ffi::Shape;
|
|
|
|
class Model;
|
|
|
|
/****************** DataNode ******************/
|
|
|
|
/*! \brief The base class of multi-modality data (text, tokens, embedding, etc). */
|
|
class DataNode : public Object {
|
|
public:
|
|
/*! \brief Get the length (equivalent number of tokens) of the data. */
|
|
virtual int GetLength() const = 0;
|
|
|
|
/*!
|
|
* \brief Compute the embedding of this data with regard to the input model.
|
|
* When the input destination pointer is not nullptr, it in-place writes the
|
|
* embedding into the input destination array at the given offset.
|
|
* Otherwise, the embeddings will be directly returned back.
|
|
* \param model The model to take embeddings from.
|
|
* \param dst The destination array of the embedding lookup.
|
|
* \param offset The token offset where the computed embeddings will be written
|
|
* into the destination array.
|
|
* \return The updated destination embedding array or the computed embeddings.
|
|
* \note When `dst` is nullptr, we require `offset` to be 0.
|
|
*/
|
|
virtual ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const = 0;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<DataNode>();
|
|
}
|
|
|
|
static constexpr const bool _type_has_method_sequal_reduce = false;
|
|
static constexpr const bool _type_has_method_shash_reduce = false;
|
|
static constexpr const uint32_t _type_child_slots = 3;
|
|
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.Data", DataNode, Object);
|
|
};
|
|
|
|
class Data : public ObjectRef {
|
|
public:
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Data, ObjectRef, DataNode);
|
|
};
|
|
|
|
/*! \brief Split the given data array into two arrays at the "split_pos" position. */
|
|
std::pair<Array<Data>, Array<Data>> SplitData(const Array<Data>& original_data, int total_length,
|
|
int split_pos);
|
|
|
|
/****************** TextDataNode ******************/
|
|
|
|
/*! \brief The class of text data, containing a text string. */
|
|
class TextDataNode : public DataNode {
|
|
public:
|
|
/*! \brief The text string. */
|
|
tvm::ffi::String text;
|
|
|
|
int GetLength() const final;
|
|
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<TextDataNode>();
|
|
}
|
|
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TextData", TextDataNode, DataNode);
|
|
};
|
|
|
|
class TextData : public Data {
|
|
public:
|
|
explicit TextData(String text);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TextData, Data, TextDataNode);
|
|
};
|
|
|
|
/****************** TokenDataNode ******************/
|
|
|
|
/*! \brief The class of token data, containing a list of token ids. */
|
|
class TokenDataNode : public DataNode {
|
|
public:
|
|
/*! \brief The token ids. */
|
|
Shape token_ids;
|
|
|
|
int GetLength() const final;
|
|
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<TokenDataNode>();
|
|
}
|
|
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TokenData", TokenDataNode, DataNode);
|
|
};
|
|
|
|
class TokenData : public Data {
|
|
public:
|
|
explicit TokenData(Shape token_ids);
|
|
|
|
explicit TokenData(std::vector<int32_t> token_ids);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TokenData, Data, TokenDataNode);
|
|
};
|
|
|
|
/****************** ImageDataNode ******************/
|
|
|
|
/*! \brief The class of image data, containing a 3D array of pixel values. */
|
|
class ImageDataNode : public DataNode {
|
|
public:
|
|
/*! \brief The pixel values. */
|
|
Tensor image;
|
|
int embed_size;
|
|
|
|
int GetLength() const final;
|
|
ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<ImageDataNode>();
|
|
}
|
|
|
|
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.ImageData", ImageDataNode, DataNode);
|
|
};
|
|
|
|
class ImageData : public Data {
|
|
public:
|
|
explicit ImageData(Tensor image, int embed_size);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ImageData, Data, ImageDataNode);
|
|
};
|
|
|
|
/****************** SampleResult ******************/
|
|
|
|
// The pair of a token id and its probability in sampling.
|
|
using TokenProbPair = std::pair<int32_t, float>;
|
|
|
|
/*!
|
|
* \brief The class of sampler's sampling result.
|
|
* It's not a TVM object since it will not be used directly on Python side.
|
|
*/
|
|
struct SampleResult {
|
|
/*! \brief The token id and probability of the sampled token. */
|
|
TokenProbPair sampled_token_id;
|
|
/*! \brief The token id and probability of the tokens with top probabilities. */
|
|
std::vector<TokenProbPair> top_prob_tokens;
|
|
|
|
/*! \brief Get the sampled token id. */
|
|
int32_t GetTokenId() const;
|
|
|
|
/*!
|
|
* \brief Get the logprob JSON string of this token with regard
|
|
* to OpenAI API at https://platform.openai.com/docs/api-reference/chat/object.
|
|
* \param tokenizer The tokenizer for token table lookup.
|
|
* \param logprob A boolean indicating if need to return log probability.
|
|
* \return A JSON string that conforms to the logprob spec in OpenAI API.
|
|
*/
|
|
std::string GetLogProbJSON(const Tokenizer& tokenizer, bool logprob) const;
|
|
};
|
|
|
|
/****************** RequestStreamOutput ******************/
|
|
|
|
/*!
|
|
* \brief The generated delta request output that is streamed back
|
|
* through callback stream function.
|
|
*
|
|
* \note: This output object corresponds to parallel generated outputs when n != 1.
|
|
*
|
|
* For example, if n=2, then group_delta_token_ids[0] matches to the output stream 0
|
|
* and group_delta_token_ids[1] matches to the output stream 1
|
|
*/
|
|
class RequestStreamOutputObj : public Object {
|
|
public:
|
|
/*! \brief The id of the request that the function is invoked for. */
|
|
String request_id;
|
|
/*!
|
|
* \brief The new generated token ids since the last callback invocation
|
|
* for the input request.
|
|
*/
|
|
std::vector<std::vector<int64_t>> group_delta_token_ids;
|
|
/*! \brief The logprobs JSON strings of the new generated tokens since last invocation. */
|
|
std::optional<std::vector<std::vector<String>>> group_delta_logprob_json_strs;
|
|
/*!
|
|
* \brief The finish reason of the request when it is finished,
|
|
* of None if the request has not finished yet.
|
|
*/
|
|
std::vector<Optional<String>> group_finish_reason;
|
|
/*!
|
|
* \brief The usage field of the response, this is global to all streams.
|
|
*/
|
|
Optional<String> request_final_usage_json_str;
|
|
|
|
/*!
|
|
* \brief The extra prefix string of all requests.
|
|
*/
|
|
std::vector<String> group_extra_prefix_string;
|
|
|
|
std::atomic<bool> unpacked = false;
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<RequestStreamOutputObj>();
|
|
}
|
|
|
|
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.RequestStreamOutput", RequestStreamOutputObj, Object);
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference to RequestStreamOutputObj.
|
|
* \sa RequestStreamOutputObj
|
|
*/
|
|
class RequestStreamOutput : public ObjectRef {
|
|
public:
|
|
explicit RequestStreamOutput(
|
|
String request_id, std::vector<std::vector<int64_t>> group_delta_token_ids,
|
|
std::optional<std::vector<std::vector<String>>> group_delta_logprob_json_strs,
|
|
std::vector<Optional<String>> group_finish_reason,
|
|
std::vector<String> group_extra_prefix_string);
|
|
|
|
static RequestStreamOutput Usage(String request_id, String request_final_usage_json_str);
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestStreamOutput, ObjectRef,
|
|
RequestStreamOutputObj);
|
|
};
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|
|
|
|
#endif // MLC_LLM_SERVE_DATA_H_
|