wehub-resource-sync
145 lines
6.1 KiB
C++
145 lines
6.1 KiB
C++
/*!
|
|
* Copyright (c) 2023-2025 by Contributors
|
|
* \file serve/engine_actions/batch_prefill_base.h
|
|
*/
|
|
|
|
#include "../config.h"
|
|
#include "../model.h"
|
|
#include "action.h"
|
|
#include "action_commons.h"
|
|
|
|
namespace mlc {
|
|
namespace llm {
|
|
namespace serve {
|
|
|
|
/*!
|
|
* \brief The base action of that prefills requests in the `waiting_queue` of
|
|
* the engine state.
|
|
*/
|
|
class BatchPrefillBaseActionObj : public EngineActionObj {
|
|
protected:
|
|
/*! \brief The class of request state entry and its maximum allowed length for prefill. */
|
|
struct PrefillInput {
|
|
RequestStateEntry rsentry;
|
|
int max_prefill_length = 0;
|
|
int num_child_to_activate = 0;
|
|
bool is_decode = false;
|
|
};
|
|
|
|
BatchPrefillBaseActionObj(Array<Model> models, EngineConfig engine_config,
|
|
std::vector<tvm::ffi::json::Object> model_configs,
|
|
Optional<EventTraceRecorder> trace_recorder);
|
|
|
|
/*!
|
|
* \brief Find one or multiple request state entries to run prefill.
|
|
* \param estate The engine state.
|
|
* \return The request entries to prefill, together with their input lengths.
|
|
*/
|
|
std::vector<PrefillInput> GetRequestStateEntriesToPrefill(EngineState estate);
|
|
|
|
/*! \brief Check if the input requests can be prefilled under conditions. */
|
|
bool CanPrefill(EngineState estate, int num_prefill_rsentries, int total_input_length,
|
|
int num_required_pages, int num_available_pages, int current_total_seq_len,
|
|
int num_running_rsentries, KVStateKind kv_state_kind,
|
|
bool sliding_window_enabled);
|
|
|
|
/*!
|
|
* \brief Chunk the input of the given RequestModelState for prefill
|
|
* with regard to the provided maximum allowed prefill length.
|
|
* Return the list of input for prefill and the total prefill length.
|
|
* The `inputs` field of the given `mstate` will be mutated to exclude
|
|
* the returned input.
|
|
* \param mstate The RequestModelState whose input data is to be chunked.
|
|
* \param max_prefill_length The maximum allowed prefill length for the mstate.
|
|
* \return The list of input for prefill and the total prefill length.
|
|
*/
|
|
std::pair<Array<Data>, int> ChunkPrefillInputData(const RequestModelState& mstate,
|
|
int max_prefill_length);
|
|
|
|
/*!
|
|
* \brief Update status of request states from pending to alive and collect request state entries
|
|
* from the prefill input.
|
|
* \param prefill_inputs The prefill input.
|
|
* \param estate The engine state.
|
|
* \param[out] request_ids The array to store the request ids of the request state entries.
|
|
* \param[out] rstates_of_entries The vector to store the request state entries.
|
|
* \param[out] status_before_prefill The vector to store the status of the request state entries
|
|
* before prefill.
|
|
*/
|
|
void UpdateRequestToAlive(const std::vector<PrefillInput>& prefill_inputs,
|
|
const EngineState& estate, Array<String>* request_ids,
|
|
std::vector<RequestState>* rstates_of_entries,
|
|
std::vector<RequestStateStatus>* status_before_prefill);
|
|
|
|
/*!
|
|
* \brief Remove the request from waiting queue if all its request states are now alive and have
|
|
* no remaining chunked inputs.
|
|
* \param prefill_inputs The prefill input.
|
|
* \param estate The engine state.
|
|
* \param rstates_of_entries The request state entries for each prefill input.
|
|
* \return The processed requests.
|
|
*/
|
|
std::vector<Request> RemoveProcessedRequests(const std::vector<PrefillInput>& prefill_inputs,
|
|
const EngineState& estate,
|
|
const std::vector<RequestState>& rstates_of_entries);
|
|
|
|
/*!
|
|
* \brief Update the committed tokens of states. If a request is first-time prefilled, set the
|
|
* prefill finish time.
|
|
* \param rsentries_for_sample The request state entries for sample.
|
|
* \param rsentry_activated The activation status of the request state entries.
|
|
* \param sample_results The sample results.
|
|
*/
|
|
void UpdateRequestStateEntriesWithSampleResults(
|
|
const std::vector<RequestStateEntry>& rsentries_for_sample,
|
|
const std::vector<bool>& rsentry_activated, const std::vector<SampleResult>& sample_results);
|
|
|
|
/*!
|
|
* \brief Get the concatenated Shape of RequestModelState input data, return empty Shape if
|
|
* there is untokenized data.
|
|
* \param mstate The RequestModelState whose input data is to be concatenated.
|
|
* \return The concatenate Shape.
|
|
*/
|
|
std::vector<int32_t> GetConcatPrefillInputData(const RequestModelState& mstate);
|
|
|
|
/*!
|
|
* \brief Pop the prefix tokens of the RequestModelState input data array.
|
|
* \param mstate The RequestModelState to be popped.
|
|
* \param num_tokens The number of prefix tokens to be popped.
|
|
*/
|
|
void PopPrefillInputData(const RequestModelState& mstate, size_t num_tokens);
|
|
|
|
/*!
|
|
* \brief Match the request state entry with prefix cache, to skip prefilling common prefix
|
|
* tokens. If the request state entry is not added to KVCache yet, this method will add/fork the
|
|
* request in the KVCache, depending on the matching result from prefix cache.
|
|
* \param estate The engine state.
|
|
* \param[in, out] input The prefill input to be matched and updated.
|
|
* \return The matched length in prefix cache.
|
|
*/
|
|
virtual int MatchPrefixCache(EngineState estate, PrefillInput* input) = 0;
|
|
|
|
/*! \brief The models to run prefill in. */
|
|
Array<Model> models_;
|
|
/*! \brief The engine config. */
|
|
EngineConfig engine_config_;
|
|
/*! \brief The KV state kind. */
|
|
KVStateKind kv_state_kind_;
|
|
/*! \brief The sliding window size of each model. */
|
|
std::vector<int> sliding_window_sizes_;
|
|
/*! \brief Event trace recorder. */
|
|
Optional<EventTraceRecorder> trace_recorder_;
|
|
};
|
|
|
|
/*!
|
|
* \brief A utility function to check whether there is enough spare space in
|
|
* KV cache for the number of required pages and total input length.
|
|
*/
|
|
bool HasPrefillSpace(int num_required_pages, bool sliding_window_enabled, int new_batch_size,
|
|
int num_available_pages, int current_total_seq_len, int total_input_length,
|
|
int max_total_sequence_length);
|
|
|
|
} // namespace serve
|
|
} // namespace llm
|
|
} // namespace mlc
|