mlc-ai--mlc-llm/cpp/serve/request_state.h

/*!
 *  Copyright (c) 2023-2025 by Contributors
 * \file serve/request_state.h
 * \brief The data structure maintaining the generation states of user requests.
 */
#ifndef MLC_LLM_SERVE_REQUEST_STATE_H_
#define MLC_LLM_SERVE_REQUEST_STATE_H_

#include <tvm/ffi/container/array.h>
#include <tvm/ffi/object.h>
#include <tvm/runtime/tensor.h>
#include <xgrammar/xgrammar.h>

#include <optional>

#include "../support/random.h"
#include "../tokenizers/streamer.h"
#include "config.h"
#include "metrics.h"
#include "request.h"

namespace mlc {
namespace llm {
namespace serve {

using namespace tvm::runtime;
using tvm::ffi::Object;
using tvm::ffi::ObjectRef;

/*!
 * \brief The state of a request with regard to some single model.
 * \details In MLC LLM, the serving engine may leverage multiple models
 * to fulfill a user generation request (e.g., use speculation decoding).
 * For each request, we isolate its states (e.g. the generated tokens)
 * on each model. This is to say, we use RequestModelState to store
 * the state of a user request on a single model (rather than all models).
 */
class RequestModelStateNode : public Object {
 public:
  /*! \brief The request that this state corresponds to. */
  Request request;
  /*!
   * \brief The internal request id of this state.
   * It is the **physical index** of the request in the running request queue.
   * If the request is on hold (not in the running queue), the request id
   * should be -1.
   */
  int64_t internal_id = -1;
  /*! \brief The corresponding model id of this state. */
  int model_id = -1;
  /*!
   * \brief The committed generated token ids and related probability info.
   * A token is "committed" means it will no longer be updated (or changed).
   */
  std::vector<SampleResult> committed_tokens;
  /*! \brief The list of input data yet for the model to prefill. */
  Array<Data> inputs;
  /*! \brief The list of prefilled input data, used to notify prefix cache. */
  std::vector<Data> prefilled_inputs;
  /*! \brief The number of tokens already cached in prefix cache. */
  int64_t cached_committed_tokens = 0;
  /*! \brief The number of tokens that is already prefilled from the inputs. */
  int64_t num_prefilled_tokens = 0;
  /*! \brief The number of tokens that need to be processed in the next decoding. */
  int num_tokens_for_next_decode = 0;
  /*! \brief Whether retokenization is needed in the next decoding. When the jump-forward decoding
   * is enabled, retokenization is needed after every jump-forward and decoding action. */
  bool require_retokenization_in_next_decode = false;

  // NOTE: The following fields are reserved for future speculative inference
  // settings, and are produced by the speculative small models.
  /*!
   * \brief The draft generated token ids and related probability info,
   * which are usually generated by "small" speculative models.
   * These tokens will be fed to a "large" model to determine the final
   * result of speculation.
   */
  std::vector<SampleResult> draft_output_tokens;
  /*! \brief The storage slots for the associated states of draft tokens. */
  std::vector<int> draft_token_slots;
  /*! \brief The parent indices of the draft tokens. */
  std::vector<int64_t> draft_token_parent_idx;
  /*! \brief The first child indices of the draft tokens. */
  std::vector<int64_t> draft_token_first_child_idx;

  /*! \brief The appeared committed and draft tokens and their occurrence times. */
  std::unordered_map<int32_t, int32_t> appeared_token_ids;

  /*!
   * \brief The current state of the generated token matching the grammar. Used in grammar-guided
   * generation, otherwise it's std::nullopt.
   */
  std::optional<xgrammar::GrammarMatcher> grammar_matcher;

  /*! \brief Return the total length of the input data. */
  int GetInputLength() const;
  /*!
   * \brief Return whether the next token bitmask is required, i.e. the grammar-guided generation is
   * enabled.
   */
  bool RequireNextTokenBitmask();
  /*!
   * \brief Find the next token bitmask and store it in the given DLTensor.
   * \param bitmask The DLTensor to store the next token bitmask. The bitmask should be a tensor
   * with dtype uint32_t and shape (ceildiv(vocab_size, 32),).
   */
  void GetNextTokenBitmask(DLTensor* bitmask);
  /*! \brief Commit a new token into committed_tokens. Does not effect the kv cache. Update
   * appeared_token_ids and the grammar state. */
  void CommitToken(SampleResult sampled_token);
  /*! \brief Roll back the last tokens back from committed_tokens. Does not effect the kv cache.
   * Also roll back appeared_token_ids and the grammar state. */
  void RollbackTokens(int count);

