[Bug] KVCache takes up too much memory when running mlc_llm.serve.server

[gxmlfx]

🐛 Bug

KVCache takes up too much memory when running mlc_llm.serve.server, but memory usage is normal when using cli or mlc_llm.gradio

To Reproduce

Steps to reproduce the behavior:

• Download model from https://huggingface.co/mlc-ai
• Compile the binary (I tried to reduce context_window_size and prefill_chunk_size, but it did not work to reduce the memory taken)
• Start Rest API server by

  python -m mlc_llm.serve.server --model /home/orangepi/mlc-llm/dist/Llama-2-7b-chat-hf-q4f16_1-MLC --model-lib-path /home/orangepi/mlc-llm/dist/libs/Llama-2-7b-chat-hf-q4f16_1-opencl-1024.so --device opencl:0 --host 192.168.1.108

  I tried reduce prefill-chunk-size here, still no help.

• Output from executing the code in a conda environment

  [2024-03-25 15:35:00] INFO engine.py:238: Estimated KVCacheConfig "max_total_sequence_length": 21275.
  [2024-03-25 15:35:00] INFO engine.py:243: Estimated total single GPU memory usage: 14361.14 MB (Parameters: 3615.13 MB. KVCache: 10689.56 MB. Temporary buffer: 56.44 MB)

Expected behavior

That's the KVCache takes when using Gradio.

INFO model_metadata.py:96: Total memory usage: 4077.14 MB (Parameters: 3615.13 MB. KVCache: 0.00 MB. Temporary buffer: 462.01 MB)

Environment

• Platform (OpenCL):
• Operating system (Ubuntu.):
• Device (RK3588)
• How you installed MLC-LLM (source):
• How you installed TVM-Unity (source):
• Python version (3.11):
• TVM Unity Hash Tag (python -c "import tvm; print('\n'.join(f'{k}: {v}' for k, v in tvm.support.libinfo().items()))", applicable if you compile models):

  USE_NVTX: OFF
  USE_GTEST: AUTO
  SUMMARIZE: OFF
  USE_IOS_RPC: OFF
  USE_MSC: OFF
  USE_ETHOSU: OFF
  CUDA_VERSION: NOT-FOUND
  USE_LIBBACKTRACE: AUTO
  DLPACK_PATH: 3rdparty/dlpack/include
  USE_TENSORRT_CODEGEN: OFF
  USE_THRUST: OFF
  USE_TARGET_ONNX: OFF
  USE_AOT_EXECUTOR: ON
  BUILD_DUMMY_LIBTVM: OFF
  USE_CUDNN: OFF
  USE_TENSORRT_RUNTIME: OFF
  USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR: OFF
  USE_CCACHE: AUTO
  USE_ARM_COMPUTE_LIB: OFF
  USE_CPP_RTVM: OFF
  USE_OPENCL_GTEST: /path/to/opencl/gtest
  USE_MKL: OFF
  USE_PT_TVMDSOOP: OFF
  MLIR_VERSION: NOT-FOUND
  USE_CLML: OFF
  USE_STACKVM_RUNTIME: OFF
  USE_GRAPH_EXECUTOR_CUDA_GRAPH: OFF
  ROCM_PATH: /opt/rocm
  USE_DNNL: OFF
  USE_VITIS_AI: OFF
  USE_MLIR: OFF
  USE_RCCL: OFF
  USE_LLVM: /home/orangepi/clang+llvm-17.0.2-aarch64-linux-gnu/bin/llvm-config
  USE_VERILATOR: OFF
  USE_TF_TVMDSOOP: OFF
  USE_THREADS: ON
  USE_MSVC_MT: OFF
  BACKTRACE_ON_SEGFAULT: OFF
  USE_GRAPH_EXECUTOR: ON
  USE_NCCL: OFF
  USE_ROCBLAS: OFF
  GIT_COMMIT_HASH: 1d4da926c726e2700593c7f62006545bda6a46f9
  USE_VULKAN: OFF
  USE_RUST_EXT: OFF
  USE_CUTLASS: OFF
  USE_CPP_RPC: OFF
  USE_HEXAGON: OFF
  USE_CUSTOM_LOGGING: OFF
  USE_UMA: OFF
  USE_FALLBACK_STL_MAP: OFF
  USE_SORT: ON
  USE_RTTI: ON
  GIT_COMMIT_TIME: 2024-03-12 09:50:16 -0400
  USE_HEXAGON_SDK: /path/to/sdk
  USE_BLAS: none
  USE_ETHOSN: OFF
  USE_LIBTORCH: OFF
  USE_RANDOM: ON
  USE_CUDA: OFF
  USE_COREML: OFF
  USE_AMX: OFF
  BUILD_STATIC_RUNTIME: OFF
  USE_CMSISNN: OFF
  USE_KHRONOS_SPIRV: OFF
  USE_CLML_GRAPH_EXECUTOR: OFF
  USE_TFLITE: OFF
  USE_HEXAGON_GTEST: /path/to/hexagon/gtest
  PICOJSON_PATH: 3rdparty/picojson
  USE_OPENCL_ENABLE_HOST_PTR: OFF
  INSTALL_DEV: OFF
  USE_PROFILER: ON
  USE_NNPACK: OFF
  LLVM_VERSION: 17.0.2
  USE_MRVL: OFF
  USE_OPENCL: ON
  COMPILER_RT_PATH: 3rdparty/compiler-rt
  RANG_PATH: 3rdparty/rang/include
  USE_SPIRV_KHR_INTEGER_DOT_PRODUCT: OFF
  USE_OPENMP: none
  USE_BNNS: OFF
  USE_CUBLAS: OFF
  USE_METAL: OFF
  USE_MICRO_STANDALONE_RUNTIME: OFF
  USE_HEXAGON_EXTERNAL_LIBS: OFF
  USE_ALTERNATIVE_LINKER: AUTO
  USE_BYODT_POSIT: OFF
  USE_HEXAGON_RPC: OFF
  USE_MICRO: OFF
  DMLC_PATH: 3rdparty/dmlc-core/include
  INDEX_DEFAULT_I64: ON
  USE_RELAY_DEBUG: OFF
  USE_RPC: ON
  USE_TENSORFLOW_PATH: none
  TVM_CLML_VERSION:
  USE_MIOPEN: OFF
  USE_ROCM: OFF
  USE_PAPI: OFF
  USE_CURAND: OFF
  TVM_CXX_COMPILER_PATH: /usr/bin/c++
  HIDE_PRIVATE_SYMBOLS: ON

