Merge branch 'mlc-ai:main' into main

CMakeLists.txt

@@ -54,15 +54,15 @@ set(USE_PROFILER OFF)
 set(USE_GTEST OFF)
 set(USE_LIBBACKTRACE OFF)
 set(BUILD_DUMMY_LIBTVM ON)
-if (NOT DEFINED TVM_HOME)
-  if(DEFINED ENV{TVM_HOME})
-    set(TVM_HOME "$ENV{TVM_HOME}")
+if (NOT DEFINED TVM_SOURCE_DIR)
+  if(DEFINED ENV{TVM_SOURCE_DIR})
+    set(TVM_SOURCE_DIR "$ENV{TVM_SOURCE_DIR}")
   else()
-    set(TVM_HOME 3rdparty/tvm)
-  endif(DEFINED ENV{TVM_HOME})
-endif (NOT DEFINED TVM_HOME)
-message(STATUS "TVM_HOME: ${TVM_HOME}")
-add_subdirectory(${TVM_HOME} tvm EXCLUDE_FROM_ALL)
+    set(TVM_SOURCE_DIR 3rdparty/tvm)
+  endif(DEFINED ENV{TVM_SOURCE_DIR})
+endif (NOT DEFINED TVM_SOURCE_DIR)
+message(STATUS "TVM_SOURCE_DIR: ${TVM_SOURCE_DIR}")
+add_subdirectory(${TVM_SOURCE_DIR} tvm EXCLUDE_FROM_ALL)
 
 set(MLC_LLM_RUNTIME_LINKER_LIB "")
 set(TOKENZIER_CPP_PATH 3rdparty/tokenizers-cpp)
@@ -74,10 +74,10 @@ add_library(mlc_llm_objs OBJECT ${MLC_LLM_SRCS})
 
 set(
   MLC_LLM_INCLUDES
-  ${TVM_HOME}/include
-  ${TVM_HOME}/3rdparty/dlpack/include
-  ${TVM_HOME}/3rdparty/dmlc-core/include
-  ${TVM_HOME}/3rdparty/picojson
+  ${TVM_SOURCE_DIR}/include
+  ${TVM_SOURCE_DIR}/3rdparty/dlpack/include
+  ${TVM_SOURCE_DIR}/3rdparty/dmlc-core/include
+  ${TVM_SOURCE_DIR}/3rdparty/picojson
 )
 
 set(MLC_LLM_COMPILE_DEFS ${MLC_LLM_COMPILE_DEFS} DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
@@ -100,7 +100,7 @@ target_link_libraries(mlc_llm PRIVATE tokenizers_cpp)
 
 find_library(
   FLASH_ATTN_LIBRARY flash_attn
-  HINTS ${TVM_HOME}/*/3rdparty/libflash_attn/src
+  HINTS ${TVM_SOURCE_DIR}/*/3rdparty/libflash_attn/src
 )
 
 if (FLASH_ATTN_LIBRARY STREQUAL "FLASH_ATTN_LIBRARY-NOTFOUND")

android/MLCEngineExample/mlc-package-config.json

@@ -7,4 +7,4 @@
             "model_id": "phi-2-q4f16_1"
         }
     ]
-}
+}

android/mlc4j/CMakeLists.txt

@@ -10,10 +10,10 @@ set(MLC_LLM_BINARY_DIR mlc_llm)
 set(MLC_LLM_COMPILE_DEFS TVM_LOG_CUSTOMIZE=1)
 add_subdirectory(${MLC_LLM_DIR} ${MLC_LLM_BINARY_DIR} EXCLUDE_FROM_ALL)
 
-if (NOT DEFINED TVM_HOME)
-  set(TVM_HOME ${MLC_LLM_DIR}/3rdparty/tvm)
-endif (NOT DEFINED TVM_HOME)
-message(STATUS "TVM_HOME: ${TVM_HOME}")
+if (NOT DEFINED TVM_SOURCE_DIR)
+  set(TVM_SOURCE_DIR ${MLC_LLM_DIR}/3rdparty/tvm)
+endif (NOT DEFINED TVM_SOURCE_DIR)
+message(STATUS "TVM_SOURCE_DIR: ${TVM_SOURCE_DIR}")
 
 find_package(Java REQUIRED)
 find_package(JNI REQUIRED)
