/*!
 *  Copyright (c) 2023-2025 by Contributors
 * \file serve/engine.cc
 * \brief The implementation for runtime module of serving engine module in MLC LLM.
 */
#include "engine.h"

#include <dlpack/dlpack.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/support/cuda/nvtx.h>
#include <xgrammar/xgrammar.h>

#include <cstdlib>
#include <functional>
#include <numeric>
#include <optional>
#include <tuple>
#include <unordered_set>

#include "../support/json_parser.h"
#include "../support/module_vtable.h"
#include "../support/result.h"
#include "../support/threading_backend.h"
#include "../support/utils.h"
#include "../tokenizers/tokenizers.h"
#include "engine_actions/action.h"
#include "engine_actions/action_commons.h"
#include "engine_state.h"
#include "event_trace_recorder.h"
#include "logit_processor.h"
#include "model.h"
#include "request.h"
#include "request_state.h"
#include "sampler/sampler.h"

namespace mlc {
namespace llm {
namespace serve {

using tvm::Device;
using namespace tvm::runtime;
using tvm::ffi::Function;
using tvm::support::NVTXScopedRange;

class EngineModule;

// get tokenizer info from model config
inline std::optional<TokenizerInfo> GetTokenizerInfo(const tvm::ffi::json::Object& model_config) {
  if (model_config.count("tokenizer_info") == 0) {
    LOG(WARNING) << "Tokenizer info not found in mlc-chat-config.json. "
                 << "Trying to automatically detect the tokenizer info";
    return std::nullopt;
  }
  const tvm::ffi::json::Object& tokenizer_info_obj =
      model_config.at("tokenizer_info").cast<tvm::ffi::json::Object>();
  auto info = tvm::ffi::make_object<TokenizerInfoNode>();
  if (tokenizer_info_obj.count("token_postproc_method")) {
    info->token_postproc_method =
        tokenizer_info_obj.at("token_postproc_method").cast<std::string>();
  }
  if (tokenizer_info_obj.count("prepend_space_in_encode")) {
    info->prepend_space_in_encode = tokenizer_info_obj.at("prepend_space_in_encode").cast<bool>();
  }
  if (tokenizer_info_obj.count("strip_space_in_decode")) {
    info->strip_space_in_decode = tokenizer_info_obj.at("strip_space_in_decode").cast<bool>();
  }
  return TokenizerInfo(info);
}

inline std::pair<std::optional<std::string>, int> GetEnvSocketHostPort() {
  char* host_str = std::getenv("MLC_SOCKET_HOST");
  char* port_str = std::getenv("MLC_SOCKET_PORT");
  if (host_str == nullptr || port_str == nullptr) {
    return {std::nullopt, -1};
  }
  std::string host(host_str);
  if (host.empty()) {
    return {std::nullopt, -1};
  }
  return {host, std::atoi(port_str)};
}

// string back error node
void StreamBackErrorImpl(Request request, FRequestStreamCallback request_stream_callback,
                         String finish_reason) {
  // If the request input length exceeds the maximum allowed single sequence length,
  // invoke callback and do not process the request.
  Array<RequestStreamOutput> output{RequestStreamOutput(
      request->id, std::vector<std::vector<int64_t>>(request->generation_cfg->n), std::nullopt,
      std::vector<Optional<String>>(request->generation_cfg->n, finish_reason),
      std::vector<String>(request->generation_cfg->n))};
  // NOTE: Invariant requirement
  // always stream back final usage
  // otherwise frontend may have issues deciding
  String dummy_usage = ("{ \"prompt_tokens\": 0, \"completion_tokens\": 0, \"total_tokens\": 0 }");
  output.push_back(RequestStreamOutput::Usage(request->id, dummy_usage));
  if (request_stream_callback != nullptr) {
    request_stream_callback(output);
  }
}

void AbortRequestImpl(EngineState estate, const Array<Model>& models, const String& request_id,
                      String finish_reason) {
  auto it_rstate = estate->request_states.find(request_id);
  if (it_rstate == estate->request_states.end()) {
    // The request to abort does not exist.
    return;
  }

  RequestState rstate = it_rstate->second;
  Request request = rstate->entries[0]->request;

  // - Check if the request is running or pending.
  auto it_running = std::find(estate->running_queue.begin(), estate->running_queue.end(), request);
  auto it_waiting = std::find(estate->waiting_queue.begin(), estate->waiting_queue.end(), request);

  estate->request_states.erase(request->id);
  if (it_running != estate->running_queue.end()) {
    // The request to abort is in running queue
    estate->running_queue.erase(it_running);

    for (int i = static_cast<int>(rstate->entries.size()) - 1; i >= 0; --i) {
      if (estate->prefix_cache->HasSequence(rstate->entries[i]->mstates[0]->internal_id)) {
        estate->prefix_cache->RecycleSequence(rstate->entries[i]->mstates[0]->internal_id,
                                              /*lazy=*/false);
      } else {
        if (rstate->entries[i]->status != RequestStateStatus::kAlive) {
          estate->id_manager.RecycleId(rstate->entries[i]->mstates[0]->internal_id);
          continue;
        }
        RemoveRequestFromModel(estate, rstate->entries[i]->mstates[0]->internal_id, models);
        estate->id_manager.RecycleId(rstate->entries[i]->mstates[0]->internal_id);
      }
    }
  }
  if (it_waiting != estate->waiting_queue.end()) {
    // The request to abort is in waiting queue
    estate->waiting_queue.erase(it_waiting);
  }
  // Todo: abortion when the request is not in either queue?

  // Send a callback to notice the abortion.
  StreamBackErrorImpl(request, estate->request_stream_callback_, finish_reason);
  estate->running_rsentries_changed = true;
}

/*!
 *  \brief This a mock engine that always echo back the inputs
 *   and attaches the generation config to usage.extra
 *
 * \note: mock engine test cannot replace real engine test.
 *
 * It only tests that parameters are converted and
 * passed correctly to the backend.
 */
class MockEchoEngineImpl : public Engine {
 public:
  static Result<EngineCreationOutput> Create(const std::string& engine_config_json_str,
                                             FRequestStreamCallback request_stream_callback,
                                             const tvm::ffi::json::Object& model_config) {
    using TResult = Result<EngineCreationOutput>;
    // set dummy values
    InferrableEngineConfig inferrable_config;
    inferrable_config.max_num_sequence = 32;
    inferrable_config.max_total_sequence_length = 32 * 4096;
    inferrable_config.max_single_sequence_length = 4096;
    inferrable_config.prefill_chunk_size = 1024;
    inferrable_config.max_history_size = 1024;
    tvm::ffi::String err;
    auto config_json = tvm::ffi::json::Parse(engine_config_json_str, &err);
    if (!err.empty()) {
      return TResult::Error(err);
    }
    EngineConfig engine_config = EngineConfig::FromJSONAndInferredConfig(
        config_json.cast<tvm::ffi::json::Object>(), inferrable_config);

    auto n = std::make_unique<MockEchoEngineImpl>();
    n->request_stream_callback_ = request_stream_callback;
    n->tokenizer_ = Tokenizer::FromPath(engine_config->model, GetTokenizerInfo(model_config));
    // - Get the default generation config from the first model.
    GenerationConfig default_generation_cfg =
        GenerationConfig::GetDefaultFromModelConfig(model_config);
    return TResult::Ok({std::move(n), std::move(engine_config), std::move(default_generation_cfg)});
  }