• Any other relevant information:

Additional context

[MasterJH5574]

Hi @gxmlfx, this is intentional because the automatic estimation wants to maximize the KV cache capacity for serving. To limit the GPU memory taken by KV cache, you can manually specify --max-total-seq-length with a smaller value. Or you can set the environment variable MLC_GPU_SIZE_BYTES to the number of bytes you want the server to use in total (including parameters, KV cache, temporary buffers, etc.).

https://github.com/mlc-ai/mlc-llm/blob/f04cd3e9e81bcd3c02015df6fe0f0eaa9ffd8453/python/mlc_llm/serve/engine.py#L208-L215

Thank you for the feedback! We will provide detailed documentation regarding this.

[gxmlfx]

Hi @gxmlfx, this is intentional because the automatic estimation wants to maximize the KV cache capacity for serving. To limit the GPU memory taken by KV cache, you can manually specify --max-total-seq-length with a smaller value. Or you can set the environment variable MLC_GPU_SIZE_BYTES to the number of bytes you want the server to use in total (including parameters, KV cache, temporary buffers, etc.).您好，这是有意为之的，因为自动估计希望最大化 KV 缓存容量来提供服务。为了限制KV缓存占用的GPU内存，您可以手动指定较小的值 --max-total-seq-length 。或者你可以将环境变量 MLC_GPU_SIZE_BYTES 设置为你希望服务器总共使用的字节数（包括参数、KV缓存、临时缓冲区等）。

https://github.com/mlc-ai/mlc-llm/blob/f04cd3e9e81bcd3c02015df6fe0f0eaa9ffd8453/python/mlc_llm/serve/engine.py#L208-L215

Thank you for the feedback! We will provide detailed documentation regarding this.感谢您的反馈！我们将提供与此相关的详细文档。

@MasterJH5574 Thanks! That works well, I changed int(gpu_size_bytes) 0.90 to int(gpu_size_bytes) 0.50, in https://github.com/mlc-ai/mlc-llm/blob/f04cd3e9e81bcd3c02015df6fe0f0eaa9ffd8453/python/mlc_llm/serve/engine.py#L217-L225 Appriciate for upcoming new document.