Merge branch 'main' into feat/mrope-stage1

Author: MagellaX

.github/ISSUE_TEMPLATE/bug-report.md

@@ -31,11 +31,11 @@ Steps to reproduce the behavior:
  - Operating system (e.g. Ubuntu/Windows/MacOS/...):
  - Device (e.g. iPhone 12 Pro, PC+RTX 3090, ...)
  - How you installed MLC-LLM (`conda`, source):
- - How you installed TVM-Unity (`pip`, source):
+ - How you installed TVM (`pip`, source):
  - Python version (e.g. 3.10):
  - GPU driver version (if applicable):
  - CUDA/cuDNN version (if applicable):
- - 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):
+ - TVM 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):
  - Any other relevant information:
 
 ## Additional context

3rdparty/tokenizers-cpp

@@ -24,6 +24,8 @@ if(NOT CMAKE_BUILD_TYPE)
 endif(NOT CMAKE_BUILD_TYPE)
 
 option(MLC_HIDE_PRIVATE_SYMBOLS "Hide private symbols" ON)
+option(MLC_LLM_BUILD_PYTHON_MODULE "Build Python module with scikit-build-core"
+       OFF)
 
 if(MLC_LLM_INSTALL_STATIC_LIB)
   set(BUILD_STATIC_RUNTIME ON)
@@ -184,3 +186,32 @@ else()
     RUNTIME DESTINATION bin
     LIBRARY DESTINATION lib${LIB_SUFFIX})
 endif()
+
+# Python package installation configuration This section ensures that all
+# necessary files are installed for the Python wheel
+if(MLC_LLM_BUILD_PYTHON_MODULE)
+  message(STATUS "Configuring Python package installation")
+
+  # Set RPATH for mlc_llm and mlc_llm_module to find other libraries relatively
+  if(APPLE)
+    # macOS uses @loader_path
+    set_target_properties(mlc_llm PROPERTIES INSTALL_RPATH "@loader_path")
+    set_target_properties(mlc_llm_module PROPERTIES INSTALL_RPATH
+                                                    "@loader_path")
+  elseif(LINUX)
+    # Linux uses $ORIGIN
+    set_target_properties(mlc_llm PROPERTIES INSTALL_RPATH "\$ORIGIN")
+    set_target_properties(mlc_llm_module PROPERTIES INSTALL_RPATH "\$ORIGIN")
+  endif()
+
+  # Install compiled shared libraries
+  install(TARGETS mlc_llm DESTINATION ".")
+  install(TARGETS mlc_llm_module DESTINATION ".")
+  install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/cpp/" DESTINATION "cpp/")
+  install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/web/" DESTINATION "web/")
+  install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/README.md"
+                "${CMAKE_CURRENT_SOURCE_DIR}/LICENSE"
+                "${CMAKE_CURRENT_SOURCE_DIR}/NOTICE" DESTINATION ".")
+
+  message(STATUS "Python package installation configured")
+endif()

3rdparty/tvm