  void Reset() final {}

  bool Empty() final { return request_map_.empty(); }

  void SetRequestStreamCallback(FRequestStreamCallback request_stream_callback) final {
    request_stream_callback_ = request_stream_callback;
  }

  FRequestStreamCallback GetRequestStreamCallback() final { return request_stream_callback_; }

  void AddRequest(Request request) final {
    // precompute the stream back results and store them in the request_map
    request = Request::FromUntokenized(request, tokenizer_);
    std::vector<RequestStreamOutput> outputs;
    int64_t completion_tokens = 0;
    int64_t prompt_tokens = 0;

    for (Data input : request->inputs) {
      // only stream back token data
      if (auto* token_data = input.as<TokenDataNode>()) {
        for (int64_t token_id : token_data->token_ids) {
          prompt_tokens += 1;
          completion_tokens += 1;
          if (request->generation_cfg->max_tokens == -1 ||
              completion_tokens <= request->generation_cfg->max_tokens) {
            outputs.push_back(RequestStreamOutput(
                request->id,
                std::vector<std::vector<int64_t>>(request->generation_cfg->n, {token_id}),
                std::nullopt,
                std::vector<Optional<String>>(request->generation_cfg->n, std::nullopt),
                std::vector<String>(request->generation_cfg->n)));
          }
        }
      }
    }

    // output go beyond max tokens
    String finish_reason = "stop";
    if (request->generation_cfg->max_tokens != -1 &&
        prompt_tokens > request->generation_cfg->max_tokens) {
      finish_reason = "length";
    }
    std::vector<std::vector<int64_t>> group_delta_token_ids;

    // correct the last output with right finish reason
    if (outputs.size() > 0) {
      group_delta_token_ids = outputs.back()->group_delta_token_ids;
      outputs.pop_back();
    }
    outputs.push_back(RequestStreamOutput(
        request->id, group_delta_token_ids, std::nullopt,
        std::vector<Optional<String>>(request->generation_cfg->n, finish_reason),
        std::vector<String>(request->generation_cfg->n)));

    // attach usage and config
    tvm::ffi::json::Object usage;
    usage.Set("prompt_tokens", static_cast<int64_t>(prompt_tokens));
    usage.Set("completion_tokens",
              static_cast<int64_t>(completion_tokens * request->generation_cfg->n));
    usage.Set("total_tokens",
              static_cast<int64_t>(prompt_tokens + completion_tokens * request->generation_cfg->n));
    usage.Set("extra", request->generation_cfg->AsJSON());
    // NOTE: Invariant requirement
    // always stream back final usage
    // otherwise frontend may have issues deciding termination
    outputs.push_back(RequestStreamOutput::Usage(request->id, tvm::ffi::json::Stringify(usage)));
    // reverse the stream back so we can just pop back and get out
    std::reverse(outputs.begin(), outputs.end());

    request_map_[request->id] = MockRequestState{request, std::move(outputs)};
  }

  void AbortRequest(const String& request_id) {
    auto it = request_map_.find(request_id);
    if (it == request_map_.end()) return;
    Request request = it->second.request;

    // If the request input length exceeds the maximum allowed single sequence length,
    // invoke callback and do not process the request.
    Array<RequestStreamOutput> output{RequestStreamOutput(
        request_id, std::vector<std::vector<int64_t>>(request->generation_cfg->n), std::nullopt,
        std::vector<Optional<String>>(request->generation_cfg->n, String("abort")),
        std::vector<String>(request->generation_cfg->n))};
    // NOTE: Invariant requirement
    // always stream back final usage
    // otherwise frontend may have issues deciding
    String dummy_usage =
        ("{ \"prompt_tokens\": 0, \"completion_tokens\": 0, \"total_tokens\": 0 }");
    output.push_back(RequestStreamOutput::Usage(request->id, dummy_usage));
    request_map_.erase(it);
    if (request_stream_callback_ != nullptr) {
      request_stream_callback_(output);
    }
  }

  void AbortAllRequests() final {
    // avoid deletion during iteraton
    std::vector<String> request_ids;
    for (const auto& kv : request_map_) {
      request_ids.push_back(kv.first);
    }
    for (String req_id : request_ids) {
      AbortRequest(req_id);
    }
  }

  void Step() final {
    Array<RequestStreamOutput> outputs;
    std::vector<String> finished_request_ids;
    for (auto& kv : request_map_) {
      MockRequestState& state = kv.second;
      TVM_FFI_ICHECK_GE(state.reversed_outputs.size(), 2);
      if (state.reversed_outputs.size() == 2) {
        outputs.push_back(state.reversed_outputs.back());
        state.reversed_outputs.pop_back();
        outputs.push_back(state.reversed_outputs.back());
        finished_request_ids.push_back(kv.first);
      } else {
        outputs.push_back(state.reversed_outputs.back());
        state.reversed_outputs.pop_back();
      }
    }
    for (String req_id : finished_request_ids) {
      request_map_.erase(req_id);
    }
    if (request_stream_callback_ != nullptr) {
      request_stream_callback_(outputs);
    }
  }