@@ -25,29 +25,29 @@ include(UseJava)
 
 
 file(GLOB_RECURSE javasources
-    ${TVM_HOME}/jvm/core/src/main/java/org/apache/tvm/*.java
+    ${TVM_SOURCE_DIR}/jvm/core/src/main/java/org/apache/tvm/*.java
     ${ANDROID_DIR}/src/java/*.java
 )
 set(JNI_HEADER ${CMAKE_BINARY_DIR}/jni_header)
 add_jar(tvm4j_core ${javasources} GENERATE_NATIVE_HEADERS tvm4jheaders DESTINATION ${JNI_HEADER})
 
 add_custom_command(
-  TARGET tvm4j_core POST_BUILD 
+  TARGET tvm4j_core POST_BUILD
   COMMAND ${CMAKE_COMMAND} -E copy ${JNI_HEADER}/org_apache_tvm_LibInfo.h ${JNI_HEADER}/org_apache_tvm_native_c_api.h
 )
 
 add_library(model_android STATIC IMPORTED)
 set_target_properties(model_android PROPERTIES IMPORTED_LOCATION ${ANDROID_BIN_DIR}/lib/libmodel_android.a)
 
-add_library(tvm4j_runtime_packed SHARED ${TVM_HOME}/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc)
+add_library(tvm4j_runtime_packed SHARED ${TVM_SOURCE_DIR}/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc)
 
 target_include_directories(tvm4j_runtime_packed PUBLIC
   ${JNI_INCLUDE_DIRS}
   ${JNI_HEADER}
   ${ANDROID_DIR}/src/cpp
-  ${TVM_HOME}/3rdparty/dlpack/include
-  ${TVM_HOME}/3rdparty/dmlc-core/include
-  ${TVM_HOME}/include
+  ${TVM_SOURCE_DIR}/3rdparty/dlpack/include
+  ${TVM_SOURCE_DIR}/3rdparty/dmlc-core/include
+  ${TVM_SOURCE_DIR}/include
 )
 
 target_link_libraries(tvm4j_runtime_packed

android/mlc4j/prepare_libs.py

@@ -71,7 +71,7 @@ def run_cmake_install():
     subprocess.run(cmd, check=True, env=os.environ)
 
 
-def main(mlc_llm_home: Path):
+def main(mlc_source_dir: Path):
     # - Setup rust.
     subprocess.run(["rustup", "target", "add", "aarch64-linux-android"], check=True, env=os.environ)
 
@@ -81,13 +81,13 @@ def main(mlc_llm_home: Path):
     logger.info('Entering "%s" for MLC LLM and tvm4j build.', os.path.abspath(build_path))
     os.chdir(build_path)
     # Generate config.cmake if TVM Home is set.
-    if "TVM_HOME" in os.environ:
-        logger.info('Set TVM_HOME to "%s"', os.environ["TVM_HOME"])
+    if "TVM_SOURCE_DIR" in os.environ:
+        logger.info('Set TVM_SOURCE_DIR to "%s"', os.environ["TVM_SOURCE_DIR"])
         with open("config.cmake", "w", encoding="utf-8") as file:
-            print("set(TVM_HOME ${%s})" % os.environ["TVM_HOME"], file=file)
+            print("set(TVM_SOURCE_DIR ${%s})" % os.environ["TVM_SOURCE_DIR"], file=file)
 
     # - Run cmake, build and install
-    run_cmake(mlc_llm_home / "android" / "mlc4j")
+    run_cmake(mlc_source_dir / "android" / "mlc4j")
     run_cmake_build()
     run_cmake_install()
 
@@ -98,11 +98,11 @@ def main(mlc_llm_home: Path):
     parser.add_argument(
         "--mlc-llm-home",
         type=Path,
-        default=os.environ.get("MLC_LLM_HOME", None),
+        default=os.environ.get("MLC_LLM_SOURCE_DIR", None),
         help="The path to MLC LLM source",
     )
     parsed = parser.parse_args()
-    if parsed.mlc_llm_home is None:
-        parsed.mlc_llm_home = Path(os.path.abspath(os.path.curdir)).parent.parent
-    os.environ["MLC_LLM_HOME"] = str(parsed.mlc_llm_home)
-    main(parsed.mlc_llm_home)
+    if parsed.mlc_source_dir is None:
+        parsed.mlc_source_dir = Path(os.path.abspath(os.path.curdir)).parent.parent
+    os.environ["MLC_LLM_SOURCE_DIR"] = str(parsed.mlc_source_dir)
+    main(parsed.mlc_source_dir)

ci/task/test_model_compile.sh

@@ -20,10 +20,11 @@ elif [[ ${GPU} == metal ]]; then
 elif [[ ${GPU} == wasm* ]]; then
 	TARGET=wasm
 	pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly
-	export TVM_HOME=$(dirname $(python -c 'import tvm; print(tvm.__file__)'))
-	export MLC_LLM_HOME=$(pwd)
-	cd $TVM_HOME/web/ && make -j${NUM_THREADS} && cd -
-	cd $MLC_LLM_HOME/web/ && make -j${NUM_THREADS} && cd -
+	export TVM_SOURCE_DIR=$(dirname $(python -c 'import tvm; print(tvm.__file__)'))
+	export TVM_HOME=${TVM_SOURCE_DIR}
+	export MLC_LLM_SOURCE_DIR=$(pwd)
+	cd $TVM_SOURCE_DIR/web/ && make -j${NUM_THREADS} && cd -
+	cd $MLC_LLM_SOURCE_DIR/web/ && make -j${NUM_THREADS} && cd -
 elif [[ ${GPU} == ios ]]; then
 	TARGET=ios
 	pip install --pre -U --force-reinstal -f https://mlc.ai/wheels mlc-ai-nightly

cmake/gen_cmake_config.py

@@ -6,12 +6,12 @@
     tvm_home = ""  # pylint: disable=invalid-name
 
     tvm_home = input(
-        "Enter TVM_HOME in absolute path. If not specified, 3rdparty/tvm will be used by default: "
+        "Enter TVM_SOURCE_DIR in absolute path. If not specified, 3rdparty/tvm will be used by default: "
     )
     if len(tvm_home) == 0:
         tvm_home = "3rdparty/tvm"  # pylint: disable=invalid-name
 
-    cmake_config_str = f"set(TVM_HOME {tvm_home})\n"
+    cmake_config_str = f"set(TVM_SOURCE_DIR {tvm_home})\n"
     cmake_config_str += "set(CMAKE_BUILD_TYPE RelWithDebInfo)\n"
     backends = [
         Backend("CUDA", "USE_CUDA", "Use CUDA? (y/n): "),

cpp/json_ffi/json_ffi_engine.cc

@@ -85,7 +85,7 @@ bool JSONFFIEngine::AddRequest(std::string request_json_str, std::string request
                                   request.top_logprobs, request.logit_bias, request.seed,
                                   request.ignore_eos, request.max_tokens, std::move(stop_strs),
                                   conv_template_.stop_token_ids, /*response_format=*/std::nullopt,
-                                  this->default_generation_cfg_json_str_);
+                                  request.debug_config, this->default_generation_cfg_json_str_);
 
   Request engine_request(request_id, inputs, generation_cfg);
   this->engine_->AddRequest(engine_request);

cpp/json_ffi/openai_api_protocol.h

@@ -13,6 +13,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "../serve/config.h"
 #include "../support/result.h"
 #include "picojson.h"
 
@@ -127,6 +128,8 @@ class RequestResponseFormat {
   std::optional<std::string> json_schema = std::nullopt;
 };
 
+using serve::DebugConfig;
+
 class ChatCompletionRequest {
  public:
   std::vector<ChatCompletionMessage> messages;
@@ -148,6 +151,7 @@ class ChatCompletionRequest {
   std::optional<std::string> user = std::nullopt;
   bool ignore_eos = false;
   //   RequestResponseFormat response_format; //TODO: implement this
+  std::optional<DebugConfig> debug_config = std::nullopt;
 
   /*! \brief Parse and create a ChatCompletionRequest instance from the given JSON string. */
   static Result<ChatCompletionRequest> FromJSON(const std::string& json_str);

cpp/llm_chat.cc

@@ -26,8 +26,8 @@
 #include "./metadata/model.h"
 #include "./serve/config.h"
 #include "./support/load_bytes_from_file.h"
+#include "./support/random.h"
 #include "conversation.h"
-#include "random.h"
 #include "tokenizers.h"
 
 namespace mlc {