wehub-resource-sync
134 lines
3.9 KiB
C++
134 lines
3.9 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/engine_state.h
|
|
*/
|
|
#ifndef MLC_LLM_SERVE_ENGINE_STATE_H_
|
|
#define MLC_LLM_SERVE_ENGINE_STATE_H_
|
|
|
|
#include <tvm/ffi/cast.h>
|
|
#include <tvm/ffi/string.h>
|
|
|
|
#include "config.h"
|
|
#include "metrics.h"
|
|
#include "prefix_cache.h"
|
|
#include "request.h"
|
|
#include "request_state.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
using namespace tvm::runtime;
|
|
using tvm::ffi::GetRef;
|
|
using tvm::ffi::Object;
|
|
using tvm::ffi::ObjectRef;
|
|
|
|
typedef TypedFunction<void(Array<RequestStreamOutput>)> FRequestStreamCallback;
|
|
|
|
/*! \brief The manager of internal id for requests in engine. */
|
|
struct EngineInternalIDManager {
|
|
std::vector<int64_t> available_ids;
|
|
int64_t id_cnt = 0;
|
|
|
|
/*! \brief Return an unused id. */
|
|
int64_t GetNewId() {
|
|
if (!available_ids.empty()) {
|
|
int64_t id = available_ids.back();
|
|
available_ids.pop_back();
|
|
return id;
|
|
} else {
|
|
return id_cnt++;
|
|
}
|
|
}
|
|
|
|
/*! \brief Recycle an id. */
|
|
void RecycleId(int64_t id) { available_ids.push_back(id); }
|
|
|
|
/*! \brief Reset the manager. */
|
|
void Reset() {
|
|
available_ids.clear();
|
|
id_cnt = 0;
|
|
}
|
|
};
|
|
|
|
/*! \brief The data structures used in the action post-process. */
|
|
struct ActionPostProcessWorkspace {
|
|
std::vector<RequestStateEntry> finished_rsentries;
|
|
Array<RequestStreamOutput> callback_delta_outputs;
|
|
};
|
|
|
|
/*!
|
|
* \brief The state of the running engine.
|
|
* It contains the requests and their states submitted to the Engine.
|
|
*/
|
|
class EngineStateObj : public Object {
|
|
public:
|
|
/*! \brief The requests being processed. */
|
|
std::vector<Request> running_queue;
|
|
/*! \brief The requests that have not started for process yet. */
|
|
std::vector<Request> waiting_queue;
|
|
/*! \brief The states of all requests. */
|
|
std::unordered_map<String, RequestState> request_states;
|
|
/*! \brief The internal id manager. */
|
|
EngineInternalIDManager id_manager;
|
|
/*! \brief Runtime metrics. */
|
|
EngineMetrics metrics;
|
|
/*! \brief The prefix cache. */
|
|
PrefixCache prefix_cache{nullptr};
|
|
/*! \brief A boolean flag denoting whether the running request state entry list has changed. */
|
|
bool running_rsentries_changed = true;
|
|
/*!
|
|
* \brief The current engine speculative decoding draft length.
|
|
* The length may change across time under the auto speculative decoding mode.
|
|
* Value 0 means undefined. It must have a positive value for speculative decoding to
|
|
* properly work.
|
|
*/
|
|
int spec_draft_length = 0;
|
|
/*! \brief A boolean flag denoting whether the engine is in disaggregation mode. */
|
|
bool disaggregation = false;
|
|
// Request stream callback function
|
|
FRequestStreamCallback request_stream_callback_;
|
|
/*!
|
|
* \brief The post-process data structures.
|
|
* We make it a workspace to avoid repetitive memory allocation/free in the action post process.
|
|
*/
|
|
ActionPostProcessWorkspace postproc_workspace;
|
|
|
|
/*! \brief Reset the engine state and clear the metrics. */
|
|
void Reset();
|
|
/*! \brief Get the request state of the given request. */
|
|
RequestState GetRequestState(Request request);
|
|
/*! \brief Return the running request state entries*/
|
|
const std::vector<RequestStateEntry>& GetRunningRequestStateEntries();
|
|
|
|
static void RegisterReflection() {
|
|
namespace refl = tvm::ffi::reflection;
|
|
refl::ObjectDef<EngineStateObj>();
|
|
}
|
|
|
|
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_FINAL("mlc.serve.EngineState", EngineStateObj, Object);
|
|
|
|
private:
|
|
std::vector<RequestStateEntry> cached_running_rsentries_;
|
|
};
|
|
|
|
/*!
|
|
* \brief Managed reference of EngineStateObj.
|
|
* \sa EngineStateObj
|
|
*/
|
|
class EngineState : public ObjectRef {
|
|
public:
|
|
explicit EngineState();
|
|
|
|
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(EngineState, ObjectRef, EngineStateObj);
|
|
};
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|
|
|
|
#endif // MLC_LLM_SERVE_ENGINE_STATE_H_
|