  /************** Debug/Profile **************/

  /*! \brief Internal engine metrics. */
  String JSONMetrics() final { return "{}"; }

  /*! \brief Call the given global function on all workers. Only for debug purpose. */
  void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) final {}

 private:
  struct MockRequestState {
    Request request;
    std::vector<RequestStreamOutput> reversed_outputs;
  };

  // internal tokenizer
  // keep for future usage, in case we want to echo back the tokens
  Tokenizer tokenizer_;
  // callback stream
  FRequestStreamCallback request_stream_callback_;
  // active requests
  std::unordered_map<String, MockRequestState> request_map_;
};

/********************** Engine Impl **********************/

/*! \brief The implementation of Engine. */
class EngineImpl : public Engine {
  friend class EngineModule;

 public:
  /********************** Engine Management **********************/

  static Result<EngineCreationOutput> Create(const std::string& engine_config_json_str,
                                             DLDevice device,
                                             FRequestStreamCallback request_stream_callback,
                                             Optional<EventTraceRecorder> trace_recorder) {
    using TResult = Result<EngineCreationOutput>;
    std::unique_ptr<EngineImpl> n = std::make_unique<EngineImpl>();

    // - Read the models and model libs from the EngineConfig JSON string.
    Result<std::vector<std::pair<std::string, std::string>>> models_and_model_libs_res =
        EngineConfig::GetModelsAndModelLibsFromJSONString(engine_config_json_str);
    if (models_and_model_libs_res.IsErr()) {
      return TResult::Error(models_and_model_libs_res.UnwrapErr());
    }
    std::vector<std::pair<std::string, std::string>> models_and_model_libs =
        models_and_model_libs_res.Unwrap();

    int num_model = models_and_model_libs.size();
    TVM_FFI_ICHECK_GE(num_model, 1);
    // - Initialize singleton states inside the engine.
    n->estate_->Reset();
    n->estate_->request_stream_callback_ = std::move(request_stream_callback);
    n->trace_recorder_ = trace_recorder;
    n->device_ = device;
    // - Load model config, create a shared disco session when tensor
    // parallelism is enabled.
    std::vector<std::string> model_libs;
    std::vector<tvm::ffi::json::Object> model_configs;
    model_libs.reserve(num_model);
    model_configs.reserve(num_model);
    for (int i = 0; i < num_model; ++i) {
      const auto& [model_str, model_lib] = models_and_model_libs[i];
      Result<tvm::ffi::json::Object> model_config_res = Model::LoadModelConfig(model_str);
      if (model_config_res.IsErr()) {
        return TResult::Error("Model " + std::to_string(i) +
                              " has invalid mlc-chat-config.json: " + model_config_res.UnwrapErr());
      }
      model_libs.push_back(model_lib);
      model_configs.push_back(model_config_res.Unwrap());
    }

    // kick in mock path so we don't have to load in models
    if (models_and_model_libs[0].second == "mock://echo") {
      return MockEchoEngineImpl::Create(engine_config_json_str,
                                        n->estate_->request_stream_callback_, model_configs[0]);
    }

    auto [session, num_shards, model_num_pipeline_stages] =
        n->CreateDiscoSession(model_libs, model_configs, device);

    // - Initialize each model independently.
    n->models_.clear();
    for (int i = 0; i < num_model; ++i) {
      const auto& [model_str, model_lib] = models_and_model_libs[i];
      Model model = Model::Create(model_lib, model_str, model_configs[i], device, session,
                                  num_shards, model_num_pipeline_stages[i],
                                  /*trace_enabled=*/trace_recorder.has_value());
      n->models_.push_back(model);
    }
    // - Initialize NVSHMEM
    n->estate_->disaggregation = n->models_[0]->GetMetadata().disaggregation;
    if (n->estate_->disaggregation) {
      LOG(INFO) << "Initializing NVSHMEM";
      char* nvshmem_init_config_json_char = std::getenv("MLC_NVSHMEM_INIT_CONFIG_JSON_STR");
      TVM_FFI_ICHECK(nvshmem_init_config_json_char != nullptr)
          << "The environment variables MLC_NVSHMEM_INIT_CONFIG_JSON_STR should be set.";
      std::string f_name = "runtime.disco.nvshmem.init_nvshmem_wrapper";
      if (session != nullptr) {
        n->DebugCallFuncOnAllAllWorker(f_name, String(nvshmem_init_config_json_char));
      } else {
        static Function func = Function::GetGlobalRequired(f_name);
        func(String(nvshmem_init_config_json_char));
      }
      LOG(INFO) << "NVSHMEM initialized successfully.";
    }

    // - Automatically infer the missing fields in EngineConfig JSON strings
    // and get the final EngineConfig.
    Result<EngineConfig> engine_config_res =
        n->AutoDecideEngineConfig(engine_config_json_str, model_configs);
    if (engine_config_res.IsErr()) {
      return TResult::Error(engine_config_res.UnwrapErr());
    }
    EngineConfig engine_config = engine_config_res.Unwrap();
    {
      if (engine_config->prefix_cache_mode == PrefixCacheMode::kRadix) {
        n->estate_->prefix_cache = PrefixCache::CreateRadixPrefixCache(
            static_cast<size_t>(engine_config->prefix_cache_max_num_recycling_seqs),
            [engine_ptr = n.get()](int64_t seq_id) {
              RemoveRequestFromModel(engine_ptr->estate_, seq_id, engine_ptr->models_);
              engine_ptr->estate_->id_manager.RecycleId(seq_id);
            });
      } else if (engine_config->prefix_cache_mode == PrefixCacheMode::kDisable) {
        n->estate_->prefix_cache = PrefixCache::CreateNoPrefixCache();
      } else {
        LOG(FATAL) << "Unsupported prefix cache mode: "
                   << static_cast<int>(engine_config->prefix_cache_mode);
      }
      if (engine_config->speculative_mode != SpeculativeMode::kDisable &&
          engine_config->prefill_mode == PrefillMode::kHybrid) {
        engine_config->prefill_mode = PrefillMode::kChunked;
        LOG(WARNING)
            << "Hybrid prefill mode fallbacks to chunked prefill, due to speculative mode is "
               "enabled and not implemented with hybrid prefill yet.";
      }
    }
    // - Load model weights, create KV cache and workspace.
    n->model_workspaces_.clear();
    for (const Model& model : n->models_) {
      model->LoadParams();
      model->SetMaxNumSequence(engine_config->max_num_sequence);
      model->SetPrefillChunkSize(engine_config->prefill_chunk_size);
      model->CreateKVCache(engine_config->kv_cache_page_size, engine_config->max_num_sequence,
                           engine_config->max_total_sequence_length,
                           engine_config->prefill_chunk_size, engine_config->max_history_size,
                           engine_config->prefix_cache_max_num_recycling_seqs);
      n->model_workspaces_.push_back(
          ModelWorkspace{model->AllocEmbeddingTensor(), model->AllocHiddenStatesTensor()});
    }
    // - Initialize tokenizer and grammar
    n->tokenizer_ = Tokenizer::FromPath(engine_config->model, GetTokenizerInfo(model_configs[0]));
    n->token_table_ = n->tokenizer_->PostProcessedTokenTable();
    n->cached_grammar_compiler_ = xgrammar::CachedGrammarCompiler(n->token_table_);
    // - Create the logit processor and sampler, and
    // the DraftTokenWorkspaceManager for speculative decoding.
    int max_num_tokens = engine_config->max_num_sequence;
    DraftTokenWorkspaceManager draft_token_workspace_manager{nullptr};
    if (engine_config->speculative_mode != SpeculativeMode::kDisable) {
      // multiply max num_tokens by two so we can do ping-pong swaping during draft/verify process
      draft_token_workspace_manager =
          n->models_[0]->CreateDraftTokenWorkspaceManager(max_num_tokens * 2);
      draft_token_workspace_manager->AllocWorkspace(
          &n->model_workspaces_[0],
          /*require_hidden_states=*/engine_config->speculative_mode == SpeculativeMode::kEagle);
    }
    LogitProcessor logit_processor =
        n->models_[0]->CreateLogitProcessor(max_num_tokens, trace_recorder);
    Sampler sampler = n->models_[0]->CreateSampler(
        max_num_tokens, static_cast<int>(n->models_.size()), trace_recorder);
    // - Initialize engine actions that represent state transitions.
    if (engine_config->speculative_mode != SpeculativeMode::kDisable) {
      n->estate_->spec_draft_length = engine_config->spec_draft_length;
    }
    n->actions_ =
        CreateEngineActions(n->models_, engine_config, model_configs, n->model_workspaces_,
                            logit_processor, sampler, draft_token_workspace_manager, n->tokenizer_,
                            n->trace_recorder_, n->estate_->request_stream_callback_, device);
    n->draft_token_workspace_manager_ = draft_token_workspace_manager;
    // - Automatically set the threading backend max concurrency.
    n->engine_config_ = engine_config;
    n->SetThreadMaxConcurrency();
    // - Get the default generation config from the first model.
    GenerationConfig default_generation_cfg =
        GenerationConfig::GetDefaultFromModelConfig(model_configs[0]);
    return TResult::Ok({std::move(n), std::move(engine_config), std::move(default_generation_cfg)});
  }