@@ -51,7 +51,7 @@ class AppViewModel(application: Application) : AndroidViewModel(application) {
     companion object {
         const val AppConfigFilename = "mlc-app-config.json"
         const val ModelConfigFilename = "mlc-chat-config.json"
-        const val ParamsConfigFilename = "ndarray-cache.json"
+        const val ParamsConfigFilename = "tensor-cache.json"
         const val ModelUrlSuffix = "resolve/main/"
     }
 

CMakeLists.txt

@@ -57,14 +57,14 @@ target_include_directories(
   PUBLIC ${JNI_INCLUDE_DIRS}
          ${JNI_HEADER}
          ${ANDROID_DIR}/src/cpp
-         ${TVM_SOURCE_DIR}/3rdparty/dlpack/include
+         ${TVM_SOURCE_DIR}/3rdparty/tvm-ffi/3rdparty/dlpack/include
          ${TVM_SOURCE_DIR}/3rdparty/dmlc-core/include
          ${TVM_SOURCE_DIR}/3rdparty/OpenCL-Headers
          ${TVM_SOURCE_DIR}/3rdparty/picojson
          ${TVM_SOURCE_DIR}/include
          ${TVM_SOURCE_DIR}/src
-         ${TVM_SOURCE_DIR}/ffi/include
-         ${TVM_SOURCE_DIR}/ffi/src)
+         ${TVM_SOURCE_DIR}/3rdparty/tvm-ffi/include
+         ${TVM_SOURCE_DIR}/3rdparty/tvm-ffi/src)
 target_compile_definitions(tvm4j_runtime_packed PUBLIC ${MLC_LLM_COMPILE_DEFS})
 target_compile_definitions(
   tvm4j_runtime_packed

android/MLCChat/app/src/main/java/ai/mlc/mlcchat/AppViewModel.kt

@@ -6,37 +6,41 @@
 #include <dmlc/logging.h>
 #include <dmlc/thread_local.h>
 
+#include <ffi/backtrace.cc>
 #include <ffi/container.cc>
 #include <ffi/dtype.cc>
 #include <ffi/error.cc>
+#include <ffi/extra/env_c_api.cc>
+#include <ffi/extra/env_context.cc>
+#include <ffi/extra/library_module.cc>
+#include <ffi/extra/library_module_dynamic_lib.cc>
+#include <ffi/extra/library_module_system_lib.cc>
+#include <ffi/extra/module.cc>
 #include <ffi/function.cc>
 #include <ffi/object.cc>
-#include <ffi/traceback.cc>
 #include <runtime/cpu_device_api.cc>
 #include <runtime/device_api.cc>
 #include <runtime/file_utils.cc>
-#include <runtime/library_module.cc>
 #include <runtime/logging.cc>
 #include <runtime/memory/memory_manager.cc>
 #include <runtime/module.cc>
-#include <runtime/ndarray.cc>
 #include <runtime/nvtx.cc>
 #include <runtime/opencl/opencl_device_api.cc>
 #include <runtime/opencl/opencl_module.cc>
 #include <runtime/opencl/opencl_wrapper/opencl_wrapper.cc>
 #include <runtime/profiling.cc>
 #include <runtime/source_utils.cc>
-#include <runtime/system_library.cc>
+#include <runtime/tensor.cc>
 #include <runtime/thread_pool.cc>
 #include <runtime/threading_backend.cc>
 #include <runtime/vm/attn_backend.cc>
 #include <runtime/vm/builtin.cc>
 #include <runtime/vm/bytecode.cc>
 #include <runtime/vm/executable.cc>
 #include <runtime/vm/kv_state.cc>
-#include <runtime/vm/ndarray_cache_support.cc>
 #include <runtime/vm/paged_kv_cache.cc>
 #include <runtime/vm/rnn_state.cc>
+#include <runtime/vm/tensor_cache_support.cc>
 #include <runtime/vm/vm.cc>
 #include <runtime/workspace_pool.cc>
 

android/mlc4j/CMakeLists.txt

@@ -125,7 +125,6 @@ stage('Build') {
           sh(script: "ls -alh", label: 'Show work directory')
           sh(script: "${pkg_cuda} conda env export --name py312", label: 'Checkout version')
           sh(script: "${pkg_cuda} -j 8 -v \$HOME/.ccache /ccache conda run -n py312 ./ci/task/build_lib.sh", label: 'Build MLC LLM runtime')
-          sh(script: "${pkg_cuda} -j 8 conda run -n py312 ./ci/task/build_wheel.sh", label: 'Build MLC LLM wheel')
           sh(script: "${pkg_cuda} -j 1 conda run -n py312 ./ci/task/build_clean.sh", label: 'Clean up after build')
           sh(script: "ls -alh ./wheels/", label: 'Build artifact')
           pack_lib('mlc_wheel_cuda', 'wheels/*.whl')
@@ -139,7 +138,6 @@ stage('Build') {
     //       sh(script: "ls -alh", label: 'Show work directory')
     //       sh(script: "${pkg_rocm} conda env export --name py38", label: 'Checkout version')
     //       sh(script: "${pkg_rocm} -j 8 conda run -n py38 ./ci/task/build_lib.sh", label: 'Build MLC LLM runtime')
-    //       sh(script: "${pkg_rocm} -j 8 conda run -n py38 ./ci/task/build_wheel.sh", label: 'Build MLC LLM wheel')
     //       sh(script: "${pkg_rocm} -j 1 conda run -n py38 ./ci/task/build_clean.sh", label: 'Clean up after build')
     //       sh(script: "ls -alh ./wheels/", label: 'Build artifact')
     //       pack_lib('mlc_wheel_rocm', 'wheels/*.whl')
@@ -153,7 +151,6 @@ stage('Build') {
           sh(script: "ls -alh", label: 'Show work directory')
           sh(script: "conda env export --name mlc-llm-ci", label: 'Checkout version')
           sh(script: "NUM_THREADS=6 GPU=metal conda run -n mlc-llm-ci ./ci/task/build_lib.sh", label: 'Build MLC LLM runtime')
-          sh(script: "NUM_THREADS=6 GPU=metal conda run -n mlc-llm-ci ./ci/task/build_wheel.sh", label: 'Build MLC LLM wheel')
           sh(script: "NUM_THREADS=6 GPU=metal conda run -n mlc-llm-ci ./ci/task/build_clean.sh", label: 'Clean up after build')
           sh(script: "ls -alh ./wheels/", label: 'Build artifact')
           pack_lib('mlc_wheel_metal', 'wheels/*.whl')
@@ -167,7 +164,6 @@ stage('Build') {
           sh(script: "ls -alh", label: 'Show work directory')
           sh(script: "${pkg_cpu} conda env export --name py312", label: 'Checkout version')
           sh(script: "${pkg_cpu} -j 8 conda run -n py312 ./ci/task/build_lib.sh", label: 'Build MLC LLM runtime')
-          sh(script: "${pkg_cpu} -j 8 conda run -n py312 ./ci/task/build_wheel.sh", label: 'Build MLC LLM wheel')
           sh(script: "${pkg_cpu} -j 1 conda run -n py312 ./ci/task/build_clean.sh", label: 'Clean up after build')
           sh(script: "ls -alh ./wheels/", label: 'Build artifact')
           pack_lib('mlc_wheel_vulkan', 'wheels/*.whl')

android/mlc4j/src/cpp/tvm_runtime.h

@@ -6,6 +6,4 @@ set -x
 : ${GPU:="cpu"}
 
 rm -rf ${WORKSPACE_CWD}/build/ \
-    ${WORKSPACE_CWD}/python/dist/ \
-    ${WORKSPACE_CWD}/python/build/ \
-    ${WORKSPACE_CWD}/python/mlc_llm.egg-info
+    ${WORKSPACE_CWD}/dist/

ci/jenkinsfile.groovy

@@ -9,14 +9,18 @@ export CCACHE_NOHASHDIR=1
 export CCACHE_DIR=/ccache
 
 # Temporary workaround to install ccache.
+if [[ ${GPU} != metal ]]; then
+    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
+    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
+fi
 conda install -c conda-forge ccache
 
 if [[ ${GPU} != metal ]]; then
     source /multibuild/manylinux_utils.sh
     source /opt/rh/gcc-toolset-11/enable # GCC-11 is the hightest GCC version compatible with NVCC < 12
 fi
 
-mkdir -p $WORKSPACE_CWD/build/ && cd $WORKSPACE_CWD/build/
+mkdir -p $WORKSPACE_CWD/build
 if [[ ${GPU} == rocm* ]]; then
     echo set\(USE_VULKAN ON\) >>config.cmake
     echo set\(USE_ROCM ON\) >>config.cmake
@@ -29,17 +33,6 @@ elif [[ ${GPU} == cuda* ]]; then
     echo set\(USE_CUDA ON\) >>config.cmake
     echo set\(USE_CUBLAS ON\) >>config.cmake
     echo set\(USE_NCCL ON\) >>config.cmake
-    echo set\(USE_FLASHINFER ON\) >>config.cmake
-    echo set\(FLASHINFER_ENABLE_FP8 OFF\) >>config.cmake
-    echo set\(FLASHINFER_ENABLE_BF16 OFF\) >>config.cmake
-    echo set\(FLASHINFER_GEN_GROUP_SIZES 1 4 6 8\) >>config.cmake
-    echo set\(FLASHINFER_GEN_PAGE_SIZES 16\) >>config.cmake
-    echo set\(FLASHINFER_GEN_HEAD_DIMS 128\) >>config.cmake
-    echo set\(FLASHINFER_GEN_KV_LAYOUTS 0 1\) >>config.cmake
-    echo set\(FLASHINFER_GEN_POS_ENCODING_MODES 0 1\) >>config.cmake
-    echo set\(FLASHINFER_GEN_ALLOW_FP16_QK_REDUCTIONS "false"\) >>config.cmake
-    echo set\(FLASHINFER_GEN_CASUALS "false" "true"\) >>config.cmake
-    echo set\(USE_CUTLASS ON\) >>config.cmake
 elif [[ ${GPU} == metal ]]; then
     export CCACHE_DIR=$HOME/ci/ccache
     echo set\(USE_METAL ON\) >>config.cmake
@@ -49,4 +42,26 @@ fi
 
 cat config.cmake
 
-cmake -DCMAKE_POLICY_VERSION_MINIMUM=3.5 .. && make -j${NUM_THREADS}
+AUDITWHEEL_OPTS="--plat ${AUDITWHEEL_PLAT} -w repaired_wheels/"
+AUDITWHEEL_OPTS="--exclude libtvm --exclude libtvm_runtime --exclude libtvm_ffi --exclude libvulkan ${AUDITWHEEL_OPTS}"
+if [[ ${GPU} == rocm* ]]; then
+    AUDITWHEEL_OPTS="--exclude libamdhip64 --exclude libhsa-runtime64 --exclude librocm_smi64 --exclude librccl ${AUDITWHEEL_OPTS}"
+elif [[ ${GPU} == cuda* ]]; then
+    AUDITWHEEL_OPTS="--exclude libcuda --exclude libcudart --exclude libnvrtc --exclude libcublas --exclude libcublasLt ${AUDITWHEEL_OPTS}"
+fi
+
+rm -rf ${WORKSPACE_CWD}/dist
+cd ${WORKSPACE_CWD} && pip wheel --no-deps -w dist . -v
+
+rm -rf ${WORKSPACE_CWD}/wheels/
+if [[ ${GPU} != metal ]]; then
+    mkdir -p ${WORKSPACE_CWD}/repaired_wheels
+    rm -rf ${WORKSPACE_CWD}/repaired_wheels/*
+    auditwheel repair ${AUDITWHEEL_OPTS} dist/*.whl
+    mv ${WORKSPACE_CWD}/repaired_wheels/ ${WORKSPACE_CWD}/wheels/
+else
+    mkdir ${WORKSPACE_CWD}/wheels/
+    mv dist/*.whl ${WORKSPACE_CWD}/wheels/
+fi
+
+chown -R $ENV_USER_ID:$ENV_GROUP_ID ${WORKSPACE_CWD}/wheels/