Repository: mlc-ai/mlc-llm
File: cpp/serve/config.cc (C++)

Code before:
    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 * 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;
  }
  // Get single-card GPU size.
  TVMRetValue rv;
  DeviceAPI::Get(device)->GetAttr(device, DeviceAttrKind::kTotalGlobalMemory, &rv);
  int64_t gpu_size_bytes = rv;
  // 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);
  }
  // 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 (" +

Review comment (suggestion):
Lift this and the above into a common function TotalDetectGlobalMemory(DLDevice), so we only need to have the OpenCL change in a single place

Code after:
    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 * 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.