  void Reset() final {
    AbortAllRequests();
    estate_->Reset();
    for (Model model : models_) {
      model->Reset();
    }
  }

  bool Empty() final { return estate_->running_queue.empty() && estate_->waiting_queue.empty(); }

  String JSONMetrics() final { return tvm::ffi::json::Stringify(estate_->metrics.AsJSON(), 2); }

  FRequestStreamCallback GetRequestStreamCallback() final {
    return estate_->request_stream_callback_;
  }

  void SetRequestStreamCallback(FRequestStreamCallback request_stream_callback) final {
    estate_->request_stream_callback_ = std::move(request_stream_callback);
  }

  // string back error node
  void StreamBackError(Request request, String finish_reason) {
    StreamBackErrorImpl(request, estate_->request_stream_callback_, finish_reason);
  }

  /***************** High-level Request Management *****************/

  void HandleSpecialRequests(Request request) {
    auto special_request = request->generation_cfg->debug_config.special_request;
    switch (special_request) {
      case SpecialRequestKind::kQueryEngineMetrics: {
        Array<RequestStreamOutput> output = {
            RequestStreamOutput::Usage(request->id, estate_->metrics.AsUsageJSONStr())};
        estate_->request_stream_callback_(output);
        break;
      }
      default:
        break;
    }
  }