  /*! \brief Add a draft token into draft_output_tokens. Update appeared_token_ids. */
  void AddDraftToken(SampleResult sampled_token, int draft_token_slot, int64_t parent_idx);
  /*! \brief Remove all draft tokens from draft_output_tokens. Update appeared_token_ids. */
  void RemoveAllDraftTokens(std::vector<int>* removed_draft_token_slots = nullptr);

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RequestModelStateNode>();
  }

  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.RequestModelState", RequestModelStateNode, Object);
};

class RequestModelState : public ObjectRef {
 public:
  explicit RequestModelState(Request request, int model_id, int64_t internal_id, Array<Data> inputs,
                             const std::optional<xgrammar::CompiledGrammar>& compiled_grammar);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestModelState, ObjectRef, RequestModelStateNode);
};

struct DeltaRequestReturn {
  std::vector<int64_t> delta_token_ids;
  std::vector<String> delta_logprob_json_strs;
  Optional<String> finish_reason;
  /*! \brief The extra string to prepend the delta output. The delta output should be
   * extra_prefix_string + detokenize(delta_token_ids). */
  String extra_prefix_string = "";
};

/****************** Request States ******************/

/*!
 * \brief For each request, we maintain its "request state" in the
 * engine. Generally, the state of a request contains the information
 * of the request's generation at the current moment, including
 * the generated token ids, the grammar handler, etc.
 *
 * When a request has multiple parallel generations (e.g., the field
 * `n` of its generation config is more than 1), each generation will
 * have different states all the time.
 *
 * Therefore, to better support parallel generations, we denote the
 * state of a single generation as a "RequestStateEntry" instance,
 * and denote the state of a request's all generations using a vector,
 * named as a "RequestState" instance.
 *
 * A request's all state entries are organized as a tree structure
 * when there are parallel generations.
 * - the request input has the root status entry,
 * - each parallel generation is a child of the root.
 * This tree structure may be further extended to more complicated
 * cases in the future. As of now, for the case of `n > 1`, there
 * will be (n + 1) entries in total. In a "RequestState", the root
 * entry always has index 0. And we guarantee that the entry order
 * from the vector begin to the end is always a topological order
 * of the tree.
 */

/*! \brief Request state status. */
enum class RequestStateStatus : int {
  kPending = 0,
  kAlive = 1,
  kFinished = 2,
};

/*! \brief The data structures for each request used in the action post-process. */
struct RequestActionPostProcWorkspace {
  std::vector<RequestStreamOutput> stream_outputs;

  RequestStreamOutput GetStreamOutput();
};

// forward declare request state node.
class RequestStateNode;

/*!
 * \brief A request's state entry. It contains the state of a single
 * generation of a request, or the state of a prompt prefix of a request.
 */
class RequestStateEntryNode : public Object {
 public:
  /*! \brief The status of the request state entry. */
  RequestStateStatus status;
  /*! \brief The request that this state corresponds to. */
  Request request;
  /*!
   * \brief The idx of the parent request state entry of this state.
   * Being -1 means the state has no parent and is the foremost
   * "prefix" entry or the only entry.
   */
  int parent_idx = -1;
  /*! \brief The children indices of the request state entry. */
  std::vector<int> child_indices;

  /*!
   * \brief The state with regard to each model.
   * \sa RequestModelState
   */
  Array<RequestModelState> mstates;
  /*! \brief The random number generator of this request state entry. */
  RandomGenerator rng;
  /*! \brief The stop string handler of this request state entry. */
  StopStrHandler stop_str_handler;
  /*!
   * \brief The start position of the committed tokens in the
   * next request stream callback invocation.
   */
  int next_callback_token_pos;

  /*! \brief The extra string to prepend the output. */
  std::string extra_prefix_string;

  std::vector<int32_t> token_ids_for_prefix_cache_update;

  /*!
   * \brief Back reference to the request state.
   * Use ObjectRef to avoid circulate reference.
   */
  RequestStateNode* rstate = nullptr;

  /*!
   * \brief Get the delta token ids and the logprob JSON strings for this request to return since
   * the last time calling into this function, and return the finish reason if the request
   * generation has finished.
   * \note This function follows the destination passing style, which means it writes the
   * output into the "idx"-th slot in "delta_stream_output".
   * We adopt the destination passing style to reduce the CPU data structure allocation and
   * construction overhead.
   * \param tokenizer The tokenizer for logprob process.
   * \param max_single_sequence_length The maximum allowed single sequence length.
   * \param delta_stream_output The delta token ids to return, the logprob JSON strings
   * of each delta token id, and the optional finish reason.
   * \param idx The index denoting which slot to write results in "delta_request_return".
   */
  void GetDeltaRequestReturn(const Tokenizer& tokenizer, int64_t max_single_sequence_length,
                             RequestStreamOutput* delta_stream_output, int idx);

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RequestStateEntryNode>();
  }

  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.RequestStateEntry", RequestStateEntryNode, Object);
};

class RequestStateEntry : public ObjectRef {
 public:
  explicit RequestStateEntry(Request request, int num_models, int64_t internal_id, int rng_seed,
                             const std::vector<std::string>& token_table,
                             const std::optional<xgrammar::CompiledGrammar>& compiled_grammar,
                             int parent_idx = -1);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestStateEntry, ObjectRef, RequestStateEntryNode);
};

/*! \brief A request's state, which groups all the request state entries. */
class RequestStateNode : public Object {
 public:
  /*! \brief the request state entries */
  std::vector<RequestStateEntry> entries;
  /*! \brief tracks the request metrics. */
  RequestMetrics metrics;
  /*!
   * \brief The post-process data structures.
   * We make it a state to avoid repetitive memory allocation/free in the action post process.
   */
  RequestActionPostProcWorkspace postproc_states;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RequestStateNode>();
  }

  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.RequestState", RequestStateNode, Object);
};

class RequestState : public ObjectRef {
 public:
  /*!
   * \brief Request state constructor. We take the number of response (namely "n" in OpenAI
   * API) to pre-allocate all the data structure, in order to reduce the CPU data structure
   * allocation overhead when updating the request state.
   */
  explicit RequestState(std::vector<RequestStateEntry> entries, int num_response,
                        std::chrono::high_resolution_clock::time_point add_time_point);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestState, ObjectRef, RequestStateNode);
};

}  // namespace serve
}  // namespace llm
}  // namespace mlc

#endif  // MLC_LLM_SERVE_REQUEST_STATE_H_