/*!
 *  Copyright (c) 2023-2025 by Contributors
 * \file serve/config.cc
 */
#include "config.h"

#include <tvm/ffi/function.h>
#include <tvm/runtime/device_api.h>

#include <limits>
#include <random>

#include "../json_ffi/openai_api_protocol.h"
#include "../support/json_parser.h"
#include "../support/utils.h"
#include "data.h"

namespace mlc {
namespace llm {
namespace serve {

TVM_FFI_STATIC_INIT_BLOCK() {
  GenerationConfigNode::RegisterReflection();
  EngineConfigNode::RegisterReflection();
}

uint64_t TotalDetectGlobalMemory(DLDevice device) {
  // Get single-card GPU size.
  tvm::ffi::Any rv;
  DeviceAPI::Get(device)->GetAttr(device, DeviceAttrKind::kTotalGlobalMemory, &rv);
  int64_t gpu_size_bytes = rv.cast<int64_t>();
  // Since the memory size returned by the OpenCL runtime is smaller than the actual available
  // memory space, we set a best available space so that MLC LLM can run 7B or 8B models on Android
  // with OpenCL.
  if (device.device_type == kDLOpenCL) {
    int64_t min_size_bytes = 5LL * 1024 * 1024 * 1024;  //  Minimum size is 5 GB
    gpu_size_bytes = std::max(gpu_size_bytes, min_size_bytes);
  }
  return gpu_size_bytes;
}

/****************** ResponseFormat ******************/

Result<ResponseFormat> ResponseFormat::FromJSON(const tvm::ffi::json::Object& config) {
  using TResult = Result<ResponseFormat>;
  ResponseFormat res;
  res.type = json::LookupOrDefault<std::string>(config, "type", "text");

  std::optional<std::string> schema = json::LookupOptional<std::string>(config, "schema");
  if (schema.has_value()) {
    res.schema = schema.value();
  }

  if (res.type != "text" && res.type != "function" && res.type != "json_object") {
    return TResult::Error("Uknonwn response_format type " + res.type);
  }

  return TResult::Ok(res);
}

tvm::ffi::json::Object ResponseFormat::AsJSON() const {
  tvm::ffi::json::Object config;
  config.Set("type", type);
  if (schema.has_value()) {
    config.Set("schema", schema.value());
  }
  return config;
}

/****************** DisaggConfig ******************/

Result<DisaggConfig> DisaggConfig::FromJSON(const tvm::ffi::json::Object& config) {
  using TResult = Result<DisaggConfig>;
  DisaggConfig res;
  std::optional<std::string> kind = json::LookupOptional<std::string>(config, "kind");
  if (kind.has_value()) {
    if (kind.value() == "prepare_receive") {
      res.kind = DisaggRequestKind::kPrepareReceive;
    } else if (kind.value() == "remote_send") {
      res.kind = DisaggRequestKind::kRemoteSend;
    } else if (kind.value() == "start_generation") {
      res.kind = DisaggRequestKind::kStartGeneration;
    } else {
      return TResult::Error("Unknown disaggregation request kind " + kind.value());
    }
  }
  std::optional<std::string> kv_append_metadata_encoded =
      json::LookupOptional<std::string>(config, "kv_append_metadata");
  if (kv_append_metadata_encoded.has_value()) {
    tvm::ffi::String err;
    auto parse_result =
        tvm::ffi::json::Parse(Base64Decode(kv_append_metadata_encoded.value()), &err);
    if (!err.empty()) {
      return TResult::Error("kv_append_metadata parse error: " + std::string(err));
    }
    if (!parse_result.try_cast<tvm::ffi::json::Array>().has_value()) {
      return TResult::Error("kv_append_metadata is not array of integer.");
    }
    tvm::ffi::json::Array kv_append_metadata_arr = parse_result.cast<tvm::ffi::json::Array>();
    std::vector<Shape> kv_append_metadata;
    int ptr = 0;
    while (ptr < static_cast<int>(kv_append_metadata_arr.size())) {
      if (!kv_append_metadata_arr[ptr].try_cast<int64_t>().has_value()) {
        return TResult::Error("Invalid kv append metadata value in kv_append_metadata array");
      }
      int num_segments = kv_append_metadata_arr[ptr].cast<int64_t>();
      if (ptr + num_segments * 2 + 1 > static_cast<int>(kv_append_metadata_arr.size())) {
        return TResult::Error("Invalid kv append metadata compression in kv_append_metadata");
      }
      std::vector<int64_t> compressed_kv_append_metadata{num_segments};
      compressed_kv_append_metadata.reserve(num_segments * 2 + 1);
      for (int i = 1; i <= num_segments * 2; ++i) {
        if (!kv_append_metadata_arr[ptr + i].try_cast<int64_t>().has_value()) {
          return TResult::Error("Invalid kv append metadata value in kv_append_metadata array");
        }
        compressed_kv_append_metadata.push_back(kv_append_metadata_arr[ptr + i].cast<int64_t>());
      }
      kv_append_metadata.push_back(Shape(std::move(compressed_kv_append_metadata)));
      ptr += num_segments * 2 + 1;
    }
    res.kv_append_metadata = std::move(kv_append_metadata);
  }
  res.kv_window_begin = json::LookupOptional<int64_t>(config, "kv_window_begin");
  res.kv_window_end = json::LookupOptional<int64_t>(config, "kv_window_end");
  res.dst_group_offset = json::LookupOptional<int64_t>(config, "dst_group_offset");
  return TResult::Ok(res);
}

tvm::ffi::json::Object DisaggConfig::AsJSON() const {
  tvm::ffi::json::Object config;
  switch (kind) {
    case DisaggRequestKind::kPrepareReceive: {
      config.Set("kind", "prepare_receive");
      break;
    }
    case DisaggRequestKind::kRemoteSend: {
      config.Set("kind", "remote_send");
      break;
    }
    case DisaggRequestKind::kStartGeneration: {
      config.Set("kind", "start_generation");
      break;
    }
    default:
      break;
  }
  if (!kv_append_metadata.empty()) {
    tvm::ffi::json::Array kv_append_metadata_arr;
    for (const Shape& compressed_kv_append_metadata : kv_append_metadata) {
      for (int64_t value : compressed_kv_append_metadata) {
        kv_append_metadata_arr.push_back(value);
      }
    }
    config.Set("kv_append_metadata",
               Base64Encode(tvm::ffi::json::Stringify(kv_append_metadata_arr)));
  }
  if (kv_window_begin.has_value()) {
    config.Set("kv_window_begin", static_cast<int64_t>(kv_window_begin.value()));
  }
  if (kv_window_end.has_value()) {
    config.Set("kv_window_end", static_cast<int64_t>(kv_window_end.value()));
  }
  if (dst_group_offset.has_value()) {
    config.Set("dst_group_offset", static_cast<int64_t>(dst_group_offset.value()));
  }
  return config;
}

/****************** DebugConfig ******************/

Result<DebugConfig> DebugConfig::FromJSON(const tvm::ffi::json::Object& config) {
  using TResult = Result<DebugConfig>;
  DebugConfig res;
  res.ignore_eos = json::LookupOrDefault<bool>(config, "ignore_eos", false);
  res.pinned_system_prompt = json::LookupOrDefault<bool>(config, "pinned_system_prompt", false);
  std::string special_request = json::LookupOrDefault<std::string>(config, "special_request", "");
  if (special_request.length() != 0) {
    if (special_request == "query_engine_metrics") {
      res.special_request = SpecialRequestKind::kQueryEngineMetrics;
    } else {
      return TResult::Error("Unknown special request " + special_request);
    }
  }
  std::string grammar_execution_mode =
      json::LookupOrDefault<std::string>(config, "grammar_execution_mode", "jump_forward");
  if (grammar_execution_mode == "jump_forward") {
    res.grammar_execution_mode = GrammarExecutionMode::kJumpForward;
  } else if (grammar_execution_mode == "constraint") {
    res.grammar_execution_mode = GrammarExecutionMode::kConstraint;
  } else {
    return TResult::Error("Unknown grammar execution mode " + grammar_execution_mode);
  }
  if (auto disagg_config_obj =
          json::LookupOptional<tvm::ffi::json::Object>(config, "disagg_config")) {
    Result<DisaggConfig> disagg_config = DisaggConfig::FromJSON(disagg_config_obj.value());
    if (disagg_config.IsErr()) {
      return TResult::Error(disagg_config.UnwrapErr());
    }
    res.disagg_config = disagg_config.Unwrap();
  }
  return TResult::Ok(res);
}

/**
 * \return serialized json value of the config.
 */
tvm::ffi::json::Object DebugConfig::AsJSON() const {
  tvm::ffi::json::Object config;
  config.Set("ignore_eos", ignore_eos);
  config.Set("pinned_system_prompt", pinned_system_prompt);
  switch (special_request) {
    case SpecialRequestKind::kQueryEngineMetrics: {
      config.Set("special_request", "query_engine_metrics");
      break;
    }
    case SpecialRequestKind::kNone:
      break;
  }
  switch (grammar_execution_mode) {
    case GrammarExecutionMode::kJumpForward: {
      config.Set("grammar_execution_mode", "jump_forward");
      break;
    }
    case GrammarExecutionMode::kConstraint: {
      config.Set("grammar_execution_mode", "constraint");
      break;
    }
  }
  if (disagg_config.kind != DisaggRequestKind::kNone) {
    config.Set("disagg_config", disagg_config.AsJSON());
  }
  return config;
}

/****************** GenerationConfig ******************/

Result<GenerationConfig> GenerationConfig::Validate(GenerationConfig cfg) {
  using TResult = Result<GenerationConfig>;
  if (cfg->n <= 0) {
    return TResult::Error("\"n\" should be at least 1");
  }
  if (cfg->temperature < 0) {
    return TResult::Error("\"temperature\" should be non-negative");
  }
  if (cfg->top_p < 0 || cfg->top_p > 1) {
    return TResult::Error("\"top_p\" should be in range [0, 1]");
  }
  if (std::fabs(cfg->frequency_penalty) > 2.0) {
    return TResult::Error("frequency_penalty must be in [-2, 2]!");
  }
  if (cfg->repetition_penalty <= 0) {
    return TResult::Error("\"repetition_penalty\" must be positive");
  }
  if (cfg->top_logprobs < 0 || cfg->top_logprobs > 20) {
    return TResult::Error("At most 20 top logprob tokens are supported");
  }
  if (cfg->top_logprobs != 0 && !(cfg->logprobs)) {
    return TResult::Error("\"logprobs\" must be true to support \"top_logprobs\"");
  }
  for (const auto& item : cfg->logit_bias) {
    double bias_value = item.second;
    if (std::fabs(bias_value) > 100.0) {
      return TResult::Error("Logit bias value should be in range [-100, 100].");
    }
  }
  return TResult::Ok(cfg);
}

Result<GenerationConfig> GenerationConfig::FromJSON(const tvm::ffi::json::Object& config,
                                                    const GenerationConfig& default_config) {
  using TResult = Result<GenerationConfig>;
  ObjectPtr<GenerationConfigNode> n = tvm::ffi::make_object<GenerationConfigNode>();
  n->n = json::LookupOrDefault<int64_t>(config, "n", default_config->n);
  n->temperature =
      json::LookupOrDefault<double>(config, "temperature", default_config->temperature);
  n->top_p = json::LookupOrDefault<double>(config, "top_p", default_config->top_p);
  n->frequency_penalty =
      json::LookupOrDefault<double>(config, "frequency_penalty", default_config->frequency_penalty);
  n->presence_penalty =
      json::LookupOrDefault<double>(config, "presence_penalty", default_config->presence_penalty);
  n->repetition_penalty = json::LookupOrDefault<double>(config, "repetition_penalty",
                                                        default_config->repetition_penalty);
  n->logprobs = json::LookupOrDefault<bool>(config, "logprobs", default_config->logprobs);
  n->top_logprobs =
      json::LookupOrDefault<int64_t>(config, "top_logprobs", default_config->top_logprobs);

  std::optional<tvm::ffi::json::Object> logit_bias_obj =
      json::LookupOptional<tvm::ffi::json::Object>(config, "logit_bias");
  if (logit_bias_obj.has_value()) {
    std::vector<std::pair<int, float>> logit_bias;
    logit_bias.reserve(logit_bias_obj.value().size());
    for (auto [k, v] : logit_bias_obj.value()) {
      std::string token_id_str(k.cast<tvm::ffi::String>());
      TVM_FFI_ICHECK(v.try_cast<double>().has_value());
      double bias_value = v.cast<double>();
      logit_bias.emplace_back(std::stoi(token_id_str), bias_value);
    }
    n->logit_bias = std::move(logit_bias);
  } else {
    n->logit_bias = default_config->logit_bias;
  }

  n->seed = json::LookupOrDefault<int64_t>(config, "seed", std::random_device{}());
  // "-1" means the generation will not stop until exceeding
  // model capability or hit any stop criteria.
  n->max_tokens = json::LookupOrDefault<int64_t>(config, "max_tokens", -1);

  std::optional<tvm::ffi::json::Array> stop_strs_arr =
      json::LookupOptional<tvm::ffi::json::Array>(config, "stop_strs");
  if (stop_strs_arr.has_value()) {
    Array<String> stop_strs;
    stop_strs.reserve(stop_strs_arr.value().size());
    for (const auto& v : stop_strs_arr.value()) {
      if (!v.try_cast<std::string>().has_value()) {
        return TResult::Error("Invalid stop string in stop_strs");
      }
      stop_strs.push_back(v.cast<std::string>());
    }
    n->stop_strs = std::move(stop_strs);
  } else {
    n->stop_strs = default_config->stop_strs;
  }
  std::optional<tvm::ffi::json::Array> stop_token_ids_arr =
      json::LookupOptional<tvm::ffi::json::Array>(config, "stop_token_ids");
  if (stop_token_ids_arr.has_value()) {
    std::vector<int> stop_token_ids;
    stop_token_ids.reserve(stop_token_ids_arr.value().size());
    for (const auto& v : stop_token_ids_arr.value()) {
      if (!v.try_cast<int64_t>().has_value()) {
        return TResult::Error("Invalid stop token in stop_token_ids");
      }
      stop_token_ids.push_back(v.cast<int64_t>());
    }
    n->stop_token_ids = std::move(stop_token_ids);
  } else {
    n->stop_token_ids = default_config->stop_token_ids;
  }

  std::optional<tvm::ffi::json::Object> response_format_obj =
      json::LookupOptional<tvm::ffi::json::Object>(config, "response_format");
  if (response_format_obj.has_value()) {
    Result<ResponseFormat> response_format_res =
        ResponseFormat::FromJSON(response_format_obj.value());
    if (response_format_res.IsErr()) {
      return TResult::Error(response_format_res.UnwrapErr());
    }
    n->response_format = response_format_res.Unwrap();
  } else {
    n->response_format = default_config->response_format;
  }
  // "debug_config" is for internal usage. Not the part of OpenAI API spec.
  std::optional<tvm::ffi::json::Object> debug_config_obj =
      json::LookupOptional<tvm::ffi::json::Object>(config, "debug_config");

  if (debug_config_obj.has_value()) {
    Result<DebugConfig> debug_config_res = DebugConfig::FromJSON(debug_config_obj.value());
    if (debug_config_res.IsErr()) {
      return TResult::Error(debug_config_res.UnwrapErr());
    }
    n->debug_config = debug_config_res.Unwrap();
  }
  return Validate(GenerationConfig(n));
}

GenerationConfig GenerationConfig::GetDefaultFromModelConfig(
    const tvm::ffi::json::Object& model_config_json) {
  ObjectPtr<GenerationConfigNode> n = tvm::ffi::make_object<GenerationConfigNode>();
  n->max_tokens = -1;
  n->temperature = json::LookupOrDefault<double>(model_config_json, "temperature", n->temperature);
  n->top_p = json::LookupOrDefault<double>(model_config_json, "top_p", n->top_p);
  n->frequency_penalty =
      json::LookupOrDefault<double>(model_config_json, "frequency_penalty", n->frequency_penalty);
  n->presence_penalty =
      json::LookupOrDefault<double>(model_config_json, "presence_penalty", n->presence_penalty);
  return GenerationConfig(n);
}

tvm::ffi::json::Object GenerationConfigNode::AsJSON() const {
  tvm::ffi::json::Object config;
  config.Set("n", static_cast<int64_t>(this->n));
  config.Set("temperature", this->temperature);
  config.Set("top_p", this->top_p);
  config.Set("frequency_penalty", this->frequency_penalty);
  config.Set("presence_penalty", this->presence_penalty);
  config.Set("repetition_penalty", this->repetition_penalty);
  config.Set("logprobs", this->logprobs);
  config.Set("top_logprobs", static_cast<int64_t>(this->top_logprobs));
  config.Set("max_tokens", static_cast<int64_t>(this->max_tokens));
  config.Set("seed", static_cast<int64_t>(this->seed));

  tvm::ffi::json::Object logit_bias_obj;
  for (auto [token_id, bias] : logit_bias) {
    logit_bias_obj.Set(std::to_string(token_id), static_cast<double>(bias));
  }
  config.Set("logit_bias", logit_bias_obj);

  tvm::ffi::json::Array stop_strs_arr;
  for (String stop_str : this->stop_strs) {
    stop_strs_arr.push_back(stop_str);
  }
  config.Set("stop_strs", stop_strs_arr);

  tvm::ffi::json::Array stop_token_ids_arr;
  for (int stop_token_id : this->stop_token_ids) {
    stop_token_ids_arr.push_back(static_cast<int64_t>(stop_token_id));
  }
  config.Set("stop_token_ids", stop_token_ids_arr);

  tvm::ffi::json::Object response_format;
  response_format.Set("type", this->response_format.type);
  if (this->response_format.schema) {
    response_format.Set("schema", this->response_format.schema.value());
  } else {
    response_format.Set("schema", tvm::Any(nullptr));
  }
  config.Set("response_format", response_format);
  config.Set("debug_config", debug_config.AsJSON());
  return config;
}

/****************** EngineConfig ******************/

EngineConfig EngineConfig::FromJSONAndInferredConfig(
    const tvm::ffi::json::Object& json, const InferrableEngineConfig& inferred_config) {
  TVM_FFI_ICHECK(inferred_config.max_num_sequence.has_value());
  TVM_FFI_ICHECK(inferred_config.max_total_sequence_length.has_value());
  TVM_FFI_ICHECK(inferred_config.prefill_chunk_size.has_value());
  TVM_FFI_ICHECK(inferred_config.max_history_size.has_value());
  ObjectPtr<EngineConfigNode> n = tvm::ffi::make_object<EngineConfigNode>();

  // - Get models and model libs.
  n->model = json::Lookup<std::string>(json, "model");
  n->model_lib = json::Lookup<std::string>(json, "model_lib");
  std::vector<String> additional_models;
  std::vector<String> additional_model_libs;
  tvm::ffi::json::Array additional_models_arr = json::LookupOrDefault<tvm::ffi::json::Array>(
      json, "additional_models", tvm::ffi::json::Array());
  int num_additional_models = additional_models_arr.size();
  additional_models.reserve(num_additional_models);
  additional_model_libs.reserve(num_additional_models);
  for (int i = 0; i < num_additional_models; ++i) {
    tvm::ffi::json::Array additional_model_pair =
        json::Lookup<tvm::ffi::json::Array>(additional_models_arr, i);
    additional_models.push_back(json::Lookup<std::string>(additional_model_pair, 0));
    additional_model_libs.push_back(json::Lookup<std::string>(additional_model_pair, 1));
  }
  n->additional_models = additional_models;
  n->additional_model_libs = additional_model_libs;
  n->mode = EngineModeFromString(json::Lookup<std::string>(json, "mode"));

  // - Other fields with default value.
  n->gpu_memory_utilization = static_cast<float>(
      json::LookupOrDefault<double>(json, "gpu_memory_utilization", n->gpu_memory_utilization));
  n->kv_cache_page_size = static_cast<int>(
      json::LookupOrDefault<int64_t>(json, "kv_cache_page_size", n->kv_cache_page_size));
  n->speculative_mode = SpeculativeModeFromString(json::LookupOrDefault<std::string>(
      json, "speculative_mode", SpeculativeModeToString(n->speculative_mode)));
  n->spec_draft_length = static_cast<int>(
      json::LookupOrDefault<int64_t>(json, "spec_draft_length", n->spec_draft_length));
  n->spec_tree_width =
      static_cast<int>(json::LookupOrDefault<int64_t>(json, "spec_tree_width", n->spec_tree_width));
  n->prefill_mode = PrefillModeFromString(json::LookupOrDefault<std::string>(
      json, "prefill_mode", PrefillModeToString(n->prefill_mode)));
  n->verbose = json::LookupOrDefault<bool>(json, "verbose", n->verbose);

  // - Fields from the inferred engine config.
  n->max_num_sequence = inferred_config.max_num_sequence.value();
  n->max_total_sequence_length = inferred_config.max_total_sequence_length.value();
  if (inferred_config.max_single_sequence_length.has_value()) {
    n->max_single_sequence_length = inferred_config.max_single_sequence_length.value();
  }
  n->prefill_chunk_size = inferred_config.prefill_chunk_size.value();
  n->max_history_size = inferred_config.max_history_size.value();

  n->prefix_cache_mode = PrefixCacheModeFromString(json::LookupOrDefault<std::string>(
      json, "prefix_cache_mode", PrefixCacheModeToString(n->prefix_cache_mode)));
  n->prefix_cache_max_num_recycling_seqs = static_cast<int>(json::LookupOrDefault<int64_t>(
      json, "prefix_cache_max_num_recycling_seqs", n->max_num_sequence));
  return EngineConfig(n);
}

Result<std::vector<std::pair<std::string, std::string>>>
EngineConfig::GetModelsAndModelLibsFromJSONString(const std::string& json_str) {
  using TResult = Result<std::vector<std::pair<std::string, std::string>>>;
  tvm::ffi::String err;
  auto config_json = tvm::ffi::json::Parse(json_str, &err);
  if (!err.empty()) {
    return TResult::Error(err);
  }

  // Get the models and model libs from JSON.
  tvm::ffi::json::Object config = config_json.cast<tvm::ffi::json::Object>();
  String model = json::Lookup<std::string>(config, "model");
  String model_lib = json::Lookup<std::string>(config, "model_lib");
  tvm::ffi::json::Array additional_models_arr = json::LookupOrDefault<tvm::ffi::json::Array>(
      config, "additional_models", tvm::ffi::json::Array());

  int num_additional_models = additional_models_arr.size();
  std::vector<std::pair<std::string, std::string>> models_and_model_libs;
  models_and_model_libs.reserve(num_additional_models + 1);
  models_and_model_libs.emplace_back(model, model_lib);
  for (int i = 0; i < num_additional_models; ++i) {
    tvm::ffi::json::Array additional_model_pair =
        json::Lookup<tvm::ffi::json::Array>(additional_models_arr, i);
    models_and_model_libs.emplace_back(json::Lookup<std::string>(additional_model_pair, 0),
                                       json::Lookup<std::string>(additional_model_pair, 1));
  }
  return TResult::Ok(models_and_model_libs);
}

String EngineConfigNode::AsJSONString() const {
  tvm::ffi::json::Object config;

  // - Models and model libs
  config.Set("model", this->model);
  config.Set("model_lib", this->model_lib);
  tvm::ffi::json::Array additional_models_arr;
  additional_models_arr.reserve(this->additional_models.size());
  for (int i = 0; i < static_cast<int>(this->additional_models.size()); ++i) {
    tvm::ffi::json::Array pair;
    pair.push_back(this->additional_models[i]);
    pair.push_back(this->additional_model_libs[i]);
    additional_models_arr.push_back(pair);
  }
  config.Set("additional_models", additional_models_arr);

  // - Other fields
  config.Set("mode", EngineModeToString(this->mode));
  config.Set("gpu_memory_utilization", static_cast<double>(this->gpu_memory_utilization));
  config.Set("kv_cache_page_size", static_cast<int64_t>(this->kv_cache_page_size));
  config.Set("max_num_sequence", static_cast<int64_t>(this->max_num_sequence));
  config.Set("max_total_sequence_length", static_cast<int64_t>(this->max_total_sequence_length));
  config.Set("max_single_sequence_length", static_cast<int64_t>(this->max_single_sequence_length));
  config.Set("prefill_chunk_size", static_cast<int64_t>(this->prefill_chunk_size));
  config.Set("max_history_size", static_cast<int64_t>(this->max_history_size));
  config.Set("prefix_cache_mode", PrefixCacheModeToString(this->prefix_cache_mode));
  config.Set("prefix_cache_max_num_recycling_seqs",
             static_cast<int64_t>(this->prefix_cache_max_num_recycling_seqs));
  config.Set("speculative_mode", SpeculativeModeToString(this->speculative_mode));
  config.Set("spec_draft_length", static_cast<int64_t>(this->spec_draft_length));
  config.Set("prefill_mode", PrefillModeToString(this->prefill_mode));
  config.Set("verbose", static_cast<bool>(this->verbose));

  return tvm::ffi::json::Stringify(config, 2);
}

/****************** InferrableEngineConfig ******************/

/*! \brief The class for config limitation from models. */
struct ModelConfigLimits {
  int64_t model_compile_time_max_single_sequence_length;
  int64_t model_runtime_max_single_sequence_length;
  int64_t model_compile_time_max_prefill_chunk_size;
  int64_t model_runtime_max_prefill_chunk_size;
  int64_t model_max_sliding_window_size;
  int64_t model_max_batch_size;
};

/*! \brief Convert the bytes to megabytes, keeping 3 decimals. */
inline std::string BytesToMegabytesString(double bytes) {
  std::ostringstream os;
  os << std::setprecision(3) << std::fixed << (bytes / 1024 / 1024);
  return os.str();
}

/*!
 * \brief Get the upper bound of single sequence length, prefill size and batch size
 * from model config.
 */
Result<ModelConfigLimits> GetModelConfigLimits(
    const std::vector<tvm::ffi::json::Object>& model_configs,
    const std::vector<ModelMetadata>& model_metadata) {
  TVM_FFI_ICHECK_EQ(model_configs.size(), model_metadata.size());
  int64_t model_compile_time_max_single_sequence_length = std::numeric_limits<int64_t>::max();
  int64_t model_runtime_max_single_sequence_length = std::numeric_limits<int64_t>::max();
  int64_t model_compile_time_max_prefill_chunk_size = std::numeric_limits<int64_t>::max();
  int64_t model_runtime_max_prefill_chunk_size = std::numeric_limits<int64_t>::max();
  int64_t model_max_batch_size = std::numeric_limits<int64_t>::max();
  int64_t model_max_sliding_window_size = std::numeric_limits<int64_t>::max();
  for (int i = 0; i < static_cast<int>(model_configs.size()); ++i) {
    // - The maximum single sequence length is the minimum context window size among all models.
    int64_t runtime_context_window_size =
        json::LookupOptional<int64_t>(model_configs[i], "context_window_size").value_or(-1);
    int64_t compile_time_context_window_size = model_metadata[i].context_window_size;

    // limit runtime setting by compile time setting
    if (compile_time_context_window_size != -1) {
      if (runtime_context_window_size == -1 ||
          runtime_context_window_size > compile_time_context_window_size) {
        runtime_context_window_size = compile_time_context_window_size;
      }
    }

    if (compile_time_context_window_size != -1) {
      model_compile_time_max_single_sequence_length =
          std::min(model_compile_time_max_single_sequence_length, compile_time_context_window_size);
    }
    if (runtime_context_window_size != -1) {
      model_runtime_max_single_sequence_length =
          std::min(model_runtime_max_single_sequence_length, runtime_context_window_size);
    }
    // - The maximum prefill chunk size is the minimum prefill chunk size among all models.
    int64_t runtime_prefill_chunk_size =
        json::Lookup<int64_t>(model_configs[i], "prefill_chunk_size");
    int64_t compile_time_prefill_chunk_size = model_metadata[i].prefill_chunk_size;

    // limit runtime setting by compile time setting
    if (compile_time_prefill_chunk_size != -1) {
      if (runtime_prefill_chunk_size == -1 ||
          runtime_prefill_chunk_size > compile_time_prefill_chunk_size) {
        runtime_prefill_chunk_size = compile_time_prefill_chunk_size;
      }
    }

    if (compile_time_prefill_chunk_size != -1) {
      model_compile_time_max_prefill_chunk_size =
          std::min(model_compile_time_max_prefill_chunk_size, compile_time_prefill_chunk_size);
    }
    if (runtime_prefill_chunk_size != -1) {
      model_runtime_max_prefill_chunk_size =
          std::min(model_runtime_max_prefill_chunk_size, runtime_prefill_chunk_size);
    }
    // - The maximum batch size is the minimum max batch size among all models.
    model_max_batch_size = std::min(model_max_batch_size, model_metadata[i].max_batch_size);
    // - The maximum sliding window size is the minimum among all models.
    int64_t runtime_sliding_window_size =
        json::LookupOptional<int64_t>(model_configs[i], "sliding_window_size").value_or(-1);
    if (runtime_sliding_window_size != -1) {
      model_max_sliding_window_size =
          std::min(model_max_sliding_window_size, runtime_sliding_window_size);
    }
  }
  TVM_FFI_ICHECK_NE(model_compile_time_max_prefill_chunk_size, std::numeric_limits<int64_t>::max());
  TVM_FFI_ICHECK_NE(model_runtime_max_prefill_chunk_size, std::numeric_limits<int64_t>::max());
  TVM_FFI_ICHECK_NE(model_max_batch_size, std::numeric_limits<int64_t>::max());
  TVM_FFI_ICHECK_GT(model_compile_time_max_prefill_chunk_size, 0);
  TVM_FFI_ICHECK_GT(model_runtime_max_prefill_chunk_size, 0);
  TVM_FFI_ICHECK_GT(model_max_batch_size, 0);
  return Result<ModelConfigLimits>::Ok(
      {model_compile_time_max_single_sequence_length, model_runtime_max_single_sequence_length,
       model_compile_time_max_prefill_chunk_size, model_runtime_max_prefill_chunk_size,
       model_max_sliding_window_size, model_max_batch_size});
}

/*! \brief The class for memory usage estimation result. */
struct MemUsageEstimationResult {
  double total_memory_bytes;
  double kv_cache_memory_bytes;
  double temp_memory_bytes;
  InferrableEngineConfig inferred_config;
};

Result<MemUsageEstimationResult> EstimateMemoryUsageOnMode(
    EngineMode mode, Device device, double gpu_memory_utilization, int64_t params_bytes,
    int64_t temp_buffer_bytes,
    const std::vector<tvm::ffi::json::Object>& model_configs,  //
    const std::vector<ModelMetadata>& model_metadata,          //
    ModelConfigLimits model_config_limits,                     //
    InferrableEngineConfig init_config, bool verbose) {
  std::ostringstream os;
  InferrableEngineConfig inferred_config = init_config;
  // - 1. max_num_sequence
  if (!init_config.max_num_sequence.has_value()) {
    if (mode == EngineMode::kLocal) {
      inferred_config.max_num_sequence =
          std::min(static_cast<int64_t>(4), model_config_limits.model_max_batch_size);
    } else if (mode == EngineMode::kInteractive) {
      inferred_config.max_num_sequence = 1;
    } else {
      inferred_config.max_num_sequence = model_config_limits.model_max_batch_size;
    }
    os << "max batch size will be set to " << inferred_config.max_num_sequence.value() << ", ";
  } else {
    os << "max batch size " << inferred_config.max_num_sequence.value()
       << " is specified by user, ";
  }
  int64_t max_num_sequence = inferred_config.max_num_sequence.value();
  // - 2. max_single_sequence_length
  if (!init_config.max_single_sequence_length.has_value()) {
    inferred_config.max_single_sequence_length =
        model_config_limits.model_runtime_max_single_sequence_length;
  } else {
    inferred_config.max_single_sequence_length =
        std::min(inferred_config.max_single_sequence_length.value(),
                 model_config_limits.model_compile_time_max_single_sequence_length);
  }
  // - 3. infer the maximum total sequence length that can fit GPU memory.
  double kv_bytes_per_token = 0;
  double kv_aux_workspace_bytes = 0;
  double model_workspace_bytes = 0;
  double logit_processor_workspace_bytes = 0;
  TVM_FFI_ICHECK_EQ(model_configs.size(), model_metadata.size());
  int num_models = model_configs.size();
  for (int i = 0; i < num_models; ++i) {
    // - Read the vocab size and compile-time prefill chunk size (which affects memory allocation).
    tvm::ffi::json::Object compile_time_model_config =
        json::Lookup<tvm::ffi::json::Object>(model_configs[i], "model_config");
    int64_t vocab_size = json::Lookup<int64_t>(compile_time_model_config, "vocab_size");
    int64_t prefill_chunk_size =
        json::Lookup<int64_t>(compile_time_model_config, "prefill_chunk_size");
    // - Calculate KV cache memory usage.
    int64_t num_layers = model_metadata[i].kv_cache_metadata.num_hidden_layers;
    int64_t head_dim = model_metadata[i].kv_cache_metadata.head_dim;
    int64_t num_qo_heads = model_metadata[i].kv_cache_metadata.num_attention_heads;
    int64_t num_kv_heads = model_metadata[i].kv_cache_metadata.num_key_value_heads;
    int64_t hidden_size = head_dim * num_qo_heads;
    kv_bytes_per_token +=
        head_dim * num_kv_heads * (num_layers / model_metadata[i].pipeline_parallel_stages) * 4 +
        1.25;
    kv_aux_workspace_bytes +=
        (max_num_sequence + 1) * 88 + prefill_chunk_size * (num_qo_heads + 1) * 8 +
        prefill_chunk_size * head_dim * (num_qo_heads + num_kv_heads) * 4 + 48 * 1024 * 1024;
    model_workspace_bytes += prefill_chunk_size * 4 + max_num_sequence * 4 +
                             (prefill_chunk_size * 2 + max_num_sequence) * hidden_size * 2;
    logit_processor_workspace_bytes +=
        max_num_sequence * 20 + max_num_sequence * vocab_size * 16.125;
  }
  int64_t gpu_size_bytes = TotalDetectGlobalMemory(device);
  // Compute the maximum total sequence length under the GPU memory budget.
  int64_t model_max_total_sequence_length =
      static_cast<int>((gpu_size_bytes * gpu_memory_utilization  //
                        - params_bytes                           //
                        - temp_buffer_bytes                      //
                        - kv_aux_workspace_bytes                 //
                        - model_workspace_bytes                  //
                        - logit_processor_workspace_bytes) /
                       kv_bytes_per_token);
  if (model_max_total_sequence_length <= 0) {
    if (verbose) {
      LOG(INFO) << "temp_buffer = " << BytesToMegabytesString(temp_buffer_bytes);
      LOG(INFO) << "kv_aux workspace = " << BytesToMegabytesString(kv_aux_workspace_bytes);
      LOG(INFO) << "model workspace = " << BytesToMegabytesString(model_workspace_bytes);
      LOG(INFO) << "logit processor workspace = "
                << BytesToMegabytesString(logit_processor_workspace_bytes);
    }
    return Result<MemUsageEstimationResult>::Error(
        "Insufficient GPU memory error: "
        "The available single GPU memory is " +
        BytesToMegabytesString(gpu_size_bytes * gpu_memory_utilization) +
        " MB, "
        "which is less than the sum of model weight size (" +
        BytesToMegabytesString(params_bytes) + " MB) and temporary buffer size (" +
        BytesToMegabytesString(temp_buffer_bytes + kv_aux_workspace_bytes + model_workspace_bytes +
                               logit_processor_workspace_bytes) +
        " MB).\n"
        "1. You can set a larger \"gpu_memory_utilization\" value.\n"
        "2. If the model weight size is too large, please enable tensor parallelism by passing "
        "`--tensor-parallel-shards $NGPU` to `mlc_llm gen_config` or use quantization.\n"
        "3. If the temporary buffer size is too large, please use a smaller `--prefill-chunk-size` "
        "in `mlc_llm gen_config`.");
  }
  if (device.device_type == DLDeviceType::kDLMetal) {
    // NOTE: Metal runtime has severe performance issues with large buffers.
    // To work around the issue, we limit the KV cache capacity to 32768.
    model_max_total_sequence_length =
        std::min(model_max_total_sequence_length, static_cast<int64_t>(32768));
  }
  // Compute the total memory usage except the KV cache part.
  double total_mem_usage_except_kv_cache =
      (params_bytes + temp_buffer_bytes + kv_aux_workspace_bytes + model_workspace_bytes +
       logit_processor_workspace_bytes);

  // - 4. max_total_sequence_length
  if (!init_config.max_total_sequence_length.has_value()) {
    if (mode == EngineMode::kLocal) {
      inferred_config.max_total_sequence_length = std::min(
          {model_max_total_sequence_length, inferred_config.max_single_sequence_length.value(),
           static_cast<int64_t>(8192)});
    } else if (mode == EngineMode::kInteractive) {
      inferred_config.max_total_sequence_length = std::min(
          {model_max_total_sequence_length, inferred_config.max_single_sequence_length.value()});
    } else {
      inferred_config.max_total_sequence_length =
          inferred_config.max_single_sequence_length.value() == std::numeric_limits<int64_t>::max()
              ? model_max_total_sequence_length
              : std::min(model_max_total_sequence_length,
                         max_num_sequence * inferred_config.max_single_sequence_length.value());
    }
    os << "max KV cache token capacity will be set to "
       << inferred_config.max_total_sequence_length.value() << ", ";
  } else {
    os << "max KV cache token capacity " << inferred_config.max_total_sequence_length.value()
       << " is specified by user, ";
  }
  // - 5. prefill_chunk_size
  if (!init_config.prefill_chunk_size.has_value()) {
    if (mode == EngineMode::kLocal || mode == EngineMode::kInteractive) {
      inferred_config.prefill_chunk_size =
          std::min({model_config_limits.model_runtime_max_prefill_chunk_size,
                    inferred_config.max_total_sequence_length.value(),
                    inferred_config.max_single_sequence_length.value()});
    } else {
      inferred_config.prefill_chunk_size = model_config_limits.model_runtime_max_prefill_chunk_size;
    }
    os << "prefill chunk size will be set to " << inferred_config.prefill_chunk_size.value()
       << ". ";
  } else {
    os << "prefill chunk size " << inferred_config.prefill_chunk_size.value()
       << " is specified by user. ";
  }

  // - Print logging message
  if (verbose) {
    LOG(INFO) << "Under mode \"" << EngineModeToString(mode) << "\", " << os.str();
  }

  return Result<MemUsageEstimationResult>::Ok(
      {total_mem_usage_except_kv_cache +
           inferred_config.max_total_sequence_length.value() * kv_bytes_per_token,
       kv_bytes_per_token * inferred_config.max_total_sequence_length.value() +
           kv_aux_workspace_bytes,
       model_workspace_bytes + logit_processor_workspace_bytes + temp_buffer_bytes,
       inferred_config});
}

Result<InferrableEngineConfig> InferrableEngineConfig::InferForKVCache(
    EngineMode mode, Device device, double gpu_memory_utilization,
    const std::vector<tvm::ffi::json::Object>& model_configs,
    const std::vector<ModelMetadata>& model_metadata, InferrableEngineConfig init_config,
    bool verbose) {
  // - Check if max_history_size is not set.
  if (init_config.max_history_size.has_value() && init_config.max_history_size.value() != 0) {
    return Result<InferrableEngineConfig>::Error(
        "KV cache does not support max_history_size, while it is set to " +
        std::to_string(init_config.max_history_size.value()) + " in the input EngineConfig");
  }
  // - Get the upper bound of single sequence length, prefill size and batch size
  // from model config.
  Result<ModelConfigLimits> model_config_limits_res =
      GetModelConfigLimits(model_configs, model_metadata);

  if (model_config_limits_res.IsErr()) {
    return Result<InferrableEngineConfig>::Error(model_config_limits_res.UnwrapErr());
  }
  ModelConfigLimits model_config_limits = model_config_limits_res.Unwrap();
  // - Get total model parameter size and temporary in-function buffer
  // size in bytes on single GPU.
  int64_t params_bytes = 0;
  int64_t temp_buffer_bytes = 0;
  for (const ModelMetadata& metadata : model_metadata) {
    for (const ModelMetadata::Param& param : metadata.params) {
      int64_t param_size = (param.dtype.bits * param.dtype.lanes + 7) / 8;
      for (int64_t v : param.shape) {
        TVM_FFI_ICHECK_GE(v, 0);
        param_size *= v;
      }
      params_bytes += param_size;
    }
    params_bytes /= metadata.pipeline_parallel_stages;
    for (const auto& [func_name, temp_buffer_size] : metadata.memory_usage) {
      temp_buffer_bytes = std::max(temp_buffer_bytes, temp_buffer_size);
    }
  }
  // Magnify the temp buffer by a factor of 2 for safety.
  temp_buffer_bytes *= 2;

  // - Infer the engine config and estimate memory usage for each mode.
  Result<MemUsageEstimationResult> local_mode_estimation_result = EstimateMemoryUsageOnMode(
      EngineMode::kLocal, device, gpu_memory_utilization, params_bytes, temp_buffer_bytes,
      model_configs, model_metadata, model_config_limits, init_config, verbose);
  Result<MemUsageEstimationResult> interactive_mode_estimation_result = EstimateMemoryUsageOnMode(
      EngineMode::kInteractive, device, gpu_memory_utilization, params_bytes, temp_buffer_bytes,
      model_configs, model_metadata, model_config_limits, init_config, verbose);
  Result<MemUsageEstimationResult> server_mode_estimation_result = EstimateMemoryUsageOnMode(
      EngineMode::kServer, device, gpu_memory_utilization, params_bytes, temp_buffer_bytes,
      model_configs, model_metadata, model_config_limits, init_config, verbose);
  // - Pick the estimation result according to the mode.
  std::string mode_name;
  Result<MemUsageEstimationResult> final_estimation_result;
  if (mode == EngineMode::kLocal) {
    final_estimation_result = std::move(local_mode_estimation_result);
  } else if (mode == EngineMode::kInteractive) {
    final_estimation_result = std::move(interactive_mode_estimation_result);
  } else {
    final_estimation_result = std::move(server_mode_estimation_result);
  }
  if (final_estimation_result.IsErr()) {
    return Result<InferrableEngineConfig>::Error(final_estimation_result.UnwrapErr());
  }
  // - Print log message.
  MemUsageEstimationResult final_estimation = final_estimation_result.Unwrap();
  InferrableEngineConfig inferred_config = std::move(final_estimation.inferred_config);

  if (verbose) {
    LOG(INFO) << "The actual engine mode is \"" << EngineModeToString(mode)
              << "\". So max batch size is " << inferred_config.max_num_sequence.value()
              << ", max KV cache token capacity is "
              << inferred_config.max_total_sequence_length.value() << ", prefill chunk size is "
              << inferred_config.prefill_chunk_size.value() << ".";
    LOG(INFO) << "Estimated total single GPU memory usage: "
              << BytesToMegabytesString(final_estimation.total_memory_bytes)
              << " MB (Parameters: " << BytesToMegabytesString(params_bytes)
              << " MB. KVCache: " << BytesToMegabytesString(final_estimation.kv_cache_memory_bytes)
              << " MB. Temporary buffer: "
              << BytesToMegabytesString(final_estimation.temp_memory_bytes)
              << " MB). The actual usage might be slightly larger than the estimated number.";
  }

  inferred_config.max_history_size = 0;
  return Result<InferrableEngineConfig>::Ok(inferred_config);
}

Result<InferrableEngineConfig> InferrableEngineConfig::InferForRNNState(
    EngineMode mode, Device device, double gpu_memory_utilization,
    const std::vector<tvm::ffi::json::Object>& model_configs,
    const std::vector<ModelMetadata>& model_metadata, InferrableEngineConfig init_config,
    bool verbose) {
  // - Check max_single_sequence_length is not set.
  if (init_config.max_single_sequence_length.has_value()) {
    return Result<InferrableEngineConfig>::Error(
        "RNN state does not support max_single_sequence_length, while it is set to " +
        std::to_string(init_config.max_single_sequence_length.value()) +
        " in the input EngineConfig");
  }
  // - Get the upper bound of single sequence length, prefill size and batch size
  // from model config.
  Result<ModelConfigLimits> model_config_limits_res =
      GetModelConfigLimits(model_configs, model_metadata);
  if (model_config_limits_res.IsErr()) {
    return Result<InferrableEngineConfig>::Error(model_config_limits_res.UnwrapErr());
  }
  ModelConfigLimits model_config_limits = model_config_limits_res.Unwrap();

  std::ostringstream os;
  InferrableEngineConfig inferred_config = init_config;
  // - 1. prefill_chunk_size
  if (!init_config.prefill_chunk_size.has_value()) {
    inferred_config.prefill_chunk_size = std::min(
        model_config_limits.model_runtime_max_prefill_chunk_size, static_cast<int64_t>(4096));
    os << "prefill chunk size will be set to " << inferred_config.prefill_chunk_size.value()
       << ", ";
  } else {
    os << "prefill chunk size " << inferred_config.prefill_chunk_size.value()
       << " is specified by user, ";
  }
  // - 2. max_batch_size
  if (!init_config.max_num_sequence.has_value()) {
    inferred_config.max_num_sequence =
        mode == EngineMode::kInteractive
            ? 1
            : std::min(static_cast<int64_t>(4), model_config_limits.model_max_batch_size);
    os << "max batch size will be set to " << inferred_config.max_num_sequence.value() << ", ";
  } else {
    os << "max batch size " << inferred_config.max_num_sequence.value()
       << " is specified by user, ";
  }
  int64_t max_num_sequence = inferred_config.max_num_sequence.value();
  // - 3. max_total_sequence_length
  if (!init_config.max_total_sequence_length.has_value()) {
    inferred_config.max_total_sequence_length = 32768;
    os << "max RNN state token capacity will be set to "
       << inferred_config.max_total_sequence_length.value() << ". ";
  } else {
    os << "max RNN state token capacity " << inferred_config.max_total_sequence_length.value()
       << " is specified by user. ";
  }

  // - Extra logging message
  if (mode == EngineMode::kLocal) {
    os << "We choose small max batch size and RNN state capacity to use less GPU memory.";
  } else if (mode == EngineMode::kInteractive) {
    os << "We fix max batch size to 1 for interactive single sequence use.";
  } else {
    os << "We use as much GPU memory as possible (within the limit of gpu_memory_utilization).";
  }
  if (verbose) {
    LOG(INFO) << "Under mode \"" << EngineModeToString(mode) << "\", " << os.str();
  }

  // - Get total model parameter size and temporary in-function buffer
  // size in bytes on single GPU.
  int64_t params_bytes = 0;
  int64_t temp_buffer_bytes = 0;
  for (const ModelMetadata& metadata : model_metadata) {
    for (const ModelMetadata::Param& param : metadata.params) {
      int64_t param_size = (param.dtype.bits * param.dtype.lanes + 7) / 8;
      for (int64_t v : param.shape) {
        TVM_FFI_ICHECK_GE(v, 0);
        param_size *= v;
      }
      params_bytes += param_size;
    }
    for (const auto& [func_name, temp_buffer_size] : metadata.memory_usage) {
      temp_buffer_bytes += temp_buffer_size;
    }
  }
  // - 4. max_history_size
  double rnn_state_base_bytes = 0;  // The memory usage for rnn state when history = 1.
  double model_workspace_bytes = 0;
  double logit_processor_workspace_bytes = 0;
  TVM_FFI_ICHECK_EQ(model_configs.size(), model_metadata.size());
  int num_models = model_configs.size();
  for (int i = 0; i < num_models; ++i) {
    // - Read the vocab size and compile-time prefill chunk size (which affects memory allocation).
    tvm::ffi::json::Object compile_time_model_config =
        json::Lookup<tvm::ffi::json::Object>(model_configs[i], "model_config");
    int64_t vocab_size = json::Lookup<int64_t>(compile_time_model_config, "vocab_size");
    int64_t prefill_chunk_size =
        json::Lookup<int64_t>(compile_time_model_config, "prefill_chunk_size");
    int64_t head_size = json::Lookup<int64_t>(compile_time_model_config, "head_size");
    int64_t num_heads = json::Lookup<int64_t>(compile_time_model_config, "num_heads");
    int64_t num_layers = json::Lookup<int64_t>(compile_time_model_config, "num_hidden_layers");
    int64_t hidden_size = json::Lookup<int64_t>(compile_time_model_config, "hidden_size");
    // - Calculate RNN state memory usage.
    rnn_state_base_bytes += (max_num_sequence * hidden_size * num_layers * 2 * 2 +
                             max_num_sequence * num_heads * head_size * head_size * num_layers * 2);
    model_workspace_bytes += prefill_chunk_size * 4 + max_num_sequence * 4 +
                             (prefill_chunk_size * 2 + max_num_sequence) * hidden_size * 2;
    logit_processor_workspace_bytes +=
        max_num_sequence * 20 + max_num_sequence * vocab_size * 16.125;
  }
  int64_t gpu_size_bytes = TotalDetectGlobalMemory(device);
  // Compute the maximum history size length under the GPU memory budget.
  int64_t model_max_history_size = static_cast<int>((gpu_size_bytes * gpu_memory_utilization  //
                                                     - params_bytes                           //
                                                     - temp_buffer_bytes                      //
                                                     - model_workspace_bytes                  //
                                                     - logit_processor_workspace_bytes) /
                                                    rnn_state_base_bytes);
  if (model_max_history_size <= 0) {
    return Result<InferrableEngineConfig>::Error(
        "Insufficient GPU memory error: "
        "The available single GPU memory is " +
        BytesToMegabytesString(gpu_size_bytes * gpu_memory_utilization) +
        " MB, "
        "which is less than the sum of model weight size (" +
        BytesToMegabytesString(params_bytes) + " MB) and temporary buffer size (" +
        BytesToMegabytesString(
            (temp_buffer_bytes + model_workspace_bytes + logit_processor_workspace_bytes)) +
        " MB). "
        "If the model weight size is too large, please use quantization. "
        "If the temporary buffer size is too large, please use a smaller `--prefill-chunk-size` in "
        "`mlc_llm gen_config`.");
  }
  if (!init_config.max_history_size.has_value()) {
    inferred_config.max_history_size = model_max_history_size;
  } else {
    inferred_config.max_history_size =
        std::min(inferred_config.max_history_size.value(), model_max_history_size);
  }
  if (verbose) {
    LOG(INFO) << "The actual engine mode is \"" << EngineModeToString(mode)
              << "\". So max batch size is " << inferred_config.max_num_sequence.value()
              << ", max RNN state token capacity is "
              << inferred_config.max_total_sequence_length.value() << ", prefill chunk size is "
              << inferred_config.prefill_chunk_size.value() << ".";
    LOG(INFO) << "Estimated total single GPU memory usage: "
              << BytesToMegabytesString(params_bytes + temp_buffer_bytes +
                                        inferred_config.max_history_size.value() *
                                            rnn_state_base_bytes)
              << " MB (Parameters: " << BytesToMegabytesString(params_bytes) << " MB. RNN state: "
              << BytesToMegabytesString(inferred_config.max_history_size.value() *
                                        rnn_state_base_bytes)
              << " MB. Temporary buffer: "
              << BytesToMegabytesString(model_workspace_bytes + logit_processor_workspace_bytes +
                                        temp_buffer_bytes)
              << " MB). The actual usage might be slightly larger than the estimated number.";
  }

  return Result<InferrableEngineConfig>::Ok(inferred_config);
}

/****************** Config utils ******************/

Result<bool> ModelsUseKVCache(const std::vector<tvm::ffi::json::Object>& model_configs) {
  TVM_FFI_ICHECK_GE(model_configs.size(), 1);
  std::string model_type = json::Lookup<std::string>(model_configs[0], "model_type");
  bool use_kv_cache = model_type.find("rwkv") == std::string::npos;
  for (int i = 1; i < static_cast<int>(model_configs.size()); ++i) {
    if ((json::Lookup<std::string>(model_configs[i], "model_type").find("rwkv") ==
         std::string::npos) != use_kv_cache) {
      return Result<bool>::Error(
          "Invalid models in EngineConfig. Models must be all RNN model or none model is RNN "
          "model.");
    }
  }
  return Result<bool>::Ok(use_kv_cache);
}

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