  /*!
   * \brief Handle the given disaggregation request.
   * Return true if skipping the subsequent AddRequest process.
   */
  bool HandleDisaggRequest(Request request) {
    DisaggConfig disagg_config = request->generation_cfg->debug_config.disagg_config;
    DisaggRequestKind kind = disagg_config.kind;
    if (kind == DisaggRequestKind::kPrepareReceive) {
      // No-op.
      return false;
    } else if (kind == DisaggRequestKind::kRemoteSend) {
      int input_length = 0;
      for (Data input : request->inputs) {
        input_length += input->GetLength();
      }
      // - Truncate the inputs to the desired prefill length (specified by "end").
      int kv_window_begin = disagg_config.kv_window_begin.value_or(0);
      int kv_window_end = disagg_config.kv_window_end.value_or(input_length);
      TVM_FFI_ICHECK_GE(kv_window_begin, 0);
      if (kv_window_end < 0) {
        kv_window_end = input_length + kv_window_end;
      }
      TVM_FFI_ICHECK_LT(kv_window_end, input_length)
          << "Prefill the full input on the remote machine is not supported.";
      TVM_FFI_ICHECK_LT(kv_window_begin, kv_window_end)
          << "\"begin >= end\" is not supported by remote prefill";
      request->inputs = SplitData(request->inputs, input_length, kv_window_end).first;
      // - Check the invariant: "end - begin" equals the expanded metadata length.
      TVM_FFI_ICHECK_EQ(disagg_config.kv_append_metadata.size(), models_.size());
      for (const Shape& compressed_kv_append_metadata : disagg_config.kv_append_metadata) {
        TVM_FFI_ICHECK(!compressed_kv_append_metadata.empty());
        int num_segments = compressed_kv_append_metadata[0];
        TVM_FFI_ICHECK_EQ(compressed_kv_append_metadata.size(), num_segments * 2 + 1);
        int transmission_length = 0;
        for (int i = 0; i < num_segments; ++i) {
          transmission_length += compressed_kv_append_metadata[i * 2 + 2];
        }
        TVM_FFI_ICHECK_EQ(transmission_length, kv_window_end - kv_window_begin);
      }
      // - Override the "n" in generation config to 1.
      ObjectPtr<GenerationConfigNode> updated_generation_cfg =
          tvm::ffi::make_object<GenerationConfigNode>(*request->generation_cfg.get());
      updated_generation_cfg->n = 1;
      request->generation_cfg = GenerationConfig(updated_generation_cfg);
      return false;
    } else if (kind == DisaggRequestKind::kStartGeneration) {
      auto it_rstate = estate_->request_states.find(request->id);
      TVM_FFI_ICHECK(it_rstate != estate_->request_states.end());
      TVM_FFI_ICHECK(!it_rstate->second->entries.empty());
      request = it_rstate->second->entries[0]->request;
      TVM_FFI_ICHECK(request->generation_cfg->debug_config.disagg_config.kind ==
                     DisaggRequestKind::kPrepareReceive);
      int input_length = 0;
      for (Data input : request->inputs) {
        input_length += input->GetLength();
      }
      // - Truncate the inputs to the desired prefill length (specified by "end").
      int kv_window_begin = disagg_config.kv_window_begin.value_or(0);
      int kv_window_end = disagg_config.kv_window_end.value_or(input_length);
      TVM_FFI_ICHECK_EQ(kv_window_end, input_length);
      if (kv_window_begin < 0) {
        kv_window_begin = input_length + kv_window_begin;
      }
      TVM_FFI_ICHECK_GE(kv_window_begin, 0);
      TVM_FFI_ICHECK_LT(kv_window_begin, input_length);
      // The request is not supposed to be in running queue nor waiting queue.
      auto it_running =
          std::find(estate_->running_queue.begin(), estate_->running_queue.end(), request);
      auto it_waiting =
          std::find(estate_->waiting_queue.begin(), estate_->waiting_queue.end(), request);
      TVM_FFI_ICHECK(it_running == estate_->running_queue.end());
      TVM_FFI_ICHECK(it_waiting == estate_->waiting_queue.end());

      RequestState rstate = it_rstate->second;
      ObjectPtr<GenerationConfigNode> updated_generation_cfg =
          tvm::ffi::make_object<GenerationConfigNode>(*request->generation_cfg.get());
      // - Split the input data into two parts at the position "kv_window_begin".
      TVM_FFI_ICHECK(!request->inputs.empty());
      auto [lhs_data, rhs_data] = SplitData(request->inputs, input_length, kv_window_begin);
      if (input_length - kv_window_begin == 1 && request->generation_cfg->n == 1) {
        // - Commit the last token id to the request states.
        TVM_FFI_ICHECK_EQ(rhs_data.size(), 1);
        const auto* token_data = rhs_data.back().as<TokenDataNode>();
        TVM_FFI_ICHECK(token_data != nullptr);
        TVM_FFI_ICHECK_EQ(token_data->GetLength(), 1);
        SampleResult last_token;
        last_token.sampled_token_id = {token_data->token_ids.back(), 1.0};
        for (RequestModelState mstate : rstate->entries[0]->mstates) {
          mstate->CommitToken(last_token);
          TVM_FFI_ICHECK_EQ(mstate->committed_tokens.size(), 1);
        }
        // - Set "next_callback_token_pos" so that this token will not be streamed back to user.
        rstate->entries[0]->next_callback_token_pos = 1;
        // - Update the request input.
        request->inputs = lhs_data;
        // - Increment the max_tokens in generation config.
        if (request->generation_cfg->max_tokens != -1) {
          ++updated_generation_cfg->max_tokens;
        }
      } else {
        // Since there are multiple tokens to prefill, we add the remaining inputs
        // to the request's RequestModelStates for prefill.
        for (RequestModelState mstate : rstate->entries[0]->mstates) {
          mstate->inputs = rhs_data;
        }
        // Add to waiting queue for prefill.
        estate_->waiting_queue.insert(estate_->waiting_queue.begin(), request);
      }
      estate_->running_queue.push_back(request);
      // Erase the disaggregation request kind.
      updated_generation_cfg->debug_config.disagg_config.kind = DisaggRequestKind::kNone;
      request->generation_cfg = GenerationConfig(updated_generation_cfg);
      estate_->running_rsentries_changed = true;
      return true;
    }
    LOG(FATAL) << "Cannot reach here";
    throw;
  }

  void AddRequest(Request request) final {
    NVTXScopedRange nvtx_scope("Add request " + request->id);
    // special requests do not involve generation
    if (request->generation_cfg->debug_config.special_request != SpecialRequestKind::kNone) {
      this->HandleSpecialRequests(request);
      return;
    }

    RECORD_EVENT(trace_recorder_, request->id, "request added to engine");
    auto add_time_point = std::chrono::high_resolution_clock::now();

    // Get a request copy where all text inputs are tokenized.
    request = Request::FromUntokenized(request, tokenizer_);
    TVM_FFI_ICHECK_NE(request->prompt_tokens, -1);

    if (request->prompt_tokens >= engine_config_->max_single_sequence_length &&
        estate_->request_stream_callback_ != nullptr) {
      this->StreamBackError(request, "length");
      return;
    }

    // Handle disaggregation requests.
    if (request->generation_cfg->debug_config.disagg_config.kind != DisaggRequestKind::kNone) {
      bool return_now = this->HandleDisaggRequest(request);
      if (return_now) {
        return;
      }
    }

    // Append to the waiting queue and create the request state.
    estate_->waiting_queue.push_back(request);

    int n = request->generation_cfg->n;
    int rng_seed = request->generation_cfg->seed;
    auto compiled_grammar = GetGrammarFromResponseFormat(request->generation_cfg->response_format);

    std::vector<RequestStateEntry> rsentries;
    // Create the request state entry for the input.
    rsentries.emplace_back(request, models_.size(), estate_->id_manager.GetNewId(), rng_seed,
                           token_table_, compiled_grammar);
    if (n > 1) {
      // Then create a request state entry for each parallel generation branch.
      // We add a offset to the rng seed so that to make generations different.
      rsentries.reserve(n + 1);
      rsentries[0]->child_indices.reserve(n);
      for (int i = 0; i < n; ++i) {
        rsentries[0]->child_indices.push_back(rsentries.size());
        rsentries.emplace_back(request, models_.size(), estate_->id_manager.GetNewId(),
                               rng_seed + i + 1, token_table_, compiled_grammar,
                               /*parent_idx=*/0);
      }
    }
    RequestState rstate = RequestState(std::move(rsentries), n, add_time_point);
    for (const RequestStateEntry& rsentry : rstate->entries) {
      // Set the back reference.
      // note, we avoid cyclic reference and use raw ptr.
      rsentry->rstate = rstate.operator->();
    }
    request->rstate = rstate.operator->();
    estate_->request_states.emplace(request->id, rstate);
  }

  void AbortRequest(const String& request_id) final {
    AbortRequestImpl(estate_, models_, request_id);
  }

  void AbortAllRequests() final {
    // - Collect all the request ids.
    std::vector<String> request_ids;
    request_ids.reserve(estate_->request_states.size());
    for (const auto& kv : estate_->request_states) {
      request_ids.push_back(kv.first);
    }
    // - Abort all the requests.
    for (const String& request_id : request_ids) {
      AbortRequest(request_id);
    }
  }

  /*********************** Engine Action ***********************/

  void Step() final {
    TVM_FFI_ICHECK(estate_->request_stream_callback_ != nullptr)
        << "The request stream callback is not set. Engine cannot execute.";
    for (EngineAction action : actions_) {
      Array<Request> processed_requests;
      {
        NVTXScopedRange nvtx_scope("Action step");
        processed_requests = action->Step(estate_);
      }
      if (!processed_requests.empty()) {
        ActionStepPostProcess(processed_requests, estate_, models_, tokenizer_,
                              estate_->request_stream_callback_,
                              engine_config_->max_single_sequence_length,
                              draft_token_workspace_manager_, trace_recorder_);
        return;
      }
    }
    TVM_FFI_ICHECK(estate_->running_queue.empty())
        << "Internal assumption violated: It is expected that an engine step takes at least one "
           "action (e.g. prefill, decode, etc.) but it does not.";
  }

  /************** Utility Functions **************/
  std::tuple<Optional<Session>, int, std::vector<int>> CreateDiscoSession(
      const std::vector<std::string>& model_libs,
      const std::vector<tvm::ffi::json::Object>& model_configs, Device device) {
    const auto& base_model_config = model_configs[0];

    auto f_get_num_shards_num_stages =
        [&device](const std::string& model_lib,
                  const tvm::ffi::json::Object& model_config) -> std::pair<int, int> {
      if (!StartsWith(model_lib, "system://")) {
        Module executable = ffi::Module::LoadFromFile(model_lib);
        Optional<Function> fload_exec = executable->GetFunction("vm_load_executable");
        TVM_FFI_ICHECK(fload_exec.has_value()) << "TVM runtime cannot find vm_load_executable";
        Module local_vm = fload_exec.value()().cast<Module>();
        local_vm->GetFunction("vm_initialization")
            .value()(static_cast<int>(device.device_type), device.device_id,
                     static_cast<int>(tvm::runtime::memory::AllocatorType::kPooled),
                     static_cast<int>(kDLCPU), 0,
                     static_cast<int>(tvm::runtime::memory::AllocatorType::kPooled));
        ModelMetadata metadata = ModelMetadata::FromModule(local_vm, std::move(model_config));
        return {metadata.tensor_parallel_shards, metadata.pipeline_parallel_stages};
      } else {
        return {1, 1};
      }
    };

    int num_shards = -1;
    int max_num_stages = 1;
    std::vector<int> model_num_pipeline_stages;
    model_num_pipeline_stages.reserve(model_libs.size());
    TVM_FFI_ICHECK_EQ(model_libs.size(), model_configs.size());
    for (int i = 0; i < static_cast<int>(model_libs.size()); ++i) {
      auto [model_num_shards, model_num_stages] =
          f_get_num_shards_num_stages(model_libs[i], model_configs[i]);
      model_num_pipeline_stages.push_back(model_num_stages);
      max_num_stages = std::max(max_num_stages, model_num_stages);
      if (i == 0) {
        num_shards = model_num_shards;
      } else {
        TVM_FFI_ICHECK_EQ(model_num_shards, num_shards)
            << "Inconsistent tensor_parallel_shards values across models. Some model is compiled "
               "with tensor_parallel_shards "
            << num_shards << " and some other model is compiled with tensor_parallel_shards "
            << model_num_shards;
      }
    }

    Optional<Session> session = std::nullopt;
    int num_workers = num_shards * max_num_stages;
    if (num_workers > 1) {
#ifndef MLC_SINGLE_GPU_ONLY
      constexpr const char* f_create_process_pool = "runtime.disco.create_process_pool";
      if (!Function::GetGlobal(f_create_process_pool).has_value()) {
        LOG(FATAL) << "Cannot find process launcher `" << f_create_process_pool << "`. "
                   << "Multi-GPU inference depends on MLC LLM Python API to launch process.";
      }
      std::string ccl;
      if (device.device_type == kDLCUDA) {
        ccl = "nccl";
      } else if (device.device_type == kDLROCM) {
        ccl = "rccl";
      } else {
        LOG(FATAL) << "ValueError: Multi-GPU on device " << DLDeviceType2Str(device.device_type)
                   << " is not supported. Currently, only NCCL and RCCL are integrated.";
      }
      std::vector<int64_t> device_ids(num_workers);
      for (int i = 0; i < num_workers; ++i) {
        // device.device_id is the start of the worker 0 of this model
        device_ids[i] = device.device_id + i;
      }
      const std::string& green_text_begin = "\033[92m";
      const std::string& yellow_text_begin = "\033[93m";
      const std::string& colored_text_end = "\033[0m";
      auto [socket_host, socket_port] = GetEnvSocketHostPort();
      if (max_num_stages > 1 && socket_host.has_value()) {
        // Use SocketSession when pipeline parallelism enabled and socket host and port are set.
        TVM_FFI_ICHECK_GT(socket_port, 0)
            << "Invalid MLC socket port " << socket_port
            << ". Please set a valid port value in environment variable \"MLC_SOCKET_PORT\".";
        LOG(INFO) << "Creating MLC socket session with socket host " << socket_host.value()
                  << " and port " << socket_port;
        LOG(INFO) << "Please launch " << green_text_begin << max_num_stages - 1 << colored_text_end
                  << " remote socket node(s) with the following command to proceed:\n\t"
                  << green_text_begin << "python -m mlc_llm.cli.disco_remote_socket_session "
                  << (socket_host.value() == "0.0.0.0" ? "<YOUR_NODE_IP>" : socket_host.value())
                  << " " << socket_port << " " << num_shards << colored_text_end;
        static Function f_create_socket_sess =
            Function::GetGlobalRequired("runtime.disco.SocketSession");
        Session sess =
            f_create_socket_sess(max_num_stages, num_shards, /*num_groups=*/max_num_stages,
                                 socket_host.value(), socket_port)
                .cast<Session>();
        session = std::move(sess);
      } else {
        if (max_num_stages > 1) {
          LOG(INFO)
              << yellow_text_begin
              << "Model is enabled with \"pipeline_parallel_stages\" but the socket host/port is "
                 "not set. If you intend to run the model on multiple nodes, please set "
                 "environment variable \"MLC_SOCKET_HOST\" and \"MLC_SOCKET_PORT\" and run again."
              << colored_text_end;
        }
        // Use ProcessSession otherwise.
        session = Session::ProcessSession(num_workers, max_num_stages, f_create_process_pool,
                                          "mlc_llm.cli.worker");
      }
      session.value()->InitCCL(ccl, Shape(device_ids));
#else
      LOG(FATAL) << "MLC_SINGLE_GPU_ONLY is specified. Multi-GPU is not enabled.";
#endif  // MLC_SINGLE_GPU_ONLY
    }
    return {session, num_shards, model_num_pipeline_stages};
  }

  /************** Debug/Profile **************/

  void DebugCallFuncOnAllAllWorker(const String& func_name, Optional<String> func_args) final {
    TVM_FFI_ICHECK(!models_.empty()) << "There is no model running in Engine.";
    models_[0]->DebugCallFuncOnAllAllWorker(func_name, func_args);
  }

 private:
  Result<EngineConfig> AutoDecideEngineConfig(
      const std::string& engine_config_json_str,
      const std::vector<tvm::ffi::json::Object>& model_configs) {
    using TResult = Result<EngineConfig>;
    tvm::ffi::String err;
    auto config_json = tvm::ffi::json::Parse(engine_config_json_str, &err);
    if (!err.empty()) {
      return TResult::Error(err);
    }
    tvm::ffi::json::Object config = config_json.cast<tvm::ffi::json::Object>();
    ObjectPtr<EngineConfigNode> n = tvm::ffi::make_object<EngineConfigNode>();

    // - Get the engine mode and maximum GPU utilization for inference.
    EngineMode mode = EngineModeFromString(json::Lookup<std::string>(config, "mode"));
    double gpu_memory_utilization =
        json::LookupOrDefault<double>(config, "gpu_memory_utilization", n->gpu_memory_utilization);
    bool verbose = json::LookupOrDefault<bool>(config, "verbose", n->verbose);

    // - Get the config fields that can be automatically inferred.
    std::optional<int64_t> max_num_sequence =
        json::LookupOptional<int64_t>(config, "max_num_sequence");
    std::optional<int64_t> max_total_sequence_length =
        json::LookupOptional<int64_t>(config, "max_total_sequence_length");
    std::optional<int64_t> max_single_sequence_length =
        json::LookupOptional<int64_t>(config, "max_single_sequence_length");
    std::optional<int64_t> prefill_chunk_size =
        json::LookupOptional<int64_t>(config, "prefill_chunk_size");
    std::optional<int64_t> max_history_size =
        json::LookupOptional<int64_t>(config, "max_history_size");
    std::optional<std::string> kv_state_kind_str =
        json::LookupOptional<std::string>(config, "kv_state_kind");
    InferrableEngineConfig inferrable_cfg{max_num_sequence, max_total_sequence_length,
                                          max_single_sequence_length, prefill_chunk_size,
                                          max_history_size};

    // - Get the model metadata.
    std::vector<ModelMetadata> model_metadata;
    for (const Model& model : models_) {
      model_metadata.push_back(model->GetMetadata());
    }
    // - Select from kv cache or RNN state.
    Result<bool> use_kv_cache = ModelsUseKVCache(model_configs);
    if (use_kv_cache.IsErr()) {
      return TResult::Error(use_kv_cache.UnwrapErr());
    }
    Result<InferrableEngineConfig> inferrable_cfg_res;
    if (use_kv_cache.Unwrap()) {
      // - Infer configuration.
      inferrable_cfg_res = InferrableEngineConfig::InferForKVCache(
          mode, device_, gpu_memory_utilization, model_configs, model_metadata, inferrable_cfg,
          verbose);
    } else {
      // - Infer configuration.
      inferrable_cfg_res = InferrableEngineConfig::InferForRNNState(
          mode, device_, gpu_memory_utilization, model_configs, model_metadata, inferrable_cfg,
          verbose);
    }

    if (inferrable_cfg_res.IsErr()) {
      return TResult::Error(inferrable_cfg_res.UnwrapErr());
    }
    inferrable_cfg = inferrable_cfg_res.Unwrap();
    TVM_FFI_ICHECK(inferrable_cfg.max_num_sequence.has_value());
    TVM_FFI_ICHECK(inferrable_cfg.max_total_sequence_length.has_value());
    use_kv_cache = ModelsUseKVCache(model_configs);
    if (use_kv_cache.Unwrap()) {
      TVM_FFI_ICHECK(inferrable_cfg.max_single_sequence_length.has_value());
    }
    TVM_FFI_ICHECK(inferrable_cfg.prefill_chunk_size.has_value());
    TVM_FFI_ICHECK(inferrable_cfg.max_history_size.has_value());
    return TResult::Ok(EngineConfig::FromJSONAndInferredConfig(config, inferrable_cfg));
  }

  /*! \brief Set the maximum threading backend concurrency. */
  void SetThreadMaxConcurrency() {
    int host_cpu_usage = 1;
    for (Model model : models_) {
      host_cpu_usage += model->EstimateHostCPURequirement();
    }
    if (host_cpu_usage > 1) {
      int max_concurrency = tvm::runtime::threading::MaxConcurrency();
      tvm::runtime::threading::SetMaxConcurrency(std::min(
          std::max(max_concurrency - host_cpu_usage, 1), engine_config_->max_num_sequence));
    }
  }

  /*! \brief Create a grammar init context according to the response format. If the response format
   * is not JSON, return std::nullopt. */
  std::optional<xgrammar::CompiledGrammar> GetGrammarFromResponseFormat(
      const ResponseFormat& response_format) {
    if (response_format.type != "json_object") {
      return std::nullopt;
    } else if (!response_format.schema) {
      return cached_grammar_compiler_.GetCompiledGrammarForJSON();
    } else {
      return cached_grammar_compiler_.GetCompiledGrammarForJSONSchema(
          response_format.schema.value());
    }
  }

  // Engine state, managing requests and request states.
  EngineState estate_;
  // Configurations and singletons
  EngineConfig engine_config_;
  // internal tokenizer
  Tokenizer tokenizer_;
  std::vector<std::string> token_table_;
  // Cached grammar compiler for grammar matching.
  xgrammar::CachedGrammarCompiler cached_grammar_compiler_;
  // Models
  Array<Model> models_;
  // Device that the models run on.
  Device device_;
  // Workspace of each model.
  std::vector<ModelWorkspace> model_workspaces_;
  // Engine actions.
  Array<EngineAction> actions_;
  // Draft token workspace manager for speculative decoding.
  Optional<DraftTokenWorkspaceManager> draft_token_workspace_manager_;
  // Event trace recorder.
  Optional<EventTraceRecorder> trace_recorder_;
};

Result<EngineCreationOutput> Engine::Create(const std::string& engine_config_json_str,
                                            Device device,
                                            FRequestStreamCallback request_stream_callback,
                                            Optional<EventTraceRecorder> trace_recorder) {
  return EngineImpl::Create(engine_config_json_str, device, request_stream_callback,
                            std::move(trace_recorder));
}

/*! \brief Clear global memory manager */
void ClearGlobalMemoryManager() {
  static const char* kFunc = "vm.builtin.memory_manager.clear";
  static Function f = Function::GetGlobalRequired(kFunc);
  f();
}

class EngineModule : public ffi::ModuleObj {
 public:
  TVM_MODULE_VTABLE_BEGIN("mlc.serve.engine");
  TVM_MODULE_VTABLE_ENTRY("init", &EngineModule::Init);
  TVM_MODULE_VTABLE_ENTRY("add_request", &EngineModule::AddRequest);
  TVM_MODULE_VTABLE_ENTRY("create_request", &EngineModule::CreateRequest);
  TVM_MODULE_VTABLE_ENTRY("abort_request", &EngineModule::Abort);
  TVM_MODULE_VTABLE_ENTRY("step", &EngineModule::Step);
  TVM_MODULE_VTABLE_ENTRY("reset", &EngineModule::Reset);
  TVM_MODULE_VTABLE_ENTRY("json_metrics", &EngineModule::JSONMetrics);
  TVM_MODULE_VTABLE_ENTRY("get_request_stream_callback", &EngineModule::GetRequestStreamCallback);
  TVM_MODULE_VTABLE_ENTRY("set_request_stream_callback", &EngineModule::SetRequestStreamCallback);
  TVM_MODULE_VTABLE_END();

  /*! \brief Initialize the engine with config and other fields. */
  void Init(const std::string& engine_config_json_str, Device device,
            FRequestStreamCallback request_stream_callback,
            Optional<EventTraceRecorder> trace_recorder) {
    Result<EngineCreationOutput> output_res = Engine::Create(
        engine_config_json_str, device, request_stream_callback, std::move(trace_recorder));
    TVM_FFI_ICHECK(output_res.IsOk()) << output_res.UnwrapErr();
    EngineCreationOutput output = output_res.Unwrap();
    this->engine_ = std::move(output.reloaded_engine);
    this->default_generation_config_ = output.default_generation_cfg;
  }
  /*! \brief Construct an EngineModule. */
  static ffi::Module Create() { return ffi::Module(tvm::ffi::make_object<EngineModule>()); }
  /*! \brief Redirection to `Engine::AddRequest`. */
  void AddRequest(Request request) { return GetEngine()->AddRequest(std::move(request)); }
  /*! \brief Redirection to `Engine::AbortRequest`. */
  void Abort(const String& request_id) { return GetEngine()->AbortRequest(request_id); }
  /*! \brief Create request with given arguments and the engine default generation config. */
  Request CreateRequest(String id, Array<Data> inputs, String generation_cfg_json_str) {
    auto config = json::ParseToJSONObject(generation_cfg_json_str);
    auto gen_config = GenerationConfig::FromJSON(config, default_generation_config_);
    TVM_FFI_ICHECK(gen_config.IsOk()) << gen_config.UnwrapErr();
    return Request(std::move(id), std::move(inputs), gen_config.Unwrap());
  }
  /*! \brief Redirection to `Engine::Step`. */
  void Step() { return GetEngine()->Step(); }
  /*! \brief Redirection to `Engine::GetRequestStreamCallback`. */
  FRequestStreamCallback GetRequestStreamCallback() {
    return GetEngine()->GetRequestStreamCallback();
  }
  /*! \brief Redirection to `Engine::SetRequestStreamCallback` */
  void SetRequestStreamCallback(FRequestStreamCallback request_stream_callback) {
    GetEngine()->SetRequestStreamCallback(request_stream_callback);
  }
  /*! \brief Redirection to `Engine::Reset`. */
  void Reset() { return GetEngine()->Reset(); }

  /*! \brief Redirection to `Engine::JSONMetrics`. */
  String JSONMetrics() { return GetEngine()->JSONMetrics(); }

 private:
  Engine* GetEngine() {
    TVM_FFI_ICHECK(engine_ != nullptr) << "Engine is not initialized via init";
    return engine_.get();
  }

  std::unique_ptr<Engine> engine_ = nullptr;
  GenerationConfig default_generation_config_;
};

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("mlc.serve.create_engine", EngineModule::Create);
}

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