commit 770d92cb1f8e7e155a617699098d5ad516e1ccf0 Author: wehub-resource-sync Date: Mon Jul 13 13:23:58 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..b61ed53 --- /dev/null +++ b/.clang-format @@ -0,0 +1,8 @@ +# Run the following command to reformat a file: +# clang-format -i -style=Google +# Or use clang-format-diff to only reformat the changed lines: +# https://clang.llvm.org/docs/ClangFormat.html +BasedOnStyle: Google +DerivePointerAlignment: false +ColumnLimit: 100 +PointerAlignment: Left diff --git a/.github/ISSUE_TEMPLATE/bug-report.md b/.github/ISSUE_TEMPLATE/bug-report.md new file mode 100644 index 0000000..2633692 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug-report.md @@ -0,0 +1,43 @@ +--- +name: "🐛 Bug Report" +about: Submit a bug report to help us improve MLC-LLM +title: '[Bug] ' +labels: ['bug'] +assignees: '' + +--- + +## 🐛 Bug + + + +## To Reproduce + +Steps to reproduce the behavior: + +1. +1. +1. + + + +## Expected behavior + + + +## Environment + + - Platform (e.g. WebGPU/Vulkan/IOS/Android/CUDA): + - 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 (`pip`, source): + - Python version (e.g. 3.10): + - GPU driver version (if applicable): + - CUDA/cuDNN version (if applicable): + - 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 + + diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000..6e88ecf --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,9 @@ +blank_issues_enabled: false + +contact_links: + - name: Check the MLC-LLM Documentation + url: https://llm.mlc.ai/docs/ + about: Our documentation might provide answers to your questions. + - name: Chat on Discord + url: https://discord.gg/9Xpy2HGBuD + about: Join the Discord Server to live chat with the community. diff --git a/.github/ISSUE_TEMPLATE/documentation.md b/.github/ISSUE_TEMPLATE/documentation.md new file mode 100644 index 0000000..c89d99e --- /dev/null +++ b/.github/ISSUE_TEMPLATE/documentation.md @@ -0,0 +1,17 @@ +--- +name: "\U0001F4DA Documentation" +about: Report an issue related to https://llm.mlc.ai/docs/ +title: '[Doc] ' +labels: ['documentation'] +assignees: '' + +--- + +## 📚 Documentation + +### Suggestion + + +### Bug +- Link to the buggy documentation/tutorial: +- Description of the bug: diff --git a/.github/ISSUE_TEMPLATE/feature-request.md b/.github/ISSUE_TEMPLATE/feature-request.md new file mode 100644 index 0000000..5d92d35 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature-request.md @@ -0,0 +1,23 @@ +--- +name: "\U0001F680 Feature Request" +about: Submit a proposal/request for a new MLC-LLM feature, or an enhancement on existing features. +title: '[Feature Request] ' +labels: ['feature request'] +assignees: '' + +--- + +## 🚀 Feature + + +## Motivation + + + +## Alternatives + + + +## Additional context + + diff --git a/.github/ISSUE_TEMPLATE/general.md b/.github/ISSUE_TEMPLATE/general.md new file mode 100644 index 0000000..61527c7 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/general.md @@ -0,0 +1,12 @@ +--- +name: "❓ General Questions" +about: General questions you have about MLC-LLM. +title: '[Question] ' +labels: ['question'] +assignees: '' + +--- + +## ❓ General Questions + + diff --git a/.github/ISSUE_TEMPLATE/model-request.md b/.github/ISSUE_TEMPLATE/model-request.md new file mode 100644 index 0000000..fb48ce5 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/model-request.md @@ -0,0 +1,17 @@ +--- +name: "️️⚙️ Model Request" +about: Request a new model in MLC-LLM +title: '[Model Request] ' +labels: ['new-models'] +assignees: '' + +--- + +## ⚙️ Request New Models + +- Link to an existing implementation (e.g. Hugging Face/Github): +- Is this model architecture supported by MLC-LLM? (the list of [supported models](https://llm.mlc.ai/docs/prebuilt_models.html)) + +## Additional context + + diff --git a/.github/ISSUE_TEMPLATE/speed-report.md b/.github/ISSUE_TEMPLATE/speed-report.md new file mode 100644 index 0000000..c589969 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/speed-report.md @@ -0,0 +1,24 @@ +--- +name: " 🏎️ Speed Report" +about: Submit a speed report of an model running in MLC-LLM +title: '[Speed] ' +labels: ['performance'] +assignees: '' + +--- + +# 🏎️ Speed Report + + + +- The model code: + + +- The model configuration (e.g. quantization mode, running data type, etc.): +- Device (e.g. MacBook Pro M2, PC+RTX 3080): +- OS (if applicable): +- Encode speed (Token/s): +- Decode speed (Token/s): +- Memory usage (if applicable): + + diff --git a/.github/ISSUE_TEMPLATE/tracking.md b/.github/ISSUE_TEMPLATE/tracking.md new file mode 100644 index 0000000..d84b745 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/tracking.md @@ -0,0 +1,40 @@ +--- +name: "Tracking" +about: A tracking issue that tracks ongoing item in the project +title: '[Tracking] ' +labels: ['status: tracking'] +assignees: '' + +--- + + + + +## Overview + + + + +## Action Items + + +- [ ] + + +## Links to Related Issues and PRs + + + diff --git a/.github/workflows/documentation.yaml b/.github/workflows/documentation.yaml new file mode 100644 index 0000000..6ec3492 --- /dev/null +++ b/.github/workflows/documentation.yaml @@ -0,0 +1,39 @@ +name: Build Docs + +on: + push: + branches: + - main + +jobs: + test_linux: + name: Deploy Docs + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Configuring build Environment + run: | + sudo apt-get update + python -m pip install -U pip wheel + + - name: Setup Ruby + uses: ruby/setup-ruby@v1 + with: + ruby-version: '3.0' + + - name: Installing dependencies + run: | + python -m pip install -r docs/requirements.txt + gem install jekyll jekyll-remote-theme + + - name: Deploying on GitHub Pages + if: github.ref == 'refs/heads/main' + run: | + git remote set-url origin https://x-access-token:${{ secrets.MLC_GITHUB_TOKEN }}@github.com/$GITHUB_REPOSITORY + git config --global user.email "mlc-gh-actions-bot@nomail" + git config --global user.name "mlc-gh-actions-bot" + ./scripts/gh_deploy_site.sh diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..6695be0 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,25 @@ +name: Lint + +on: + push: + branches: [main] + pull_request: + branches: [main] + +concurrency: + group: Lint-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: true + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 + fetch-tags: true + - name: Set up uv + uses: astral-sh/setup-uv@b75a909f75acd358c2196fb9a5f1299a9a8868a4 # v6.7.0 + - name: Set up Python environment + run: uv sync --group lint --no-install-project + - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd # v3.0.1 diff --git a/.github/workflows/update-relax.yaml b/.github/workflows/update-relax.yaml new file mode 100644 index 0000000..eb78900 --- /dev/null +++ b/.github/workflows/update-relax.yaml @@ -0,0 +1,32 @@ +name: 'Relax Submodule Sync' + +on: + workflow_dispatch: + +jobs: + sync: + name: 'Relax Submodule Sync' + runs-on: ubuntu-latest + + defaults: + run: + shell: bash + + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + submodules: true + + - name: Git Sumbodule Update + run: | + git submodule update --remote 3rdparty/tvm + + - name: Commit update + env: + GITHUB_TOKEN: ${{ secrets.MLC_GITHUB_TOKEN }} + run: | + git config --global user.name 'Git bot' + git config --global user.email 'bot@noreply.github.com' + git remote set-url origin https://$GITHUB_TOKEN@github.com/mlc-ai/mlc-llm + git commit -am "Auto updated submodule references" && git push || echo "No changes to commit" diff --git a/.github/workflows/windows-build.yaml b/.github/workflows/windows-build.yaml new file mode 100644 index 0000000..0352c7e --- /dev/null +++ b/.github/workflows/windows-build.yaml @@ -0,0 +1,49 @@ +# GH actions. +# We use it to cover windows builds +# Jenkins is still the primary CI +name: Windows CI + +on: + push: + branches: + - main + pull_request: + branches: + - main + +jobs: + Windows: + runs-on: windows-latest + defaults: + run: + shell: 'cmd /C call {0}' + + steps: + - name: Git config + run: >- + git config --system core.longpaths true + - uses: actions/checkout@v3 + with: + submodules: 'recursive' + - uses: conda-incubator/setup-miniconda@v3 + with: + activate-environment: mlc-llm-build + channel-priority: strict + environment-file: ci/build-environment.yaml + auto-activate-base: false + - name: Conda info + run: | + conda info + conda list + python --version + # The wheel is built with scikit-build-core, which uses the Ninja + # generator on Windows. Ninja needs cl.exe on PATH, so activate the MSVC + # developer environment before building (otherwise CMake aborts with + # "CMAKE_C_COMPILER not set, after EnableLanguage"). + - name: Set up MSVC developer environment + uses: ilammy/msvc-dev-cmd@v1 + with: + arch: x64 + - name: Build MLC-LLM + run: >- + ci/task/build_win.bat diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f9454e5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,324 @@ +tmp/ +dist/ +params/ +debug/ +*.bak +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +.DS_Store + +*.S +# C extensions +*.so + +build/ + +*.ll +.npm +# Distribution / packaging +.Python +env/ +build/ +build-*/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +.conda/ +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Generated by python/gen_requirements.py +python/requirements/*.txt + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ +docs/_staging/ + +# PyBuilder +target/ +/target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject +*~ +*.pyc +*~ +config.mk +config.cmake +Win32 +*.dir +perf +*.wasm +.emscripten + +## IOS +DerivedData/ + +## Java +*.class +jvm/*/target/ +jvm/*/*/target/ +jvm/native/*/generated +jvm/native/src/main/native/org_apache_tvm_native_c_api.h +*.worksheet +*.idea +*.iml +*.classpath +*.project +*.settings +*/node_modules/ + +## Various settings +*.pbxuser +!default.pbxuser +*.mode1v3 +!default.mode1v3 +*.mode2v3 +!default.mode2v3 +*.perspectivev3 +!default.perspectivev3 +xcuserdata/ +.pkl_memoize_* + +.emscripten* +.m2 + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll + +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +## Other +*.moved-aside +*.xccheckout +*.xcscmblueprint +.DS_Store +tags +cscope* +*.lock + +# vim temporary files +*.swp +*.swo + +# TVM generated code +perf +.bash_history +# *.json +*.params +*.ro +*.onnx +*.h5 +synset.txt +cat.jpg +cat.png +docs.tgz +cat.png +*.mlmodel +tvm_u.* +tvm_t.* +# Mac OS X +.DS_Store + +# Jetbrain +.idea +.ipython +.jupyter +.nv +.pylint.d +.python_history +.pytest_cache +.local +cmake-build-debug + +# Visual Studio +.vs + +# Visual Studio Code +.vscode + +# tmp file +.nfs* + +# keys +*.pem +*.p12 +*.pfx +*.cer +*.crt +*.der + +# patch sentinel +patched.txt + +# Python type checking +.mypy_cache/ +.pyre/ + +# pipenv files +Pipfile +Pipfile.lock + +# conda package artifacts +conda/Dockerfile.cuda* +conda/pkg +.node_repl_history +# nix files +.envrc +*.nix + +# Docker files +.sudo_as_admin_successful + +# Downloaded models/datasets +.tvm_test_data +.dgl +.caffe2 + +# Local docs build +_docs/ +jvm/target +.config/configstore/ +.ci-py-scripts/ + +# Generated Hexagon files +src/runtime/hexagon/rpc/hexagon_rpc.h +src/runtime/hexagon/rpc/hexagon_rpc_skel.c +src/runtime/hexagon/rpc/hexagon_rpc_stub.c + +# Local tvm-site checkout +tvm-site/ + +# Generated docs files +gallery/how_to/work_with_microtvm/micro_tvmc.py + +# Test sample data files +!tests/python/ci/sample_prs/*.json + +# Used in CI to communicate between Python and Jenkins +.docker-image-names/ + +# Printed TIR code on disk +*.tir + +# GDB history file +.gdb_history + +dist diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..8939f0e --- /dev/null +++ b/.gitmodules @@ -0,0 +1,18 @@ +[submodule "3rdparty/argparse"] + path = 3rdparty/argparse + url = https://github.com/p-ranav/argparse +[submodule "3rdparty/tokenizers-cpp"] + path = 3rdparty/tokenizers-cpp + url = https://github.com/mlc-ai/tokenizers-cpp +[submodule "3rdparty/googletest"] + path = 3rdparty/googletest + url = https://github.com/google/googletest.git +[submodule "3rdparty/tvm"] + path = 3rdparty/tvm + url = https://github.com/mlc-ai/relax.git +[submodule "3rdparty/stb"] + path = 3rdparty/stb + url = https://github.com/nothings/stb.git +[submodule "3rdparty/xgrammar"] + path = 3rdparty/xgrammar + url = https://github.com/mlc-ai/xgrammar.git diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..9134353 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,68 @@ +# To use: +# +# pre-commit run -a +# +# Or: +# +# pre-commit install # (runs every time you commit in git) +# +# To update this file: +# +# pre-commit autoupdate +# +# See https://github.com/pre-commit/pre-commit +default_install_hook_types: + - pre-commit +repos: + # ---------- General file checks ---------- + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: check-added-large-files + - id: check-case-conflict + - id: check-merge-conflict + - id: check-symlinks + - id: end-of-file-fixer + exclude: ^3rdparty/ + - id: mixed-line-ending + exclude: ^3rdparty/ + - id: requirements-txt-fixer + - id: trailing-whitespace + exclude: ^3rdparty/ + - id: check-yaml + - id: check-toml + + # ---------- YAML lint ---------- + - repo: https://github.com/adrienverge/yamllint + rev: v1.35.1 + hooks: + - id: yamllint + args: + - --config-file + - .yamllint.yaml + exclude: ^\.github/ + + # ---------- TOML format ---------- + - repo: https://github.com/ComPWA/taplo-pre-commit + rev: v0.9.3 + hooks: + - id: taplo-format + + # ---------- Python: ruff ---------- + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.12.3 + hooks: + - id: ruff-check + types_or: [python, pyi, jupyter] + args: [--fix] + - id: ruff-format + types_or: [python, pyi, jupyter] + + # ---------- C/C++: clang-format ---------- + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: "v20.1.8" + hooks: + - id: clang-format + types_or: [c, c++, cuda, objective-c++] + exclude: | + (?x)^(.*cubin\.cpp$|.*fmha_cubin\.h|3rdparty/.*)$ diff --git a/.yamllint.yaml b/.yamllint.yaml new file mode 100644 index 0000000..8f9c5f5 --- /dev/null +++ b/.yamllint.yaml @@ -0,0 +1,15 @@ +extends: default + +rules: + document-start: disable + line-length: + max: 120 + level: warning + truthy: + allowed-values: + - "on" + - "off" + - "yes" + - "no" + - "true" + - "false" diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..c2c9b57 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,218 @@ +cmake_minimum_required(VERSION 3.18) +project(mlc_llm C CXX) + +include(CheckCXXCompilerFlag) +if(MSVC) + set(CMAKE_CXX_FLAGS "/fp:fast ${CMAKE_CXX_FLAGS}") +else() + set(CMAKE_CXX_FLAGS "-ffast-math ${CMAKE_CXX_FLAGS}") +endif() + +if(EXISTS ${CMAKE_BINARY_DIR}/config.cmake) + include(${CMAKE_BINARY_DIR}/config.cmake) +else() + if(EXISTS ${CMAKE_SOURCE_DIR}/config.cmake) + include(${CMAKE_SOURCE_DIR}/config.cmake) + endif() +endif() + +if(NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE + RelWithDebInfo + CACHE STRING "Build type" FORCE) + message(STATUS "Setting default build type to " ${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) +endif() + +set(MLC_VISIBILITY_FLAG "") +if(MLC_HIDE_PRIVATE_SYMBOLS) + set(HIDE_PRIVATE_SYMBOLS ON) + if(NOT MSVC) + set(MLC_VISIBILITY_FLAG "-fvisibility=hidden") + endif() + message(STATUS "Hide private symbols") +endif() + +option(BUILD_CPP_TEST "Build cpp unittests" OFF) + +set(CMAKE_CUDA_STANDARD 17) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +# tvm runtime config: minimize runtime components +set(USE_RPC OFF) +set(USE_MICRO OFF) +set(USE_GRAPH_EXECUTOR OFF) +set(USE_GRAPH_EXECUTOR_DEBUG OFF) +set(USE_AOT_EXECUTOR OFF) +set(USE_PROFILER OFF) +set(USE_GTEST OFF) +set(USE_LIBBACKTRACE OFF) +set(BUILD_DUMMY_LIBTVM ON) +if(NOT DEFINED TVM_SOURCE_DIR) + if(DEFINED ENV{TVM_SOURCE_DIR}) + set(TVM_SOURCE_DIR "$ENV{TVM_SOURCE_DIR}") + else() + 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) +add_subdirectory(${TOKENZIER_CPP_PATH} tokenizers EXCLUDE_FROM_ALL) + +set(XGRAMMAR_PATH 3rdparty/xgrammar) +tvm_file_glob(GLOB_RECURSE MLC_LLM_SRCS cpp/*.cc) +tvm_file_glob(GLOB_RECURSE XGRAMMAR_SRCS ${XGRAMMAR_PATH}/cpp/*.cc) +list(FILTER XGRAMMAR_SRCS EXCLUDE REGEX "${XGRAMMAR_PATH}/cpp/pybind/.*\\.cc") +list(APPEND MLC_LLM_SRCS ${XGRAMMAR_SRCS}) +add_library(mlc_llm_objs OBJECT ${MLC_LLM_SRCS}) + +set(MLC_LLM_INCLUDES + ${TVM_SOURCE_DIR}/include ${TVM_SOURCE_DIR}/3rdparty/dlpack/include) +set(MLC_LLM_COMPILE_DEFS ${MLC_LLM_COMPILE_DEFS} __STDC_FORMAT_MACROS=1) +set(MLC_LLM_COMPILE_DEFS ${MLC_LLM_COMPILE_DEFS} XGRAMMAR_ENABLE_LOG_DEBUG=0) + +target_compile_definitions(mlc_llm_objs PRIVATE ${MLC_LLM_COMPILE_DEFS}) +target_compile_definitions(mlc_llm_objs PRIVATE -DMLC_LLM_EXPORTS) +target_include_directories(mlc_llm_objs PRIVATE ${MLC_LLM_INCLUDES}) +target_include_directories(mlc_llm_objs PRIVATE 3rdparty/stb) +target_include_directories(mlc_llm_objs PRIVATE ${TOKENZIER_CPP_PATH}/include) +target_include_directories(mlc_llm_objs PRIVATE ${XGRAMMAR_PATH}/include) +# xgrammar still depends on picojson - use its bundled copy +target_include_directories(mlc_llm_objs + PRIVATE ${XGRAMMAR_PATH}/3rdparty/picojson) +target_link_libraries(mlc_llm_objs PRIVATE tvm_ffi_header) + +add_library(mlc_llm SHARED $) +add_library(mlc_llm_static STATIC $) +add_dependencies(mlc_llm_static tokenizers_cpp sentencepiece-static + tokenizers_c tvm_runtime tvm_runtime_extra) +set_target_properties(mlc_llm_static PROPERTIES OUTPUT_NAME mlc_llm) + +target_link_libraries(mlc_llm PUBLIC tvm_runtime) +target_link_libraries(mlc_llm PRIVATE tvm_runtime_extra) +target_link_libraries(mlc_llm PRIVATE tokenizers_cpp) + +find_library(FLASH_ATTN_LIBRARY flash_attn + HINTS ${TVM_SOURCE_DIR}/*/3rdparty/libflash_attn/src) + +if(FLASH_ATTN_LIBRARY STREQUAL "FLASH_ATTN_LIBRARY-NOTFOUND") + message( + WARNING + "Cannot find libflash_attn. The model must not have been built with --use-flash-attn-mqa option." + ) +else() + target_link_libraries(mlc_llm PUBLIC -Wl,--no-as-needed ${FLASH_ATTN_LIBRARY}) +endif() + +if(CMAKE_BUILD_TYPE STREQUAL "Debug") + target_compile_definitions(mlc_llm PRIVATE "TVM_LOG_DEBUG") + target_compile_definitions(mlc_llm_objs PRIVATE "TVM_LOG_DEBUG") + target_compile_definitions(mlc_llm_static PRIVATE "TVM_LOG_DEBUG") +endif() + +if(BUILD_CPP_TEST) + message(STATUS "Building cpp unittests") + add_subdirectory(3rdparty/googletest) + file(GLOB_RECURSE MLC_LLM_TEST_SRCS + ${PROJECT_SOURCE_DIR}/tests/cpp/*unittest.cc) + add_executable(mlc_llm_cpp_tests ${MLC_LLM_TEST_SRCS}) + target_include_directories(mlc_llm_cpp_tests PRIVATE ${MLC_LLM_INCLUDES}) + target_include_directories(mlc_llm_cpp_tests + PRIVATE ${PROJECT_SOURCE_DIR}/cpp) + target_include_directories( + mlc_llm_cpp_tests PRIVATE ${gtest_SOURCE_DIR}/include ${gtest_SOURCE_DIR}) + target_link_libraries(mlc_llm_cpp_tests PUBLIC mlc_llm gtest gtest_main) +endif(BUILD_CPP_TEST) + +if(CMAKE_SYSTEM_NAME STREQUAL "Android") + target_link_libraries(mlc_llm PRIVATE log) + target_link_libraries(tokenizers_cpp PRIVATE log) +endif() + +add_library(mlc_llm_module SHARED $) +target_link_libraries(mlc_llm_module PUBLIC tvm_runtime) +target_link_libraries(mlc_llm_module PRIVATE tvm_runtime_extra) +target_link_libraries(mlc_llm_module PRIVATE tokenizers_cpp) + +set_property( + TARGET mlc_llm_module + APPEND + PROPERTY LINK_OPTIONS "${MLC_VISIBILITY_FLAG}") +set_property( + TARGET mlc_llm + APPEND + PROPERTY LINK_OPTIONS "${MLC_VISIBILITY_FLAG}") + +find_program(CARGO_EXECUTABLE cargo) + +if(NOT CARGO_EXECUTABLE) + message(FATAL_ERROR "Cargo is not found! Please install cargo.") +endif() + +# when this option is on, we install all static lib deps into lib +if(MLC_LLM_INSTALL_STATIC_LIB) + install(TARGETS mlc_llm_static tokenizers_cpp sentencepiece-static tvm_runtime + tvm_runtime_extra + LIBRARY DESTINATION lib${LIB_SUFFIX}) + # tokenizers need special handling as it builds from rust + if(MSVC) + install(FILES ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/libtokenizers_c.lib + DESTINATION lib${LIB_SUFFIX}) + else() + install(FILES ${CMAKE_CURRENT_BINARY_DIR}/tokenizers/libtokenizers_c.a + DESTINATION lib${LIB_SUFFIX}) + endif() +else() + install( + TARGETS tvm_runtime + tvm_runtime_extra + mlc_llm + mlc_llm_module + mlc_llm_static + tokenizers_cpp + sentencepiece-static + RUNTIME_DEPENDENCY_SET + tokenizers_c + 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() diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md new file mode 100644 index 0000000..3f70fac --- /dev/null +++ b/CONTRIBUTORS.md @@ -0,0 +1,6 @@ +MLC LLM Contributors +==================== + + +## List of Contributors +- [Full List of Contributors](https://github.com/mlc-ai/mlc-llm/graphs/contributors) diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..19a51c2 --- /dev/null +++ b/NOTICE @@ -0,0 +1,3 @@ +MLC LLM + +Copyright (c) 2023-2025 by MLC LLM Contributors diff --git a/README.md b/README.md new file mode 100644 index 0000000..d5ba283 --- /dev/null +++ b/README.md @@ -0,0 +1,133 @@ +
+ +# MLC LLM + +[![Installation](https://img.shields.io/badge/docs-latest-green)](https://llm.mlc.ai/docs/) +[![License](https://img.shields.io/badge/license-apache_2-blue)](https://github.com/mlc-ai/mlc-llm/blob/main/LICENSE) +[![Join Discoard](https://img.shields.io/badge/Join-Discord-7289DA?logo=discord&logoColor=white)](https://discord.gg/9Xpy2HGBuD) +[![Related Repository: WebLLM](https://img.shields.io/badge/Related_Repo-WebLLM-fafbfc?logo=github)](https://github.com/mlc-ai/web-llm/) + +**Universal LLM Deployment Engine with ML Compilation** + +[Get Started](https://llm.mlc.ai/docs/get_started/quick_start) | [Documentation](https://llm.mlc.ai/docs) | [Blog](https://blog.mlc.ai/) + +
+ +## About + +MLC LLM is a machine learning compiler and high-performance deployment engine for large language models. The mission of this project is to enable everyone to develop, optimize, and deploy AI models natively on everyone's platforms.  + +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
AMD GPUNVIDIA GPUApple GPUIntel GPU
Linux / Win✅ Vulkan, ROCm✅ Vulkan, CUDAN/A✅ Vulkan
macOS✅ Metal (dGPU)N/A✅ Metal✅ Metal (iGPU)
Web Browser✅ WebGPU and WASM
iOS / iPadOS✅ Metal on Apple A-series GPU
Android✅ OpenCL on Adreno GPU✅ OpenCL on Mali GPU
+
+ +MLC LLM compiles and runs code on MLCEngine -- a unified high-performance LLM inference engine across the above platforms. MLCEngine provides OpenAI-compatible API available through REST server, python, javascript, iOS, Android, all backed by the same engine and compiler that we keep improving with the community. + +## Get Started + +Please visit our [documentation](https://llm.mlc.ai/docs/) to get started with MLC LLM. +- [Installation](https://llm.mlc.ai/docs/install/mlc_llm) +- [Quick start](https://llm.mlc.ai/docs/get_started/quick_start) +- [Introduction](https://llm.mlc.ai/docs/get_started/introduction) + +## Citation + +Please consider citing our project if you find it useful: + +```bibtex +@software{mlc-llm, + author = {{MLC team}}, + title = {{MLC-LLM}}, + url = {https://github.com/mlc-ai/mlc-llm}, + year = {2023-2025} +} +``` + +The underlying techniques of MLC LLM include: + +
+ References (Click to expand) + + ```bibtex + @inproceedings{tensorir, + author = {Feng, Siyuan and Hou, Bohan and Jin, Hongyi and Lin, Wuwei and Shao, Junru and Lai, Ruihang and Ye, Zihao and Zheng, Lianmin and Yu, Cody Hao and Yu, Yong and Chen, Tianqi}, + title = {TensorIR: An Abstraction for Automatic Tensorized Program Optimization}, + year = {2023}, + isbn = {9781450399166}, + publisher = {Association for Computing Machinery}, + address = {New York, NY, USA}, + url = {https://doi.org/10.1145/3575693.3576933}, + doi = {10.1145/3575693.3576933}, + booktitle = {Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2}, + pages = {804–817}, + numpages = {14}, + keywords = {Tensor Computation, Machine Learning Compiler, Deep Neural Network}, + location = {Vancouver, BC, Canada}, + series = {ASPLOS 2023} + } + + @inproceedings{metaschedule, + author = {Shao, Junru and Zhou, Xiyou and Feng, Siyuan and Hou, Bohan and Lai, Ruihang and Jin, Hongyi and Lin, Wuwei and Masuda, Masahiro and Yu, Cody Hao and Chen, Tianqi}, + booktitle = {Advances in Neural Information Processing Systems}, + editor = {S. Koyejo and S. Mohamed and A. Agarwal and D. Belgrave and K. Cho and A. Oh}, + pages = {35783--35796}, + publisher = {Curran Associates, Inc.}, + title = {Tensor Program Optimization with Probabilistic Programs}, + url = {https://proceedings.neurips.cc/paper_files/paper/2022/file/e894eafae43e68b4c8dfdacf742bcbf3-Paper-Conference.pdf}, + volume = {35}, + year = {2022} + } + + @inproceedings{tvm, + author = {Tianqi Chen and Thierry Moreau and Ziheng Jiang and Lianmin Zheng and Eddie Yan and Haichen Shen and Meghan Cowan and Leyuan Wang and Yuwei Hu and Luis Ceze and Carlos Guestrin and Arvind Krishnamurthy}, + title = {{TVM}: An Automated {End-to-End} Optimizing Compiler for Deep Learning}, + booktitle = {13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)}, + year = {2018}, + isbn = {978-1-939133-08-3}, + address = {Carlsbad, CA}, + pages = {578--594}, + url = {https://www.usenix.org/conference/osdi18/presentation/chen}, + publisher = {USENIX Association}, + month = oct, + } + ``` +
diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..0972cd4 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`mlc-ai/mlc-llm` +- 原始仓库：https://github.com/mlc-ai/mlc-llm +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/android/.gitignore b/android/.gitignore new file mode 100644 index 0000000..002b05d --- /dev/null +++ b/android/.gitignore @@ -0,0 +1,19 @@ +app/src/main/jni/*.h +app/src/main/jni/*.cc +app/src/main/obj + +*.iml +.gradle +/local.properties +/.idea/caches +/.idea/libraries +/.idea/modules.xml +/.idea/workspace.xml +/.idea/navEditor.xml +/.idea/assetWizardSettings.xml +.DS_Store +/build +/captures +.externalNativeBuild +.cxx +local.properties diff --git a/android/MLCChat/README.md b/android/MLCChat/README.md new file mode 100644 index 0000000..445d09a --- /dev/null +++ b/android/MLCChat/README.md @@ -0,0 +1,6 @@ +# MLC-LLM Android + +Checkout [Documentation page](https://llm.mlc.ai/docs/deploy/android.html) for more information. + +- run `mlc_llm package` +- open this `MLCChat/` folder as a project in Android Studio diff --git a/android/MLCChat/app/.gitignore b/android/MLCChat/app/.gitignore new file mode 100644 index 0000000..fe0ce51 --- /dev/null +++ b/android/MLCChat/app/.gitignore @@ -0,0 +1,2 @@ +/build +/src/main/libs diff --git a/android/MLCChat/app/build.gradle b/android/MLCChat/app/build.gradle new file mode 100644 index 0000000..e6f8320 --- /dev/null +++ b/android/MLCChat/app/build.gradle @@ -0,0 +1,74 @@ +plugins { + id 'com.android.application' + id 'org.jetbrains.kotlin.android' +} + +android { + namespace 'ai.mlc.mlcchat' + compileSdk 35 + + defaultConfig { + applicationId "ai.mlc.mlcchat" + minSdk 26 + targetSdk 33 + versionCode 1 + versionName "1.0" + + testInstrumentationRunner "androidx.test.runner.AndroidJUnitRunner" + vectorDrawables { + useSupportLibrary true + } + } + + buildTypes { + release { + minifyEnabled false + proguardFiles getDefaultProguardFile('proguard-android-optimize.txt'), 'proguard-rules.pro' + } + } + compileOptions { + sourceCompatibility JavaVersion.VERSION_1_8 + targetCompatibility JavaVersion.VERSION_1_8 + } + kotlinOptions { + jvmTarget = '1.8' + } + buildFeatures { + compose true + } + composeOptions { + kotlinCompilerExtensionVersion '1.4.3' + } + packagingOptions { + resources { + excludes += '/META-INF/{AL2.0,LGPL2.1}' + } + } +} + +dependencies { + implementation project(":mlc4j") + implementation 'androidx.core:core-ktx:1.10.1' + implementation 'androidx.lifecycle:lifecycle-runtime-ktx:2.6.1' + implementation 'com.github.jeziellago:compose-markdown:0.5.2' + implementation 'androidx.activity:activity-compose:1.7.1' + implementation platform('androidx.compose:compose-bom:2022.10.00') + implementation 'androidx.lifecycle:lifecycle-viewmodel-compose:2.6.1' + implementation 'androidx.compose.ui:ui' + implementation 'androidx.compose.ui:ui-graphics' + implementation 'androidx.compose.ui:ui-tooling-preview' + implementation 'androidx.compose.material3:material3:1.1.0' + implementation 'androidx.compose.material:material-icons-extended' + implementation 'androidx.appcompat:appcompat:1.6.1' + implementation 'androidx.navigation:navigation-compose:2.5.3' + implementation 'com.google.code.gson:gson:2.10.1' + implementation fileTree(dir: 'src/main/libs', include: ['*.aar', '*.jar'], exclude: []) + testImplementation 'junit:junit:4.13.2' + androidTestImplementation 'androidx.test.ext:junit:1.1.5' + androidTestImplementation 'androidx.test.espresso:espresso-core:3.5.1' + androidTestImplementation platform('androidx.compose:compose-bom:2022.10.00') + androidTestImplementation 'androidx.compose.ui:ui-test-junit4' + debugImplementation 'androidx.compose.ui:ui-tooling' + debugImplementation 'androidx.compose.ui:ui-test-manifest' + +} diff --git a/android/MLCChat/app/proguard-rules.pro b/android/MLCChat/app/proguard-rules.pro new file mode 100644 index 0000000..f1b4245 --- /dev/null +++ b/android/MLCChat/app/proguard-rules.pro @@ -0,0 +1,21 @@ +# Add project specific ProGuard rules here. +# You can control the set of applied configuration files using the +# proguardFiles setting in build.gradle. +# +# For more details, see +# http://developer.android.com/guide/developing/tools/proguard.html + +# If your project uses WebView with JS, uncomment the following +# and specify the fully qualified class name to the JavaScript interface +# class: +#-keepclassmembers class fqcn.of.javascript.interface.for.webview { +# public *; +#} + +# Uncomment this to preserve the line number information for +# debugging stack traces. +#-keepattributes SourceFile,LineNumberTable + +# If you keep the line number information, uncomment this to +# hide the original source file name. +#-renamesourcefileattribute SourceFile diff --git a/android/MLCChat/app/src/main/AndroidManifest.xml b/android/MLCChat/app/src/main/AndroidManifest.xml new file mode 100644 index 0000000..9a14f38 --- /dev/null +++ b/android/MLCChat/app/src/main/AndroidManifest.xml @@ -0,0 +1,42 @@ + + + + + + + + + + + + + + + + + + + + diff --git a/android/MLCChat/app/src/main/ic_launcher-playstore.png b/android/MLCChat/app/src/main/ic_launcher-playstore.png new file mode 100644 index 0000000..3c16fd6 Binary files /dev/null and b/android/MLCChat/app/src/main/ic_launcher-playstore.png differ diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/AppViewModel.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/AppViewModel.kt new file mode 100644 index 0000000..1e34911 --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/AppViewModel.kt @@ -0,0 +1,852 @@ +package ai.mlc.mlcchat + +import ai.mlc.mlcllm.MLCEngine +import ai.mlc.mlcllm.OpenAIProtocol +import android.app.Application +import android.content.ClipData +import android.content.ClipboardManager +import android.content.Context +import android.os.Environment +import android.widget.Toast +import androidx.compose.runtime.mutableStateOf +import androidx.compose.runtime.toMutableStateList +import androidx.lifecycle.AndroidViewModel +import androidx.lifecycle.viewModelScope +import com.google.gson.Gson +import com.google.gson.annotations.SerializedName +import kotlinx.coroutines.launch +import java.io.File +import java.io.FileOutputStream +import java.net.URL +import java.nio.channels.Channels +import java.util.UUID +import java.util.concurrent.Executors +import kotlin.concurrent.thread +import ai.mlc.mlcllm.OpenAIProtocol.ChatCompletionMessage +import ai.mlc.mlcllm.OpenAIProtocol.ChatCompletionMessageContent +import android.app.Activity +import kotlinx.coroutines.* +import android.graphics.Bitmap +import android.graphics.BitmapFactory +import android.net.Uri +import java.io.ByteArrayOutputStream +import android.util.Base64 +import android.util.Log + +class AppViewModel(application: Application) : AndroidViewModel(application) { + val modelList = emptyList().toMutableStateList() + val chatState = ChatState() + val modelSampleList = emptyList().toMutableStateList() + private var showAlert = mutableStateOf(false) + private var alertMessage = mutableStateOf("") + private var appConfig = AppConfig( + emptyList().toMutableList(), + emptyList().toMutableList() + ) + private val application = getApplication() + private val appDirFile = application.getExternalFilesDir("") + private val gson = Gson() + private val modelIdSet = emptySet().toMutableSet() + + companion object { + const val AppConfigFilename = "mlc-app-config.json" + const val ModelConfigFilename = "mlc-chat-config.json" + const val ParamsConfigFilename = "tensor-cache.json" + const val ModelUrlSuffix = "resolve/main/" + } + + init { + loadAppConfig() + } + + fun isShowingAlert(): Boolean { + return showAlert.value + } + + fun errorMessage(): String { + return alertMessage.value + } + + fun dismissAlert() { + require(showAlert.value) + showAlert.value = false + } + + fun copyError() { + require(showAlert.value) + val clipboard = + application.getSystemService(Context.CLIPBOARD_SERVICE) as ClipboardManager + clipboard.setPrimaryClip(ClipData.newPlainText("MLCChat", errorMessage())) + } + + private fun issueAlert(error: String) { + showAlert.value = true + alertMessage.value = error + } + + fun requestDeleteModel(modelId: String) { + deleteModel(modelId) + issueAlert("Model: $modelId has been deleted") + } + + + private fun loadAppConfig() { + val appConfigFile = File(appDirFile, AppConfigFilename) + val jsonString: String = if (!appConfigFile.exists()) { + application.assets.open(AppConfigFilename).bufferedReader().use { it.readText() } + } else { + appConfigFile.readText() + } + appConfig = gson.fromJson(jsonString, AppConfig::class.java) + appConfig.modelLibs = emptyList().toMutableList() + modelList.clear() + modelIdSet.clear() + modelSampleList.clear() + for (modelRecord in appConfig.modelList) { + appConfig.modelLibs.add(modelRecord.modelLib) + val modelDirFile = File(appDirFile, modelRecord.modelId) + val modelConfigFile = File(modelDirFile, ModelConfigFilename) + if (modelConfigFile.exists()) { + val modelConfigString = modelConfigFile.readText() + val modelConfig = gson.fromJson(modelConfigString, ModelConfig::class.java) + modelConfig.modelId = modelRecord.modelId + modelConfig.modelLib = modelRecord.modelLib + modelConfig.estimatedVramBytes = modelRecord.estimatedVramBytes + addModelConfig(modelConfig, modelRecord.modelUrl, true) + } else { + downloadModelConfig( + if (modelRecord.modelUrl.endsWith("/")) modelRecord.modelUrl else "${modelRecord.modelUrl}/", + modelRecord, + true + ) + } + } + } + + private fun updateAppConfig(action: () -> Unit) { + action() + val jsonString = gson.toJson(appConfig) + val appConfigFile = File(appDirFile, AppConfigFilename) + appConfigFile.writeText(jsonString) + } + + private fun addModelConfig(modelConfig: ModelConfig, modelUrl: String, isBuiltin: Boolean) { + require(!modelIdSet.contains(modelConfig.modelId)) + modelIdSet.add(modelConfig.modelId) + modelList.add( + ModelState( + modelConfig, + modelUrl + if (modelUrl.endsWith("/")) "" else "/", + File(appDirFile, modelConfig.modelId) + ) + ) + if (!isBuiltin) { + updateAppConfig { + appConfig.modelList.add( + ModelRecord( + modelUrl, + modelConfig.modelId, + modelConfig.estimatedVramBytes, + modelConfig.modelLib + ) + ) + } + } + } + + private fun deleteModel(modelId: String) { + val modelDirFile = File(appDirFile, modelId) + modelDirFile.deleteRecursively() + require(!modelDirFile.exists()) + modelIdSet.remove(modelId) + modelList.removeIf { modelState -> modelState.modelConfig.modelId == modelId } + updateAppConfig { + appConfig.modelList.removeIf { modelRecord -> modelRecord.modelId == modelId } + } + } + + private fun isModelConfigAllowed(modelConfig: ModelConfig): Boolean { + if (appConfig.modelLibs.contains(modelConfig.modelLib)) return true + viewModelScope.launch { + issueAlert("Model lib ${modelConfig.modelLib} is not supported.") + } + return false + } + + + private fun downloadModelConfig( + modelUrl: String, + modelRecord: ModelRecord, + isBuiltin: Boolean + ) { + thread(start = true) { + try { + val url = URL("${modelUrl}${ModelUrlSuffix}${ModelConfigFilename}") + val tempId = UUID.randomUUID().toString() + val tempFile = File( + application.getExternalFilesDir(Environment.DIRECTORY_DOWNLOADS), + tempId + ) + url.openStream().use { + Channels.newChannel(it).use { src -> + FileOutputStream(tempFile).use { fileOutputStream -> + fileOutputStream.channel.transferFrom(src, 0, Long.MAX_VALUE) + } + } + } + require(tempFile.exists()) + viewModelScope.launch { + try { + val modelConfigString = tempFile.readText() + val modelConfig = gson.fromJson(modelConfigString, ModelConfig::class.java) + modelConfig.modelId = modelRecord.modelId + modelConfig.modelLib = modelRecord.modelLib + modelConfig.estimatedVramBytes = modelRecord.estimatedVramBytes + if (modelIdSet.contains(modelConfig.modelId)) { + tempFile.delete() + issueAlert("${modelConfig.modelId} has been used, please consider another local ID") + return@launch + } + if (!isModelConfigAllowed(modelConfig)) { + tempFile.delete() + return@launch + } + val modelDirFile = File(appDirFile, modelConfig.modelId) + val modelConfigFile = File(modelDirFile, ModelConfigFilename) + tempFile.copyTo(modelConfigFile, overwrite = true) + tempFile.delete() + require(modelConfigFile.exists()) + addModelConfig(modelConfig, modelUrl, isBuiltin) + } catch (e: Exception) { + viewModelScope.launch { + issueAlert("Add model failed: ${e.localizedMessage}") + } + } + } + } catch (e: Exception) { + viewModelScope.launch { + issueAlert("Download model config failed: ${e.localizedMessage}") + } + } + + } + } + + inner class ModelState( + val modelConfig: ModelConfig, + private val modelUrl: String, + private val modelDirFile: File + ) { + var modelInitState = mutableStateOf(ModelInitState.Initializing) + private var paramsConfig = ParamsConfig(emptyList()) + val progress = mutableStateOf(0) + val total = mutableStateOf(1) + val id: UUID = UUID.randomUUID() + private val remainingTasks = emptySet().toMutableSet() + private val downloadingTasks = emptySet().toMutableSet() + private val maxDownloadTasks = 3 + private val gson = Gson() + + + init { + switchToInitializing() + } + + private fun switchToInitializing() { + val paramsConfigFile = File(modelDirFile, ParamsConfigFilename) + if (paramsConfigFile.exists()) { + loadParamsConfig() + switchToIndexing() + } else { + downloadParamsConfig() + } + } + + private fun loadParamsConfig() { + val paramsConfigFile = File(modelDirFile, ParamsConfigFilename) + require(paramsConfigFile.exists()) + val jsonString = paramsConfigFile.readText() + paramsConfig = gson.fromJson(jsonString, ParamsConfig::class.java) + } + + private fun downloadParamsConfig() { + thread(start = true) { + val url = URL("${modelUrl}${ModelUrlSuffix}${ParamsConfigFilename}") + val tempId = UUID.randomUUID().toString() + val tempFile = File(modelDirFile, tempId) + url.openStream().use { + Channels.newChannel(it).use { src -> + FileOutputStream(tempFile).use { fileOutputStream -> + fileOutputStream.channel.transferFrom(src, 0, Long.MAX_VALUE) + } + } + } + require(tempFile.exists()) + val paramsConfigFile = File(modelDirFile, ParamsConfigFilename) + tempFile.renameTo(paramsConfigFile) + require(paramsConfigFile.exists()) + viewModelScope.launch { + loadParamsConfig() + switchToIndexing() + } + } + } + + fun handleStart() { + switchToDownloading() + } + + fun handlePause() { + switchToPausing() + } + + fun handleClear() { + require( + modelInitState.value == ModelInitState.Downloading || + modelInitState.value == ModelInitState.Paused || + modelInitState.value == ModelInitState.Finished + ) + switchToClearing() + } + + private fun switchToClearing() { + if (modelInitState.value == ModelInitState.Paused) { + modelInitState.value = ModelInitState.Clearing + clear() + } else if (modelInitState.value == ModelInitState.Finished) { + modelInitState.value = ModelInitState.Clearing + if (chatState.modelName.value == modelConfig.modelId) { + chatState.requestTerminateChat { clear() } + } else { + clear() + } + } else { + modelInitState.value = ModelInitState.Clearing + } + } + + fun handleDelete() { + require( + modelInitState.value == ModelInitState.Downloading || + modelInitState.value == ModelInitState.Paused || + modelInitState.value == ModelInitState.Finished + ) + switchToDeleting() + } + + private fun switchToDeleting() { + if (modelInitState.value == ModelInitState.Paused) { + modelInitState.value = ModelInitState.Deleting + delete() + } else if (modelInitState.value == ModelInitState.Finished) { + modelInitState.value = ModelInitState.Deleting + if (chatState.modelName.value == modelConfig.modelId) { + chatState.requestTerminateChat { delete() } + } else { + delete() + } + } else { + modelInitState.value = ModelInitState.Deleting + } + } + + private fun switchToIndexing() { + modelInitState.value = ModelInitState.Indexing + progress.value = 0 + total.value = modelConfig.tokenizerFiles.size + paramsConfig.paramsRecords.size + for (tokenizerFilename in modelConfig.tokenizerFiles) { + val file = File(modelDirFile, tokenizerFilename) + if (file.exists()) { + ++progress.value + } else { + remainingTasks.add( + DownloadTask( + URL("${modelUrl}${ModelUrlSuffix}${tokenizerFilename}"), + file + ) + ) + } + } + for (paramsRecord in paramsConfig.paramsRecords) { + val file = File(modelDirFile, paramsRecord.dataPath) + if (file.exists()) { + ++progress.value + } else { + remainingTasks.add( + DownloadTask( + URL("${modelUrl}${ModelUrlSuffix}${paramsRecord.dataPath}"), + file + ) + ) + } + } + if (progress.value < total.value) { + switchToPaused() + } else { + switchToFinished() + } + } + + private fun switchToDownloading() { + modelInitState.value = ModelInitState.Downloading + for (downloadTask in remainingTasks) { + if (downloadingTasks.size < maxDownloadTasks) { + handleNewDownload(downloadTask) + } else { + return + } + } + } + + private fun handleNewDownload(downloadTask: DownloadTask) { + require(modelInitState.value == ModelInitState.Downloading) + require(!downloadingTasks.contains(downloadTask)) + downloadingTasks.add(downloadTask) + thread(start = true) { + val tempId = UUID.randomUUID().toString() + val tempFile = File(modelDirFile, tempId) + downloadTask.url.openStream().use { + Channels.newChannel(it).use { src -> + FileOutputStream(tempFile).use { fileOutputStream -> + fileOutputStream.channel.transferFrom(src, 0, Long.MAX_VALUE) + } + } + } + require(tempFile.exists()) + tempFile.renameTo(downloadTask.file) + require(downloadTask.file.exists()) + viewModelScope.launch { + handleFinishDownload(downloadTask) + } + } + } + + private fun handleNextDownload() { + require(modelInitState.value == ModelInitState.Downloading) + for (downloadTask in remainingTasks) { + if (!downloadingTasks.contains(downloadTask)) { + handleNewDownload(downloadTask) + break + } + } + } + + private fun handleFinishDownload(downloadTask: DownloadTask) { + remainingTasks.remove(downloadTask) + downloadingTasks.remove(downloadTask) + ++progress.value + require( + modelInitState.value == ModelInitState.Downloading || + modelInitState.value == ModelInitState.Pausing || + modelInitState.value == ModelInitState.Clearing || + modelInitState.value == ModelInitState.Deleting + ) + if (modelInitState.value == ModelInitState.Downloading) { + if (remainingTasks.isEmpty()) { + if (downloadingTasks.isEmpty()) { + switchToFinished() + } + } else { + handleNextDownload() + } + } else if (modelInitState.value == ModelInitState.Pausing) { + if (downloadingTasks.isEmpty()) { + switchToPaused() + } + } else if (modelInitState.value == ModelInitState.Clearing) { + if (downloadingTasks.isEmpty()) { + clear() + } + } else if (modelInitState.value == ModelInitState.Deleting) { + if (downloadingTasks.isEmpty()) { + delete() + } + } + } + + private fun clear() { + val files = modelDirFile.listFiles { dir, name -> + !(dir == modelDirFile && name == ModelConfigFilename) + } + require(files != null) + for (file in files) { + file.deleteRecursively() + require(!file.exists()) + } + val modelConfigFile = File(modelDirFile, ModelConfigFilename) + require(modelConfigFile.exists()) + switchToIndexing() + } + + private fun delete() { + modelDirFile.deleteRecursively() + require(!modelDirFile.exists()) + requestDeleteModel(modelConfig.modelId) + } + + private fun switchToPausing() { + modelInitState.value = ModelInitState.Pausing + } + + private fun switchToPaused() { + modelInitState.value = ModelInitState.Paused + } + + + private fun switchToFinished() { + modelInitState.value = ModelInitState.Finished + } + + fun startChat() { + chatState.requestReloadChat( + modelConfig, + modelDirFile.absolutePath, + ) + } + + } + + inner class ChatState { + val messages = emptyList().toMutableStateList() + val report = mutableStateOf("") + val modelName = mutableStateOf("") + private var modelChatState = mutableStateOf(ModelChatState.Ready) + @Synchronized get + @Synchronized set + private val engine = MLCEngine() + private var historyMessages = mutableListOf() + private var modelLib = "" + private var modelPath = "" + private val executorService = Executors.newSingleThreadExecutor() + private val viewModelScope = CoroutineScope(Dispatchers.Main + Job()) + private var imageUri: Uri? = null + private fun mainResetChat() { + imageUri = null + executorService.submit { + callBackend { engine.reset() } + historyMessages = mutableListOf() + viewModelScope.launch { + clearHistory() + switchToReady() + } + } + } + + private fun clearHistory() { + messages.clear() + report.value = "" + historyMessages.clear() + } + + + private fun switchToResetting() { + modelChatState.value = ModelChatState.Resetting + } + + private fun switchToGenerating() { + modelChatState.value = ModelChatState.Generating + } + + private fun switchToReloading() { + modelChatState.value = ModelChatState.Reloading + } + + private fun switchToReady() { + modelChatState.value = ModelChatState.Ready + } + + private fun switchToFailed() { + modelChatState.value = ModelChatState.Falied + } + + private fun callBackend(callback: () -> Unit): Boolean { + try { + callback() + } catch (e: Exception) { + viewModelScope.launch { + val stackTrace = e.stackTraceToString() + val errorMessage = e.localizedMessage + appendMessage( + MessageRole.Assistant, + "MLCChat failed\n\nStack trace:\n$stackTrace\n\nError message:\n$errorMessage" + ) + switchToFailed() + } + return false + } + return true + } + + fun requestResetChat() { + require(interruptable()) + interruptChat( + prologue = { + switchToResetting() + }, + epilogue = { + mainResetChat() + } + ) + } + + private fun interruptChat(prologue: () -> Unit, epilogue: () -> Unit) { + // prologue runs before interruption + // epilogue runs after interruption + require(interruptable()) + if (modelChatState.value == ModelChatState.Ready) { + prologue() + epilogue() + } else if (modelChatState.value == ModelChatState.Generating) { + prologue() + executorService.submit { + viewModelScope.launch { epilogue() } + } + } else { + require(false) + } + } + + fun requestTerminateChat(callback: () -> Unit) { + require(interruptable()) + interruptChat( + prologue = { + switchToTerminating() + }, + epilogue = { + mainTerminateChat(callback) + } + ) + } + + private fun mainTerminateChat(callback: () -> Unit) { + executorService.submit { + callBackend { engine.unload() } + viewModelScope.launch { + clearHistory() + switchToReady() + callback() + } + } + } + + private fun switchToTerminating() { + modelChatState.value = ModelChatState.Terminating + } + + + fun requestReloadChat(modelConfig: ModelConfig, modelPath: String) { + + if (this.modelName.value == modelConfig.modelId && this.modelLib == modelConfig.modelLib && this.modelPath == modelPath) { + return + } + require(interruptable()) + interruptChat( + prologue = { + switchToReloading() + }, + epilogue = { + mainReloadChat(modelConfig, modelPath) + } + ) + } + + private fun mainReloadChat(modelConfig: ModelConfig, modelPath: String) { + clearHistory() + this.modelName.value = modelConfig.modelId + this.modelLib = modelConfig.modelLib + this.modelPath = modelPath + executorService.submit { + viewModelScope.launch { + Toast.makeText(application, "Initialize...", Toast.LENGTH_SHORT).show() + } + if (!callBackend { + engine.unload() + engine.reload(modelPath, modelConfig.modelLib) + }) return@submit + viewModelScope.launch { + Toast.makeText(application, "Ready to chat", Toast.LENGTH_SHORT).show() + switchToReady() + } + } + } + + fun requestImageBitmap(uri: Uri?) { + require(chatable()) + switchToGenerating() + executorService.submit { + imageUri = uri + viewModelScope.launch { + report.value = "Image process is done, ask any question." + if (modelChatState.value == ModelChatState.Generating) switchToReady() + } + } + } + + fun bitmapToURL(bm: Bitmap): String { + val targetSize = 336 + val scaledBitmap = Bitmap.createScaledBitmap(bm, targetSize, targetSize, true) + + val outputStream = ByteArrayOutputStream() + scaledBitmap.compress(Bitmap.CompressFormat.JPEG, 100, outputStream) + scaledBitmap.recycle() + + val imageBytes = outputStream.toByteArray() + val imageBase64 = Base64.encodeToString(imageBytes, Base64.NO_WRAP) + return "data:image/jpg;base64,$imageBase64" + } + + fun requestGenerate(prompt: String, activity: Activity) { + require(chatable()) + switchToGenerating() + appendMessage(MessageRole.User, prompt) + appendMessage(MessageRole.Assistant, "") + var content = ChatCompletionMessageContent(text=prompt) + if (imageUri != null) { + val uri = imageUri + val bitmap = uri?.let { + activity.contentResolver.openInputStream(it)?.use { input -> + BitmapFactory.decodeStream(input) + } + } + val imageBase64URL = bitmapToURL(bitmap!!) + Log.v("requestGenerate", "image base64 url: $imageBase64URL") + val parts = listOf( + mapOf("type" to "text", "text" to prompt), + mapOf("type" to "image_url", "image_url" to imageBase64URL) + ) + content = ChatCompletionMessageContent(parts=parts) + imageUri = null + } + + executorService.submit { + historyMessages.add(ChatCompletionMessage( + role = OpenAIProtocol.ChatCompletionRole.user, + content = content + )) + + viewModelScope.launch { + val responses = engine.chat.completions.create( + messages = historyMessages, + stream_options = OpenAIProtocol.StreamOptions(include_usage = true) + ) + + var finishReasonLength = false + var streamingText = "" + + for (res in responses) { + if (!callBackend { + for (choice in res.choices) { + choice.delta.content?.let { content -> + streamingText += content.asText() + } + choice.finish_reason?.let { finishReason -> + if (finishReason == "length") { + finishReasonLength = true + } + } + } + updateMessage(MessageRole.Assistant, streamingText) + res.usage?.let { finalUsage -> + report.value = finalUsage.extra?.asTextLabel() ?: "" + } + if (finishReasonLength) { + streamingText += " [output truncated due to context length limit...]" + updateMessage(MessageRole.Assistant, streamingText) + } + }); + } + if (streamingText.isNotEmpty()) { + historyMessages.add(ChatCompletionMessage( + role = OpenAIProtocol.ChatCompletionRole.assistant, + content = streamingText + )) + streamingText = "" + } else { + if (historyMessages.isNotEmpty()) { + historyMessages.removeAt(historyMessages.size - 1) + } + } + + if (modelChatState.value == ModelChatState.Generating) switchToReady() + } + } + } + + private fun appendMessage(role: MessageRole, text: String) { + messages.add(MessageData(role, text)) + } + + + private fun updateMessage(role: MessageRole, text: String) { + messages[messages.size - 1] = MessageData(role, text) + } + + fun chatable(): Boolean { + return modelChatState.value == ModelChatState.Ready + } + + fun interruptable(): Boolean { + return modelChatState.value == ModelChatState.Ready + || modelChatState.value == ModelChatState.Generating + || modelChatState.value == ModelChatState.Falied + } + } +} + +enum class ModelInitState { + Initializing, + Indexing, + Paused, + Downloading, + Pausing, + Clearing, + Deleting, + Finished +} + +enum class ModelChatState { + Generating, + Resetting, + Reloading, + Terminating, + Ready, + Falied +} + +enum class MessageRole { + Assistant, + User +} + +data class DownloadTask(val url: URL, val file: File) + +data class MessageData(val role: MessageRole, val text: String, val id: UUID = UUID.randomUUID(), var imageUri: Uri? = null) + +data class AppConfig( + @SerializedName("model_libs") var modelLibs: MutableList, + @SerializedName("model_list") val modelList: MutableList, +) + +data class ModelRecord( + @SerializedName("model_url") val modelUrl: String, + @SerializedName("model_id") val modelId: String, + @SerializedName("estimated_vram_bytes") val estimatedVramBytes: Long?, + @SerializedName("model_lib") val modelLib: String +) + +data class ModelConfig( + @SerializedName("model_lib") var modelLib: String, + @SerializedName("model_id") var modelId: String, + @SerializedName("estimated_vram_bytes") var estimatedVramBytes: Long?, + @SerializedName("tokenizer_files") val tokenizerFiles: List, + @SerializedName("context_window_size") val contextWindowSize: Int, + @SerializedName("prefill_chunk_size") val prefillChunkSize: Int, +) + +data class ParamsRecord( + @SerializedName("dataPath") val dataPath: String +) + +data class ParamsConfig( + @SerializedName("records") val paramsRecords: List +) diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ChatView.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ChatView.kt new file mode 100644 index 0000000..b07e521 --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ChatView.kt @@ -0,0 +1,343 @@ +package ai.mlc.mlcchat + +import android.app.Activity +import android.graphics.Bitmap +import android.graphics.BitmapFactory +import androidx.compose.foundation.Image +import androidx.compose.foundation.background +import androidx.compose.foundation.gestures.detectTapGestures +import androidx.compose.foundation.layout.Arrangement +import androidx.compose.foundation.layout.Column +import androidx.compose.foundation.layout.IntrinsicSize +import androidx.compose.foundation.layout.Row +import androidx.compose.foundation.layout.aspectRatio +import androidx.compose.foundation.layout.fillMaxSize +import androidx.compose.foundation.layout.fillMaxWidth +import androidx.compose.foundation.layout.height +import androidx.compose.foundation.layout.padding +import androidx.compose.foundation.layout.widthIn +import androidx.compose.foundation.layout.wrapContentHeight +import androidx.compose.foundation.layout.wrapContentWidth +import androidx.compose.foundation.lazy.LazyColumn +import androidx.compose.foundation.lazy.items +import androidx.compose.foundation.lazy.rememberLazyListState +import androidx.compose.foundation.shape.RoundedCornerShape +import androidx.compose.foundation.text.selection.SelectionContainer +import androidx.compose.material.icons.Icons +import androidx.compose.material.icons.filled.AddAPhoto +import androidx.compose.material.icons.filled.ArrowBack +import androidx.compose.material.icons.filled.Photo +import androidx.compose.material.icons.filled.Replay +import androidx.compose.material.icons.filled.Send +import androidx.compose.material3.Divider +import androidx.compose.material3.ExperimentalMaterial3Api +import androidx.compose.material3.Icon +import androidx.compose.material3.IconButton +import androidx.compose.material3.MaterialTheme +import androidx.compose.material3.OutlinedTextField +import androidx.compose.material3.Scaffold +import androidx.compose.material3.Switch +import androidx.compose.material3.Text +import androidx.compose.material3.TopAppBar +import androidx.compose.material3.TopAppBarDefaults +import androidx.compose.runtime.Composable +import androidx.compose.runtime.getValue +import androidx.compose.runtime.mutableStateOf +import androidx.compose.runtime.remember +import androidx.compose.runtime.rememberCoroutineScope +import androidx.compose.runtime.saveable.rememberSaveable +import androidx.compose.runtime.setValue +import androidx.compose.ui.Alignment +import androidx.compose.ui.Modifier +import androidx.compose.ui.graphics.asImageBitmap +import androidx.compose.ui.input.pointer.pointerInput +import androidx.compose.ui.platform.LocalFocusManager +import androidx.compose.ui.text.style.TextAlign +import androidx.compose.ui.tooling.preview.Preview +import androidx.compose.ui.unit.dp +import androidx.navigation.NavController +import dev.jeziellago.compose.markdowntext.MarkdownText +import kotlinx.coroutines.launch + +@ExperimentalMaterial3Api +@Composable +fun ChatView( + navController: NavController, chatState: AppViewModel.ChatState, activity: Activity +) { + val localFocusManager = LocalFocusManager.current + (activity as MainActivity).chatState = chatState + Scaffold(topBar = { + TopAppBar( + title = { + Text( + text = "MLCChat: " + chatState.modelName.value.split("-")[0], + color = MaterialTheme.colorScheme.onPrimary + ) + }, + colors = TopAppBarDefaults.topAppBarColors(containerColor = MaterialTheme.colorScheme.primary), + navigationIcon = { + IconButton( + onClick = { navController.popBackStack() }, + enabled = chatState.interruptable() + ) { + Icon( + imageVector = Icons.Filled.ArrowBack, + contentDescription = "back home page", + tint = MaterialTheme.colorScheme.onPrimary + ) + } + }, + actions = { + IconButton( + onClick = { + chatState.requestResetChat() + activity.hasImage = false }, + enabled = chatState.interruptable() + ) { + Icon( + imageVector = Icons.Filled.Replay, + contentDescription = "reset the chat", + tint = MaterialTheme.colorScheme.onPrimary + ) + } + }) + }, modifier = Modifier.pointerInput(Unit) { + detectTapGestures(onTap = { + localFocusManager.clearFocus() + }) + }) { paddingValues -> + Column( + modifier = Modifier + .fillMaxSize() + .padding(paddingValues) + .padding(horizontal = 10.dp) + ) { + val lazyColumnListState = rememberLazyListState() + val coroutineScope = rememberCoroutineScope() + Text( + text = chatState.report.value, + textAlign = TextAlign.Center, + modifier = Modifier + .fillMaxWidth() + .wrapContentHeight() + .padding(top = 5.dp) + ) + Divider(thickness = 1.dp, modifier = Modifier.padding(vertical = 5.dp)) + LazyColumn( + modifier = Modifier.weight(9f), + verticalArrangement = Arrangement.spacedBy(5.dp, alignment = Alignment.Bottom), + state = lazyColumnListState + ) { + coroutineScope.launch { + lazyColumnListState.animateScrollToItem(chatState.messages.size) + } + items( + items = chatState.messages, + key = { message -> message.id }, + ) { message -> + MessageView(messageData = message, activity) + } + item { + // place holder item for scrolling to the bottom + } + } + Divider(thickness = 1.dp, modifier = Modifier.padding(top = 5.dp)) + SendMessageView(chatState = chatState, activity) + } + } +} + +@Composable +fun MessageView(messageData: MessageData, activity: Activity?) { + // default render the Assistant text as MarkdownText + var useMarkdown by remember { mutableStateOf(true) } + var localActivity : MainActivity = activity as MainActivity + SelectionContainer { + if (messageData.role == MessageRole.Assistant) { + Column { + if (messageData.text.isNotEmpty()) { + Row( + verticalAlignment = Alignment.CenterVertically, + ) { + Text( + text = "Show as Markdown", + color = MaterialTheme.colorScheme.onSecondaryContainer, + modifier = Modifier + .wrapContentWidth() + .padding(end = 8.dp) + .widthIn(max = 300.dp) + ) + Switch( + checked = useMarkdown, + onCheckedChange = { useMarkdown = it } + ) + } + } + Row( + horizontalArrangement = Arrangement.Start, + modifier = Modifier.fillMaxWidth() + ) { + if (useMarkdown) { + MarkdownText( + isTextSelectable = true, + modifier = Modifier + .wrapContentWidth() + .background( + color = MaterialTheme.colorScheme.secondaryContainer, + shape = RoundedCornerShape(5.dp) + ) + .padding(5.dp) + .widthIn(max = 300.dp), + markdown = messageData.text, + ) + } else { + Text( + text = messageData.text, + textAlign = TextAlign.Left, + color = MaterialTheme.colorScheme.onSecondaryContainer, + modifier = Modifier + .wrapContentWidth() + .background( + color = MaterialTheme.colorScheme.secondaryContainer, + shape = RoundedCornerShape(5.dp) + ) + .padding(5.dp) + .widthIn(max = 300.dp) + ) + } + } + } + } else { + Row( + horizontalArrangement = Arrangement.End, + modifier = Modifier.fillMaxWidth() + ) { + if (messageData.imageUri != null) { + val uri = messageData.imageUri + val bitmap = uri?.let { + activity.contentResolver.openInputStream(it)?.use { input -> + BitmapFactory.decodeStream(input) + } + } + val displayBitmap = bitmap?.let { Bitmap.createScaledBitmap(it, 224, 224, true) } + if (displayBitmap != null) { + Image( + displayBitmap.asImageBitmap(), + "", + modifier = Modifier + .wrapContentWidth() + .background( + color = MaterialTheme.colorScheme.secondaryContainer, + shape = RoundedCornerShape(5.dp) + ) + .padding(5.dp) + .widthIn(max = 300.dp) + ) + } + if (!localActivity.hasImage) { + localActivity.chatState.requestImageBitmap(messageData.imageUri) + } + localActivity.hasImage = true + } else { + Text( + text = messageData.text, + textAlign = TextAlign.Right, + color = MaterialTheme.colorScheme.onPrimaryContainer, + modifier = Modifier + .wrapContentWidth() + .background( + color = MaterialTheme.colorScheme.primaryContainer, + shape = RoundedCornerShape(5.dp) + ) + .padding(5.dp) + .widthIn(max = 300.dp) + ) + } + + } + } + } +} + +@ExperimentalMaterial3Api +@Composable +fun SendMessageView(chatState: AppViewModel.ChatState, activity: Activity) { + val localFocusManager = LocalFocusManager.current + val localActivity : MainActivity = activity as MainActivity + Row( + horizontalArrangement = Arrangement.spacedBy(5.dp), + verticalAlignment = Alignment.CenterVertically, + modifier = Modifier + .height(IntrinsicSize.Max) + .fillMaxWidth() + .padding(bottom = 5.dp) + ) { + var text by rememberSaveable { mutableStateOf("") } + OutlinedTextField( + value = text, + onValueChange = { text = it }, + label = { Text(text = "Input") }, + modifier = Modifier + .weight(9f), + ) + IconButton( + onClick = { + activity.takePhoto() + }, + modifier = Modifier + .aspectRatio(1f) + .weight(1f), + enabled = (chatState.chatable() && !localActivity.hasImage) + ) { + Icon( + imageVector = Icons.Filled.AddAPhoto, + contentDescription = "use camera", + ) + } + IconButton( + onClick = { + activity.pickImageFromGallery() + }, + modifier = Modifier + .aspectRatio(1f) + .weight(1f), + enabled = (chatState.chatable() && !localActivity.hasImage) + ) { + Icon( + imageVector = Icons.Filled.Photo, + contentDescription = "select image", + ) + } + IconButton( + onClick = { + localFocusManager.clearFocus() + chatState.requestGenerate(text, activity) + text = "" + }, + modifier = Modifier + .aspectRatio(1f) + .weight(1f), + enabled = (text != "" && chatState.chatable()) + ) { + Icon( + imageVector = Icons.Filled.Send, + contentDescription = "send message", + ) + } + } +} + +@Preview +@Composable +fun MessageViewPreviewWithMarkdown() { + MessageView( + messageData = MessageData( + role = MessageRole.Assistant, text = """ +# Sample Header +* Markdown +* [Link](https://example.com) +Google +""" + ), null + ) +} diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/MainActivity.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/MainActivity.kt new file mode 100644 index 0000000..b50bd7b --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/MainActivity.kt @@ -0,0 +1,157 @@ +package ai.mlc.mlcchat + +import android.Manifest +import android.content.ContentValues +import android.content.pm.PackageManager +import android.net.Uri +import android.os.Build +import android.os.Bundle +import android.provider.MediaStore +import android.util.Log +import androidx.activity.ComponentActivity +import androidx.activity.compose.setContent +import androidx.activity.result.contract.ActivityResultContracts +import androidx.annotation.RequiresApi +import androidx.compose.foundation.layout.fillMaxSize +import androidx.compose.material3.ExperimentalMaterial3Api +import androidx.compose.material3.Surface +import androidx.compose.ui.Modifier +import androidx.core.content.ContextCompat +import ai.mlc.mlcchat.ui.theme.MLCChatTheme +import java.text.SimpleDateFormat +import java.util.Date +import java.util.Locale +import java.util.UUID + +class MainActivity : ComponentActivity() { + var hasImage = false + + private val pickImageLauncher = registerForActivityResult( + ActivityResultContracts.GetContent() + ) { uri: Uri? -> + uri?.let { + Log.v("pickImageLauncher", "Selected image uri: $it") + chatState.messages.add( + MessageData( + role = MessageRole.User, + text = "", + id = UUID.randomUUID(), + imageUri = it + ) + ) + } + } + + private var cameraImageUri: Uri? = null + private val takePictureLauncher = registerForActivityResult( + ActivityResultContracts.TakePicture() + ) { success: Boolean -> + if (success && cameraImageUri != null) { + Log.v("takePictureLauncher", "Camera image uri: $cameraImageUri") + chatState.messages.add( + MessageData( + role = MessageRole.User, + text = "", + id = UUID.randomUUID(), + imageUri = cameraImageUri + ) + ) + } + } + + private val requestPermissionLauncher = + registerForActivityResult(ActivityResultContracts.RequestMultiplePermissions()) { permissions -> + permissions.entries.forEach { + Log.d("Permissions", "${it.key} = ${it.value}") + } + } + + lateinit var chatState: AppViewModel.ChatState + + @RequiresApi(Build.VERSION_CODES.TIRAMISU) + @ExperimentalMaterial3Api + override fun onCreate(savedInstanceState: Bundle?) { + super.onCreate(savedInstanceState) + + chatState = AppViewModel(this.application).ChatState() + requestNeededPermissions() + + setContent { + Surface( + modifier = Modifier.fillMaxSize() + ) { + MLCChatTheme { + NavView(this) + } + } + } + } + + private fun requestNeededPermissions() { + val permissionsToRequest = mutableListOf() + + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU) { + if (ContextCompat.checkSelfPermission( + this, + Manifest.permission.READ_MEDIA_IMAGES + ) != PackageManager.PERMISSION_GRANTED + ) { + permissionsToRequest.add(Manifest.permission.READ_MEDIA_IMAGES) + } + if (ContextCompat.checkSelfPermission( + this, + Manifest.permission.CAMERA + ) != PackageManager.PERMISSION_GRANTED + ) { + permissionsToRequest.add(Manifest.permission.CAMERA) + } + } else { + if (ContextCompat.checkSelfPermission( + this, + Manifest.permission.READ_EXTERNAL_STORAGE + ) != PackageManager.PERMISSION_GRANTED + ) { + permissionsToRequest.add(Manifest.permission.READ_EXTERNAL_STORAGE) + } + if (ContextCompat.checkSelfPermission( + this, + Manifest.permission.WRITE_EXTERNAL_STORAGE + ) != PackageManager.PERMISSION_GRANTED + ) { + permissionsToRequest.add(Manifest.permission.WRITE_EXTERNAL_STORAGE) + } + if (ContextCompat.checkSelfPermission( + this, + Manifest.permission.CAMERA + ) != PackageManager.PERMISSION_GRANTED + ) { + permissionsToRequest.add(Manifest.permission.CAMERA) + } + } + + if (permissionsToRequest.isNotEmpty()) { + requestPermissionLauncher.launch(permissionsToRequest.toTypedArray()) + } + } + + fun pickImageFromGallery() { + pickImageLauncher.launch("image/*") + } + + fun takePhoto() { + val contentValues = ContentValues().apply { + val timeFormatter = SimpleDateFormat("yyyyMMdd_HHmmss", Locale.getDefault()) + val fileName = "IMG_${timeFormatter.format(Date())}.jpg" + put(MediaStore.Images.Media.DISPLAY_NAME, fileName) + put(MediaStore.Images.Media.MIME_TYPE, "image/jpeg") + put(MediaStore.Images.Media.DATE_ADDED, System.currentTimeMillis() / 1000) + } + + cameraImageUri = contentResolver.insert( + MediaStore.Images.Media.EXTERNAL_CONTENT_URI, + contentValues + ) + + takePictureLauncher.launch(cameraImageUri) + } +} diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/NavView.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/NavView.kt new file mode 100644 index 0000000..008187c --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/NavView.kt @@ -0,0 +1,19 @@ +package ai.mlc.mlcchat + +import android.app.Activity +import androidx.compose.material3.ExperimentalMaterial3Api +import androidx.compose.runtime.Composable +import androidx.lifecycle.viewmodel.compose.viewModel +import androidx.navigation.compose.NavHost +import androidx.navigation.compose.composable +import androidx.navigation.compose.rememberNavController + +@ExperimentalMaterial3Api +@Composable +fun NavView(activity: Activity, appViewModel: AppViewModel = viewModel()) { + val navController = rememberNavController() + NavHost(navController = navController, startDestination = "home") { + composable("home") { StartView(navController, appViewModel) } + composable("chat") { ChatView(navController, appViewModel.chatState, activity) } + } +} diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/StartView.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/StartView.kt new file mode 100644 index 0000000..b650af0 --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/StartView.kt @@ -0,0 +1,250 @@ +package ai.mlc.mlcchat + +import androidx.compose.foundation.gestures.detectTapGestures +import androidx.compose.foundation.layout.Arrangement +import androidx.compose.foundation.layout.Column +import androidx.compose.foundation.layout.Row +import androidx.compose.foundation.layout.aspectRatio +import androidx.compose.foundation.layout.fillMaxSize +import androidx.compose.foundation.layout.fillMaxWidth +import androidx.compose.foundation.layout.height +import androidx.compose.foundation.layout.padding +import androidx.compose.foundation.layout.width +import androidx.compose.foundation.layout.wrapContentHeight +import androidx.compose.foundation.lazy.LazyColumn +import androidx.compose.foundation.lazy.items +import androidx.compose.foundation.text.selection.SelectionContainer +import androidx.compose.material.icons.Icons +import androidx.compose.material.icons.outlined.Chat +import androidx.compose.material.icons.outlined.Delete +import androidx.compose.material.icons.outlined.Download +import androidx.compose.material.icons.outlined.Pause +import androidx.compose.material.icons.outlined.Schedule +import androidx.compose.material3.AlertDialog +import androidx.compose.material3.Divider +import androidx.compose.material3.ExperimentalMaterial3Api +import androidx.compose.material3.Icon +import androidx.compose.material3.IconButton +import androidx.compose.material3.LinearProgressIndicator +import androidx.compose.material3.MaterialTheme +import androidx.compose.material3.OutlinedTextField +import androidx.compose.material3.Scaffold +import androidx.compose.material3.Text +import androidx.compose.material3.TextButton +import androidx.compose.material3.TopAppBar +import androidx.compose.material3.TopAppBarDefaults +import androidx.compose.runtime.Composable +import androidx.compose.runtime.getValue +import androidx.compose.runtime.mutableStateOf +import androidx.compose.runtime.saveable.rememberSaveable +import androidx.compose.runtime.setValue +import androidx.compose.ui.Alignment +import androidx.compose.ui.Modifier +import androidx.compose.ui.input.pointer.pointerInput +import androidx.compose.ui.platform.LocalFocusManager +import androidx.compose.ui.text.style.TextAlign +import androidx.compose.ui.unit.dp +import androidx.navigation.NavController + + +@ExperimentalMaterial3Api +@Composable +fun StartView( + navController: NavController, + appViewModel: AppViewModel +) { + val localFocusManager = LocalFocusManager.current + Scaffold( + topBar = { + TopAppBar( + title = { Text(text = "MLCChat", color = MaterialTheme.colorScheme.onPrimary) }, + colors = TopAppBarDefaults.topAppBarColors(containerColor = MaterialTheme.colorScheme.primary) + ) + }, + modifier = Modifier.pointerInput(Unit) { + detectTapGestures(onTap = { + localFocusManager.clearFocus() + }) + } + ) + { paddingValues -> + Column( + modifier = Modifier + .fillMaxSize() + .padding(paddingValues) + .padding(horizontal = 10.dp) + ) { + Text(text = "Model List", modifier = Modifier.padding(top = 10.dp)) + LazyColumn() { + items(items = appViewModel.modelList, + key = { modelState -> modelState.id } + ) { modelState -> + ModelView( + navController = navController, + modelState = modelState, + appViewModel = appViewModel + ) + } + } + } + if (appViewModel.isShowingAlert()) { + AlertDialog( + onDismissRequest = { appViewModel.dismissAlert() }, + onConfirmation = { appViewModel.copyError() }, + error = appViewModel.errorMessage() + ) + } + } +} + +@ExperimentalMaterial3Api +@Composable +fun AlertDialog( + onDismissRequest: () -> Unit, + onConfirmation: () -> Unit, + error: String, +) { + AlertDialog( + title = { Text(text = "Error") }, + text = { Text(text = error) }, + onDismissRequest = { onDismissRequest() }, + confirmButton = { + TextButton(onClick = { onConfirmation() }) { Text("Copy") } + }, + dismissButton = { + TextButton(onClick = { onDismissRequest() }) { Text("Dismiss") } + } + ) +} + +@Composable +fun ModelView( + navController: NavController, + modelState: AppViewModel.ModelState, + appViewModel: AppViewModel +) { + var isDeletingModel by rememberSaveable { mutableStateOf(false) } + Column( + verticalArrangement = Arrangement.SpaceBetween, + modifier = Modifier + .wrapContentHeight() + ) { + Row( + horizontalArrangement = Arrangement.spacedBy(5.dp), + verticalAlignment = Alignment.CenterVertically, + modifier = Modifier + .fillMaxWidth() + .wrapContentHeight() + ) { + Text( + text = modelState.modelConfig.modelId, + textAlign = TextAlign.Left, + modifier = Modifier + .wrapContentHeight() + .weight(8f) + ) + Divider( + modifier = Modifier + .height(20.dp) + .width(1.dp) + ) + if (modelState.modelInitState.value == ModelInitState.Paused) { + IconButton( + onClick = { modelState.handleStart() }, modifier = Modifier + .aspectRatio(1f) + .weight(1f) + ) { + Icon( + imageVector = Icons.Outlined.Download, + contentDescription = "start downloading", + ) + } + + } else if (modelState.modelInitState.value == ModelInitState.Downloading) { + IconButton( + onClick = { modelState.handlePause() }, modifier = Modifier + .aspectRatio(1f) + .weight(1f) + ) { + Icon( + imageVector = Icons.Outlined.Pause, + contentDescription = "pause downloading", + ) + } + } else if (modelState.modelInitState.value == ModelInitState.Finished) { + IconButton( + onClick = { + modelState.startChat() + navController.navigate("chat") + }, + enabled = appViewModel.chatState.interruptable(), + modifier = Modifier + .aspectRatio(1f) + .weight(1f) + ) { + Icon( + imageVector = Icons.Outlined.Chat, + contentDescription = "start chatting", + ) + } + } else { + IconButton( + enabled = false, onClick = {}, modifier = Modifier + .aspectRatio(1f) + .weight(1f) + ) { + Icon( + imageVector = Icons.Outlined.Schedule, + contentDescription = "pending", + ) + } + } + if (modelState.modelInitState.value == ModelInitState.Downloading || + modelState.modelInitState.value == ModelInitState.Paused || + modelState.modelInitState.value == ModelInitState.Finished + ) { + IconButton( + onClick = { isDeletingModel = true }, + modifier = Modifier + .aspectRatio(1f) + .weight(1f) + ) { + Icon( + imageVector = Icons.Outlined.Delete, + contentDescription = "start downloading", + tint = MaterialTheme.colorScheme.error + ) + } + } + } + LinearProgressIndicator( + progress = modelState.progress.value.toFloat() / modelState.total.value, + modifier = Modifier.fillMaxWidth() + ) + if (isDeletingModel) { + Row( + horizontalArrangement = Arrangement.End, + verticalAlignment = Alignment.CenterVertically, + modifier = Modifier + .fillMaxWidth() + .wrapContentHeight() + ) { + TextButton(onClick = { isDeletingModel = false }) { + Text(text = "cancel") + } + TextButton(onClick = { + isDeletingModel = false + modelState.handleClear() + }) { + Text(text = "clear data", color = MaterialTheme.colorScheme.error) + } + TextButton(onClick = { + isDeletingModel = false + modelState.handleDelete() + }) { + Text(text = "delete model", color = MaterialTheme.colorScheme.error) + } + } + } + } +} diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Color.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Color.kt new file mode 100644 index 0000000..5bef131 --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Color.kt @@ -0,0 +1,44 @@ +package ai.mlc.mlcchat.ui.theme + +import androidx.compose.ui.graphics.Color + +val Blue10 = Color(0xFF000F5E) +val Blue20 = Color(0xFF001E92) +val Blue30 = Color(0xFF002ECC) +val Blue40 = Color(0xFF1546F6) +val Blue80 = Color(0xFFB8C3FF) +val Blue90 = Color(0xFFDDE1FF) + +val DarkBlue10 = Color(0xFF00036B) +val DarkBlue20 = Color(0xFF000BA6) +val DarkBlue30 = Color(0xFF1026D3) +val DarkBlue40 = Color(0xFF3648EA) +val DarkBlue80 = Color(0xFFBBC2FF) +val DarkBlue90 = Color(0xFFDEE0FF) + +val Yellow10 = Color(0xFF261900) +val Yellow20 = Color(0xFF402D00) +val Yellow30 = Color(0xFF5C4200) +val Yellow40 = Color(0xFF7A5900) +val Yellow80 = Color(0xFFFABD1B) +val Yellow90 = Color(0xFFFFDE9C) + +val Red10 = Color(0xFF410001) +val Red20 = Color(0xFF680003) +val Red30 = Color(0xFF930006) +val Red40 = Color(0xFFBA1B1B) +val Red80 = Color(0xFFFFB4A9) +val Red90 = Color(0xFFFFDAD4) + +val Grey10 = Color(0xFF191C1D) +val Grey20 = Color(0xFF2D3132) +val Grey80 = Color(0xFFC4C7C7) +val Grey90 = Color(0xFFE0E3E3) +val Grey95 = Color(0xFFEFF1F1) +val Grey99 = Color(0xFFFBFDFD) + +val BlueGrey30 = Color(0xFF45464F) +val BlueGrey50 = Color(0xFF767680) +val BlueGrey60 = Color(0xFF90909A) +val BlueGrey80 = Color(0xFFC6C5D0) +val BlueGrey90 = Color(0xFFE2E1EC) diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Theme.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Theme.kt new file mode 100644 index 0000000..dd2de4c --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Theme.kt @@ -0,0 +1,107 @@ +package ai.mlc.mlcchat.ui.theme + +import android.app.Activity +import android.os.Build +import androidx.compose.foundation.isSystemInDarkTheme +import androidx.compose.material3.MaterialTheme +import androidx.compose.material3.darkColorScheme +import androidx.compose.material3.dynamicDarkColorScheme +import androidx.compose.material3.dynamicLightColorScheme +import androidx.compose.material3.lightColorScheme +import androidx.compose.runtime.Composable +import androidx.compose.runtime.SideEffect +import androidx.compose.ui.graphics.Color +import androidx.compose.ui.graphics.toArgb +import androidx.compose.ui.platform.LocalContext +import androidx.compose.ui.platform.LocalView +import androidx.core.view.WindowCompat + +private val DarkColorScheme = darkColorScheme( + primary = Blue80, + onPrimary = Blue20, + primaryContainer = Blue30, + onPrimaryContainer = Blue90, + inversePrimary = Blue40, + secondary = DarkBlue80, + onSecondary = DarkBlue20, + secondaryContainer = DarkBlue30, + onSecondaryContainer = DarkBlue90, + tertiary = Yellow80, + onTertiary = Yellow20, + tertiaryContainer = Yellow30, + onTertiaryContainer = Yellow90, + error = Red80, + onError = Red20, + errorContainer = Red30, + onErrorContainer = Red90, + background = Grey10, + onBackground = Grey90, + surface = Grey10, + onSurface = Grey80, + inverseSurface = Grey90, + inverseOnSurface = Grey20, + surfaceVariant = BlueGrey30, + onSurfaceVariant = BlueGrey80, + outline = BlueGrey60 +) + +private val LightColorScheme = lightColorScheme( + primary = Blue40, + onPrimary = Color.White, + primaryContainer = Blue90, + onPrimaryContainer = Blue10, + inversePrimary = Blue80, + secondary = DarkBlue40, + onSecondary = Color.White, + secondaryContainer = DarkBlue90, + onSecondaryContainer = DarkBlue10, + tertiary = Yellow40, + onTertiary = Color.White, + tertiaryContainer = Yellow90, + onTertiaryContainer = Yellow10, + error = Red40, + onError = Color.White, + errorContainer = Red90, + onErrorContainer = Red10, + background = Grey99, + onBackground = Grey10, + surface = Grey99, + onSurface = Grey10, + inverseSurface = Grey20, + inverseOnSurface = Grey95, + surfaceVariant = BlueGrey90, + onSurfaceVariant = BlueGrey30, + outline = BlueGrey50 +) + +@Composable +fun MLCChatTheme( + darkTheme: Boolean = isSystemInDarkTheme(), + // Dynamic color is available on Android 12+ + dynamicColor: Boolean = true, + content: @Composable () -> Unit +) { + val colorScheme = when { + dynamicColor && Build.VERSION.SDK_INT >= Build.VERSION_CODES.S -> { + val context = LocalContext.current + if (darkTheme) dynamicDarkColorScheme(context) else dynamicLightColorScheme(context) + } + + darkTheme -> DarkColorScheme + else -> LightColorScheme + } + val view = LocalView.current + if (!view.isInEditMode) { + SideEffect { + val window = (view.context as Activity).window + window.statusBarColor = colorScheme.primary.toArgb() + WindowCompat.getInsetsController(window, view).isAppearanceLightStatusBars = darkTheme + } + } + + MaterialTheme( + colorScheme = colorScheme, + typography = Typography, + content = content + ) +} diff --git a/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Type.kt b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Type.kt new file mode 100644 index 0000000..56c6cf1 --- /dev/null +++ b/android/MLCChat/app/src/main/java/ai/mlc/mlcchat/ui/theme/Type.kt @@ -0,0 +1,34 @@ +package ai.mlc.mlcchat.ui.theme + +import androidx.compose.material3.Typography +import androidx.compose.ui.text.TextStyle +import androidx.compose.ui.text.font.FontFamily +import androidx.compose.ui.text.font.FontWeight +import androidx.compose.ui.unit.sp + +// Set of Material typography styles to start with +val Typography = Typography( + bodyLarge = TextStyle( + fontFamily = FontFamily.Default, + fontWeight = FontWeight.Normal, + fontSize = 16.sp, + lineHeight = 24.sp, + letterSpacing = 0.5.sp + ) + /* Other default text styles to override + titleLarge = TextStyle( + fontFamily = FontFamily.Default, + fontWeight = FontWeight.Normal, + fontSize = 22.sp, + lineHeight = 28.sp, + letterSpacing = 0.sp + ), + labelSmall = TextStyle( + fontFamily = FontFamily.Default, + fontWeight = FontWeight.Medium, + fontSize = 11.sp, + lineHeight = 16.sp, + letterSpacing = 0.5.sp + ) + */ +) diff --git a/android/MLCChat/app/src/main/res/drawable/ic_android_black_24dp.xml b/android/MLCChat/app/src/main/res/drawable/ic_android_black_24dp.xml new file mode 100644 index 0000000..fe51230 --- /dev/null +++ b/android/MLCChat/app/src/main/res/drawable/ic_android_black_24dp.xml @@ -0,0 +1,5 @@ + + + diff --git a/android/MLCChat/app/src/main/res/drawable/mlc_logo_108.xml b/android/MLCChat/app/src/main/res/drawable/mlc_logo_108.xml new file mode 100644 index 0000000..d5307e0 --- /dev/null +++ b/android/MLCChat/app/src/main/res/drawable/mlc_logo_108.xml @@ -0,0 +1,11 @@ + + + diff --git a/android/MLCChat/app/src/main/res/values/colors.xml b/android/MLCChat/app/src/main/res/values/colors.xml new file mode 100644 index 0000000..ca1931b --- /dev/null +++ b/android/MLCChat/app/src/main/res/values/colors.xml @@ -0,0 +1,10 @@ + + + #FFBB86FC + #FF6200EE + #FF3700B3 + #FF03DAC5 + #FF018786 + #FF000000 + #FFFFFFFF + diff --git a/android/MLCChat/app/src/main/res/values/strings.xml b/android/MLCChat/app/src/main/res/values/strings.xml new file mode 100644 index 0000000..f85e3fa --- /dev/null +++ b/android/MLCChat/app/src/main/res/values/strings.xml @@ -0,0 +1,3 @@ + + MLCChat + diff --git a/android/MLCChat/app/src/main/res/values/themes.xml b/android/MLCChat/app/src/main/res/values/themes.xml new file mode 100644 index 0000000..045e465 --- /dev/null +++ b/android/MLCChat/app/src/main/res/values/themes.xml @@ -0,0 +1,6 @@ + + + + + + + + + + + + + + + + + + + diff --git a/docs/_static/img/project-structure.svg b/docs/_static/img/project-structure.svg new file mode 100644 index 0000000..e4ad7db --- /dev/null +++ b/docs/_static/img/project-structure.svg @@ -0,0 +1,1189 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/_static/img/project-workflow.svg b/docs/_static/img/project-workflow.svg new file mode 100644 index 0000000..eac1313 --- /dev/null +++ b/docs/_static/img/project-workflow.svg @@ -0,0 +1,1173 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/community/faq.rst b/docs/community/faq.rst new file mode 100644 index 0000000..0778f5a --- /dev/null +++ b/docs/community/faq.rst @@ -0,0 +1,16 @@ +.. _FAQ: + +Frequently Asked Questions +========================== + +This is a list of Frequently Asked Questions (FAQ) about the MLC-LLM. Feel free to suggest new entries! + +... How can I customize the temperature, and repetition penalty of models? + Please check our :ref:`configure-mlc-chat-json` tutorial. + +... What's the quantization algorithm MLC-LLM using? + Please check our :doc:`/compilation/configure_quantization` tutorial. + +... Why do I encounter an error ``free(): invalid pointer, Aborted (core dumped)`` at the end of model compilation? + This happens if you compiled TVM from source and didn't hide LLVM symbols in cmake configurations. + Please follow our instructions in :ref:`Building TVM from Source ` tutorial to compile TVM which hides LLVM symbols, or use our pre-built MLC-LLM :doc:`pip wheels <../install/mlc_llm>`. diff --git a/docs/community/guideline.rst b/docs/community/guideline.rst new file mode 100644 index 0000000..467ffe6 --- /dev/null +++ b/docs/community/guideline.rst @@ -0,0 +1,125 @@ +.. _community_guide: + +Community Guideline +=================== + +.. contents:: + :depth: 2 + :local: + +Welcome to the MLC-LLM community! Just like you, all of us are in awe of the immense power of large language models. +Our goal for MLC-LLM is to foster a project that is driven by an open-source community, working together to democratize +this technology and make it accessible across various devices. We are thrilled to have you as part of our +community and eagerly anticipate your valuable contributions. + + +.. _community_discussion: + +Participate in Community Discussions +------------------------------------ + +We encourage open discussions. If you encounter a bug or have a feature request, please file an issue in MLC-LLM's +GitHub `issue tracker `__. You are encouraged to tag the issue with labels +such as "bug," "feature request," or "iOS" so that the relevant developers can quickly notice your concern. + +Additionally, we have set up a `discord server `__ for online discussions. +While we encourage participation in the Discord server, we also recommend creating a GitHub issue even if the +topic has been discussed there. This ensures that the discussion is archived and searchable for future reference. + +Before submitting an issue, we kindly ask you to check our :doc:`/community/faq` to see if your question has already been answered. + +.. _contribute-to-mlc-llm: + +Contribute to MLC-LLM +--------------------- + +.. _fork-and-create-pull-requests: + +Fork and Create Pull Requests +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Ready to contribute to MLC-LLM? Awesome! We are excited to see you are ready to contribute your code. +The standard way to make changes to MLC-LLM code base is through creating a `pull-request `__, +and we will review your code and merge it to the code base when it is ready. + +The first step to becoming a developer is to `fork `__ the repository to your own +github account, you will notice a repository under ``https://github.com/username/mlc-llm`` where ``username`` is your github user name. + +You can clone your fork to your local machine and commit changes, or edit the contents of your fork (in the case you are just fixing typos) +on GitHub directly. Once your update is complete, you can click the ``contribute`` button and open a pull request to the main repository. + +.. _contribute-new-models: + +Contribute New Models to MLC-LLM +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +* If you have compiled a model using our :doc:`/compilation/compile_models` tutorial for an existing model architecture, please upload your models to the internet (e.g., Hugging Face) by following :ref:`distribute-compiled-models` tutorial. + +* If you add a new model variant to MLC-LLM by following our :doc:`/compilation/define_new_models` tutorial. + Please create a pull request to add your model architecture (currently model architectures are placed under + `relax_models `__ folder). + +.. _coding-styles: + +Coding Styles +^^^^^^^^^^^^^ + +For python codes, we generally follow the `PEP8 style guide `__. +The python comments follow `NumPy style `__ python docstrings. +To make things easy, you can use `black `__ to automatically format your python code. + +.. code:: bash + + pip install black + black your_python_file.py + +For C++ codes, we generally follow the `Google C++ style guide `__. +The C++ comments should be `Doxygen compatible `__. +Fo your convenience, you can use `clang-format `__ to automatically format your C++ code. + +.. code:: bash + + clang-format -i your_cpp_file.cpp + +.. _general-development-process: + +General Development Process +--------------------------- + +Everyone in the community is welcome to send patches, documents, and propose new directions to the project. +The key guideline here is to enable everyone in the community to get involved and participate in the decision and development. +We encourage public discussion in different channels, so that everyone in the community can participate +and get informed in developments. + +Code reviews are one of the key ways to ensure the quality of the code. High-quality code reviews prevent technical debt +for long-term and are crucial to the success of the project. A pull request needs to be reviewed before it gets merged. +A committer who has the expertise of the corresponding area would moderate the pull request and merge the code when +it is ready. The corresponding committer could request multiple reviewers who are familiar with the area of the code. +We encourage contributors to request code reviews themselves and help review each other's code -- remember everyone +is volunteering their time to the community, high-quality code review itself costs as much as the actual code +contribution, you could get your code quickly reviewed if you do others the same favor. + +The community should strive to reach a consensus on technical decisions through discussion. We expect committers to +moderate technical discussions in a diplomatic way, and provide suggestions with clear technical reasoning when necessary. + + +.. _roles-committers: + +Committers +^^^^^^^^^^ + +Committers are individuals who are granted with write access to the project. A committer is usually responsible for +a certain area or several areas of the code where they oversee the code review process. +The area of contribution can take all forms, including code contributions and code reviews, documents, education, and outreach. +The review of pull requests will be assigned to the committers who recently contribute to the area this PR belongs to. +Committers are essential for a high quality and healthy project. The community actively looks for new committers +from contributors. Each existing committer can nominate new committers to MLC projects. + +.. _roles-contributors: + +Contributors +^^^^^^^^^^^^ +We also welcome contributors if you are not ready to be a committer yet. Everyone who contributes to +the project (in the form of code, bugfix, documentation, tutorials, etc) is a contributor. +We maintain a `page `__ to acknowledge contributors, +please let us know if you contribute to the project and if your name is not included in the list. diff --git a/docs/compilation/compile_models.rst b/docs/compilation/compile_models.rst new file mode 100644 index 0000000..e8a4b06 --- /dev/null +++ b/docs/compilation/compile_models.rst @@ -0,0 +1,1092 @@ +.. _compile-model-libraries: + +Compile Model Libraries +======================= + +To run a model with MLC LLM in any platform, we need: + +1. **Model weights** converted to MLC format (e.g. `RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC `__.) +2. **Model library** that comprises the inference logic + +This page describes how to compile a model library with MLC LLM. Model compilation optimizes +the model inference for a given platform, allowing users bring their own new model +architecture, use different quantization modes, and customize the overall model +optimization flow. + + + +Notably, in many cases you do not need to explicit call compile. + +- If you are using the Python API, you can skip specifying ``model_lib`` and + the system will JIT compile the library. + +- If you are building iOS/android package, checkout :ref:`package-libraries-and-weights`, + which provides a simpler high-level command that leverages the compile behind the scheme. + + +This page is still helpful to understand the compilation flow behind the scheme, +or be used to explicit create model libraries. +We compile ``RedPajama-INCITE-Chat-3B-v1`` with ``q4f16_1`` as an example for all platforms. + +.. note:: + Before you proceed, make sure you followed :ref:`install-tvm`, a required + backend to compile models with MLC LLM. + + Please also follow the instructions in :ref:`deploy-cli` / :ref:`deploy-python-engine` to obtain + the CLI app / Python API that can be used to chat with the compiled model. + + +.. contents:: Table of Contents + :depth: 1 + :local: + +1. Verify Installation +---------------------- + +**Step 1. Verify mlc_llm** + +We use the python package ``mlc_llm`` to compile models. This can be installed by +following :ref:`install-mlc-packages`, either by building from source, or by +installing the prebuilt package. Verify ``mlc_llm`` installation in command line via: + +.. code:: bash + + $ mlc_llm --help + # You should see help information with this line + usage: MLC LLM Command Line Interface. [-h] {compile,convert_weight,gen_config} + +.. note:: + If it runs into error ``command not found: mlc_llm``, try ``python -m mlc_llm --help``. + +**Step 2. Verify TVM** + +To compile models, you also need to follow :ref:`install-tvm`. +Here we verify ``tvm`` quickly with command line (for full verification, see :ref:`tvm-validate`): + +.. code:: bash + + $ python -c "import tvm; print(tvm.__file__)" + /some-path/lib/python3.13/site-packages/tvm/__init__.py + +1. Clone from HF and convert_weight +----------------------------------- + +This replicates :ref:`convert-weights-via-MLC`, see that page for more details. + +You can be under the mlc-llm repo, or your own working directory. Note that all platforms +can share the same compiled/quantized weights. + +.. code:: shell + + # Create directory + mkdir -p dist/models && cd dist/models + # Clone HF weights + git lfs install + git clone https://huggingface.co/togethercomputer/RedPajama-INCITE-Chat-3B-v1 + cd ../.. + # Convert weight + mlc_llm convert_weight ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + +2. Generate mlc-chat-config and compile +--------------------------------------- + +A model library is specified by: + + - The model architecture (e.g. ``llama-2``, ``gpt-neox``) + - Quantization (e.g. ``q4f16_1``, ``q0f32``) + - Metadata (e.g. ``context_window_size``, ``sliding_window_size``, ``prefill-chunk-size``), which affects memory planning + - Platform (e.g. ``cuda``, ``webgpu``, ``iOS``) + +All these knobs are specified in ``mlc-chat-config.json`` generated by ``gen_config``. + +.. code:: shell + + # Create output directory for the model library compiled + mkdir dist/libs + +.. tabs:: + + .. group-tab:: Linux - CUDA + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device cuda -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-cuda.so + + + .. group-tab:: Metal + + For M-chip Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-metal.so + + Cross-Compiling for Intel Mac on M-chip Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device metal:x86-64 -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-metal_x86_64.dylib + + For Intel Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-metal_x86_64.dylib + + + .. group-tab:: Vulkan + + For Linux: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-vulkan.so + + For Windows: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-vulkan.dll + + .. group-tab:: iOS/iPadOS + + You need a Mac to compile models for it. + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ --quantization q4f16_1 \ + --conv-template redpajama_chat --context-window-size 768 \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device iphone -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-iphone.tar + + .. note:: + If it runs into error + + .. code:: text + + Compilation error: + xcrun: error: unable to find utility "metal", not a developer tool or in PATH + xcrun: error: unable to find utility "metallib", not a developer tool or in PATH + + , please check and make sure you have Command Line Tools for Xcode installed correctly. + You can use ``xcrun metal`` to validate: when it prints ``metal: error: no input files``, it means the Command Line Tools for Xcode is installed and can be found, and you can proceed with the model compiling. + + .. group-tab:: Android + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ --quantization q4f16_1 \ + --conv-template redpajama_chat --context-window-size 768 \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device android -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-android.tar + + .. group-tab:: WebGPU + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device webgpu -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-webgpu.wasm + + .. note:: + To compile for webgpu, you need to build from source when installing ``mlc_llm``. Besides, you also need to follow :ref:`install-web-build`. + Otherwise, it would run into error + + .. code:: text + + RuntimeError: Cannot find libraries: wasm_runtime.bc + + .. note:: + For webgpu, when compiling larger models like ``Llama-2-7B``, you may want to add ``--prefill-chunk-size 1024`` or lower ``--context-window-size`` to decrease memory usage. + Otherwise, you may run into issues like: + + .. code:: text + + TypeError: Failed to execute 'createBuffer' on 'GPUDevice': Failed to read the 'size' property from + 'GPUBufferDescriptor': Value is outside the 'unsigned long long' value range. + +.. note:: + + For the ``conv-template``, `conversation_template.py `__ + contains a full list of conversation templates that MLC provides. If the model you are adding + requires a new conversation template, you would need to add your own. + Follow `this PR `__ as an example. + However, adding your own template would require you :ref:`build mlc_llm from source ` + in order for it to be recognized by the runtime. + + For more details, please see :ref:`configure-mlc-chat-json`. + +3. Verify output and chat +------------------------- + +By executing the compile command above, we generate the model weights, model lib, and a chat config. +We can check the output with the commands below: + +.. tabs:: + + .. group-tab:: Linux - CUDA + + .. code:: shell + + ~/mlc-llm > ls dist/libs + RedPajama-INCITE-Chat-3B-v1-q4f16_1-cuda.so # ===> the model library + + ~/mlc-llm > ls dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + + We can now chat with the model using the command line interface (CLI) app or the Python API. + + .. code:: shell + + python + >>> from mlc_llm import MLCEngine + >>> engine = MLCEngine(model="./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + ... model_lib="./dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-cuda.so") + >>> engine.chat.completions.create( + ... messages=[{"role": "user", "content": "hello"}] + ... ) + ChatCompletionResponse( + choices=[ChatCompletionResponseChoice( + message=ChatCompletionMessage( + content="Hi! How can I assist you today?", role='assistant' + ) + )], + ... + ) + + .. group-tab:: Metal + + .. code:: shell + + ~/mlc-llm > ls dist/libs + RedPajama-INCITE-Chat-3B-v1-q4f16_1-metal.so # ===> the model library (will be -metal_x86_64.dylib for Intel Mac) + + ~/mlc-llm > ls dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + + We can now chat with the model using the command line interface (CLI) app or the Python API. + + .. code:: shell + + python + >>> from mlc_llm import MLCEngine + >>> engine = MLCEngine(model="./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + ... model_lib="./dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-metal.so") + >>> engine.chat.completions.create( + ... messages=[{"role": "user", "content": "hello"}] + ... ) + ChatCompletionResponse( + choices=[ChatCompletionResponseChoice( + message=ChatCompletionMessage( + content="Hi! How can I assist you today?", role='assistant' + ) + )], + ... + ) + + + .. group-tab:: Vulkan + + .. code:: shell + + ~/mlc-llm > ls dist/libs + RedPajama-INCITE-Chat-3B-v1-q4f16_1-vulkan.so # ===> the model library (will be .dll for Windows) + + ~/mlc-llm > ls dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + + We can now chat with the model using the command line interface (CLI) app or the Python API. + + .. code:: shell + + python + >>> from mlc_llm import MLCEngine + >>> engine = MLCEngine(model="./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + ... model_lib="./dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-vulkan.so") + >>> engine.chat.completions.create( + ... messages=[{"role": "user", "content": "hello"}] + ... ) + ChatCompletionResponse( + choices=[ChatCompletionResponseChoice( + message=ChatCompletionMessage( + content="Hi! How can I assist you today?", role='assistant' + ) + )], + ... + ) + + .. group-tab:: iOS/iPadOS + + .. code:: shell + + ~/mlc-llm > ls dist/libs + RedPajama-INCITE-Chat-3B-v1-q4f16_1-iphone.tar # ===> the model library + + ~/mlc-llm > ls dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + + The model lib ``dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-iphone.tar`` + will be packaged as a static library into the iOS app. Checkout :ref:`deploy-ios` for more details. + + .. group-tab:: Android + + .. code:: shell + + ~/mlc-llm > ls dist/libs + RedPajama-INCITE-Chat-3B-v1-q4f16_1-android.tar # ===> the model library + + ~/mlc-llm > ls dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + + The model lib ``dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-android.tar`` + will be packaged as a static library into the android app. Checkout :ref:`deploy-android` for more details. + + .. group-tab:: WebGPU + + .. code:: shell + + ~/mlc-llm > ls dist/libs + RedPajama-INCITE-Chat-3B-v1-q4f16_1-webgpu.wasm # ===> the model library + + ~/mlc-llm > ls dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + + To use this in WebGPU runtime, checkout :ref:`webllm-runtime`. + +Compile Commands for More Models +-------------------------------- + +This section lists compile commands for more models that you can try out. Note that this can be easily +generalized to any model variant, as long as mlc-llm supports the architecture. + +.. tabs:: + + .. tab:: Model: Llama-2-7B + + Please `request for access `_ to the Llama-2 weights from Meta first. + After granted access, first create directory ``dist/models`` and download the model to the directory. + For example, you can run the following code: + + .. code:: shell + + mkdir -p dist/models && cd dist/models + git lfs install + git clone https://huggingface.co/meta-llama/Llama-2-7b-chat-hf + cd ../.. + + Then convert the HF weights into MLC-compatible weights. Note that all platforms + can share the same compiled/quantized weights. + + .. code:: shell + + mlc_llm convert_weight ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC + + Afterwards, run the following command to generate mlc config and compile the model. + + .. code:: shell + + # Create output directory for the model library compiled + mkdir dist/libs + + .. tabs:: + + .. tab:: Target: CUDA + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --conv-template llama-2 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device cuda -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-cuda.so + + .. tab:: Metal + + For M-chip Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --conv-template llama-2 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-metal.so + + Cross-Compiling for Intel Mac on M-chip Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device metal:x86-64 -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-metal_x86_64.dylib + + For Intel Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --conv-template llama-2 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-metal_x86_64.dylib + + .. tab:: Vulkan + + For Linux: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --conv-template llama-2 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-vulkan.so + + For Windows: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --conv-template llama-2 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-vulkan.dll + + .. tab:: WebGPU + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --context-window-size 2048 --conv-template llama-2 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device webgpu -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-webgpu.wasm + + .. note:: + To compile for webgpu, you need to build from source when installing ``mlc_llm``. Besides, you also need to follow :ref:`install-web-build`. + Otherwise, it would run into error + + .. code:: text + + RuntimeError: Cannot find libraries: wasm_runtime.bc + + .. tab:: iPhone/iPad + + You need a Mac to compile models for it. + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --conv-template llama-2 --context-window-size 768 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device iphone -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-iphone.tar + + .. tab:: Android + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Llama-2-7b-chat-hf/ --quantization q4f16_1 \ + --conv-template llama-2 --context-window-size 768 -o dist/Llama-2-7b-chat-hf-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device android -o dist/libs/Llama-2-7b-chat-hf-q4f16_1-android.tar + + .. tab:: Mistral-7B-Instruct-v0.2 + + Note that Mistral uses sliding window attention (SWA). Thus, instead of specifying + ``context-window-size``, we specify ``sliding-window-size``. + + First create directory ``dist/models`` and download the model to the directory. + For example, you can run the following code: + + .. code:: shell + + mkdir -p dist/models && cd dist/models + git lfs install + git clone https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2 + cd ../.. + + Then convert the HF weights into MLC-compatible weights. Note that all platforms + can share the same compiled/quantized weights. + + .. code:: shell + + mlc_llm convert_weight ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC + + Afterwards, run the following command to generate mlc config and compile the model. + + .. code:: shell + + # Create output directory for the model library compiled + mkdir dist/libs + + .. tabs:: + + .. tab:: Target: CUDA + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --conv-template mistral_default -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device cuda -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-cuda.so + + .. tab:: Metal + + For M-chip Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --conv-template mistral_default -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-metal.so + + + For Intel Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --conv-template mistral_default -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-metal_x86_64.dylib + + .. tab:: Vulkan + + For Linux: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --conv-template mistral_default -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-vulkan.so + + For Windows: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --conv-template mistral_default -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-vulkan.dll + + .. tab:: WebGPU + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --prefill-chunk-size 1024 --conv-template mistral_default \ + -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device webgpu -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-webgpu.wasm + + .. note:: + To compile for webgpu, you need to build from source when installing ``mlc_llm``. Besides, you also need to follow :ref:`install-web-build`. + Otherwise, it would run into error + + .. code:: text + + RuntimeError: Cannot find libraries: wasm_runtime.bc + + .. note:: + For webgpu, when compiling larger models like ``Llama-2-7B``, you may want to add ``--prefill-chunk-size 1024`` or lower ``--context-window-size`` to decrease memory usage. + Otherwise, you may run into issues like: + + .. code:: text + + TypeError: Failed to execute 'createBuffer' on 'GPUDevice': Failed to read the 'size' property from + 'GPUBufferDescriptor': Value is outside the 'unsigned long long' value range. + + .. tab:: iPhone/iPad + + You need a Mac to compile models for it. + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --conv-template mistral_default --sliding-window-size 1024 --prefill-chunk-size 128 \ + -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device iphone -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-iphone.tar + + .. tab:: Android + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/Mistral-7B-Instruct-v0.2/ --quantization q4f16_1 \ + --conv-template mistral_default --sliding-window-size 1024 --prefill-chunk-size 128 -o dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/Mistral-7B-Instruct-v0.2-q4f16_1-MLC/mlc-chat-config.json \ + --device android -o dist/libs/Mistral-7B-Instruct-v0.2-q4f16_1-android.tar + + .. tab:: Other models + + First create directory ``dist/models`` and download the model to the directory. + For example, you can run the following code: + + .. code:: shell + + mkdir -p dist/models && cd dist/models + git lfs install + git clone https://huggingface.co/DISTRIBUTOR/HF_MODEL + cd ../.. + + Then convert the HF weights into MLC-compatible weights. Note that all platforms + can share the same compiled/quantized weights. + + .. code:: shell + + mlc_llm convert_weight ./dist/models/HF_MODEL/ --quantization q4f16_1 -o dist/OUTPUT-MLC + + Afterwards, run the following command to generate mlc config and compile the model. + + .. code:: shell + + # Create output directory for the model library compiled + mkdir dist/libs + + .. tabs:: + + .. tab:: Target: CUDA + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device cuda -o dist/libs/OUTPUT-cuda.so + + .. tab:: Metal + + For M-chip Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device metal -o dist/libs/OUTPUT-metal.so + + + For Intel Mac: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device metal -o dist/libs/OUTPUT-metal_x86_64.dylib + + .. tab:: Vulkan + + For Linux: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device vulkan -o dist/libs/OUTPUT-vulkan.so + + For Windows: + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device vulkan -o dist/libs/OUTPUT-vulkan.dll + + .. tab:: WebGPU + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device webgpu -o dist/libs/OUTPUT-webgpu.wasm + + .. note:: + To compile for webgpu, you need to build from source when installing ``mlc_llm``. Besides, you also need to follow :ref:`install-web-build`. + Otherwise, it would run into error + + .. code:: text + + RuntimeError: Cannot find libraries: wasm_runtime.bc + + .. note:: + For webgpu, when compiling larger models like ``Llama-2-7B``, you may want to add ``--prefill-chunk-size 1024`` or lower ``--context-window-size`` to decrease memory usage. + Otherwise, you may run into issues like: + + .. code:: text + + TypeError: Failed to execute 'createBuffer' on 'GPUDevice': Failed to read the 'size' property from + 'GPUBufferDescriptor': Value is outside the 'unsigned long long' value range. + + .. tab:: iPhone/iPad + + You need a Mac to compile models for it. + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE \ + --context-window-size 768 -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device iphone -o dist/libs/OUTPUT-iphone.tar + + .. tab:: Android + + .. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/HF_MODEL/ --quantization q4f16_1 --conv-template CONV_TEMPLATE \ + --context-window-size 768 -o dist/OUTPUT-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/OUTPUT-MLC/mlc-chat-config.json --device android -o dist/libs/OUTPUT-android.tar + +For each model and each backend, the above only provides the most recommended build command (which is the most optimized). +You can also try with different argument values (e.g., different quantization modes, context window size, etc.), +whose build results affect runtime memory requirement, and it is possible that they may not run as +fast and robustly as the provided one when running the model. + +.. note:: + Uing 3-bit quantization usually can be overly aggressive and only works for limited settings. + If you encounter issues where the compiled model does not perform as expected, + consider utilizing a higher number of bits for quantization (e.g., 4-bit quantization). + +If you are interested in distributing the model besides local execution, please checkout :ref:`distribute-compiled-models`. + + +.. _compile-command-specification: + +Compile Command Specification +----------------------------- + +As you have seen in the section above, the model compilation is split into three steps: convert weights, generate +``mlc-chat-config.json``, and compile the model. This section describes the list of options that can be used +during compilation. + +1. Convert Weight +^^^^^^^^^^^^^^^^^ + +Weight conversion command follows the pattern below: + +.. code:: text + + mlc_llm convert_weight \ + CONFIG \ + --quantization QUANTIZATION_MODE \ + [--model-type MODEL_TYPE] \ + [--device DEVICE] \ + [--source SOURCE] \ + [--source-format SOURCE_FORMAT] \ + --output OUTPUT + +Note that ``CONFIG`` is a positional argument. Arguments wrapped with ``[ ]`` are optional. + +--CONFIG It can be one of the following: + + 1. Path to a HuggingFace model directory that contains a ``config.json`` or + 2. Path to ``config.json`` in HuggingFace format, or + 3. The name of a pre-defined model architecture. + + A ``config.json`` file in HuggingFace format defines the model architecture, including the vocabulary + size, the number of layers, the hidden size, number of attention heads, etc. + Example: https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json. + + A HuggingFace directory often contains a ``config.json`` which defines the model architecture, + the non-quantized model weights in PyTorch or SafeTensor format, tokenizer configurations, + as well as an optional ``generation_config.json`` provides additional default configuration for + text generation. + Example: https://huggingface.co/codellama/CodeLlama-7b-hf/tree/main. + + For existing pre-defined model architecture, see ``MODEL_PRESETS`` + `here `_. + +--quantization QUANTIZATION_MODE The quantization mode we use to compile. + + See :ref:`quantization_mode` for more information. + Available options are: ``q0f16``, ``q0f32``, ``q3f16_1``, ``q4f16_1``, ``q4f32_1``, and + ``q4f16_awq``. + + We encourage you to use 4-bit quantization, as the text generated by 3-bit + quantized models may have bad quality depending on the model. + +--model-type MODEL_TYPE Model architecture such as "llama". If not set, it is inferred from ``config.json``. + +--device DEVICE The device used to do quantization such as "cuda" or "cuda:0". Will detect from + local available GPUs if not specified. + +--source SOURCE The path to original model weight, infer from ``config`` if missing. + +--source-format SOURCE_FORMAT The format of source model weight, infer from ``config`` if missing. + +--output OUTPUT The output directory to save the quantized model weight. + Will create ``params_shard_*.bin`` and ```tensor-cache.json``` in this directory. + +2. Generate MLC Chat Config +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +In order to compile a model, we first need to generate the ``mlc-chat-config.json``. This file contains specifications +like ``context-window-size`` and ``sliding-window-size``, among others that can alter the model compiled. We also process +tokenizers in this step. + +Config generation command follows the pattern below: + +.. code:: text + + mlc_llm gen_config \ + CONFIG \ + --quantization QUANTIZATION_MODE \ + [--model-type MODEL_TYPE] \ + --conv-template CONV_TEMPLATE \ + [--context-window-size CONTEXT_WINDOW_SIZE] \ + [--sliding-window-size SLIDING_WINDOW_SIZE] \ + [--prefill-chunk-size PREFILL_CHUNK_SIZE] \ + [--tensor-parallel-shard TENSOR_PARALLEL_SHARDS] \ + --output OUTPUT + +Note that ``CONFIG`` is a positional argument. Arguments wrapped with ``[ ]`` are optional. + +--CONFIG It can be one of the following: + + 1. Path to a HuggingFace model directory that contains a ``config.json`` or + 2. Path to ``config.json`` in HuggingFace format, or + 3. The name of a pre-defined model architecture. + + A ``config.json`` file in HuggingFace format defines the model architecture, including the vocabulary + size, the number of layers, the hidden size, number of attention heads, etc. + Example: https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json. + + A HuggingFace directory often contains a ``config.json`` which defines the model architecture, + the non-quantized model weights in PyTorch or SafeTensor format, tokenizer configurations, + as well as an optional ``generation_config.json`` provides additional default configuration for + text generation. + Example: https://huggingface.co/codellama/CodeLlama-7b-hf/tree/main. + + For existing pre-defined model architecture, see ``MODEL_PRESETS`` + `here `_. + +--quantization QUANTIZATION_MODE The quantization mode we use to compile. + + See :ref:`quantization_mode` for more information. + Available options are: ``q0f16``, ``q0f32``, ``q3f16_1``, ``q4f16_1``, ``q4f32_1``, and + ``q4f16_awq``. + + We encourage you to use 4-bit quantization, as the text generated by 3-bit + quantized models may have bad quality depending on the model. + +--model-type MODEL_TYPE Model architecture such as "llama". If not set, it is inferred from ``config.json``. + +--conv-template CONV_TEMPLATE Conversation template. It depends on how the model is tuned. Use "LM" for vanilla base model + For existing pre-defined templates, see ``CONV_TEMPLATES`` + `here `_. + +--context-window-size CONTEXT_WINDOW_SIZE Option to provide the maximum sequence length supported by the model. + This is usually explicitly shown as context length or context window in the model card. + If this option is not set explicitly, by default, + it will be determined by ``context_window_size`` or ``max_position_embeddings`` in ``config.json``, + and the latter is usually inaccurate for some models. + +--sliding-window-size SLIDING_WINDOW (Experimental) The sliding window size in sliding window attention (SWA). + This optional field overrides the ``sliding_window`` in ``config.json`` for + those models that use SWA. Currently only useful when compiling mistral-based models. + This flag subjects to future refactoring. + +--prefill-chunk-size PREFILL_CHUNK_SIZE (Experimental) The chunk size during prefilling. By default, + the chunk size is the same as ``context_window_size`` or ``sliding_window_size``. + This flag subjects to future refactoring. + +--tensor-parallel-shard TENSOR_PARALLEL_SHARDS Number of shards to split the model into in tensor parallelism multi-gpu inference. + +--output OUTPUT The output directory for generated configurations, including `mlc-chat-config.json` and tokenizer configuration. + +3. Compile Model Library +^^^^^^^^^^^^^^^^^^^^^^^^ + +After generating ``mlc-chat-config.json``, we can compile the model into a model library (files ending in ``.so``, ``.tar``, etc. that contains +the inference logic of a model). + +Model compilation command follows the pattern below: + +.. code:: text + + mlc_llm compile \ + MODEL \ + [--quantization QUANTIZATION_MODE] \ + [--model-type MODEL_TYPE] \ + [--device DEVICE] \ + [--host HOST] \ + [--opt OPT] \ + [--system-lib-prefix SYSTEM_LIB_PREFIX] \ + --output OUTPUT \ + [--overrides OVERRIDES] + +Note that ``MODEL`` is a positional argument. Arguments wrapped with ``[ ]`` are optional. + +--MODEL A path to ``mlc-chat-config.json``, or an MLC model directory that contains ``mlc-chat-config.json``. + +--quantization QUANTIZATION_MODE The quantization mode we use to compile. If unprovided, will infer from ``MODEL``. + + See :ref:`quantization_mode` for more information. + Available options are: ``q0f16``, ``q0f32``, ``q3f16_1``, ``q4f16_1``, ``q4f32_1``, and + ``q4f16_awq``. + + We encourage you to use 4-bit quantization, as the text generated by 3-bit + quantized models may have bad quality depending on the model. + +--model-type MODEL_TYPE Model architecture such as "llama". If not set, it is inferred from ``mlc-chat-config.json``. + +--device DEVICE The GPU device to compile the model to. If not set, it is inferred from GPUs available locally. + +--host HOST The host LLVM triple to compile the model to. If not set, it is inferred from the local CPU and OS. + Examples of the LLVM triple: + + 1) iPhones: arm64-apple-ios; + 2) ARM64 Android phones: aarch64-linux-android; + 3) WebAssembly: wasm32-unknown-unknown-wasm; + 4) Windows: x86_64-pc-windows-msvc; + 5) ARM macOS: arm64-apple-darwin. + +--opt OPT Optimization flags. MLC LLM maintains a predefined set of optimization flags, + denoted as ``O0``, ``O1``, ``O2``, ``O3``, where ``O0`` means no optimization, ``O2`` + means majority of them, and ``O3`` represents extreme optimization that could + potentially break the system. + + Meanwhile, optimization flags could be explicitly specified via details knobs, e.g. + ``--opt="cutlass_attn=1;cutlass_norm=0;cublas_gemm=0;cudagraph=0"``. + +--system-lib-prefix SYSTEM_LIB_PREFIX Adding a prefix to all symbols exported. Similar to ``objcopy --prefix-symbols``. + This is useful when compiling multiple models into a single library to avoid symbol + conflicts. Different from objcopy, this takes no effect for shared library. + + +--output OUTPUT The path to the output file. The suffix determines if the output file is a shared library or + objects. Available suffixes: + + 1) Linux: .so (shared), .tar (objects); + 2) macOS: .dylib (shared), .tar (objects); + 3) Windows: .dll (shared), .tar (objects); + 4) Android, iOS: .tar (objects); + 5) Web: .wasm (web assembly). + +--overrides OVERRIDES Model configuration override. Configurations to override ``mlc-chat-config.json``. Supports + ``context_window_size``, ``prefill_chunk_size``, ``sliding_window``, ``max_batch_size`` and + ``tensor_parallel_shards``. Meanwhile, model config could be explicitly specified via details + knobs, e.g. ``--overrides "context_window_size=1024;prefill_chunk_size=128"``. diff --git a/docs/compilation/configure_quantization.rst b/docs/compilation/configure_quantization.rst new file mode 100644 index 0000000..c489cb1 --- /dev/null +++ b/docs/compilation/configure_quantization.rst @@ -0,0 +1,89 @@ +Configure Quantization +====================== + +Quantization Algorithm +---------------------- + +The default quantization algorithm used in MLC-LLM is grouping quantization method discussed in the papers `The case for 4-bit precision: k-bit Inference Scaling Laws `__ and `LUT-GEMM: Quantized Matrix Multiplication based on LUTs for Efficient Inference in Large-Scale Generative Language Models `__. + +.. _quantization_mode: + +Quantization Mode +----------------- + +In MLC-LLM we use a short code that indicates the quantization mode to use. MLC-LLM supports both +weight-only quantization and weight-activation quantization. + +For the weight-only quantization, he format of the code is ``qAfB(_id)``, where ``A`` represents the number +of bits for storing weights and ``B`` represents the number of bits for storing activations. +The ``_id`` is an integer identifier to distinguish different quantization algorithms (e.g. symmetric, non-symmetric, AWQ, etc). + +Currently, available options are: ``q0f16``, ``q0f32``, ``q3f16_1``, ``q4f16_1``, ``q4f32_1``, and ``q4f16_awq`` (not stable). + +For the weight-activation quantization, currently MLC-LLM supports FP8 quantization on CUDA. +The available options are: ``e4m3_e4m3_f16`` and ``e5m2_e5m2_f16``. In these modes, both weights and activations are quantized to FP8 format. +The output of each layer is in higher precision (FP16) and then requantized to FP8. + +.. _calibration: + +Calibration +----------- + +For ``e4m3_e4m3_f16`` quantization, we need to calibrate the quantization parameters for the activations. +The calibration process is done by running the following command: + +1. Compile the calibration model +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +We use the same compilation workflow to compile the model in calibration mode. +The only difference is that we need to specify the quantization mode as ``e4m3_e4m3_f16_calibrate``. + +.. code-block:: bash + + mlc_llm gen_config \ + \ + --quantization e4m3_e4m3_f16_max_calibrate \ + --output + + mlc_llm convert_weights \ + \ + --quantization e4m3_e4m3_f16_max_calibrate \ + --output + + mlc_llm compile \ + \ + --output + +2. Run the calibration model +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +We will run the calibration model on the dataset such as ShareGPT to collect the statistics of the +activations. The calibration model will updates the quantization parameters in the weights file +in-place. We turn off the cuda graph as it is not yet supported in the calibration process. + +.. code-block:: bash + + mlc_llm calibrate \ + \ + --model-lib \ + --dataset \ + --num-calibration-samples \ + --opt "cudagraph=0" + --output + +3. Compile the quantized model for inference. +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +After the calibration process, we can compile the model for inference. In this step, we only need +to generate the configuration file using the desired quantization format and compile the model. +Weights are already quantized and calibrated in the previous steps and do not need to be converted again. + +.. code-block:: bash + + mlc_llm gen_config \ + \ + --quantization e4m3_e4m3_f16 \ + --output + mlc_llm compile \ + \ + --output diff --git a/docs/compilation/convert_weights.rst b/docs/compilation/convert_weights.rst new file mode 100644 index 0000000..16f36f5 --- /dev/null +++ b/docs/compilation/convert_weights.rst @@ -0,0 +1,157 @@ +.. _convert-weights-via-MLC: + +Convert Model Weights +===================== + +To run a model with MLC LLM, +we need to convert model weights into MLC format (e.g. `RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC `_.) +This page walks us through the process of adding a model variant with ``mlc_llm convert_weight``, which +takes a huggingface model as input and converts/quantizes into MLC-compatible weights. + +Specifically, we add RedPjama-INCITE-**Instruct**-3B-v1, while MLC already +provides a model library for RedPjama-INCITE-**Chat**-3B-v1, which we can reuse. + +This can be extended to, e.g.: + +- Add ``OpenHermes-Mistral`` when MLC already supports Mistral +- Add ``Llama-2-uncensored`` when MLC already supports Llama-2 + +.. note:: + Before you proceed, make sure you followed :ref:`install-tvm`, a required + backend to compile models with MLC LLM. + + Please also follow the instructions in :ref:`deploy-cli` / :ref:`deploy-python-engine` to obtain + the CLI app / Python API that can be used to chat with the compiled model. + + +.. contents:: Table of Contents + :depth: 1 + :local: + +.. _verify_installation_for_compile: + +1. Verify installation +---------------------- + +**Step 1. Verify mlc_llm** + +We use the python package ``mlc_llm`` to compile models. This can be installed by +following :ref:`install-mlc-packages`, either by building from source, or by +installing the prebuilt package. Verify ``mlc_llm`` installation in command line via: + +.. code:: bash + + $ mlc_llm --help + # You should see help information with this line + usage: MLC LLM Command Line Interface. [-h] {compile,convert_weight,gen_config} + +.. note:: + If it runs into error ``command not found: mlc_llm``, try ``python -m mlc_llm --help``. + +**Step 2. Verify TVM** + +To compile models, you also need to follow :ref:`install-tvm`. +Here we verify ``tvm`` quickly with command line (for full verification, see :ref:`tvm-validate`): + +.. code:: bash + + $ python -c "import tvm; print(tvm.__file__)" + /some-path/lib/python3.13/site-packages/tvm/__init__.py + + +1. Clone from HF and convert_weight +----------------------------------- + +You can be under the mlc-llm repo, or your own working directory. Note that all platforms +can share the same compiled/quantized weights. See :ref:`compile-command-specification` +for specification of ``convert_weight``. + +.. code:: shell + + # Create directory + mkdir -p dist/models && cd dist/models + # Clone HF weights + git lfs install + git clone https://huggingface.co/togethercomputer/RedPajama-INCITE-Instruct-3B-v1 + cd ../.. + # Convert weight + mlc_llm convert_weight ./dist/models/RedPajama-INCITE-Instruct-3B-v1/ \ + --quantization q4f16_1 \ + -o dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC + +.. _generate_mlc_chat_config: + +2. Generate MLC Chat Config +--------------------------- + +Use ``mlc_llm gen_config`` to generate ``mlc-chat-config.json`` and process tokenizers. +See :ref:`compile-command-specification` for specification of ``gen_config``. + +.. code:: shell + + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Instruct-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC/ + + +.. note:: + The file ``mlc-chat-config.json`` is crucial in both model compilation + and runtime chatting. Here we only care about the latter case. + + You can **optionally** customize + ``dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC/mlc-chat-config.json`` (checkout :ref:`configure-mlc-chat-json` for more detailed instructions). + You can also simply use the default configuration. + + `conversation_template `__ + directory contains a full list of conversation templates that MLC provides. If the model you are adding + requires a new conversation template, you would need to add your own. + Follow `this PR `__ as an example. However, + adding your own template would require you :ref:`build mlc_llm from source ` in order for it + to be recognized by the runtime. + +By now, you should have the following files. + +.. code:: shell + + ~/mlc-llm > ls dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + +.. _distribute-compiled-models: + +(Optional) 3. Upload weights to HF +---------------------------------- + +Optionally, you can upload what we have to huggingface. + +.. code:: shell + + # First, please create a repository on Hugging Face. + # With the repository created, run + git lfs install + git clone https://huggingface.co/my-huggingface-account/my-redpajama3b-weight-huggingface-repo + cd my-redpajama3b-weight-huggingface-repo + cp path/to/mlc-llm/dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC/* . + git add . && git commit -m "Add redpajama-3b instruct model weights" + git push origin main + +This would result in something like `RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC +`_, but +for **Instruct** instead of **Chat**. + +Good job, you have successfully distributed the model you compiled. +Next, we will talk about how we can consume the model weights in applications. + +Download the Distributed Models +------------------------------- + +You can now use the existing mlc tools such as chat/serve/package with the converted weights. + +.. code:: shell + + mlc_llm chat HF://my-huggingface-account/my-redpajama3b-weight-huggingface-repo diff --git a/docs/compilation/define_new_models.rst b/docs/compilation/define_new_models.rst new file mode 100644 index 0000000..289666e --- /dev/null +++ b/docs/compilation/define_new_models.rst @@ -0,0 +1,25 @@ +Define New Model Architectures +============================== + +This page guides you how to add a new model architecture in MLC. + +This notebook (runnable in Colab) should contain all necessary information to add a model in +MLC LLM: +https://github.com/mlc-ai/notebooks/blob/main/mlc-llm/tutorial_add_new_model_architecture_in_tvm_nn_module.ipynb + +In the notebook, we leverage ``tvm.nn.module`` to define a model in MLC LLM. We also use ``JIT`` +(just-in-time compilation) to debug the implementation. + +You can also refer to the PRs below on specific examples of adding a model architecture in MLC LLM: + +- `GPTNeoX PR `_ +- `GPT-2 PR `_ +- `Mistral PR `_ + +.. note:: + + When adding a model variant that has + its architecture already supported in mlc-llm , you **only need to convert weights** + (e.g. adding ``CodeLlama`` when MLC supports ``llama-2``; adding ``OpenHermes Mistral`` + when MLC supports ``mistral``). On the other hand, a new model architecture + (or inference logic) requires more work (following the tutorial above). diff --git a/docs/compilation/package_libraries_and_weights.rst b/docs/compilation/package_libraries_and_weights.rst new file mode 100644 index 0000000..16955e9 --- /dev/null +++ b/docs/compilation/package_libraries_and_weights.rst @@ -0,0 +1,219 @@ +.. _package-libraries-and-weights: + +Package Libraries and Weights +============================= + +When we want to build LLM applications with MLC LLM (e.g., iOS/Android apps), +usually we need to build static model libraries and app binding libraries, +and sometimes bundle model weights into the app. +MLC LLM provides a tool for fast model library and weight packaging: ``mlc_llm package``. + +This page briefly introduces how to use ``mlc_llm package`` for packaging. +Tutorials :ref:`deploy-ios` and :ref:`deploy-android` contain detailed examples and instructions +on using this packaging tool for iOS and Android deployment. + +----- + +Introduction +------------ + +To use ``mlc_llm package``, we must clone the source code of `MLC LLM `_ +and `install the MLC LLM and TVM package `_. +Depending on the app we build, there might be some other dependencies, which are described in +corresponding :ref:`iOS ` and :ref:`Android ` tutorials. + +After cloning, the basic usage of ``mlc_llm package`` is as the following. + +.. code:: bash + + export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm + cd /path/to/app # The app root directory which contains "mlc-package-config.json". + # E.g., "ios/MLCChat" or "android/MLCChat" + mlc_llm package + +**The package command reads from the JSON file** ``mlc-package-config.json`` **under the current directory.** +The output of this command is a directory ``dist/``, +which contains the packaged model libraries (under ``dist/lib/``) and weights (under ``dist/bundle/``). +This directory contains all necessary data for the app build. +Depending on the app we build, the internal structure of ``dist/lib/`` may be different. + +.. code:: + + dist + ├── lib + │ └── ... + └── bundle + └── ... + +The input ``mlc-package-config.json`` file specifies + +* the device (e.g., iPhone or Android) to package model libraries and weights for, +* the list of models to package. + +Below is an example ``mlc-package-config.json`` file: + +.. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/Mistral-7B-Instruct-v0.2-q3f16_1-MLC", + "model_id": "Mistral-7B-Instruct-v0.2-q3f16_1", + "estimated_vram_bytes": 3316000000, + "bundle_weight": true, + "overrides": { + "context_window_size": 512 + } + }, + { + "model": "HF://mlc-ai/gemma-2b-it-q4f16_1-MLC", + "model_id": "gemma-2b-q4f16_1", + "estimated_vram_bytes": 3000000000, + "overrides": { + "prefill_chunk_size": 128 + } + } + ] + } + +This example ``mlc-package-config.json`` specifies "iphone" as the target device. +In the ``model_list``, + +* ``model`` points to the Hugging Face repository which contains the pre-converted model weights. Apps will download model weights from the Hugging Face URL. +* ``model_id`` is a unique model identifier. +* ``estimated_vram_bytes`` is an estimation of the vRAM the model takes at runtime. +* ``"bundle_weight": true`` means the model weights of the model will be bundled into the app when building. +* ``overrides`` specifies some model config parameter overrides. + + +Below is a more detailed specification of the ``mlc-package-config.json`` file. +Each entry in ``"model_list"`` of the JSON file has the following fields: + +``model`` + (Required) The path to the MLC-converted model to be built into the app. + + Usually it is a Hugging Face URL (e.g., ``"model": "HF://mlc-ai/phi-2-q4f16_1-MLC"```) that contains the pre-converted model weights. + For iOS, it can also be a path to a local model directory which contains converted model weights (e.g., ``"model": "../dist/gemma-2b-q4f16_1"``). + Please check out :ref:`convert-weights-via-MLC` if you want to build local model into the app. + +``model_id`` + (Required) A unique local identifier to identify the model. + It can be an arbitrary one. + +``estimated_vram_bytes`` + (Required) Estimated requirements of vRAM to run the model. + +``bundle_weight`` + (Optional) A boolean flag indicating whether to bundle model weights into the app. + If this field is set to true, the ``mlc_llm package`` command will copy the model weights + to ``dist/bundle/$model_id``. + +``overrides`` + (Optional) A dictionary to override the default model context window size (to limit the KV cache size) and prefill chunk size (to limit the model temporary execution memory). + Example: + + .. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1", + "estimated_vram_bytes": 2960000000, + "overrides": { + "context_window_size": 512, + "prefill_chunk_size": 128 + } + } + ] + } + +``model_lib`` + (Optional) A string specifying the system library prefix to use for the model. + Usually this is used when you want to build multiple model variants with the same architecture into the app. + **This field does not affect any app functionality.** + The ``"model_lib_path_for_prepare_libs"`` introduced below is also related. + Example: + + .. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1", + "estimated_vram_bytes": 2960000000, + "model_lib": "gpt_neox_q4f16_1" + } + ] + } + + +Besides ``model_list`` in ``MLCChat/mlc-package-config.json``, +you can also **optionally** specify a dictionary of ``"model_lib_path_for_prepare_libs"``, +**if you want to use model libraries that are manually compiled**. +The keys of this dictionary should be the ``model_lib`` that specified in model list, +and the values of this dictionary are the paths (absolute, or relative) to the manually compiled model libraries. +The model libraries specified in ``"model_lib_path_for_prepare_libs"`` will be built into the app when running ``mlc_llm package``. +Example: + +.. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1", + "estimated_vram_bytes": 2960000000, + "model_lib": "gpt_neox_q4f16_1" + } + ], + "model_lib_path_for_prepare_libs": { + "gpt_neox_q4f16_1": "../../dist/lib/RedPajama-INCITE-Chat-3B-v1-q4f16_1-iphone.tar" + } + } + +Compilation Cache +----------------- +``mlc_llm package`` leverage a local JIT cache to avoid repetitive compilation of the same input. +It also leverages a local cache to download weights from remote. These caches +are shared across the entire project. Sometimes it is helpful to force rebuild when +we have a new compiler update or when something goes wrong with the cached library. +You can do so by setting the environment variable ``MLC_JIT_POLICY=REDO`` + +.. code:: bash + + MLC_JIT_POLICY=REDO mlc_llm package + +Arguments of ``mlc_llm package`` +-------------------------------- + +Command ``mlc_llm package`` can optionally take the arguments below: + +``--package-config`` + A path to ``mlc-package-config.json`` which contains the device and model specification. + By default, it is the ``mlc-package-config.json`` under the current directory. + +``--mlc-llm-source-dir`` + The path to MLC LLM source code (cloned from https://github.com/mlc-ai/mlc-llm). + By default, it is the ``$MLC_LLM_SOURCE_DIR`` environment variable. + If neither ``$MLC_LLM_SOURCE_DIR`` or ``--mlc-llm-source-dir`` is specified, error will be reported. + +``--output`` / ``-o`` + The output directory of ``mlc_llm package`` command. + By default, it is ``dist/`` under the current directory. + + +Summary and What to Do Next +--------------------------- + +In this page, we introduced the ``mlc_llm package`` command for fast model library and weight packaging. + +* It takes input file ``mlc-package-config.json`` which contains the device and model specification for packaging. +* It outputs directory ``dist/``, which contains packaged libraries under ``dist/lib/`` and model weights under ``dist/bundle/``. + +Next, please feel free to check out the :ref:`iOS ` and :ref:`Android ` tutorials for detailed examples of using ``mlc_llm package``. diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..90a9a0c --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,100 @@ +import os +import sys + +import tlcpack_sphinx_addon + +# -- General configuration ------------------------------------------------ + +sys.path.insert(0, os.path.abspath("../python")) +sys.path.insert(0, os.path.abspath("../")) +autodoc_mock_imports = ["torch"] + +# General information about the project. +project = "mlc-llm" +author = "MLC LLM Contributors" +copyright = "2023-2025, %s" % author + +# Version information. + +version = "0.1.0" +release = "0.1.0" + +extensions = [ + "sphinx_tabs.tabs", + "sphinx_toolbox.collapse", + "sphinxcontrib.httpdomain", + "sphinx.ext.autodoc", + "sphinx.ext.napoleon", + "sphinx_reredirects", +] + +redirects = {"get_started/try_out": "../index.html#getting-started"} + +source_suffix = [".rst"] + +language = "en" + +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "sphinx" + +# A list of ignored prefixes for module index sorting. +# If true, `todo` and `todoList` produce output, else they produce nothing. +todo_include_todos = False + +# -- Options for HTML output ---------------------------------------------- + +# The theme is set by the make target +import sphinx_rtd_theme # noqa: E402 + +html_theme = "sphinx_rtd_theme" +html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] + +templates_path = [] + +html_static_path = [] + +footer_copyright = "© 2023-2025 MLC LLM" +footer_note = " " + +html_logo = "_static/img/mlc-logo-with-text-landscape.svg" + +html_theme_options = { + "logo_only": True, +} + +header_links = [ + ("Home", "https://llm.mlc.ai/"), + ("Github", "https://github.com/mlc-ai/mlc-llm"), + ("Discord Server", "https://discord.gg/9Xpy2HGBuD"), +] + +header_dropdown = { + "name": "Other Resources", + "items": [ + ("MLC Course", "https://mlc.ai/"), + ("MLC Blog", "https://blog.mlc.ai/"), + ("Web LLM", "https://webllm.mlc.ai/"), + ], +} + +html_context = { + "footer_copyright": footer_copyright, + "footer_note": footer_note, + "header_links": header_links, + "header_dropdown": header_dropdown, + "display_github": True, + "github_user": "mlc-ai", + "github_repo": "mlc-llm", + "github_version": "main/docs/", + "theme_vcs_pageview_mode": "edit", + # "header_logo": "/path/to/logo", + # "header_logo_link": "", + # "version_selecter": "", +} + + +# add additional overrides +templates_path += [tlcpack_sphinx_addon.get_templates_path()] +html_static_path += [tlcpack_sphinx_addon.get_static_path()] diff --git a/docs/deploy/android.rst b/docs/deploy/android.rst new file mode 100644 index 0000000..f7d58ec --- /dev/null +++ b/docs/deploy/android.rst @@ -0,0 +1,367 @@ +.. _deploy-android: + +Android SDK +=========== + +.. contents:: Table of Contents + :local: + :depth: 2 + +Demo App +-------- + +The demo APK below is built for Samsung S23 with Snapdragon 8 Gen 2 chip. + +.. image:: https://seeklogo.com/images/D/download-android-apk-badge-logo-D074C6882B-seeklogo.com.png + :width: 135 + :target: https://github.com/mlc-ai/binary-mlc-llm-libs/releases/download/Android-09262024/mlc-chat.apk + +Prerequisite +------------ + +**Rust** (`install `__) is needed to cross-compile HuggingFace tokenizers to Android. Make sure rustc, cargo, and rustup are available in ``$PATH``. + +**Android Studio** (`install `__) with NDK and CMake. To install NDK and CMake, on the Android Studio welcome page, click "Projects → SDK Manager → SDK Tools". If you have already installed NDK in your development environment, please update your NDK to avoid build android package fail(`#2696 `__). The current demo Android APK is built with NDK 27.0.11718014. Once you have installed or updated the NDK, set up the following environment variables: + + +- ``ANDROID_NDK`` so that ``$ANDROID_NDK/build/cmake/android.toolchain.cmake`` is available. +- ``TVM_NDK_CC`` that points to NDK's clang compiler. + +.. code-block:: bash + + # Example on macOS + ANDROID_NDK: $HOME/Library/Android/sdk/ndk/27.0.11718014 + TVM_NDK_CC: $ANDROID_NDK/toolchains/llvm/prebuilt/darwin-x86_64/bin/aarch64-linux-android24-clang + # Example on Linux + ANDROID_NDK: $HOME/Android/Sdk/ndk/27.0.11718014 + TVM_NDK_CC: $ANDROID_NDK/toolchains/llvm/prebuilt/linux-x86_64/bin/aarch64-linux-android24-clang + # Example on Windows + ANDROID_NDK: %HOME%/AppData/Local/Android/Sdk/ndk/27.0.11718014 + TVM_NDK_CC: %ANDROID_NDK%/toolchains/llvm/prebuilt/windows-x86_64/bin/aarch64-linux-android24-clang + +**JDK**, such as OpenJDK >= 17, to compile Java bindings of TVM runtime. +We strongly recommend setting the ``JAVA_HOME`` to the JDK bundled with Android Studio. +e.g. +``export JAVA_HOME=/Applications/Android\ Studio.app/Contents/jbr/Contents/Home`` for macOS. +``export JAVA_HOME=/opt/android-studio/jbr`` for Linux. +Using Android Studio's JBR bundle as recommended `here https://developer.android.com/build/jdks` +will reduce the chances of potential errors in JNI compilation. +Set up the following environment variable: + +- ``export JAVA_HOME=/path/to/java_home`` you can then cross check and make sure ``$JAVA_HOME/bin/java`` exists. + +Please ensure that the JDK versions for Android Studio and JAVA_HOME are the same. + +**TVM runtime** is placed under `3rdparty/tvm `__ in MLC LLM, so there is no need to install anything extra. Set up the following environment variable: + +- ``export TVM_SOURCE_DIR=/path/to/mlc-llm/3rdparty/tvm``. + +Please follow :doc:`/install/mlc_llm` to obtain a binary build of mlc_llm package. Note that this +is independent from mlc-llm source code that we use for android package build in the following up section. +Once you installed this package, you do not need to build mlc llm from source. + +.. note:: + ❗ Whenever using Python, it is highly recommended to use **conda** to manage an isolated Python environment to avoid missing dependencies, incompatible versions, and package conflicts. + +Check if **environment variable** are properly set as the last check. One way to ensure this is to place them in ``$HOME/.zshrc``, ``$HOME/.bashrc`` or environment management tools. + +.. code-block:: bash + + source $HOME/.cargo/env # Rust + export ANDROID_NDK=... # Android NDK toolchain + export TVM_NDK_CC=... # Android NDK clang + export JAVA_HOME=... # Java + export TVM_SOURCE_DIR=... # TVM runtime + +Additional Guides for Windows Users +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Building under Windows for Android is still experimental; please make sure you +first finish the above guides, then read and follow the instructions in this section +If you are using Windows, make sure you use conda to install cmake and Ninja. + +.. code-block:: bash + + conda install -c conda-forge cmake ninja git git-lfs zstd + +Windows Java findings have issues with environment variables that come with space. +Make sure you get a copy of Java in a path without space. The simplest way to do that +is to copy the Android Studio's JBR bundle to a directory without any space. +If your Android studio's installation is at ``C:\Program Files\Android\Android Studio\`` +you can try to do the following + +.. code-block:: bash + + cp -r "C:\Program Files\Android\Android Studio\jbr" C:\any-path-without-space + set JAVA_HOME=C:\any-path-without-space + +You can continue the next steps after you have set these steps correctly. + +Build Android App from Source +----------------------------- + +This section shows how we can build the app from the source. + +Step 1. Install Build Dependencies +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +First and foremost, please clone the `MLC LLM GitHub repository `_. +After cloning, go to the ``android/`` directory. + +.. code:: bash + + git clone https://github.com/mlc-ai/mlc-llm.git + cd mlc-llm + git submodule update --init --recursive + cd android + + +.. _android-build-runtime-and-model-libraries: + +Step 2. Build Runtime and Model Libraries +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The models to be built for the Android app are specified in ``MLCChat/mlc-package-config.json``: +in the ``model_list``, ``model`` points to the Hugging Face repository which + +* ``model`` points to the Hugging Face repository which contains the pre-converted model weights. The Android app will download model weights from the Hugging Face URL. +* ``model_id`` is a unique model identifier. +* ``estimated_vram_bytes`` is an estimation of the vRAM the model takes at runtime. +* ``"bundle_weight": true`` means the model weights of the model will be bundled into the app when building. +* ``overrides`` specifies some model config parameter overrides. + + +We have a one-line command to build and prepare all the model libraries: + +.. code:: bash + + cd /path/to/MLCChat # e.g., "android/MLCChat" + export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm # has to be absolute path, ../.. does not work + mlc_llm package + +This command mainly executes the following two steps: + +1. **Compile models.** We compile each model in ``model_list`` of ``MLCChat/mlc-package-config.json`` into a binary model library. +2. **Build runtime and tokenizer.** In addition to the model itself, a lightweight runtime and tokenizer are required to actually run the LLM. + +The command creates a ``./dist/`` directory that contains the runtime and model build output. +Please make sure all the following files exist in ``./dist/``. + +.. code:: + + dist + └── lib + └── mlc4j + ├── build.gradle + ├── output + │ ├── arm64-v8a + │ │ └── libtvm4j_runtime_packed.so + │ └── tvm4j_core.jar + └── src + ├── cpp + │ └── tvm_runtime.h + └── main + ├── AndroidManifest.xml + ├── assets + │ └── mlc-app-config.json + └── java + └── ... + +The model execution logic in mobile GPUs is incorporated into ``libtvm4j_runtime_packed.so``, +while ``tvm4j_core.jar`` is a lightweight (~60 kb) `Java binding `_ +to it. ``dist/lib/mlc4j`` is a gradle subproject that you should include in your app +so the Android project can reference the mlc4j (MLC LLM java library). +This library packages the dependent model libraries and necessary runtime to execute the model. + +.. code:: + + include ':mlc4j' + project(':mlc4j').projectDir = file('dist/lib/mlc4j') + + +.. note:: + + We leverage a local JIT cache to avoid repetitive compilation of the same input. + However, sometimes it is helpful to force rebuild when we have a new compiler update + or when something goes wrong with the cached library. + You can do so by setting the environment variable ``MLC_JIT_POLICY=REDO`` + + .. code:: bash + + MLC_JIT_POLICY=REDO mlc_llm package + + +Step 3. Build Android App +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Open folder ``./android/MLCChat`` as an Android Studio Project. +Connect your Android device to your machine. +In the menu bar of Android Studio, click **"Build → Make Project"**. +Once the build is finished, click **"Run → Run 'app'"** and you will see the app launched on your phone. + +.. note:: + ❗ This app cannot be run in an emulator and thus a physical phone is required, because MLC LLM needs an actual mobile GPU to meaningfully run at an accelerated speed. + + +Customize the App +----------------- + +We can customize the models built in the Android app by customizing `MLCChat/mlc-package-config.json `_. +We introduce each field of the JSON file here. + +Each entry in ``"model_list"`` of the JSON file has the following fields: + +``model`` + (Required) The path to the MLC-converted model to be built into the app. + It is a Hugging Face URL (e.g., ``"model": "HF://mlc-ai/phi-2-q4f16_1-MLC"```) that contains + the pre-converted model weights. + +``model_id`` + (Required) A unique local identifier to identify the model. + It can be an arbitrary one. + +``estimated_vram_bytes`` + (Required) Estimated requirements of vRAM to run the model. + +``bundle_weight`` + (Optional) A boolean flag indicating whether to bundle model weights into the app. See :ref:`android-bundle-model-weights` below. + +``overrides`` + (Optional) A dictionary to override the default model context window size (to limit the KV cache size) and prefill chunk size (to limit the model temporary execution memory). + Example: + + .. code:: json + + { + "device": "android", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "estimated_vram_bytes": 1948348579, + "overrides": { + "context_window_size": 512, + "prefill_chunk_size": 128 + } + } + ] + } + +``model_lib`` + (Optional) A string specifying the system library prefix to use for the model. + Usually this is used when you want to build multiple model variants with the same architecture into the app. + **This field does not affect any app functionality.** + The ``"model_lib_path_for_prepare_libs"`` introduced below is also related. + Example: + + .. code:: json + + { + "device": "android", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "estimated_vram_bytes": 1948348579, + "model_lib": "gpt_neox_q4f16_1" + } + ] + } + + +Besides ``model_list`` in ``MLCChat/mlc-package-config.json``, +you can also **optionally** specify a dictionary of ``"model_lib_path_for_prepare_libs"``, +**if you want to use model libraries that are manually compiled**. +The keys of this dictionary should be the ``model_lib`` that specified in model list, +and the values of this dictionary are the paths (absolute, or relative) to the manually compiled model libraries. +The model libraries specified in ``"model_lib_path_for_prepare_libs"`` will be built into the app when running ``mlc_llm package``. +Example: + +.. code:: json + + { + "device": "android", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "estimated_vram_bytes": 1948348579, + "model_lib": "gpt_neox_q4f16_1" + } + ], + "model_lib_path_for_prepare_libs": { + "gpt_neox_q4f16_1": "../../dist/lib/RedPajama-INCITE-Chat-3B-v1-q4f16_1-android.tar" + } + } + +.. _android-bundle-model-weights: + +Bundle Model Weights +-------------------- + +Instructions have been provided to build an Android App with MLC LLM in previous sections, +but it requires run-time weight downloading from HuggingFace, +as configured in ``MLCChat/mlc-package-config.json``. +However, it could be desirable to bundle weights together into the app to avoid downloading over the network. +In this section, we provide a simple ADB-based walkthrough that hopefully helps with further development. + +**Enable weight bundle**. +Set the field ``"bundle_weight": true`` for any model you want to bundle weights +in ``MLCChat/mlc-package-config.json``, and run ``mlc_llm package`` again. +Below is an example: + +.. code:: json + + { + "device": "android", + "model_list": [ + { + "model": "HF://mlc-ai/gemma-2b-it-q4f16_1-MLC", + "model_id": "gemma-2b-q4f16_1-MLC", + "estimated_vram_bytes": 3000000000, + "bundle_weight": true + } + ] + } + +The outcome of running ``mlc_llm package`` should be as follows: + +.. code:: + + dist + ├── bundle + │ ├── gemma-2b-q4f16_1 # The model weights that will be bundled into the app. + │ └── mlc-app-config.json + └── ... + + +**Generating APK**. Enter Android Studio, and click **"Build → Generate Signed Bundle/APK"** to build an APK for release. If it is the first time you generate an APK, you will need to create a key according to `the official guide from Android `_. +This APK will be placed under ``android/MLCChat/app/release/app-release.apk``. + +**Install ADB and USB debugging**. Enable "USB debugging" in the developer mode in your phone settings. +In "SDK manager - SDK Tools", install `Android SDK Platform-Tools `_. +Add the path to platform-tool path to the environment variable ``PATH`` (on macOS, it is ``$HOME/Library/Android/sdk/platform-tools``). +Run the following commands, and if ADB is installed correctly, your phone will appear as a device: + +.. code-block:: bash + + adb devices + +**Install the APK and weights to your phone**. +Run the commands below to install the app, and push the local weights to the app data directory on your device. +Once it finishes, you can start the MLCChat app on your device. +The models with ``bundle_weight`` set to true will have their weights already on device. + +.. code-block:: bash + + cd /path/to/MLCChat # e.g., "android/MLCChat" + python bundle_weight.py --apk-path app/release/app-release.apk + +Known issues +------------ + +One known issue that has been observed on Android devices equipped with Adreno GPUs is that model formats ending with a ``_1`` suffix cause a ~20-50 seconds system UI freeze that occurs at prefill stage (initialization before the first inference; the issue does not happen on any subsequent inference of a given model instance). +It has been observed that models with a ``_0`` suffix do not experience this issue. +The two suffixes denote the layouts of weights in the models that differ by a transpose operation. +In case you encounter the freeze issue, the workaround to avoid this problem is to use a model in the ``_0`` weight layout. +For more details, please consult `issue #3363 `_. diff --git a/docs/deploy/cli.rst b/docs/deploy/cli.rst new file mode 100644 index 0000000..c100f65 --- /dev/null +++ b/docs/deploy/cli.rst @@ -0,0 +1,90 @@ +.. _deploy-cli: + +CLI +=============== + +MLC Chat CLI is the command line tool to run MLC-compiled LLMs out of the box interactively. + +.. contents:: Table of Contents + :local: + :depth: 2 + +Install MLC-LLM Package +------------------------ + +Chat CLI is a part of the MLC-LLM package. +To use the chat CLI, first install MLC LLM by following the instructions :ref:`here `. +Once you have install the MLC-LLM package, you can run the following command to check if the installation was successful: + +.. code:: bash + + mlc_llm chat --help + +You should see serve help message if the installation was successful. + +Quick Start +------------ + +This section provides a quick start guide to work with MLC-LLM chat CLI. +To launch the CLI session, run the following command: + +.. code:: bash + + mlc_llm chat MODEL [--model-lib PATH-TO-MODEL-LIB] + +where ``MODEL`` is the model folder after compiling with :ref:`MLC-LLM build process `. Information about other arguments can be found in the next section. + +Once the chat CLI is ready, you can enter the prompt to interact with the model. + +.. code:: + + You can use the following special commands: + /help print the special commands + /exit quit the cli + /stats print out stats of last request (token/sec) + /metrics print out full engine metrics + /reset restart a fresh chat + /set [overrides] override settings in the generation config. For example, + `/set temperature=0.5;top_p=0.8;seed=23;max_tokens=100;stop=str1,str2` + Note: Separate stop words in the `stop` option with commas (,). + Multi-line input: Use escape+enter to start a new line. + + >>> What's the meaning of life? + The meaning of life is a philosophical and metaphysical question related to the purpose or significance of life or existence in general... + +Run CLI with Multi-GPU +---------------------- + +If you want to enable tensor parallelism to run LLMs on multiple GPUs, please specify argument ``--overrides "tensor_parallel_shards=$NGPU"``. For example, + +.. code:: shell + + mlc_llm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC --overrides "tensor_parallel_shards=2" + + +The ``mlc_llm chat`` Command +---------------------------- + +We provide the list of chat CLI interface for reference. + +.. code:: bash + + mlc_llm chat MODEL [--model-lib PATH-TO-MODEL-LIB] [--device DEVICE] [--overrides OVERRIDES] + + +MODEL The model folder after compiling with MLC-LLM build process. The parameter + can either be the model name with its quantization scheme + (e.g. ``Llama-2-7b-chat-hf-q4f16_1``), or a full path to the model + folder. In the former case, we will use the provided name to search + for the model folder over possible paths. + +--model-lib A field to specify the full path to the model library file to use (e.g. a ``.so`` file). +--device The description of the device to run on. User should provide a string in the + form of ``device_name:device_id`` or ``device_name``, where ``device_name`` is one of + ``cuda``, ``metal``, ``vulkan``, ``rocm``, ``opencl``, ``auto`` (automatically detect the + local device), and ``device_id`` is the device id to run on. The default value is ``auto``, + with the device id set to 0 for default. +--overrides Model configuration override. Supports overriding + ``context_window_size``, ``prefill_chunk_size``, ``sliding_window_size``, ``attention_sink_size``, + and ``tensor_parallel_shards``. The overrides could be explicitly + specified via details knobs, e.g. --overrides ``context_window_size=1024;prefill_chunk_size=128``. diff --git a/docs/deploy/ide_integration.rst b/docs/deploy/ide_integration.rst new file mode 100644 index 0000000..89a9edb --- /dev/null +++ b/docs/deploy/ide_integration.rst @@ -0,0 +1,179 @@ +.. _deploy-ide-integration: + +IDE Integration +=============== + +.. contents:: Table of Contents + :local: + :depth: 2 + +MLC LLM has now support for code completion on multiple IDEs. This means you can easily integrate an LLM with coding capabilities with your IDE through the MLC LLM :ref:`deploy-rest-api`. Here we provide a step-by-step guide on how to do this. + +Convert Your Model Weights +-------------------------- + +To run a model with MLC LLM in any platform, you need to convert your model weights to the MLC format (e.g. `CodeLlama-7b-hf-q4f16_1-MLC `__). You can always refer to :ref:`convert-weights-via-MLC` for in-depth details on how to convert your model weights. If you are using your own model weights, i.e., you finetuned the model on your personal codebase, it is important to follow these steps to convert the respective weights properly. However, it is also possible to download precompiled weights from the original models, available in the MLC format. See the full list of all precompiled weights `here `__. + +**Example:** + +.. code:: bash + + # convert model weights + mlc_llm convert_weight ./dist/models/CodeLlama-7b-hf \ + --quantization q4f16_1 \ + -o ./dist/CodeLlama-7b-hf-q4f16_1-MLC + +Compile Your Model +------------------ + +Compiling the model architecture is the crucial step to optimize inference for a given platform. However, compilation relies on multiple settings that will impact the runtime. This configuration is specified inside the ``mlc-chat-config.json`` file, which can be generated by the ``gen_config`` command. You can learn more about the ``gen_config`` command `here `__. + +**Example:** + +.. code:: bash + + # generate mlc-chat-config.json + mlc_llm gen_config ./dist/models/CodeLlama-7b-hf \ + --quantization q4f16_1 --conv-template LM \ + -o ./dist/CodeLlama-7b-hf-q4f16_1-MLC + +.. note:: + Make sure to set the ``--conv-template`` flag to ``LM``. This template is specifically tailored to perform vanilla LLM completion, generally adopted by code completion models. + +After generating the MLC model configuration file, we are all set to compile and create the model library. You can learn more about the ``compile`` command `here `__ + +**Example:** + +.. tabs:: + + .. group-tab:: Linux - CUDA + + .. code:: bash + + # compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/CodeLlama-7b-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device cuda -o ./dist/libs/CodeLlama-7b-hf-q4f16_1-cuda.so + + .. group-tab:: Metal + + For M-chip Mac: + + .. code:: bash + + # compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/CodeLlama-7b-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o ./dist/libs/CodeLlama-7b-hf-q4f16_1-metal.so + + Cross-Compiling for Intel Mac on M-chip Mac: + + .. code:: bash + + # compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/CodeLlama-7b-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device metal:x86-64 -o ./dist/libs/CodeLlama-7b-hf-q4f16_1-metal_x86_64.dylib + + For Intel Mac: + + .. code:: bash + + # compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/CodeLlama-7b-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device metal -o ./dist/libs/CodeLlama-7b-hf-q4f16_1-metal_x86_64.dylib + + .. group-tab:: Vulkan + + For Linux: + + .. code:: bash + + # compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/CodeLlama-7b-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o ./dist/libs/CodeLlama-7b-hf-q4f16_1-vulkan.so + + For Windows: + + .. code:: bash + + # compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/CodeLlama-7b-hf-q4f16_1-MLC/mlc-chat-config.json \ + --device vulkan -o ./dist/libs/CodeLlama-7b-hf-q4f16_1-vulkan.dll + +.. note:: + The generated model library can be shared across multiple model variants, as long as the architecture and number of parameters does not change, e.g., same architecture, but different weights (your finetuned model). + +Setting up the Inference Entrypoint +----------------------------------- + +You can now locally deploy your compiled model with the MLC serve module. To find more details about the MLC LLM API visit our :ref:`deploy-rest-api` page. + +**Example:** + +.. code:: bash + + python -m mlc_llm.serve.server \ + --model dist/CodeLlama-7b-hf-q4f16_1-MLC \ + --model-lib ./dist/libs/CodeLlama-7b-hf-q4f16_1-cuda.so + +Configure the IDE Extension +--------------------------- + +After deploying the LLM we can easily connect the IDE with the MLC Rest API. In this guide, we will be using the Hugging Face Code Completion extension `llm-ls `__ which has support across multiple IDEs (e.g., `vscode `__, `intellij `__ and `nvim `__) to connect to an external OpenAI compatible API (i.e., our MLC LLM :ref:`deploy-rest-api`). + +After installing the extension on your IDE, open the ``settings.json`` extension configuration file: + +.. figure:: /_static/img/ide_code_settings.png + :width: 450 + :align: center + :alt: settings.json + +| + +Then, make sure to replace the following settings with the respective values: + +.. code:: javascript + + "llm.modelId": "dist/CodeLlama-7b-hf-q4f16_1-MLC" + "llm.url": "http://127.0.0.1:8000/v1/completions" + "llm.backend": "openai" + +This will enable the extension to send OpenAI compatible requests to the MLC Serve API. Also, feel free to tune the API parameters. Please refer to our :ref:`deploy-rest-api` documentation for more details about these API parameters. + +.. code:: javascript + + "llm.requestBody": { + "best_of": 1, + "frequency_penalty": 0.0, + "presence_penalty": 0.0, + "logprobs": false, + "top_logprobs": 0, + "logit_bias": null, + "max_tokens": 128, + "seed": null, + "stop": null, + "suffix": null, + "temperature": 1.0, + "top_p": 1.0 + } + +The llm-ls extension supports a variety of different model code completion templates. Choose the one that best matches your model, i.e., the template with the correct tokenizer and Fill in the Middle tokens. + +.. figure:: /_static/img/ide_code_templates.png + :width: 375 + :align: center + :alt: llm-ls templates + +| + +After everything is all set, the extension will be ready to use the responses from the MLC Serve API to provide off-the-shelf code completion on your IDE. + +.. figure:: /_static/img/code_completion.png + :width: 700 + :align: center + :alt: IDE Code Completion + +| + +Conclusion +---------- + +Please, let us know if you have any questions. Feel free to open an issue on the `MLC LLM repo `__! diff --git a/docs/deploy/ios.rst b/docs/deploy/ios.rst new file mode 100644 index 0000000..b81a934 --- /dev/null +++ b/docs/deploy/ios.rst @@ -0,0 +1,413 @@ +.. _deploy-ios: + +iOS Swift SDK +============= + +.. contents:: Table of Contents + :local: + :depth: 2 + +The MLC LLM iOS app can be installed in two ways: through the pre-built package or by building from the source. +If you are an iOS user looking to try out the models, the pre-built package is recommended. If you are a +developer seeking to integrate new features into the package, building the iOS package from the source is required. + +Use Pre-built iOS App +--------------------- +The MLC Chat app is now available in App Store at no cost. You can download and explore it by simply clicking the button below: + + .. image:: https://developer.apple.com/assets/elements/badges/download-on-the-app-store.svg + :width: 135 + :target: https://apps.apple.com/us/app/mlc-chat/id6448482937 + + +Build iOS App from Source +------------------------- + +This section shows how we can build the app from the source. + +Step 1. Install Build Dependencies +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +First and foremost, please clone the `MLC LLM GitHub repository `_. +After cloning, go to the ``ios/`` directory. + +.. code:: bash + + git clone https://github.com/mlc-ai/mlc-llm.git + cd mlc-llm + git submodule update --init --recursive + cd ./ios + + +Please follow :doc:`/install/mlc_llm` to obtain a binary build of mlc_llm package. Note that this +is independent from the above source code that we use for iOS package build. +You do not need to build mlc_llm for your host and we can use the prebuilt package for that purpose. + +We also need to have the following build dependencies: + +* CMake >= 3.24, +* Git and Git-LFS, +* `Rust and Cargo `_, which are required by Hugging Face's tokenizer. + +.. _ios-build-runtime-and-model-libraries: + +Step 2. Build Runtime and Model Libraries +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The models to be built for the iOS app are specified in ``MLCChat/mlc-package-config.json``: +in the ``model_list``, + +* ``model`` points to the Hugging Face repository which contains the pre-converted model weights. The iOS app will download model weights from the Hugging Face URL. +* ``model_id`` is a unique model identifier. +* ``estimated_vram_bytes`` is an estimation of the vRAM the model takes at runtime. +* ``"bundle_weight": true`` means the model weights of the model will be bundled into the app when building. +* ``overrides`` specifies some model config parameter overrides. + + +We have a one-line command to build and prepare all the model libraries: + +.. code:: bash + + cd /path/to/MLCChat # e.g., "ios/MLCChat" + export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm # e.g., "../.." + mlc_llm package + +This command mainly executes the following two steps: + +1. **Compile models.** We compile each model in ``model_list`` of ``MLCChat/mlc-package-config.json`` into a binary model library. +2. **Build runtime and tokenizer.** In addition to the model itself, a lightweight runtime and tokenizer are required to actually run the LLM. + +The command creates a ``./dist/`` directory that contains the runtime and model build output. +Please make sure ``dist/`` follows the structure below, except the optional model weights. + +.. code:: + + dist + ├── bundle # The directory for mlc-app-config.json (and optionally model weights) + │ │ # that will be bundled into the iOS app. + │ ├── mlc-app-config.json # The app config JSON file. + │ └── [optional model weights] + └── lib + ├── libmlc_llm.a # A lightweight interface to interact with LLM, tokenizer, and TVM runtime. + ├── libmodel_iphone.a # The compiled model lib. + ├── libsentencepiece.a # SentencePiece tokenizer + ├── libtokenizers_cpp.a # Huggingface tokenizer. + └── libtvm_runtime.a # TVM runtime. + + +.. note:: + + We leverage a local JIT cache to avoid repetitive compilation of the same input. + However, sometimes it is helpful to force rebuild when we have a new compiler update + or when something goes wrong with the cached library. + You can do so by setting the environment variable ``MLC_JIT_POLICY=REDO`` + + .. code:: bash + + MLC_JIT_POLICY=REDO mlc_llm package + +.. _ios-bundle-model-weights: + +Step 3. (Optional) Bundle model weights into the app +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +By default, we download the model weights from Hugging Face when running the app. +**As an option,**, we bundle model weights into the app: +set the field ``"bundle_weight": true`` for any model you want to bundle weights +in ``MLCChat/mlc-package-config.json``, and run ``mlc_llm package`` again. +Below is an example: + +.. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/gemma-2b-it-q4f16_1-MLC", + "model_id": "gemma-2b-q4f16_1", + "estimated_vram_bytes": 3000000000, + "overrides": { + "prefill_chunk_size": 128 + }, + "bundle_weight": true + } + ] + } + +The outcome of running ``mlc_llm package`` should be as follows: + +.. code:: + + dist + ├── bundle + │ ├── gemma-2b-q4f16_1 # The model weights that will be bundled into the app. + │ └── mlc-app-config.json + └── ... + +.. _ios-build-app: + +Step 4. Build iOS App +^^^^^^^^^^^^^^^^^^^^^ + +Open ``./ios/MLCChat/MLCChat.xcodeproj`` using Xcode. Note that you will need an +Apple Developer Account to use Xcode, and you may be prompted to use +your own developer team credential and product bundle identifier. + +Ensure that all the necessary dependencies and configurations are +correctly set up in the Xcode project. + +Once you have made the necessary changes, build the iOS app using Xcode. +If you have an Apple Silicon Mac, you can select target "My Mac (designed for iPad)" +to run on your Mac. You can also directly run it on your iPad or iPhone. + +.. image:: https://raw.githubusercontent.com/mlc-ai/web-data/main/images/mlc-llm/tutorials/xcode-build.jpg + :align: center + :width: 60% + +| + +Customize the App +----------------- + +We can customize the models built in the iOS app by customizing `MLCChat/mlc-package-config.json `_. +We introduce each field of the JSON file here. + +Each entry in ``"model_list"`` of the JSON file has the following fields: + +``model`` + (Required) The path to the MLC-converted model to be built into the app. + + It can be either a Hugging Face URL (e.g., ``"model": "HF://mlc-ai/phi-2-q4f16_1-MLC"```), or a path to a local model directory which contains converted model weights (e.g., ``"model": "../dist/gemma-2b-q4f16_1"``). Please check out :ref:`convert-weights-via-MLC` if you want to build local model into the app. + + *Note: the local path (if relative) is relative to the* ``ios/`` *directory.* + +``model_id`` + (Required) A unique local identifier to identify the model. + It can be an arbitrary one. + +``estimated_vram_bytes`` + (Required) Estimated requirements of vRAM to run the model. + +``bundle_weight`` + (Optional) A boolean flag indicating whether to bundle model weights into the app. See :ref:`ios-bundle-model-weights`. + +``overrides`` + (Optional) A dictionary to override the default model context window size (to limit the KV cache size) and prefill chunk size (to limit the model temporary execution memory). + Example: + + .. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1", + "estimated_vram_bytes": 2960000000, + "overrides": { + "context_window_size": 512, + "prefill_chunk_size": 128 + } + } + ] + } + +``model_lib`` + (Optional) A string specifying the system library prefix to use for the model. + Usually this is used when you want to build multiple model variants with the same architecture into the app. + **This field does not affect any app functionality.** + The ``"model_lib_path_for_prepare_libs"`` introduced below is also related. + Example: + + .. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1", + "estimated_vram_bytes": 2960000000, + "model_lib": "gpt_neox_q4f16_1" + } + ] + } + + +Besides ``model_list`` in ``MLCChat/mlc-package-config.json``, +you can also **optionally** specify a dictionary of ``"model_lib_path_for_prepare_libs"``, +**if you want to use model libraries that are manually compiled**. +The keys of this dictionary should be the ``model_lib`` that specified in model list, +and the values of this dictionary are the paths (absolute, or relative) to the manually compiled model libraries. +The model libraries specified in ``"model_lib_path_for_prepare_libs"`` will be built into the app when running ``mlc_llm package``. +Example: + +.. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC", + "model_id": "RedPajama-INCITE-Chat-3B-v1-q4f16_1", + "estimated_vram_bytes": 2960000000, + "model_lib": "gpt_neox_q4f16_1" + } + ], + "model_lib_path_for_prepare_libs": { + "gpt_neox_q4f16_1": "../../dist/lib/RedPajama-INCITE-Chat-3B-v1-q4f16_1-iphone.tar" + } + } + + +Bring Your Own Model +-------------------- + +This section introduces how to build your own model into the iOS app. +We use the example of `NeuralHermes `_ model, which a variant of Mistral model. + +.. note:: + + This section largely replicates :ref:`convert-weights-via-MLC`. + See that page for more details. Note that the weights are shared across + all platforms in MLC. + +**Step 1. Clone from HF and convert_weight** + +You can be under the mlc-llm repo, or your own working directory. Note that all platforms +can share the same compiled/quantized weights. See :ref:`compile-command-specification` +for specification of ``convert_weight``. + +.. code:: shell + + # Create directory + mkdir -p dist/models && cd dist/models + # Clone HF weights + git lfs install + git clone https://huggingface.co/mlabonne/NeuralHermes-2.5-Mistral-7B + cd ../.. + # Convert weight + mlc_llm convert_weight ./dist/models/NeuralHermes-2.5-Mistral-7B/ \ + --quantization q4f16_1 \ + -o dist/NeuralHermes-2.5-Mistral-7B-q3f16_1-MLC + +**Step 2. Generate MLC Chat Config** + +Use ``mlc_llm gen_config`` to generate ``mlc-chat-config.json`` and process tokenizers. +See :ref:`compile-command-specification` for specification of ``gen_config``. + +.. code:: shell + + mlc_llm gen_config ./dist/models/NeuralHermes-2.5-Mistral-7B/ \ + --quantization q3f16_1 --conv-template neural_hermes_mistral \ + -o dist/NeuralHermes-2.5-Mistral-7B-q3f16_1-MLC + +For the ``conv-template``, `conversation_template.py `__ +contains a full list of conversation templates that MLC provides. + +If the model you are adding requires a new conversation template, you would need to add your own. +Follow `this PR `__ as an example. +We look up the template to use with the ``conv_template`` field in ``mlc-chat-config.json``. + +For more details, please see :ref:`configure-mlc-chat-json`. + +**Step 3. Upload weights to HF** + +.. code:: shell + + # First, please create a repository on Hugging Face. + # With the repository created, run + git lfs install + git clone https://huggingface.co/my-huggingface-account/my-mistral-weight-huggingface-repo + cd my-mistral-weight-huggingface-repo + cp path/to/mlc-llm/dist/NeuralHermes-2.5-Mistral-7B-q3f16_1-MLC/* . + git add . && git commit -m "Add mistral model weights" + git push origin main + +After successfully following all steps, you should end up with a Huggingface repo similar to +`NeuralHermes-2.5-Mistral-7B-q3f16_1-MLC `__, +which includes the converted/quantized weights, the ``mlc-chat-config.json``, and tokenizer files. + + +**Step 4. Register in Model List** + +Finally, we add the model into the ``model_list`` of +`MLCChat/mlc-package-config.json `_ by specifying the Hugging Face link as ``model``: + +.. code:: json + + { + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/NeuralHermes-2.5-Mistral-7B-q3f16_1-MLC", + "model_id": "Mistral-7B-Instruct-v0.2-q3f16_1", + "estimated_vram_bytes": 3316000000, + } + ] + } + + +Now, go through :ref:`ios-build-runtime-and-model-libraries` and :ref:`ios-build-app` again. +The app will use the ``NeuralHermes-Mistral`` model you just added. + + +Build Apps with MLC Swift API +----------------------------- + +We also provide a Swift package that you can use to build +your own app. The package is located under ``ios/MLCSwift``. + +- First, create ``mlc-package-config.json`` in your project folder. + You do so by copying the files in MLCChat folder. + Run ``mlc_llm package``. + This will give us the necessary libraries under ``/path/to/project/dist``. +- Under "Build phases", add ``/path/to/project/dist/bundle`` this will copying + this folder into your app to include bundled weights and configs. +- Add ``ios/MLCSwift`` package to your app in Xcode. + Under "Frameworks, Libraries, and Embedded Content", click add package dependencies + and add local package that points to ``ios/MLCSwift``. +- Finally, we need to add the libraries dependencies. Under build settings: + + - Add library search path ``/path/to/project/dist/lib``. + - Add the following items to "other linker flags". + + .. code:: + + -Wl,-all_load + -lmodel_iphone + -lmlc_llm -ltvm_runtime + -ltokenizers_cpp + -lsentencepiece + -ltokenizers_c + + +You can then import the `MLCSwift` package into your app. +The following code shows an illustrative example of how to use the chat module. + +.. code:: swift + + import MLCSwift + + func runExample() async { + let engine = MLCEngine() + let modelPath = "/path/to/model/weights" + let modelLib = "model-lib-name" + + await engine.reload(modelPath: modelPath, modelLib: modelLib) + + // run chat completion as in OpenAI API style + for await res in await engine.chat.completions.create( + messages: [ + ChatCompletionMessage( + role: .user, + content: "What is the meaning of life?" + ) + ] + ) { + print(res.choices[0].delta.content!.asText()) + } + } + +Checkout `MLCEngineExample `_ +for a minimal starter example. diff --git a/docs/deploy/mlc_chat_config.rst b/docs/deploy/mlc_chat_config.rst new file mode 100644 index 0000000..4222a2c --- /dev/null +++ b/docs/deploy/mlc_chat_config.rst @@ -0,0 +1,199 @@ +.. _configure-mlc-chat-json: + +Customize MLC Chat Config +========================= + +``mlc-chat-config.json`` is required for both compile-time and runtime, hence serving two purposes: + +1. Specify how we compile a model (shown in :ref:`compile-model-libraries`), and +2. Specify conversation behavior in runtime. + +**This page focuses on the second purpose.** We explain the components of a chat +configuration and how to customize them by modifying the file. Additionally, +the runtimes also provide APIs to optionally override some of the configurations. + +In runtime, this file is stored under the directory of each compiled model +(e.g. `RedPajama chat config `__). + + +.. _struct-mlc-chat-conv: + +Structure of MLCChat Configuration +---------------------------------- + +Below is the ``mlc-chat-config.json`` file corresponding to Llama2 model: + +.. code:: json + + // mlc-chat-config.json + { + // 1. Metadata used to specify how to compile a model + "model_type": "llama", + "quantization": "q4f16_1", + "version": "0.1.0", + "model_config": { + "hidden_size": 4096, + "intermediate_size": 11008, + // more fields here... + }, + "vocab_size": 32000, + "context_window_size": 4096, + "sliding_window_size": -1, + "prefill_chunk_size": 4096, + "tensor_parallel_shards": 1, + + // 2. Tokenizer-related fields + "pad_token_id": 0, + "bos_token_id": 1, + "eos_token_id": 2, + "tokenizer_files": [ + "tokenizer.model", + "tokenizer.json", + "tokenizer_config.json" + ] + + // 3. Conversation template related fields + "conv_template": { + "name": "llama-2", + "system_template": "[INST] <>\n{system_message}\n<>\n\n ", + "system_message": "You are a helpful, respectful and honest assistant.", + // more fields here... + }, + + // 4. Chat related fields that affect runtime behavior + "temperature": 0.6, + "repetition_penalty": 1.0, + "top_p": 0.9 + } + +.. note:: + Fields in the first part of ``mlc-chat-config.json`` (e.g. ``context-window-size``) + is only for compile-time. Changing them during runtime may lead to unexpected behavior. + +**As shown above, the file is divided into three parts. We focus on the third part, which +can be customized to change the behavior of the model.** + +``conv_template`` + .. note:: + Legacy ``mlc-chat-config.json`` may specify a string for this field to look up a registered conversation + template. It will be deprecated in the future. Re-generate config using the latest version of mlc_llm + to make sure this field is a complete JSON object. + + The conversation template that this chat uses. For more information, please refer to :ref:`conversation structure `. + +``temperature`` + The temperature applied to logits before sampling. The default value is ``0.7``. A higher temperature encourages more diverse outputs, while a lower temperature produces more deterministic outputs. + +``repetition_penalty`` + The repetition penalty controls the likelihood of the model generating repeated texts. The default value is set to ``1.0``, indicating that no repetition penalty is applied. Increasing the value reduces the likelihood of repeat text generation. However, setting a high ``repetition_penalty`` may result in the model generating meaningless texts. The ideal choice of repetition penalty may vary among models. + + For more details on how repetition penalty controls text generation, please check out the `CTRL paper `_. + +``top_p`` + This parameter determines the set of tokens from which we sample during decoding. The default value is set to ``0.95``. At each step, we select tokens from the minimal set that has a cumulative probability exceeding the ``top_p`` parameter. + + For additional information on top-p sampling, please refer to this `blog post `_. + + +.. _struct-conv: + +Conversation Structure +^^^^^^^^^^^^^^^^^^^^^^ + +MLC-LLM provided a set of pre-defined conversation templates, which you can directly use by +specifying ``--conv-template [name]`` when generating config. Below is a list (not complete) of +supported conversation templates: + +- ``llama-2`` +- ``mistral_default`` +- ``chatml`` +- ``phi-2`` +- ... + +Please refer to `conversation_template `_ directory for the full list of supported templates and their implementations. + +Below is a generic structure of a JSON conversation configuration (we use vicuna as an example): + +.. code:: json + + // mlc-chat-config.json + { + // ... + "conv_template": { + "name": "llama-2", + "system_template": "[INST] <>\n{system_message}\n<>\n\n ", + "system_message": "You are a helpful, respectful and honest assistant.", + "roles": { + "user": "[INST]", + "assistant": "[/INST]", + "tool": "[INST]" + }, + "role_templates": { + "user": "{user_message}", + "assistant": "{assistant_message}", + "tool": "{tool_message}" + }, + "messages": [], + "seps": [ + " " + ], + "role_content_sep": " ", + "role_empty_sep": " ", + "stop_str": [ + "[INST]" + ], + "stop_token_ids": [ + 2 + ], + "function_string": "", + "use_function_calling": false + } + } + +``name`` + Name of the conversation. +``system_template`` + The system prompt template, it optionally contains the system + message placeholder, and the placeholder will be replaced with + the system message below. +``system_message`` + The content of the system prompt (without the template format). +``system_prefix_token_ids`` + The system token ids to be prepended at the beginning of tokenized + generated prompt. +``roles`` + The conversation roles +``role_templates`` + The roles prompt template, it optionally contains the defaults + message placeholders and will be replaced by actual content +``messages`` + The conversation history messages. + Each message is a pair of strings, denoting "(role, content)". + The content can be None. +``seps`` + An array of strings indicating the separators to be used after a user + message and a model message respectively. +``role_content_sep`` + The separator between the role and the content in a message. +``role_empty_sep`` + The separator between the role and empty contents. +``stop_str`` + When the ``stop_str`` is encountered, the model will stop generating output. +``stop_token_ids`` + A list of token IDs that act as stop tokens. +``function_string`` + The function calling string. +``use_function_calling`` + Whether using function calling or not, helps check for output message format in API call. + + +Given a conversation template, the corresponding prompt generated out +from it is in the following format: + +.. code:: text + + <><><><><> + <><><><> + ... + <><><><> + <><> diff --git a/docs/deploy/python_engine.rst b/docs/deploy/python_engine.rst new file mode 100644 index 0000000..d2f149a --- /dev/null +++ b/docs/deploy/python_engine.rst @@ -0,0 +1,296 @@ +.. _deploy-python-engine: + +Python API +========== + +.. note:: + This page introduces the Python API with MLCEngine in MLC LLM. + +.. contents:: Table of Contents + :local: + :depth: 2 + + +MLC LLM provides Python API through classes :class:`mlc_llm.MLCEngine` and :class:`mlc_llm.AsyncMLCEngine` +which **support full OpenAI API completeness** for easy integration into other Python projects. + +This page introduces how to use the engines in MLC LLM. +The Python API is a part of the MLC-LLM package, which we have prepared pre-built pip wheels via +the :ref:`installation page `. + + +Verify Installation +------------------- + +.. code:: bash + + python -c "from mlc_llm import MLCEngine; print(MLCEngine)" + +You are expected to see the output of ````. + +If the command above results in error, follow :ref:`install-mlc-packages` to install prebuilt pip +packages or build MLC LLM from source. + + +Run MLCEngine +------------- + +:class:`mlc_llm.MLCEngine` provides the interface of OpenAI chat completion synchronously. +:class:`mlc_llm.MLCEngine` does not batch concurrent request due to the synchronous design, +and please use :ref:`AsyncMLCEngine ` for request batching process. + +**Stream Response.** In :ref:`quick-start` and :ref:`introduction-to-mlc-llm`, +we introduced the basic use of :class:`mlc_llm.MLCEngine`. + +.. code:: python + + from mlc_llm import MLCEngine + + # Create engine + model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" + engine = MLCEngine(model) + + # Run chat completion in OpenAI API. + for response in engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=True, + ): + for choice in response.choices: + print(choice.delta.content, end="", flush=True) + print("\n") + + engine.terminate() + +This code example first creates an :class:`mlc_llm.MLCEngine` instance with the 8B Llama-3 model. +**We design the Python API** :class:`mlc_llm.MLCEngine` **to align with OpenAI API**, +which means you can use :class:`mlc_llm.MLCEngine` in the same way of using +`OpenAI's Python package `_ +for both synchronous and asynchronous generation. + +**Non-stream Response.** The code example above uses the synchronous chat completion +interface and iterate over all the stream responses. +If you want to run without streaming, you can run + +.. code:: python + + response = engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=False, + ) + print(response) + +Please refer to `OpenAI's Python package `_ +and `OpenAI chat completion API `_ +for the complete chat completion interface. + +.. note:: + + If you want to enable tensor parallelism to run LLMs on multiple GPUs, + please specify argument ``model_config_overrides`` in MLCEngine constructor. + For example, + + .. code:: python + + from mlc_llm import MLCEngine + from mlc_llm.serve.config import EngineConfig + + model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" + engine = MLCEngine( + model, + engine_config=EngineConfig(tensor_parallel_shards=2), + ) + + +.. _python-engine-async-llm-engine: + +Run AsyncMLCEngine +------------------ + +:class:`mlc_llm.AsyncMLCEngine` provides the interface of OpenAI chat completion with +asynchronous features. +**We recommend using** :class:`mlc_llm.AsyncMLCEngine` **to batch concurrent request for better throughput.** + +**Stream Response.** The core use of :class:`mlc_llm.AsyncMLCEngine` for stream responses is as follows. + +.. code:: python + + async for response in await engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=True, + ): + for choice in response.choices: + print(choice.delta.content, end="", flush=True) + +.. collapse:: The collapsed is a complete runnable example of AsyncMLCEngine in Python. + + .. code:: python + + import asyncio + from typing import Dict + + from mlc_llm.serve import AsyncMLCEngine + + model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" + prompts = [ + "Write a three-day travel plan to Pittsburgh.", + "What is the meaning of life?", + ] + + + async def test_completion(): + # Create engine + async_engine = AsyncMLCEngine(model=model) + + num_requests = len(prompts) + output_texts: Dict[str, str] = {} + + async def generate_task(prompt: str): + async for response in await async_engine.chat.completions.create( + messages=[{"role": "user", "content": prompt}], + model=model, + stream=True, + ): + if response.id not in output_texts: + output_texts[response.id] = "" + output_texts[response.id] += response.choices[0].delta.content + + tasks = [asyncio.create_task(generate_task(prompts[i])) for i in range(num_requests)] + await asyncio.gather(*tasks) + + # Print output. + for request_id, output in output_texts.items(): + print(f"Output of request {request_id}:\n{output}\n") + + async_engine.terminate() + + + asyncio.run(test_completion()) + +| + +**Non-stream Response.** Similarly, :class:`mlc_llm.AsyncEngine` provides the non-stream response +interface. + +.. code:: python + + response = await engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=False, + ) + print(response) + +Please refer to `OpenAI's Python package `_ +and `OpenAI chat completion API `_ +for the complete chat completion interface. + +.. note:: + + If you want to enable tensor parallelism to run LLMs on multiple GPUs, + please specify argument ``model_config_overrides`` in AsyncMLCEngine constructor. + For example, + + .. code:: python + + from mlc_llm import AsyncMLCEngine + from mlc_llm.serve.config import EngineConfig + + model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" + engine = AsyncMLCEngine( + model, + engine_config=EngineConfig(tensor_parallel_shards=2), + ) + + +Engine Mode +----------- + +To ease the engine configuration, the constructors of :class:`mlc_llm.MLCEngine` and +:class:`mlc_llm.AsyncMLCEngine` have an optional argument ``mode``, +which falls into one of the three options ``"local"``, ``"interactive"`` or ``"server"``. +The default mode is ``"local"``. + +Each mode denotes a pre-defined configuration of the engine to satisfy different use cases. +The choice of the mode controls the request concurrency of the engine, +as well as engine's KV cache token capacity (or in other words, the maximum +number of tokens that the engine's KV cache can hold), +and further affects the GPU memory usage of the engine. + +In short, + +- mode ``"local"`` uses low request concurrency and low KV cache capacity, which is suitable for cases where **concurrent requests are not too many, and the user wants to save GPU memory usage**. +- mode ``"interactive"`` uses 1 as the request concurrency and low KV cache capacity, which is designed for **interactive use cases** such as chats and conversations. +- mode ``"server"`` uses as much request concurrency and KV cache capacity as possible. This mode aims to **fully utilize the GPU memory for large server scenarios** where concurrent requests may be many. + +**For system benchmark, please select mode** ``"server"``. +Please refer to :ref:`python-engine-api-reference` for detailed documentation of the engine mode. + + +Deploy Your Own Model with Python API +------------------------------------- + +The :ref:`introduction page ` introduces how we can deploy our +own models with MLC LLM. +This section introduces how you can use the model weights you convert and the model library you build +in :class:`mlc_llm.MLCEngine` and :class:`mlc_llm.AsyncMLCEngine`. + +We use the `Phi-2 `_ as the example model. + +**Specify Model Weight Path.** Assume you have converted the model weights for your own model, +you can construct a :class:`mlc_llm.MLCEngine` as follows: + +.. code:: python + + from mlc_llm import MLCEngine + + model = "models/phi-2" # Assuming the converted phi-2 model weights are under "models/phi-2" + engine = MLCEngine(model) + + +**Specify Model Library Path.** Further, if you build the model library on your own, +you can use it in :class:`mlc_llm.MLCEngine` by passing the library path through argument ``model_lib``. + +.. code:: python + + from mlc_llm import MLCEngine + + model = "models/phi-2" + model_lib = "models/phi-2/lib.so" # Assuming the phi-2 model library is built at "models/phi-2/lib.so" + engine = MLCEngine(model, model_lib=model_lib) + + +The same applies to :class:`mlc_llm.AsyncMLCEngine`. + + +.. _python-engine-api-reference: + +API Reference +------------- + +The :class:`mlc_llm.MLCEngine` and :class:`mlc_llm.AsyncMLCEngine` classes provide the following constructors. + +The MLCEngine and AsyncMLCEngine have full OpenAI API completeness. +Please refer to `OpenAI's Python package `_ +and `OpenAI chat completion API `_ +for the complete chat completion interface. + +.. currentmodule:: mlc_llm + +.. autoclass:: MLCEngine + :members: + :exclude-members: evaluate + :undoc-members: + :show-inheritance: + + .. automethod:: __init__ + +.. autoclass:: AsyncMLCEngine + :members: + :exclude-members: evaluate + :undoc-members: + :show-inheritance: + + .. automethod:: __init__ diff --git a/docs/deploy/rest.rst b/docs/deploy/rest.rst new file mode 100644 index 0000000..82f5291 --- /dev/null +++ b/docs/deploy/rest.rst @@ -0,0 +1,435 @@ +.. _deploy-rest-api: + +REST API +======== + +.. contents:: Table of Contents + :local: + :depth: 2 + +We provide `REST API `_ +for a user to interact with MLC-LLM in their own programs. + +Install MLC-LLM Package +------------------------ + +SERVE is a part of the MLC-LLM package, installation instruction for which can be found :ref:`here `. Once you have install the MLC-LLM package, you can run the following command to check if the installation was successful: + +.. code:: bash + + mlc_llm serve --help + +You should see serve help message if the installation was successful. + +Quick Start +------------ + +This section provides a quick start guide to work with MLC-LLM REST API. To launch a server, run the following command: + +.. code:: bash + + mlc_llm serve MODEL [--model-lib PATH-TO-MODEL-LIB] + +where ``MODEL`` is the model folder after compiling with :ref:`MLC-LLM build process `. Information about other arguments can be found under :ref:`Launch the server ` section. + +Once you have launched the Server, you can use the API in your own program to send requests. Below is an example of using the API to interact with MLC-LLM in Python without Streaming (suppose the server is running on ``http://127.0.0.1:8080/``): + +.. code:: bash + + import requests + + # Get a response using a prompt without streaming + payload = { + "model": "./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/", + "messages": [ + {"role": "user", "content": "Write a haiku about apples."}, + ], + "stream": False, + # "n": 1, + "max_tokens": 300, + } + r = requests.post("http://127.0.0.1:8080/v1/chat/completions", json=payload) + choices = r.json()["choices"] + for choice in choices: + print(f"{choice['message']['content']}\n") + +Run CLI with Multi-GPU +---------------------- + +If you want to enable tensor parallelism to run LLMs on multiple GPUs, please specify argument ``--overrides "tensor_parallel_shards=$NGPU"``. For example, + +.. code:: shell + + mlc_llm serve HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC --overrides "tensor_parallel_shards=2" + +------------------------------------------------ + + +.. _rest_launch_server: + + +Launch the Server +----------------- + +To launch the MLC Server for MLC-LLM, run the following command in your terminal. + +.. code:: bash + + mlc_llm serve MODEL [--model-lib PATH-TO-MODEL-LIB] [--device DEVICE] [--mode MODE] \ + [--additional-models ADDITIONAL-MODELS] \ + [--speculative-mode SPECULATIVE-MODE] \ + [--overrides OVERRIDES] \ + [--enable-tracing] \ + [--host HOST] \ + [--port PORT] \ + [--allow-credentials] \ + [--allowed-origins ALLOWED_ORIGINS] \ + [--allowed-methods ALLOWED_METHODS] \ + [--allowed-headers ALLOWED_HEADERS] + + +MODEL The model folder after compiling with MLC-LLM build process. The parameter + can either be the model name with its quantization scheme + (e.g. ``Llama-2-7b-chat-hf-q4f16_1``), or a full path to the model + folder. In the former case, we will use the provided name to search + for the model folder over possible paths. + +--model-lib A field to specify the full path to the model library file to use (e.g. a ``.so`` file). +--device The description of the device to run on. User should provide a string in the + form of ``device_name:device_id`` or ``device_name``, where ``device_name`` is one of + ``cuda``, ``metal``, ``vulkan``, ``rocm``, ``opencl``, ``auto`` (automatically detect the + local device), and ``device_id`` is the device id to run on. The default value is ``auto``, + with the device id set to 0 for default. +--mode The engine mode in MLC LLM. + We provide three preset modes: ``local``, ``interactive`` and ``server``. + The default mode is ``local``. + + The choice of mode decides the values of "max_num_sequence", "max_total_sequence_length" + and "prefill_chunk_size" when they are not explicitly specified. + + 1. Mode "local" refers to the local server deployment which has low + request concurrency. So the max batch size will be set to 4, and max + total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + + 2. Mode "interactive" refers to the interactive use of server, which + has at most 1 concurrent request. So the max batch size will be set to 1, + and max total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + + 3. Mode "server" refers to the large server use case which may handle + many concurrent request and want to use GPU memory as much as possible. + In this mode, we will automatically infer the largest possible max batch + size and max total sequence length. + + You can manually specify arguments "max_num_sequence", "max_total_seq_length" and + "prefill_chunk_size" via ``--overrides`` to override the automatic inferred values. + For example: ``--overrides "max_num_sequence=32;max_total_seq_length=4096"``. +--additional-models The model paths and (optional) model library paths of additional models (other + than the main model). + + When engine is enabled with speculative decoding, additional models are needed. + **We only support one additional model for speculative decoding now.** + The way of specifying the additional model is: + ``--additional-models model_path_1`` or + ``--additional-models model_path_1,model_lib_1``. + + When the model lib of a model is not given, JIT model compilation will be activated + to compile the model automatically. +--speculative-mode The speculative decoding mode. Right now four options are supported: + + - ``disable``, where speculative decoding is not enabled, + + - ``small_draft``, denoting the normal speculative decoding (small draft) style, + + - ``eagle``, denoting the eagle-style speculative decoding. + + - ``medusa``, denoting the medusa-style speculative decoding. +--overrides Overriding extra configurable fields of EngineConfig. + + Supporting fields that can be be overridden: ``tensor_parallel_shards``, ``max_num_sequence``, + ``max_total_seq_length``, ``prefill_chunk_size``, ``max_history_size``, ``gpu_memory_utilization``, + ``spec_draft_length``, ``prefix_cache_max_num_recycling_seqs``, ``context_window_size``, + ``sliding_window_size``, ``attention_sink_size``. + + Please check out the documentation of EngineConfig in ``mlc_llm/serve/config.py`` + for detailed docstring of each field. + Example: ``--overrides "max_num_sequence=32;max_total_seq_length=4096;tensor_parallel_shards=2"`` +--enable-tracing A boolean indicating if to enable event logging for requests. +--host The host at which the server should be started, defaults to ``127.0.0.1``. +--port The port on which the server should be started, defaults to ``8000``. +--allow-credentials A flag to indicate whether the server should allow credentials. If set, the server will + include the ``CORS`` header in the response +--allowed-origins Specifies the allowed origins. It expects a JSON list of strings, with the default value being ``["*"]``, allowing all origins. +--allowed-methods Specifies the allowed methods. It expects a JSON list of strings, with the default value being ``["*"]``, allowing all methods. +--allowed-headers Specifies the allowed headers. It expects a JSON list of strings, with the default value being ``["*"]``, allowing all headers. + +You can access ``http://127.0.0.1:PORT/docs`` (replace ``PORT`` with the port number you specified) to see the list of +supported endpoints. + +API Endpoints +------------- + +The REST API provides the following endpoints: + +.. http:get:: /v1/models + +------------------------------------------------ + + Get a list of models available for MLC-LLM. + +**Example** + +.. code:: bash + + import requests + + url = "http://127.0.0.1:8000/v1/models" + headers = {"accept": "application/json"} + + response = requests.get(url, headers=headers) + + if response.status_code == 200: + print("Response:") + print(response.json()) + else: + print("Error:", response.status_code) + + +.. http:post:: /v1/chat/completions + +------------------------------------------------ + + Get a response from MLC-LLM using a prompt, either with or without streaming. + +**Chat Completion Request Object** + +- **messages** (*List[ChatCompletionMessage]*, required): A sequence of messages that have been exchanged in the conversation so far. Each message in the conversation is represented by a `ChatCompletionMessage` object, which includes the following fields: + - **content** (*Optional[Union[str, List[Dict[str, str]]]]*): The text content of the message or structured data in case of tool-generated messages. + - **role** (*Literal["system", "user", "assistant", "tool"]*): The role of the message sender, indicating whether the message is from the system, user, assistant, or a tool. + - **name** (*Optional[str]*): An optional name for the sender of the message. + - **tool_calls** (*Optional[List[ChatToolCall]]*): A list of calls to external tools or functions made within this message, applicable when the role is `tool`. + - **tool_call_id** (*Optional[str]*): A unique identifier for the tool call, relevant when integrating external tools or services. + +- **model** (*str*, required): The model to be used for generating responses. + +- **frequency_penalty** (*float*, optional, default=0.0): Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model’s likelihood to repeat tokens. + +- **presence_penalty** (*float*, optional, default=0.0): Positive values penalize new tokens if they are already present in the text so far, decreasing the model’s likelihood to repeat tokens. + +- **logprobs** (*bool*, optional, default=False): Indicates whether to include log probabilities for each token in the response. + +- **top_logprobs** (*int*, optional, default=0): An integer ranging from 0 to 20. It determines the number of tokens, most likely to appear at each position, to be returned. Each token is accompanied by a log probability. If this parameter is used, 'logprobs' must be set to true. + +- **logit_bias** (*Optional[Dict[int, float]]*): Allows specifying biases for or against specific tokens during generation. + +- **max_tokens** (*Optional[int]*): The maximum number of tokens to generate in the response(s). + +- **n** (*int*, optional, default=1): Number of responses to generate for the given prompt. + +- **seed** (*Optional[int]*): A seed for deterministic generation. Using the same seed and inputs will produce the same output. + +- **stop** (*Optional[Union[str, List[str]]]*): One or more strings that, if encountered, will cause generation to stop. + +- **stream** (*bool*, optional, default=False): If `True`, responses are streamed back as they are generated. + +- **temperature** (*float*, optional, default=1.0): Controls the randomness of the generation. Lower values lead to less random completions. + +- **top_p** (*float*, optional, default=1.0): Nucleus sampling parameter that controls the diversity of the generated responses. + +- **tools** (*Optional[List[ChatTool]]*): Specifies external tools or functions that can be called as part of the chat. + +- **tool_choice** (*Optional[Union[Literal["none", "auto"], Dict]]*): Controls how tools are selected for use in responses. + +- **user** (*Optional[str]*): An optional identifier for the user initiating the request. + +- **response_format** (*RequestResponseFormat*, optional): Specifies the format of the response. Can be either "text" or "json_object", with optional schema definition for JSON responses. + +**Returns** + +- If `stream` is `False`, a `ChatCompletionResponse` object containing the generated response(s). +- If `stream` is `True`, a stream of `ChatCompletionStreamResponse` objects, providing a real-time feed of generated responses. + + +**ChatCompletionResponseChoice** + +- **finish_reason** (*Optional[Literal["stop", "length", "tool_calls", "error"]]*, optional): The reason the completion process was terminated. It can be due to reaching a stop condition, the maximum length, output of tool calls, or an error. + +- **index** (*int*, required, default=0): Indicates the position of this choice within the list of choices. + +- **message** (*ChatCompletionMessage*, required): The message part of the chat completion, containing the content of the chat response. + +- **logprobs** (*Optional[LogProbs]*, optional): Optionally includes log probabilities for each output token + +**ChatCompletionStreamResponseChoice** + +- **finish_reason** (*Optional[Literal["stop", "length", "tool_calls"]]*, optional): Specifies why the streaming completion process ended. Valid reasons are "stop", "length", and "tool_calls". + +- **index** (*int*, required, default=0): Indicates the position of this choice within the list of choices. + +- **delta** (*ChatCompletionMessage*, required): Represents the incremental update or addition to the chat completion message in the stream. + +- **logprobs** (*Optional[LogProbs]*, optional): Optionally includes log probabilities for each output token + +**ChatCompletionResponse** + +- **id** (*str*, required): A unique identifier for the chat completion session. + +- **choices** (*List[ChatCompletionResponseChoice]*, required): A collection of `ChatCompletionResponseChoice` objects, representing the potential responses generated by the model. + +- **created** (*int*, required, default=current time): The UNIX timestamp representing when the response was generated. + +- **model** (*str*, required): The name of the model used to generate the chat completions. + +- **system_fingerprint** (*str*, required): A system-generated fingerprint that uniquely identifies the computational environment. + +- **object** (*Literal["chat.completion"]*, required, default="chat.completion"): A string literal indicating the type of object, here always "chat.completion". + +- **usage** (*UsageInfo*, required, default=empty `UsageInfo` object): Contains information about the API usage for this specific request. + +**ChatCompletionStreamResponse** + +- **id** (*str*, required): A unique identifier for the streaming chat completion session. + +- **choices** (*List[ChatCompletionStreamResponseChoice]*, required): A list of `ChatCompletionStreamResponseChoice` objects, each representing a part of the streaming chat response. + +- **created** (*int*, required, default=current time): The creation time of the streaming response, represented as a UNIX timestamp. + +- **model** (*str*, required): Specifies the model that was used for generating the streaming chat completions. + +- **system_fingerprint** (*str*, required): A unique identifier for the system generating the streaming completions. + +- **object** (*Literal["chat.completion.chunk"]*, required, default="chat.completion.chunk"): A literal indicating that this object represents a chunk of a streaming chat completion. + +------------------------------------------------ + + +**Example** + +Below is an example of using the API to interact with MLC-LLM in Python with Streaming. + +.. code:: bash + + import requests + import json + + # Get a response using a prompt with streaming + payload = { + "model": "./dist/Llama-2-7b-chat-hf-q4f16_1-MLC/", + "messages": [{"role": "user", "content": "Write a haiku"}], + "stream": True, + } + with requests.post("http://127.0.0.1:8080/v1/chat/completions", json=payload, stream=True) as r: + for chunk in r.iter_content(chunk_size=None): + chunk = chunk.decode("utf-8") + if "[DONE]" in chunk[6:]: + break + response = json.loads(chunk[6:]) + content = response["choices"][0]["delta"].get("content", "") + print(content, end="", flush=True) + print("\n") + +------------------------------------------------ + +There is also support for function calling similar to OpenAI (https://platform.openai.com/docs/guides/function-calling). Below is an example on how to use function calling in Python. + +.. code:: bash + + import requests + import json + + tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + }, + } + ] + + payload = { + "model": "./dist/gorilla-openfunctions-v1-q4f16_1-MLC/", + "messages": [ + { + "role": "user", + "content": "What is the current weather in Pittsburgh, PA in fahrenheit?", + } + ], + "stream": False, + "tools": tools, + } + + r = requests.post("http://127.0.0.1:8080/v1/chat/completions", json=payload) + print(f"{r.json()['choices'][0]['message']['tool_calls'][0]['function']}\n") + + # Output: {'name': 'get_current_weather', 'arguments': {'location': 'Pittsburgh, PA', 'unit': 'fahrenheit'}} + +------------------------------------------------ + +Function Calling with streaming is also supported. Below is an example on how to use function calling with streaming in Python. + +.. code:: bash + + import requests + import json + + tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + }, + } + ] + + payload = { + "model": "./dist/gorilla-openfunctions-v1-q4f16_1-MLC/", + "messages": [ + { + "role": "user", + "content": "What is the current weather in Pittsburgh, PA and Tokyo, JP in fahrenheit?", + } + ], + "stream": True, + "tools": tools, + } + + with requests.post("http://127.0.0.1:8080/v1/chat/completions", json=payload, stream=True) as r: + for chunk in r.iter_content(chunk_size=None): + chunk = chunk.decode("utf-8") + if "[DONE]" in chunk[6:]: + break + response = json.loads(chunk[6:]) + content = response["choices"][0]["delta"].get("content", "") + print(f"{content}", end="", flush=True) + print("\n") + + # Output: ["get_current_weather(location='Pittsburgh,PA',unit='fahrenheit')", "get_current_weather(location='Tokyo,JP',unit='fahrenheit')"] + + +.. note:: + The API is a uniform interface that supports multiple languages. You can also utilize these functionalities in languages other than Python. diff --git a/docs/deploy/webllm.rst b/docs/deploy/webllm.rst new file mode 100644 index 0000000..61c85d5 --- /dev/null +++ b/docs/deploy/webllm.rst @@ -0,0 +1,370 @@ +.. _webllm-runtime: + +WebLLM Javascript SDK +===================== + +.. contents:: Table of Contents + :local: + :depth: 2 + +`WebLLM `_ is a high-performance in-browser LLM +inference engine, aiming to be the backend of AI-powered web applications and agents. + +It provides a specialized runtime for the web backend of MLCEngine, leverages +`WebGPU `_ for local acceleration, offers OpenAI-compatible API, +and provides built-in support for web workers to separate heavy computation from the UI flow. + +Please checkout the `WebLLM repo `__ on how to use WebLLM to build +web application in Javascript/Typescript. Here we only provide a high-level idea and discuss how to +use MLC-LLM to compile your own model to run with WebLLM. + +Getting Started +--------------- + +To get started, try out `WebLLM Chat `__, which provides a great example +of integrating WebLLM into a full web application. + +A WebGPU-compatible browser is needed to run WebLLM-powered web applications. +You can download the latest Google Chrome and use `WebGPU Report `__ +to verify the functionality of WebGPU on your browser. + +WebLLM is available as an `npm package `_ and is +also CDN-delivered. Try a simple chatbot example in +`this JSFiddle example `__ without setup. + +You can also checkout `existing examples `__ +on more advanced usage of WebLLM such as JSON mode, streaming, and more. + +Model Records in WebLLM +----------------------- + +Each of the model in `WebLLM Chat `__ is registered as an instance of +``ModelRecord`` and can be accessed at +`webllm.prebuiltAppConfig.model_list `__. + +Looking at the most straightforward example `get-started `__, +there are two ways to run a model. + +One can either use the prebuilt model by simply calling ``reload()`` with the ``model_id``: + +.. code:: typescript + + const selectedModel = "Llama-3-8B-Instruct-q4f32_1-MLC"; + const engine = await webllm.CreateMLCEngine(selectedModel); + +Or one can specify their own model to run by creating a model record: + +.. code:: typescript + + const appConfig: webllm.AppConfig = { + model_list: [ + { + model: "https://huggingface.co/mlc-ai/Llama-3-8B-Instruct-q4f32_1-MLC", + model_id: "Llama-3-8B-Instruct-q4f32_1-MLC", + model_lib: + webllm.modelLibURLPrefix + + webllm.modelVersion + + "/Llama-3-8B-Instruct-q4f32_1-ctx4k_cs1k-webgpu.wasm", + }, + // Add your own models here... + ], + }; + const selectedModel = "Llama-3-8B-Instruct-q4f32_1-MLC"; + const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine( + selectedModel, + { appConfig: appConfig }, + ); + +Looking at the code above, we find that, just like any other platforms supported by MLC-LLM, to +run a model on WebLLM, you need: + +1. **Model weights** converted to MLC format (e.g. `Llama-3-8B-Instruct-q4f32_1-MLC + `_.): downloaded through the url ``ModelRecord.model`` +2. **Model library** that comprises the inference logic (see repo `binary-mlc-llm-libs `__): downloaded through the url ``ModelRecord.model_lib``. + +In sections below, we walk you through two examples on how to add your own model besides the ones in +`webllm.prebuiltAppConfig.model_list `__. +Before proceeding, please verify installation of ``mlc_llm`` and ``tvm``. + +Verify Installation for Adding Models +------------------------------------- + +**Step 1. Verify mlc_llm** + +We use the python package ``mlc_llm`` to compile models. This can be installed by +following :ref:`install-mlc-packages`, either by building from source, or by +installing the prebuilt package. Verify ``mlc_llm`` installation in command line via: + +.. code:: bash + + $ mlc_llm --help + # You should see help information with this line + usage: MLC LLM Command Line Interface. [-h] {compile,convert_weight,gen_config} + +.. note:: + If it runs into error ``command not found: mlc_llm``, try ``python -m mlc_llm --help``. + +**Step 2. Verify TVM** + +To compile models, you also need to follow :ref:`install-tvm`. +Here we verify ``tvm`` quickly with command line (for full verification, see :ref:`tvm-validate`): + +.. code:: bash + + $ python -c "import tvm; print(tvm.__file__)" + /some-path/lib/python3.13/site-packages/tvm/__init__.py + + +.. _webllm-add-model-variant: + +Bring Your Own Model Variant +---------------------------- + +In cases where the model you are adding is simply a variant of an existing +model, we only need to convert weights and reuse existing model library. For instance: + +- Adding ``OpenMistral`` when MLC supports ``Mistral`` +- Adding a ``Llama3`` fine-tuned on a domain-specific task when MLC supports ``Llama3`` + + +In this section, we walk you through adding ``WizardMath-7B-V1.1-q4f16_1`` to the +`get-started `__ example. +According to the model's ``config.json`` on `its Huggingface repo `_, +it reuses the Mistral model architecture. + +.. note:: + + This section largely replicates :ref:`convert-weights-via-MLC`. + See that page for more details. Note that the weights are shared across + all platforms in MLC. + +**Step 1 Clone from HF and convert_weight** + +You can be under the mlc-llm repo, or your own working directory. Note that all platforms +can share the same compiled/quantized weights. See :ref:`compile-command-specification` +for specification of ``convert_weight``. + +.. code:: shell + + # Create directory + mkdir -p dist/models && cd dist/models + # Clone HF weights + git lfs install + git clone https://huggingface.co/WizardLM/WizardMath-7B-V1.1 + cd ../.. + # Convert weight + mlc_llm convert_weight ./dist/models/WizardMath-7B-V1.1/ \ + --quantization q4f16_1 \ + -o dist/WizardMath-7B-V1.1-q4f16_1-MLC + +**Step 2 Generate MLC Chat Config** + +Use ``mlc_llm gen_config`` to generate ``mlc-chat-config.json`` and process tokenizers. +See :ref:`compile-command-specification` for specification of ``gen_config``. + +.. code:: shell + + mlc_llm gen_config ./dist/models/WizardMath-7B-V1.1/ \ + --quantization q4f16_1 --conv-template wizard_coder_or_math \ + -o dist/WizardMath-7B-V1.1-q4f16_1-MLC/ + +For the ``conv-template``, `conversation_template.py `__ +contains a full list of conversation templates that MLC provides. You can also manually modify the ``mlc-chat-config.json`` to +add your customized conversation template. + +**Step 3 Upload weights to HF** + +.. code:: shell + + # First, please create a repository on Hugging Face. + # With the repository created, run + git lfs install + git clone https://huggingface.co/my-huggingface-account/my-wizardMath-weight-huggingface-repo + cd my-wizardMath-weight-huggingface-repo + cp path/to/mlc-llm/dist/WizardMath-7B-V1.1-q4f16_1-MLC/* . + git add . && git commit -m "Add wizardMath model weights" + git push origin main + +After successfully following all steps, you should end up with a Huggingface repo similar to +`WizardMath-7B-V1.1-q4f16_1-MLC `__, +which includes the converted/quantized weights, the ``mlc-chat-config.json``, and tokenizer files. + + +**Step 4 Register as a ModelRecord** + +Finally, we modify the code snippet for +`get-started `__ +pasted above. + +We simply specify the Huggingface link as ``model``, while reusing the ``model_lib`` for +``Mistral-7B``. + +.. code:: typescript + + const appConfig: webllm.AppConfig = { + model_list: [ + { + model: "https://huggingface.co/mlc-ai/WizardMath-7B-V1.1-q4f16_1-MLC", + model_id: "WizardMath-7B-V1.1-q4f16_1-MLC", + model_lib: + webllm.modelLibURLPrefix + + webllm.modelVersion + + "/Mistral-7B-Instruct-v0.3-q4f16_1-ctx4k_cs1k-webgpu.wasm", + }, + // Add your own models here... + ], + }; + + const selectedModel = "WizardMath-7B-V1.1-q4f16_1" + const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine( + selectedModel, + { appConfig: appConfig }, + ); + +Now, running the ``get-started`` example will use the ``WizardMath`` model you just added. +See `get-started's README `__ +on how to run it. + + +Bring Your Own Model Library +---------------------------- + +A model library is specified by: + + - The model architecture (e.g. ``llama-3``, ``gpt-neox``, ``phi-3``) + - Quantization (e.g. ``q4f16_1``, ``q0f32``) + - Metadata (e.g. ``context_window_size``, ``sliding_window_size``, ``prefill-chunk-size``), which affects memory planning (currently only ``prefill-chunk-size`` affects the compiled model) + - Platform (e.g. ``cuda``, ``webgpu``, ``iOS``) + +In cases where the model you want to run is not compatible with the provided MLC +prebuilt model libraries (e.g. having a different quantization, a different +metadata spec, or even a different model architecture), you need to build your +own model library. + +In this section, we walk you through adding ``RedPajama-INCITE-Chat-3B-v1`` to the +`get-started `__ example. + +This section largely replicates :ref:`compile-model-libraries`. See that page for +more details, specifically the ``WebGPU`` option. + +**Step 0. Install dependencies** + +To compile model libraries for webgpu, you need to :ref:`build mlc_llm from source `. +Besides, you also need to follow :ref:`install-web-build`. Otherwise, it would run into error: + +.. code:: text + + RuntimeError: Cannot find libraries: wasm_runtime.bc + +**Step 1. Clone from HF and convert_weight** + +You can be under the mlc-llm repo, or your own working directory. Note that all platforms +can share the same compiled/quantized weights. + +.. code:: shell + + # Create directory + mkdir -p dist/models && cd dist/models + # Clone HF weights + git lfs install + git clone https://huggingface.co/togethercomputer/RedPajama-INCITE-Chat-3B-v1 + cd ../.. + # Convert weight + mlc_llm convert_weight ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + +**Step 2. Generate mlc-chat-config and compile** + +A model library is specified by: + + - The model architecture (e.g. ``llama-2``, ``gpt-neox``) + - Quantization (e.g. ``q4f16_1``, ``q0f32``) + - Metadata (e.g. ``context_window_size``, ``sliding_window_size``, ``prefill-chunk-size``), which affects memory planning + - Platform (e.g. ``cuda``, ``webgpu``, ``iOS``) + +All these knobs are specified in ``mlc-chat-config.json`` generated by ``gen_config``. + +.. code:: shell + + # 1. gen_config: generate mlc-chat-config.json and process tokenizers + mlc_llm gen_config ./dist/models/RedPajama-INCITE-Chat-3B-v1/ \ + --quantization q4f16_1 --conv-template redpajama_chat \ + -o dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/ + # 2. compile: compile model library with specification in mlc-chat-config.json + mlc_llm compile ./dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC/mlc-chat-config.json \ + --device webgpu -o dist/libs/RedPajama-INCITE-Chat-3B-v1-q4f16_1-webgpu.wasm + +.. note:: + When compiling larger models like ``Llama-3-8B``, you may want to add ``--prefill_chunk_size 1024`` + to decrease memory usage. Otherwise, during runtime, you may run into issues like: + + .. code:: text + + TypeError: Failed to execute 'createBuffer' on 'GPUDevice': Failed to read the 'size' property from + 'GPUBufferDescriptor': Value is outside the 'unsigned long long' value range. + + +**Step 3. Distribute model library and model weights** + +After following the steps above, you should end up with: + +.. code:: shell + + ~/mlc-llm > ls dist/libs + RedPajama-INCITE-Chat-3B-v1-q4f16_1-webgpu.wasm # ===> the model library + + ~/mlc-llm > ls dist/RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC + mlc-chat-config.json # ===> the chat config + tensor-cache.json # ===> the model weight info + params_shard_0.bin # ===> the model weights + params_shard_1.bin + ... + tokenizer.json # ===> the tokenizer files + tokenizer_config.json + +Upload the ``RedPajama-INCITE-Chat-3B-v1-q4f16_1-webgpu.wasm`` to a github repository (for us, +it is in `binary-mlc-llm-libs `__). Then +upload the ``RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC`` to a Huggingface repo: + +.. code:: shell + + # First, please create a repository on Hugging Face. + # With the repository created, run + git lfs install + git clone https://huggingface.co/my-huggingface-account/my-redpajama3b-weight-huggingface-repo + cd my-redpajama3b-weight-huggingface-repo + cp path/to/mlc-llm/dist/RedPajama-INCITE-Instruct-3B-v1-q4f16_1-MLC/* . + git add . && git commit -m "Add redpajama-3b instruct model weights" + git push origin main + +This would result in something like `RedPajama-INCITE-Chat-3B-v1-q4f16_1-MLC +`_. + +**Step 4. Register as a ModelRecord** + +Finally, we are able to run the model we added in WebLLM's `get-started `__: + +.. code:: typescript + + const myAppConfig: AppConfig = { + model_list: [ + // Other records here omitted... + { + "model": "https://huggingface.co/my-hf-account/my-redpajama3b-weight-huggingface-repo/resolve/main/", + "model_id": "RedPajama-INCITE-Instruct-3B-v1", + "model_lib": "https://raw.githubusercontent.com/my-gh-account/my-repo/main/RedPajama-INCITE-Chat-3B-v1-q4f16_1-webgpu.wasm", + "required_features": ["shader-f16"], + }, + ] + } + + const selectedModel = "RedPajama-INCITE-Instruct-3B-v1"; + const engine: webllm.MLCEngineInterface = await webllm.CreateMLCEngine( + selectedModel, + { appConfig: appConfig }, + ); + +Now, running the ``get-started`` example will use the ``RedPajama`` model you just added. +See `get-started's README `__ +on how to run it. diff --git a/docs/get_started/introduction.rst b/docs/get_started/introduction.rst new file mode 100644 index 0000000..1123596 --- /dev/null +++ b/docs/get_started/introduction.rst @@ -0,0 +1,358 @@ +.. _introduction-to-mlc-llm: + +Introduction to MLC LLM +======================= + +.. contents:: Table of Contents + :local: + :depth: 2 + +MLC LLM is a machine learning compiler and high-performance deployment +engine for large language models. The mission of this project is to enable everyone to develop, +optimize, and deploy AI models natively on everyone's platforms.  + +This page is a quick tutorial to introduce how to try out MLC LLM, and the steps to +deploy your own models with MLC LLM. + +Installation +------------ + +:ref:`MLC LLM ` is available via pip. +It is always recommended to install it in an isolated conda virtual environment. + +To verify the installation, activate your virtual environment, run + +.. code:: bash + + python -c "import mlc_llm; print(mlc_llm.__path__)" + +You are expected to see the installation path of MLC LLM Python package. + + +Chat CLI +-------- + +As the first example, we try out the chat CLI in MLC LLM with 4-bit quantized 8B Llama-3 model. +You can run MLC chat through a one-liner command: + +.. code:: bash + + mlc_llm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC + +It may take 1-2 minutes for the first time running this command. +After waiting, this command launch a chat interface where you can enter your prompt and chat with the model. + +.. code:: + + You can use the following special commands: + /help print the special commands + /exit quit the cli + /stats print out the latest stats (token/sec) + /reset restart a fresh chat + /set [overrides] override settings in the generation config. For example, + `/set temperature=0.5;max_gen_len=100;stop=end,stop` + Note: Separate stop words in the `stop` option with commas (,). + Multi-line input: Use escape+enter to start a new line. + + user: What's the meaning of life + assistant: + What a profound and intriguing question! While there's no one definitive answer, I'd be happy to help you explore some perspectives on the meaning of life. + + The concept of the meaning of life has been debated and... + + +The figure below shows what run under the hood of this chat CLI command. +For the first time running the command, there are three major phases. + +- **Phase 1. Pre-quantized weight download.** This phase automatically downloads pre-quantized Llama-3 model from `Hugging Face `_ and saves it to your local cache directory. +- **Phase 2. Model compilation.** This phase automatically optimizes the Llama-3 model to accelerate model inference on GPU with techniques of machine learning compilation in `Apache TVM `_ compiler, and generate the binary model library that enables the execution language models on your local GPU. +- **Phase 3. Chat runtime.** This phase consumes the model library built in phase 2 and the model weights downloaded in phase 1, launches a platform-native chat runtime to drive the execution of Llama-3 model. + +We cache the pre-quantized model weights and compiled model library locally. +Therefore, phase 1 and 2 will only execute **once** over multiple runs. + +.. figure:: /_static/img/project-workflow.svg + :width: 700 + :align: center + :alt: Project Workflow + + Workflow in MLC LLM + +.. note:: + + If you want to enable tensor parallelism to run LLMs on multiple GPUs, + please specify argument ``--overrides "tensor_parallel_shards=$NGPU"``. + For example, + + .. code:: shell + + mlc_llm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC --overrides "tensor_parallel_shards=2" + +.. _introduction-to-mlc-llm-python-api: + +Python API +---------- + +In the second example, we run the Llama-3 model with the chat completion Python API of MLC LLM. +You can save the code below into a Python file and run it. + +.. code:: python + + from mlc_llm import MLCEngine + + # Create engine + model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" + engine = MLCEngine(model) + + # Run chat completion in OpenAI API. + for response in engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=True, + ): + for choice in response.choices: + print(choice.delta.content, end="", flush=True) + print("\n") + + engine.terminate() + +.. figure:: https://raw.githubusercontent.com/mlc-ai/web-data/main/images/mlc-llm/tutorials/python-engine-api.jpg + :width: 500 + :align: center + + MLC LLM Python API + +This code example first creates an :class:`mlc_llm.MLCEngine` instance with the 4-bit quantized Llama-3 model. +**We design the Python API** :class:`mlc_llm.MLCEngine` **to align with OpenAI API**, +which means you can use :class:`mlc_llm.MLCEngine` in the same way of using +`OpenAI's Python package `_ +for both synchronous and asynchronous generation. + +In this code example, we use the synchronous chat completion interface and iterate over +all the stream responses. +If you want to run without streaming, you can run + +.. code:: python + + response = engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=False, + ) + print(response) + +You can also try different arguments supported in `OpenAI chat completion API `_. +If you would like to do concurrent asynchronous generation, you can use :class:`mlc_llm.AsyncMLCEngine` instead. + +.. note:: + + If you want to enable tensor parallelism to run LLMs on multiple GPUs, + please specify argument ``model_config_overrides`` in MLCEngine constructor. + For example, + + .. code:: python + + from mlc_llm import MLCEngine + from mlc_llm.serve.config import EngineConfig + + model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" + engine = MLCEngine( + model, + engine_config=EngineConfig(tensor_parallel_shards=2), + ) + + +REST Server +----------- + +For the third example, we launch a REST server to serve the 4-bit quantized Llama-3 model +for OpenAI chat completion requests. The server can be launched in command line with + +.. code:: bash + + mlc_llm serve HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC + +The server is hooked at ``http://127.0.0.1:8000`` by default, and you can use ``--host`` and ``--port`` +to set a different host and port. +When the server is ready (showing ``INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)``), +we can open a new shell and send a cURL request via the following command: + +.. code:: bash + + curl -X POST \ + -H "Content-Type: application/json" \ + -d '{ + "model": "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC", + "messages": [ + {"role": "user", "content": "Hello! Our project is MLC LLM. What is the name of our project?"} + ] + }' \ + http://127.0.0.1:8000/v1/chat/completions + +The server will process this request and send back the response. +Similar to :ref:`introduction-to-mlc-llm-python-api`, you can pass argument ``"stream": true`` +to request for stream responses. + +.. note:: + + If you want to enable tensor parallelism to run LLMs on multiple GPUs, + please specify argument ``--overrides "tensor_parallel_shards=$NGPU"``. + For example, + + .. code:: shell + + mlc_llm serve HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC --overrides "tensor_parallel_shards=2" + +.. _introduction-deploy-your-own-model: + +Deploy Your Own Model +--------------------- + +So far we have been using pre-converted models weights from Hugging Face. +This section introduces the core workflow regarding how you can *run your own models with MLC LLM*. + +We use the `Phi-2 `_ as the example model. +Assuming the Phi-2 model is downloaded and placed under ``models/phi-2``, +there are two major steps to prepare your own models. + +- **Step 1. Generate MLC config.** The first step is to generate the configuration file of MLC LLM. + + .. code:: bash + + export LOCAL_MODEL_PATH=models/phi-2 # The path where the model resides locally. + export MLC_MODEL_PATH=dist/phi-2-MLC/ # The path where to place the model processed by MLC. + export QUANTIZATION=q0f16 # The choice of quantization. + export CONV_TEMPLATE=phi-2 # The choice of conversation template. + mlc_llm gen_config $LOCAL_MODEL_PATH \ + --quantization $QUANTIZATION \ + --conv-template $CONV_TEMPLATE \ + -o $MLC_MODEL_PATH + + The config generation command takes in the local model path, the target path of MLC output, + the conversation template name in MLC and the quantization name in MLC. + Here the quantization ``q0f16`` means float16 without quantization, + and the conversation template ``phi-2`` is the Phi-2 model's template in MLC. + + If you want to enable tensor parallelism on multiple GPUs, add argument + ``--tensor-parallel-shards $NGPU`` to the config generation command. + + - `The full list of supported quantization in MLC `_. You can try different quantization methods with MLC LLM. Typical quantization methods are ``q4f16_1`` for 4-bit group quantization, ``q4f16_ft`` for 4-bit FasterTransformer format quantization. + - `The full list of conversation template in MLC `_. + +- **Step 2. Convert model weights.** In this step, we convert the model weights to MLC format. + + .. code:: bash + + mlc_llm convert_weight $LOCAL_MODEL_PATH \ + --quantization $QUANTIZATION \ + -o $MLC_MODEL_PATH + + This step consumes the raw model weights and converts them to for MLC format. + The converted weights will be stored under ``$MLC_MODEL_PATH``, + which is the same directory where the config file generated in Step 1 resides. + +Now, we can try to run your own model with chat CLI: + +.. code:: bash + + mlc_llm chat $MLC_MODEL_PATH + +For the first run, model compilation will be triggered automatically to optimize the +model for GPU accelerate and generate the binary model library. +The chat interface will be displayed after model JIT compilation finishes. +You can also use this model in Python API, MLC serve and other use scenarios. + +(Optional) Compile Model Library +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +In previous sections, model libraries are compiled when the :class:`mlc_llm.MLCEngine` launches, +which is what we call "JIT (Just-in-Time) model compilation". +In some cases, it is beneficial to explicitly compile the model libraries. +We can deploy LLMs with reduced dependencies by shipping the library for deployment without going through compilation. +It will also enable advanced options such as cross-compiling the libraries for web and mobile deployments. + + +Below is an example command of compiling model libraries in MLC LLM: + +.. code:: bash + + export MODEL_LIB=$MLC_MODEL_PATH/lib.so # ".dylib" for Intel Macs. + # ".dll" for Windows. + # ".wasm" for web. + # ".tar" for iPhone/Android. + mlc_llm compile $MLC_MODEL_PATH -o $MODEL_LIB + +At runtime, we need to specify this model library path to use it. For example, + +.. code:: bash + + # For chat CLI + mlc_llm chat $MLC_MODEL_PATH --model-lib $MODEL_LIB + # For REST server + mlc_llm serve $MLC_MODEL_PATH --model-lib $MODEL_LIB + +.. code:: python + + from mlc_llm import MLCEngine + + # For Python API + model = "models/phi-2" + model_lib = "models/phi-2/lib.so" + engine = MLCEngine(model, model_lib=model_lib) + +:ref:`compile-model-libraries` introduces the model compilation command in detail, +where you can find instructions and example commands to compile model to different +hardware backends, such as WebGPU, iOS and Android. + +Universal Deployment +-------------------- + +MLC LLM is a high-performance universal deployment solution for large language models, +to enable native deployment of any large language models with native APIs with compiler acceleration +So far, we have gone through several examples running on a local GPU environment. +The project supports multiple kinds of GPU backends. + +You can use `--device` option in compilation and runtime to pick a specific GPU backend. +For example, if you have an NVIDIA or AMD GPU, you can try to use the option below +to run chat through the vulkan backend. Vulkan-based LLM applications run in less typical +environments (e.g. SteamDeck). + +.. code:: bash + + mlc_llm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC --device vulkan + +The same core LLM runtime engine powers all the backends, enabling the same model to be deployed across backends as +long as they fit within the memory and computing budget of the corresponding hardware backend. +We also leverage machine learning compilation to build backend-specialized optimizations to +get out the best performance on the targetted backend when possible, and reuse key insights and optimizations +across backends we support. + +Please checkout the what to do next sections below to find out more about different deployment scenarios, +such as WebGPU-based browser deployment, mobile and other settings. + +Summary and What to Do Next +--------------------------- + +To briefly summarize this page, + +- We went through three examples (chat CLI, Python API, and REST server) of MLC LLM, +- we introduced how to convert model weights for your own models to run with MLC LLM, and (optionally) how to compile your models. +- We also discussed the universal deployment capability of MLC LLM. + +Next, please feel free to check out the pages below for quick start examples and more detailed information +on specific platforms + +- :ref:`Quick start examples ` for Python API, chat CLI, REST server, web browser, iOS and Android. +- Depending on your use case, check out our API documentation and tutorial pages: + + - :ref:`webllm-runtime` + - :ref:`deploy-rest-api` + - :ref:`deploy-cli` + - :ref:`deploy-python-engine` + - :ref:`deploy-ios` + - :ref:`deploy-android` + - :ref:`deploy-ide-integration` + +- :ref:`Convert model weight to MLC format `, if you want to run your own models. +- :ref:`Compile model libraries `, if you want to deploy to web/iOS/Android or control the model optimizations. +- Report any problem or ask any question: open new issues in our `GitHub repo `_. diff --git a/docs/get_started/quick_start.rst b/docs/get_started/quick_start.rst new file mode 100644 index 0000000..445106b --- /dev/null +++ b/docs/get_started/quick_start.rst @@ -0,0 +1,190 @@ +.. _quick-start: + +Quick Start +=========== + +Examples +-------- + +To begin with, try out MLC LLM support for int4-quantized Llama3 8B. +It is recommended to have at least 6GB free VRAM to run it. + +.. tabs:: + + .. tab:: Python + + **Install MLC LLM**. :ref:`MLC LLM ` is available via pip. + It is always recommended to install it in an isolated conda virtual environment. + + **Run chat completion in Python.** The following Python script showcases the Python API of MLC LLM: + + .. code:: python + + from mlc_llm import MLCEngine + + # Create engine + model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" + engine = MLCEngine(model) + + # Run chat completion in OpenAI API. + for response in engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=True, + ): + for choice in response.choices: + print(choice.delta.content, end="", flush=True) + print("\n") + + engine.terminate() + + .. Todo: link the colab notebook when ready: + + **Documentation and tutorial.** Python API reference and its tutorials are :ref:`available online `. + + .. figure:: https://raw.githubusercontent.com/mlc-ai/web-data/main/images/mlc-llm/tutorials/python-engine-api.jpg + :width: 600 + :align: center + + MLC LLM Python API + + .. tab:: REST Server + + **Install MLC LLM**. :ref:`MLC LLM ` is available via pip. + It is always recommended to install it in an isolated conda virtual environment. + + **Launch a REST server.** Run the following command from command line to launch a REST server at ``http://127.0.0.1:8000``. + + .. code:: shell + + mlc_llm serve HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC + + **Send requests to server.** When the server is ready (showing ``INFO: Uvicorn running on http://127.0.0.1:8000 (Press CTRL+C to quit)``), + open a new shell and send a request via the following command: + + .. code:: shell + + curl -X POST \ + -H "Content-Type: application/json" \ + -d '{ + "model": "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC", + "messages": [ + {"role": "user", "content": "Hello! Our project is MLC LLM. What is the name of our project?"} + ] + }' \ + http://127.0.0.1:8000/v1/chat/completions + + **Documentation and tutorial.** Check out :ref:`deploy-rest-api` for the REST API reference and tutorial. + Our REST API has complete OpenAI API support. + + .. figure:: https://raw.githubusercontent.com/mlc-ai/web-data/main/images/mlc-llm/tutorials/python-serve-request.jpg + :width: 600 + :align: center + + Send HTTP request to REST server in MLC LLM + + .. tab:: Command Line + + **Install MLC LLM**. :ref:`MLC LLM ` is available via pip. + It is always recommended to install it in an isolated conda virtual environment. + + For Windows/Linux users, make sure to have latest :ref:`Vulkan driver ` installed. + + **Run in command line**. + + .. code:: bash + + mlc_llm chat HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC + + + If you are using windows/linux/steamdeck and would like to use vulkan, + we recommend installing necessary vulkan loader dependency via conda + to avoid vulkan not found issues. + + .. code:: bash + + conda install -c conda-forge gcc libvulkan-loader + + + .. tab:: Web Browser + + `WebLLM `__. MLC LLM generates performant code for WebGPU and WebAssembly, + so that LLMs can be run locally in a web browser without server resources. + + **Download pre-quantized weights**. This step is self-contained in WebLLM. + + **Download pre-compiled model library**. WebLLM automatically downloads WebGPU code to execute. + + **Check browser compatibility**. The latest Google Chrome provides WebGPU runtime and `WebGPU Report `__ as a useful tool to verify WebGPU capabilities of your browser. + + .. figure:: https://blog.mlc.ai/img/redpajama/web.gif + :width: 300 + :align: center + + MLC LLM on Web + + .. tab:: iOS + + **Install MLC Chat iOS**. It is available on AppStore: + + .. image:: https://developer.apple.com/assets/elements/badges/download-on-the-app-store.svg + :width: 135 + :target: https://apps.apple.com/us/app/mlc-chat/id6448482937 + + | + + **Note**. The larger model might take more VRAM, try start with smaller models first. + + **Tutorial and source code**. The source code of the iOS app is fully `open source `__, + and a :ref:`tutorial ` is included in documentation. + + .. figure:: https://blog.mlc.ai/img/redpajama/ios.gif + :width: 300 + :align: center + + MLC Chat on iOS + + .. tab:: Android + + **Install MLC Chat Android**. A prebuilt is available as an APK: + + .. image:: https://seeklogo.com/images/D/download-android-apk-badge-logo-D074C6882B-seeklogo.com.png + :width: 135 + :target: https://github.com/mlc-ai/binary-mlc-llm-libs/releases/download/Android-09262024/mlc-chat.apk + + | + + **Note**. The larger model might take more VRAM, try start with smaller models first. + The demo is tested on + + - Samsung S23 with Snapdragon 8 Gen 2 chip + - Redmi Note 12 Pro with Snapdragon 685 + - Google Pixel phones + + **Tutorial and source code**. The source code of the android app is fully `open source `__, + and a :ref:`tutorial ` is included in documentation. + + .. figure:: https://blog.mlc.ai/img/android/android-recording.gif + :width: 300 + :align: center + + MLC LLM on Android + + +What to Do Next +--------------- + +- Check out :ref:`introduction-to-mlc-llm` for the introduction of a complete workflow in MLC LLM. +- Depending on your use case, check out our API documentation and tutorial pages: + + - :ref:`webllm-runtime` + - :ref:`deploy-rest-api` + - :ref:`deploy-cli` + - :ref:`deploy-python-engine` + - :ref:`deploy-ios` + - :ref:`deploy-android` + - :ref:`deploy-ide-integration` + +- :ref:`convert-weights-via-MLC`, if you want to run your own models. +- :ref:`compile-model-libraries`, if you want to deploy to web/iOS/Android or control the model optimizations. +- Report any problem or ask any question: open new issues in our `GitHub repo `_. diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 0000000..28dfa92 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,89 @@ +👋 Welcome to MLC LLM +===================== + +`Discord `_ | `GitHub `_ + + + + +MLC LLM is a machine learning compiler and high-performance deployment +engine for large language models. The mission of this project is to enable +everyone to develop, optimize, and deploy AI models natively on everyone's platforms.  + +Quick Start +----------- + +Check out :ref:`quick-start` for quick start examples of using MLC LLM. + +Introduction to MLC LLM +----------------------- + +Check out :ref:`introduction-to-mlc-llm` for the introduction and tutorial of a complete workflow in MLC LLM. + + +.. toctree:: + :maxdepth: 1 + :caption: Get Started + :hidden: + + get_started/quick_start.rst + get_started/introduction.rst + +.. toctree:: + :maxdepth: 1 + :caption: Build and Deploy Apps + :hidden: + + deploy/webllm.rst + deploy/rest.rst + deploy/cli.rst + deploy/python_engine.rst + deploy/ios.rst + deploy/android.rst + deploy/ide_integration.rst + deploy/mlc_chat_config.rst + +.. toctree:: + :maxdepth: 1 + :caption: Compile Models + :hidden: + + compilation/convert_weights.rst + compilation/compile_models.rst + compilation/package_libraries_and_weights.rst + compilation/define_new_models.rst + compilation/configure_quantization.rst + +.. toctree:: + :maxdepth: 1 + :caption: Dependency Installation + :hidden: + + install/tvm.rst + install/mlc_llm.rst + install/conda.rst + install/gpu.rst + install/emcc.rst + +.. toctree:: + :maxdepth: 1 + :caption: Microserving API + :hidden: + + microserving/tutorial.rst + +.. toctree:: + :maxdepth: 1 + :caption: Community + :hidden: + + community/guideline.rst + community/faq.rst + + +.. toctree:: + :maxdepth: 1 + :caption: Privacy + :hidden: + + privacy.rst diff --git a/docs/install/conda.rst b/docs/install/conda.rst new file mode 100644 index 0000000..305c0f6 --- /dev/null +++ b/docs/install/conda.rst @@ -0,0 +1,66 @@ +Install Conda +============= + +MLC LLM does not depend on, but generally recommends conda as a generic dependency manager, primarily because it creates unified cross-platform experience to make windows/Linux/macOS development equally easy. Moreover, conda is python-friendly and provides all the python packages needed for MLC LLM, such as numpy. + +.. contents:: Table of Contents + :depth: 2 + + +Install Miniconda +----------------- + +**Use installer.** Miniconda, a minimal distribution of conda, comes with out-of-box installer across Windows/macOS/Linux. Please refer to its `official website `_ link for detailed instructions. + +**Set libmamba as the dependency solver.** The default dependency solver in conda could be slow in certain scenarios, and it is always recommended to upgrade it to libmamba, a faster solver. + +.. code-block:: bash + :caption: Set libmamba as the default solver + + # update conda + conda update --yes -n base -c defaults conda + # install `conda-libmamba-solver` + conda install --yes -n base conda-libmamba-solver + # set it as the default solver + conda config --set solver libmamba + +.. note:: + Conda is a generic dependency manager, which is not necessarily related to any Python distributions. + In fact, some of our tutorials recommends to use conda to install cmake, git and rust for its unified experience across OS platforms. + + +Validate installation +--------------------- + +**Step 1. Check conda-arch mismatch.** Nowadays macOS runs on two different architectures: arm64 and x86_64, which could particularly lead to many misuses in MLC LLM, where the error message hints about "architecture mismatch". Use the following command to make sure particular conda architecture is installed accordingly: + +.. code-block:: bash + :caption: Check conda architecture + + >>> conda info | grep platform + # for arm mac + platform : osx-arm64 + # for x86 mac + platform : osx-64 + +**Step 2. Check conda virtual environment.** If you have installed python in your conda virtual environment, make sure conda, Python and pip are all from this environment: + +.. code-block:: bash + :caption: Check conda virtual environment (macOS, Linux) + + >>> echo $CONDA_PREFIX + /.../miniconda3/envs/mlc-doc-venv + >>> which python + /.../miniconda3/envs/mlc-doc-venv/bin/python + >>> which pip + /.../miniconda3/envs/mlc-doc-venv/bin/pip + +.. code-block:: bat + :caption: Check conda virtual environment (Windows) + + >>> echo $Env:CONDA_PREFIX + \...\miniconda3\envs\mlc-doc-venv + >>> Get-Command python.exe + \...\miniconda3\envs\mlc-doc-venv\bin\python.exe + >>> Get-Command pip.exe + \...\miniconda3\envs\mlc-doc-venv\bin\pip.exe diff --git a/docs/install/emcc.rst b/docs/install/emcc.rst new file mode 100644 index 0000000..b08ddcb --- /dev/null +++ b/docs/install/emcc.rst @@ -0,0 +1,80 @@ +.. _install-web-build: + +Install Wasm Build Environment +============================== + +This page describes the steps to setup build environment for WebAssembly and WebGPU builds. + +Step 1: Install EMSDK +--------------------- + +Emscripten is an LLVM-based compiler that compiles C/C++ source code to WebAssembly. +We need to install emscripten for webgpu build. + +- Please follow the installation instruction `here `__ + to install the latest emsdk. +- Source path/to/emsdk_env.sh so emcc is reachable from PATH and the command emcc works. + +Validate that emcc is accessible in shell + +.. code:: bash + + emcc --version + +.. note:: + We recently found that using the latest ``emcc`` version may run into issues during runtime. Use + ``./emsdk install 3.1.56`` instead of ``./emsdk install latest`` for now as a workaround. + + The error may look like + + .. code:: text + + Init error, LinkError: WebAssembly.instantiate(): Import #6 module="wasi_snapshot_preview1" + function="proc_exit": function import requires a callable + + +Step 2: Set TVM_SOURCE_DIR and MLC_LLM_SOURCE_DIR +------------------------------------------------- + +We need to set a path to a tvm source in order to build tvm runtime. +Note that you do not need to build TVM from the source. The source here is only used to build the web runtime component. +Set environment variable in your shell startup profile in to point to ``3rdparty/tvm`` (if preferred, you could also +point to your own TVM address if you installed TVM from source). + +Besides, we also need to set ``MLC_LLM_SOURCE_DIR`` so that we can locate ``mlc_wasm_runtime.bc`` when compiling a model library wasm. + +.. code:: bash + + export TVM_SOURCE_DIR=/path/to/3rdparty/tvm + export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm + + +Step 3: Prepare Wasm Runtime +---------------------------- + +First, we need to obtain a copy of the mlc-llm source code for the setup script + +.. code:: bash + + git clone https://github.com/mlc-ai/mlc-llm.git --recursive + cd mlc-llm + +Now we can prepare wasm runtime using the script in mlc-llm repo + +.. code:: bash + + ./web/prep_emcc_deps.sh + +We can then validate the outcome + +.. code:: bash + + >>> echo ${TVM_SOURCE_DIR} + + /path/set/in/step2 + + >>> ls -l ${TVM_SOURCE_DIR}/web/dist/wasm/*.bc + + tvmjs_support.bc + wasm_runtime.bc + webgpu_runtime.bc diff --git a/docs/install/gpu.rst b/docs/install/gpu.rst new file mode 100644 index 0000000..5c21558 --- /dev/null +++ b/docs/install/gpu.rst @@ -0,0 +1,201 @@ +GPU Drivers and SDKs +==================== + +.. contents:: Table of Contents + :depth: 2 + +MLC LLM is a universal deployment solution that allows efficient CPU/GPU code generation without AutoTVM-based performance tuning. This section focuses on generic GPU environment setup and troubleshooting. + +CUDA +---- + +CUDA is required to compile and run models with CUDA backend. + +Installation +^^^^^^^^^^^^ + +If you have a NVIDIA GPU and you want to use models compiled with CUDA +backend, you should install CUDA, which can be downloaded from +`here `__. + +Validate Installation +^^^^^^^^^^^^^^^^^^^^^ + +To verify you have correctly installed CUDA runtime and NVIDIA driver, run ``nvidia-smi`` in command line and see if you can get the GPU information. + +ROCm +---- + +ROCm is required to compile and run models with ROCm backend. + +Installation +^^^^^^^^^^^^ + +Right now MLC LLM only supports ROCm 6.1/6.2. +If you have AMD GPU and you want to use models compiled with ROCm +backend, you should install ROCm from `here `__. + +Validate Installation +^^^^^^^^^^^^^^^^^^^^^ + +To verify you have correctly installed ROCm, run ``rocm-smi`` in command line. +If you see the list of AMD devices printed out in a table, it means the ROCm is correctly installed. + +.. _vulkan_driver: + +Vulkan Driver +------------- + +Installation +^^^^^^^^^^^^ + +To run pre-trained models (e.g. pulled from MLC-AI's Hugging Face repository) compiled with Vulkan backend, you are expected to install Vulkan driver on your machine. + +Please check `this +page `__ and find the +Vulkan driver according to your GPU vendor. + +AMD Radeon and Radeon PRO +######################### + +For AMD Radeon and Radeon PRO users, please download AMD's drivers from official website (`Linux `__ / `Windows `__). +For Linux users, after you installed the ``amdgpu-install`` package, you can follow the instructions in its `documentation `__ to install +the driver. We recommend you installing ROCr OpenCL and PRO Vulkan (proprietary) for best performance, which can be done by running the following command: + +.. code:: bash + + amdgpu-install --usecase=graphics,opencl --opencl=rocr --vulkan=pro --no-32 + +Validate Installation +^^^^^^^^^^^^^^^^^^^^^ + +To verify whether Vulkan installation is successful or not, you are encouraged to install ``vulkaninfo``, below are the instructions to install ``vulkaninfo`` on different platforms: + +.. tabs :: + + .. code-tab :: bash Ubuntu/Debian + + sudo apt-get update + sudo apt-get install vulkan-tools + + .. code-tab :: bash Windows + + # It comes with your GPU driver + + .. code-tab :: bash Fedora + + sudo dnf install vulkan-tools + + .. code-tab :: bash Arch Linux + + sudo pacman -S vulkan-tools + # Arch Linux has maintained an awesome wiki page for Vulkan which you can refer to for troubleshooting: https://wiki.archlinux.org/title/Vulkan + + .. code-tab :: bash Other Distributions + + # Please install Vulkan SDK for your platform + # https://vulkan.lunarg.com/sdk/home + + +After installation, you can run ``vulkaninfo`` in command line and see if you can get the GPU information. + +.. note:: + WSL support for Windows is work-in-progress at the moment. Please do not use WSL on Windows to run Vulkan. + +Vulkan SDK +---------- + +Vulkan SDK is required for compiling models to Vulkan backend. To build TVM compiler from source, you will need to install Vulkan SDK as a dependency, but our :doc:`pre-built wheels <../install/mlc_llm>` already ships with Vulkan SDK. + +Check Vulkan SDK installation guide according to your platform: + +.. tabs :: + + .. tab :: Windows + + `Getting Started with the Windows Tarball Vulkan SDK `__ + + .. tab :: Linux + + For Ubuntu user, please check + `Getting Started with the Ubuntu Vulkan SDK `__ + + For other Linux distributions, please check + `Getting Started with the Linux Tarball Vulkan SDK `__ + + .. tab :: Mac + + `Getting Started with the macOS Vulkan SDK `__ + +Please refer to installation and setup page for next steps to build TVM from source. + +OpenCL SDK +---------- + +OpenCL SDK is only required when you want to build your own models for OpenCL backend. Please refer to `OpenCL's Github Repository `__ for installation guide of OpenCL-SDK. + +Orange Pi 5 (RK3588 based SBC) +------------------------------ + +OpenCL SDK and Mali GPU driver is required to compile and run models for OpenCL backend. + +Installation +^^^^^^^^^^^^ + +* Download and install the Ubuntu 22.04 for your board from `here `__ + +* Download and install ``libmali-g610.so`` + +.. code-block:: bash + + cd /usr/lib && sudo wget https://github.com/JeffyCN/mirrors/raw/libmali/lib/aarch64-linux-gnu/libmali-valhall-g610-g6p0-x11-wayland-gbm.so + +* Check if file ``mali_csffw.bin`` exist under path ``/lib/firmware``, if not download it with command: + +.. code-block:: bash + + cd /lib/firmware && sudo wget https://github.com/JeffyCN/mirrors/raw/libmali/firmware/g610/mali_csffw.bin + +* Download OpenCL ICD loader and manually add libmali to ICD + +.. code-block:: bash + + sudo apt update + sudo apt install mesa-opencl-icd + sudo mkdir -p /etc/OpenCL/vendors + echo "/usr/lib/libmali-valhall-g610-g6p0-x11-wayland-gbm.so" | sudo tee /etc/OpenCL/vendors/mali.icd + +* Download and install ``libOpenCL`` + +.. code-block:: bash + + sudo apt install ocl-icd-opencl-dev + +* Download and install dependencies for Mali OpenCL + +.. code-block:: bash + + sudo apt install libxcb-dri2-0 libxcb-dri3-0 libwayland-client0 libwayland-server0 libx11-xcb1 + +* Download and install clinfo to check if OpenCL successfully installed + +.. code-block:: bash + + sudo apt install clinfo + +Validate Installation +^^^^^^^^^^^^^^^^^^^^^ + +To verify you have correctly installed OpenCL runtime and Mali GPU driver, run ``clinfo`` in command line and see if you can get the GPU information. +You are expect to see the following information: + +.. code-block:: bash + + $ clinfo + arm_release_ver: g13p0-01eac0, rk_so_ver: 3 + Number of platforms 2 + Platform Name ARM Platform + Platform Vendor ARM + Platform Version OpenCL 2.1 v1.g6p0-01eac0.2819f9d4dbe0b5a2f89c835d8484f9cd + Platform Profile FULL_PROFILE + ... diff --git a/docs/install/mlc_llm.rst b/docs/install/mlc_llm.rst new file mode 100644 index 0000000..bcca6a2 --- /dev/null +++ b/docs/install/mlc_llm.rst @@ -0,0 +1,246 @@ +.. _install-mlc-packages: + +Install MLC LLM Python Package +============================== + +.. contents:: Table of Contents + :local: + :depth: 2 + +MLC LLM Python Package can be installed directly from a prebuilt developer package, or built from source. + +Option 1. Prebuilt Package +-------------------------- + +We provide nightly built pip wheels for MLC-LLM via pip. +Select your operating system/compute platform and run the command in your terminal: + +.. note:: + ❗ Whenever using Python, it is highly recommended to use **conda** to manage an isolated Python environment to avoid missing dependencies, incompatible versions, and package conflicts. + Please make sure your conda environment has Python and pip installed. + +.. tabs:: + + .. tab:: Linux + + .. tabs:: + + .. tab:: CPU + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cpu mlc-ai-nightly-cpu + + .. tab:: CUDA 12.8 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cu128 mlc-ai-nightly-cu128 + + .. tab:: CUDA 13.0 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cu130 mlc-ai-nightly-cu130 + + .. tab:: ROCm 6.1 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-rocm61 mlc-ai-nightly-rocm61 + + .. tab:: ROCm 6.2 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-rocm62 mlc-ai-nightly-rocm62 + + .. tab:: Vulkan + + Supported in all Linux packages. Checkout the following instructions + to install the latest vulkan loader to avoid vulkan not found issue. + + .. code-block:: bash + + conda install -c conda-forge gcc libvulkan-loader + + .. note:: + We need git-lfs in the system, you can install it via + + .. code-block:: bash + + conda install -c conda-forge git-lfs + + If encountering issues with GLIBC not found, please install the latest glibc in conda: + + .. code-block:: bash + + conda install -c conda-forge libstdcxx-ng + + Besides, we would recommend using Python 3.13; so if you are creating a new environment, + you could use the following command: + + .. code-block:: bash + + conda create --name mlc-prebuilt python=3.13 + + .. tab:: macOS + + .. tabs:: + + .. tab:: CPU + Metal + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cpu mlc-ai-nightly-cpu + + .. note:: + + Always check if conda is installed properly in macOS using the command below: + + .. code-block:: bash + + conda info | grep platform + + It should return "osx-64" for Mac with Intel chip, and "osx-arm64" for Mac with Apple chip. + We need git-lfs in the system, you can install it via + + .. code-block:: bash + + conda install -c conda-forge git-lfs + + .. tab:: Windows + + .. tabs:: + + .. tab:: CPU + Vulkan + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-llm-nightly-cpu mlc-ai-nightly-cpu + + .. note:: + Please make sure your conda environment comes with python and pip. + Make sure you also install the following packages, + vulkan loader, clang, git and git-lfs to enable proper automatic download + and jit compilation. + + .. code-block:: bash + + conda install -c conda-forge clang libvulkan-loader git-lfs git + + If encountering the error below: + + .. code-block:: bash + + FileNotFoundError: Could not find module 'path\to\site-packages\tvm\tvm.dll' (or one of its dependencies). Try using the full path with constructor syntax. + + It is likely `zstd`, a dependency to LLVM, was missing. Please use the command below to get it installed: + + .. code-block:: bash + + conda install zstd + + +Then you can verify installation in command line: + +.. code-block:: bash + + python -c "import mlc_llm; print(mlc_llm)" + # Prints out: + +| + +.. _mlcchat_build_from_source: + +Option 2. Build from Source +--------------------------- + +We also provide options to build mlc runtime libraries ``mlc_llm`` from source. +This step is useful when you want to make modification or obtain a specific version of mlc runtime. + + +**Step 1. Set up build dependency.** To build from source, you need to ensure that the following build dependencies are satisfied: + +* CMake >= 3.24 +* Git +* `Rust and Cargo `_, required by Hugging Face's tokenizer +* One of the GPU runtimes: + + * CUDA >= 11.8 (NVIDIA GPUs) + * Metal (Apple GPUs) + * Vulkan (NVIDIA, AMD, Intel GPUs) + +.. code-block:: bash + :caption: Set up build dependencies in Conda + + # make sure to start with a fresh environment + conda env remove -n mlc-chat-venv + # create the conda environment with build dependency + conda create -n mlc-chat-venv -c conda-forge \ + "cmake>=3.24" \ + rust \ + git \ + python=3.13 + # enter the build environment + conda activate mlc-chat-venv + +.. note:: + For runtime, :doc:`TVM ` compiler is not a dependency for MLCChat CLI or Python API. Only TVM's runtime is required, which is automatically included in `3rdparty/tvm `_. + However, if you would like to compile your own models, you need to follow :doc:`TVM `. + +**Step 2. Configure and build.** A standard git-based workflow is recommended to download MLC LLM, after which you can specify build requirements with our lightweight config generation tool: + +.. code-block:: bash + :caption: Configure and build + + # clone from GitHub + git clone --recursive https://github.com/mlc-ai/mlc-llm.git && cd mlc-llm/ + # create build directory + mkdir -p build && cd build + # generate build configuration + python ../cmake/gen_cmake_config.py + # build mlc_llm libraries + cmake .. && make -j $(nproc) && cd .. + +**Step 3. Install via Python.** We recommend that you install ``mlc_llm`` as a Python package, giving you +access to ``mlc_llm.compile``, ``mlc_llm.MLCEngine``, and the CLI. +There are two ways to do so: + + .. tabs :: + + .. code-tab :: bash Install via environment variable + + export MLC_LLM_SOURCE_DIR=/path-to-mlc-llm + export PYTHONPATH=$MLC_LLM_SOURCE_DIR/python:$PYTHONPATH + alias mlc_llm="python -m mlc_llm" + + .. code-tab :: bash Install via pip local project + + conda activate your-own-env + which python # make sure python is installed, expected output: path_to_conda/envs/your-own-env/bin/python + cd /path-to-mlc-llm/python + pip install -e . + +**Step 4. Validate installation.** You may validate if MLC libarires and mlc_llm CLI is compiled successfully using the following command: + +.. code-block:: bash + :caption: Validate installation + + # expected to see `libmlc_llm.so` and `libtvm_runtime.so` + ls -l ./build/ + # expected to see help message + mlc_llm chat -h + +Finally, you can verify installation in command line. You should see the path you used to build from source with: + +.. code:: bash + + python -c "import mlc_llm; print(mlc_llm)" diff --git a/docs/install/tvm.rst b/docs/install/tvm.rst new file mode 100644 index 0000000..e24a87c --- /dev/null +++ b/docs/install/tvm.rst @@ -0,0 +1,306 @@ +.. _install-tvm: + +Install TVM Compiler +========================== + +.. contents:: Table of Contents + :local: + :depth: 2 + +`TVM Unity `__, the latest development in Apache TVM, is required to build MLC LLM. Its features include: + +- High-performance CPU/GPU code generation instantly without tuning; +- Dynamic shape and symbolic shape tracking by design; +- Supporting both inference and training; +- Productive python-first compiler implementation. As a concrete example, MLC LLM compilation is implemented in pure python using its API. + +TVM can be installed directly from a prebuilt developer package, or built from source. + +.. _tvm-prebuilt-package: + +Option 1. Prebuilt Package +-------------------------- + +A nightly prebuilt Python package of Apache TVM is provided. + +.. note:: + ❗ Whenever using Python, it is highly recommended to use **conda** to manage an isolated Python environment to avoid missing dependencies, incompatible versions, and package conflicts. + +.. tabs:: + + .. tab:: Linux + + .. tabs:: + + .. tab:: CPU + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-cpu + + .. tab:: CUDA 12.8 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-cu128 + + .. tab:: CUDA 13.0 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-cu130 + + .. tab:: ROCm 6.1 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-rocm61 + + .. tab:: ROCm 6.2 + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-rocm62 + + .. tab:: Vulkan + + Supported in all Linux packages. + + .. note:: + + If encountering issues with GLIBC not found, please install the latest glibc in conda: + + .. code-block:: bash + + conda install -c conda-forge libstdcxx-ng + + .. tab:: macOS + + .. tabs:: + + .. tab:: CPU + Metal + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-cpu + + .. note:: + + Always check if conda is installed properly in macOS using the command below: + + .. code-block:: bash + + conda info | grep platform + + It should return "osx-64" for Mac with Intel chip, and "osx-arm64" for Mac with Apple chip. + + .. tab:: Windows + + .. tabs:: + + .. tab:: CPU + Vulkan + + .. code-block:: bash + + conda activate your-environment + python -m pip install --pre -U -f https://mlc.ai/wheels mlc-ai-nightly-cpu + + .. note:: + Make sure you also install vulkan loader and clang to avoid vulkan + not found error or clang not found(needed for jit compile) + + .. code-block:: bash + + conda install -c conda-forge clang libvulkan-loader + + If encountering the error below: + + .. code-block:: bash + + FileNotFoundError: Could not find module 'path\to\site-packages\tvm\tvm.dll' (or one of its dependencies). Try using the full path with constructor syntax. + + It is likely `zstd`, a dependency to LLVM, was missing. Please use the command below to get it installed: + + .. code-block:: bash + + conda install zstd + +.. _tvm-build-from-source: + +Option 2. Build from Source +--------------------------- + +While it is generally recommended to always use the prebuilt TVM, if you require more customization, you may need to build it from source. **NOTE.** this should only be attempted if you are familiar with the intricacies of C++, CMake, LLVM, Python, and other related systems. + +.. collapse:: Details + + **Step 1. Set up build dependency.** To build from source, you need to ensure that the following build dependencies are met: + + - CMake >= 3.24 + - LLVM >= 15 + - For please install LLVM>=17 for ROCm 6.1 and LLVM>=18 for ROCm 6.2. + - Git + - (Optional) CUDA >= 11.8 (targeting NVIDIA GPUs) + - (Optional) Metal (targeting Apple GPUs such as M1 and M2) + - (Optional) Vulkan (targeting NVIDIA, AMD, Intel and mobile GPUs) + - (Optional) OpenCL (targeting NVIDIA, AMD, Intel and mobile GPUs) + + .. note:: + - To target NVIDIA GPUs, either CUDA or Vulkan is required (CUDA is recommended); + - For AMD and Intel GPUs, Vulkan is necessary; + - When targeting Apple (macOS, iOS, iPadOS), Metal is a mandatory dependency; + - Some Android devices only support OpenCL, but most of them support Vulkan. + + To easiest way to manage dependency is via conda, which maintains a set of toolchains including LLVM across platforms. To create the environment of those build dependencies, one may simply use: + + .. code-block:: bash + :caption: Set up build dependencies in conda + + # make sure to start with a fresh environment + conda env remove -n tvm-build-venv + # create the conda environment with build dependency + conda create -n tvm-build-venv -c conda-forge \ + "llvmdev>=15" \ + "cmake>=3.24" \ + git \ + python=3.13 + # enter the build environment + conda activate tvm-build-venv + + **Step 2. Configure and build.** Standard git-based workflow are recommended to download Apache TVM, and then specify build requirements in ``config.cmake``: + + .. code-block:: bash + :caption: Download TVM from GitHub + + # clone from GitHub + git clone --recursive https://github.com/apache/tvm.git && cd tvm + # create the build directory + rm -rf build && mkdir build && cd build + # specify build requirements in `config.cmake` + cp ../cmake/config.cmake . + + We want to specifically tweak the following flags by appending them to the end of the configuration file: + + .. code-block:: bash + :caption: Configure build in ``config.cmake`` + + # controls default compilation flags + echo "set(CMAKE_BUILD_TYPE RelWithDebInfo)" >> config.cmake + # LLVM is a must dependency + echo "set(USE_LLVM \"llvm-config --ignore-libllvm --link-static\")" >> config.cmake + echo "set(HIDE_PRIVATE_SYMBOLS ON)" >> config.cmake + # GPU SDKs, turn on if needed + echo "set(USE_CUDA OFF)" >> config.cmake + echo "set(USE_ROCM OFF)" >> config.cmake + echo "set(USE_METAL OFF)" >> config.cmake + echo "set(USE_VULKAN OFF)" >> config.cmake + echo "set(USE_OPENCL OFF)" >> config.cmake + # Below are options for CUDA, turn on if needed + # CUDA_ARCH is the cuda compute capability of your GPU. + # Examples: 89 for 4090, 90a for H100/H200, 100a for B200. + # Reference: https://developer.nvidia.com/cuda-gpus + echo "set(CMAKE_CUDA_ARCHITECTURES YOUR_CUDA_COMPUTE_CAPABILITY_HERE)" >> config.cmake + echo "set(USE_CUBLAS ON)" >> config.cmake + echo "set(USE_CUTLASS ON)" >> config.cmake + echo "set(USE_THRUST ON)" >> config.cmake + echo "set(USE_NVTX ON)" >> config.cmake + # Below is the option for ROCM, turn on if needed + echo "set(USE_HIPBLAS ON)" >> config.cmake + + .. note:: + ``HIDE_PRIVATE_SYMBOLS`` is a configuration option that enables the ``-fvisibility=hidden`` flag. This flag helps prevent potential symbol conflicts between TVM and PyTorch. These conflicts arise due to the frameworks shipping LLVMs of different versions. + + `CMAKE_BUILD_TYPE `_ controls default compilation flag: + + - ``Debug`` sets ``-O0 -g`` + - ``RelWithDebInfo`` sets ``-O2 -g -DNDEBUG`` (recommended) + - ``Release`` sets ``-O3 -DNDEBUG`` + + Once ``config.cmake`` is edited accordingly, kick off build with the commands below: + + .. code-block:: bash + :caption: Build ``libtvm`` using cmake and cmake + + cmake .. && make -j $(nproc) && cd .. + + A success build should produce ``libtvm`` and ``libtvm_runtime`` under ``/path-tvm/build/`` directory. + + Leaving the build environment ``tvm-build-venv``, there are two ways to install the successful build into your environment: + + .. tabs :: + + .. code-tab :: bash Install via environment variable + + export PYTHONPATH=/path-to-tvm/python:$PYTHONPATH + + .. code-tab :: bash Install via pip local project + + conda activate your-own-env + conda install python # make sure python is installed + cd /path-to-tvm/python + pip install -e . + +.. `|` adds a blank line + +| + +.. _tvm-validate: + +Validate TVM Installation +------------------------- + +Using a compiler infrastructure with multiple language bindings could be error-prone. +Therefore, it is highly recommended to validate TVM installation before use. + +**Step 1. Locate TVM Python package.** The following command can help confirm that TVM is properly installed as a python package and provide the location of the TVM python package: + +.. code-block:: bash + + >>> python -c "import tvm; print(tvm.__file__)" + /some-path/lib/python3.13/site-packages/tvm/__init__.py + +**Step 2. Confirm which TVM library is used.** When maintaining multiple build or installation of TVM, it becomes important to double check if the python package is using the proper ``libtvm`` with the following command: + +.. code-block:: bash + + >>> python -c "import tvm; print(tvm.base._LIB)" + + +**Step 3. Reflect TVM build option.** Sometimes when downstream application fails, it could likely be some mistakes with a wrong TVM commit, or wrong build flags. To find it out, the following commands will be helpful: + +.. code-block:: bash + + >>> python -c "import tvm; print('\n'.join(f'{k}: {v}' for k, v in tvm.support.libinfo().items()))" + ... # Omitted less relevant options + GIT_COMMIT_HASH: 4f6289590252a1cf45a4dc37bce55a25043b8338 + HIDE_PRIVATE_SYMBOLS: ON + USE_LLVM: llvm-config --link-static + LLVM_VERSION: 15.0.7 + USE_VULKAN: OFF + USE_CUDA: OFF + CUDA_VERSION: NOT-FOUND + USE_OPENCL: OFF + USE_METAL: ON + USE_ROCM: OFF + +.. note:: + ``GIT_COMMIT_HASH`` indicates the exact commit of the TVM build, and it can be found on GitHub via ``https://github.com/mlc-ai/relax/commit/$GIT_COMMIT_HASH``. + +**Step 4. Check device detection.** Sometimes it could be helpful to understand if TVM could detect your device at all with the following commands: + +.. code-block:: bash + + >>> python -c "import tvm; print(tvm.metal().exist)" + True # or False + >>> python -c "import tvm; print(tvm.cuda().exist)" + False # or True + >>> python -c "import tvm; print(tvm.vulkan().exist)" + False # or True + +Please note that the commands above verify the presence of an actual device on the local machine for the TVM runtime (not the compiler) to execute properly. However, TVM compiler can perform compilation tasks without requiring a physical device. As long as the necessary toolchain, such as NVCC, is available, TVM supports cross-compilation even in the absence of an actual device. diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 0000000..51d3652 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=. +set BUILDDIR=_build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/microserving/tutorial.rst b/docs/microserving/tutorial.rst new file mode 100644 index 0000000..cfae7cb --- /dev/null +++ b/docs/microserving/tutorial.rst @@ -0,0 +1,205 @@ +Implement LLM Cross-engine Orchestration Patterns +====================================================================== + +In this tutorial, we will introduce how to implement LLM cross-engine +orchestration patterns, like prefill-decode disaggregation, in MLC-LLM +via microserving API. Aiming to make disaggregated serving programmable, +MicroServing provides a new RISC-style approach to design LLM serving +API at sub-request level. It enables programmable cross-engine serving +patterns in a few lines of python code. For more information of +microserving API, check out +https://blog.mlc.ai/2025/01/07/microserving-llm-engines. + +Below is an example of prefill-decode disaggregation implementation. An +LLM cross-engine orchestration pattern is implemented in a router, which +dispatches original OpenAI-style completion requests to a chain of +microserving API calls. In this code example, we create a subclass of +Router (which includes wrappers for calling microserving APIs), and +override ``translate_request`` function. The ``translate_request`` +function takes in a request and a unique identifier of the request +(``request_id``), and returns an AsyncGenerator of response. We launch +the CustomRouter and 2 engines, each of which has tensor parallel degree +2. Engine 0 is prefill engine and engine 1 is decode engine. + +.. code:: python + + from mlc_llm.router import Router + from mlc_llm.protocol import openai_api_protocol + from typing import Any, AsyncGenerator + from mlc_llm.serve.entrypoints import microserving_entrypoints + from mlc_llm.interface.router import serve + + import aiohttp + + class CustomRouter(Router): + async def translate_request(self, request: openai_api_protocol.CompletionRequest, request_id: str) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + pass + + + serve( + model="/path/to/model", # replace this with actual path + model_lib="/path/to/model_lib", # replace this with actual path + router_host="127.0.0.1", + router_port=9123, + endpoint_hosts=["127.0.0.1", "127.0.0.1"], + endpoint_ports=[9124,9125], + endpoint_num_gpus=[2,2], + enable_prefix_cache=False, + router_type=CustomRouter, + ) + +In the ``translate_request`` function, we first assign ``request_id`` to +request.user, and later the request id will be passed as an argument to +the microserving API. + +.. code:: python + + # we will pass request_id as an argument in microserving API calls + request.user = request_id + + +Next, call ``prep_recv`` on the decode engine to prepare KV entries for +receiving from remote. ``end=-1`` means that we will let the prefill +engine prefill all except the last token, which makes sure that the +prefill engine does not need sampling logic. ``prep_recv`` returns +address to receive KV from remote and matched prefix length. For +simplicity, we do not enable prefix cache in the tutorial, so we only +need the kv address here. + +.. code:: python + + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=3 * 3600), trust_env=True + ) as session: + decode_start = len(request.prompt) -1 + # 1. Ask decode engine to prepare KV entries to receive from prefill engine + prep_recv_request = microserving_entrypoints.PrepRecvRequest( + **request.model_dump(), end=decode_start + ) + ( + kv_addr_info, + _, + ) = await self.send_prepare_receive( + session=session, + request=prep_recv_request, + server_url=self.server_urls[1], # engine 0 is prefill, engine 1 is decode. Here is decode engine + ) + +Then, call ``remote_send`` on the prefill engine to compute and send KV +to decode engine. ``recv_rank=self.device_id_starts[1]`` means that we +are sending KV to engine 1 (decode engine). + +.. code:: python + + + # 2. Ask prefill engine to send KV to decode engine + remote_send_request = microserving_entrypoints.RemoteSendRequest( + **request.model_dump(), + begin=0, + end=decode_start, + kv_addr_info=kv_addr_info, + recv_rank=self.device_id_starts[1], # the rank of decode engine + ) + await self.send_remote_send( + session=session, + request=remote_send_request, + server_url=self.server_urls[0], # prefill engine + ) + +Finally, call ``start_generate`` on the decode engine to start +generating tokens. ``begin=decode_start`` means we will prefill the last +token in the prompt and start decoding. Notably, the decode process of +the request may be preempted. In such case, we yield None, so that the +router will rerun the ``translate_request`` function. + +.. code:: python + + # 3. Start decoding + start_generate_request = microserving_entrypoints.StartGenerateRequest( + **request.model_dump(), + begin=decode_start, + ) + async for response in self.send_start_generate( + session=session, + request=start_generate_request, + server_url=self.server_urls[1], + ): + if len(response.choices) > 0: + finish_reason = response.choices[0].finish_reason + if finish_reason == "preempt": + yield None + yield response + +Bringing everything together, the complete code is as below: + +.. code:: python + + from mlc_llm.router import Router + from mlc_llm.protocol import openai_api_protocol + from typing import Any, AsyncGenerator + from mlc_llm.serve.entrypoints import microserving_entrypoints + from mlc_llm.interface.router import serve + + import aiohttp + class CustomRouter(Router): + async def translate_request(self, request: openai_api_protocol.CompletionRequest, request_id: str) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + # we will pass request_id as an argument in microserving API calls + request.user = request_id + + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=3 * 3600), trust_env=True + ) as session: + decode_start = len(request.prompt) -1 + # 1. Ask decode engine to prepare KV entries to receive from prefill engine + prep_recv_request = microserving_entrypoints.PrepRecvRequest( + **request.model_dump(), end=decode_start + ) + ( + kv_addr_info, + _, + ) = await self.send_prepare_receive( + session=session, + request=prep_recv_request, + server_url=self.server_urls[1], # engine 0 is prefill, engine 1 is decode. Here is decode engine + ) + # 2. Ask prefill engine to send KV to decode engine + remote_send_request = microserving_entrypoints.RemoteSendRequest( + **request.model_dump(), + begin=0, + end=decode_start, + kv_addr_info=kv_addr_info, + recv_rank=self.device_id_starts[1], # the rank of decode engine + ) + await self.send_remote_send( + session=session, + request=remote_send_request, + server_url=self.server_urls[0], # prefill engine + ) + # 3. Start decoding + start_generate_request = microserving_entrypoints.StartGenerateRequest( + **request.model_dump(), + begin=decode_start, + ) + async for response in self.send_start_generate( + session=session, + request=start_generate_request, + server_url=self.server_urls[1], + ): + if len(response.choices) > 0: + finish_reason = response.choices[0].finish_reason + if finish_reason == "preempt": + yield None + yield response + + + serve( + model="/path/to/model", # replace this with actual path + model_lib="/path/to/model_lib", # replace this with actual path + router_host="127.0.0.1", + router_port=9123, + endpoint_hosts=["127.0.0.1", "127.0.0.1"], + endpoint_ports=[9124,9125], + endpoint_num_gpus=[2,2], + enable_prefix_cache=False, + router_type=CustomRouter, + ) diff --git a/docs/privacy.rst b/docs/privacy.rst new file mode 100644 index 0000000..cdd3c91 --- /dev/null +++ b/docs/privacy.rst @@ -0,0 +1,5 @@ +MLC Chat App Privacy +==================== + +MLC Chat run all generation locally. +All data stays in users' device and is not collected by the app. diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000..55ff467 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,20 @@ +--find-links https://mlc.ai/wheels +fastapi +ml_dtypes>=0.5.1 +mlc-ai-nightly-cpu +openai +prompt_toolkit +pydantic +safetensors +shortuuid +sphinx == 5.2.3 +sphinx-reredirects==0.1.2 +sphinx-rtd-theme +sphinx-tabs == 3.4.1 +sphinx-toolbox == 3.4.0 +sphinxcontrib-napoleon==0.7 +sphinxcontrib_httpdomain==1.8.1 +tiktoken +tlcpack-sphinx-addon==0.2.2 +torch +uvicorn diff --git a/examples/python/microserving/custom_router.py b/examples/python/microserving/custom_router.py new file mode 100644 index 0000000..f78da33 --- /dev/null +++ b/examples/python/microserving/custom_router.py @@ -0,0 +1,81 @@ +"""Microserving customized router example.""" + +from collections.abc import AsyncGenerator +from typing import Any + +import aiohttp + +from mlc_llm.interface.router import serve +from mlc_llm.protocol import openai_api_protocol +from mlc_llm.router import Router +from mlc_llm.serve.entrypoints import microserving_entrypoints + + +class CustomRouter(Router): + """A customized router class in Microserving.""" + + async def translate_request( + self, request: openai_api_protocol.CompletionRequest, request_id: str + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + # we will pass request_id as an argument in microserving API calls + request.user = request_id + + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=3 * 3600), trust_env=True + ) as session: + decode_start = len(request.prompt) - 1 + # 1. Ask decode engine to prepare KV entries to receive from prefill engine + prep_recv_request = microserving_entrypoints.PrepRecvRequest( + **request.model_dump(), end=decode_start + ) + ( + kv_addr_info, + _, + ) = await self.send_prepare_receive( + session=session, + request=prep_recv_request, + server_url=self.server_urls[ + 1 + ], # engine 0 is prefill, engine 1 is decode. Here is decode engine + ) + # 2. Ask prefill engine to send KV to decode engine + remote_send_request = microserving_entrypoints.RemoteSendRequest( + **request.model_dump(), + begin=0, + end=decode_start, + kv_addr_info=kv_addr_info, + recv_rank=self.device_id_starts[1], # the rank of decode engine + ) + await self.send_remote_send( + session=session, + request=remote_send_request, + server_url=self.server_urls[0], # prefill engine + ) + # 3. Start decoding + start_generate_request = microserving_entrypoints.StartGenerateRequest( + **request.model_dump(), + begin=decode_start, + ) + async for response in self.send_start_generate( + session=session, + request=start_generate_request, + server_url=self.server_urls[1], + ): + if len(response.choices) > 0: + finish_reason = response.choices[0].finish_reason + if finish_reason == "preempt": + yield None + yield response + + +serve( + model="/path/to/model", # replace this with actual path + model_lib="/path/to/model_lib.so", # replace this with actual path + router_host="127.0.0.1", + router_port=9123, + endpoint_hosts=["127.0.0.1", "127.0.0.1"], + endpoint_ports=[9124, 9125], + endpoint_num_gpus=[2, 2], + enable_prefix_cache=False, + router_type=CustomRouter, +) diff --git a/examples/python/sample_mlc_engine.py b/examples/python/sample_mlc_engine.py new file mode 100644 index 0000000..cf54106 --- /dev/null +++ b/examples/python/sample_mlc_engine.py @@ -0,0 +1,19 @@ +"""MLC Engine Python example.""" + +from mlc_llm import MLCEngine + +# Create engine +model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC" +engine = MLCEngine(model) + +# Run chat completion in OpenAI API. +for response in engine.chat.completions.create( + messages=[{"role": "user", "content": "What is the meaning of life?"}], + model=model, + stream=True, +): + for choice in response.choices: + print(choice.delta.content, end="", flush=True) +print("\n") + +engine.terminate() diff --git a/examples/rest/nodejs/README.MD b/examples/rest/nodejs/README.MD new file mode 100755 index 0000000..7bcdd2d --- /dev/null +++ b/examples/rest/nodejs/README.MD @@ -0,0 +1,21 @@ +# Node/Javascript/Typescript Access Examples for mlc_llm REST APIs + +Please make sure you are running v18.17.x of node (and npm v9.6.7) -- v20.x currently has some compatibility problems with typescript used in the langchain example. + +First install dependencies. + +`npm i` + +Copy `dotenv.exmaple` to `.env`. + +To run JS chat completion (both streaming and non-streaming) example: + +`node sample_client.js` + +To run OpenAI (chat completion streaming and non-streaming, and legacy completion) example: + +`node sample_openai.js` + +To run LangchainJS Typescript example: + +`npm run example` diff --git a/examples/rest/nodejs/dotenv.example b/examples/rest/nodejs/dotenv.example new file mode 100755 index 0000000..4bc5e03 --- /dev/null +++ b/examples/rest/nodejs/dotenv.example @@ -0,0 +1,2 @@ +OPENAI_API_KEY="none" +OPENAI_API_BASE="http://127.0.0.1:8000/v1" diff --git a/examples/rest/nodejs/package.json b/examples/rest/nodejs/package.json new file mode 100755 index 0000000..2a3ebf2 --- /dev/null +++ b/examples/rest/nodejs/package.json @@ -0,0 +1,40 @@ +{ + "name": "mlc-llm-js-examples", + "version": "1.0.0", + "description": "", + "main": "index.js", + "type": "module", + "license": "AGPL-version-3.0", + "private": false, + "engines": { + "node": ">= 14.0.0", + "npm": ">= 6.0.0" + }, + "homepage": "", + "repository": { + "type": "git", + "url": "" + }, + "bugs": "", + "keywords": [], + "author": { + "name": "", + "email": "", + "url": "" + }, + "contributors": [], + "scripts": { + "example": "ts-node --esm ./sample_langchain.ts" + }, + "dependencies": { + "@types/node": "^20.4.4", + "dotenv": "^16.3.1", + "langchain": "^0.0.117", + "needle": "^3.2.0", + "openai": "^3.3.0", + "typescript": "^5.1.6" + }, + "devDependencies": { + "ts-node": "^10.9.1" + } +} diff --git a/examples/rest/nodejs/sample_client.js b/examples/rest/nodejs/sample_client.js new file mode 100755 index 0000000..6f52a27 --- /dev/null +++ b/examples/rest/nodejs/sample_client.js @@ -0,0 +1,72 @@ +import request from 'needle'; + +( async () => { +const color = { + PURPLE : '\x1b[95m', + CYAN : '\x1b[96m', + DARKCYAN : '\x1b[36m', + BLUE : '\x1b[94m', + GREEN : '\x1b[92m', + YELLOW : '\x1b[93m', + RED : '\x1b[91m', + BOLD : '\x1b[1m', + UNDERLINE : '\x1b[4m', + END : '\x1b[0m' +}; + +let payload = { + model : 'vicuna-v1-7b', + messages: [{"role": "user", "content": "Write a haiku"}], + stream: false +}; + +const print = ( str ) => { + process.stdout.write(str); +}; + +const newline = () => { + print('\n'); +} + +newline(); +print(color.BOLD + "Without streaming:" + color.END); +newline(); + +let r = await request("post", "http://127.0.0.1:8000/v1/chat/completions", payload, {json: true}); + +print(color.GREEN + r.body.choices[0].message.content + color.END); +print('\n'); +// Reset the chat +r = await request("post", "http://127.0.0.1:8000/v1/chat/completions", payload, {json: true}); +print(color.BOLD + "Reset chat" + color.END); +newline(); + +// Get a response using a prompt with streaming + +payload = { + "model": "vicuna-v1-7b", + "messages": [{"role": "user", "content": "Write a haiku"}], + "stream": true +} + +print( color.BOLD + "With streaming:" + color.END); +newline(); +r = request.post( "http://127.0.0.1:8000/v1/chat/completions", payload, {json: true}) +.on('readable', function() { + let jsData = ''; + let data = ''; + while (data = this.read()) { + const chunk = data.toString().substring(6); + if (chunk.trim() === "[DONE]") break; + jsData = JSON.parse(chunk); + print(color.GREEN + jsData.choices[0].delta.content + color.END); + } +}) +.on('done', async function () { + newline(); + let txtresp = await request("get", "http://127.0.0.1:8000/stats"); + print(color.BOLD + "Runtime stats:" + color.END + txtresp.body); + +}) + +})() diff --git a/examples/rest/nodejs/sample_langchain.ts b/examples/rest/nodejs/sample_langchain.ts new file mode 100644 index 0000000..e4c9417 --- /dev/null +++ b/examples/rest/nodejs/sample_langchain.ts @@ -0,0 +1,74 @@ +import { OpenAI } from "langchain/llms/openai"; +import { BufferWindowMemory } from "langchain/memory"; +import { LLMChain } from "langchain/chains"; +import { PromptTemplate } from "langchain/prompts"; +import {TextLoader } from "langchain/document_loaders/fs/text"; +import { loadQAStuffChain } from "langchain/chains"; + +const color = { + PURPLE : '\x1b[95m', + CYAN : '\x1b[96m', + DARKCYAN : '\x1b[36m', + BLUE : '\x1b[94m', + GREEN : '\x1b[92m', + YELLOW : '\x1b[93m', + RED : '\x1b[91m', + BOLD : '\x1b[1m', + UNDERLINE : '\x1b[4m', + END : '\x1b[0m' +}; + +function print(str: string) { + process.stdout.write(str); +} + +const newline = () => { + print('\n'); +} + + const chat = new OpenAI( { + openAIApiKey: "empty", + temperature: 0 + }, { + basePath: 'http://127.0.0.1:8000/v1' + }); + +// Conversational LLMChain example + const memory = new BufferWindowMemory({ memoryKey: "history", k: 1 }); + + const template = `The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know. + + Current conversation: + {history} + Human: {human_input} + AI:`; + + + const prompt = PromptTemplate.fromTemplate(template); + let chain = new LLMChain({ llm: chat, prompt, memory }); + + let input = "Write a poem about Pittsburgh."; + print(color.BOLD + input + "..." + color.END); + newline(); + let res = await chain.call({ human_input: input }); + newline(); + print(color.GREEN + res.text + color.END); + newline(); + input = "What does it mean?"; + print(color.BOLD + input + "..." + color.END); + newline(); + res = await chain.call({ human_input: input }); + newline(); + print(color.GREEN + res.text + color.END); + newline(); + +// Question and answer stuff chain example with text loader +const loader = new TextLoader('../resources/linux.txt'); +const documents = await loader.load(); +const schain = loadQAStuffChain(chat); +const query = "When was Linux released?"; +newline(); newline(); +print(color.BOLD + "Query: " + color.END + color.BLUE + query + color.END); +newline(); +const result = await schain.call({ input_documents: documents, question: query}); +print(color.BOLD + "Response: " + color.END + color.GREEN + result.text + color.END); diff --git a/examples/rest/nodejs/sample_openai.js b/examples/rest/nodejs/sample_openai.js new file mode 100755 index 0000000..9a7631b --- /dev/null +++ b/examples/rest/nodejs/sample_openai.js @@ -0,0 +1,77 @@ +import { Configuration, OpenAIApi } from "openai"; +import dotenv from "dotenv"; +dotenv.config(); + +( async () => { + +const configuration = new Configuration({ + apiKey: process.env.OPENAI_API_KEY, + basePath : process.env.OPENAI_API_BASE +}) +const openai = new OpenAIApi(configuration); +let model = "vicuna-v1-7b" + +const color = { + PURPLE : '\x1b[95m', + CYAN : '\x1b[96m', + DARKCYAN : '\x1b[36m', + BLUE : '\x1b[94m', + GREEN : '\x1b[92m', + YELLOW : '\x1b[93m', + RED : '\x1b[91m', + BOLD : '\x1b[1m', + UNDERLINE : '\x1b[4m', + END : '\x1b[0m' +}; + +const print = ( str ) => { + process.stdout.write(str); +}; + +const newline = () => { + print('\n'); +} + +// Chat completion example without streaming +newline(); +print(color.BOLD + "OpenAI chat completion example without streaming:" + color.END); +newline(); + +let completion = await openai.createChatCompletion({ + model: model, + messages: [{"role": "user", "content": "Write a poem about OpenAI"}] +}); + + +print(color.GREEN + completion.data.choices[0].message.content + color.END) +newline(); newline(); + + +// Chat completion example with streaming +// (raw implementation since npm module does not support it yet - it will have support in upcoming 4.x) + +print(color.BOLD + "OpenAI chat completion example with streaming:" + color.END); +newline(); +completion = await openai.createChatCompletion({ + model: model, + messages: [{"role": "user", "content": "Write a poem about OpenAI"}], + stream: true, +}, {responseType: 'stream'}); + +completion.data.on('data', async (data) => { + const parsed = JSON.parse(data.toString().substring(6)); + print(color.GREEN + parsed.choices[0].delta.content + color.END); +}); + +completion.data.on('close', async () => { + newline(); newline(); + + // Completion example + print(color.BOLD + "OpenAI completion example:" + color.END) + newline(); + let res = await openai.createCompletion({ prompt: "Write a poem about OpenAI", model: model}); + print(color.GREEN + res.data.choices[0].text + color.END); + newline(); newline(); + + }); +})() diff --git a/examples/rest/nodejs/tsconfig.json b/examples/rest/nodejs/tsconfig.json new file mode 100755 index 0000000..f31ac99 --- /dev/null +++ b/examples/rest/nodejs/tsconfig.json @@ -0,0 +1,14 @@ +{ + "compilerOptions": { + "target": "es2020", /* Set the JavaScript language version for emitted JavaScript and include compatible library declarations. */ + "lib": ["es2020"], /* Specify a set of bundled library declaration files that describe the target runtime environment. */ + "module": "nodenext", /* Specify what module code is generated. */ + "rootDir": "src", /* Specify the root folder within your source files. */ + "outDir": "./dist", /* Specify an output folder for all emitted files. */ + "esModuleInterop": true, /* Emit additional JavaScript to ease support for importing CommonJS modules. This enables 'allowSyntheticDefaultImports' for type compatibility. */ + "forceConsistentCasingInFileNames": true, /* Ensure that casing is correct in imports. */ + "strict": true, /* Enable all strict type-checking options. */ + "noImplicitAny": true, /* Enable error reporting for expressions and declarations with an implied 'any' type. */ + "skipLibCheck": true /* Skip type checking all .d.ts files. */ + } +} diff --git a/examples/rest/python/sample_client.py b/examples/rest/python/sample_client.py new file mode 100644 index 0000000..ff25627 --- /dev/null +++ b/examples/rest/python/sample_client.py @@ -0,0 +1,51 @@ +import json + +import requests + + +class color: + PURPLE = "\033[95m" + CYAN = "\033[96m" + DARKCYAN = "\033[36m" + BLUE = "\033[94m" + GREEN = "\033[92m" + YELLOW = "\033[93m" + RED = "\033[91m" + BOLD = "\033[1m" + UNDERLINE = "\033[4m" + END = "\033[0m" + + +# Get a response using a prompt without streaming +payload = { + "model": "vicuna-v1-7b", + "messages": [{"role": "user", "content": "Write a haiku"}], + "stream": False, +} +r = requests.post("http://127.0.0.1:8000/v1/chat/completions", json=payload) +print( + f"{color.BOLD}Without streaming:{color.END}\n{color.GREEN}{r.json()['choices'][0]['message']['content']}{color.END}\n" # noqa: E501 +) + +# Reset the chat +r = requests.post("http://127.0.0.1:8000/chat/reset", json=payload) +print(f"{color.BOLD}Reset chat:{color.END} {str(r)}\n") + +# Get a response using a prompt with streaming +payload = { + "model": "vicuna-v1-7b", + "messages": [{"role": "user", "content": "Write a haiku"}], + "stream": True, +} +with requests.post("http://127.0.0.1:8000/v1/chat/completions", json=payload, stream=True) as r: + print(f"{color.BOLD}With streaming:{color.END}") + for chunk in r: + if chunk[6:].decode("utf-8").strip() == "[DONE]": + break + content = json.loads(chunk[6:])["choices"][0]["delta"].get("content", "") + print(f"{color.GREEN}{content}{color.END}", end="", flush=True) + print("\n") + +# Get the latest runtime stats +r = requests.get("http://127.0.0.1:8000/stats") +print(f"{color.BOLD}Runtime stats:{color.END} {r.json()}\n") diff --git a/examples/rest/python/sample_langchain.py b/examples/rest/python/sample_langchain.py new file mode 100644 index 0000000..4db474d --- /dev/null +++ b/examples/rest/python/sample_langchain.py @@ -0,0 +1,168 @@ +from langchain import LLMChain, PromptTemplate +from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler +from langchain.chains import RetrievalQA +from langchain.chains.question_answering import load_qa_chain +from langchain.chat_models import ChatOpenAI +from langchain.document_loaders import ( + DirectoryLoader, + TextLoader, + UnstructuredRSTLoader, +) +from langchain.llms import OpenAI +from langchain.memory import ConversationBufferWindowMemory +from langchain.text_splitter import CharacterTextSplitter +from langchain.vectorstores import Chroma + +# Note that Langchain support for embedding documents using MLC is currently blocked on +# https://github.com/langchain-ai/langchain/pull/7815 +# We have subclassed `OpenAIEmbeddings` in the meantime to get around this dependency. +from mlc_llm.contrib.embeddings.openai import MLCEmbeddings + +# First set the following in your environment: +# export OPENAI_API_BASE=http://127.0.0.1:8000/v1 +# export OPENAI_API_KEY=EMPTY + +# Note that Langchain does not currently support Pydantic v2: +# https://github.com/langchain-ai/langchain/issues/6841 +# Please ensure that your `pydantic` version is < 2.0 + + +class color: + PURPLE = "\033[95m" + CYAN = "\033[96m" + DARKCYAN = "\033[36m" + BLUE = "\033[94m" + GREEN = "\033[92m" + YELLOW = "\033[93m" + RED = "\033[91m" + BOLD = "\033[1m" + UNDERLINE = "\033[4m" + END = "\033[0m" + + +def llm_chain_example(): + template = """ + {history} + USER: {human_input} + ASSISTANT:""" + + prompt = PromptTemplate(input_variables=["history", "human_input"], template=template) + + llm_chain = LLMChain( + llm=ChatOpenAI(streaming=True, callbacks=[StreamingStdOutCallbackHandler()]), + prompt=prompt, + verbose=True, + memory=ConversationBufferWindowMemory(human_prefix="USER", ai_prefix="ASSISTANT"), + ) + + llm_chain.predict(human_input="Write a short poem about Pittsburgh.") + llm_chain.predict(human_input="What does the poem mean?") + + +def load_qa_chain_example(): + loader = TextLoader("../resources/linux.txt") + documents = loader.load() + chain = load_qa_chain(llm=OpenAI(), chain_type="stuff", verbose=False) + query = "When was Linux released?" + print(f"{color.BOLD}Query:{color.END} {color.BLUE} {query}{color.END}") + print( + f"{color.BOLD}Response:{color.END} {color.GREEN}{chain.run(input_documents=documents, question=query)}{color.END}" # noqa: E501 + ) + + +def retrieval_qa_sotu_example(): + prompt_template = """Use only the following pieces of context to answer the question at the end. Don't use any other knowledge. + + {context} + + USER: {question} + ASSISTANT:""" # noqa: E501 + + PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"]) + + loader = TextLoader("../resources/state_of_the_union.txt") + documents = loader.load() + + text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=100) + texts = text_splitter.split_documents(documents) + # print(texts) + embeddings = MLCEmbeddings(deployment="text-embedding-ada-002", embedding_ctx_length=None) + db = Chroma.from_documents(documents=texts, embedding=embeddings) + retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 2}) + qa = RetrievalQA.from_chain_type( + llm=OpenAI(), + chain_type="stuff", + retriever=retriever, + return_source_documents=True, + chain_type_kwargs={"prompt": PROMPT}, + ) + questions = [ + "What is the American Rescue Plan?", + "What did the president say about Ketanji Brown Jackson?", + "Who is mentioned in the speech?", + "To whom is the speech addressed?", + "Tell me more about the Made in America campaign.", + ] + + for qn in questions: + print(f"{color.BOLD}QUESTION:{color.END} {qn}") + res = qa({"query": qn}) + print(f"{color.BOLD}RESPONSE:{color.END} {color.GREEN}{res['result']}{color.END}") + print( + f"{color.BOLD}SOURCE:{color.END} {color.BLUE}{repr(res['source_documents'][0].page_content)}{color.END}" # noqa: E501 + ) + print() + + +def retrieval_qa_mlc_docs_example(): + prompt_template = """Use only the following pieces of context to answer the question at the end. Don't use any other knowledge. + + {context} + + USER: {question} + ASSISTANT:""" # noqa: E501 + + PROMPT = PromptTemplate(template=prompt_template, input_variables=["context", "question"]) + + loader = DirectoryLoader( + "../../../docs", + glob="*/*.rst", + show_progress=True, + loader_cls=UnstructuredRSTLoader, + loader_kwargs={"mode": "single"}, + ) + documents = loader.load() + text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100) + texts = text_splitter.split_documents(documents) + embeddings = MLCEmbeddings(deployment="text-embedding-ada-002", embedding_ctx_length=None) + db = Chroma.from_documents(collection_name="abc", documents=texts, embedding=embeddings) + retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 3}) + qa = RetrievalQA.from_chain_type( + llm=OpenAI(), + chain_type="stuff", + retriever=retriever, + return_source_documents=True, + chain_type_kwargs={"prompt": PROMPT}, + ) + while True: + qn = input(f"{color.BOLD}QUESTION:{color.END} ") + res = qa({"query": qn}) + print(f"{color.BOLD}RESPONSE:{color.END} {color.GREEN}{res['result']}{color.END}") + print( + f"{color.BOLD}SOURCE:{color.END} {color.BLUE}{repr(res['source_documents'][0].page_content)}{color.END}" # noqa: E501 + ) + print() + + # Some example questions: + # - What is the chat config? + # - What is temperature? + # - What are the REST API endpoints? + # - What are the available quantization options? + + +# Uncomment one of the following lines to try out the corresponding demo: + +# llm_chain_example() +# load_qa_chain_example() +# retrieval_qa_sotu_example() +# retrieval_qa_mlc_docs_example() diff --git a/examples/rest/python/sample_openai.py b/examples/rest/python/sample_openai.py new file mode 100644 index 0000000..f7f7269 --- /dev/null +++ b/examples/rest/python/sample_openai.py @@ -0,0 +1,42 @@ +import openai + +openai.api_key = "None" +openai.api_base = "http://127.0.0.1:8000/v1" + +model = "vicuna-v1-7b" + + +class color: + PURPLE = "\033[95m" + CYAN = "\033[96m" + DARKCYAN = "\033[36m" + BLUE = "\033[94m" + GREEN = "\033[92m" + YELLOW = "\033[93m" + RED = "\033[91m" + BOLD = "\033[1m" + UNDERLINE = "\033[4m" + END = "\033[0m" + + +# Chat completion example without streaming +print(f"{color.BOLD}OpenAI chat completion example without streaming:{color.END}\n") +completion = openai.ChatCompletion.create( + model=model, messages=[{"role": "user", "content": "Write a poem about OpenAI"}] +) +print(f"{color.GREEN}{completion.choices[0].message.content}{color.END}\n\n") + +# Chat completion example with streaming +print(f"{color.BOLD}OpenAI chat completion example with streaming:{color.END}\n") +res = openai.ChatCompletion.create( + model=model, messages=[{"role": "user", "content": "Write a poem about OpenAI"}], stream=True +) +for chunk in res: + content = chunk["choices"][0]["delta"].get("content", "") + print(f"{color.GREEN}{content}{color.END}", end="", flush=True) +print("\n") + +# Completion example +print(f"{color.BOLD}OpenAI completion example:{color.END}\n") +res = openai.Completion.create(prompt="Write a poem about OpenAI", model=model) +print(f"{color.GREEN}{res.choices[0].text}{color.END}\n\n") diff --git a/examples/rest/resources/linux.txt b/examples/rest/resources/linux.txt new file mode 100644 index 0000000..e56a218 --- /dev/null +++ b/examples/rest/resources/linux.txt @@ -0,0 +1,23 @@ +Linux is a family of open-source Unix-like operating systems based on the Linux kernel, an operating system kernel first released on September 17, 1991, by Linus Torvalds. Linux is typically packaged as a Linux distribution, which includes the kernel and supporting system software and libraries, many of which are provided by the GNU Project. Many Linux distributions use the word "Linux" in their name, but the Free Software Foundation uses the name "GNU/Linux" to emphasize the importance of GNU software, causing some controversy. + +Popular Linux distributions include Debian, Fedora Linux, and Ubuntu, the latter of which itself consists of many different distributions and modifications, including Lubuntu and Xubuntu. Commercial distributions include Red Hat Enterprise Linux and SUSE Linux Enterprise. Desktop Linux distributions include a windowing system such as X11 or Wayland, and a desktop environment such as GNOME or KDE Plasma. Distributions intended for servers may omit graphics altogether, or include a solution stack such as LAMP. Because Linux is freely redistributable, anyone may create a distribution for any purpose. + +Linux was originally developed for personal computers based on the Intel x86 architecture, but has since been ported to more platforms than any other operating system. Because of the dominance of the Linux-based Android on smartphones, Linux, including Android, has the largest installed base of all general-purpose operating systems, as of May 2022. Although Linux is, as of November 2022, used by only around 2.6 percent of desktop computers, the Chromebook, which runs the Linux kernel-based ChromeOS, dominates the US K–12 education market and represents nearly 20 percent of sub-$300 notebook sales in the US. Linux is the leading operating system on servers (over 96.4% of the top 1 million web servers' operating systems are Linux), leads other big iron systems such as mainframe computers, and is used on all of the world's 500 fastest supercomputers (since November 2017, having gradually displaced all competitors). + +Linux also runs on embedded systems, i.e. devices whose operating system is typically built into the firmware and is highly tailored to the system. This includes routers, automation controls, smart home devices, video game consoles, televisions (Samsung and LG Smart TVs), automobiles (Tesla, Audi, Mercedes-Benz, Hyundai and Toyota), and spacecraft (Falcon 9 rocket, Dragon crew capsule and the Perseverance rover). + +Linux is one of the most prominent examples of free and open-source software collaboration. The source code may be used, modified and distributed commercially or non-commercially by anyone under the terms of its respective licenses, such as the GNU General Public License (GPL). The Linux kernel, for example, is licensed under the GPLv2, with an exception for system calls that allows code that calls the kernel via system calls not to be licensed under the GPL. + +The Unix operating system was conceived and implemented in 1969, at AT&T's Bell Labs, in the United States by Ken Thompson, Dennis Ritchie, Douglas McIlroy, and Joe Ossanna. First released in 1971, Unix was written entirely in assembly language, as was common practice at the time. In 1973, in a key pioneering approach, it was rewritten in the C programming language by Dennis Ritchie (with the exception of some hardware and I/O routines). The availability of a high-level language implementation of Unix made its porting to different computer platforms easier. + +Due to an earlier antitrust case forbidding it from entering the computer business, AT&T licensed the operating system's source code as a trade secret to anyone who asked. As a result, Unix grew quickly and became widely adopted by academic institutions and businesses. In 1984, AT&T divested itself of its regional operating companies, and was released from its obligation not to enter the computer business; freed of that obligation, Bell Labs began selling Unix as a proprietary product, where users were not legally allowed to modify it. + +Onyx Systems began selling early microcomputer-based Unix workstations in 1980. Later, Sun Microsystems, founded as a spin-off of a student project at Stanford University, also began selling Unix-based desktop workstations in 1982. While Sun workstations did not utilize commodity PC hardware, for which Linux was later originally developed, it represented the first successful commercial attempt at distributing a primarily single-user microcomputer that ran a Unix operating system. + +With Unix increasingly "locked in" as a proprietary product, the GNU Project, started in 1983 by Richard Stallman, had the goal of creating a "complete Unix-compatible software system" composed entirely of free software. Work began in 1984. Later, in 1985, Stallman started the Free Software Foundation and wrote the GNU General Public License (GNU GPL) in 1989. By the early 1990s, many of the programs required in an operating system (such as libraries, compilers, text editors, a command-line shell, and a windowing system) were completed, although low-level elements such as device drivers, daemons, and the kernel, called GNU Hurd, were stalled and incomplete. + +MINIX was created by Andrew S. Tanenbaum, a computer science professor, and released in 1987 as a minimal Unix-like operating system targeted at students and others who wanted to learn operating system principles. Although the complete source code of MINIX was freely available, the licensing terms prevented it from being free software until the licensing changed in April 2000. + +Although not released until 1992, due to legal complications, development of 386BSD, from which NetBSD, OpenBSD and FreeBSD descended, predated that of Linux. + +Linus Torvalds has stated on separate occasions that if the GNU kernel or 386BSD had been available at the time (1991), he probably would not have created Linux. diff --git a/examples/rest/resources/state_of_the_union.txt b/examples/rest/resources/state_of_the_union.txt new file mode 100644 index 0000000..7cb2a02 --- /dev/null +++ b/examples/rest/resources/state_of_the_union.txt @@ -0,0 +1,723 @@ +Madam Speaker, Madam Vice President, our First Lady and Second Gentleman. Members of Congress and the Cabinet. Justices of the Supreme Court. My fellow Americans. + +Last year COVID-19 kept us apart. This year we are finally together again. + +Tonight, we meet as Democrats Republicans and Independents. But most importantly as Americans. + +With a duty to one another to the American people to the Constitution. + +And with an unwavering resolve that freedom will always triumph over tyranny. + +Six days ago, Russia’s Vladimir Putin sought to shake the foundations of the free world thinking he could make it bend to his menacing ways. But he badly miscalculated. + +He thought he could roll into Ukraine and the world would roll over. Instead he met a wall of strength he never imagined. + +He met the Ukrainian people. + +From President Zelenskyy to every Ukrainian, their fearlessness, their courage, their determination, inspires the world. + +Groups of citizens blocking tanks with their bodies. Everyone from students to retirees teachers turned soldiers defending their homeland. + +In this struggle as President Zelenskyy said in his speech to the European Parliament “Light will win over darkness.” The Ukrainian Ambassador to the United States is here tonight. + +Let each of us here tonight in this Chamber send an unmistakable signal to Ukraine and to the world. + +Please rise if you are able and show that, Yes, we the United States of America stand with the Ukrainian people. + +Throughout our history we’ve learned this lesson when dictators do not pay a price for their aggression they cause more chaos. + +They keep moving. + +And the costs and the threats to America and the world keep rising. + +That’s why the NATO Alliance was created to secure peace and stability in Europe after World War 2. + +The United States is a member along with 29 other nations. + +It matters. American diplomacy matters. American resolve matters. + +Putin’s latest attack on Ukraine was premeditated and unprovoked. + +He rejected repeated efforts at diplomacy. + +He thought the West and NATO wouldn’t respond. And he thought he could divide us at home. Putin was wrong. We were ready. Here is what we did. + +We prepared extensively and carefully. + +We spent months building a coalition of other freedom-loving nations from Europe and the Americas to Asia and Africa to confront Putin. + +I spent countless hours unifying our European allies. We shared with the world in advance what we knew Putin was planning and precisely how he would try to falsely justify his aggression. + +We countered Russia’s lies with truth. + +And now that he has acted the free world is holding him accountable. + +Along with twenty-seven members of the European Union including France, Germany, Italy, as well as countries like the United Kingdom, Canada, Japan, Korea, Australia, New Zealand, and many others, even Switzerland. + +We are inflicting pain on Russia and supporting the people of Ukraine. Putin is now isolated from the world more than ever. + +Together with our allies –we are right now enforcing powerful economic sanctions. + +We are cutting off Russia’s largest banks from the international financial system. + +Preventing Russia’s central bank from defending the Russian Ruble making Putin’s $630 Billion “war fund” worthless. + +We are choking off Russia’s access to technology that will sap its economic strength and weaken its military for years to come. + +Tonight I say to the Russian oligarchs and corrupt leaders who have bilked billions of dollars off this violent regime no more. + +The U.S. Department of Justice is assembling a dedicated task force to go after the crimes of Russian oligarchs. + +We are joining with our European allies to find and seize your yachts your luxury apartments your private jets. We are coming for your ill-begotten gains. + +And tonight I am announcing that we will join our allies in closing off American air space to all Russian flights – further isolating Russia – and adding an additional squeeze –on their economy. The Ruble has lost 30% of its value. + +The Russian stock market has lost 40% of its value and trading remains suspended. Russia’s economy is reeling and Putin alone is to blame. + +Together with our allies we are providing support to the Ukrainians in their fight for freedom. Military assistance. Economic assistance. Humanitarian assistance. + +We are giving more than $1 Billion in direct assistance to Ukraine. + +And we will continue to aid the Ukrainian people as they defend their country and to help ease their suffering. + +Let me be clear, our forces are not engaged and will not engage in conflict with Russian forces in Ukraine. + +Our forces are not going to Europe to fight in Ukraine, but to defend our NATO Allies – in the event that Putin decides to keep moving west. + +For that purpose we’ve mobilized American ground forces, air squadrons, and ship deployments to protect NATO countries including Poland, Romania, Latvia, Lithuania, and Estonia. + +As I have made crystal clear the United States and our Allies will defend every inch of territory of NATO countries with the full force of our collective power. + +And we remain clear-eyed. The Ukrainians are fighting back with pure courage. But the next few days weeks, months, will be hard on them. + +Putin has unleashed violence and chaos. But while he may make gains on the battlefield – he will pay a continuing high price over the long run. + +And a proud Ukrainian people, who have known 30 years of independence, have repeatedly shown that they will not tolerate anyone who tries to take their country backwards. + +To all Americans, I will be honest with you, as I’ve always promised. A Russian dictator, invading a foreign country, has costs around the world. + +And I’m taking robust action to make sure the pain of our sanctions is targeted at Russia’s economy. And I will use every tool at our disposal to protect American businesses and consumers. + +Tonight, I can announce that the United States has worked with 30 other countries to release 60 Million barrels of oil from reserves around the world. + +America will lead that effort, releasing 30 Million barrels from our own Strategic Petroleum Reserve. And we stand ready to do more if necessary, unified with our allies. + +These steps will help blunt gas prices here at home. And I know the news about what’s happening can seem alarming. + +But I want you to know that we are going to be okay. + +When the history of this era is written Putin’s war on Ukraine will have left Russia weaker and the rest of the world stronger. + +While it shouldn’t have taken something so terrible for people around the world to see what’s at stake now everyone sees it clearly. + +We see the unity among leaders of nations and a more unified Europe a more unified West. And we see unity among the people who are gathering in cities in large crowds around the world even in Russia to demonstrate their support for Ukraine. + +In the battle between democracy and autocracy, democracies are rising to the moment, and the world is clearly choosing the side of peace and security. + +This is a real test. It’s going to take time. So let us continue to draw inspiration from the iron will of the Ukrainian people. + +To our fellow Ukrainian Americans who forge a deep bond that connects our two nations we stand with you. + +Putin may circle Kyiv with tanks, but he will never gain the hearts and souls of the Ukrainian people. + +He will never extinguish their love of freedom. He will never weaken the resolve of the free world. + +We meet tonight in an America that has lived through two of the hardest years this nation has ever faced. + +The pandemic has been punishing. + +And so many families are living paycheck to paycheck, struggling to keep up with the rising cost of food, gas, housing, and so much more. + +I understand. + +I remember when my Dad had to leave our home in Scranton, Pennsylvania to find work. I grew up in a family where if the price of food went up, you felt it. + +That’s why one of the first things I did as President was fight to pass the American Rescue Plan. + +Because people were hurting. We needed to act, and we did. + +Few pieces of legislation have done more in a critical moment in our history to lift us out of crisis. + +It fueled our efforts to vaccinate the nation and combat COVID-19. It delivered immediate economic relief for tens of millions of Americans. + +Helped put food on their table, keep a roof over their heads, and cut the cost of health insurance. + +And as my Dad used to say, it gave people a little breathing room. + +And unlike the $2 Trillion tax cut passed in the previous administration that benefitted the top 1% of Americans, the American Rescue Plan helped working people—and left no one behind. + +And it worked. It created jobs. Lots of jobs. + +In fact—our economy created over 6.5 Million new jobs just last year, more jobs created in one year +than ever before in the history of America. + +Our economy grew at a rate of 5.7% last year, the strongest growth in nearly 40 years, the first step in bringing fundamental change to an economy that hasn’t worked for the working people of this nation for too long. + +For the past 40 years we were told that if we gave tax breaks to those at the very top, the benefits would trickle down to everyone else. + +But that trickle-down theory led to weaker economic growth, lower wages, bigger deficits, and the widest gap between those at the top and everyone else in nearly a century. + +Vice President Harris and I ran for office with a new economic vision for America. + +Invest in America. Educate Americans. Grow the workforce. Build the economy from the bottom up +and the middle out, not from the top down. + +Because we know that when the middle class grows, the poor have a ladder up and the wealthy do very well. + +America used to have the best roads, bridges, and airports on Earth. + +Now our infrastructure is ranked 13th in the world. + +We won’t be able to compete for the jobs of the 21st Century if we don’t fix that. + +That’s why it was so important to pass the Bipartisan Infrastructure Law—the most sweeping investment to rebuild America in history. + +This was a bipartisan effort, and I want to thank the members of both parties who worked to make it happen. + +We’re done talking about infrastructure weeks. + +We’re going to have an infrastructure decade. + +It is going to transform America and put us on a path to win the economic competition of the 21st Century that we face with the rest of the world—particularly with China. + +As I’ve told Xi Jinping, it is never a good bet to bet against the American people. + +We’ll create good jobs for millions of Americans, modernizing roads, airports, ports, and waterways all across America. + +And we’ll do it all to withstand the devastating effects of the climate crisis and promote environmental justice. + +We’ll build a national network of 500,000 electric vehicle charging stations, begin to replace poisonous lead pipes—so every child—and every American—has clean water to drink at home and at school, provide affordable high-speed internet for every American—urban, suburban, rural, and tribal communities. + +4,000 projects have already been announced. + +And tonight, I’m announcing that this year we will start fixing over 65,000 miles of highway and 1,500 bridges in disrepair. + +When we use taxpayer dollars to rebuild America – we are going to Buy American: buy American products to support American jobs. + +The federal government spends about $600 Billion a year to keep the country safe and secure. + +There’s been a law on the books for almost a century +to make sure taxpayers’ dollars support American jobs and businesses. + +Every Administration says they’ll do it, but we are actually doing it. + +We will buy American to make sure everything from the deck of an aircraft carrier to the steel on highway guardrails are made in America. + +But to compete for the best jobs of the future, we also need to level the playing field with China and other competitors. + +That’s why it is so important to pass the Bipartisan Innovation Act sitting in Congress that will make record investments in emerging technologies and American manufacturing. + +Let me give you one example of why it’s so important to pass it. + +If you travel 20 miles east of Columbus, Ohio, you’ll find 1,000 empty acres of land. + +It won’t look like much, but if you stop and look closely, you’ll see a “Field of dreams,” the ground on which America’s future will be built. + +This is where Intel, the American company that helped build Silicon Valley, is going to build its $20 billion semiconductor “mega site”. + +Up to eight state-of-the-art factories in one place. 10,000 new good-paying jobs. + +Some of the most sophisticated manufacturing in the world to make computer chips the size of a fingertip that power the world and our everyday lives. + +Smartphones. The Internet. Technology we have yet to invent. + +But that’s just the beginning. + +Intel’s CEO, Pat Gelsinger, who is here tonight, told me they are ready to increase their investment from +$20 billion to $100 billion. + +That would be one of the biggest investments in manufacturing in American history. + +And all they’re waiting for is for you to pass this bill. + +So let’s not wait any longer. Send it to my desk. I’ll sign it. + +And we will really take off. + +And Intel is not alone. + +There’s something happening in America. + +Just look around and you’ll see an amazing story. + +The rebirth of the pride that comes from stamping products “Made In America.” The revitalization of American manufacturing. + +Companies are choosing to build new factories here, when just a few years ago, they would have built them overseas. + +That’s what is happening. Ford is investing $11 billion to build electric vehicles, creating 11,000 jobs across the country. + +GM is making the largest investment in its history—$7 billion to build electric vehicles, creating 4,000 jobs in Michigan. + +All told, we created 369,000 new manufacturing jobs in America just last year. + +Powered by people I’ve met like JoJo Burgess, from generations of union steelworkers from Pittsburgh, who’s here with us tonight. + +As Ohio Senator Sherrod Brown says, “It’s time to bury the label “Rust Belt.” + +It’s time. + +But with all the bright spots in our economy, record job growth and higher wages, too many families are struggling to keep up with the bills. + +Inflation is robbing them of the gains they might otherwise feel. + +I get it. That’s why my top priority is getting prices under control. + +Look, our economy roared back faster than most predicted, but the pandemic meant that businesses had a hard time hiring enough workers to keep up production in their factories. + +The pandemic also disrupted global supply chains. + +When factories close, it takes longer to make goods and get them from the warehouse to the store, and prices go up. + +Look at cars. + +Last year, there weren’t enough semiconductors to make all the cars that people wanted to buy. + +And guess what, prices of automobiles went up. + +So—we have a choice. + +One way to fight inflation is to drive down wages and make Americans poorer. + +I have a better plan to fight inflation. + +Lower your costs, not your wages. + +Make more cars and semiconductors in America. + +More infrastructure and innovation in America. + +More goods moving faster and cheaper in America. + +More jobs where you can earn a good living in America. + +And instead of relying on foreign supply chains, let’s make it in America. + +Economists call it “increasing the productive capacity of our economy.” + +I call it building a better America. + +My plan to fight inflation will lower your costs and lower the deficit. + +17 Nobel laureates in economics say my plan will ease long-term inflationary pressures. Top business leaders and most Americans support my plan. And here’s the plan: + +First – cut the cost of prescription drugs. Just look at insulin. One in ten Americans has diabetes. In Virginia, I met a 13-year-old boy named Joshua Davis. + +He and his Dad both have Type 1 diabetes, which means they need insulin every day. Insulin costs about $10 a vial to make. + +But drug companies charge families like Joshua and his Dad up to 30 times more. I spoke with Joshua’s mom. + +Imagine what it’s like to look at your child who needs insulin and have no idea how you’re going to pay for it. + +What it does to your dignity, your ability to look your child in the eye, to be the parent you expect to be. + +Joshua is here with us tonight. Yesterday was his birthday. Happy birthday, buddy. + +For Joshua, and for the 200,000 other young people with Type 1 diabetes, let’s cap the cost of insulin at $35 a month so everyone can afford it. + +Drug companies will still do very well. And while we’re at it let Medicare negotiate lower prices for prescription drugs, like the VA already does. + +Look, the American Rescue Plan is helping millions of families on Affordable Care Act plans save $2,400 a year on their health care premiums. Let’s close the coverage gap and make those savings permanent. + +Second – cut energy costs for families an average of $500 a year by combatting climate change. + +Let’s provide investments and tax credits to weatherize your homes and businesses to be energy efficient and you get a tax credit; double America’s clean energy production in solar, wind, and so much more; lower the price of electric vehicles, saving you another $80 a month because you’ll never have to pay at the gas pump again. + +Third – cut the cost of child care. Many families pay up to $14,000 a year for child care per child. + +Middle-class and working families shouldn’t have to pay more than 7% of their income for care of young children. + +My plan will cut the cost in half for most families and help parents, including millions of women, who left the workforce during the pandemic because they couldn’t afford child care, to be able to get back to work. + +My plan doesn’t stop there. It also includes home and long-term care. More affordable housing. And Pre-K for every 3- and 4-year-old. + +All of these will lower costs. + +And under my plan, nobody earning less than $400,000 a year will pay an additional penny in new taxes. Nobody. + +The one thing all Americans agree on is that the tax system is not fair. We have to fix it. + +I’m not looking to punish anyone. But let’s make sure corporations and the wealthiest Americans start paying their fair share. + +Just last year, 55 Fortune 500 corporations earned $40 billion in profits and paid zero dollars in federal income tax. + +That’s simply not fair. That’s why I’ve proposed a 15% minimum tax rate for corporations. + +We got more than 130 countries to agree on a global minimum tax rate so companies can’t get out of paying their taxes at home by shipping jobs and factories overseas. + +That’s why I’ve proposed closing loopholes so the very wealthy don’t pay a lower tax rate than a teacher or a firefighter. + +So that’s my plan. It will grow the economy and lower costs for families. + +So what are we waiting for? Let’s get this done. And while you’re at it, confirm my nominees to the Federal Reserve, which plays a critical role in fighting inflation. + +My plan will not only lower costs to give families a fair shot, it will lower the deficit. + +The previous Administration not only ballooned the deficit with tax cuts for the very wealthy and corporations, it undermined the watchdogs whose job was to keep pandemic relief funds from being wasted. + +But in my administration, the watchdogs have been welcomed back. + +We’re going after the criminals who stole billions in relief money meant for small businesses and millions of Americans. + +And tonight, I’m announcing that the Justice Department will name a chief prosecutor for pandemic fraud. + +By the end of this year, the deficit will be down to less than half what it was before I took office. + +The only president ever to cut the deficit by more than one trillion dollars in a single year. + +Lowering your costs also means demanding more competition. + +I’m a capitalist, but capitalism without competition isn’t capitalism. + +It’s exploitation—and it drives up prices. + +When corporations don’t have to compete, their profits go up, your prices go up, and small businesses and family farmers and ranchers go under. + +We see it happening with ocean carriers moving goods in and out of America. + +During the pandemic, these foreign-owned companies raised prices by as much as 1,000% and made record profits. + +Tonight, I’m announcing a crackdown on these companies overcharging American businesses and consumers. + +And as Wall Street firms take over more nursing homes, quality in those homes has gone down and costs have gone up. + +That ends on my watch. + +Medicare is going to set higher standards for nursing homes and make sure your loved ones get the care they deserve and expect. + +We’ll also cut costs and keep the economy going strong by giving workers a fair shot, provide more training and apprenticeships, hire them based on their skills not degrees. + +Let’s pass the Paycheck Fairness Act and paid leave. + +Raise the minimum wage to $15 an hour and extend the Child Tax Credit, so no one has to raise a family in poverty. + +Let’s increase Pell Grants and increase our historic support of HBCUs, and invest in what Jill—our First Lady who teaches full-time—calls America’s best-kept secret: community colleges. + +And let’s pass the PRO Act when a majority of workers want to form a union—they shouldn’t be stopped. + +When we invest in our workers, when we build the economy from the bottom up and the middle out together, we can do something we haven’t done in a long time: build a better America. + +For more than two years, COVID-19 has impacted every decision in our lives and the life of the nation. + +And I know you’re tired, frustrated, and exhausted. + +But I also know this. + +Because of the progress we’ve made, because of your resilience and the tools we have, tonight I can say +we are moving forward safely, back to more normal routines. + +We’ve reached a new moment in the fight against COVID-19, with severe cases down to a level not seen since last July. + +Just a few days ago, the Centers for Disease Control and Prevention—the CDC—issued new mask guidelines. + +Under these new guidelines, most Americans in most of the country can now be mask free. + +And based on the projections, more of the country will reach that point across the next couple of weeks. + +Thanks to the progress we have made this past year, COVID-19 need no longer control our lives. + +I know some are talking about “living with COVID-19”. Tonight – I say that we will never just accept living with COVID-19. + +We will continue to combat the virus as we do other diseases. And because this is a virus that mutates and spreads, we will stay on guard. + +Here are four common sense steps as we move forward safely. + +First, stay protected with vaccines and treatments. We know how incredibly effective vaccines are. If you’re vaccinated and boosted you have the highest degree of protection. + +We will never give up on vaccinating more Americans. Now, I know parents with kids under 5 are eager to see a vaccine authorized for their children. + +The scientists are working hard to get that done and we’ll be ready with plenty of vaccines when they do. + +We’re also ready with anti-viral treatments. If you get COVID-19, the Pfizer pill reduces your chances of ending up in the hospital by 90%. + +We’ve ordered more of these pills than anyone in the world. And Pfizer is working overtime to get us 1 Million pills this month and more than double that next month. + +And we’re launching the “Test to Treat” initiative so people can get tested at a pharmacy, and if they’re positive, receive antiviral pills on the spot at no cost. + +If you’re immunocompromised or have some other vulnerability, we have treatments and free high-quality masks. + +We’re leaving no one behind or ignoring anyone’s needs as we move forward. + +And on testing, we have made hundreds of millions of tests available for you to order for free. + +Even if you already ordered free tests tonight, I am announcing that you can order more from covidtests.gov starting next week. + +Second – we must prepare for new variants. Over the past year, we’ve gotten much better at detecting new variants. + +If necessary, we’ll be able to deploy new vaccines within 100 days instead of many more months or years. + +And, if Congress provides the funds we need, we’ll have new stockpiles of tests, masks, and pills ready if needed. + +I cannot promise a new variant won’t come. But I can promise you we’ll do everything within our power to be ready if it does. + +Third – we can end the shutdown of schools and businesses. We have the tools we need. + +It’s time for Americans to get back to work and fill our great downtowns again. People working from home can feel safe to begin to return to the office. + +We’re doing that here in the federal government. The vast majority of federal workers will once again work in person. + +Our schools are open. Let’s keep it that way. Our kids need to be in school. + +And with 75% of adult Americans fully vaccinated and hospitalizations down by 77%, most Americans can remove their masks, return to work, stay in the classroom, and move forward safely. + +We achieved this because we provided free vaccines, treatments, tests, and masks. + +Of course, continuing this costs money. + +I will soon send Congress a request. + +The vast majority of Americans have used these tools and may want to again, so I expect Congress to pass it quickly. + +Fourth, we will continue vaccinating the world. + +We’ve sent 475 Million vaccine doses to 112 countries, more than any other nation. + +And we won’t stop. + +We have lost so much to COVID-19. Time with one another. And worst of all, so much loss of life. + +Let’s use this moment to reset. Let’s stop looking at COVID-19 as a partisan dividing line and see it for what it is: A God-awful disease. + +Let’s stop seeing each other as enemies, and start seeing each other for who we really are: Fellow Americans. + +We can’t change how divided we’ve been. But we can change how we move forward—on COVID-19 and other issues we must face together. + +I recently visited the New York City Police Department days after the funerals of Officer Wilbert Mora and his partner, Officer Jason Rivera. + +They were responding to a 9-1-1 call when a man shot and killed them with a stolen gun. + +Officer Mora was 27 years old. + +Officer Rivera was 22. + +Both Dominican Americans who’d grown up on the same streets they later chose to patrol as police officers. + +I spoke with their families and told them that we are forever in debt for their sacrifice, and we will carry on their mission to restore the trust and safety every community deserves. + +I’ve worked on these issues a long time. + +I know what works: Investing in crime preventionand community police officers who’ll walk the beat, who’ll know the neighborhood, and who can restore trust and safety. + +So let’s not abandon our streets. Or choose between safety and equal justice. + +Let’s come together to protect our communities, restore trust, and hold law enforcement accountable. + +That’s why the Justice Department required body cameras, banned chokeholds, and restricted no-knock warrants for its officers. + +That’s why the American Rescue Plan provided $350 Billion that cities, states, and counties can use to hire more police and invest in proven strategies like community violence interruption—trusted messengers breaking the cycle of violence and trauma and giving young people hope. + +We should all agree: The answer is not to Defund the police. The answer is to FUND the police with the resources and training they need to protect our communities. + +I ask Democrats and Republicans alike: Pass my budget and keep our neighborhoods safe. + +And I will keep doing everything in my power to crack down on gun trafficking and ghost guns you can buy online and make at home—they have no serial numbers and can’t be traced. + +And I ask Congress to pass proven measures to reduce gun violence. Pass universal background checks. Why should anyone on a terrorist list be able to purchase a weapon? + +Ban assault weapons and high-capacity magazines. + +Repeal the liability shield that makes gun manufacturers the only industry in America that can’t be sued. + +These laws don’t infringe on the Second Amendment. They save lives. + +The most fundamental right in America is the right to vote – and to have it counted. And it’s under assault. + +In state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. + +We cannot let this happen. + +Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. + +Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. + +One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. + +And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence. + +A former top litigator in private practice. A former federal public defender. And from a family of public school educators and police officers. A consensus builder. Since she’s been nominated, she’s received a broad range of support—from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. + +And if we are to advance liberty and justice, we need to secure the Border and fix the immigration system. + +We can do both. At our border, we’ve installed new technology like cutting-edge scanners to better detect drug smuggling. + +We’ve set up joint patrols with Mexico and Guatemala to catch more human traffickers. + +We’re putting in place dedicated immigration judges so families fleeing persecution and violence can have their cases heard faster. + +We’re securing commitments and supporting partners in South and Central America to host more refugees and secure their own borders. + +We can do all this while keeping lit the torch of liberty that has led generations of immigrants to this land—my forefathers and so many of yours. + +Provide a pathway to citizenship for Dreamers, those on temporary status, farm workers, and essential workers. + +Revise our laws so businesses have the workers they need and families don’t wait decades to reunite. + +It’s not only the right thing to do—it’s the economically smart thing to do. + +That’s why immigration reform is supported by everyone from labor unions to religious leaders to the U.S. Chamber of Commerce. + +Let’s get it done once and for all. + +Advancing liberty and justice also requires protecting the rights of women. + +The constitutional right affirmed in Roe v. Wade—standing precedent for half a century—is under attack as never before. + +If we want to go forward—not backward—we must protect access to health care. Preserve a woman’s right to choose. And let’s continue to advance maternal health care in America. + +And for our LGBTQ+ Americans, let’s finally get the bipartisan Equality Act to my desk. The onslaught of state laws targeting transgender Americans and their families is wrong. + +As I said last year, especially to our younger transgender Americans, I will always have your back as your President, so you can be yourself and reach your God-given potential. + +While it often appears that we never agree, that isn’t true. I signed 80 bipartisan bills into law last year. From preventing government shutdowns to protecting Asian-Americans from still-too-common hate crimes to reforming military justice. + +And soon, we’ll strengthen the Violence Against Women Act that I first wrote three decades ago. It is important for us to show the nation that we can come together and do big things. + +So tonight I’m offering a Unity Agenda for the Nation. Four big things we can do together. + +First, beat the opioid epidemic. + +There is so much we can do. Increase funding for prevention, treatment, harm reduction, and recovery. + +Get rid of outdated rules that stop doctors from prescribing treatments. And stop the flow of illicit drugs by working with state and local law enforcement to go after traffickers. + +If you’re suffering from addiction, know you are not alone. I believe in recovery, and I celebrate the 23 million Americans in recovery. + +Second, let’s take on mental health. Especially among our children, whose lives and education have been turned upside down. + +The American Rescue Plan gave schools money to hire teachers and help students make up for lost learning. + +I urge every parent to make sure your school does just that. And we can all play a part—sign up to be a tutor or a mentor. + +Children were also struggling before the pandemic. Bullying, violence, trauma, and the harms of social media. + +As Frances Haugen, who is here with us tonight, has shown, we must hold social media platforms accountable for the national experiment they’re conducting on our children for profit. + +It’s time to strengthen privacy protections, ban targeted advertising to children, demand tech companies stop collecting personal data on our children. + +And let’s get all Americans the mental health services they need. More people they can turn to for help, and full parity between physical and mental health care. + +Third, support our veterans. + +Veterans are the best of us. + +I’ve always believed that we have a sacred obligation to equip all those we send to war and care for them and their families when they come home. + +My administration is providing assistance with job training and housing, and now helping lower-income veterans get VA care debt-free. + +Our troops in Iraq and Afghanistan faced many dangers. + +One was stationed at bases and breathing in toxic smoke from “burn pits” that incinerated wastes of war—medical and hazard material, jet fuel, and more. + +When they came home, many of the world’s fittest and best trained warriors were never the same. + +Headaches. Numbness. Dizziness. + +A cancer that would put them in a flag-draped coffin. + +I know. + +One of those soldiers was my son Major Beau Biden. + +We don’t know for sure if a burn pit was the cause of his brain cancer, or the diseases of so many of our troops. + +But I’m committed to finding out everything we can. + +Committed to military families like Danielle Robinson from Ohio. + +The widow of Sergeant First Class Heath Robinson. + +He was born a soldier. Army National Guard. Combat medic in Kosovo and Iraq. + +Stationed near Baghdad, just yards from burn pits the size of football fields. + +Heath’s widow Danielle is here with us tonight. They loved going to Ohio State football games. He loved building Legos with their daughter. + +But cancer from prolonged exposure to burn pits ravaged Heath’s lungs and body. + +Danielle says Heath was a fighter to the very end. + +He didn’t know how to stop fighting, and neither did she. + +Through her pain she found purpose to demand we do better. + +Tonight, Danielle—we are. + +The VA is pioneering new ways of linking toxic exposures to diseases, already helping more veterans get benefits. + +And tonight, I’m announcing we’re expanding eligibility to veterans suffering from nine respiratory cancers. + +I’m also calling on Congress: pass a law to make sure veterans devastated by toxic exposures in Iraq and Afghanistan finally get the benefits and comprehensive health care they deserve. + +And fourth, let’s end cancer as we know it. + +This is personal to me and Jill, to Kamala, and to so many of you. + +Cancer is the #2 cause of death in America–second only to heart disease. + +Last month, I announced our plan to supercharge +the Cancer Moonshot that President Obama asked me to lead six years ago. + +Our goal is to cut the cancer death rate by at least 50% over the next 25 years, turn more cancers from death sentences into treatable diseases. + +More support for patients and families. + +To get there, I call on Congress to fund ARPA-H, the Advanced Research Projects Agency for Health. + +It’s based on DARPA—the Defense Department project that led to the Internet, GPS, and so much more. + +ARPA-H will have a singular purpose—to drive breakthroughs in cancer, Alzheimer’s, diabetes, and more. + +A unity agenda for the nation. + +We can do this. + +My fellow Americans—tonight , we have gathered in a sacred space—the citadel of our democracy. + +In this Capitol, generation after generation, Americans have debated great questions amid great strife, and have done great things. + +We have fought for freedom, expanded liberty, defeated totalitarianism and terror. + +And built the strongest, freest, and most prosperous nation the world has ever known. + +Now is the hour. + +Our moment of responsibility. + +Our test of resolve and conscience, of history itself. + +It is in this moment that our character is formed. Our purpose is found. Our future is forged. + +Well I know this nation. + +We will meet the test. + +To protect freedom and liberty, to expand fairness and opportunity. + +We will save democracy. + +As hard as these times have been, I am more optimistic about America today than I have been my whole life. + +Because I see the future that is within our grasp. + +Because I know there is simply nothing beyond our capacity. + +We are the only nation on Earth that has always turned every crisis we have faced into an opportunity. + +The only nation that can be defined by a single word: possibilities. + +So on this night, in our 245th year as a nation, I have come to report on the State of the Union. + +And my report is this: the State of the Union is strong—because you, the American people, are strong. + +We are stronger today than we were a year ago. + +And we will be stronger a year from now than we are today. + +Now is our moment to meet and overcome the challenges of our time. + +And we will, as one people. + +One America. + +The United States of America. + +May God bless you all. May God protect our troops. diff --git a/ios/.gitignore b/ios/.gitignore new file mode 100644 index 0000000..f75e367 --- /dev/null +++ b/ios/.gitignore @@ -0,0 +1,3 @@ +xuserdata +MLCSwift/tvm_home +*~ diff --git a/ios/MLCChat/MLCChat.xcodeproj/project.pbxproj b/ios/MLCChat/MLCChat.xcodeproj/project.pbxproj new file mode 100644 index 0000000..24b8dc0 --- /dev/null +++ b/ios/MLCChat/MLCChat.xcodeproj/project.pbxproj @@ -0,0 +1,560 @@ +// !$*UTF8*$! +{ + archiveVersion = 1; + classes = { + }; + objectVersion = 60; + objects = { + +/* Begin PBXBuildFile section */ + 1453A4CF2A1354B9001B909F /* StartView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1453A4CA2A1354B9001B909F /* StartView.swift */; }; + 1453A4D02A1354B9001B909F /* ModelView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1453A4CB2A1354B9001B909F /* ModelView.swift */; }; + 1453A4D12A1354B9001B909F /* AppState.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1453A4CC2A1354B9001B909F /* AppState.swift */; }; + 1453A4D22A1354B9001B909F /* ModelConfig.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1453A4CD2A1354B9001B909F /* ModelConfig.swift */; }; + 1453A4D32A1354B9001B909F /* ModelState.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1453A4CE2A1354B9001B909F /* ModelState.swift */; }; + A773CC652A5DC98200467BFE /* ImageProcessing.swift in Sources */ = {isa = PBXBuildFile; fileRef = A773CC642A5DC98200467BFE /* ImageProcessing.swift */; }; + AEC27EFA2A85C2AC00254E67 /* ParamsConfig.swift in Sources */ = {isa = PBXBuildFile; fileRef = AEC27EF92A85C2AC00254E67 /* ParamsConfig.swift */; }; + AEC27EFC2A85C3B000254E67 /* AppConfig.swift in Sources */ = {isa = PBXBuildFile; fileRef = AEC27EFB2A85C3B000254E67 /* AppConfig.swift */; }; + AEC27F022A86337E00254E67 /* Constants.swift in Sources */ = {isa = PBXBuildFile; fileRef = AEC27F012A86337E00254E67 /* Constants.swift */; }; + B08647022C6D0293001A8B5E /* MarkdownUI in Frameworks */ = {isa = PBXBuildFile; productRef = B08647012C6D0293001A8B5E /* MarkdownUI */; }; + C04105DD2BEBBEA6005A434D /* MLCSwift in Frameworks */ = {isa = PBXBuildFile; productRef = C04105DC2BEBBEA6005A434D /* MLCSwift */; }; + C0D643B329F99A7F004DDAA4 /* MLCChatApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = C0D643B229F99A7F004DDAA4 /* MLCChatApp.swift */; }; + C0D643B729F99A80004DDAA4 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C0D643B629F99A80004DDAA4 /* Assets.xcassets */; }; + C0D643BA29F99A80004DDAA4 /* Preview Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C0D643B929F99A80004DDAA4 /* Preview Assets.xcassets */; }; + C0D643C429F99B07004DDAA4 /* ChatView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C0D643C229F99B07004DDAA4 /* ChatView.swift */; }; + C0D643C829F99B34004DDAA4 /* MessageView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C0D643C729F99B34004DDAA4 /* MessageView.swift */; }; + C0DDBDF62A39103F00E9D060 /* ChatState.swift in Sources */ = {isa = PBXBuildFile; fileRef = C0D643C029F99B07004DDAA4 /* ChatState.swift */; }; + F3C280002BEB16ED00F1E016 /* bundle in CopyFiles */ = {isa = PBXBuildFile; fileRef = F3C27FFF2BEB16ED00F1E016 /* bundle */; }; +/* End PBXBuildFile section */ + +/* Begin PBXCopyFilesBuildPhase section */ + C06A74F129F9A78000BC4BE6 /* CopyFiles */ = { + isa = PBXCopyFilesBuildPhase; + buildActionMask = 2147483647; + dstPath = ""; + dstSubfolderSpec = 7; + files = ( + F3C280002BEB16ED00F1E016 /* bundle in CopyFiles */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; + C0D643CF29F99C5D004DDAA4 /* Embed Libraries */ = { + isa = PBXCopyFilesBuildPhase; + buildActionMask = 2147483647; + dstPath = ""; + dstSubfolderSpec = 10; + files = ( + ); + name = "Embed Libraries"; + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXCopyFilesBuildPhase section */ + +/* Begin PBXFileReference section */ + 1453A4CA2A1354B9001B909F /* StartView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = StartView.swift; sourceTree = ""; }; + 1453A4CB2A1354B9001B909F /* ModelView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ModelView.swift; sourceTree = ""; }; + 1453A4CC2A1354B9001B909F /* AppState.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AppState.swift; sourceTree = ""; }; + 1453A4CD2A1354B9001B909F /* ModelConfig.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ModelConfig.swift; sourceTree = ""; }; + 1453A4CE2A1354B9001B909F /* ModelState.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ModelState.swift; sourceTree = ""; }; + A773CC642A5DC98200467BFE /* ImageProcessing.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ImageProcessing.swift; sourceTree = ""; }; + AEC27EF92A85C2AC00254E67 /* ParamsConfig.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ParamsConfig.swift; sourceTree = ""; }; + AEC27EFB2A85C3B000254E67 /* AppConfig.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppConfig.swift; sourceTree = ""; }; + AEC27F012A86337E00254E67 /* Constants.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Constants.swift; sourceTree = ""; }; + C06A74E629F9A1DF00BC4BE6 /* MLCChat.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = MLCChat.entitlements; sourceTree = ""; }; + C0D643AF29F99A7F004DDAA4 /* MLCChat.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = MLCChat.app; sourceTree = BUILT_PRODUCTS_DIR; }; + C0D643B229F99A7F004DDAA4 /* MLCChatApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MLCChatApp.swift; sourceTree = ""; }; + C0D643B629F99A80004DDAA4 /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = ""; }; + C0D643B929F99A80004DDAA4 /* Preview Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = "Preview Assets.xcassets"; sourceTree = ""; }; + C0D643C029F99B07004DDAA4 /* ChatState.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ChatState.swift; sourceTree = ""; }; + C0D643C229F99B07004DDAA4 /* ChatView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = ChatView.swift; sourceTree = ""; }; + C0D643C729F99B34004DDAA4 /* MessageView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = MessageView.swift; sourceTree = ""; }; + C0DDBE0B2A3BA6F800E9D060 /* MLCSwift */ = {isa = PBXFileReference; lastKnownFileType = wrapper; path = MLCSwift; sourceTree = ""; }; + F3C27FFF2BEB16ED00F1E016 /* bundle */ = {isa = PBXFileReference; lastKnownFileType = folder; name = bundle; path = dist/bundle; sourceTree = ""; }; +/* End PBXFileReference section */ + +/* Begin PBXFrameworksBuildPhase section */ + C0D643AC29F99A7F004DDAA4 /* Frameworks */ = { + isa = PBXFrameworksBuildPhase; + buildActionMask = 2147483647; + files = ( + C04105DD2BEBBEA6005A434D /* MLCSwift in Frameworks */, + B08647022C6D0293001A8B5E /* MarkdownUI in Frameworks */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXFrameworksBuildPhase section */ + +/* Begin PBXGroup section */ + AEC27EF82A85C29000254E67 /* Models */ = { + isa = PBXGroup; + children = ( + 1453A4CD2A1354B9001B909F /* ModelConfig.swift */, + AEC27EF92A85C2AC00254E67 /* ParamsConfig.swift */, + AEC27EFB2A85C3B000254E67 /* AppConfig.swift */, + ); + path = Models; + sourceTree = ""; + }; + AEC27EFF2A85EE2800254E67 /* States */ = { + isa = PBXGroup; + children = ( + 1453A4CE2A1354B9001B909F /* ModelState.swift */, + 1453A4CC2A1354B9001B909F /* AppState.swift */, + C0D643C029F99B07004DDAA4 /* ChatState.swift */, + ); + path = States; + sourceTree = ""; + }; + AEC27F002A86306800254E67 /* Views */ = { + isa = PBXGroup; + children = ( + A773CC642A5DC98200467BFE /* ImageProcessing.swift */, + 1453A4CB2A1354B9001B909F /* ModelView.swift */, + 1453A4CA2A1354B9001B909F /* StartView.swift */, + C0D643C729F99B34004DDAA4 /* MessageView.swift */, + C0D643C229F99B07004DDAA4 /* ChatView.swift */, + ); + path = Views; + sourceTree = ""; + }; + AEC27F032A86338800254E67 /* Common */ = { + isa = PBXGroup; + children = ( + AEC27F012A86337E00254E67 /* Constants.swift */, + ); + path = Common; + sourceTree = ""; + }; + C0D643A629F99A7F004DDAA4 = { + isa = PBXGroup; + children = ( + F3C27FFF2BEB16ED00F1E016 /* bundle */, + C0DDBDF02A39068900E9D060 /* Packages */, + C0D643B129F99A7F004DDAA4 /* MLCChat */, + C0D643B029F99A7F004DDAA4 /* Products */, + C0D643C929F99BDA004DDAA4 /* Frameworks */, + ); + sourceTree = ""; + }; + C0D643B029F99A7F004DDAA4 /* Products */ = { + isa = PBXGroup; + children = ( + C0D643AF29F99A7F004DDAA4 /* MLCChat.app */, + ); + name = Products; + sourceTree = ""; + }; + C0D643B129F99A7F004DDAA4 /* MLCChat */ = { + isa = PBXGroup; + children = ( + AEC27F032A86338800254E67 /* Common */, + AEC27EF82A85C29000254E67 /* Models */, + AEC27EFF2A85EE2800254E67 /* States */, + AEC27F002A86306800254E67 /* Views */, + C06A74E629F9A1DF00BC4BE6 /* MLCChat.entitlements */, + C0D643B229F99A7F004DDAA4 /* MLCChatApp.swift */, + C0D643B629F99A80004DDAA4 /* Assets.xcassets */, + C0D643B829F99A80004DDAA4 /* Preview Content */, + ); + path = MLCChat; + sourceTree = ""; + }; + C0D643B829F99A80004DDAA4 /* Preview Content */ = { + isa = PBXGroup; + children = ( + C0D643B929F99A80004DDAA4 /* Preview Assets.xcassets */, + ); + path = "Preview Content"; + sourceTree = ""; + }; + C0D643C929F99BDA004DDAA4 /* Frameworks */ = { + isa = PBXGroup; + children = ( + ); + name = Frameworks; + sourceTree = ""; + }; + C0DDBDF02A39068900E9D060 /* Packages */ = { + isa = PBXGroup; + children = ( + C0DDBE0B2A3BA6F800E9D060 /* MLCSwift */, + ); + name = Packages; + sourceTree = ""; + }; +/* End PBXGroup section */ + +/* Begin PBXNativeTarget section */ + C0D643AE29F99A7F004DDAA4 /* MLCChat */ = { + isa = PBXNativeTarget; + buildConfigurationList = C0D643BD29F99A80004DDAA4 /* Build configuration list for PBXNativeTarget "MLCChat" */; + buildPhases = ( + C0D643AB29F99A7F004DDAA4 /* Sources */, + C0D643AC29F99A7F004DDAA4 /* Frameworks */, + C0D643AD29F99A7F004DDAA4 /* Resources */, + C0D643CF29F99C5D004DDAA4 /* Embed Libraries */, + C06A74F129F9A78000BC4BE6 /* CopyFiles */, + ); + buildRules = ( + ); + dependencies = ( + ); + name = MLCChat; + packageProductDependencies = ( + C04105DC2BEBBEA6005A434D /* MLCSwift */, + B08647012C6D0293001A8B5E /* MarkdownUI */, + ); + productName = MLCChat; + productReference = C0D643AF29F99A7F004DDAA4 /* MLCChat.app */; + productType = "com.apple.product-type.application"; + }; +/* End PBXNativeTarget section */ + +/* Begin PBXProject section */ + C0D643A729F99A7F004DDAA4 /* Project object */ = { + isa = PBXProject; + attributes = { + BuildIndependentTargetsInParallel = 1; + LastSwiftUpdateCheck = 1430; + LastUpgradeCheck = 1430; + TargetAttributes = { + C0D643AE29F99A7F004DDAA4 = { + CreatedOnToolsVersion = 14.3; + LastSwiftMigration = 1430; + }; + }; + }; + buildConfigurationList = C0D643AA29F99A7F004DDAA4 /* Build configuration list for PBXProject "MLCChat" */; + compatibilityVersion = "Xcode 14.0"; + developmentRegion = en; + hasScannedForEncodings = 0; + knownRegions = ( + en, + Base, + ); + mainGroup = C0D643A629F99A7F004DDAA4; + packageReferences = ( + C04105DB2BEBBEA6005A434D /* XCLocalSwiftPackageReference "../MLCSwift" */, + B08647002C6D0293001A8B5E /* XCRemoteSwiftPackageReference "swift-markdown-ui" */, + ); + productRefGroup = C0D643B029F99A7F004DDAA4 /* Products */; + projectDirPath = ""; + projectRoot = ""; + targets = ( + C0D643AE29F99A7F004DDAA4 /* MLCChat */, + ); + }; +/* End PBXProject section */ + +/* Begin PBXResourcesBuildPhase section */ + C0D643AD29F99A7F004DDAA4 /* Resources */ = { + isa = PBXResourcesBuildPhase; + buildActionMask = 2147483647; + files = ( + C0D643BA29F99A80004DDAA4 /* Preview Assets.xcassets in Resources */, + C0D643B729F99A80004DDAA4 /* Assets.xcassets in Resources */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXResourcesBuildPhase section */ + +/* Begin PBXSourcesBuildPhase section */ + C0D643AB29F99A7F004DDAA4 /* Sources */ = { + isa = PBXSourcesBuildPhase; + buildActionMask = 2147483647; + files = ( + A773CC652A5DC98200467BFE /* ImageProcessing.swift in Sources */, + 1453A4D12A1354B9001B909F /* AppState.swift in Sources */, + C0D643B329F99A7F004DDAA4 /* MLCChatApp.swift in Sources */, + C0DDBDF62A39103F00E9D060 /* ChatState.swift in Sources */, + C0D643C429F99B07004DDAA4 /* ChatView.swift in Sources */, + 1453A4D32A1354B9001B909F /* ModelState.swift in Sources */, + C0D643C829F99B34004DDAA4 /* MessageView.swift in Sources */, + 1453A4D22A1354B9001B909F /* ModelConfig.swift in Sources */, + AEC27EFA2A85C2AC00254E67 /* ParamsConfig.swift in Sources */, + AEC27EFC2A85C3B000254E67 /* AppConfig.swift in Sources */, + AEC27F022A86337E00254E67 /* Constants.swift in Sources */, + 1453A4D02A1354B9001B909F /* ModelView.swift in Sources */, + 1453A4CF2A1354B9001B909F /* StartView.swift in Sources */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXSourcesBuildPhase section */ + +/* Begin XCBuildConfiguration section */ + C0D643BB29F99A80004DDAA4 /* Debug */ = { + isa = XCBuildConfiguration; + buildSettings = { + ALWAYS_SEARCH_USER_PATHS = NO; + CLANG_ANALYZER_NONNULL = YES; + CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; + CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; + CLANG_ENABLE_MODULES = YES; + CLANG_ENABLE_OBJC_ARC = YES; + CLANG_ENABLE_OBJC_WEAK = YES; + CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; + CLANG_WARN_BOOL_CONVERSION = YES; + CLANG_WARN_COMMA = YES; + CLANG_WARN_CONSTANT_CONVERSION = YES; + CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; + CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; + CLANG_WARN_DOCUMENTATION_COMMENTS = YES; + CLANG_WARN_EMPTY_BODY = YES; + CLANG_WARN_ENUM_CONVERSION = YES; + CLANG_WARN_INFINITE_RECURSION = YES; + CLANG_WARN_INT_CONVERSION = YES; + CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; + CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; + CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; + CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; + CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; + CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; + CLANG_WARN_STRICT_PROTOTYPES = YES; + CLANG_WARN_SUSPICIOUS_MOVE = YES; + CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; + CLANG_WARN_UNREACHABLE_CODE = YES; + CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; + COPY_PHASE_STRIP = NO; + DEBUG_INFORMATION_FORMAT = dwarf; + ENABLE_STRICT_OBJC_MSGSEND = YES; + ENABLE_TESTABILITY = YES; + GCC_C_LANGUAGE_STANDARD = gnu11; + GCC_DYNAMIC_NO_PIC = NO; + GCC_NO_COMMON_BLOCKS = YES; + GCC_OPTIMIZATION_LEVEL = 0; + GCC_PREPROCESSOR_DEFINITIONS = ( + "DEBUG=1", + "$(inherited)", + ); + GCC_WARN_64_TO_32_BIT_CONVERSION = YES; + GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; + GCC_WARN_UNDECLARED_SELECTOR = YES; + GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; + GCC_WARN_UNUSED_FUNCTION = YES; + GCC_WARN_UNUSED_VARIABLE = YES; + IPHONEOS_DEPLOYMENT_TARGET = 16.0; + MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; + MTL_FAST_MATH = YES; + ONLY_ACTIVE_ARCH = YES; + SDKROOT = iphoneos; + SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG; + SWIFT_OPTIMIZATION_LEVEL = "-Onone"; + }; + name = Debug; + }; + C0D643BC29F99A80004DDAA4 /* Release */ = { + isa = XCBuildConfiguration; + buildSettings = { + ALWAYS_SEARCH_USER_PATHS = NO; + CLANG_ANALYZER_NONNULL = YES; + CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; + CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; + CLANG_ENABLE_MODULES = YES; + CLANG_ENABLE_OBJC_ARC = YES; + CLANG_ENABLE_OBJC_WEAK = YES; + CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; + CLANG_WARN_BOOL_CONVERSION = YES; + CLANG_WARN_COMMA = YES; + CLANG_WARN_CONSTANT_CONVERSION = YES; + CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; + CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; + CLANG_WARN_DOCUMENTATION_COMMENTS = YES; + CLANG_WARN_EMPTY_BODY = YES; + CLANG_WARN_ENUM_CONVERSION = YES; + CLANG_WARN_INFINITE_RECURSION = YES; + CLANG_WARN_INT_CONVERSION = YES; + CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; + CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; + CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; + CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; + CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; + CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; + CLANG_WARN_STRICT_PROTOTYPES = YES; + CLANG_WARN_SUSPICIOUS_MOVE = YES; + CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; + CLANG_WARN_UNREACHABLE_CODE = YES; + CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; + COPY_PHASE_STRIP = NO; + DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; + ENABLE_NS_ASSERTIONS = NO; + ENABLE_STRICT_OBJC_MSGSEND = YES; + GCC_C_LANGUAGE_STANDARD = gnu11; + GCC_NO_COMMON_BLOCKS = YES; + GCC_WARN_64_TO_32_BIT_CONVERSION = YES; + GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; + GCC_WARN_UNDECLARED_SELECTOR = YES; + GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; + GCC_WARN_UNUSED_FUNCTION = YES; + GCC_WARN_UNUSED_VARIABLE = YES; + IPHONEOS_DEPLOYMENT_TARGET = 16.0; + MTL_ENABLE_DEBUG_INFO = NO; + MTL_FAST_MATH = YES; + SDKROOT = iphoneos; + SWIFT_COMPILATION_MODE = wholemodule; + SWIFT_OPTIMIZATION_LEVEL = "-O"; + VALIDATE_PRODUCT = YES; + }; + name = Release; + }; + C0D643BE29F99A80004DDAA4 /* Debug */ = { + isa = XCBuildConfiguration; + buildSettings = { + ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; + ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; + CLANG_ENABLE_MODULES = YES; + CODE_SIGN_ENTITLEMENTS = MLCChat/MLCChat.entitlements; + CODE_SIGN_IDENTITY = "Apple Development"; + CODE_SIGN_STYLE = Automatic; + CURRENT_PROJECT_VERSION = 1; + DEVELOPMENT_ASSET_PATHS = "\"MLCChat/Preview Content\""; + DEVELOPMENT_TEAM = 3FR42MXLK9; + ENABLE_PREVIEWS = YES; + GENERATE_INFOPLIST_FILE = YES; + "HEADER_SEARCH_PATHS[arch=*]" = ""; + INFOPLIST_FILE = MLCChat/Info.plist; + INFOPLIST_KEY_LSApplicationCategoryType = "public.app-category.productivity"; + INFOPLIST_KEY_NSCameraUsageDescription = "This app requires usage of camera to function properly."; + INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES; + INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; + INFOPLIST_KEY_UILaunchScreen_Generation = YES; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + IPHONEOS_DEPLOYMENT_TARGET = 17.0; + LD_RUNPATH_SEARCH_PATHS = ( + "$(inherited)", + "@executable_path/Frameworks", + ); + LIBRARY_SEARCH_PATHS = ( + "$(inherited)", + "$(PROJECT_DIR)/dist/lib", + ); + MARKETING_VERSION = 1.6; + OTHER_LDFLAGS = ( + "-Wl,-all_load", + "-lmodel_iphone", + "-lmlc_llm", + "-ltvm_ffi_static", + "-ltvm_runtime", + "-ltokenizers_cpp", + "-lsentencepiece", + "-ltokenizers_c", + ); + PRODUCT_BUNDLE_IDENTIFIER = mlc.Chat; + PRODUCT_NAME = "$(TARGET_NAME)"; + PROVISIONING_PROFILE_SPECIFIER = ""; + SWIFT_EMIT_LOC_STRINGS = YES; + SWIFT_OBJC_BRIDGING_HEADER = ""; + SWIFT_OPTIMIZATION_LEVEL = "-Onone"; + SWIFT_VERSION = 5.0; + TARGETED_DEVICE_FAMILY = "1,2"; + }; + name = Debug; + }; + C0D643BF29F99A80004DDAA4 /* Release */ = { + isa = XCBuildConfiguration; + buildSettings = { + ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; + ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; + CLANG_ENABLE_MODULES = YES; + CODE_SIGN_ENTITLEMENTS = MLCChat/MLCChat.entitlements; + CODE_SIGN_IDENTITY = "Apple Development"; + CODE_SIGN_STYLE = Automatic; + CURRENT_PROJECT_VERSION = 1; + DEVELOPMENT_ASSET_PATHS = "\"MLCChat/Preview Content\""; + DEVELOPMENT_TEAM = 3FR42MXLK9; + ENABLE_PREVIEWS = YES; + GENERATE_INFOPLIST_FILE = YES; + "HEADER_SEARCH_PATHS[arch=*]" = ""; + INFOPLIST_FILE = MLCChat/Info.plist; + INFOPLIST_KEY_LSApplicationCategoryType = "public.app-category.productivity"; + INFOPLIST_KEY_NSCameraUsageDescription = "This app requires usage of camera to function properly."; + INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES; + INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; + INFOPLIST_KEY_UILaunchScreen_Generation = YES; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + IPHONEOS_DEPLOYMENT_TARGET = 17.0; + LD_RUNPATH_SEARCH_PATHS = ( + "$(inherited)", + "@executable_path/Frameworks", + ); + LIBRARY_SEARCH_PATHS = ( + "$(inherited)", + "$(PROJECT_DIR)/dist/lib", + ); + MARKETING_VERSION = 1.6; + OTHER_LDFLAGS = ( + "-Wl,-all_load", + "-lmodel_iphone", + "-lmlc_llm", + "-ltvm_ffi_static", + "-ltvm_runtime", + "-ltokenizers_cpp", + "-lsentencepiece", + "-ltokenizers_c", + ); + PRODUCT_BUNDLE_IDENTIFIER = mlc.Chat; + PRODUCT_NAME = "$(TARGET_NAME)"; + PROVISIONING_PROFILE_SPECIFIER = ""; + SWIFT_EMIT_LOC_STRINGS = YES; + SWIFT_OBJC_BRIDGING_HEADER = ""; + SWIFT_VERSION = 5.0; + TARGETED_DEVICE_FAMILY = "1,2"; + }; + name = Release; + }; +/* End XCBuildConfiguration section */ + +/* Begin XCConfigurationList section */ + C0D643AA29F99A7F004DDAA4 /* Build configuration list for PBXProject "MLCChat" */ = { + isa = XCConfigurationList; + buildConfigurations = ( + C0D643BB29F99A80004DDAA4 /* Debug */, + C0D643BC29F99A80004DDAA4 /* Release */, + ); + defaultConfigurationIsVisible = 0; + defaultConfigurationName = Release; + }; + C0D643BD29F99A80004DDAA4 /* Build configuration list for PBXNativeTarget "MLCChat" */ = { + isa = XCConfigurationList; + buildConfigurations = ( + C0D643BE29F99A80004DDAA4 /* Debug */, + C0D643BF29F99A80004DDAA4 /* Release */, + ); + defaultConfigurationIsVisible = 0; + defaultConfigurationName = Release; + }; +/* End XCConfigurationList section */ + +/* Begin XCLocalSwiftPackageReference section */ + C04105DB2BEBBEA6005A434D /* XCLocalSwiftPackageReference "../MLCSwift" */ = { + isa = XCLocalSwiftPackageReference; + relativePath = ../MLCSwift; + }; +/* End XCLocalSwiftPackageReference section */ + +/* Begin XCRemoteSwiftPackageReference section */ + B08647002C6D0293001A8B5E /* XCRemoteSwiftPackageReference "swift-markdown-ui" */ = { + isa = XCRemoteSwiftPackageReference; + repositoryURL = "https://github.com/gonzalezreal/swift-markdown-ui"; + requirement = { + kind = upToNextMajorVersion; + minimumVersion = 2.4.0; + }; + }; +/* End XCRemoteSwiftPackageReference section */ + +/* Begin XCSwiftPackageProductDependency section */ + B08647012C6D0293001A8B5E /* MarkdownUI */ = { + isa = XCSwiftPackageProductDependency; + package = B08647002C6D0293001A8B5E /* XCRemoteSwiftPackageReference "swift-markdown-ui" */; + productName = MarkdownUI; + }; + C04105DC2BEBBEA6005A434D /* MLCSwift */ = { + isa = XCSwiftPackageProductDependency; + productName = MLCSwift; + }; +/* End XCSwiftPackageProductDependency section */ + }; + rootObject = C0D643A729F99A7F004DDAA4 /* Project object */; +} diff --git a/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/contents.xcworkspacedata b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/contents.xcworkspacedata new file mode 100644 index 0000000..919434a --- /dev/null +++ b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/contents.xcworkspacedata @@ -0,0 +1,7 @@ + + + + + diff --git a/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist new file mode 100644 index 0000000..3d4c1e5 --- /dev/null +++ b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist @@ -0,0 +1,8 @@ + + + + + IDEDidComputeMac32BitWarning + + + diff --git a/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings new file mode 100644 index 0000000..0c67376 --- /dev/null +++ b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings @@ -0,0 +1,5 @@ + + + + + diff --git a/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved new file mode 100644 index 0000000..ea432d3 --- /dev/null +++ b/ios/MLCChat/MLCChat.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved @@ -0,0 +1,33 @@ +{ + "originHash" : "1fe9890dccb6a9581bfc88793e2dc9fc7b4589153f379b73d5ac1114daef8442", + "pins" : [ + { + "identity" : "networkimage", + "kind" : "remoteSourceControl", + "location" : "https://github.com/gonzalezreal/NetworkImage", + "state" : { + "revision" : "2849f5323265386e200484b0d0f896e73c3411b9", + "version" : "6.0.1" + } + }, + { + "identity" : "swift-cmark", + "kind" : "remoteSourceControl", + "location" : "https://github.com/swiftlang/swift-cmark", + "state" : { + "revision" : "5d9bdaa4228b381639fff09403e39a04926e2dbe", + "version" : "0.7.1" + } + }, + { + "identity" : "swift-markdown-ui", + "kind" : "remoteSourceControl", + "location" : "https://github.com/gonzalezreal/swift-markdown-ui", + "state" : { + "revision" : "5f613358148239d0292c0cef674a3c2314737f9e", + "version" : "2.4.1" + } + } + ], + "version" : 3 +} diff --git a/ios/MLCChat/MLCChat.xcodeproj/xcshareddata/xcschemes/MLCChat.xcscheme b/ios/MLCChat/MLCChat.xcodeproj/xcshareddata/xcschemes/MLCChat.xcscheme new file mode 100644 index 0000000..dba8e2d --- /dev/null +++ b/ios/MLCChat/MLCChat.xcodeproj/xcshareddata/xcschemes/MLCChat.xcscheme @@ -0,0 +1,81 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/ios/MLCChat/MLCChat/Assets.xcassets/AccentColor.colorset/Contents.json b/ios/MLCChat/MLCChat/Assets.xcassets/AccentColor.colorset/Contents.json new file mode 100644 index 0000000..eb87897 --- /dev/null +++ b/ios/MLCChat/MLCChat/Assets.xcassets/AccentColor.colorset/Contents.json @@ -0,0 +1,11 @@ +{ + "colors" : [ + { + "idiom" : "universal" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCChat/MLCChat/Assets.xcassets/AppIcon.appiconset/Contents.json b/ios/MLCChat/MLCChat/Assets.xcassets/AppIcon.appiconset/Contents.json new file mode 100644 index 0000000..7324dc2 --- /dev/null +++ b/ios/MLCChat/MLCChat/Assets.xcassets/AppIcon.appiconset/Contents.json @@ -0,0 +1,14 @@ +{ + "images" : [ + { + "filename" : "mlc-logo.png", + "idiom" : "universal", + "platform" : "ios", + "size" : "1024x1024" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCChat/MLCChat/Assets.xcassets/AppIcon.appiconset/mlc-logo.png b/ios/MLCChat/MLCChat/Assets.xcassets/AppIcon.appiconset/mlc-logo.png new file mode 100644 index 0000000..ab37f8a Binary files /dev/null and b/ios/MLCChat/MLCChat/Assets.xcassets/AppIcon.appiconset/mlc-logo.png differ diff --git a/ios/MLCChat/MLCChat/Assets.xcassets/Contents.json b/ios/MLCChat/MLCChat/Assets.xcassets/Contents.json new file mode 100644 index 0000000..73c0059 --- /dev/null +++ b/ios/MLCChat/MLCChat/Assets.xcassets/Contents.json @@ -0,0 +1,6 @@ +{ + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCChat/MLCChat/Common/Constants.swift b/ios/MLCChat/MLCChat/Common/Constants.swift new file mode 100644 index 0000000..c606271 --- /dev/null +++ b/ios/MLCChat/MLCChat/Common/Constants.swift @@ -0,0 +1,11 @@ +// +// Constants.swift +// MLCChat +// + +struct Constants { + static let prebuiltModelDir = "bundle" + static let appConfigFileName = "bundle/mlc-app-config.json" + static let modelConfigFileName = "mlc-chat-config.json" + static let paramsConfigFileName = "tensor-cache.json" +} diff --git a/ios/MLCChat/MLCChat/Info.plist b/ios/MLCChat/MLCChat/Info.plist new file mode 100644 index 0000000..0c67376 --- /dev/null +++ b/ios/MLCChat/MLCChat/Info.plist @@ -0,0 +1,5 @@ + + + + + diff --git a/ios/MLCChat/MLCChat/MLCChat.entitlements b/ios/MLCChat/MLCChat/MLCChat.entitlements new file mode 100644 index 0000000..1250534 --- /dev/null +++ b/ios/MLCChat/MLCChat/MLCChat.entitlements @@ -0,0 +1,10 @@ + + + + + com.apple.developer.kernel.extended-virtual-addressing + + com.apple.developer.kernel.increased-memory-limit + + + diff --git a/ios/MLCChat/MLCChat/MLCChatApp.swift b/ios/MLCChat/MLCChat/MLCChatApp.swift new file mode 100644 index 0000000..fcefd6f --- /dev/null +++ b/ios/MLCChat/MLCChat/MLCChatApp.swift @@ -0,0 +1,28 @@ +// +// MLCChatApp.swift +// MLCChat +// +// Created by Tianqi Chen on 4/26/23. +// + +import SwiftUI + +@main +struct MLCChatApp: App { + @StateObject private var appState = AppState() + + init() { + UITableView.appearance().separatorStyle = .none + UITableView.appearance().tableFooterView = UIView() + } + + var body: some Scene { + WindowGroup { + StartView() + .environmentObject(appState) + .task { + appState.loadAppConfigAndModels() + } + } + } +} diff --git a/ios/MLCChat/MLCChat/Models/AppConfig.swift b/ios/MLCChat/MLCChat/Models/AppConfig.swift new file mode 100644 index 0000000..69867b0 --- /dev/null +++ b/ios/MLCChat/MLCChat/Models/AppConfig.swift @@ -0,0 +1,28 @@ +// +// AppConfig.swift +// MLCChat +// + +struct AppConfig: Codable { + struct ModelRecord: Codable { + let modelPath: String? + let modelURL: String? + let modelLib: String + let estimatedVRAMReq: Int + let modelID: String + + enum CodingKeys: String, CodingKey { + case modelPath = "model_path" + case modelURL = "model_url" + case modelLib = "model_lib" + case estimatedVRAMReq = "estimated_vram_bytes" + case modelID = "model_id" + } + } + + var modelList: [ModelRecord] + + enum CodingKeys: String, CodingKey { + case modelList = "model_list" + } +} diff --git a/ios/MLCChat/MLCChat/Models/ModelConfig.swift b/ios/MLCChat/MLCChat/Models/ModelConfig.swift new file mode 100644 index 0000000..4ed8819 --- /dev/null +++ b/ios/MLCChat/MLCChat/Models/ModelConfig.swift @@ -0,0 +1,18 @@ +// +// ModelConfig.swift +// MLCChat +// + +struct ModelConfig: Decodable { + let tokenizerFiles: [String] + var modelLib: String? + var modelID: String? + var estimatedVRAMReq: Int? + + enum CodingKeys: String, CodingKey { + case tokenizerFiles = "tokenizer_files" + case modelLib = "model_lib" + case modelID = "model_id" + case estimatedVRAMReq = "estimated_vram_req" + } +} diff --git a/ios/MLCChat/MLCChat/Models/ParamsConfig.swift b/ios/MLCChat/MLCChat/Models/ParamsConfig.swift new file mode 100644 index 0000000..2635afa --- /dev/null +++ b/ios/MLCChat/MLCChat/Models/ParamsConfig.swift @@ -0,0 +1,12 @@ +// +// ParamsConfig.swift +// MLCChat +// + +struct ParamsConfig: Decodable { + struct ParamsRecord: Decodable { + let dataPath: String + } + + let records: [ParamsRecord] +} diff --git a/ios/MLCChat/MLCChat/Preview Content/Preview Assets.xcassets/Contents.json b/ios/MLCChat/MLCChat/Preview Content/Preview Assets.xcassets/Contents.json new file mode 100644 index 0000000..73c0059 --- /dev/null +++ b/ios/MLCChat/MLCChat/Preview Content/Preview Assets.xcassets/Contents.json @@ -0,0 +1,6 @@ +{ + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCChat/MLCChat/States/AppState.swift b/ios/MLCChat/MLCChat/States/AppState.swift new file mode 100644 index 0000000..b5b060b --- /dev/null +++ b/ios/MLCChat/MLCChat/States/AppState.swift @@ -0,0 +1,274 @@ +// +// AppState.swift +// MLCChat +// +// Created by Yaxing Cai on 5/13/23. +// + +import Foundation + +final class AppState: ObservableObject { + @Published var models = [ModelState]() + @Published var chatState = ChatState() + + @Published var alertMessage = "" // TODO: Should move out + @Published var alertDisplayed = false // TODO: Should move out + + private var appConfig: AppConfig? + private var modelIDs = Set() + + private let fileManager: FileManager = FileManager.default + private lazy var cacheDirectoryURL: URL = { + fileManager.urls(for: .cachesDirectory, in: .userDomainMask)[0] + }() + + private let jsonDecoder = JSONDecoder() + private let jsonEncoder = JSONEncoder() + + func loadAppConfigAndModels() { + appConfig = loadAppConfig() + // Can't do anything without a valid app config + guard let appConfig else { + return + } + loadModelsConfig(modelList: appConfig.modelList) + } + + func requestDeleteModel(modelID: String) { + // model dir should have been deleted in ModelState + assert(!fileManager.fileExists(atPath: cacheDirectoryURL.appending(path: modelID).path())) + modelIDs.remove(modelID) + models.removeAll(where: {$0.modelConfig.modelID == modelID}) + updateAppConfig { + appConfig?.modelList.removeAll(where: {$0.modelID == modelID}) + } + } +} + +private extension AppState { + func loadAppConfig() -> AppConfig? { + // models in cache to download + var appConfigFileURL = cacheDirectoryURL.appending(path: Constants.appConfigFileName) + if !fileManager.fileExists(atPath: appConfigFileURL.path()) { + appConfigFileURL = Bundle.main.bundleURL.appending(path: Constants.appConfigFileName) + } + assert(fileManager.fileExists(atPath: appConfigFileURL.path())) + + do { + let fileHandle = try FileHandle(forReadingFrom: appConfigFileURL) + let data = fileHandle.readDataToEndOfFile() + + let appConfig = try jsonDecoder.decode(AppConfig.self, from: data) + return appConfig + } catch { + showAlert(message: "Failed to load app config: \(error.localizedDescription)") + return nil + } + } + + func loadModelsConfig(modelList: [AppConfig.ModelRecord]) { + for model in modelList { + if model.modelPath != nil { + // local model + let modelDir = Bundle.main.bundleURL.appending(path: Constants.prebuiltModelDir).appending(path: model.modelPath!) + let modelConfigURL = modelDir.appending(path: Constants.modelConfigFileName) + if fileManager.fileExists(atPath: modelConfigURL.path()) { + if let modelConfig = loadModelConfig( + modelConfigURL: modelConfigURL, + modelLib: model.modelLib, + modelID: model.modelID, + estimatedVRAMReq: model.estimatedVRAMReq + ) { + addModelConfig( + modelConfig: modelConfig, + modelPath: model.modelPath!, + modelURL: nil, + isBuiltin: true + ) + } else { + showAlert(message: "Failed to load prebuilt model: \(model.modelPath!)") + } + } else { + showAlert(message: "Prebuilt mlc-chat-config.json file not found: \(model.modelPath!)") + } + } else if model.modelURL != nil { + // remote model + let modelConfigFileURL = cacheDirectoryURL + .appending(path: model.modelID) + .appending(path: Constants.modelConfigFileName) + if fileManager.fileExists(atPath: modelConfigFileURL.path()) { + if let modelConfig = loadModelConfig( + modelConfigURL: modelConfigFileURL, + modelLib: model.modelLib, + modelID: model.modelID, + estimatedVRAMReq: model.estimatedVRAMReq + ) { + addModelConfig( + modelConfig: modelConfig, + modelPath: nil, + modelURL: URL(string: model.modelURL!), + isBuiltin: true + ) + } + } else { + downloadConfig( + modelURL: URL(string: model.modelURL!), + modelLib: model.modelLib, + modelID: model.modelID, + estimatedVRAMReq: model.estimatedVRAMReq, + isBuiltin: true + ) + } + } else { + showAlert(message: "Path or URL should be provided in app config: \(model.modelID)") + } + } + } + + func loadModelConfig(modelConfigURL: URL, modelLib: String, modelID: String, estimatedVRAMReq: Int) -> ModelConfig? { + do { + assert(fileManager.fileExists(atPath: modelConfigURL.path())) + let fileHandle = try FileHandle(forReadingFrom: modelConfigURL) + let data = fileHandle.readDataToEndOfFile() + var modelConfig = try jsonDecoder.decode(ModelConfig.self, from: data) + modelConfig.modelLib = modelLib + modelConfig.modelID = modelID + modelConfig.estimatedVRAMReq = estimatedVRAMReq + return modelConfig + } catch { + showAlert(message: "Failed to resolve model config: \(error.localizedDescription)") + } + return nil + } + + func showAlert(message: String) { + DispatchQueue.main.async { [weak self] in + guard let self = self else { return } + if !self.alertDisplayed { + self.alertMessage = message + self.alertDisplayed = true + } else { + self.alertMessage.append("\n" + message) + } + } + } + + func downloadConfig(modelURL: URL?, modelLib: String, modelID: String, estimatedVRAMReq: Int, isBuiltin: Bool) { + guard let modelConfigURL = modelURL?.appending(path: "resolve").appending(path: "main").appending(path: Constants.modelConfigFileName) else { + return + } + + let downloadTask = URLSession.shared.downloadTask(with: modelConfigURL) { + [weak self] urlOrNil, responseOrNil, errorOrNil in + guard let self else { + return + } + if let error = errorOrNil { + self.showAlert(message: "Failed to download model config: \(error.localizedDescription)") + return + } + guard let fileUrl = urlOrNil else { + self.showAlert(message: "Failed to download model config") + return + } + + // cache temp file to avoid being deleted by system automatically + let tempName = UUID().uuidString + let tempFileURL = self.cacheDirectoryURL.appending(path: tempName) + + do { + try self.fileManager.moveItem(at: fileUrl, to: tempFileURL) + } catch { + self.showAlert(message: "Failed to cache downloaded file: \(error.localizedDescription)") + return + } + + do { + guard let modelConfig = loadModelConfig( + modelConfigURL: tempFileURL, + modelLib: modelLib, + modelID: modelID, + estimatedVRAMReq: estimatedVRAMReq + ) else { + try fileManager.removeItem(at: tempFileURL) + return + } + + if modelIDs.contains(modelConfig.modelID!) { + try fileManager.removeItem(at: tempFileURL) + return + } + + let modelBaseUrl = cacheDirectoryURL.appending(path: modelConfig.modelID!) + try fileManager.createDirectory(at: modelBaseUrl, withIntermediateDirectories: true) + let modelConfigUrl = modelBaseUrl.appending(path: Constants.modelConfigFileName) + try fileManager.moveItem(at: tempFileURL, to: modelConfigUrl) + assert(fileManager.fileExists(atPath: modelConfigUrl.path())) + assert(!fileManager.fileExists(atPath: tempFileURL.path())) + addModelConfig( + modelConfig: modelConfig, + modelPath: nil, + modelURL: modelURL, + isBuiltin: isBuiltin + ) + } catch { + showAlert(message: "Failed to import model: \(error.localizedDescription)") + } + } + downloadTask.resume() + } + + func addModelConfig(modelConfig: ModelConfig, modelPath: String?, modelURL: URL?, isBuiltin: Bool) { + assert(!modelIDs.contains(modelConfig.modelID!)) + modelIDs.insert(modelConfig.modelID!) + let modelBaseURL: URL + + // model_id dir should exist + if modelURL == nil { + // prebuilt model in bundle + modelBaseURL = Bundle.main.bundleURL.appending(path: Constants.prebuiltModelDir).appending(path: modelPath!) + } else { + // download model in cache + modelBaseURL = cacheDirectoryURL.appending(path: modelConfig.modelID!) + } + assert(fileManager.fileExists(atPath: modelBaseURL.path())) + + // mlc-chat-config.json should exist + let modelConfigURL = modelBaseURL.appending(path: Constants.modelConfigFileName) + assert(fileManager.fileExists(atPath: modelConfigURL.path())) + + let model = ModelState(modelConfig: modelConfig, modelLocalBaseURL: modelBaseURL, startState: self, chatState: chatState) + model.checkModelDownloadState(modelURL: modelURL) + + // addModelConfig is not called from main thread, update to models needs to be performed on main + DispatchQueue.main.async { [weak self] in + guard let self = self else { return } + models.append(model) + } + + if modelURL != nil && !isBuiltin { + updateAppConfig { + appConfig?.modelList.append( + AppConfig.ModelRecord( + modelPath: nil, + modelURL: modelURL!.absoluteString, + modelLib: modelConfig.modelLib!, + estimatedVRAMReq: modelConfig.estimatedVRAMReq!, + modelID: modelConfig.modelID! + ) + ) + } + } + } + + func updateAppConfig(action: () -> Void) { + action() + let appConfigURL = cacheDirectoryURL.appending(path: Constants.appConfigFileName) + do { + let data = try jsonEncoder.encode(appConfig) + try data.write(to: appConfigURL, options: Data.WritingOptions.atomic) + } catch { + print(error.localizedDescription) + } + } +} diff --git a/ios/MLCChat/MLCChat/States/ChatState.swift b/ios/MLCChat/MLCChat/States/ChatState.swift new file mode 100644 index 0000000..47f5784 --- /dev/null +++ b/ios/MLCChat/MLCChat/States/ChatState.swift @@ -0,0 +1,360 @@ +// +// ChatState.swift +// LLMChat +// + +import Foundation +import MLCSwift + +enum MessageRole { + case user + case assistant +} + +extension MessageRole { + var isUser: Bool { self == .user } +} + +struct MessageData: Hashable { + let id = UUID() + var role: MessageRole + var message: String +} + +final class ChatState: ObservableObject { + fileprivate enum ModelChatState { + case generating + case resetting + case reloading + case terminating + case ready + case failed + case pendingImageUpload + case processingImage + } + + @Published var displayMessages = [MessageData]() + @Published var infoText = "" + @Published var displayName = "" + // this is a legacy UI option for upload image + // TODO(mlc-team) support new UI for image processing + @Published var legacyUseImage = false + + private let modelChatStateLock = NSLock() + private var modelChatState: ModelChatState = .ready + + // the new mlc engine + private let engine = MLCEngine() + // history messages + private var historyMessages = [ChatCompletionMessage]() + + // streaming text that get updated + private var streamingText = "" + + private var modelLib = "" + private var modelPath = "" + var modelID = "" + + init() { + } + + var isInterruptible: Bool { + return getModelChatState() == .ready + || getModelChatState() == .generating + || getModelChatState() == .failed + || getModelChatState() == .pendingImageUpload + } + + var isChattable: Bool { + return getModelChatState() == .ready + } + + var isUploadable: Bool { + return getModelChatState() == .pendingImageUpload + } + + var isResettable: Bool { + return getModelChatState() == .ready + || getModelChatState() == .generating + } + + func requestResetChat() { + assert(isResettable) + interruptChat(prologue: { + switchToResetting() + }, epilogue: { [weak self] in + self?.mainResetChat() + }) + } + + // reset the chat if we switch to background + // during generation to avoid permission issue + func requestSwitchToBackground() { + if (getModelChatState() == .generating) { + self.requestResetChat() + } + } + + + func requestTerminateChat(callback: @escaping () -> Void) { + assert(isInterruptible) + interruptChat(prologue: { + switchToTerminating() + }, epilogue: { [weak self] in + self?.mainTerminateChat(callback: callback) + }) + } + + func requestReloadChat(modelID: String, modelLib: String, modelPath: String, estimatedVRAMReq: Int, displayName: String) { + if (isCurrentModel(modelID: modelID)) { + return + } + assert(isInterruptible) + interruptChat(prologue: { + switchToReloading() + }, epilogue: { [weak self] in + self?.mainReloadChat(modelID: modelID, + modelLib: modelLib, + modelPath: modelPath, + estimatedVRAMReq: estimatedVRAMReq, + displayName: displayName) + }) + } + + + func requestGenerate(prompt: String) { + assert(isChattable) + switchToGenerating() + appendMessage(role: .user, message: prompt) + appendMessage(role: .assistant, message: "") + + Task { + self.historyMessages.append( + ChatCompletionMessage(role: .user, content: prompt) + ) + var finishReasonLength = false + var finalUsageTextLabel = "" + + for await res in await engine.chat.completions.create( + messages: self.historyMessages, + stream_options: StreamOptions(include_usage: true) + ) { + for choice in res.choices { + if let content = choice.delta.content { + self.streamingText += content.asText() + } + if let finish_reason = choice.finish_reason { + if finish_reason == "length" { + finishReasonLength = true + } + } + } + if let finalUsage = res.usage { + finalUsageTextLabel = finalUsage.extra?.asTextLabel() ?? "" + } + if getModelChatState() != .generating { + break + } + + var updateText = self.streamingText + if finishReasonLength { + updateText += " [output truncated due to context length limit...]" + } + + let newText = updateText + DispatchQueue.main.async { + self.updateMessage(role: .assistant, message: newText) + } + } + + // record history messages + if !self.streamingText.isEmpty { + self.historyMessages.append( + ChatCompletionMessage(role: .assistant, content: self.streamingText) + ) + // stream text can be cleared + self.streamingText = "" + } else { + self.historyMessages.removeLast() + } + + // if we exceed history + // we can try to reduce the history and see if it can fit + if (finishReasonLength) { + let windowSize = self.historyMessages.count + assert(windowSize % 2 == 0) + let removeEnd = ((windowSize + 3) / 4) * 2 + self.historyMessages.removeSubrange(0.. Bool { + return self.modelID == modelID + } +} + +private extension ChatState { + func getModelChatState() -> ModelChatState { + modelChatStateLock.lock() + defer { modelChatStateLock.unlock() } + return modelChatState + } + + func setModelChatState(_ newModelChatState: ModelChatState) { + modelChatStateLock.lock() + modelChatState = newModelChatState + modelChatStateLock.unlock() + } + + func appendMessage(role: MessageRole, message: String) { + displayMessages.append(MessageData(role: role, message: message)) + } + + func updateMessage(role: MessageRole, message: String) { + displayMessages[displayMessages.count - 1] = MessageData(role: role, message: message) + } + + func clearHistory() { + displayMessages.removeAll() + infoText = "" + historyMessages.removeAll() + streamingText = "" + } + + func switchToResetting() { + setModelChatState(.resetting) + } + + func switchToGenerating() { + setModelChatState(.generating) + } + + func switchToReloading() { + setModelChatState(.reloading) + } + + func switchToReady() { + setModelChatState(.ready) + } + + func switchToTerminating() { + setModelChatState(.terminating) + } + + func switchToFailed() { + setModelChatState(.failed) + } + + func switchToPendingImageUpload() { + setModelChatState(.pendingImageUpload) + } + + func switchToProcessingImage() { + setModelChatState(.processingImage) + } + + func interruptChat(prologue: () -> Void, epilogue: @escaping () -> Void) { + assert(isInterruptible) + if getModelChatState() == .ready + || getModelChatState() == .failed + || getModelChatState() == .pendingImageUpload { + prologue() + epilogue() + } else if getModelChatState() == .generating { + prologue() + DispatchQueue.main.async { + epilogue() + } + } else { + assert(false) + } + } + + func mainResetChat() { + Task { + await engine.reset() + self.historyMessages = [] + self.streamingText = "" + + DispatchQueue.main.async { + self.clearHistory() + self.switchToReady() + } + } + } + + func mainTerminateChat(callback: @escaping () -> Void) { + Task { + await engine.unload() + DispatchQueue.main.async { + self.clearHistory() + self.modelID = "" + self.modelLib = "" + self.modelPath = "" + self.displayName = "" + self.legacyUseImage = false + self.switchToReady() + callback() + } + } + } + + func mainReloadChat(modelID: String, modelLib: String, modelPath: String, estimatedVRAMReq: Int, displayName: String) { + clearHistory() + self.modelID = modelID + self.modelLib = modelLib + self.modelPath = modelPath + self.displayName = displayName + + Task { + DispatchQueue.main.async { + self.appendMessage(role: .assistant, message: "[System] Initalize...") + } + + await engine.unload() + let vRAM = os_proc_available_memory() + if (vRAM < estimatedVRAMReq) { + let requiredMemory = String ( + format: "%.1fMB", Double(estimatedVRAMReq) / Double(1 << 20) + ) + let errorMessage = ( + "Sorry, the system cannot provide \(requiredMemory) VRAM as requested to the app, " + + "so we cannot initialize this model on this device." + ) + DispatchQueue.main.sync { + self.displayMessages.append(MessageData(role: MessageRole.assistant, message: errorMessage)) + self.switchToFailed() + } + return + } + await engine.reload( + modelPath: modelPath, modelLib: modelLib + ) + + // run a simple prompt with empty content to warm up system prompt + // helps to start things before user start typing + for await _ in await engine.chat.completions.create( + messages: [ChatCompletionMessage(role: .user, content: "")], + max_tokens: 1 + ) {} + + // TODO(mlc-team) run a system message prefill + DispatchQueue.main.async { + self.updateMessage(role: .assistant, message: "[System] Ready to chat") + self.switchToReady() + } + + } + } +} diff --git a/ios/MLCChat/MLCChat/States/ModelState.swift b/ios/MLCChat/MLCChat/States/ModelState.swift new file mode 100644 index 0000000..d893a63 --- /dev/null +++ b/ios/MLCChat/MLCChat/States/ModelState.swift @@ -0,0 +1,414 @@ +// +// ModelState.swift +// MLCChat +// + +import Foundation + +final class ModelState: ObservableObject, Identifiable { + enum ModelDownloadState { + case initializing + case indexing + case paused + case downloading + case pausing + case verifying + case finished + case failed + case clearing + case deleting + } + + fileprivate struct DownloadTask: Hashable { + let remoteURL: URL + let localURL: URL + } + + @Published var modelConfig: ModelConfig + @Published var modelDownloadState: ModelDownloadState = .initializing + @Published var progress: Int = 0 + @Published var total: Int = 1 + + private var modelLocalBaseURL: URL + private var startState: AppState + private var chatState: ChatState + + private let fileManager: FileManager = FileManager.default + private let decoder = JSONDecoder() + private var paramsConfig: ParamsConfig? + private var modelRemoteBaseURL: URL? + private var remainingTasks: Set = Set() + private var downloadingTasks: Set = Set() + private var maxDownloadingTasks: Int = 3 + + init(modelConfig: ModelConfig, + modelLocalBaseURL: URL, + startState: AppState, + chatState: ChatState) { + self.modelConfig = modelConfig + self.modelLocalBaseURL = modelLocalBaseURL + self.startState = startState + self.chatState = chatState + } + + func checkModelDownloadState(modelURL: URL?) { + createModelFolderIfNeeded() + + guard let modelURL else { + switchToVerifying() + return + } + + modelRemoteBaseURL = modelURL.appending(path: "resolve").appending(path: "main") + + // create local params dir + let paramsConfigURL = modelLocalBaseURL.appending(path: Constants.paramsConfigFileName) + if fileManager.fileExists(atPath: paramsConfigURL.path()) { + // tensor-cache.json already downloaded + loadParamsConfig() + switchToIndexing() + } else { + // download tensor-cache.json + downloadParamsConfig() + } + } + + func startChat(chatState: ChatState) { + chatState.requestReloadChat( + modelID: modelConfig.modelID!, + modelLib: modelConfig.modelLib!, + modelPath: modelLocalBaseURL.path(), + estimatedVRAMReq: modelConfig.estimatedVRAMReq!, + displayName: modelConfig.modelID!.components(separatedBy: "-")[0] + ) + } + + func handleStart() { + // start downloading + switchToDownloading() + } + + func handlePause() { + // pause downloading + switchToPausing() + } + + func handleClear() { + assert(modelDownloadState == .downloading || modelDownloadState == .paused || modelDownloadState == .finished) + switchToClearing() + } + + func handleDelete() { + assert(modelDownloadState == .downloading || modelDownloadState == .paused || modelDownloadState == .finished || modelDownloadState == .failed) + switchToDeleting() + } +} + +private extension ModelState { + func createModelFolderIfNeeded() { + if !fileManager.fileExists(atPath: modelLocalBaseURL.path()) { + do { + try fileManager.createDirectory(at: modelLocalBaseURL, withIntermediateDirectories: true) + } catch { + print(error.localizedDescription) + } + } + } + + func loadParamsConfig() { + let paramsConfigURL = modelLocalBaseURL.appending(path: Constants.paramsConfigFileName) + assert(fileManager.fileExists(atPath: paramsConfigURL.path())) + do { + let fileHandle = try FileHandle(forReadingFrom: paramsConfigURL) + let data = fileHandle.readDataToEndOfFile() + paramsConfig = try self.decoder.decode(ParamsConfig.self, from: data) + } catch { + print(error.localizedDescription) + } + } + + func downloadParamsConfig() { + guard let modelRemoteBaseURL else { + return + } + + let paramsConfigURL = modelLocalBaseURL.appending(path: Constants.paramsConfigFileName) + let downloadTask = URLSession.shared.downloadTask(with: modelRemoteBaseURL.appending(path: Constants.paramsConfigFileName)) { + [weak self] urlOrNil, responseOrNil, errorOrNil in + guard let self else { return } + guard let fileURL = urlOrNil else { return } + do { + try? self.fileManager.removeItem(at: paramsConfigURL) + try self.fileManager.moveItem(at: fileURL, to: paramsConfigURL) + DispatchQueue.main.async { + self.loadParamsConfig() + self.switchToIndexing() + } + } catch { + print(error.localizedDescription) + } + } + downloadTask.resume() + } + + func switchToIndexing() { + guard let paramsConfig, let modelRemoteBaseURL else { + return + } + + modelDownloadState = .indexing + progress = 0 + total = modelConfig.tokenizerFiles.count + paramsConfig.records.count + + // collect tokenizer download tasks + for tokenizerFile in modelConfig.tokenizerFiles { + let remoteURL = modelRemoteBaseURL.appending(path: tokenizerFile) + let localURL = modelLocalBaseURL.appending(path: tokenizerFile) + + if fileManager.fileExists(atPath: localURL.path()) { + progress += 1 + } else { + remainingTasks.insert(DownloadTask(remoteURL: remoteURL, localURL: localURL)) + } + } + + // collect params download tasks + for paramsRecord in paramsConfig.records { + let remoteURL = modelRemoteBaseURL.appending(path: paramsRecord.dataPath) + let localURL = modelLocalBaseURL.appending(path: paramsRecord.dataPath) + + if fileManager.fileExists(atPath: localURL.path()) { + progress += 1 + } else { + remainingTasks.insert(DownloadTask(remoteURL: remoteURL, localURL: localURL)) + } + } + + if progress < total { + switchToPaused() + } else { + switchToFinished() + } + } + + func handleNewDownload(downloadTask: DownloadTask) { + // start one download task + assert(downloadingTasks.count < maxDownloadingTasks) + let task = URLSession.shared.downloadTask(with: downloadTask.remoteURL) { + [weak self] urlOrNil, responseOrNil, errorOrNil in + guard let self else { return } + guard let fileUrl = urlOrNil else { + DispatchQueue.main.async { + self.handleCancelDownload(downloadTask: downloadTask) + } + return + } + + do { + try self.fileManager.createDirectory(at: downloadTask.localURL.deletingLastPathComponent(), withIntermediateDirectories: true) + try? self.fileManager.removeItem(at: downloadTask.localURL) + try self.fileManager.moveItem(at: fileUrl, to: downloadTask.localURL) + } catch { + print(error.localizedDescription) + } + DispatchQueue.main.async { + self.handleFinishDownload(downloadTask: downloadTask) + } + } + downloadingTasks.insert(downloadTask) + task.resume() + } + + func handleFinishDownload(downloadTask: DownloadTask) { + // update the finished download task + remainingTasks.remove(downloadTask) + downloadingTasks.remove(downloadTask) + progress += 1 + assert(modelDownloadState == .downloading || + modelDownloadState == .pausing || + modelDownloadState == .clearing || + modelDownloadState == .deleting + ) + if modelDownloadState == .downloading { + if remainingTasks.isEmpty && downloadingTasks.isEmpty { + switchToFinished() + } else { + handleNextDownload() + } + } else if modelDownloadState == .pausing && downloadingTasks.isEmpty { + switchToPaused() + } else if modelDownloadState == .clearing && downloadingTasks.isEmpty { + clear() + } else if modelDownloadState == .deleting && downloadingTasks.isEmpty { + delete() + } + } + + func handleCancelDownload(downloadTask: DownloadTask) { + // withdraw the failed download task + assert(modelDownloadState == .downloading || modelDownloadState == .pausing) + downloadingTasks.remove(downloadTask) + if modelDownloadState == .downloading { + handleNextDownload() + } else if modelDownloadState == .pausing && downloadingTasks.count == 0 { + switchToPaused() + } + } + + func handleNextDownload() { + // start next download task + assert(modelDownloadState == .downloading) + for downloadTask in remainingTasks { + if !downloadingTasks.contains(downloadTask) { + handleNewDownload(downloadTask: downloadTask) + break + } + } + } + + func switchToPaused() { + modelDownloadState = .paused + } + + func switchToPausing() { + modelDownloadState = .pausing + } + + func switchToVerifying() { + modelDownloadState = .verifying + + let paramsConfigURL = modelLocalBaseURL.appending(path: Constants.paramsConfigFileName) + guard fileManager.fileExists(atPath: paramsConfigURL.path()) else { + switchToFailed() + return + } + + loadParamsConfig() + guard let paramsConfig else { + switchToFailed() + return + } + progress = 0 + total = modelConfig.tokenizerFiles.count + paramsConfig.records.count + + if !verifyTokenizers() { + switchToFailed() + return + } + + if !verifyParams() { + switchToFailed() + return + } + + switchToFinished() + } + + func verifyTokenizers() -> Bool { + for tokenizerFile in modelConfig.tokenizerFiles { + let localURL = modelLocalBaseURL.appending(path: tokenizerFile) + + if !fileManager.fileExists(atPath: localURL.path()) { + switchToFailed() + return false + } + progress += 1 + } + return true + } + + func verifyParams() -> Bool { + guard let paramsConfig else { + return false + } + + for paramsRecord in paramsConfig.records { + let localUrl = modelLocalBaseURL.appending(path: paramsRecord.dataPath) + + if !fileManager.fileExists(atPath: localUrl.path()) { + switchToFailed() + return false + } + + progress += 1 + } + return true + } + + func switchToClearing() { + if modelDownloadState == .paused { + modelDownloadState = .clearing + clear() + } else if modelDownloadState == .finished { + if chatState.modelID == modelConfig.modelID { + chatState.requestTerminateChat { [weak self] in + self?.clear() + } + } else { + clear() + } + } else { + modelDownloadState = .clearing + } + } + + func switchToDeleting() { + if modelDownloadState == .paused || modelDownloadState == .failed { + modelDownloadState = .deleting + delete() + } else if modelDownloadState == .finished { + if chatState.modelID == modelConfig.modelID { + chatState.requestTerminateChat { [weak self] in + self?.delete() + } + } else { + delete() + } + } else { + modelDownloadState = .deleting + } + } + + func switchToFinished() { + modelDownloadState = .finished + } + + func switchToFailed() { + modelDownloadState = .failed + } + + func switchToDownloading() { + modelDownloadState = .downloading + for downloadTask in remainingTasks { + if downloadingTasks.count < maxDownloadingTasks { + handleNewDownload(downloadTask: downloadTask) + } else { + return + } + } + } + + func clear() { + do { + let fileURLs = try fileManager.contentsOfDirectory(at: modelLocalBaseURL, includingPropertiesForKeys: nil) + for fileURL in fileURLs where fileURL.lastPathComponent != Constants.modelConfigFileName { + try fileManager.removeItem(at: fileURL) + assert(!fileManager.fileExists(atPath: fileURL.path())) + } + assert(fileManager.fileExists(atPath: modelLocalBaseURL.appending(path: Constants.modelConfigFileName).path())) + switchToIndexing() + } catch { + print(error.localizedDescription) + } + } + + func delete() { + do { + try fileManager.removeItem(at: modelLocalBaseURL) + assert(!fileManager.fileExists(atPath: modelLocalBaseURL.path())) + startState.requestDeleteModel(modelID: modelConfig.modelID!) // TODO: can it decouple? + } catch { + print(error.localizedDescription) + } + } +} diff --git a/ios/MLCChat/MLCChat/Views/ChatView.swift b/ios/MLCChat/MLCChat/Views/ChatView.swift new file mode 100644 index 0000000..9a568cf --- /dev/null +++ b/ios/MLCChat/MLCChat/Views/ChatView.swift @@ -0,0 +1,181 @@ +// +// ChatView.swift +// MLCChat +// + +import SwiftUI +import GameController + +struct ChatView: View { + @EnvironmentObject private var chatState: ChatState + @Environment(\.scenePhase) var scenePhase + @State private var inputMessage: String = "" + @FocusState private var inputIsFocused: Bool + @Environment(\.dismiss) private var dismiss + @Namespace private var messagesBottomID + + // vision-related properties + @State private var showActionSheet: Bool = false + @State private var showImagePicker: Bool = false + @State private var imageConfirmed: Bool = false + @State private var imageSourceType: UIImagePickerController.SourceType = .photoLibrary + @State private var image: UIImage? + + var body: some View { + VStack { + modelInfoView + messagesView + uploadImageView + messageInputView + } + .navigationBarTitle("MLC Chat: \(chatState.displayName)", displayMode: .inline) + .navigationBarBackButtonHidden() + .onChange(of: scenePhase) { oldPhase, newPhase in + if newPhase == .background { + self.chatState.requestSwitchToBackground() + } + } + .toolbar { + ToolbarItem(placement: .navigationBarLeading) { + Button { + dismiss() + } label: { + Image(systemName: "chevron.backward") + } + .buttonStyle(.borderless) + .disabled(!chatState.isInterruptible) + } + ToolbarItem(placement: .navigationBarTrailing) { + Button("Reset") { + image = nil + imageConfirmed = false + chatState.requestResetChat() + } + .padding() + .disabled(!chatState.isResettable) + } + } + + } +} + +private extension ChatView { + var modelInfoView: some View { + Text(chatState.infoText) + .multilineTextAlignment(.center) + .opacity(0.5) + .listRowSeparator(.hidden) + } + + var messagesView: some View { + ScrollViewReader { scrollViewProxy in + ScrollView { + VStack { + let messageCount = chatState.displayMessages.count + let hasSystemMessage = messageCount > 0 && chatState.displayMessages[0].role == MessageRole.assistant + let startIndex = hasSystemMessage ? 1 : 0 + + // display the system message + if hasSystemMessage { + MessageView(role: chatState.displayMessages[0].role, message: chatState.displayMessages[0].message, isMarkdownSupported: false) + } + + // display image + if let image, imageConfirmed { + ImageView(image: image) + } + + // display conversations + ForEach(chatState.displayMessages[startIndex...], id: \.id) { message in + MessageView(role: message.role, message: message.message) + } + HStack { EmptyView() } + .id(messagesBottomID) + } + } + .onChange(of: chatState.displayMessages) { _ in + withAnimation { + scrollViewProxy.scrollTo(messagesBottomID, anchor: .bottom) + } + } + } + } + + @ViewBuilder + var uploadImageView: some View { + if chatState.legacyUseImage && !imageConfirmed { + if image == nil { + Button("Upload picture to chat") { + showActionSheet = true + } + .actionSheet(isPresented: $showActionSheet) { + ActionSheet(title: Text("Choose from"), buttons: [ + .default(Text("Photo Library")) { + showImagePicker = true + imageSourceType = .photoLibrary + }, + .default(Text("Camera")) { + showImagePicker = true + imageSourceType = .camera + }, + .cancel() + ]) + } + .sheet(isPresented: $showImagePicker) { + ImagePicker(image: $image, + showImagePicker: $showImagePicker, + imageSourceType: imageSourceType) + } + .disabled(!chatState.isUploadable) + } else { + VStack { + if let image { + Image(uiImage: image) + .resizable() + .frame(width: 300, height: 300) + + HStack { + Button("Undo") { + self.image = nil + } + .padding() + + Button("Submit") { + imageConfirmed = true + } + .padding() + } + } + } + } + } + } + + var messageInputView: some View { + HStack { + TextField("Inputs...", text: $inputMessage, axis: .vertical) + .textFieldStyle(RoundedBorderTextFieldStyle()) + .frame(minHeight: CGFloat(30)) + .focused($inputIsFocused) + .onSubmit { + let isKeyboardConnected = GCKeyboard.coalesced != nil + if isKeyboardConnected { + send() + } + } + Button("Send") { + send() + } + .bold() + .disabled(!(chatState.isChattable && inputMessage != "")) + } + .frame(minHeight: CGFloat(70)) + .padding() + } + + func send() { + inputIsFocused = false + chatState.requestGenerate(prompt: inputMessage) + inputMessage = "" + } +} diff --git a/ios/MLCChat/MLCChat/Views/ImageProcessing.swift b/ios/MLCChat/MLCChat/Views/ImageProcessing.swift new file mode 100644 index 0000000..3d7260e --- /dev/null +++ b/ios/MLCChat/MLCChat/Views/ImageProcessing.swift @@ -0,0 +1,66 @@ +// +// ImageProcessing.swift +// MLCChat +// +// Created by Kathryn Chen on 7/8/23. +// + +import Foundation +import SwiftUI +import UIKit + +// adapted from Mohammad Azam: https://github.com/azamsharp/SwiftUICamera +// delegate task to the coordinator to produce the image +struct ImagePicker : UIViewControllerRepresentable { + typealias UIViewControllerType = UIImagePickerController + typealias Coordinator = ImagePickerCoordinator + + @Binding var image: UIImage? + @Binding var showImagePicker: Bool + var imageSourceType: UIImagePickerController.SourceType = .photoLibrary + + func makeCoordinator() -> ImagePicker.Coordinator { + return ImagePickerCoordinator(image: $image, showImagePicker: $showImagePicker) + } + + func makeUIViewController(context: UIViewControllerRepresentableContext) -> UIImagePickerController { + let picker = UIImagePickerController() + picker.sourceType = imageSourceType + picker.delegate = context.coordinator + return picker + } + + func updateUIViewController(_ uiViewController: UIImagePickerController, context: UIViewControllerRepresentableContext) {} +} + +// image picker coordinator handling selecting from library or taking a photo +class ImagePickerCoordinator: NSObject, UINavigationControllerDelegate, UIImagePickerControllerDelegate { + @Binding var image: UIImage? + @Binding var showImagePicker: Bool + + init(image: Binding, showImagePicker: Binding) { + _image = image + _showImagePicker = showImagePicker + } + + func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) { + if let optionalImage = info[UIImagePickerController.InfoKey.originalImage] as? UIImage { + image = optionalImage + showImagePicker = false + } + } + + func imagePickerControllerDidCancel(_ picker: UIImagePickerController) { + showImagePicker = false + } +} + +// resize the input image to given width and height +func resizeImage(image: UIImage, width: Int, height: Int) -> UIImage { + let shape = CGSize(width: width, height: height) + UIGraphicsBeginImageContextWithOptions(shape, true, 0.0) + image.draw(in: CGRect(x: 0, y: 0, width: width, height: height)) + let resizedImage: UIImage? = UIGraphicsGetImageFromCurrentImageContext() + UIGraphicsEndImageContext() + return resizedImage ?? image +} diff --git a/ios/MLCChat/MLCChat/Views/MessageView.swift b/ios/MLCChat/MLCChat/Views/MessageView.swift new file mode 100644 index 0000000..692f4fd --- /dev/null +++ b/ios/MLCChat/MLCChat/Views/MessageView.swift @@ -0,0 +1,107 @@ +// +// MessageView.swift +// MLCChat +// + +import SwiftUI +import MarkdownUI + +struct MessageView: View { + let role: MessageRole; + let message: String + let isMarkdownSupported: Bool + + @State private var showMarkdown: Bool + + init(role: MessageRole, message: String, isMarkdownSupported: Bool = true) { + self.role = role + self.message = message + self.isMarkdownSupported = isMarkdownSupported + _showMarkdown = State(initialValue: isMarkdownSupported) + } + var body: some View { + let textColor = role.isUser ? Color.white : Color(UIColor.label) + let background = role.isUser ? Color.blue : Color(UIColor.secondarySystemBackground) + + HStack { + if role.isUser { + Spacer() + Text(message) + .padding(10) + .foregroundColor(textColor) + .background(background) + .cornerRadius(10) + .textSelection(.enabled) + } + if !role.isUser { + VStack(alignment: .leading) { + // Toggle switch to show/hide Markdown + if(isMarkdownSupported){ + Toggle(isOn: $showMarkdown) { + Text("Show as Markdown") + .font(.footnote) + .foregroundColor(.blue) + } + .padding(.bottom, 10) + } + + // Conditionally display Text or Markdown + if showMarkdown { + Markdown { + message + } + .markdownTheme(.gitHub) + .padding(10) + .foregroundColor(textColor) + .background(background) + .cornerRadius(10) + .textSelection(.enabled) + } else { + Text(message) + .padding(10) + .foregroundColor(textColor) + .background(background) + .cornerRadius(10) + .textSelection(.enabled) + } + } + Spacer() + } + } + .padding() + .listRowSeparator(.hidden) + } +} + +struct ImageView: View { + let image: UIImage + + var body: some View { + let background = Color.blue + HStack { + Spacer() + Image(uiImage: image) + .resizable() + .frame(width: 150, height: 150) + .padding(15) + .background(background) + .cornerRadius(20) + } + .padding() + .listRowSeparator(.hidden) + } +} + +struct MessageView_Previews: PreviewProvider { + static var previews: some View { + NavigationView { + VStack (spacing: 0){ + ScrollView { + MessageView(role: MessageRole.user, message: "Message 1") + MessageView(role: MessageRole.assistant, message: "Message 2") + MessageView(role: MessageRole.user, message: "Message 3") + } + } + } + } +} diff --git a/ios/MLCChat/MLCChat/Views/ModelView.swift b/ios/MLCChat/MLCChat/Views/ModelView.swift new file mode 100644 index 0000000..02acd3f --- /dev/null +++ b/ios/MLCChat/MLCChat/Views/ModelView.swift @@ -0,0 +1,97 @@ +// +// ModelView.swift +// MLCChat +// +// Created by Yaxing Cai on 5/14/23. +// + +import SwiftUI + +struct ModelView: View { + @EnvironmentObject private var modelState: ModelState + @EnvironmentObject private var chatState: ChatState + @Binding var isRemoving: Bool + + @State private var isShowingDeletionConfirmation: Bool = false + + var body: some View { + VStack(alignment: .leading) { + if (modelState.modelDownloadState == .finished) { + NavigationLink(destination: + ChatView() + .environmentObject(chatState) + .onAppear { + modelState.startChat(chatState: chatState) + } + ) { + HStack { + Text(modelState.modelConfig.modelID!) + Spacer() + if chatState.isCurrentModel(modelID: modelState.modelConfig.modelID!) { + Image(systemName: "checkmark").foregroundColor(.blue) + } + } + } + .buttonStyle(.borderless) + } else { + Text(modelState.modelConfig.modelID!).opacity(0.5) + } + HStack{ + if modelState.modelDownloadState != .finished || isRemoving { + ProgressView(value: Double(modelState.progress) / Double(modelState.total)) + .progressViewStyle(.linear) + } + + if (modelState.modelDownloadState == .paused) { + Button { + modelState.handleStart() + } label: { + Image(systemName: "icloud.and.arrow.down") + } + .buttonStyle(.borderless) + } else if (modelState.modelDownloadState == .downloading) { + Button { + modelState.handlePause() + } label: { + Image(systemName: "stop.circle") + } + .buttonStyle(.borderless) + } else if (modelState.modelDownloadState == .failed) { + Image(systemName: "exclamationmark.triangle") + .foregroundColor(.red) + } + + if isRemoving { + Button(role: .destructive) { + isShowingDeletionConfirmation = true + } label: { + Image(systemName: "trash") + } + .confirmationDialog("Delete Model", isPresented: $isShowingDeletionConfirmation) { + Button("Delete Model", role: .destructive) { + modelState.handleDelete() + } + .disabled( + modelState.modelDownloadState != .downloading && + modelState.modelDownloadState != .paused && + modelState.modelDownloadState != .finished && + modelState.modelDownloadState != .failed) + Button("Clear Data") { + modelState.handleClear() + } + .disabled( + modelState.modelDownloadState != .downloading && + modelState.modelDownloadState != .paused && + modelState.modelDownloadState != .finished) + Button("Cancel", role: .cancel) { + isShowingDeletionConfirmation = false + } + } message: { + Text("Delete model will delete the all files with model config, and delete the entry in list. \n Clear model will keep the model config only, and keep the entry in list for future re-downloading.") + } + .buttonStyle(.borderless) + } + } + } + } +} diff --git a/ios/MLCChat/MLCChat/Views/StartView.swift b/ios/MLCChat/MLCChat/Views/StartView.swift new file mode 100644 index 0000000..87e585d --- /dev/null +++ b/ios/MLCChat/MLCChat/Views/StartView.swift @@ -0,0 +1,46 @@ +// +// DownloadView.swift +// MLCChat +// +// Created by Yaxing Cai on 5/11/23. +// + +import SwiftUI + +struct StartView: View { + @EnvironmentObject private var appState: AppState + @State private var isAdding: Bool = false + @State private var isRemoving: Bool = false + @State private var inputModelUrl: String = "" + + var body: some View { + NavigationStack { + List{ + Section(header: Text("Models")) { + ForEach(appState.models) { modelState in + ModelView(isRemoving: $isRemoving) + .environmentObject(modelState) + .environmentObject(appState.chatState) + } + if !isRemoving { + Button("Edit model") { + isRemoving = true + } + .buttonStyle(.borderless) + } else { + Button("Cancel edit model") { + isRemoving = false + } + .buttonStyle(.borderless) + } + } + } + .navigationTitle("MLC Chat") + .alert("Error", isPresented: $appState.alertDisplayed) { + Button("OK") { } + } message: { + Text(appState.alertMessage) + } + } + } +} diff --git a/ios/MLCChat/README.md b/ios/MLCChat/README.md new file mode 100644 index 0000000..f4f4820 --- /dev/null +++ b/ios/MLCChat/README.md @@ -0,0 +1,6 @@ +# MLC Chat App + +Checkout [Documentation page](https://llm.mlc.ai/docs/deploy/ios.html) for more information. + +- run `mlc_llm package` +- open the Xcode project diff --git a/ios/MLCChat/mlc-package-config.json b/ios/MLCChat/mlc-package-config.json new file mode 100644 index 0000000..8e20043 --- /dev/null +++ b/ios/MLCChat/mlc-package-config.json @@ -0,0 +1,49 @@ +{ + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/Llama-3.2-3B-Instruct-q4f16_1-MLC", + "model_id": "Llama-3.2-3B-Instruct-q4f16_1-MLC", + "estimated_vram_bytes": 3000000000, + "overrides": { + "prefill_chunk_size": 128, + "context_window_size": 2048 + }, + "bundle_weight": true + }, + { + "model": "HF://mlc-ai/gemma-2-2b-it-q4f16_1-MLC", + "model_id": "gemma-2-2b-q4f16_1-MLC", + "estimated_vram_bytes": 3000000000, + "overrides": { + "prefill_chunk_size": 128 + } + }, + { + "model": "HF://mlc-ai/Phi-3.5-mini-instruct-q4f16_1-MLC", + "model_id": "Phi-3.5-mini-instruct-q4f16_1-MLC", + "estimated_vram_bytes": 3043000000, + "overrides": { + "prefill_chunk_size": 128 + } + }, + { + "model": "HF://mlc-ai/Qwen3-0.6B-q0f16-MLC", + "model_id": "Qwen3-0.6B-q0f16-MLC", + "estimated_vram_bytes": 3000000000, + "overrides": { + "prefill_chunk_size": 128, + "context_window_size": 2048 + } + }, + { + "model": "HF://mlc-ai/Qwen3-1.7B-q4f16_1-MLC", + "model_id": "Qwen3-1.7B-q4f16_1-MLC", + "estimated_vram_bytes": 3000000000, + "overrides": { + "prefill_chunk_size": 128, + "context_window_size": 2048 + } + } + ] +} diff --git a/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.pbxproj b/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.pbxproj new file mode 100644 index 0000000..7ccdde6 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.pbxproj @@ -0,0 +1,422 @@ +// !$*UTF8*$! +{ + archiveVersion = 1; + classes = { + }; + objectVersion = 60; + objects = { + +/* Begin PBXBuildFile section */ + C04105DF2BEBC61B005A434D /* MLCSwift in Frameworks */ = {isa = PBXBuildFile; productRef = C04105DE2BEBC61B005A434D /* MLCSwift */; }; + C07094522BEBC6C4005C29FC /* bundle in Copy Files */ = {isa = PBXBuildFile; fileRef = C07094512BEBC6C4005C29FC /* bundle */; }; + C0B37B892BE8226A00B2F80B /* MLCEngineExampleApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = C0B37B882BE8226A00B2F80B /* MLCEngineExampleApp.swift */; }; + C0B37B8B2BE8226A00B2F80B /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C0B37B8A2BE8226A00B2F80B /* ContentView.swift */; }; + C0B37B8D2BE8226B00B2F80B /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C0B37B8C2BE8226B00B2F80B /* Assets.xcassets */; }; + C0B37B902BE8226B00B2F80B /* Preview Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C0B37B8F2BE8226B00B2F80B /* Preview Assets.xcassets */; }; + C0B37B982BE8234D00B2F80B /* MLCSwift in Frameworks */ = {isa = PBXBuildFile; productRef = C0B37B972BE8234D00B2F80B /* MLCSwift */; }; +/* End PBXBuildFile section */ + +/* Begin PBXCopyFilesBuildPhase section */ + C0B37B992BE8255600B2F80B /* Copy Files */ = { + isa = PBXCopyFilesBuildPhase; + buildActionMask = 12; + dstPath = ""; + dstSubfolderSpec = 7; + files = ( + C07094522BEBC6C4005C29FC /* bundle in Copy Files */, + ); + name = "Copy Files"; + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXCopyFilesBuildPhase section */ + +/* Begin PBXFileReference section */ + C07094512BEBC6C4005C29FC /* bundle */ = {isa = PBXFileReference; lastKnownFileType = folder; name = bundle; path = dist/bundle; sourceTree = ""; }; + C0B37B852BE8226A00B2F80B /* MLCEngineExample.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = MLCEngineExample.app; sourceTree = BUILT_PRODUCTS_DIR; }; + C0B37B882BE8226A00B2F80B /* MLCEngineExampleApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MLCEngineExampleApp.swift; sourceTree = ""; }; + C0B37B8A2BE8226A00B2F80B /* ContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = ""; }; + C0B37B8C2BE8226B00B2F80B /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = ""; }; + C0B37B8F2BE8226B00B2F80B /* Preview Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = "Preview Assets.xcassets"; sourceTree = ""; }; + C0B37C0C2BE8349300B2F80B /* MLCEngineExample.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = MLCEngineExample.entitlements; sourceTree = ""; }; +/* End PBXFileReference section */ + +/* Begin PBXFrameworksBuildPhase section */ + C0B37B822BE8226A00B2F80B /* Frameworks */ = { + isa = PBXFrameworksBuildPhase; + buildActionMask = 2147483647; + files = ( + C0B37B982BE8234D00B2F80B /* MLCSwift in Frameworks */, + C04105DF2BEBC61B005A434D /* MLCSwift in Frameworks */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXFrameworksBuildPhase section */ + +/* Begin PBXGroup section */ + C0B37B7C2BE8226A00B2F80B = { + isa = PBXGroup; + children = ( + C07094512BEBC6C4005C29FC /* bundle */, + C0B37B872BE8226A00B2F80B /* MLCEngineExample */, + C0B37B862BE8226A00B2F80B /* Products */, + ); + sourceTree = ""; + }; + C0B37B862BE8226A00B2F80B /* Products */ = { + isa = PBXGroup; + children = ( + C0B37B852BE8226A00B2F80B /* MLCEngineExample.app */, + ); + name = Products; + sourceTree = ""; + }; + C0B37B872BE8226A00B2F80B /* MLCEngineExample */ = { + isa = PBXGroup; + children = ( + C0B37C0C2BE8349300B2F80B /* MLCEngineExample.entitlements */, + C0B37B882BE8226A00B2F80B /* MLCEngineExampleApp.swift */, + C0B37B8A2BE8226A00B2F80B /* ContentView.swift */, + C0B37B8C2BE8226B00B2F80B /* Assets.xcassets */, + C0B37B8E2BE8226B00B2F80B /* Preview Content */, + ); + path = MLCEngineExample; + sourceTree = ""; + }; + C0B37B8E2BE8226B00B2F80B /* Preview Content */ = { + isa = PBXGroup; + children = ( + C0B37B8F2BE8226B00B2F80B /* Preview Assets.xcassets */, + ); + path = "Preview Content"; + sourceTree = ""; + }; +/* End PBXGroup section */ + +/* Begin PBXNativeTarget section */ + C0B37B842BE8226A00B2F80B /* MLCEngineExample */ = { + isa = PBXNativeTarget; + buildConfigurationList = C0B37B932BE8226B00B2F80B /* Build configuration list for PBXNativeTarget "MLCEngineExample" */; + buildPhases = ( + C0B37B812BE8226A00B2F80B /* Sources */, + C0B37B822BE8226A00B2F80B /* Frameworks */, + C0B37B832BE8226A00B2F80B /* Resources */, + C0B37B992BE8255600B2F80B /* Copy Files */, + ); + buildRules = ( + ); + dependencies = ( + ); + name = MLCEngineExample; + packageProductDependencies = ( + C0B37B972BE8234D00B2F80B /* MLCSwift */, + C04105DE2BEBC61B005A434D /* MLCSwift */, + ); + productName = MLCEngineExample; + productReference = C0B37B852BE8226A00B2F80B /* MLCEngineExample.app */; + productType = "com.apple.product-type.application"; + }; +/* End PBXNativeTarget section */ + +/* Begin PBXProject section */ + C0B37B7D2BE8226A00B2F80B /* Project object */ = { + isa = PBXProject; + attributes = { + BuildIndependentTargetsInParallel = 1; + LastSwiftUpdateCheck = 1530; + LastUpgradeCheck = 1530; + TargetAttributes = { + C0B37B842BE8226A00B2F80B = { + CreatedOnToolsVersion = 15.3; + }; + }; + }; + buildConfigurationList = C0B37B802BE8226A00B2F80B /* Build configuration list for PBXProject "MLCEngineExample" */; + compatibilityVersion = "Xcode 14.0"; + developmentRegion = en; + hasScannedForEncodings = 0; + knownRegions = ( + en, + Base, + ); + mainGroup = C0B37B7C2BE8226A00B2F80B; + packageReferences = ( + C0B37B962BE8234D00B2F80B /* XCLocalSwiftPackageReference "../MLCSwift" */, + ); + productRefGroup = C0B37B862BE8226A00B2F80B /* Products */; + projectDirPath = ""; + projectRoot = ""; + targets = ( + C0B37B842BE8226A00B2F80B /* MLCEngineExample */, + ); + }; +/* End PBXProject section */ + +/* Begin PBXResourcesBuildPhase section */ + C0B37B832BE8226A00B2F80B /* Resources */ = { + isa = PBXResourcesBuildPhase; + buildActionMask = 2147483647; + files = ( + C0B37B902BE8226B00B2F80B /* Preview Assets.xcassets in Resources */, + C0B37B8D2BE8226B00B2F80B /* Assets.xcassets in Resources */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXResourcesBuildPhase section */ + +/* Begin PBXSourcesBuildPhase section */ + C0B37B812BE8226A00B2F80B /* Sources */ = { + isa = PBXSourcesBuildPhase; + buildActionMask = 2147483647; + files = ( + C0B37B8B2BE8226A00B2F80B /* ContentView.swift in Sources */, + C0B37B892BE8226A00B2F80B /* MLCEngineExampleApp.swift in Sources */, + ); + runOnlyForDeploymentPostprocessing = 0; + }; +/* End PBXSourcesBuildPhase section */ + +/* Begin XCBuildConfiguration section */ + C0B37B912BE8226B00B2F80B /* Debug */ = { + isa = XCBuildConfiguration; + buildSettings = { + ALWAYS_SEARCH_USER_PATHS = NO; + ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES; + CLANG_ANALYZER_NONNULL = YES; + CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; + CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; + CLANG_ENABLE_MODULES = YES; + CLANG_ENABLE_OBJC_ARC = YES; + CLANG_ENABLE_OBJC_WEAK = YES; + CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; + CLANG_WARN_BOOL_CONVERSION = YES; + CLANG_WARN_COMMA = YES; + CLANG_WARN_CONSTANT_CONVERSION = YES; + CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; + CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; + CLANG_WARN_DOCUMENTATION_COMMENTS = YES; + CLANG_WARN_EMPTY_BODY = YES; + CLANG_WARN_ENUM_CONVERSION = YES; + CLANG_WARN_INFINITE_RECURSION = YES; + CLANG_WARN_INT_CONVERSION = YES; + CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; + CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; + CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; + CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; + CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; + CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; + CLANG_WARN_STRICT_PROTOTYPES = YES; + CLANG_WARN_SUSPICIOUS_MOVE = YES; + CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; + CLANG_WARN_UNREACHABLE_CODE = YES; + CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; + COPY_PHASE_STRIP = NO; + DEBUG_INFORMATION_FORMAT = dwarf; + ENABLE_STRICT_OBJC_MSGSEND = YES; + ENABLE_TESTABILITY = YES; + ENABLE_USER_SCRIPT_SANDBOXING = YES; + GCC_C_LANGUAGE_STANDARD = gnu17; + GCC_DYNAMIC_NO_PIC = NO; + GCC_NO_COMMON_BLOCKS = YES; + GCC_OPTIMIZATION_LEVEL = 0; + GCC_PREPROCESSOR_DEFINITIONS = ( + "DEBUG=1", + "$(inherited)", + ); + GCC_WARN_64_TO_32_BIT_CONVERSION = YES; + GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; + GCC_WARN_UNDECLARED_SELECTOR = YES; + GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; + GCC_WARN_UNUSED_FUNCTION = YES; + GCC_WARN_UNUSED_VARIABLE = YES; + IPHONEOS_DEPLOYMENT_TARGET = 17.4; + LOCALIZATION_PREFERS_STRING_CATALOGS = YES; + MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; + MTL_FAST_MATH = YES; + ONLY_ACTIVE_ARCH = YES; + SDKROOT = iphoneos; + SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)"; + SWIFT_OPTIMIZATION_LEVEL = "-Onone"; + }; + name = Debug; + }; + C0B37B922BE8226B00B2F80B /* Release */ = { + isa = XCBuildConfiguration; + buildSettings = { + ALWAYS_SEARCH_USER_PATHS = NO; + ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES; + CLANG_ANALYZER_NONNULL = YES; + CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; + CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; + CLANG_ENABLE_MODULES = YES; + CLANG_ENABLE_OBJC_ARC = YES; + CLANG_ENABLE_OBJC_WEAK = YES; + CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; + CLANG_WARN_BOOL_CONVERSION = YES; + CLANG_WARN_COMMA = YES; + CLANG_WARN_CONSTANT_CONVERSION = YES; + CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; + CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; + CLANG_WARN_DOCUMENTATION_COMMENTS = YES; + CLANG_WARN_EMPTY_BODY = YES; + CLANG_WARN_ENUM_CONVERSION = YES; + CLANG_WARN_INFINITE_RECURSION = YES; + CLANG_WARN_INT_CONVERSION = YES; + CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; + CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; + CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; + CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; + CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; + CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; + CLANG_WARN_STRICT_PROTOTYPES = YES; + CLANG_WARN_SUSPICIOUS_MOVE = YES; + CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; + CLANG_WARN_UNREACHABLE_CODE = YES; + CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; + COPY_PHASE_STRIP = NO; + DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; + ENABLE_NS_ASSERTIONS = NO; + ENABLE_STRICT_OBJC_MSGSEND = YES; + ENABLE_USER_SCRIPT_SANDBOXING = YES; + GCC_C_LANGUAGE_STANDARD = gnu17; + GCC_NO_COMMON_BLOCKS = YES; + GCC_WARN_64_TO_32_BIT_CONVERSION = YES; + GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; + GCC_WARN_UNDECLARED_SELECTOR = YES; + GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; + GCC_WARN_UNUSED_FUNCTION = YES; + GCC_WARN_UNUSED_VARIABLE = YES; + IPHONEOS_DEPLOYMENT_TARGET = 17.4; + LOCALIZATION_PREFERS_STRING_CATALOGS = YES; + MTL_ENABLE_DEBUG_INFO = NO; + MTL_FAST_MATH = YES; + SDKROOT = iphoneos; + SWIFT_COMPILATION_MODE = wholemodule; + VALIDATE_PRODUCT = YES; + }; + name = Release; + }; + C0B37B942BE8226B00B2F80B /* Debug */ = { + isa = XCBuildConfiguration; + buildSettings = { + ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; + ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; + CODE_SIGN_ENTITLEMENTS = MLCEngineExample/MLCEngineExample.entitlements; + CODE_SIGN_STYLE = Automatic; + CURRENT_PROJECT_VERSION = 1; + DEVELOPMENT_ASSET_PATHS = "\"MLCEngineExample/Preview Content\""; + DEVELOPMENT_TEAM = 3FR42MXLK9; + ENABLE_PREVIEWS = YES; + GENERATE_INFOPLIST_FILE = YES; + INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES; + INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; + INFOPLIST_KEY_UILaunchScreen_Generation = YES; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + IPHONEOS_DEPLOYMENT_TARGET = 16.0; + LD_RUNPATH_SEARCH_PATHS = ( + "$(inherited)", + "@executable_path/Frameworks", + ); + LIBRARY_SEARCH_PATHS = "${PROJECT_DIR}/dist/lib"; + MARKETING_VERSION = 1.0; + OTHER_LDFLAGS = ( + "-Wl,-all_load", + "-lmodel_iphone", + "-lmlc_llm", + "-ltvm_runtime", + "-ltokenizers_cpp", + "-lsentencepiece", + "-ltokenizers_c", + ); + PRODUCT_BUNDLE_IDENTIFIER = mlc.MLCEngineExample; + PRODUCT_NAME = "$(TARGET_NAME)"; + SWIFT_EMIT_LOC_STRINGS = YES; + SWIFT_VERSION = 5.0; + TARGETED_DEVICE_FAMILY = "1,2"; + }; + name = Debug; + }; + C0B37B952BE8226B00B2F80B /* Release */ = { + isa = XCBuildConfiguration; + buildSettings = { + ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; + ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; + CODE_SIGN_ENTITLEMENTS = MLCEngineExample/MLCEngineExample.entitlements; + CODE_SIGN_STYLE = Automatic; + CURRENT_PROJECT_VERSION = 1; + DEVELOPMENT_ASSET_PATHS = "\"MLCEngineExample/Preview Content\""; + DEVELOPMENT_TEAM = 3FR42MXLK9; + ENABLE_PREVIEWS = YES; + GENERATE_INFOPLIST_FILE = YES; + INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES; + INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES; + INFOPLIST_KEY_UILaunchScreen_Generation = YES; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight"; + IPHONEOS_DEPLOYMENT_TARGET = 16.0; + LD_RUNPATH_SEARCH_PATHS = ( + "$(inherited)", + "@executable_path/Frameworks", + ); + LIBRARY_SEARCH_PATHS = "${PROJECT_DIR}/dist/lib"; + MARKETING_VERSION = 1.0; + OTHER_LDFLAGS = ( + "-Wl,-all_load", + "-lmodel_iphone", + "-lmlc_llm", + "-ltvm_runtime", + "-ltokenizers_cpp", + "-lsentencepiece", + "-ltokenizers_c", + ); + PRODUCT_BUNDLE_IDENTIFIER = mlc.MLCEngineExample; + PRODUCT_NAME = "$(TARGET_NAME)"; + SWIFT_EMIT_LOC_STRINGS = YES; + SWIFT_VERSION = 5.0; + TARGETED_DEVICE_FAMILY = "1,2"; + }; + name = Release; + }; +/* End XCBuildConfiguration section */ + +/* Begin XCConfigurationList section */ + C0B37B802BE8226A00B2F80B /* Build configuration list for PBXProject "MLCEngineExample" */ = { + isa = XCConfigurationList; + buildConfigurations = ( + C0B37B912BE8226B00B2F80B /* Debug */, + C0B37B922BE8226B00B2F80B /* Release */, + ); + defaultConfigurationIsVisible = 0; + defaultConfigurationName = Release; + }; + C0B37B932BE8226B00B2F80B /* Build configuration list for PBXNativeTarget "MLCEngineExample" */ = { + isa = XCConfigurationList; + buildConfigurations = ( + C0B37B942BE8226B00B2F80B /* Debug */, + C0B37B952BE8226B00B2F80B /* Release */, + ); + defaultConfigurationIsVisible = 0; + defaultConfigurationName = Release; + }; +/* End XCConfigurationList section */ + +/* Begin XCLocalSwiftPackageReference section */ + C0B37B962BE8234D00B2F80B /* XCLocalSwiftPackageReference "../MLCSwift" */ = { + isa = XCLocalSwiftPackageReference; + relativePath = ../MLCSwift; + }; +/* End XCLocalSwiftPackageReference section */ + +/* Begin XCSwiftPackageProductDependency section */ + C04105DE2BEBC61B005A434D /* MLCSwift */ = { + isa = XCSwiftPackageProductDependency; + productName = MLCSwift; + }; + C0B37B972BE8234D00B2F80B /* MLCSwift */ = { + isa = XCSwiftPackageProductDependency; + productName = MLCSwift; + }; +/* End XCSwiftPackageProductDependency section */ + }; + rootObject = C0B37B7D2BE8226A00B2F80B /* Project object */; +} diff --git a/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.xcworkspace/contents.xcworkspacedata b/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.xcworkspace/contents.xcworkspacedata new file mode 100644 index 0000000..919434a --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.xcworkspace/contents.xcworkspacedata @@ -0,0 +1,7 @@ + + + + + diff --git a/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist b/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist new file mode 100644 index 0000000..3d4c1e5 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist @@ -0,0 +1,8 @@ + + + + + IDEDidComputeMac32BitWarning + + + diff --git a/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/AccentColor.colorset/Contents.json b/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/AccentColor.colorset/Contents.json new file mode 100644 index 0000000..eb87897 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/AccentColor.colorset/Contents.json @@ -0,0 +1,11 @@ +{ + "colors" : [ + { + "idiom" : "universal" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/AppIcon.appiconset/Contents.json b/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/AppIcon.appiconset/Contents.json new file mode 100644 index 0000000..13613e3 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/AppIcon.appiconset/Contents.json @@ -0,0 +1,13 @@ +{ + "images" : [ + { + "idiom" : "universal", + "platform" : "ios", + "size" : "1024x1024" + } + ], + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/Contents.json b/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/Contents.json new file mode 100644 index 0000000..73c0059 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample/Assets.xcassets/Contents.json @@ -0,0 +1,6 @@ +{ + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCEngineExample/MLCEngineExample/ContentView.swift b/ios/MLCEngineExample/MLCEngineExample/ContentView.swift new file mode 100644 index 0000000..650cd38 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample/ContentView.swift @@ -0,0 +1,21 @@ +// This is a minimum example App to interact with MLC Engine +// +// for a complete example, take a look at the MLCChat + +import SwiftUI + +struct ContentView: View { + @EnvironmentObject private var appState: AppState + // simply display text on the app + var body: some View { + HStack { + Text(appState.displayText) + Spacer() + } + .padding() + } +} + +#Preview { + ContentView() +} diff --git a/ios/MLCEngineExample/MLCEngineExample/MLCEngineExample.entitlements b/ios/MLCEngineExample/MLCEngineExample/MLCEngineExample.entitlements new file mode 100644 index 0000000..1250534 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample/MLCEngineExample.entitlements @@ -0,0 +1,10 @@ + + + + + com.apple.developer.kernel.extended-virtual-addressing + + com.apple.developer.kernel.increased-memory-limit + + + diff --git a/ios/MLCEngineExample/MLCEngineExample/MLCEngineExampleApp.swift b/ios/MLCEngineExample/MLCEngineExample/MLCEngineExampleApp.swift new file mode 100644 index 0000000..9dac7ad --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample/MLCEngineExampleApp.swift @@ -0,0 +1,80 @@ +// This is a minimum example App to interact with MLC Engine +// This app is mainly created with minimalism in mind for +// example and quick testing purposes. +// +// To build this app, select target My Mac(Designed for iPad) and run +// Make sure you run "mlc_llm package" first with "MLCChat" +// replaced by "MLCEngineExample" +// to ensure the "dist/bundle" folder populates with the right model file +// and we have the model lib packaged correctly +import Foundation +import SwiftUI + +import MLCSwift + +class AppState: ObservableObject { + // the MLC engine instance + private let engine = MLCEngine() + // obtain the local path to store models + // this that stores the model files in the dist folder + private let bundleURL = Bundle.main.bundleURL.appending(path: "bundle") + // model path, this must match a builtin + // file name in prepare_params.sh + private let modelPath = "Llama-3-8B-Instruct-q3f16_1-MLC" + // model lib identifier of within the packaged library + // make sure we run "mlc_llm package" + private let modelLib = "llama_q3f16_1" + + // this is a message to be displayed in app + @Published var displayText = "" + + public func runExample() { + // MLCEngine is a actor that can be called in an async context + Task { + let modelLocalPath = bundleURL.appending(path: modelPath).path() + // Step 0: load the engine + await engine.reload(modelPath: modelLocalPath, modelLib: modelLib) + + // run chat completion as in OpenAI API style + for await res in await engine.chat.completions.create( + messages: [ + ChatCompletionMessage( + role: .user, + content: "What is the meaning of life?" + ) + ], + stream_options: StreamOptions(include_usage: true) + ) { + // publish at main event loop + DispatchQueue.main.async { + // parse the result content in structured form + // and stream back to the display + if let finalUsage = res.usage { + self.displayText += "\n" + (finalUsage.extra?.asTextLabel() ?? "") + } else { + self.displayText += res.choices[0].delta.content!.asText() + } + } + } + } + } +} + + +@main +struct MLCEngineExampleApp: App { + private let appState = AppState() + + init() { + // we simply run test + // please checkout output in console + appState.runExample() + } + + var body: some Scene { + WindowGroup { + ContentView() + .environmentObject(appState) + } + } +} diff --git a/ios/MLCEngineExample/MLCEngineExample/Preview Content/Preview Assets.xcassets/Contents.json b/ios/MLCEngineExample/MLCEngineExample/Preview Content/Preview Assets.xcassets/Contents.json new file mode 100644 index 0000000..73c0059 --- /dev/null +++ b/ios/MLCEngineExample/MLCEngineExample/Preview Content/Preview Assets.xcassets/Contents.json @@ -0,0 +1,6 @@ +{ + "info" : { + "author" : "xcode", + "version" : 1 + } +} diff --git a/ios/MLCEngineExample/README.md b/ios/MLCEngineExample/README.md new file mode 100644 index 0000000..273b691 --- /dev/null +++ b/ios/MLCEngineExample/README.md @@ -0,0 +1,8 @@ +# MLCEngine Example + +Minimal example of MLCSwift API. + +Checkout [Documentation page](https://llm.mlc.ai/docs/deploy/ios.html) for more information. + +- run `mlc_llm package` +- open the Xcode project diff --git a/ios/MLCEngineExample/mlc-package-config.json b/ios/MLCEngineExample/mlc-package-config.json new file mode 100644 index 0000000..1d84dce --- /dev/null +++ b/ios/MLCEngineExample/mlc-package-config.json @@ -0,0 +1,12 @@ +{ + "device": "iphone", + "model_list": [ + { + "model": "HF://mlc-ai/Llama-3-8B-Instruct-q3f16_1-MLC", + "model_id": "Llama-3-8B-Instruct-q3f16_1-MLC", + "estimated_vram_bytes": 3316000000, + "bundle_weight": true, + "model_lib": "llama_q3f16_1" + } + ] +} diff --git a/ios/MLCSwift/Package.swift b/ios/MLCSwift/Package.swift new file mode 100644 index 0000000..a84755c --- /dev/null +++ b/ios/MLCSwift/Package.swift @@ -0,0 +1,32 @@ +// swift-tools-version:5.5 +// The swift-tools-version declares the minimum version of Swift required to build this package. + +import PackageDescription + +let package = Package( + name: "MLCSwift", + products: [ + .library( + name: "MLCSwift", + targets: ["MLCEngineObjC", "MLCSwift"] + ) + ], + dependencies: [], + targets: [ + .target( + name: "MLCEngineObjC", + path: "Sources/ObjC", + cxxSettings: [ + .headerSearchPath("../../tvm_home/include"), + .headerSearchPath("../../tvm_home/3rdparty/tvm-ffi/include"), + .headerSearchPath("../../tvm_home/3rdparty/tvm-ffi/3rdparty/dlpack/include") + ] + ), + .target( + name: "MLCSwift", + dependencies: ["MLCEngineObjC"], + path: "Sources/Swift" + ) + ], + cxxLanguageStandard: .cxx17 +) diff --git a/ios/MLCSwift/README.md b/ios/MLCSwift/README.md new file mode 100644 index 0000000..3a7c2b5 --- /dev/null +++ b/ios/MLCSwift/README.md @@ -0,0 +1,4 @@ +# MLCSwift + +This is a simple swift package that exposes the chat module to swift. +Checkout our [documentation](https://llm.mlc.ai/docs/) for more examples. diff --git a/ios/MLCSwift/Sources/ObjC/LLMEngine.mm b/ios/MLCSwift/Sources/ObjC/LLMEngine.mm new file mode 100644 index 0000000..d32ec8f --- /dev/null +++ b/ios/MLCSwift/Sources/ObjC/LLMEngine.mm @@ -0,0 +1,121 @@ +// +// LLMEngine.mm +// LLMEngine +// +#import +#import +#include + +#include "LLMEngine.h" + +#define TVM_USE_LIBBACKTRACE 0 + +#include +#include +#include +#include +#include + +using namespace tvm::runtime; +using tvm::ffi::Function; +using tvm::ffi::Module; +using tvm::ffi::Optional; +using tvm::ffi::String; +using tvm::ffi::TypedFunction; + +@implementation JSONFFIEngine { + // Internal c++ classes + // internal module backed by JSON FFI + Optional json_ffi_engine_; + // member functions + Function init_background_engine_func_; + Function unload_func_; + Function reload_func_; + Function reset_func_; + Function chat_completion_func_; + Function abort_func_; + Function run_background_loop_func_; + Function run_background_stream_back_loop_func_; + Function exit_background_loop_func_; +} + +- (instancetype)init { + if (self = [super init]) { + // load chat module + Function f_json_ffi_create = Function::GetGlobalRequired("mlc.json_ffi.CreateJSONFFIEngine"); + json_ffi_engine_ = f_json_ffi_create().cast(); + init_background_engine_func_ = + json_ffi_engine_.value()->GetFunction("init_background_engine").value_or(Function(nullptr)); + reload_func_ = json_ffi_engine_.value()->GetFunction("reload").value_or(Function(nullptr)); + unload_func_ = json_ffi_engine_.value()->GetFunction("unload").value_or(Function(nullptr)); + reset_func_ = json_ffi_engine_.value()->GetFunction("reset").value_or(Function(nullptr)); + chat_completion_func_ = + json_ffi_engine_.value()->GetFunction("chat_completion").value_or(Function(nullptr)); + abort_func_ = json_ffi_engine_.value()->GetFunction("abort").value_or(Function(nullptr)); + run_background_loop_func_ = + json_ffi_engine_.value()->GetFunction("run_background_loop").value_or(Function(nullptr)); + run_background_stream_back_loop_func_ = json_ffi_engine_.value() + ->GetFunction("run_background_stream_back_loop") + .value_or(Function(nullptr)); + exit_background_loop_func_ = + json_ffi_engine_.value()->GetFunction("exit_background_loop").value_or(Function(nullptr)); + + TVM_FFI_ICHECK(init_background_engine_func_ != nullptr); + TVM_FFI_ICHECK(reload_func_ != nullptr); + TVM_FFI_ICHECK(unload_func_ != nullptr); + TVM_FFI_ICHECK(reset_func_ != nullptr); + TVM_FFI_ICHECK(chat_completion_func_ != nullptr); + TVM_FFI_ICHECK(abort_func_ != nullptr); + TVM_FFI_ICHECK(run_background_loop_func_ != nullptr); + TVM_FFI_ICHECK(run_background_stream_back_loop_func_ != nullptr); + TVM_FFI_ICHECK(exit_background_loop_func_ != nullptr); + } + return self; +} + +- (void)initBackgroundEngine:(void (^)(NSString*))streamCallback { + TypedFunction internal_stream_callback([streamCallback](String value) { + streamCallback([NSString stringWithUTF8String:value.c_str()]); + }); + int device_type = kDLMetal; + int device_id = 0; + init_background_engine_func_(device_type, device_id, internal_stream_callback); +} + +- (void)reload:(NSString*)engineConfigJson { + std::string engine_config = engineConfigJson.UTF8String; + reload_func_(engine_config); +} + +- (void)unload { + unload_func_(); +} + +- (void)reset { + reset_func_(); +} + +- (void)chatCompletion:(NSString*)requestJSON requestID:(NSString*)requestID { + std::string request_json = requestJSON.UTF8String; + std::string request_id = requestID.UTF8String; + chat_completion_func_(request_json, request_id); +} + +- (void)abort:(NSString*)requestID { + std::string request_id = requestID.UTF8String; + abort_func_(request_id); +} + +- (void)runBackgroundLoop { + run_background_loop_func_(); +} + +- (void)runBackgroundStreamBackLoop { + run_background_stream_back_loop_func_(); +} + +- (void)exitBackgroundLoop { + exit_background_loop_func_(); +} + +@end diff --git a/ios/MLCSwift/Sources/ObjC/include/LLMEngine.h b/ios/MLCSwift/Sources/ObjC/include/LLMEngine.h new file mode 100644 index 0000000..22fc4ef --- /dev/null +++ b/ios/MLCSwift/Sources/ObjC/include/LLMEngine.h @@ -0,0 +1,32 @@ +// +// Use this file to import your target's public headers that you would like to expose to Swift. +// LLM Chat Module +// +// Exposed interface of Object-C, enables swift binding. +#import +#import + +/** + * This is an internal Raw JSON FFI Engine that redirects request to internal JSON FFI Engine in C++ + */ +@interface JSONFFIEngine : NSObject + +- (void)initBackgroundEngine:(void (^)(NSString*))streamCallback; + +- (void)reload:(NSString*)engineConfig; + +- (void)unload; + +- (void)reset; + +- (void)chatCompletion:(NSString*)requestJSON requestID:(NSString*)requestID; + +- (void)abort:(NSString*)requestID; + +- (void)runBackgroundLoop; + +- (void)runBackgroundStreamBackLoop; + +- (void)exitBackgroundLoop; + +@end diff --git a/ios/MLCSwift/Sources/Swift/LLMEngine.swift b/ios/MLCSwift/Sources/Swift/LLMEngine.swift new file mode 100644 index 0000000..1342bf6 --- /dev/null +++ b/ios/MLCSwift/Sources/Swift/LLMEngine.swift @@ -0,0 +1,232 @@ +import Foundation +import MLCEngineObjC +import os + +class BackgroundWorker : Thread { + private var task: ()->Void; + + public init(task: @escaping () -> Void) { + self.task = task + } + + public override func main() { + self.task(); + } +} + +@available(iOS 14.0.0, *) +public class MLCEngine { + struct RequestState { + let request: ChatCompletionRequest + let continuation: AsyncStream.Continuation + + init( + request: ChatCompletionRequest, + continuation: AsyncStream.Continuation + ) { + self.request = request + self.continuation = continuation + } + } + + // internal engine state + // that maintains logger and continuations + // we decouple it from MLCEngine + // and explicitly pass in jsonFFIEngine + // so there is no cyclic dependency + // when we capture things + actor EngineState { + public let logger = Logger() + private var requestStateMap = Dictionary() + + // completion function + func chatCompletion( + jsonFFIEngine: JSONFFIEngine, + request: ChatCompletionRequest + ) -> AsyncStream { + let encoder = JSONEncoder() + let data = try! encoder.encode(request) + let jsonRequest = String(data: data, encoding: .utf8)! + // generate a UUID for the request + let requestID = UUID().uuidString + let stream = AsyncStream(ChatCompletionStreamResponse.self) { continuation in + continuation.onTermination = { termination in + if termination == .cancelled { + jsonFFIEngine.abort(requestID); + } + } + // store continuation map for further callbacks + self.requestStateMap[requestID] = RequestState( + request: request, continuation: continuation + ) + // start invoking engine for completion + jsonFFIEngine.chatCompletion(jsonRequest, requestID: requestID) + } + return stream + } + + func streamCallback(result: String?) { + var responses: [ChatCompletionStreamResponse] = [] + + let decoder = JSONDecoder() + do { + responses = try decoder.decode([ChatCompletionStreamResponse].self, from: result!.data(using: .utf8)!) + } catch let lastError { + logger.error("Swift json parsing error: error=\(lastError), jsonsrc=\(result!)") + } + + // dispatch to right request ID + for res in responses { + if let requestState = self.requestStateMap[res.id] { + // final chunk always come with usage + if let finalUsage = res.usage { + if let include_usage = requestState.request.stream_options?.include_usage { + if include_usage { + requestState.continuation.yield(res) + } + } + requestState.continuation.finish() + self.requestStateMap.removeValue(forKey: res.id) + } else { + requestState.continuation.yield(res) + } + } + } + // Todo(mlc-team): check the last error in engine and report if there's any + } + } + + public class Completions { + private let jsonFFIEngine: JSONFFIEngine + private let state: EngineState + + init(jsonFFIEngine: JSONFFIEngine, state: EngineState) { + self.jsonFFIEngine = jsonFFIEngine + self.state = state + } + + private func create( + request: ChatCompletionRequest + ) async -> AsyncStream { + return await state.chatCompletion(jsonFFIEngine: jsonFFIEngine, request: request) + } + + // offer a direct convenient method to pass in messages + public func create( + messages: [ChatCompletionMessage], + model: Optional = nil, + frequency_penalty: Optional = nil, + presence_penalty: Optional = nil, + logprobs: Bool = false, + top_logprobs: Int = 0, + logit_bias: Optional<[Int : Float]> = nil, + max_tokens: Optional = nil, + n: Int = 1, + seed: Optional = nil, + stop: Optional<[String]> = nil, + stream: Bool = true, + stream_options: Optional = nil, + temperature: Optional = nil, + top_p: Optional = nil, + tools: Optional<[ChatTool]> = nil, + user: Optional = nil, + response_format: Optional = nil + ) async -> AsyncStream { + if !stream { + state.logger.error("Only stream=true is supported in MLCSwift") + } + let request = ChatCompletionRequest( + messages: messages, + model: model, + frequency_penalty: frequency_penalty, + presence_penalty: presence_penalty, + logprobs: logprobs, + top_logprobs: top_logprobs, + logit_bias: logit_bias, + max_tokens: max_tokens, + n: n, + seed: seed, + stop: stop, + stream: stream, + stream_options: stream_options, + temperature: temperature, + top_p: top_p, + tools: tools, + user: user, + response_format: response_format + ) + return await self.create(request: request) + } + } + + public class Chat { + public let completions: Completions + + init(jsonFFIEngine: JSONFFIEngine, state: EngineState) { + self.completions = Completions( + jsonFFIEngine: jsonFFIEngine, + state: state + ) + } + } + + private let state : EngineState; + private let jsonFFIEngine: JSONFFIEngine; + public let chat : Chat; + private var threads = Array(); + + public init() { + let state_ = EngineState(); + let jsonFFIEngine_ = JSONFFIEngine(); + + self.chat = Chat(jsonFFIEngine: jsonFFIEngine_, state: state_) + self.jsonFFIEngine = jsonFFIEngine_ + self.state = state_ + + // note: closure do not capture self + jsonFFIEngine_.initBackgroundEngine { + [state_](result : String?) -> Void in + state_.streamCallback(result: result) + } + let backgroundWorker = BackgroundWorker { [jsonFFIEngine_] in + Thread.setThreadPriority(1) + jsonFFIEngine_.runBackgroundLoop() + } + let backgroundStreamBackWorker = BackgroundWorker { + [jsonFFIEngine_] in + jsonFFIEngine_.runBackgroundStreamBackLoop() + } + // set background worker to be high QoS so it gets higher p for gpu + backgroundWorker.qualityOfService = QualityOfService.userInteractive + threads.append(backgroundWorker) + threads.append(backgroundStreamBackWorker) + backgroundWorker.start() + backgroundStreamBackWorker.start() + } + + deinit { + jsonFFIEngine.exitBackgroundLoop() + } + + // The following functions do not have to be async for now + // But to be safe and consistent with chat.completions.create + // and for future API changes we keep them as async calls + public func reload(modelPath: String, modelLib: String) async { + let engineConfig = """ + { + "model": "\(modelPath)", + "model_lib": "system://\(modelLib)", + "mode": "interactive" + } + """ + jsonFFIEngine.reload(engineConfig) + } + + public func reset() async { + jsonFFIEngine.reset() + } + + public func unload() async { + jsonFFIEngine.unload() + } +} diff --git a/ios/MLCSwift/Sources/Swift/OpenAIProtocol.swift b/ios/MLCSwift/Sources/Swift/OpenAIProtocol.swift new file mode 100644 index 0000000..e0c5698 --- /dev/null +++ b/ios/MLCSwift/Sources/Swift/OpenAIProtocol.swift @@ -0,0 +1,289 @@ +// Protocol definition of OpenAI API +import Foundation + +// Protocols for v1/chat/completions +// API reference: https://platform.openai.com/docs/api-reference/chat/create + +public struct TopLogProbs : Codable { + public var token: String + public var logprob: Float + public var bytes: Optional<[Int]> +} + +public struct LogProbsContent : Codable { + public var token: String + public var logprob: Float + public var bytes: Optional<[Int]> = nil + public var top_logprobs: [TopLogProbs] = [] +} + +public struct LogProbs : Codable { + public var content: [LogProbsContent] = [] +} + +public struct ChatFunction : Codable { + public var name: String + public var description: Optional = nil + public var parameters: [String: String] + + public init( + name: String, + description: Optional = nil, + parameters: [String : String] + ) { + self.name = name + self.description = description + self.parameters = parameters + } +} + +public struct ChatTool : Codable { + public var type: String = "function" + public let function: ChatFunction + + public init(type: String, function: ChatFunction) { + self.type = type + self.function = function + } +} + +public struct ChatFunctionCall : Codable { + public var name: String + // NOTE: arguments shold be dict str to any codable + // for now only allow string output due to typing issues + public var arguments: Optional<[String: String]> = nil + + public init(name: String, arguments: Optional<[String : String]> = nil) { + self.name = name + self.arguments = arguments + } +} + +public struct ChatToolCall : Codable { + public var id: String = UUID().uuidString + public var type: String = "function" + public var function: ChatFunctionCall + + public init( + id: String = UUID().uuidString, + type: String = "function", + function: ChatFunctionCall + ) { + self.id = id + self.type = type + self.function = function + } +} + +public enum ChatCompletionRole: String, Codable { + case system = "system" + case user = "user" + case assistant = "assistant" + case tool = "tool" +} + +public enum ChatCompletionMessageContent: Codable { + case text(String) + case parts([[String: String]]) + + public init(from decoder: Decoder) throws { + let container = try decoder.singleValueContainer() + if let text = try? container.decode(String.self) { + self = .text(text) + } else { + let parts = try container.decode([[String: String]].self) + self = .parts(parts) + } + } + + public func encode(to encoder: Encoder) throws { + var container = encoder.singleValueContainer() + switch self { + case .text(let text): try container.encode(text) + case .parts(let parts): try container.encode(parts) + } + } + + public func asText() -> String { + switch (self) { + case .text(let text): return text + case .parts(let parts): + var res = "" + for item in parts { + if item["type"]! == "text" { + res += item["text"]! + } + } + return res + } + } +} + +public struct ChatCompletionMessage: Codable { + public var role: ChatCompletionRole + public var content: Optional = nil + public var name: Optional = nil + public var tool_calls: Optional<[ChatToolCall]> = nil + public var tool_call_id: Optional = nil + + // more complicated content construction + public init( + role: ChatCompletionRole, + content: Optional<[[String : String]]> = nil, + name: Optional = nil, + tool_calls: Optional<[ChatToolCall]> = nil, + tool_call_id: Optional = nil + ) { + self.role = role + if let cvalue = content { + self.content = .parts(cvalue) + } else { + self.content = nil + } + self.name = name + self.tool_calls = tool_calls + self.tool_call_id = tool_call_id + } + + // convenient method to construct content from string + public init( + role: ChatCompletionRole, + content: String, + name: Optional = nil, + tool_calls: Optional<[ChatToolCall]> = nil, + tool_call_id: Optional = nil + ) { + self.role = role + self.content = .text(content) + self.name = name + self.tool_calls = tool_calls + self.tool_call_id = tool_call_id + } +} + +public struct ChatCompletionStreamResponseChoice: Codable { + public var finish_reason: Optional = nil + public var index: Int + public var delta: ChatCompletionMessage + public var lobprobs: Optional = nil +} + +public struct CompletionUsageExtra: Codable { + public var prefill_tokens_per_s: Optional = nil + public var decode_tokens_per_s: Optional = nil + public var num_prefill_tokens: Optional = nil + + public func asTextLabel() -> String { + var outputText = "" + if let prefill_tokens_per_s = self.prefill_tokens_per_s { + outputText += "prefill: " + outputText += String(format: "%.1f", prefill_tokens_per_s) + outputText += " tok/s" + } + if let decode_tokens_per_s = self.decode_tokens_per_s { + if !outputText.isEmpty { + outputText += ", " + } + outputText += "decode: " + outputText += String(format: "%.1f", decode_tokens_per_s) + outputText += " tok/s" + } + return outputText + } +} + +public struct CompletionUsage: Codable { + public var prompt_tokens: Int + public var completion_tokens: Int + public var total_tokens: Int + public var extra: Optional +} + +public struct ChatCompletionStreamResponse: Codable { + public var id : String + public var choices: [ChatCompletionStreamResponseChoice] = [] + public var created: Optional = nil + public var model: Optional = nil + public var system_fingerprint: String + public var object: Optional = nil + public var usage: Optional = nil +} + +public struct ResponseFormat: Codable { + public var type: String + public var schema: Optional = nil + + public init(type: String, schema: Optional = nil) { + self.type = type + self.schema = schema + } +} + +public struct StreamOptions: Codable { + public var include_usage: Bool = false + + public init(include_usage: Bool) { + self.include_usage = include_usage + } +} + +public struct ChatCompletionRequest: Codable { + public var messages: [ChatCompletionMessage] + public var model: Optional = nil + public var frequency_penalty: Optional = nil + public var presence_penalty: Optional = nil + public var logprobs: Bool = false + public var top_logprobs: Int = 0 + public var logit_bias: Optional<[Int: Float]> = nil + public var max_tokens: Optional = nil + public var n: Int = 1 + public var seed: Optional = nil + public var stop: Optional<[String]> = nil + public var stream: Bool = true + public var stream_options: Optional = nil + public var temperature: Optional = nil + public var top_p: Optional = nil + public var tools: Optional<[ChatTool]> = nil + public var user: Optional = nil + public var response_format: Optional = nil + + public init( + messages: [ChatCompletionMessage], + model: Optional = nil, + frequency_penalty: Optional = nil, + presence_penalty: Optional = nil, + logprobs: Bool = false, + top_logprobs: Int = 0, + logit_bias: Optional<[Int : Float]> = nil, + max_tokens: Optional = nil, + n: Int = 1, + seed: Optional = nil, + stop: Optional<[String]> = nil, + stream: Bool = true, + stream_options: Optional = nil, + temperature: Optional = nil, + top_p: Optional = nil, + tools: Optional<[ChatTool]> = nil, + user: Optional = nil, + response_format: Optional = nil + ) { + self.messages = messages + self.model = model + self.frequency_penalty = frequency_penalty + self.presence_penalty = presence_penalty + self.logprobs = logprobs + self.top_logprobs = top_logprobs + self.logit_bias = logit_bias + self.max_tokens = max_tokens + self.n = n + self.seed = seed + self.stop = stop + self.stream = stream + self.stream_options = stream_options + self.temperature = temperature + self.top_p = top_p + self.tools = tools + self.user = user + self.response_format = response_format + } +} diff --git a/ios/README.md b/ios/README.md new file mode 100644 index 0000000..39f0e0b --- /dev/null +++ b/ios/README.md @@ -0,0 +1,3 @@ +# MLC-LLM iOS + +[Documentation page](https://llm.mlc.ai/docs/deploy/ios.html) diff --git a/ios/prepare_libs.sh b/ios/prepare_libs.sh new file mode 100755 index 0000000..bd1d89b --- /dev/null +++ b/ios/prepare_libs.sh @@ -0,0 +1,141 @@ +# Command to prepare the mlc llm static libraries +# This command will be invoked by the "mlc_llm package" command +function help { + echo -e "OPTION:" + echo -e " -s, --simulator Build for Simulator" + echo -e " -a, --arch x86_64 | arm64 Simulator arch " + echo -e " -c, --catalyst Build for Mac Catalyst (arm64 only)" + echo -e " --deployment-target VERSION Mac Catalyst deployment target (default: 18.0)" + echo -e " -h, --help Prints this help\n" +} + +MLC_LLM_SOURCE_DIR="${MLC_LLM_SOURCE_DIR:-..}" +is_simulator="false" +is_catalyst="false" +arch="arm64" +deployment_target="18.0" + +# rustup is required to install iOS target stdlibs, and we need to make sure +# the rustup-managed cargo/rustc are used during cross compilation. +if [ -d "${HOME}/.cargo/bin" ]; then + export PATH="${HOME}/.cargo/bin:${PATH}" +fi +if ! command -v rustup >/dev/null 2>&1; then + echo "error: rustup is required to build iOS static libraries." >&2 + echo "Install rustup and retry, e.g.:" >&2 + echo " curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y" >&2 + exit 1 +fi + +# Args while-loop +while [ "$1" != "" ]; +do + case $1 in + -s | --simulator ) is_simulator="true" + ;; + -c | --catalyst ) is_catalyst="true" + ;; + -a | --arch ) shift + arch=$1 + ;; + --deployment-target ) shift + deployment_target=$1 + ;; + -h | --help ) help + exit + ;; + *) + echo "$script: illegal option $1" + usage + exit 1 # error + ;; + esac + shift +done + +set -euxo pipefail + +sysroot="iphoneos" +type="Release" +build_dir="build" + +if [ "$is_catalyst" = "true" ]; then + if [ "$is_simulator" = "true" ]; then + echo "error: --simulator is not supported with --catalyst." >&2 + exit 1 + fi + if [ "$arch" != "x86_64" ]; then + arch="arm64" + fi + sysroot="macosx" + build_dir="build-maccatalyst-$arch" +fi + +if [ "$is_simulator" = "true" ]; then + if [ "$arch" = "arm64" ]; then + # iOS simulator on Apple processors + rustup target add aarch64-apple-ios-sim + else + # iOS simulator on x86 processors + rustup target add x86_64-apple-ios + fi + sysroot="iphonesimulator" + type="Debug" +else + # iOS devices + rustup target add aarch64-apple-ios + if [ "$is_catalyst" = "true" ]; then + if [ "$arch" = "x86_64" ]; then + rustup target add x86_64-apple-ios-macabi + else + rustup target add aarch64-apple-ios-macabi + fi + fi +fi + +mkdir -p "$build_dir" && cd "$build_dir" + +cmake_args=( + -DCMAKE_BUILD_TYPE="$type" + -DCMAKE_BUILD_WITH_INSTALL_NAME_DIR=ON + -DCMAKE_SKIP_INSTALL_ALL_DEPENDENCY=ON + -DCMAKE_INSTALL_PREFIX=. + -DCMAKE_CXX_FLAGS="-O3" + -DMLC_LLM_INSTALL_STATIC_LIB=ON + -DUSE_METAL=ON + -DTVM_FFI_USE_LIBBACKTRACE=OFF + -DTVM_FFI_BACKTRACE_ON_SEGFAULT=OFF +) + +if [ "$is_catalyst" = "true" ]; then + toolchain="$MLC_LLM_SOURCE_DIR/3rdparty/tokenizers-cpp/sentencepiece/cmake/ios.toolchain.cmake" + if [ "$arch" = "x86_64" ]; then + platform="MAC_CATALYST" + else + platform="MAC_CATALYST_ARM64" + fi + cmake_args+=( + -DCMAKE_TOOLCHAIN_FILE="$toolchain" + -DPLATFORM="$platform" + -DDEPLOYMENT_TARGET="$deployment_target" + -DENABLE_BITCODE=OFF + ) +else + cmake_args+=( + -DCMAKE_SYSTEM_NAME=iOS + -DCMAKE_SYSTEM_VERSION=14.0 + -DCMAKE_OSX_SYSROOT="$sysroot" + -DCMAKE_OSX_ARCHITECTURES="$arch" + -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 + ) +fi + +cmake "$MLC_LLM_SOURCE_DIR" "${cmake_args[@]}" + + +cmake --build . --config release --target mlc_llm_static -j +cmake --build . --target install --config release -j +cd .. + +rm -rf $MLC_LLM_SOURCE_DIR/ios/MLCSwift/tvm_home +ln -s $MLC_LLM_SOURCE_DIR/3rdparty/tvm $MLC_LLM_SOURCE_DIR/ios/MLCSwift/tvm_home diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..973f0df --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,177 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[project] +name = "mlc_llm" +# Note: Call version.py to update the version before building the wheel +version = "0.20.0.dev0" +description = "MLC LLM: an universal LLM deployment engine via ML compilation." + +authors = [{ name = "MLC LLM Contributors" }] +readme = "README.md" +license = { text = "Apache 2.0" } +classifiers = [ + "License :: OSI Approved :: Apache Software License", + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", +] +keywords = ["machine learning"] +requires-python = ">=3.9" + +dependencies = [ + "apache-tvm-ffi", + "datasets", + "fastapi", + "flashinfer-python; sys_platform == 'linux'", + "ml_dtypes>=0.5.1", + "openai", + "pandas", + "prompt_toolkit", + "requests", + "safetensors", + "sentencepiece", + "shortuuid", + "tiktoken", + "torch", + "tqdm", + "transformers", + "uvicorn", +] + +[project.urls] +Homepage = "https://llm.mlc.ai/" +Documentation = "https://llm.mlc.ai/docs/" +Repository = "https://github.com/mlc-ai/mlc-llm" +"Bug Tracker" = "https://github.com/mlc-ai/mlc-llm/issues" + + +[build-system] +requires = ["scikit-build-core>=0.10.0"] +build-backend = "scikit_build_core.build" + +[tool.scikit-build] +# Point to the root CMakeLists.txt +cmake.source-dir = "." +cmake.build-type = "Release" + +# Configure the wheel to be Python version-agnostic +wheel.py-api = "py3" + +# Build configuration +build-dir = "build" + +# CMake configuration - ensure proper installation paths +cmake.args = ["-DMLC_LLM_BUILD_PYTHON_MODULE=ON"] + +# Wheel configuration +wheel.packages = ["python/mlc_llm"] +wheel.install-dir = "mlc_llm" + +# Source distribution configuration +sdist.include = [ + # Build files + "/CMakeLists.txt", + "/pyproject.toml", + "/cmake/**/*", + "/3rdparty/**/*", + + # Source code + "/src/**/*.cc", + "/src/**/*.h", + + # Python source + "/python/mlc_llm/**/*.py", + + # Documentation and metadata + "/docs/**/*", + "/LICENSE", + "/README.md", + "/NOTICE", + + # Tests + "/tests/**/*", + + "/.pre-commit-config.yaml", +] + +sdist.exclude = [ + "**/.git", + "**/.github", + "**/__pycache__", + "**/*.pyc", + "build", + "dist", + "3rdparty/tvm", + "**/3rdparty/*/docs", + "**/3rdparty/*/media", + "**/3rdparty/*/examples", + "**/3rdparty/*/test", +] + +# Logging +logging.level = "INFO" + + +[tool.ruff] +include = [ + "python/**/*.py", + "tests/**/*.py", + "examples/**/*.py", + "docs/**/*.py", + "cmake/**/*.py", + "version.py", +] +line-length = 100 +indent-width = 4 +target-version = "py39" +exclude = ["3rdparty", "build", "dist", ".venv", ".ruff_cache"] + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "F", # pyflakes + "I", # isort + "UP", # pyupgrade + "RUF", # ruff-specific rules +] +ignore = [ + "UP028", # yield-in-for-loop + "UP030", # format-literals + "UP031", # printf-string-formatting + "RUF010", # explicit-f-string-type-conversion — !s/!r/!a is obscure; prefer explicit str()/repr() + "RUF043", # pytest-raises-ambiguous-pattern + "RUF059", # unused-unpacked-variable +] +fixable = ["ALL"] +unfixable = [] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["E402", "F401", "F403", "F405"] + +[tool.ruff.lint.isort] +known-first-party = ["mlc_llm"] + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" +skip-magic-trailing-comma = false +line-ending = "auto" + +[dependency-groups] +lint = ["pre-commit"] diff --git a/python/mlc_llm/__init__.py b/python/mlc_llm/__init__.py new file mode 100644 index 0000000..ee7da3a --- /dev/null +++ b/python/mlc_llm/__init__.py @@ -0,0 +1,20 @@ +"""MLC Chat python package. + +MLC Chat is the app runtime of MLC LLM. +""" + +from tvm import register_global_func + +from . import protocol, serve +from .libinfo import __version__ +from .serve import AsyncMLCEngine, MLCEngine + + +@register_global_func("runtime.disco.create_socket_session_local_workers", override=True) +def _create_socket_session_local_workers(num_workers): + """Create the local session for each distributed node over socket session.""" + from tvm.runtime.disco import ( + ProcessSession, + ) + + return ProcessSession(num_workers, num_groups=1, entrypoint="mlc_llm.cli.worker") diff --git a/python/mlc_llm/__main__.py b/python/mlc_llm/__main__.py new file mode 100644 index 0000000..c281b4a --- /dev/null +++ b/python/mlc_llm/__main__.py @@ -0,0 +1,67 @@ +"""Entrypoint of all CLI commands from MLC LLM""" + +import sys + +from mlc_llm.support import logging +from mlc_llm.support.argparse import ArgumentParser + +logging.enable_logging() + + +def main(): + """Entrypoint of all CLI commands from MLC LLM""" + parser = ArgumentParser("MLC LLM Command Line Interface.") + parser.add_argument( + "subcommand", + type=str, + choices=[ + "compile", + "convert_weight", + "gen_config", + "chat", + "serve", + "package", + "calibrate", + "router", + ], + help="Subcommand to to run. (choices: %(choices)s)", + ) + parsed = parser.parse_args(sys.argv[1:2]) + if parsed.subcommand == "compile": + from mlc_llm.cli import compile as cli + + cli.main(sys.argv[2:]) + elif parsed.subcommand == "convert_weight": + from mlc_llm.cli import convert_weight as cli + + cli.main(sys.argv[2:]) + elif parsed.subcommand == "gen_config": + from mlc_llm.cli import gen_config as cli + + cli.main(sys.argv[2:]) + elif parsed.subcommand == "chat": + from mlc_llm.cli import chat as cli + + cli.main(sys.argv[2:]) + elif parsed.subcommand == "serve": + from mlc_llm.cli import serve as cli + + cli.main(sys.argv[2:]) + elif parsed.subcommand == "package": + from mlc_llm.cli import package as cli + + cli.main(sys.argv[2:]) + elif parsed.subcommand == "calibrate": + from mlc_llm.cli import calibrate as cli + + cli.main(sys.argv[2:]) + elif parsed.subcommand == "router": + from mlc_llm.cli import router as cli + + cli.main(sys.argv[2:]) + else: + raise ValueError(f"Unknown subcommand {parsed.subcommand}") + + +if __name__ == "__main__": + main() diff --git a/python/mlc_llm/base.py b/python/mlc_llm/base.py new file mode 100644 index 0000000..d6d3942 --- /dev/null +++ b/python/mlc_llm/base.py @@ -0,0 +1,45 @@ +"""Load MLC LLM library and _ffi_api functions.""" + +import ctypes +import os +import sys + +import tvm +import tvm.base + +from . import libinfo + +SKIP_LOADING_MLCLLM_SO = os.environ.get("SKIP_LOADING_MLCLLM_SO", "0") + + +def _load_mlc_llm_lib(): + """Load MLC LLM lib""" + if sys.platform.startswith("win32") and sys.version_info >= (3, 8): + for path in libinfo.get_dll_directories(): + os.add_dll_directory(path) + lib_name = "mlc_llm" if tvm.base._RUNTIME_ONLY else "mlc_llm_module" + lib_path = libinfo.find_lib_path(lib_name, optional=False) + return ctypes.CDLL(lib_path[0]), lib_path[0] + + +@tvm.register_global_func("mlc.debug_cuda_profiler_start") +def _debug_cuda_profiler_start() -> None: + """Start cuda profiler.""" + import cuda + import cuda.cudart + + cuda.cudart.cudaProfilerStart() + + +@tvm.register_global_func("mlc.debug_cuda_profiler_stop") +def _debug_cuda_profiler_stop() -> None: + """Stop cuda profiler.""" + import cuda + import cuda.cudart + + cuda.cudart.cudaProfilerStop() + + +# only load once here +if SKIP_LOADING_MLCLLM_SO == "0": + _LIB, _LIB_PATH = _load_mlc_llm_lib() diff --git a/python/mlc_llm/bench/__init__.py b/python/mlc_llm/bench/__init__.py new file mode 100644 index 0000000..f8fc6a6 --- /dev/null +++ b/python/mlc_llm/bench/__init__.py @@ -0,0 +1 @@ +"""Subdirectory of bench.""" diff --git a/python/mlc_llm/bench/__main__.py b/python/mlc_llm/bench/__main__.py new file mode 100644 index 0000000..d63242e --- /dev/null +++ b/python/mlc_llm/bench/__main__.py @@ -0,0 +1,403 @@ +"""MLC LLM benchmark main entrance""" + +import functools +import json +import random +from typing import Any, Dict, List, Optional, Tuple # noqa: UP035 + +import numpy as np +import requests +from transformers import AutoTokenizer + +import mlc_llm +from mlc_llm.bench.api_endpoint import SUPPORTED_BACKENDS, create_api_endpoint +from mlc_llm.bench.dataset import SUPPORTED_DATASET, Dataset, create_dataset +from mlc_llm.bench.request_processor import ( + MetricAnalyzer, + RequestProcessor, + create_pipelines, +) +from mlc_llm.bench.request_record import ( + RequestRecord, + convert_reports_to_df, + generate_metrics_summary, + pretty_print_report, +) +from mlc_llm.cli.serve import EngineConfigOverride +from mlc_llm.serve import EngineConfig +from mlc_llm.support import argparse, logging + +logging.enable_logging() +logger = logging.getLogger(__name__) + + +def _parse_num_concurrent_requests(num_str: Optional[str]) -> Optional[List[int]]: # noqa: UP006 + if num_str is None: + return None + numbers = num_str.split(",") + if any(not number.isdigit() for number in numbers): + raise ValueError(f"Unrecognized num_concurrent_requests list: {numbers}") + return list(int(number) for number in numbers) + + +def _parse_request_rate(request_rate_str: Optional[str]) -> Optional[List[np.float32]]: # noqa: UP006 + if request_rate_str is None: + return None + request_rates = request_rate_str.split(",") + results = [] + for rate_str in request_rates: + request_rate = float(rate_str) + if request_rate <= 0: + raise ValueError(f"Invalid request rate {request_rate}") + results.append(np.float32(request_rate)) + return results + + +def _parse_mlc_engine_config(config_str: Optional[str]) -> EngineConfig: + if config_str is None: + return None + engine_config_override = EngineConfigOverride.from_str(config_str) + return EngineConfig( + tensor_parallel_shards=engine_config_override.tensor_parallel_shards, + max_num_sequence=engine_config_override.max_num_sequence, + max_total_sequence_length=engine_config_override.max_total_seq_length, + prefill_chunk_size=engine_config_override.prefill_chunk_size, + sliding_window_size=engine_config_override.sliding_window_size, + attention_sink_size=engine_config_override.attention_sink_size, + max_history_size=engine_config_override.max_history_size, + gpu_memory_utilization=engine_config_override.gpu_memory_utilization, + spec_draft_length=engine_config_override.spec_draft_length, + prefill_mode=engine_config_override.prefill_mode, + prefix_cache_max_num_recycling_seqs=engine_config_override.prefix_cache_max_num_recycling_seqs, + prefix_cache_mode=engine_config_override.prefix_cache_mode, + ) + + +def _launch_mlc_server(args: argparse.argparse.Namespace): + return mlc_llm.serve.PopenServer( + model=args.tokenizer, + mode="server", + model_lib=args.mlc_model_lib, + enable_tracing=False, + host=args.host, + port=args.port, + engine_config=args.mlc_engine_config, + ) + + +def run_pipeline( + pipeline: RequestProcessor, + dataset: Dataset, + tokenizer: AutoTokenizer, + args: argparse.argparse.Namespace, +) -> Tuple[Dict[str, Any], List[RequestRecord]]: # noqa: UP006 + """Run the pipeline with the given dataset and args. Return the benchmark report dict.""" + random.seed(args.seed) + np.random.seed(args.seed) + request_records = dataset.generate_request_records( + args.input_len, + args.output_len, + args.input_len_std, + args.output_len_std, + ) + request_records = pipeline(request_records) + num_total_requests = ( + args.num_requests if not args.per_gpu_workload else args.num_requests * args.num_gpus + ) + assert len(request_records) == num_total_requests + sorted_requests: List[RequestRecord] = [None] * num_total_requests # noqa: UP006 + for request_record in request_records: + assert request_record.request_id is not None + assert sorted_requests[request_record.request_id] is None + sorted_requests[request_record.request_id] = request_record + + request_records = MetricAnalyzer(tokenizer)(request_records) + report = generate_metrics_summary(request_records, num_total_requests, args.num_gpus) + return report, sorted_requests + + +def query_mlc_server_metrics(host: str, port: int): + """Try to get the MLC server metrics whenever it exists.""" + try: + r = requests.post(f"http://{host}:{port}/debug/dump_engine_metrics", json={}, timeout=10) + if r.status_code == 200: + print(f"MLC server metrics: {r.json()}") + except Exception: + pass + + +def main(args: argparse.argparse.Namespace): + """Main benchmark entrance.""" + mlc_server = None + if args.mlc_model_lib: + mlc_server = _launch_mlc_server(args) + if args.num_requests <= 0: + raise ValueError("Number of requests to benchmark must be positive.") + + def _main(): + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) + dataset = create_dataset(args, tokenizer) + f_create_api_endpoint = functools.partial(create_api_endpoint, args) + pipelines = create_pipelines(args, f_create_api_endpoint, dataset) + reports = [] + alltime_records = {} + for i, pipeline in enumerate(pipelines): + report, request_records = run_pipeline(pipeline, dataset, tokenizer, args) + exec_feature = ( + json.dumps(report["exec_feature"]) + if report["exec_feature"] is not None + else f"pipeline{i}" + ) + alltime_records[exec_feature] = [ + request_record.model_dump() for request_record in request_records + ] + reports.append(report) + pretty_print_report(report) + query_mlc_server_metrics(args.host, args.port) + + # Construct data frame + df = convert_reports_to_df(reports) + print(df) + df.to_csv(args.output, index=False) + logger.info("Benchmark results dumped to file %s", args.output) + if args.debug_dump: + debug_dump_filepath = ( + args.output[:-4] if args.output.endswith(".csv") else args.output + ) + "_debug_dump.log" + with open(debug_dump_filepath, "w", encoding="utf-8") as file: + json.dump(alltime_records, file, indent=4) + logger.info("Debug log dumped to file %s", debug_dump_filepath) + + if mlc_server is not None: + with mlc_server: + _main() + else: + _main() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser("MLC LLM benchmark") + + parser.add_argument( + "--dataset", + type=str, + choices=SUPPORTED_DATASET, + help=f"The benchmark dataset kind. Supporting {SUPPORTED_DATASET}", + ) + parser.add_argument( + "--dataset-path", + type=str, + help="The dataset file path.", + ) + parser.add_argument( + "--api-endpoint", + type=str, + choices=SUPPORTED_BACKENDS, + default="openai", + help="The API endpoint API for benchmarking.", + ) + parser.add_argument( + "--tokenizer", + type=str, + required=True, + help="The path of the tokenizer directory.", + ) + parser.add_argument( + "--num-gpus", + type=int, + required=True, + help="The number of GPUs used by the server. " + "We need this to better analyze the throughput per GPU.", + ) + parser.add_argument( + "--num-requests", + type=int, + required=True, + help="The number of requests for benchmark.", + ) + parser.add_argument( + "--num-warmup-requests", + type=int, + help="The number of requests for warmup. " + "It is optional when fixing the number of concurrent requests, and is required otherwise.", + ) + parser.add_argument( + "--per-gpu-workload", + default=False, + action="store_true", + help='When set to True, the specified "num_concurrent_requests"/"request_rate" ' + "denote the workload **per GPU**, which means that the real values of " + '"num_concurrent_requests"/"request_rate" used in benchmark' + 'will be multiplied by "num_gpus".', + ) + parser.add_argument( + "--num-concurrent-requests", + type=_parse_num_concurrent_requests, + help="The number(s) of concurrent requests to benchmark. " + 'It can be either one integer or a list of integer separated by commas(","). ' + "When specified, for each integer, the benchmark keeps these many consistent " + "number of concurrently running requests.", + ) + parser.add_argument( + "--request-rate", + type=_parse_request_rate, + help="The request rate(s) denoting the number of new requests each second. " + 'It can be either one float number (or "inf") or a list of numbers separated ' + 'by commas(","). ' + "When specified, the benchmark sends these many new requests each second. " + 'If it is "inf", all requests will be sent together at once.', + ) + parser.add_argument( + "--replay-timestamp-scale", + type=float, + help="The timestamp scale when replaying the timestamps in a dataset. " + 'The dataset replay mode is enabled when neither "--num-concurrent-requests" and ' + '"--request-rate" is specified. ' + "The scale is 1 by default in the replay mode.", + ) + parser.add_argument( + "--input-len", + type=int, + help="The benchmark request average input length. Default to None, " + "which means the request input length depends on the dataset being used.", + ) + parser.add_argument( + "--input-len-std", + type=float, + default=0, + help="The benchmark request input length standard deviation. Default to 0.", + ) + parser.add_argument( + "--output-len", + type=int, + help="The benchmark request average output length. Default to None, " + "which means the request output length depends on the dataset being used.", + ) + parser.add_argument( + "--output-len-std", + type=float, + default=0, + help="The benchmark request output length standard deviation. Default to 0.", + ) + parser.add_argument( + "--stream", + type=bool, + default=True, + help="Whether to benchmark stream responses. " + "When not enabled, metrics such as time-to-first-token (TTFT) will not be available. " + "Default to True.", + ) + parser.add_argument( + # NOTE: The current implementation of server metrics still has some issues that need fixes, + # which makes it not work to include server metrics. + "--include-server-metrics", + action="store_true", + help="Whether to also benchmark the server side request metrics. " + "This option is only available when benchmarking MLC server.", + ) + parser.add_argument( + "--host", + type=str, + required=True, + help="The host address of the backend API.", + ) + parser.add_argument( + "--port", + type=int, + required=True, + help="The port of the backend API.", + ) + parser.add_argument( + "--timeout", + type=float, + default=3 * 60 * 60, + help="The timeout limit of each request.", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="The random number seed. Default to 0.", + ) + parser.add_argument( + "--temperature", + type=float, + default=1.0, + help="The temperature value for logit adjustment. Default to 1.", + ) + parser.add_argument( + "--top-p", + type=float, + default=1.0, + help="The top-p value for sampling. Default to 1.", + ) + parser.add_argument( + "--ignore-eos", + default=False, + action="store_true", + help='Whether to set the "ignore_eos" field.', + ) + parser.add_argument( + "--apply-chat-template", + default=False, + action="store_true", + help="Whether to apply chat template to the request input text. " + 'It is not supported when "--input-len" is specified.', + ) + parser.add_argument( + "--num-process-workers", + type=int, + help="The number of parallel process workers to send the requests.", + ) + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Whether to disable showing progress bar with tqdm during benchmarking.", + ) + parser.add_argument( + "--max-schedule-gap", + type=float, + default=0.5, + help="The maximum allowed delay between the scheduled time in seconds.", + ) + parser.add_argument( + "--mlc-model-lib", + type=str, + help="The model lib path when benchmarking MLC serve. " + "When specified, the server is automatic launched and no external server launch is needed.", + ) + parser.add_argument( + "--mlc-engine-config", + type=_parse_mlc_engine_config, + help="The engine config used when launch MLC server.", + ) + parser.add_argument( + "--cuda-profile", + default=False, + action="store_true", + help="Whether to enable cuda profile on server. " + "The --mlc-model-lib path should be provided when enabling this option.", + ) + parser.add_argument( + "--debug-dump", + default=False, + action="store_true", + help="Whether to dump all request record raw data to file.", + ) + parser.add_argument( + "--multi-round", + default=False, + action="store_true", + help="Whether to chat like multi round conversion with history log each request. " + "Only enabled when benchmarked with fixed concurrent request mode." + "The --num-concurrent-requests should be provided when enabling this option.", + ) + parser.add_argument( + "--output", + "-o", + type=str, + default="mlc_benchmark.csv", + help="The path of the output file where to dump the benchmark results.", + ) + + main(parser.parse_args()) diff --git a/python/mlc_llm/bench/api_endpoint.py b/python/mlc_llm/bench/api_endpoint.py new file mode 100644 index 0000000..d021ddb --- /dev/null +++ b/python/mlc_llm/bench/api_endpoint.py @@ -0,0 +1,452 @@ +"""MLC LLM bench backends""" + +import argparse +import json +import os +import time +import traceback +from typing import Optional + +from typing_extensions import Self + +from mlc_llm.bench.request_record import Metrics, RequestRecord, ServerMetrics +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +class APIEndPoint: + """Manages the sending of requests to a specified API endpoint and gathers + inference statistics. + """ + + def __init__(self, include_server_metrics: bool = False) -> None: + self.include_server_metrics = include_server_metrics + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, exc_type, exc_value, tb) -> None: + pass + + async def __call__(self, request: RequestRecord) -> RequestRecord: + raise NotImplementedError() + + +class OpenAIChatEndPoint(APIEndPoint): + """The backend of sending HTTP requests in OpenAI API through "v1/chat/completions".""" + + def __init__( + self, + host: str, + port: int, + timeout: Optional[float] = None, + include_server_metrics: bool = False, + ) -> None: + super().__init__(include_server_metrics=include_server_metrics) + + import aiohttp + + self.timeout = timeout + self.client: aiohttp.ClientSession = None + self.url = f"http://{host}:{port}/v1/chat/completions" + self.headers = {"Content-Type": "application/json"} + if os.getenv("MLC_LLM_API_KEY"): + self.headers["Authorization"] = f"Bearer {os.getenv('MLC_LLM_API_KEY')}" + + async def __aenter__(self) -> Self: + import aiohttp + + self.client = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(self.timeout)) + return self + + async def __aexit__(self, exc_type, exc_value, tb) -> None: + await self.client.close() + + async def __call__(self, request_record: RequestRecord) -> RequestRecord: + payload = request_record.chat_cmpl.model_dump() + if self.timeout is not None and "timeout" not in payload: + payload["timeout"] = self.timeout + if self.include_server_metrics: + if "stream_options" not in payload or payload["stream_options"] is None: + payload["stream_options"] = {"include_usage": True} + else: + payload["stream_options"]["include_usage"] = True + if ( + request_record.chat_cmpl.debug_config is not None + and request_record.chat_cmpl.debug_config.ignore_eos + ): + payload["ignore_eos"] = True + + generated_text = "" + first_chunk_output_str = "" + time_to_first_token_s = None + start_time = time.monotonic() + server_metrics = None + + try: + async with self.client.post(self.url, json=payload, headers=self.headers) as response: + assert response.status == 200, await response.text() + if payload["stream"]: + async for chunk in response.content: + chunk = chunk.strip() + if not chunk or chunk == b"\n": + continue + # Get rid of the prefix "data: " and suffix "\n" + raw_data = chunk[6:].strip() + if raw_data == b"[DONE]": + continue + data = json.loads(raw_data) + if not data["choices"]: + continue + delta = data["choices"][0]["delta"] + content = delta.get("content", None) + if content is not None and not time_to_first_token_s: + time_to_first_token_s = time.monotonic() - start_time + first_chunk_output_str = content + if self.include_server_metrics and data["usage"] is not None: + # fmt: off + server_metrics = ServerMetrics( + input_tokens=data["usage"]["extra"]["prompt_tokens"], + prefill_tokens=data["usage"]["extra"]["prefill_tokens"], + output_tokens=data["usage"]["extra"]["completion_tokens"], + end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"], + prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"], + inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"], + time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"], # noqa: E501 + time_to_first_token_s=data["usage"]["extra"]["ttft_s"], + ) + # fmt: on + + if content is not None: + generated_text += content + else: + data = await response.json() + generated_text = data["choices"][0]["message"]["content"] + if self.include_server_metrics and data["usage"] is not None: + # fmt: off + server_metrics = ServerMetrics( + input_tokens=data["usage"]["extra"]["prompt_tokens"], + prefill_tokens=data["usage"]["extra"]["prefill_tokens"], + output_tokens=data["usage"]["extra"]["completion_tokens"], + end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"], + prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"], + inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"], + time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"], # noqa: E501 + time_to_first_token_s=data["usage"]["extra"]["ttft_s"], + ) + # fmt: on + except Exception: + error_msg = "API endpoint errored when sending request: " + traceback.format_exc() + logger.info(error_msg) + finish_time = time.monotonic() + request_record.output_str = generated_text + request_record.first_chunk_output_str = first_chunk_output_str + request_record.metrics = Metrics( + success=False, + start_time=start_time, + finish_time=finish_time, + end_to_end_latency_s=finish_time - start_time, + input_tokens=request_record.metrics.input_tokens, + time_to_first_token_s=time_to_first_token_s, + server_metrics=server_metrics, + exec_feature=request_record.metrics.exec_feature, + ) + request_record.error_msg = error_msg + return request_record + + finish_time = time.monotonic() + request_record.output_str = generated_text + request_record.first_chunk_output_str = first_chunk_output_str + success = True + error_msg = None + if len(generated_text) == 0: + success = False + error_msg = "Empty generated text." + request_record.metrics = Metrics( + success=success, + start_time=start_time, + finish_time=finish_time, + end_to_end_latency_s=finish_time - start_time, + input_tokens=request_record.metrics.input_tokens, + time_to_first_token_s=time_to_first_token_s, + server_metrics=server_metrics, + exec_feature=request_record.metrics.exec_feature, + ) + request_record.error_msg = error_msg + return request_record + + +class OpenAIEndPoint(APIEndPoint): + """The backend of sending HTTP requests in OpenAI API through "v1/completions".""" + + def __init__( + self, + host: str, + port: int, + timeout: Optional[float] = None, + include_server_metrics: bool = False, + no_debug_config: bool = False, + ) -> None: + super().__init__(include_server_metrics=include_server_metrics) + + import aiohttp + + self.timeout = timeout + self.client: aiohttp.ClientSession = None + self.url = f"http://{host}:{port}/v1/completions" + self.headers = {"Content-Type": "application/json"} + if os.getenv("MLC_LLM_API_KEY"): + self.headers["Authorization"] = f"Bearer {os.getenv('MLC_LLM_API_KEY')}" + assert not include_server_metrics, ( + '"include_server_metrics" only works for "openai-chat" endpoint for now' + ) + self.no_debug_config = no_debug_config + + async def __aenter__(self) -> Self: + import aiohttp + + self.client = aiohttp.ClientSession() + return self + + async def __aexit__(self, exc_type, exc_value, tb) -> None: + await self.client.close() + + async def __call__(self, request_record: RequestRecord) -> RequestRecord: + assert len(request_record.chat_cmpl.messages) == 1, ( + 'Endpoint "openai" does not support system prompt and multi-round conversation.' + ) + assert isinstance(request_record.chat_cmpl.messages[0].content, str) + payload = { + "model": request_record.chat_cmpl.model, + "prompt": request_record.chat_cmpl.messages[0].content, + "temperature": request_record.chat_cmpl.temperature, + "top_p": request_record.chat_cmpl.top_p, + "max_tokens": request_record.chat_cmpl.max_tokens, + "stream": True, + } + if self.timeout is not None and "timeout" not in payload: + payload["timeout"] = self.timeout + if ( + request_record.chat_cmpl.debug_config is not None + and request_record.chat_cmpl.debug_config.ignore_eos + ): + payload["ignore_eos"] = True + if not self.no_debug_config: + payload["debug_config"] = {"ignore_eos": True} + + generated_text = "" + first_chunk_output_str = "" + time_to_first_token_s = None + start_time = time.monotonic() + + try: + async with self.client.post( + self.url, json=payload, headers=self.headers, timeout=3600 + ) as response: + assert response.status == 200, await response.text() + if payload["stream"]: + async for chunk in response.content: + chunk = chunk.strip() + if not chunk or chunk == b"\n": + continue + # Get rid of the prefix "data: " and suffix "\n" + raw_data = chunk[6:].strip() + if raw_data == b"[DONE]": + continue + data = json.loads(raw_data) + if not data["choices"]: + continue + content = data["choices"][0]["text"] + if content is not None and not time_to_first_token_s: + time_to_first_token_s = time.monotonic() - start_time + first_chunk_output_str = content + if content is not None: + generated_text += content + else: + data = await response.json() + generated_text = data["choices"][0]["message"]["content"] + except Exception: + error_msg = "API endpoint errored when sending request: " + traceback.format_exc() + logger.info(error_msg) + finish_time = time.monotonic() + request_record.output_str = generated_text + request_record.first_chunk_output_str = first_chunk_output_str + request_record.metrics = Metrics( + success=False, + start_time=start_time, + finish_time=finish_time, + end_to_end_latency_s=finish_time - start_time, + input_tokens=request_record.metrics.input_tokens, + time_to_first_token_s=time_to_first_token_s, + server_metrics=None, + exec_feature=request_record.metrics.exec_feature, + ) + request_record.error_msg = error_msg + return request_record + + finish_time = time.monotonic() + request_record.output_str = generated_text + request_record.first_chunk_output_str = first_chunk_output_str + success = True + error_msg = None + if len(generated_text) == 0: + success = False + error_msg = "Empty generated text." + request_record.metrics = Metrics( + success=success, + start_time=start_time, + finish_time=finish_time, + end_to_end_latency_s=finish_time - start_time, + input_tokens=request_record.metrics.input_tokens, + time_to_first_token_s=time_to_first_token_s, + server_metrics=None, + exec_feature=request_record.metrics.exec_feature, + ) + request_record.error_msg = error_msg + return request_record + + +class TensorRTLLMEndPoint(APIEndPoint): + """The backend of sending HTTP requests in TensorRT-LLM API.""" + + def __init__(self, host: str, port: int, timeout: Optional[float] = None) -> None: + super().__init__(include_server_metrics=False) + + import aiohttp + + self.timeout = timeout + self.client: aiohttp.ClientSession = None + self.url_stream = f"http://{host}:{port}/v2/models/ensemble/generate_stream" + self.url_no_stream = f"http://{host}:{port}/v2/models/ensemble/generate" + + async def __aenter__(self) -> Self: + import aiohttp + + self.client = aiohttp.ClientSession() + return self + + async def __aexit__(self, exc_type, exc_value, tb) -> None: + await self.client.close() + + async def __call__(self, request_record: RequestRecord) -> RequestRecord: + assert len(request_record.chat_cmpl.messages) == 1 + assert isinstance(request_record.chat_cmpl.messages[0].content, str) + payload = { + "accumulate_tokens": True, + "text_input": request_record.chat_cmpl.messages[0].content, + "temperature": ( + max(request_record.chat_cmpl.temperature, 1e-5) + if request_record.chat_cmpl.temperature + else 1 + ), + "top_p": request_record.chat_cmpl.top_p if request_record.chat_cmpl.top_p else 1, + "max_tokens": request_record.chat_cmpl.max_tokens, + "stream": request_record.chat_cmpl.stream, + } + if ( + request_record.chat_cmpl.debug_config is not None + and request_record.chat_cmpl.debug_config.ignore_eos + ): + payload["min_length"] = payload["max_tokens"] + if self.timeout is not None and "timeout" not in payload: + payload["timeout"] = self.timeout + + generated_text = "" + first_chunk_output_str = "" + url = self.url_stream if request_record.chat_cmpl.stream else self.url_no_stream + time_to_first_token_s = None + start_time = time.monotonic() + + try: + async with self.client.post(url, json=payload) as response: + assert response.status == 200, await response.text() + if payload["stream"]: + async for chunk in response.content: + chunk = chunk.strip() + if not chunk or chunk == b"\n": + continue + # Get rid of the prefix "data:" and suffix "\n" + raw_data = chunk[5:].strip() + data = json.loads(raw_data) + delta = data["text_output"] + if delta is None: + continue + + if not time_to_first_token_s: + time_to_first_token_s = time.monotonic() - start_time + first_chunk_output_str = delta + generated_text += delta + else: + data = await response.json() + generated_text = data["text_output"] + except Exception: + error_msg = "API endpoint errored when sending request: " + traceback.format_exc() + logger.info(error_msg) + finish_time = time.monotonic() + request_record.output_str = generated_text + request_record.first_chunk_output_str = first_chunk_output_str + request_record.metrics = Metrics( + success=False, + start_time=start_time, + finish_time=finish_time, + end_to_end_latency_s=finish_time - start_time, + input_tokens=request_record.metrics.input_tokens, + time_to_first_token_s=time_to_first_token_s, + exec_feature=request_record.metrics.exec_feature, + ) + request_record.error_msg = error_msg + return request_record + + finish_time = time.monotonic() + request_record.output_str = generated_text + request_record.first_chunk_output_str = first_chunk_output_str + success = True + error_msg = None + if len(generated_text) == 0: + success = False + error_msg = "Empty generated text." + request_record.metrics = Metrics( + success=success, + start_time=start_time, + finish_time=finish_time, + end_to_end_latency_s=finish_time - start_time, + input_tokens=request_record.metrics.input_tokens, + time_to_first_token_s=time_to_first_token_s, + exec_feature=request_record.metrics.exec_feature, + ) + request_record.error_msg = error_msg + return request_record + + +# Todo: APIEndPoint with AsyncOpenAI Python interface +# class OpenAIPythonEndPoint(APIEndPoint): +# pass + +SUPPORTED_BACKENDS = [ + "openai", + "openai-chat", + "mlc", + "sglang", + "tensorrt-llm", + "vllm", +] + + +def create_api_endpoint(args: argparse.Namespace) -> APIEndPoint: + """Create an API endpoint instance with regard to the specified endpoint kind.""" + if args.api_endpoint in ["openai", "mlc", "sglang"]: + return OpenAIEndPoint(args.host, args.port, args.timeout, args.include_server_metrics) + if args.api_endpoint == "vllm": + return OpenAIEndPoint( + args.host, + args.port, + args.timeout, + include_server_metrics=False, + no_debug_config=True, + ) + if args.api_endpoint == "openai-chat": + return OpenAIChatEndPoint(args.host, args.port, args.timeout, args.include_server_metrics) + if args.api_endpoint == "tensorrt-llm": + return TensorRTLLMEndPoint(args.host, args.port, args.timeout) + raise ValueError(f'Unrecognized endpoint "{args.api_endpoint}"') diff --git a/python/mlc_llm/bench/dataset.py b/python/mlc_llm/bench/dataset.py new file mode 100644 index 0000000..c9d4774 --- /dev/null +++ b/python/mlc_llm/bench/dataset.py @@ -0,0 +1,878 @@ +"""MLC LLM benchmark dataset classes""" + +import argparse +import json +import random +from datetime import datetime +from typing import ClassVar, Dict, List, Optional, Tuple # noqa: UP035 + +import numpy as np +import pandas as pd +from datasets import load_dataset +from transformers import AutoTokenizer + +from mlc_llm.bench.request_record import GroupedRequestRecord, Metrics, RequestRecord +from mlc_llm.protocol.openai_api_protocol import ( + ChatCompletionMessage, + ChatCompletionRequest, + DebugConfig, +) + + +class Dataset: + """The dataset base class.""" + + # We set a truncation limit of 100k. + truncate_length = int(1e5) + # For some that datasets (e.g., dataset that has shared common prefix), + # we need fake warmup requests to avoid prefilling common prefixes to the engine. + require_fake_warmup: bool = False + # Whether the dataset contains timestamps already. + # If the dataset comes with timestamps, the benchmark can just replay + # the requests according to their timestamps. + timestamp_available: bool = False + + def generate_request_records( + self, + input_len: Optional[int], + output_len: Optional[int], + input_len_std: float = 0.0, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + """Get the raw unprocessed request records of the dataset.""" + raise NotImplementedError() + + +class ShareGPTDataset(Dataset): + """The dataset class for ShareGPT dataset.""" + + _tokenized_dataset: List[Tuple[str, List[int], int]] # noqa: UP006 + apply_chat_template: bool + + def __init__( + self, dataset_path: str, tokenizer: AutoTokenizer, apply_chat_template: bool + ) -> None: + self.apply_chat_template = apply_chat_template + with open(dataset_path, encoding="utf-8") as f: + raw_dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + _dataset = [ + (data["conversations"][0]["value"], data["conversations"][1]["value"]) + for data in raw_dataset + if len(data["conversations"]) >= 2 and data["conversations"][0]["from"] == "human" + ] + # Tokenize the prompts and completions. + self.tokenizer = tokenizer + prompts = [prompt for prompt, _ in _dataset] + if apply_chat_template: + assert getattr(tokenizer, "chat_template", None) is not None, ( + '"--apply-chat-template" is set but the tokenizer does not have chat template.' + ) + prompts = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + for prompt in prompts + ] + + prompt_token_ids = list( + tokenizer( + prompts, + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + ) + completions = [completion for _, completion in _dataset] + completion_token_ids = tokenizer( + completions, + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + self._tokenized_dataset: List[Tuple[str, List[int], int]] = [] # noqa: UP006 + for i in range(len(_dataset)): + if ( + len(prompt_token_ids[i]) < 4 + or len(completion_token_ids[i]) < 4 + or len(prompt_token_ids[i]) + len(completion_token_ids[i]) + >= min(tokenizer.model_max_length, 8192) + ): + # Filter out sequences that are too short or too long + continue + self._tokenized_dataset.append( + (prompts[i], prompt_token_ids[i], len(completion_token_ids[i])) + ) + + def generate_request_records( + self, + input_len: Optional[int], + output_len: Optional[int], + input_len_std: float = 0.0, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + if self.apply_chat_template: + assert input_len is None, ( + '"--apply-chat-template" is not supported when "--input-len" is specified.' + ) + + request_records = [] + for prompt, input_token_ids, output_length in self._tokenized_dataset: + input_length = len(input_token_ids) + # If the request does not have enough length, discard it. + if input_len is not None and input_length < input_len + 4 * input_len_std: + continue + + if input_len is not None: + input_length = round( + float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0]) + ) + input_token_ids = input_token_ids[:input_length] + input_truncated = True + else: + input_truncated = False + if output_len is not None: + output_length = round( + float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0]) + ) + elif output_length <= 1: + continue + request_records.append( + RequestRecord( + chat_cmpl=ChatCompletionRequest( + messages=[ + { + "role": "user", + "content": ( + self.tokenizer.decode(input_token_ids) + if input_truncated + else prompt + ), + } + ], + model="", + max_tokens=output_length, + ), + metrics=Metrics( + success=False, + start_time=0, + finish_time=0, + end_to_end_latency_s=0, + input_tokens=len(input_token_ids), + ), + ) + ) + return request_records + + +class LoogleDataset(Dataset): + """The dataset class for Loogle dataset.""" + + task2prompt: ClassVar[Dict[str, str]] = { # noqa: UP006 + "shortdep_qa": "Please answer the question based on the long texts below. \n{input}\nQuestion: {Q}\nAnswer: ", # noqa: E501 + "longdep_qa": "Please answer the question based on the long texts below. \n{input}\nQuestion: {Q}\nAnswer: ", # noqa: E501 + "longdep_summarization": "Please generate a summary of the below paper. \n{input}\n Summarization: ", # noqa: E501 + "shortdep_cloze": "Please fill in the clozes based on the given long texts below. Each of the placeholder '' in the question could be an entity of Person, Location or Organiocation. The same masks represent the same entity. Output a json format answer, for example: {{'': 'Bob', '': 'Gorrosion Magazine','': 'Bethel Horizon'}}\n{input}\n Question: {Q} What are the masked entities? \nAnswer:", # noqa: E501 + } + require_fake_warmup: bool = True + + def __init__(self, tokenizer: AutoTokenizer, testset_name: str) -> None: + raw_dataset = load_dataset("bigainlco/LooGLE", testset_name, split="test") + self.tokenizer = tokenizer + self.dataset = [] + self.prompt_format = self.task2prompt[testset_name] + prompts = [] + generate_lens = [] + questions = [] + for data in raw_dataset: + prompt = data["input"] + prompts.append(prompt) + qa_pairs = eval(data["qa_pairs"]) + questions.append([j["Q"] for j in qa_pairs]) + generate_lens.append( + [len(tokenizer.encode(j["A"], add_special_tokens=False)) for j in qa_pairs] + ) + prompt_token_ids = tokenizer( + prompts, + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + for prompt, prompt_token_id, question, generate_len in zip( + prompts, prompt_token_ids, questions, generate_lens + ): + self.dataset.append((prompt, prompt_token_id, question, generate_len)) + + def generate_request_records( + self, + input_len: Optional[int], + output_len: Optional[int], + input_len_std: float = 0.0, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + request_records = [] + for prompt, input_token_ids, questions, generate_lens in self.dataset: + input_length = round(float(np.random.normal(loc=input_len, scale=input_len_std))) + if len(input_token_ids) > input_length: + input_token_ids = input_token_ids[:input_length] + prompt = self.tokenizer.decode(input_token_ids) + grouped_request_records = [] + for question, generate_len in zip(questions, generate_lens): + json_obj = {"input": prompt, "Q": question} + full_prompt = self.prompt_format.format(**json_obj) + + output_length = ( + round(float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0])) + if output_len is not None + else generate_len + ) + grouped_request_records.append( + RequestRecord( + chat_cmpl=ChatCompletionRequest( + messages=[ + { + "role": "user", + "content": full_prompt, + } + ], + model="", + max_tokens=output_length, + ), + metrics=Metrics( + success=False, + start_time=0, + finish_time=0, + end_to_end_latency_s=0, + input_tokens=len(input_token_ids), + ), + ) + ) + request_records.append( + GroupedRequestRecord( + # Create a dummy ChatCompletionRequest. + chat_cmpl=ChatCompletionRequest(messages=[]), + records=grouped_request_records, + ) + ) + return request_records + + +class LLMPerfDataset(Dataset): + """The dataset class for LLMPerf dataset.""" + + def __init__(self, dataset_path: str, num_requests: int, tokenizer: AutoTokenizer) -> None: + self.tokenizer = tokenizer + self.num_requests = num_requests + + with open(dataset_path, encoding="utf-8") as f: + untokenized_data = f.readlines() + # Tokenize the prompts and completions. + tokenized_data = tokenizer( + untokenized_data, + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + tokenized_data_lengths = [len(tokens) for tokens in tokenized_data] + self.dataset: List[Tuple[str, List[int], int]] = list( # noqa: UP006 + zip(untokenized_data, tokenized_data, tokenized_data_lengths) + ) + + def generate_request_records( + self, + input_len: Optional[int] = None, + output_len: Optional[int] = None, + input_len_std: float = 250, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + if input_len is None or input_len < 40: + input_len = 550 + if output_len is None: + output_len = 150 + + request_records = [] + for _ in range(self.num_requests): + input_length = round(float(np.random.normal(loc=input_len, scale=input_len_std))) + output_length = round(float(np.random.normal(loc=output_len, scale=output_len_std))) + + prompt = ( + "Randomly stream lines from the following text " + f"with {output_length} output tokens. " + "Don't generate eos tokens:\n\n" + ) + + remaining_token_length = input_length - len( + self.tokenizer.encode(prompt, add_special_tokens=False) + ) + + random.shuffle(self.dataset) + + while remaining_token_length > 0: + for text, tokens, token_length in self.dataset: + if remaining_token_length < token_length: + prompt += self.tokenizer.decode(tokens[:remaining_token_length]) + else: + prompt += text + + remaining_token_length -= token_length + if remaining_token_length < 0: + break + + request_records.append( + RequestRecord( + chat_cmpl=ChatCompletionRequest( + messages=[{"role": "user", "content": prompt}], + model="", + max_tokens=output_length, + debug_config=DebugConfig(ignore_eos=True), + ), + metrics=Metrics( + success=False, + start_time=0, + finish_time=0, + end_to_end_latency_s=0, + input_tokens=input_length, + ), + ) + ) + return request_records + + +class JSONModeEvalDataset(Dataset): + """The dataset class for JSON dataset.""" + + def __init__(self, tokenizer: AutoTokenizer) -> None: + raw_dataset = load_dataset("NousResearch/json-mode-eval") + self.tokenizer = tokenizer + self.dataset = [] + for data in raw_dataset["train"]: + messages = data["prompt"] + schema = { + "type": "json_object", + "schema": data["schema"], + } + num_tokens = 0 + for message in messages: + num_tokens += len( + self.tokenizer.encode(message["content"], add_special_tokens=False) + ) + self.dataset.append((messages, schema, num_tokens)) + + def generate_request_records( + self, + input_len: Optional[int], + output_len: Optional[int], + input_len_std: float = 0.0, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + request_records = [] + for messages, schema, num_tokens in self.dataset: + # If the request does not have enough length, discard it. + if input_len is not None and num_tokens < input_len + 4 * input_len_std: + continue + + if output_len is not None: + output_length = max( + round(np.random.normal(loc=output_len, scale=output_len_std)), 1 + ) + else: + output_length = None + request_records.append( + RequestRecord( + chat_cmpl=ChatCompletionRequest( + messages=[ + ChatCompletionMessage(content=message["content"], role=message["role"]) + for message in messages + ], + model="", + max_tokens=output_length, + response_format=schema, + ), + metrics=Metrics( + success=False, + start_time=0, + finish_time=0, + end_to_end_latency_s=0, + input_tokens=num_tokens, + ), + ) + ) + return request_records + + +class ReActDataset(Dataset): + """The dataset class for replaying a given ReAct trace for benchmark purpose. + It is not an actual ReAct agent implementation. + """ + + _dataset: List[List[Tuple[str, int, int]]] # noqa: UP006 + require_fake_warmup: bool = True + prefix: str = """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: +(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. +(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage. +(3) Finish[answer], which returns the answer and finishes the task. +Here are some examples. +Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into? +Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area. +Action 1: Search[Colorado orogeny] +Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas. +Thought 2: It does not mention the eastern sector. So I need to look up eastern sector. +Action 2: Lookup[eastern sector] +Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny. +Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range. +Action 3: Search[High Plains] +Observation 3: High Plains refers to one of two distinct land regions: +Thought 4: I need to instead search High Plains (United States). +Action 4: Search[High Plains (United States)] +Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3] +Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft. +Action 5: Finish[1,800 to 7,000 ft] +Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who? +Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after. +Action 1: Search[Milhouse] +Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening. +Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after". +Action 2: Lookup[named after] +Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. +Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon. +Action 3: Finish[Richard Nixon] +Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture? +Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups. +Action 1: Search[Adam Clayton Powell] +Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito']. +Thought 2: To find the documentary, I can search Adam Clayton Powell (film). +Action 2: Search[Adam Clayton Powell (film)] +Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg. +The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience. +Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture. +Action 3: Finish[The Saimaa Gesture] +Question: What profession does Nicholas Ray and Elia Kazan have in common? +Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common. +Action 1: Search[Nicholas Ray] +Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause. +Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions. +Action 2: Search[Elia Kazan] +Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor. +Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor. +Action 3: Finish[director, screenwriter, actor] +Question: Which magazine was started first Arthur's Magazine or First for Women? +Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first. +Action 1: Search[Arthur's Magazine] +Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. +Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next. +Action 2: Search[First for Women] +Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. +Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first. +Action 3: Finish[Arthur's Magazine] +Question: Were Pavel Urysohn and Leonid Levin known for the same type of work? +Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same. +Action 1: Search[Pavel Urysohn] +Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory. +Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work. +Action 2: Search[Leonid Levin] +Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. +Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. +Action 3: Finish[yes] +""" # noqa: E501, RUF001 + + def __init__(self, dataset_path: str, tokenizer: AutoTokenizer) -> None: + raw_entries: List[Dict] = [] # noqa: UP006 + with open(dataset_path) as fin: + for line in fin: + line_content = json.loads(line) + raw_entries += list({"question": k, "triplets": v} for k, v in line_content.items()) + + self._dataset = [] + max_rounds = 0 + for raw_entry in raw_entries: + processed_entry = [] + question = raw_entry["question"] + triplets = raw_entry["triplets"] + seq = self.prefix + question + max_rounds = max(max_rounds, len(triplets) + 1) + output_lengths: List[int] = [] # noqa: UP006 + for i, triplet in enumerate(triplets): + output_lengths.append( + len( + tokenizer( + triplet["thought"] + + "\nAction " + + str(i + 1) + + ": " + + triplet["action"] + + "\n", + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + ) + ) + + for i in range(1, len(triplets) + 2): + seq += "Thought " + str(i) + ":" + input_len = len( + tokenizer( + seq, + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + ) + output_length = ( + output_lengths[i - 1] + if i <= len(triplets) + else int(sum(output_lengths) / len(triplets)) + ) + processed_entry.append((seq, input_len, output_length)) + if i != len(triplets) + 1: + seq += ( + triplets[i - 1]["thought"] + + "\nAction " + + str(i) + + ": " + + triplets[i - 1]["action"] + + "\nObservation " + + str(i) + + ": " + + triplets[i - 1]["observation"] + + "\n" + ) + self._dataset.append(processed_entry) + + def generate_request_records( + self, + input_len: Optional[int], + output_len: Optional[int], + input_len_std: float = 0.0, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + if input_len is not None or output_len is not None: + raise ValueError("ReAct dataset does not support specifying input/output length.") + + request_records = [] + for processed_entries in self._dataset: + grouped_request_records = [] + for prompt, input_length, output_length in processed_entries: + grouped_request_records.append( + RequestRecord( + chat_cmpl=ChatCompletionRequest( + messages=[{"role": "user", "content": prompt}], + model="", + max_tokens=output_length, + ), + metrics=Metrics( + success=False, + start_time=0, + finish_time=0, + end_to_end_latency_s=0, + input_tokens=input_length, + ), + ) + ) + request_records.append( + GroupedRequestRecord( + # Create a dummy ChatCompletionRequest. + chat_cmpl=ChatCompletionRequest(messages=[]), + records=grouped_request_records, + ) + ) + return request_records + + +class WildChatDataset(Dataset): + """The dataset class for WildChat dataset.""" + + apply_chat_template: bool + + def __init__(self, tokenizer: AutoTokenizer, apply_chat_template: bool) -> None: + raw_dataset = load_dataset("allenai/WildChat", split="train") + self.tokenizer = tokenizer + self.apply_chat_template = apply_chat_template + + # Filter out the conversations with less than 2 turns. + _dataset = [ + (entry["conversation"][0]["content"], entry["conversation"][1]["content"]) + for entry in raw_dataset + if len(entry["conversation"]) >= 2 + and entry["conversation"][0]["role"] == "user" + and entry["conversation"][1]["role"] == "assistant" + ] + + prompts = [] + completions = [] + for prompt, completion in _dataset: + prompts.append(prompt) + completions.append(completion) + if apply_chat_template: + assert getattr(tokenizer, "chat_template", None) is not None, ( + '"--apply-chat-template" is set but the tokenizer does not have chat template.' + ) + prompts = [ + tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + for prompt in prompts + ] + + prompt_token_ids = list( + tokenizer( + prompts, + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + ) + completion_token_ids = tokenizer( + completions, + truncation=True, + max_length=min(tokenizer.model_max_length, self.truncate_length), + add_special_tokens=False, + ).input_ids + self._tokenized_dataset: List[Tuple[str, List[int], int]] = [] # noqa: UP006 + for i in range(len(_dataset)): + if len(prompt_token_ids[i]) < 4 or len(completion_token_ids[i]) < 4: + # Filter out sequences that are too short + continue + self._tokenized_dataset.append( + (prompts[i], prompt_token_ids[i], len(completion_token_ids[i])) + ) + + def generate_request_records( + self, + input_len: Optional[int], + output_len: Optional[int], + input_len_std: float = 0.0, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + if self.apply_chat_template: + assert input_len is None, ( + '"--apply-chat-template" is not supported when "--input-len" is specified.' + ) + + request_records = [] + for prompt, input_token_ids, output_length in self._tokenized_dataset: + input_length = len(input_token_ids) + # If the request does not have enough length, discard it. + if input_len is not None and input_length < input_len + 4 * input_len_std: + continue + + if input_len is not None: + input_length = round( + float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0]) + ) + input_token_ids = input_token_ids[:input_length] + input_truncated = True + else: + input_truncated = False + if output_len is not None: + output_length = round( + float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0]) + ) + elif output_length <= 1: + continue + request_records.append( + RequestRecord( + chat_cmpl=ChatCompletionRequest( + messages=[ + { + "role": "user", + "content": ( + self.tokenizer.decode(input_token_ids) + if input_truncated + else prompt + ), + } + ], + model="", + max_tokens=output_length, + ), + metrics=Metrics( + success=False, + start_time=0, + finish_time=0, + end_to_end_latency_s=0, + input_tokens=len(input_token_ids), + ), + ) + ) + return request_records + + +class AzureLLMInferenceDataset(Dataset): + """The dataset class for AzureLLMInference dataset. + Reference: https://github.com/Azure/AzurePublicDataset + """ + + timestamp_available: bool = True + + def __init__(self, dataset_path: str, tokenizer: AutoTokenizer) -> None: + df = pd.read_csv(dataset_path) + self.tokenizer = tokenizer + + # Filter out the conversations with less than 2 turns. + self.dataset = [ + ( + entry["TIMESTAMP"], + min( + entry["ContextTokens"], + tokenizer.model_max_length, + self.truncate_length, + ), + min( + entry["GeneratedTokens"], + tokenizer.model_max_length, + self.truncate_length, + ), + ) + for _, entry in df.iterrows() + if entry["ContextTokens"] >= 4 and entry["GeneratedTokens"] >= 4 + ] + + def generate_request_records( + self, + input_len: Optional[int], + output_len: Optional[int], + input_len_std: float = 0.0, + output_len_std: float = 0.0, + ) -> List[RequestRecord]: # noqa: UP006 + time_fmt = "%Y-%m-%d %H:%M:%S.%f" + start_time = datetime.strptime(self.dataset[0][0][:-1], time_fmt) + request_records = [] + for timestamp, input_length, output_length in self.dataset: + # If the request does not have enough length, discard it. + if input_len is not None and input_length < input_len + 4 * input_len_std: + continue + + if input_len is not None: + input_length = round( + float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0]) + ) + if output_len is not None: + output_length = round( + float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0]) + ) + elif output_length <= 1: + continue + + prompt_token_ids = [ + random.randint(0, self.tokenizer.vocab_size - 1) for _ in range(input_length) + ] + while True: + # Adjust the token ids until the retokenization on the decoded string + # matches the required input length. + prompt = self.tokenizer.decode(prompt_token_ids) + retokenized_token_ids = self.tokenizer.encode(prompt, add_special_tokens=False) + if len(retokenized_token_ids) < input_length: + prompt_token_ids = retokenized_token_ids + [ + random.randint(0, self.tokenizer.vocab_size - 1) + for _ in range(input_length - len(retokenized_token_ids)) + ] + elif len(retokenized_token_ids) > input_length: + prompt_token_ids = retokenized_token_ids[:input_length] + else: + break + + time_diff = (datetime.strptime(timestamp[:-1], time_fmt) - start_time).total_seconds() + request_records.append( + RequestRecord( + chat_cmpl=ChatCompletionRequest( + messages=[{"role": "user", "content": prompt}], + model="", + max_tokens=output_length, + ), + timestamp=time_diff, + metrics=Metrics( + success=False, + start_time=0, + finish_time=0, + end_to_end_latency_s=0, + input_tokens=input_length, + ), + ) + ) + return request_records + + +SUPPORTED_DATASET = [ + "sharegpt", + "llmperf", + "json-mode-eval", + "loogle", + "react", + "wildchat", + "azure-llm-inference", +] + + +def create_dataset(args: argparse.Namespace, tokenizer: AutoTokenizer) -> Dataset: + """Create a dataset instance with regard to the specified dataset kind and file path.""" + if args.dataset_path is not None and not isinstance(args.dataset_path, str): + raise TypeError(f"Invalid dataset path {args.dataset_path}. Please use a string.") + if args.dataset is None and args.dataset_path is not None: + # Auto-detect the dataset kind by looking into the dataset path. + if "sharegpt" in args.dataset_path.lower(): + args.dataset = "sharegpt" + else: + raise ValueError( + f"Unable to detect the dataset kind from dataset path {args.dataset_path}. " + 'Please specify the dataset kind via "--dataset".' + ) + if args.dataset == "sharegpt": + if args.dataset_path is None: + raise ValueError( + 'ShareGPT dataset requires dataset path. Please specify it with "--dataset-path".' + ) + return ShareGPTDataset(args.dataset_path, tokenizer, args.apply_chat_template) + if args.dataset == "llmperf": + if args.dataset_path is None: + raise ValueError( + 'LLMPerf dataset requires dataset path. Please specify it with "--dataset-path".' + ) + assert args.apply_chat_template is False, ( + "LLMPerf dataset does not support applying chat template" + ) + return LLMPerfDataset( + args.dataset_path, + (args.num_requests + args.num_warmup_requests) * 4, + tokenizer, + ) + if args.dataset == "json-mode-eval": + assert args.apply_chat_template is False, ( + "JSON mode evaluation does not support applying chat template" + ) + return JSONModeEvalDataset(tokenizer) + if args.dataset == "loogle": + if args.dataset_path is None: + raise ValueError( + 'Loogle dataset requires a testset name. Please specify it with "--dataset-path".' + ) + assert args.apply_chat_template is False, ( + "Loogle dataset does not support applying chat template" + ) + return LoogleDataset(tokenizer, testset_name=args.dataset_path) + if args.dataset == "react": + if args.dataset_path is None: + raise ValueError( + 'ReAct dataset requires dataset path. Please specify it with "--dataset-path".' + ) + assert args.apply_chat_template is False, ( + "ReAct dataset does not support applying chat template" + ) + return ReActDataset(args.dataset_path, tokenizer) + if args.dataset == "wildchat": + return WildChatDataset(tokenizer, args.apply_chat_template) + if args.dataset == "azure-llm-inference": + if args.dataset_path is None: + raise ValueError( + "AzureLLMInference dataset requires dataset path. " + 'Please specify it with "--dataset-path".' + ) + assert args.apply_chat_template is False, ( + "AzureLLMInference dataset does not support applying chat template" + ) + return AzureLLMInferenceDataset(args.dataset_path, tokenizer) + raise ValueError(f"Unrecognized dataset {args.dataset}") diff --git a/python/mlc_llm/bench/evaluation/gsm8k.py b/python/mlc_llm/bench/evaluation/gsm8k.py new file mode 100644 index 0000000..ffdd468 --- /dev/null +++ b/python/mlc_llm/bench/evaluation/gsm8k.py @@ -0,0 +1,315 @@ +"""Eval GSM8K with MLCEngine.""" + +import argparse +import asyncio +import json +import random +import re +from datetime import datetime +from pathlib import Path +from typing import List, Literal, Optional # noqa: UP035 + +import tqdm + +from mlc_llm import AsyncMLCEngine + +DEVICES = ["cuda", "rocm", "metal", "vulkan"] +ANSWER_TRIGGER = "The answer is" +INVALID_ANS = "[invalid]" + + +def extract_answer(text: str, regex: re.Pattern, select_index: int) -> str: + """Extract the answer from the text.""" + match_all = regex.findall(text) + if len(match_all) == 0: + return INVALID_ANS + match = match_all[select_index] + if isinstance(match, tuple): + match = next(m for m in match if m) + match_str: str = match.strip() + match_str = match_str.lstrip("$").rstrip(".").replace(",", "") + return match_str + + +def extract_ground_truth(text: str) -> str: + """Extract the ground truth from the text.""" + return extract_answer(text, re.compile(r"#### (\-?[0-9\.\,]+)"), 0) + + +def strict_extract_answer(text: str) -> str: + """Strictly extract the answer from the text.""" + return extract_answer(text, re.compile(r"The answer is \$?(\-?[0-9\.\,]+)."), 0) + + +def flexible_extract_answer(text: str) -> str: + """Extract the last number from the text.""" + return extract_answer(text, re.compile(r"(-?[$0-9.,]{2,})|(-?[0-9]+)"), -1) + + +def create_few_shot_prompt(n_shot: int, use_cot: bool, random_order=False) -> str: + """ + Create a prompt for the few-shot learning task. + + Note + ---- + The examples are taken from the paper https://arxiv.org/pdf/2201.11903.pdf page 35. + """ + question, chain, answer = [], [], [] + + question.append( + "There are 15 trees in the grove. " + "Grove workers will plant trees in the grove today. " + "After they are done, there will be 21 trees. " + "How many trees did the grove workers plant today?" + ) + chain.append( + "There are 15 trees originally. " + "Then there were 21 trees after some more were planted. " + "So there must have been 21 - 15 = 6." + ) + answer.append("6") + + question.append( + "If there are 3 cars in the parking lot and 2 more cars arrive, " + "how many cars are in the parking lot?" + ) + chain.append("There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5.") + answer.append("5") + + question.append( + "Leah had 32 chocolates and her sister had 42. If they ate 35, " + "how many pieces do they have left in total?" + ) + chain.append( + "Originally, Leah had 32 chocolates. " + "Her sister had 42. So in total they had 32 + 42 = 74. " + "After eating 35, they had 74 - 35 = 39." + ) + answer.append("39") + + question.append( + "Jason had 20 lollipops. He gave Denny some lollipops. Now Jason " + "has 12 lollipops. How many lollipops did Jason give to Denny?" + ) + chain.append( + "Jason started with 20 lollipops. Then he had 12 after giving some " + "to Denny. So he gave Denny 20 - 12 = 8." + ) + answer.append("8") + + question.append( + "Shawn has five toys. For Christmas, he got two toys each from his " + "mom and dad. How many toys does he have now?" + ) + chain.append( + "Shawn started with 5 toys. If he got 2 toys each from his mom and " + "dad, then that is 4 more toys. 5 + 4 = 9." + ) + answer.append("9") + + question.append( + "There were nine computers in the server room. Five more computers " + "were installed each day, from monday to thursday. " + "How many computers are now in the server room?" + ) + chain.append( + "There were originally 9 computers. For each of 4 days, 5 more " + "computers were added. So 5 * 4 = 20 computers were added. " + "9 + 20 is 29." + ) + answer.append("29") + + question.append( + "Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On " + "wednesday, he lost 2 more. " + "How many golf balls did he have at the end of wednesday?" + ) + chain.append( + "Michael started with 58 golf balls. After losing 23 on tuesday, " + "he had 58 - 23 = 35. After losing 2 more, " + "he had 35 - 2 = 33 golf balls." + ) + answer.append("33") + + question.append( + "Olivia has $23. She bought five bagels for $3 each. How much money does she have left?" + ) + chain.append( + "Olivia had 23 dollars. " + "5 bagels for 3 dollars each will be 5 x 3 = 15 dollars. " + "So she has 23 - 15 dollars left. 23 - 15 is 8." + ) + answer.append("8") + + index_list = list(range(len(question))) + if random_order: + random.shuffle(index_list) + + prompt = "" + for i in index_list[:n_shot]: + if use_cot: + prompt += f"Q: {question[i]}\nA: {chain[i]} {ANSWER_TRIGGER} {answer[i]}.\n\n" + else: + prompt += f"Question: {question[i]}\nAnswer: {ANSWER_TRIGGER} {answer[i]}.\n\n" + return prompt + + +def create_prompt(question: str, n_shot: int, use_cot: bool, random_order: bool = False) -> str: + """Create a prompt for the few-shot learning task.""" + prompt = create_few_shot_prompt(n_shot, use_cot, random_order) + if use_cot: + prompt += f"Q: {question}\nA:" + else: + prompt += f"Question: {question}\nAnswer:" + return prompt + + +def parse_args(): + """Parse command line arguments.""" + + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, required=True) + parser.add_argument( + "--dataset", type=Path, required=True, help="Path to GSM8K test dataset home." + ) + parser.add_argument("--device", type=str, choices=["auto", *DEVICES], default="auto") + parser.add_argument("--model-lib", type=str, default=None) + parser.add_argument("--n-shot", type=int, default=8) + parser.add_argument("--disable_cot", action="store_true", default=False) + parser.add_argument("-bs", "--batch-size", type=int, default=16) + parser.add_argument("--log-dir", type=Path, default=None) + return parser.parse_args() + + +async def send_request( + async_engine: AsyncMLCEngine, + prompts: List[str], # noqa: UP006 + semaphore: asyncio.Semaphore, +): + """Send the calibration requests to the engine.""" + tasks = [] + + async def generate_task(prompt): + async with semaphore: + return await async_engine.completions.create( + prompt=prompt, + stream=False, + max_tokens=512, + stop=["Q:", "Question:"], + temperature=0.0, + ) + + for prompt in prompts: + task = asyncio.create_task(generate_task(prompt)) + tasks.append(task) + + return await tqdm.asyncio.tqdm.gather(*tasks) + + +async def evaluate( + model: str, + device: str, + dataset: Path, + model_lib: Optional[str], + n_shot: int, + use_cot: bool, + batch_size: int, + log_dir: Optional[Path], +): + """Evaluate GSM8K for the model.""" + mode: Literal["local", "interactive", "server"] = ( + "server" if batch_size > 4 else "interactive" if batch_size == 1 else "local" + ) + async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode=mode) + + with open(dataset / "test.jsonl", encoding="utf-8") as file: + tests = [json.loads(line) for line in file] + + prompts = [create_prompt(test["question"], n_shot, use_cot) for test in tests] + responses = await send_request(async_engine, prompts, asyncio.Semaphore(batch_size)) + assert len(responses) == len(tests) + + num_strict_correct, num_flexible_correct = 0, 0 + num_tests = len(tests) + logs = [] + + for response, test in zip(responses, tests): + response_text = response.choices[0].text.strip() + gt_answer = extract_ground_truth(test["answer"]) + assert gt_answer != INVALID_ANS + strict_answer = strict_extract_answer(response_text) + flexible_answer = flexible_extract_answer(response_text) + + if gt_answer == strict_extract_answer(response_text): + # If the answer is exactly the same as the response, then it is correct + num_strict_correct += 1 + num_flexible_correct += 1 + + elif gt_answer == flexible_extract_answer(response_text): + # Try flexible extract if the strict match fails + num_flexible_correct += 1 + + logs.append( + { + "question": test["question"], + "response": response_text, + "ground_truth": gt_answer, + "strict_answer": strict_answer, + "flexible_answer": flexible_answer, + "strict_match": gt_answer == strict_answer, + "flexible_match": gt_answer == flexible_answer, + } + ) + + results = { + "config": { + "model": model, + "device": device, + "model_lib": model_lib, + "n_shot": n_shot, + "use_cot": use_cot, + }, + "results": { + "strict_match": num_strict_correct, + "flexible_match": num_flexible_correct, + "total": num_tests, + }, + } + print( + f"Strict Matching Accuracy: {num_strict_correct} / {num_tests} = " + f"{num_strict_correct / num_tests * 100:.2f}%" + ) + print( + f"Flexible Matching Accuracy: {num_flexible_correct} / {num_tests} = " + f"{num_flexible_correct / num_tests * 100:.2f}%" + ) + + if log_dir: + with open(log_dir / "summary.json", "w", encoding="utf-8") as f: + json.dump(results, f, indent=2) + with open(log_dir / "logs.json", "w", encoding="utf-8") as f: + json.dump(logs, f, indent=2) + + +if __name__ == "__main__": + args = parse_args() + start_time = datetime.now() + log_dir: Optional[Path] = None + if args.log_dir is not None: + time_dir = start_time.strftime("%Y-%m-%d_%H-%M-%S") + log_dir = args.log_dir / time_dir + log_dir.mkdir(parents=True, exist_ok=True) + asyncio.run( + evaluate( + model=args.model, + device=args.device, + dataset=args.dataset, + model_lib=args.model_lib, + n_shot=args.n_shot, + use_cot=not args.disable_cot, + batch_size=args.batch_size, + log_dir=log_dir, + ) + ) + end_time = datetime.now() + print(f"Time used: {end_time - start_time}") diff --git a/python/mlc_llm/bench/evaluation/mmlu.py b/python/mlc_llm/bench/evaluation/mmlu.py new file mode 100644 index 0000000..bcb0b01 --- /dev/null +++ b/python/mlc_llm/bench/evaluation/mmlu.py @@ -0,0 +1,245 @@ +"""Eval MMLU with MLCEngine.""" + +import argparse +import asyncio +import csv +import json +import string +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional # noqa: UP035 + +import numpy as np +import tqdm + +from mlc_llm import AsyncMLCEngine + +SUBJECTS = [ + "abstract_algebra", + "anatomy", + "astronomy", + "business_ethics", + "clinical_knowledge", + "college_biology", + "college_chemistry", + "college_computer_science", + "college_mathematics", + "college_medicine", + "college_physics", + "computer_security", + "conceptual_physics", + "econometrics", + "electrical_engineering", + "elementary_mathematics", + "formal_logic", + "global_facts", + "high_school_biology", + "high_school_chemistry", + "high_school_computer_science", + "high_school_european_history", + "high_school_geography", + "high_school_government_and_politics", + "high_school_macroeconomics", + "high_school_mathematics", + "high_school_microeconomics", + "high_school_physics", + "high_school_psychology", + "high_school_statistics", + "high_school_us_history", + "high_school_world_history", + "human_aging", + "human_sexuality", + "international_law", + "jurisprudence", + "logical_fallacies", + "machine_learning", + "management", + "marketing", + "medical_genetics", + "miscellaneous", + "moral_disputes", + "moral_scenarios", + "nutrition", + "philosophy", + "prehistory", + "professional_accounting", + "professional_law", + "professional_medicine", + "professional_psychology", + "public_relations", + "security_studies", + "sociology", + "us_foreign_policy", + "virology", + "world_religions", +] +PADDING_LEN = max(len(subject) for subject in SUBJECTS) +DEVICES = ["cuda", "rocm", "metal", "vulkan"] +PROMPT_TEMPLATE = string.Template("$Q\nA. $A\nB. $B\nC. $C\nD. $D\nAnswer:") + + +def parse_args(): + """Parse command line arguments.""" + + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=str, required=True) + parser.add_argument( + "--dataset", type=Path, required=True, help="Path to MMLU test dataset home." + ) + parser.add_argument("--device", type=str, choices=["auto", *DEVICES], default="auto") + parser.add_argument("--model-lib", type=str, default=None) + parser.add_argument("-s", "--subject", nargs="+", type=str, choices=SUBJECTS, default=SUBJECTS) + parser.add_argument("-bs", "--batch-size", type=int, default=16) + parser.add_argument("--log-dir", type=Path, default=None) + return parser.parse_args() + + +async def send_request( + async_engine: AsyncMLCEngine, + prompts: List[str], # noqa: UP006 + semaphore: asyncio.Semaphore, + subject: str, +): + """Send the calibration requests to the engine.""" + tasks = [] + + async def generate_task(prompt): + async with semaphore: + return await async_engine.completions.create( + prompt=prompt, + stream=False, + max_tokens=1, + temperature=1.0, + logprobs=True, + top_logprobs=5, + ) + + for prompt in prompts: + task = asyncio.create_task(generate_task(prompt)) + tasks.append(task) + + return await tqdm.asyncio.tqdm.gather( + *tasks, + desc=f"Running {subject.ljust(PADDING_LEN)}", + bar_format="{desc} {percentage:3.0f}%|{bar}{r_bar}", + ) + + +async def evaluate( + model: str, + device: str, + dataset: Path, + model_lib: Optional[str], + subjects: List[str], # noqa: UP006 + semaphore: asyncio.Semaphore, + log_dir: Optional[Path], +): + """Evaluate MMLU for the model.""" + async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode="server") + + results: Dict[str, Any] = {} # noqa: UP006 + for subject in subjects: + with open(dataset / "test" / f"{subject}_test.csv", encoding="utf-8") as csvfile: + tests = list(csv.reader(csvfile, delimiter=",", quotechar='"')) + assert all(len(test) == 6 for test in tests) + + logs = [] + num_correct = 0 + prompts = [ + PROMPT_TEMPLATE.substitute(Q=test[0], A=test[1], B=test[2], C=test[3], D=test[4]) + for test in tests + ] + responses = await send_request(async_engine, prompts, semaphore, subject) + + assert len(responses) == len(tests) + for response, test in zip(responses, tests): + token_logprobs = {} + logprobs = response.choices[0].logprobs.content[0].top_logprobs + for logprob in logprobs: + if logprob.token not in token_logprobs: + token_logprobs[logprob.token] = logprob.logprob + + abcd_logprobs = {} + for choice in ["A", "B", "C", "D"]: + abcd_logprobs[choice] = token_logprobs[choice] if choice in token_logprobs else -100 + + pred = {0: "A", 1: "B", 2: "C", 3: "D"}[int(np.argmax(list(abcd_logprobs.values())))] + num_correct += pred == test[5] + + logs.append( + { + "Question": { + "Q": test[0], + "A": test[1], + "B": test[2], + "C": test[3], + "D": test[4], + }, + "Answer": test[5], + "Response": { + "pred": pred, + "logprobs": list(abcd_logprobs.values()), + }, + } + ) + + results[subject] = { + "correct": num_correct, + "total": len(tests), + "accuracy": num_correct / len(tests), + } + + if log_dir: + with open(log_dir / "subjects" / f"{subject}.json", "w", encoding="utf-8") as f: + json.dump(logs, f, indent=2) + + total_correct, total_tests = 0, 0 + for subject, v in results.items(): + num_correct, num_tests, accuracy = v["correct"], v["total"], v["accuracy"] + print(f"{subject}: {num_correct} / {num_tests} = {accuracy * 100:.2f}%") + total_correct += num_correct + total_tests += num_tests + + total_accuracy = total_correct / total_tests + results["total"] = { + "correct": total_correct, + "total": total_tests, + "accuracy": total_accuracy, + } + print(f"Total accuracy: {total_correct} / {total_tests} = {total_accuracy * 100:.2f}%") + + if log_dir: + results = { + "config": { + "model": model, + "device": device, + "model_lib": model_lib, + "subjects": subjects, + }, + "results": results, + } + with open(log_dir / "summary.json", "w", encoding="utf-8") as f: + json.dump(results, f, indent=2) + + +if __name__ == "__main__": + args = parse_args() + start_time = datetime.now() + log_dir: Optional[Path] = None + if args.log_dir is not None: + time_dir = start_time.strftime("%Y-%m-%d_%H-%M-%S") + log_dir = args.log_dir / time_dir + (log_dir / "subjects").mkdir(parents=True, exist_ok=True) + asyncio.run( + evaluate( + model=args.model, + device=args.device, + dataset=args.dataset, + model_lib=args.model_lib, + subjects=args.subject, + semaphore=asyncio.Semaphore(args.batch_size), + log_dir=log_dir, + ) + ) + end_time = datetime.now() + print(f"Time used: {end_time - start_time}") diff --git a/python/mlc_llm/bench/request_processor.py b/python/mlc_llm/bench/request_processor.py new file mode 100644 index 0000000..4e4f836 --- /dev/null +++ b/python/mlc_llm/bench/request_processor.py @@ -0,0 +1,740 @@ +"""MLC LLM Bench Request""" + +import argparse +import asyncio +import concurrent.futures +import copy +import os +import random +import time +from typing import Any, Callable, Dict, List, Optional # noqa: UP035 + +import numpy as np +import requests +from tqdm import tqdm +from transformers import AutoTokenizer + +from mlc_llm.bench.api_endpoint import APIEndPoint +from mlc_llm.bench.dataset import Dataset +from mlc_llm.bench.request_record import GroupedRequestRecord, RequestRecord +from mlc_llm.protocol.openai_api_protocol import ( + ChatCompletionMessage, + ChatCompletionRequest, + DebugConfig, +) +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +class RequestProcessor: + """The request processor base class. + Each processor can take a list of RequestRecord, applying the process, + and returning the processed RequestRecord in the end. + """ + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + raise NotImplementedError() + + +class LogMessage(RequestProcessor): + """The processor that prints the logger message.""" + + def __init__(self, message: str) -> None: + self.message = message + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + logger.info(self.message) + return request_records + + +class SampleRequests(RequestProcessor): + """The processor that samples requests out from the given request list.""" + + def __init__(self, num_requests: int, take_first_x_requests: bool = False) -> None: + self.num_requests = num_requests + # If `take_first_x_requests` is True, the first `num_requests` requests + # are returned and sampling will not happen. + self.take_first_x_requests = take_first_x_requests + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + assert len(request_records) > 0, "Empty input request record." + + # We expect the input request records to be all grouped or all plain. + if isinstance(request_records[0], GroupedRequestRecord): + assert all(isinstance(record, GroupedRequestRecord) for record in request_records) + return self._sample_from_grouped_request_records(request_records) + + assert all(not isinstance(record, GroupedRequestRecord) for record in request_records) + return self._sample_from_plain_request_records(request_records) + + def _sample_from_plain_request_records( + self, + request_records: List[RequestRecord], # noqa: UP006 + ) -> List[RequestRecord]: # noqa: UP006 + samples: List[RequestRecord] = [] # noqa: UP006 + if self.take_first_x_requests: + if len(request_records) < self.num_requests: + raise ValueError( + f"Insufficient requests. Requiring {self.num_requests} requests " + f"but only {len(request_records)} are available." + ) + samples = copy.deepcopy(list(request_records[: self.num_requests])) + else: + while len(samples) < self.num_requests: + # Create a new list so that the in-place shuffle does not mutate the input list. + records = list(request_records) + random.shuffle(records) + samples += copy.deepcopy(records) + samples = samples[: self.num_requests] + for i, record in enumerate(samples): + record.request_id = i + return samples + + def _sample_from_grouped_request_records( + self, + grouped_request_records: List[GroupedRequestRecord], # noqa: UP006 + ) -> List[RequestRecord]: # noqa: UP006 + num_total_available_requests = sum( + len(record.records) for record in grouped_request_records + ) + if self.num_requests > num_total_available_requests: + raise ValueError( + "Due to the existence of shared common prefixes, we do not allow " + "benchmarking with requests more than the available requests in the dataset. " + f"The required number of requests {self.num_requests} exceeds the " + f"number of total available requests {num_total_available_requests}." + ) + + # Create a new list so that the in-place shuffle does not mutate the input list. + records = list(grouped_request_records) + if not self.take_first_x_requests: + random.shuffle(records) + remaining = self.num_requests + samples: List[RequestRecord] = [] # noqa: UP006 + for grouped_request_record in grouped_request_records: + num_used_requests = min(len(grouped_request_record.records), remaining) + samples += grouped_request_record.records[:num_used_requests] + remaining -= num_used_requests + if remaining == 0: + break + for i, record in enumerate(samples): + record.request_id = i + return samples + + +class AttachModelName(RequestProcessor): + """The processor that attaches model name to requests.""" + + def __init__(self, model: str) -> None: + self.model = model + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + for request_record in request_records: + request_record.chat_cmpl.model = self.model + return request_records + + +class AttachRequestRateTimestamp(RequestProcessor): + """The processor that applies timestamps to the requests.""" + + def __init__(self, request_rate: np.float32) -> None: + self.request_rate = request_rate + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + timestamp = 0.0 + for request_record in request_records: + assert request_record.timestamp is None, "The request record already has a timestamp" + request_record.timestamp = timestamp + timestamp += float(np.random.exponential(1.0 / self.request_rate)) + return request_records + + +class AttachExecutionFeature(RequestProcessor): + """The processor that attaches execution features to all requests""" + + def __init__(self, exec_feature: Dict[str, Any]) -> None: # noqa: UP006 + self.exec_feature = exec_feature + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + for request_record in request_records: + assert request_record.metrics is not None + request_record.metrics.exec_feature = self.exec_feature + return request_records + + +class AttachStreamFlag(RequestProcessor): + """The processor that attaches the stream flag to the requests.""" + + def __init__(self, stream: Optional[bool]) -> None: + self.stream = stream + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + if self.stream is None: + return request_records + for request_record in request_records: + request_record.chat_cmpl.stream = self.stream + return request_records + + +class AttachSamplingOptions(RequestProcessor): + """The processor that attaches the stream flag to the requests.""" + + def __init__(self, temperature: float, top_p: float, ignore_eos: bool) -> None: + self.temperature = temperature + self.top_p = top_p + self.ignore_eos = ignore_eos + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + for request_record in request_records: + request_record.chat_cmpl.temperature = self.temperature + request_record.chat_cmpl.top_p = self.top_p + request_record.chat_cmpl.frequency_penalty = 0.0 + request_record.chat_cmpl.presence_penalty = 0.0 + request_record.chat_cmpl.tool_choice = "none" + if self.ignore_eos: + request_record.chat_cmpl.debug_config = DebugConfig(ignore_eos=True) + return request_records + + +class ScaleTimestamp(RequestProcessor): + """Scale the timestamp of requests by the given scale factor.""" + + def __init__(self, timestamp_scale: float): + self.timestamp_scale = timestamp_scale + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + for request_record in request_records: + if request_record.timestamp is None: + raise ValueError( + f"The timestamp of request {request_record} has not been initialized." + ) + request_record.timestamp *= self.timestamp_scale + return request_records + + +class MetricAnalyzer(RequestProcessor): + """The processor that analyzes the raw benchmark results and computes more detailed metrics.""" + + def __init__(self, tokenizer: AutoTokenizer) -> None: + self.tokenizer = tokenizer + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + updated_records = [] + for request_record in request_records: + metrics = request_record.metrics + if not metrics.success: + assert request_record.error_msg is not None + continue + + metrics.output_tokens = len( + self.tokenizer.encode(request_record.output_str, add_special_tokens=False) + ) + first_chunk_output_tokens = len( + self.tokenizer.encode( + request_record.first_chunk_output_str, add_special_tokens=False + ) + ) + if metrics.output_tokens <= first_chunk_output_tokens: + metrics.success = False + request_record.error_msg = ( + f"Total output token num ({metrics.output_tokens}) equals " + f'the first chunk output token. Output text "{request_record.output_str}", ' + f'first chunk output text "{request_record.first_chunk_output_str}"' + ) + continue + assert metrics.input_tokens > 0, "Invalid prompt tokens" + metrics.inter_token_latency_s = metrics.end_to_end_latency_s / metrics.output_tokens + if metrics.time_to_first_token_s is None: + metrics.time_to_first_token_s = 0 + metrics.time_per_output_token_s = ( + metrics.end_to_end_latency_s - metrics.time_to_first_token_s + ) / (metrics.output_tokens - first_chunk_output_tokens) + updated_records.append(request_record) + return updated_records + + +class WarmupAndRun(RequestProcessor): + """The processor that runs warmup first and then runs the benchmark with the given pipeline.""" + + def __init__( + self, + num_warmup_requests: int, + num_benchmark_requests: int, + pipeline: RequestProcessor, + cuda_profile_url: Optional[str], + fake_warmup: bool = False, + ) -> None: + self.num_warmup_requests = num_warmup_requests + self.num_benchmark_requests = num_benchmark_requests + self.pipeline = pipeline + self.cuda_profile_url = cuda_profile_url + self.fake_warmup = fake_warmup + + def generate_fake_warmup_requests( + self, num_warmup_requests: int, example_request: RequestRecord + ) -> List[RequestRecord]: # noqa: UP006 + records = [] + for _ in range(num_warmup_requests): + record = copy.deepcopy(example_request) + record.chat_cmpl = ChatCompletionRequest( + messages=[ + { + "role": "user", + "content": "Please output arbitrary coherent sentences. Do not output eos token.", # noqa: E501 + } + ], + model="", + max_tokens=128, + ) + records.append(record) + return records + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + # Warmup + if self.fake_warmup: + assert len(request_records) == self.num_benchmark_requests + benchmark_requests = request_records + example_request = benchmark_requests[0] + warmup_requests = self.generate_fake_warmup_requests( + self.num_warmup_requests, example_request=example_request + ) + else: + assert len(request_records) == self.num_warmup_requests + self.num_benchmark_requests + benchmark_requests = request_records[: -self.num_warmup_requests] + warmup_requests = request_records[-self.num_warmup_requests :] + for request_record in warmup_requests: + request_record.timestamp = 0 if request_record.timestamp is not None else None + warmup_requests = self._process_warmup_requests(warmup_requests) + logger.info("Warmup with %d request(s)...", self.num_warmup_requests) + self.pipeline(warmup_requests) + + # Then run benchmark + if self.cuda_profile_url is not None: + cuda_profiler_start_url = self.cuda_profile_url + "/debug/cuda_profiler_start" + cuda_profiler_start_response = requests.post(cuda_profiler_start_url, timeout=60) + assert cuda_profiler_start_response.status_code == 200 + logger.info("Warmup finished. Start benchmarking...") + updated_request_records = self.pipeline(benchmark_requests) + if self.cuda_profile_url is not None: + cuda_profiler_stop_url = self.cuda_profile_url + "/debug/cuda_profiler_stop" + cuda_profiler_stop_response = requests.post(cuda_profiler_stop_url, timeout=60) + assert cuda_profiler_stop_response.status_code == 200 + + return updated_request_records + + def _process_warmup_requests(self, warmup_requests: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + if len(warmup_requests) == 0: + return warmup_requests + # NOTE: to warm up the server for as more different batch sizes as possible, + # we usese 128 output tokens for the first request and use two more tokens + # for every followup request. + # Setting a high temperature and top-p to avoid early stop as much as possible. + warmup_requests[0].chat_cmpl.max_tokens = 128 + for i in range(1, len(warmup_requests)): + warmup_requests[i].chat_cmpl.max_tokens = ( + warmup_requests[i - 1].chat_cmpl.max_tokens + 1 + ) + warmup_requests[i].chat_cmpl.temperature = 2.0 + warmup_requests[i].chat_cmpl.top_p = 1.0 + return warmup_requests + + +class SequentialProcessor(RequestProcessor): + """The processor that sequentially applies a list of processors in order.""" + + processors: List[RequestProcessor] # noqa: UP006 + + def __init__(self, *processors: RequestProcessor) -> None: + self.processors = list(processors) + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + for processor in self.processors: + request_records = processor(request_records) + return request_records + + +class Executor(RequestProcessor): + """The executor base class, denoting the kind of benchmark mode.""" + + def __init__( + self, + f_create_api_endpoint: Callable[[], APIEndPoint], + num_processes: int, + disable_tqdm: bool, + ) -> None: + self.f_create_api_endpoint = f_create_api_endpoint + self.disable_tqdm = disable_tqdm + self.num_processes = num_processes + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + raise NotImplementedError() + + +class FixedConcurrentRequestExecutor(Executor): + """The benchmark executor of fixing the number of concurrent requests.""" + + def __init__( + self, + f_create_api_endpoint: Callable[[], APIEndPoint], + num_processes: Optional[int], + disable_tqdm: bool, + num_concurrent_requests: int, + multi_round: bool, + ) -> None: + if num_processes is None: + # We assign each process at most 32 concurrent requests to send + # so that the asyncio pressure will not be too much. + num_processes = min((num_concurrent_requests + 31) // 32, 10) + super().__init__(f_create_api_endpoint, num_processes, disable_tqdm) + self.num_concurrent_requests = num_concurrent_requests + self.multi_round = multi_round + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + partitions: List[List[RequestRecord]] = [ # noqa: UP006 + request_records[slice(i, len(request_records), self.num_processes)] + for i in range(self.num_processes) + ] + # Package "tokenizers" reports warnings with multiprocessing. + # We disable "TOKENIZERS_PARALLELISM" to depress the warnings. + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + pbar = None if self.disable_tqdm else tqdm(total=len(request_records)) + with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_processes) as pool: + futures = [ + pool.submit( + FixedConcurrentRequestExecutor._process_task, + self.f_create_api_endpoint, + partition, + self.num_concurrent_requests // self.num_processes + + int(i < self.num_concurrent_requests % self.num_processes), + self.multi_round, + ) + for i, partition in enumerate(partitions) + ] + results: List[RequestRecord] = [] # noqa: UP006 + for i, future in enumerate(concurrent.futures.as_completed(futures)): + results.extend(future.result()) + if pbar is not None: + pbar.update(len(partitions[i])) + + return results + + @staticmethod + def _process_task( + f_create_api_endpoint: Callable[[], APIEndPoint], + request_records: List[RequestRecord], # noqa: UP006 + num_concurrent_requests: int, + multi_round: bool, + ) -> List[RequestRecord]: # noqa: UP006 + if len(request_records) == 0: + return [] + chat_history: List[List[ChatCompletionMessage]] = [ # noqa: UP006 + [] for _ in range(num_concurrent_requests) + ] + + async def process_task_impl( + f_create_api_endpoint: Callable[[], APIEndPoint], + request_records: List[RequestRecord], # noqa: UP006 + num_concurrent_requests: int, + multi_round: bool, + ) -> List[RequestRecord]: # noqa: UP006 + api_endpoint = f_create_api_endpoint() + updated_request_records: List[RequestRecord] = [None for _ in request_records] # noqa: UP006 + async with api_endpoint: + num_sent_request = 0 + + async def _task(i: int) -> None: + nonlocal num_sent_request + while True: + if num_sent_request == len(request_records): + break + idx = num_sent_request + num_sent_request += 1 + request = request_records[idx] + + if multi_round: + request.chat_cmpl.messages = ( + chat_history[i] + request.chat_cmpl.messages + ) + + updated_request_records[idx] = await api_endpoint(request) + + if multi_round: + chat_history[i] = [ + *updated_request_records[idx].chat_cmpl.messages, + ChatCompletionMessage( + content=updated_request_records[idx].output_str, + role="assistant", + ), + ] + + tasks = [asyncio.create_task(_task(i)) for i in range(num_concurrent_requests)] + await asyncio.gather(*tasks) + + return updated_request_records + + return asyncio.run( + process_task_impl( + f_create_api_endpoint, + request_records, + num_concurrent_requests, + multi_round, + ) + ) + + +class FixTimestampExecutor(Executor): + """The benchmark executor of fixing the timestamps of sending requests.""" + + def __init__( + self, + f_create_api_endpoint: Callable[[], APIEndPoint], + num_processes: Optional[int], + disable_tqdm: bool, + max_schedule_gap: float, + num_requests: int, + ) -> None: + if num_processes is None: + # We assign each process at most 32 requests to send + # so that the asyncio pressure will not be too much. + num_processes = min((num_requests + 31) // 32, 10) + super().__init__(f_create_api_endpoint, num_processes, disable_tqdm) + self.max_schedule_gap = max_schedule_gap + self.num_requests = num_requests + + def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]: # noqa: UP006 + assert len(request_records) > 0 + assert all(request_record.timestamp is not None for request_record in request_records) + # Sort the request records in timestamp ascending order before partitioning. + request_records.sort(key=lambda request_record: request_record.timestamp) + base_timestamp = request_records[0].timestamp + partitions: List[List[RequestRecord]] = [ # noqa: UP006 + request_records[slice(i, len(request_records), self.num_processes)] + for i in range(self.num_processes) + ] + base_sys_time = time.time() + # Package "tokenizers" reports warnings with multiprocessing. + # We disable "TOKENIZERS_PARALLELISM" to depress the warnings. + os.environ["TOKENIZERS_PARALLELISM"] = "false" + + pbar = None if self.disable_tqdm else tqdm(total=len(request_records)) + with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_processes) as pool: + futures = [ + pool.submit( + FixTimestampExecutor._process_task, + self.f_create_api_endpoint, + partition, + base_timestamp, + base_sys_time, + self.max_schedule_gap, + ) + for partition in partitions + ] + results: List[RequestRecord] = [] # noqa: UP006 + for i, future in enumerate(concurrent.futures.as_completed(futures)): + results.extend(future.result()) + if pbar is not None: + pbar.update(len(partitions[i])) + + return results + + @staticmethod + def _process_task( + f_create_api_endpoint: Callable[[], APIEndPoint], + request_records: List[RequestRecord], # noqa: UP006 + base_timestamp: float, + base_sys_time: float, + max_schedule_gap: float, + ) -> List[RequestRecord]: # noqa: UP006 + if len(request_records) == 0: + return [] + + async def process_task_impl( + f_create_api_endpoint: Callable[[], APIEndPoint], + request_records: List[RequestRecord], # noqa: UP006 + base_timestamp: float, + base_sys_time: float, + max_schedule_gap: float, + ) -> List[RequestRecord]: # noqa: UP006 + api_endpoint = f_create_api_endpoint() + loop = asyncio.get_running_loop() + # Get the delta time to convert system time to the loop time. + # We must use the system time `time.time()` which is consistent across processes. + loop_sys_delta_time = loop.time() - time.time() + updated_request_records: List[RequestRecord] = [] # noqa: UP006 + async with api_endpoint: + + async def _task(request_record: RequestRecord) -> None: + updated_request_records.append(await api_endpoint(request_record)) + + tasks = [] + for request_record in request_records: + launch_time = ( + (request_record.timestamp - base_timestamp) + + (base_sys_time + max_schedule_gap) + + loop_sys_delta_time + ) + loop.call_at( + launch_time, + lambda record: tasks.append(asyncio.create_task(_task(record))), + request_record, + ) + # Sleep to allow runs of other scheduled tasks if any. + await asyncio.sleep(max(launch_time - loop.time() - max_schedule_gap, 0)) + + # Sleep until all the tasks are launched. + await asyncio.sleep(launch_time - loop.time() + max_schedule_gap) + # Wait for all tasks to be scheduled + assert len(tasks) == len(request_records) + await asyncio.gather(*tasks) + + assert len(updated_request_records) == len(request_records) + return updated_request_records + + return asyncio.run( + process_task_impl( + f_create_api_endpoint, + request_records, + base_timestamp, + base_sys_time, + max_schedule_gap, + ) + ) + + +def create_pipelines( + args: argparse.Namespace, + f_create_api_endpoint: Callable[[], APIEndPoint], + dataset: Dataset, +) -> List[RequestProcessor]: # noqa: UP006 + """Creating request processing pipelines with regard to the specified args.""" + cuda_profile_url = f"http://{args.host}:{args.port}" if args.cuda_profile else None + pipelines: List[RequestProcessor] = [] # noqa: UP006 + if args.num_concurrent_requests is not None: + if args.request_rate is not None: + raise ValueError( + 'Both "num_concurrent_requests" and "request_rate" are specified. ' + "Please specify only one of them." + ) + if args.replay_timestamp_scale is not None: + raise ValueError( + "Dataset replay is unsupported when fixing number of concurrent requests." + ) + for num_concurrent_requests in args.num_concurrent_requests: + num_warmup_requests = ( + args.num_warmup_requests + if args.num_warmup_requests is not None + else num_concurrent_requests + ) + pipelines.append( + SequentialProcessor( + LogMessage(f"Fixing number of concurrent requests: {num_concurrent_requests}"), + SampleRequests(args.num_requests + num_warmup_requests), + AttachModelName(args.tokenizer), + AttachStreamFlag(args.stream), + AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos), + AttachExecutionFeature({"num_concurrent_requests": num_concurrent_requests}), + WarmupAndRun( + num_warmup_requests=num_warmup_requests, + num_benchmark_requests=args.num_requests, + pipeline=FixedConcurrentRequestExecutor( + f_create_api_endpoint, + args.num_process_workers, + args.disable_tqdm, + num_concurrent_requests, + args.multi_round, + ), + cuda_profile_url=cuda_profile_url, + fake_warmup=dataset.require_fake_warmup, + ), + ) + ) + return pipelines + if args.request_rate is not None: + if args.num_warmup_requests is None: + raise ValueError( + "Please specify the number of warmup requests via " + '"--num-warmup-requests" when fixing request rate.' + ) + if args.replay_timestamp_scale is not None: + raise ValueError("Dataset replay is unsupported when fixing request rates.") + num_total_requests = int( + args.num_requests if not args.per_gpu_workload else args.num_requests * args.num_gpus + ) + if dataset.require_fake_warmup: + num_samples = num_total_requests + else: + num_samples = num_total_requests + args.num_warmup_requests + return [ + SequentialProcessor( + LogMessage(f"Fixing request rate: {request_rate}"), + SampleRequests(num_samples), + AttachModelName(args.tokenizer), + AttachRequestRateTimestamp( + request_rate if not args.per_gpu_workload else request_rate * args.num_gpus + ), + AttachStreamFlag(args.stream), + AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos), + AttachExecutionFeature({"request_rate": float(request_rate)}), + WarmupAndRun( + num_warmup_requests=args.num_warmup_requests, + num_benchmark_requests=num_total_requests, + pipeline=FixTimestampExecutor( + f_create_api_endpoint, + args.num_process_workers, + args.disable_tqdm, + args.max_schedule_gap, + args.num_requests, + ), + cuda_profile_url=cuda_profile_url, + fake_warmup=dataset.require_fake_warmup, + ), + ) + for request_rate in args.request_rate + ] + + # Default: dataset replay mode + # The dataset must come with timestamps. + if not dataset.timestamp_available: + raise ValueError( + "The dataset does not have timestamps, so dataset replay is unsupported. " + 'Please specify one of "num_concurrent_requests" ' + 'and "request_rate".' + ) + if args.per_gpu_workload: + raise ValueError("Fixing per-GPU workload is not compatible with dataset replay.") + if args.num_warmup_requests is None: + raise ValueError( + "Please specify the number of warmup requests via " + '"--num-warmup-requests" for dataset replay.' + ) + timestamp_scale = args.replay_timestamp_scale or 1.0 + if dataset.require_fake_warmup: + num_samples = args.num_requests + else: + num_samples = args.num_requests + args.num_warmup_requests + return [ + SequentialProcessor( + LogMessage(f"Dataset replay with time scaling of {timestamp_scale}"), + SampleRequests(num_samples, take_first_x_requests=True), + AttachModelName(args.tokenizer), + ScaleTimestamp(timestamp_scale), + AttachStreamFlag(args.stream), + AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos), + AttachExecutionFeature({"timestamp_scale": timestamp_scale}), + WarmupAndRun( + num_warmup_requests=args.num_warmup_requests, + num_benchmark_requests=args.num_requests, + pipeline=FixTimestampExecutor( + f_create_api_endpoint, + args.num_process_workers, + args.disable_tqdm, + args.max_schedule_gap, + args.num_requests, + ), + cuda_profile_url=cuda_profile_url, + fake_warmup=dataset.require_fake_warmup, + ), + ) + ] diff --git a/python/mlc_llm/bench/request_record.py b/python/mlc_llm/bench/request_record.py new file mode 100644 index 0000000..f23448f --- /dev/null +++ b/python/mlc_llm/bench/request_record.py @@ -0,0 +1,278 @@ +"""MLC LLM Bench Request""" + +from typing import Any, Dict, List, Optional, Tuple, Union # noqa: UP035 + +import pandas as pd +from pydantic import BaseModel + +from mlc_llm.protocol.openai_api_protocol import ChatCompletionRequest +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +class ServerMetrics(BaseModel): + """The metrics from the server side.""" + + input_tokens: int + prefill_tokens: int + output_tokens: int + end_to_end_latency_s: float + prefill_tokens_per_s: float + inter_token_latency_s: float + time_per_output_token_s: float + time_to_first_token_s: Optional[float] = None + + +class Metrics(BaseModel): + """The list of metric keys""" + + success: bool + start_time: float + finish_time: float + end_to_end_latency_s: float + + input_tokens: Optional[int] = None + output_tokens: Optional[int] = None + inter_token_latency_s: Optional[float] = None + time_per_output_token_s: Optional[float] = None + time_to_first_token_s: Optional[float] = None + server_metrics: Optional[ServerMetrics] = None + + exec_feature: Optional[Dict[str, Any]] = None # noqa: UP006 + + +class RequestRecord(BaseModel): + """The request records collected from LLM inference requests.""" + + request_id: Optional[int] = None + chat_cmpl: ChatCompletionRequest + output_str: Optional[str] = None + first_chunk_output_str: str = "" + timestamp: Optional[float] = None + metrics: Optional[Metrics] = None + error_msg: Optional[str] = None + + +class GroupedRequestRecord(RequestRecord): + """The data structure for request record groups. + For datasets that have common prefix sharing, the request records + that share a same common prefix will be wrapped in a GroupedRequestRecord + at the beginning. + """ + + records: List[RequestRecord] # noqa: UP006 + + +def generate_metrics_summary( + request_records: List[RequestRecord], # noqa: UP006 + num_total_requests: int, + num_gpus: int, +) -> Dict[str, Any]: # noqa: UP006 + """Computes summary statistics across all metrics collected. + Return a dictionary as the report. + """ + num_completed_requests = len(request_records) + assert num_completed_requests <= num_total_requests + request_metrics = [record.metrics for record in request_records] + duration = ( + max(metrics.finish_time for metrics in request_metrics) + - min(metrics.start_time for metrics in request_metrics) + if num_completed_requests > 0 + else 1e-5 + ) + + report = _compute_metrics_statistics(request_metrics) + report["num_gpus"] = num_gpus + report["duration"] = duration + report["num_total_requests"] = num_total_requests + report["num_completed_requests"] = num_completed_requests + report["request_throughput"] = num_completed_requests / duration + + total_input_tokens = sum(metric.input_tokens for metric in request_metrics) + total_output_tokens = sum(metric.output_tokens for metric in request_metrics) + report["total_input_tokens"] = total_input_tokens + report["total_output_tokens"] = total_output_tokens + report["input_token_throughput"] = total_input_tokens / duration + report["input_token_throughput_per_gpu"] = report["input_token_throughput"] / num_gpus + report["output_token_throughput"] = total_output_tokens / duration + report["output_token_throughput_per_gpu"] = report["output_token_throughput"] / num_gpus + + # Generate the server metrics statistics + server_metrics = [metric.server_metrics for metric in request_metrics if metric.server_metrics] + server_report = _compute_metrics_statistics(server_metrics) + if server_report is not None and len(server_report) > 0: + report["server_metrics"] = server_report + + report = { + "exec_feature": ( + request_records[0].metrics.exec_feature if num_completed_requests > 0 else None + ), + **report, + } + return report + + +def _compute_metrics_statistics( + metrics: List[Union[Metrics, ServerMetrics]], # noqa: UP006 +) -> Dict[str, Any]: # noqa: UP006 + """ + Compute the statistics of the metrics. + + Parameters + ---------- + metrics : List[Union[Metrics, ServerMetrics]] + The list of metrics to get the statistics. + + Returns + ------- + report : Dict + The statistics of the metrics. + """ + if not metrics: + return {} + + report: Dict = {} # noqa: UP006 + df = pd.DataFrame([metric.model_dump() for metric in metrics]) + for key, _ in metrics[0].model_fields.items(): + if key in [ + "success", + "start_time", + "finish_time", + "server_metrics", + "exec_feature", + ]: + continue + if key in df.columns: + series = df[key].dropna() + report[key] = { + "quantiles": { + f"p{int(q * 100)}": v + for q, v in series.quantile([0.25, 0.5, 0.75, 0.9, 0.95, 0.99]).items() + }, + "mean": series.mean(), + "min": series.min(), + "max": series.max(), + "stddev": series.std(), + } + return report + + +def convert_reports_to_df(reports: List[Dict[str, Any]]) -> pd.DataFrame: # noqa: UP006 + """Convert benchmark reports to pandas DataFrame.""" + + def _flatten_dict(d: Dict[str, Any], parent_key: str = "") -> Dict[str, Any]: # noqa: UP006 + items: List[Tuple[str, Any]] = [] # noqa: UP006 + for key, value in d.items(): + new_key = f"{parent_key}.{key}" if parent_key != "" else key + if isinstance(value, dict): + items.extend(_flatten_dict(value, new_key).items()) + else: + items.append((new_key, value)) + return dict(items) + + return pd.DataFrame([_flatten_dict(report) for report in reports]) + + +def pretty_print_report(report: Dict[str, Any]) -> None: # noqa: UP006 + """Pretty print the metrics report.""" + + def _print(report: Dict[str, Any], server_metrics: bool): # noqa: UP006 + # fmt: off + title = "Benchmark Result" + if server_metrics: + title += " (server side)" + print(f" {title} ".center(50, "=")) + if not server_metrics: + print(f"{'Total requests:':<40} {report['num_total_requests']:<10}") + print(f"{'Completed requests:':<40} {report['num_completed_requests']:<10}") + print(f"{'Duration (s):':<40} {report['duration']:<10.2f}") + print(f"{'Num GPUs:':<40} {report['num_gpus']:<10}") + print(f"{'Total input tokens:':<40} {report['total_input_tokens']:<10}") + print(f"{'Total output tokens:':<40} {report['total_output_tokens']:<10}") + print(f"{'Request throughput (req/s):':<40} {report['request_throughput']:<10.2f}") + print(f"{'Input token throughput (tok/s):':<40} {report['input_token_throughput']:<10.2f}") # noqa: E501 + print(f"{'Input token throughput per GPU (tok/s):':<40} {report['input_token_throughput_per_gpu']:<10.2f}") # noqa: E501 + print(f"{'Output token throughput (tok/s):':<40} {report['output_token_throughput']:<10.2f}") # noqa: E501 + print(f"{'Output token throughput per GPU (tok/s):':<40} {report['output_token_throughput_per_gpu']:<10.2f}") # noqa: E501 + + if report["num_completed_requests"] == 0: + return + ttft = report["time_to_first_token_s"] + print(" Time to First Token (TTFT, ms) ".center(50, "-")) + print(f"{'Mean:':<40} {ttft['mean'] * 1000:<10.2f}") + print(f"{'Stddev:':<40} {ttft['stddev'] * 1000:<10.2f}") + print(f"{'P25:':<40} {ttft['quantiles']['p25'] * 1000:<10.2f}") + print(f"{'P50:':<40} {ttft['quantiles']['p50'] * 1000:<10.2f}") + print(f"{'P75:':<40} {ttft['quantiles']['p75'] * 1000:<10.2f}") + print(f"{'P90:':<40} {ttft['quantiles']['p90'] * 1000:<10.2f}") + print(f"{'P95:':<40} {ttft['quantiles']['p95'] * 1000:<10.2f}") + print(f"{'P99:':<40} {ttft['quantiles']['p99'] * 1000:<10.2f}") + print(f"{'Min:':<40} {ttft['min'] * 1000:<10.2f}") + print(f"{'Max:':<40} {ttft['max'] * 1000:<10.2f}") + + tpot = report["time_per_output_token_s"] + print(" Time per Output Token (TPOT, ms) ".center(50, "-")) + print(f"{'Mean:':<40} {tpot['mean'] * 1000:<10.2f}") + print(f"{'Stddev:':<40} {tpot['stddev'] * 1000:<10.2f}") + print(f"{'P25:':<40} {tpot['quantiles']['p25'] * 1000:<10.2f}") + print(f"{'P50:':<40} {tpot['quantiles']['p50'] * 1000:<10.2f}") + print(f"{'P75:':<40} {tpot['quantiles']['p75'] * 1000:<10.2f}") + print(f"{'P90:':<40} {tpot['quantiles']['p90'] * 1000:<10.2f}") + print(f"{'P95:':<40} {tpot['quantiles']['p95'] * 1000:<10.2f}") + print(f"{'P99:':<40} {tpot['quantiles']['p99'] * 1000:<10.2f}") + print(f"{'Min:':<40} {tpot['min'] * 1000:<10.2f}") + print(f"{'Max:':<40} {tpot['max'] * 1000:<10.2f}") + + itl = report["inter_token_latency_s"] + print(" Inter-Token Latency (ms) ".center(50, "-")) + print(f"{'Mean:':<40} {itl['mean'] * 1000:<10.2f}") + print(f"{'Stddev:':<40} {itl['stddev'] * 1000:<10.2f}") + print(f"{'P25:':<40} {itl['quantiles']['p25'] * 1000:<10.2f}") + print(f"{'P50:':<40} {itl['quantiles']['p50'] * 1000:<10.2f}") + print(f"{'P75:':<40} {itl['quantiles']['p75'] * 1000:<10.2f}") + print(f"{'P90:':<40} {itl['quantiles']['p90'] * 1000:<10.2f}") + print(f"{'P95:':<40} {itl['quantiles']['p95'] * 1000:<10.2f}") + print(f"{'P99:':<40} {itl['quantiles']['p99'] * 1000:<10.2f}") + print(f"{'Min:':<40} {itl['min'] * 1000:<10.2f}") + print(f"{'Max:':<40} {itl['max'] * 1000:<10.2f}") + + e2e_latency = report["end_to_end_latency_s"] + print(" End-to-End Latency (ms) ".center(50, "-")) + print(f"{'Mean:':<40} {e2e_latency['mean'] * 1000:<10.2f}") + print(f"{'Stddev:':<40} {e2e_latency['stddev'] * 1000:<10.2f}") + print(f"{'P25:':<40} {e2e_latency['quantiles']['p25'] * 1000:<10.2f}") + print(f"{'P50:':<40} {e2e_latency['quantiles']['p50'] * 1000:<10.2f}") + print(f"{'P75:':<40} {e2e_latency['quantiles']['p75'] * 1000:<10.2f}") + print(f"{'P90:':<40} {e2e_latency['quantiles']['p90'] * 1000:<10.2f}") + print(f"{'P95:':<40} {e2e_latency['quantiles']['p95'] * 1000:<10.2f}") + print(f"{'P99:':<40} {e2e_latency['quantiles']['p99'] * 1000:<10.2f}") + print(f"{'Min:':<40} {e2e_latency['min'] * 1000:<10.2f}") + print(f"{'Max:':<40} {e2e_latency['max'] * 1000:<10.2f}") + + input_tokens = report["input_tokens"] + print(" Input Tokens ".center(50, "-")) + print(f"{'Mean:':<40} {input_tokens['mean']:<1}") + print(f"{'Stddev:':<40} {input_tokens['stddev']:<1}") + print(f"{'P25:':<40} {input_tokens['quantiles']['p25']:<1}") + print(f"{'P50:':<40} {input_tokens['quantiles']['p50']:<1}") + print(f"{'P95:':<40} {input_tokens['quantiles']['p95']:<1}") + print(f"{'Min:':<40} {input_tokens['min']:<1}") + print(f"{'Max:':<40} {input_tokens['max']:<1}") + + output_tokens = report["output_tokens"] + print(" Output Tokens ".center(50, "-")) + print(f"{'Mean:':<40} {output_tokens['mean']:<1}") + print(f"{'Stddev:':<40} {output_tokens['stddev']:<1}") + print(f"{'P25:':<40} {output_tokens['quantiles']['p25']:<1}") + print(f"{'P50:':<40} {output_tokens['quantiles']['p50']:<1}") + print(f"{'P95:':<40} {output_tokens['quantiles']['p95']:<1}") + print(f"{'Min:':<40} {output_tokens['min']:<1}") + print(f"{'Max:':<40} {output_tokens['max']:<1}") + + print("=" * 50) + + # fmt: on + _print(report, server_metrics=False) + if "server_metrics" in report: + _print(report["server_metrics"], server_metrics=True) diff --git a/python/mlc_llm/cli/__init__.py b/python/mlc_llm/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/cli/calibrate.py b/python/mlc_llm/cli/calibrate.py new file mode 100644 index 0000000..9903171 --- /dev/null +++ b/python/mlc_llm/cli/calibrate.py @@ -0,0 +1,80 @@ +"""Command line entrypoint of calibration.""" + +from mlc_llm.interface.calibrate import calibrate +from mlc_llm.interface.help import HELP +from mlc_llm.support.argparse import ArgumentParser + +from .serve import EngineConfigOverride + + +def main(argv): + """Main entrypoint for calibration.""" + parser = ArgumentParser("MLC LLM Calibration CLI") + parser.add_argument( + "model", + type=str, + help=HELP["model"] + " (required)", + ) + parser.add_argument( + "--device", + type=str, + default="auto", + help=HELP["device_deploy"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--model-lib", + type=str, + default=None, + help=HELP["model_lib"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--output", + "-o", + type=str, + required=True, + help=HELP["output_calibration"] + " (required)", + ) + # Download dataset from + # https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json + parser.add_argument( + "--dataset", + type=str, + required=True, + help=HELP["calibration_dataset"] + " (required)", + ) + + parser.add_argument( + "--num-calibration-samples", + type=int, + default=16, + help=HELP["num_calibration_samples"] + ' (default: "%(default)s")', + ) + + parser.add_argument( + "--seed", + type=int, + default=0, + help=HELP["seed_calibrate"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--overrides", + type=EngineConfigOverride.from_str, + default="", + help=HELP["overrides_serve"], + ) + + parsed = parser.parse_args(argv) + calibrate( + model=parsed.model, + device=parsed.device, + model_lib=parsed.model_lib, + output=parsed.output, + dataset=parsed.dataset, + num_calibration_samples=parsed.num_calibration_samples, + max_num_sequence=parsed.overrides.max_num_sequence, + max_total_sequence_length=parsed.overrides.max_total_seq_length, + prefill_chunk_size=parsed.overrides.prefill_chunk_size, + max_history_size=parsed.overrides.max_history_size, + gpu_memory_utilization=parsed.overrides.gpu_memory_utilization, + seed=parsed.seed, + ) diff --git a/python/mlc_llm/cli/chat.py b/python/mlc_llm/cli/chat.py new file mode 100644 index 0000000..cb2d089 --- /dev/null +++ b/python/mlc_llm/cli/chat.py @@ -0,0 +1,41 @@ +"""Command line entrypoint of chat.""" + +from mlc_llm.interface.chat import ModelConfigOverride, chat +from mlc_llm.interface.help import HELP +from mlc_llm.support.argparse import ArgumentParser + + +def main(argv): + """Parse command line arguments and call `mlc_llm.interface.chat`.""" + parser = ArgumentParser("MLC LLM Chat CLI") + + parser.add_argument( + "model", + type=str, + help=HELP["model"] + " (required)", + ) + parser.add_argument( + "--device", + type=str, + default="auto", + help=HELP["device_deploy"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--model-lib", + type=str, + default=None, + help=HELP["model_lib"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--overrides", + type=ModelConfigOverride.from_str, + default="", + help=HELP["modelconfig_overrides"] + ' (default: "%(default)s")', + ) + parsed = parser.parse_args(argv) + chat( + model=parsed.model, + device=parsed.device, + model_lib=parsed.model_lib, + overrides=parsed.overrides, + ) diff --git a/python/mlc_llm/cli/check_device.py b/python/mlc_llm/cli/check_device.py new file mode 100644 index 0000000..d2b3a14 --- /dev/null +++ b/python/mlc_llm/cli/check_device.py @@ -0,0 +1,34 @@ +"""Check if a device exists.""" + +import os +import sys + +from tvm.runtime import Device +from tvm.runtime import device as as_device + + +def _check_device(device: Device) -> bool: + try: + return bool(device.exist) + except Exception: + return False + + +def main(): + """Entrypoint for device check.""" + device_str = sys.argv[1] + device_ids = [] + i = 0 + while True: + if _check_device(as_device(device_str, i)): + device_ids.append(i) + i += 1 + if device_str in ["cpu", "llvm"] and i > os.cpu_count() / 2: + break + else: + break + print(f"check_device:{','.join(str(i) for i in device_ids)}") + + +if __name__ == "__main__": + main() diff --git a/python/mlc_llm/cli/compile.py b/python/mlc_llm/cli/compile.py new file mode 100644 index 0000000..233160c --- /dev/null +++ b/python/mlc_llm/cli/compile.py @@ -0,0 +1,152 @@ +"""Command line entrypoint of compilation.""" + +import argparse +import json +import re +from functools import partial +from pathlib import Path +from typing import Union + +from mlc_llm.interface.compile import ( + ModelConfigOverride, + OptimizationFlags, + compile, +) +from mlc_llm.interface.help import HELP +from mlc_llm.model import MODELS +from mlc_llm.quantization import QUANTIZATION +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.support.auto_config import ( + detect_mlc_chat_config, + detect_model_type, + detect_quantization, +) +from mlc_llm.support.auto_target import detect_system_lib_prefix, detect_target_and_host + + +def main(argv): + """Parse command line arguments and call `mlc_llm.compiler.compile`.""" + + def _parse_output(path: Union[str, Path]) -> Path: + path = Path(path) + if path.is_dir(): + raise argparse.ArgumentTypeError(f"Output cannot be a directory: {path}") + parent = path.parent + if not parent.is_dir(): + raise argparse.ArgumentTypeError(f"Directory does not exist: {parent}") + return path + + def _parse_dir(path: Union[str, Path], auto_create: bool = False) -> Path: + path = Path(path) + if not auto_create and not path.is_dir(): + raise argparse.ArgumentTypeError(f"Directory does not exist: {path}") + if auto_create and not path.is_dir(): + path.mkdir(parents=True) + return path + + def _check_system_lib_prefix(prefix: str) -> str: + pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$" + if prefix == "" or re.match(pattern, prefix): + return prefix + raise argparse.ArgumentTypeError( + "Invalid prefix. It should only consist of " + "numbers (0-9), alphabets (A-Z, a-z) and underscore (_)." + ) + + parser = ArgumentParser("mlc_llm compile") + parser.add_argument( + "model", + type=detect_mlc_chat_config, + help=HELP["model"] + " (required)", + ) + parser.add_argument( + "--quantization", + type=str, + choices=list(QUANTIZATION.keys()), + help=HELP["quantization"] + + " (default: look up mlc-chat-config.json, choices: %(choices)s)", + ) + parser.add_argument( + "--model-type", + type=str, + default="auto", + choices=["auto", *list(MODELS.keys())], + help=HELP["model_type"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--device", + type=str, + default="auto", + help=HELP["device_compile"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--host", + type=str, + default="auto", + help=HELP["host"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--enable-subgroups", + action="store_true", + help=HELP["enable_subgroups"], + ) + parser.add_argument( + "--opt", + type=OptimizationFlags.from_str, + default="O2", + help=HELP["opt"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--system-lib-prefix", + type=str, + default="auto", + help=HELP["system_lib_prefix"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--output", + "-o", + type=_parse_output, + required=True, + help=HELP["output_compile"] + " (required)", + ) + parser.add_argument( + "--overrides", + type=ModelConfigOverride.from_str, + default="", + help=HELP["overrides"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--debug-dump", + type=partial(_parse_dir, auto_create=True), + default=None, + help=HELP["debug_dump"] + " (default: %(default)s)", + ) + parsed = parser.parse_args(argv) + target, build_func = detect_target_and_host( + parsed.device, + parsed.host, + enable_subgroups=parsed.enable_subgroups, + ) + parsed.model_type = detect_model_type(parsed.model_type, parsed.model) + parsed.quantization = detect_quantization(parsed.quantization, parsed.model) + parsed.system_lib_prefix = detect_system_lib_prefix( + parsed.device, + parsed.system_lib_prefix, + parsed.model_type.name, + parsed.quantization.name, + ) + with open(parsed.model, encoding="utf-8") as config_file: + config = json.load(config_file) + + compile( + config=config, + quantization=parsed.quantization, + model_type=parsed.model_type, + target=target, + opt=parsed.opt, + build_func=build_func, + system_lib_prefix=parsed.system_lib_prefix, + output=parsed.output, + overrides=parsed.overrides, + debug_dump=parsed.debug_dump, + ) diff --git a/python/mlc_llm/cli/convert_weight.py b/python/mlc_llm/cli/convert_weight.py new file mode 100644 index 0000000..b489f3e --- /dev/null +++ b/python/mlc_llm/cli/convert_weight.py @@ -0,0 +1,112 @@ +"""Command line entrypoint of weight conversion.""" + +import argparse +from pathlib import Path +from typing import Union + +from mlc_llm.interface.convert_weight import convert_weight +from mlc_llm.interface.help import HELP +from mlc_llm.model import MODELS +from mlc_llm.quantization import QUANTIZATION +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.support.auto_config import detect_config, detect_model_type +from mlc_llm.support.auto_device import detect_device +from mlc_llm.support.auto_weight import detect_weight + + +def main(argv): + """Parse command line argumennts and apply quantization.""" + + def _parse_source(path: Union[str, Path], config_path: Path) -> Path: + if path == "auto": + return config_path.parent + path = Path(path) + if not path.exists(): + raise argparse.ArgumentTypeError(f"Model source does not exist: {path}") + return path + + def _parse_output(path: Union[str, Path]) -> Path: + path = Path(path) + if not path.is_dir(): + path.mkdir(parents=True, exist_ok=True) + return path + + def _parse_lora_adapter(path: Union[str, Path]) -> Path: + path = Path(path) + if not path.exists() or not path.is_dir(): + raise argparse.ArgumentTypeError(f"LoRA adapter directory does not exist: {path}") + return path + + parser = ArgumentParser("MLC AutoLLM Quantization Framework") + parser.add_argument( + "config", + type=detect_config, + help=HELP["config"] + " (required)", + ) + parser.add_argument( + "--quantization", + type=str, + required=True, + choices=list(QUANTIZATION.keys()), + help=HELP["quantization"] + " (required, choices: %(choices)s)", + ) + parser.add_argument( + "--model-type", + type=str, + default="auto", + choices=["auto", *list(MODELS.keys())], + help=HELP["model_type"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--device", + default="auto", + type=detect_device, + help=HELP["device_quantize"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--source", + type=str, + default="auto", + help=HELP["source"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--source-format", + type=str, + choices=["auto", "huggingface-torch", "huggingface-safetensor", "awq"], + default="auto", + help=HELP["source_format"] + ' (default: "%(default)s", choices: %(choices)s")', + ) + parser.add_argument( + "--output", + "-o", + type=_parse_output, + required=True, + help=HELP["output_quantize"] + " (required)", + ) + parser.add_argument( + "--lora-adapter", + type=_parse_lora_adapter, + default=None, + help=( + "Path to a LoRA adapter directory in PEFT format. " + "When provided, adapter weights are merged into the base model before quantization." + ), + ) + + parsed = parser.parse_args(argv) + parsed.source, parsed.source_format = detect_weight( + _parse_source(parsed.source, parsed.config), + parsed.config, + parsed.source_format, + ) + model = detect_model_type(parsed.model_type, parsed.config) + convert_weight( + config=parsed.config, + quantization=QUANTIZATION[parsed.quantization], + model=model, + device=parsed.device, + source=parsed.source, + source_format=parsed.source_format, + output=parsed.output, + lora_adapter=parsed.lora_adapter, + ) diff --git a/python/mlc_llm/cli/delivery.py b/python/mlc_llm/cli/delivery.py new file mode 100644 index 0000000..735da54 --- /dev/null +++ b/python/mlc_llm/cli/delivery.py @@ -0,0 +1,452 @@ +"""Continuous model delivery for MLC LLM models.""" + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union # noqa: UP035 + +from huggingface_hub import HfApi, snapshot_download +from huggingface_hub.utils import HfHubHTTPError +from pydantic import BaseModel, Field, ValidationError + +from mlc_llm.support import logging +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.support.style import bold, green, red + +logger = logging.getLogger(__name__) + +GEN_CONFIG_OPTIONAL_ARGS = [ + "context_window_size", + "sliding_window_size", + "prefill_chunk_size", + "attention_sink_size", + "tensor_parallel_shards", + "pipeline_parallel_stages", +] + +T = TypeVar("T", bound="BaseModel") + + +class OverrideConfigs(BaseModel): + """ + The class that specifies the override configurations. + """ + + context_window_size: Optional[int] = None + sliding_window_size: Optional[int] = None + prefill_chunk_size: Optional[int] = None + attention_sink_size: Optional[int] = None + tensor_parallel_shards: Optional[int] = None + pipeline_parallel_stages: Optional[int] = None + + +class ModelDeliveryTask(BaseModel): + """ + Example: + { + "model_id": "Phi-3-mini-128k-instruct", + "model": "HF://microsoft/Phi-3-mini-128k-instruct", + "conv_template": "phi-3", + "quantization": ["q3f16_1"], + "overrides": { + "q3f16_1": { + "context_window_size": 512 + } + } + } + """ + + model_id: str + model: str + conv_template: str + quantization: Union[List[str], str] = Field(default_factory=list) # noqa: UP006 + overrides: Dict[str, OverrideConfigs] = Field(default_factory=dict) # noqa: UP006 + destination: Optional[str] = None + gen_config_only: Optional[bool] = False + + +class ModelDeliveryList(BaseModel): + """ + The class that specifies the model delivery list. + """ + + tasks: List[ModelDeliveryTask] # noqa: UP006 + # For delivered log, the default destination and quantization fields are optional + default_destination: Optional[str] = None + default_quantization: List[str] = Field(default_factory=list) # noqa: UP006 + default_overrides: Dict[str, OverrideConfigs] = Field(default_factory=dict) # noqa: UP006 + + @classmethod + def from_json(cls: Type[T], json_dict: Dict[str, Any]) -> T: # noqa: UP006 + """ + Convert from a json dictionary. + """ + try: + return ModelDeliveryList.model_validate(json_dict) + except ValidationError as e: + logger.error("Error validating ModelDeliveryList: %s", e) + raise e + + def to_json(self) -> Dict[str, Any]: # noqa: UP006 + """ + Convert to a json dictionary. + """ + return self.model_dump(exclude_none=True) + + +def _clone_repo(model: Union[str, Path], hf_local_dir: Optional[str]) -> str: + if isinstance(model, Path): + if not model.exists(): + raise ValueError(f"Invalid model source: {model}") + return str(model) + prefixes, mlc_prefix = ["HF://", "https://huggingface.co/"], "" + mlc_prefix = next(p for p in prefixes if model.startswith(p)) + if mlc_prefix: + repo_name = model[len(mlc_prefix) :] + model_name = repo_name.split("/")[-1] + if hf_local_dir: + hf_local_dir = os.path.join(hf_local_dir, model_name) + logger.info("[HF] Downloading model to %s", hf_local_dir) + return snapshot_download(repo_id=repo_name, local_dir=hf_local_dir) + result = Path(model) + if result.exists(): + return model + raise ValueError(f"Invalid model source: {model}") + + +def _run_quantization( + model_info: ModelDeliveryTask, + repo: str, + api: HfApi, + output_dir: str, +) -> bool: + logger.info("[HF] Creating repo https://huggingface.co/%s", repo) + try: + api.create_repo(repo_id=repo, private=False) + except HfHubHTTPError as error: + if error.response.status_code != 409: + raise + logger.info("[HF] Repo already exists. Skipping creation.") + succeeded = True + log_path = Path(output_dir) / "logs.txt" + with log_path.open("a", encoding="utf-8") as log_file: + assert isinstance(model_info.quantization, str) + logger.info("[MLC] Processing in directory: %s", output_dir) + # Required arguments + cmd = [ + sys.executable, + "-m", + "mlc_llm", + "gen_config", + model_info.model, + "--quantization", + model_info.quantization, + "--conv-template", + model_info.conv_template, + "--output", + output_dir, + ] + # Optional arguments + for optional_arg in GEN_CONFIG_OPTIONAL_ARGS: + optional_arg_val = getattr(model_info, optional_arg, None) + if optional_arg_val is not None: + # e.g. --context-window-size 4096 + cmd += ["--" + optional_arg.replace("_", "-"), str(optional_arg_val)] + + print(" ".join(cmd), file=log_file, flush=True) + subprocess.run(cmd, check=True, stdout=log_file, stderr=subprocess.STDOUT, env=os.environ) + if not model_info.gen_config_only: + cmd = [ + sys.executable, + "-m", + "mlc_llm", + "convert_weight", + str(model_info.model), + "--quantization", + model_info.quantization, + "--output", + output_dir, + ] + print(" ".join(cmd), file=log_file, flush=True) + subprocess.run( + cmd, + check=False, + stdout=log_file, + stderr=subprocess.STDOUT, + env=os.environ, + ) + logger.info("[MLC] Complete!") + if not (Path(output_dir) / "tensor-cache.json").exists() and not model_info.gen_config_only: + logger.error( + "[%s] Model %s. Quantization %s. No weights metadata found.", + red("FAILED"), + model_info.model_id, + model_info.quantization, + ) + succeeded = False + logger.info("[HF] Uploading to: https://huggingface.co/%s", repo) + for _retry in range(10): + try: + api.upload_folder( + folder_path=output_dir, + repo_id=repo, + ignore_patterns=["logs.txt"], + ) + except Exception as exc: + logger.error("[%s] %s. Retrying...", red("FAILED"), exc) + else: + break + else: + raise RuntimeError("Failed to upload to HuggingFace Hub with 10 retries") + return succeeded + + +def _get_current_log(log: str) -> ModelDeliveryList: + log_path = Path(log) + if not log_path.exists(): + with log_path.open("w", encoding="utf-8") as o_f: + current_log = ModelDeliveryList(tasks=[]) + json.dump(current_log.to_json(), o_f, indent=4) + else: + with log_path.open("r", encoding="utf-8") as i_f: + current_log = ModelDeliveryList.from_json(json.load(i_f)) + return current_log + + +def _generate_model_delivery_diff( + spec: ModelDeliveryList, log: ModelDeliveryList +) -> ModelDeliveryList: + diff_tasks = [] + default_quantization = spec.default_quantization + default_overrides = spec.default_overrides + + for task in spec.tasks: + model_id = task.model_id + conv_template = task.conv_template + quantization = task.quantization + overrides = {**default_overrides, **task.overrides} + + logger.info( + "Checking task: %s %s %s %s", + model_id, + conv_template, + quantization, + overrides, + ) + log_tasks = [t for t in log.tasks if t.model_id == model_id] + delivered_quantizations = set() + gen_config_only = set() + + for log_task in log_tasks: + log_quantization = log_task.quantization + assert isinstance(log_quantization, str) + log_override = log_task.overrides.get(log_quantization, OverrideConfigs()) + override = overrides.get(log_quantization, OverrideConfigs()) + if log_override == override: + if log_task.conv_template == conv_template: + delivered_quantizations.add(log_quantization) + else: + gen_config_only.add(log_quantization) + + all_quantizations = set(default_quantization) | set(quantization) + quantization_diff = all_quantizations - set(delivered_quantizations) + + if quantization_diff: + for q in quantization_diff: + logger.info("Adding task %s %s %s to the diff.", model_id, conv_template, q) + task_copy = task.model_copy() + task_copy.quantization = [q] + task_copy.overrides = {q: overrides.get(q, OverrideConfigs())} + task_copy.gen_config_only = task_copy.gen_config_only or q in gen_config_only + diff_tasks.append(task_copy) + else: + logger.info("Task %s %s %s is up-to-date.", model_id, conv_template, quantization) + + diff_config = spec.model_copy() + diff_config.default_quantization = [] + diff_config.default_overrides = {} + diff_config.tasks = diff_tasks + + logger.info( + "Model delivery diff: %s", + diff_config.model_dump_json(indent=4, exclude_none=True), + ) + + return diff_config + + +def _main( + username: str, + api: HfApi, + spec: ModelDeliveryList, + log: str, + hf_local_dir: Optional[str], + output: str, + dry_run: bool, +): + delivery_diff = _generate_model_delivery_diff(spec, _get_current_log(log)) + if dry_run: + logger.info("Dry run. No actual delivery.") + return + + failed_cases: List[Tuple[str, str]] = [] # noqa: UP006 + delivered_log = _get_current_log(log) + for task_index, task in enumerate(delivery_diff.tasks, 1): + logger.info( + bold("[{task_index}/{total_tasks}] Processing model: ").format( + task_index=task_index, + total_tasks=len(delivery_diff.tasks), + ) + + green(task.model_id) + ) + model = _clone_repo(task.model, hf_local_dir) + + quantizations = [] + + if delivery_diff.default_quantization: + quantizations += delivery_diff.default_quantization + + if task.quantization: + if isinstance(task.quantization, str): + quantizations.append(task.quantization) + else: + quantizations += task.quantization + + default_destination = ( + delivery_diff.default_destination or "{username}/{model_id}-{quantization}-MLC" + ) + for quantization in quantizations: + repo = default_destination.format( + username=username, + model_id=task.model_id, + quantization=quantization, + ) + model_info = ModelDeliveryTask( + model=model, + quantization=quantization, + destination=repo, + **task.model_dump(exclude_none=True, exclude={"model", "quantization"}), + ) + logger.info("Model info: %s", model_info.model_dump_json(indent=4)) + output_dir = os.path.join( + output, f"{model_info.model_id}-{model_info.quantization}-MLC" + ) + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + result = _run_quantization( + model_info=model_info, + repo=repo, + api=api, + output_dir=output_dir, + ) + if not result: + failed_cases.append( + (task.model_id, quantization), + ) + else: + delivered_log.tasks = [ + task + for task in delivered_log.tasks + if task.model_id != model_info.model_id + or task.quantization != model_info.quantization + ] + delivered_log.tasks.append(model_info) + if failed_cases: + logger.info("Total %s %s:", len(failed_cases), red("failures")) + for model_id, quantization in failed_cases: + logger.info(" Model %s. Quantization %s.", model_id, quantization) + + delivered_log.tasks.sort(key=lambda task: task.model_id) + logger.info("Writing log to %s", log) + with open(log, "w", encoding="utf-8") as o_f: + json.dump(delivered_log.to_json(), o_f, indent=4) + + +def main(): + """Entry point.""" + + def _load_spec(path_spec: str) -> ModelDeliveryList: + path = Path(path_spec) + if not path.exists(): + raise argparse.ArgumentTypeError(f"Spec file does not exist: {path}") + with path.open("r", encoding="utf-8") as i_f: + return ModelDeliveryList.from_json(json.load(i_f)) + + def _get_default_hf_token() -> str: + # Try to get the token from the environment variable + hf_token = os.getenv("HF_TOKEN") + if hf_token: + logger.info("HF token found in environment variable HF_TOKEN") + return hf_token + + # If not found, look for the token in the default cache folder + token_file_path = os.path.expanduser("~/.cache/huggingface/token") + if os.path.exists(token_file_path): + with open(token_file_path, encoding="utf-8") as token_file: + hf_token = token_file.read().strip() + if hf_token: + logger.info("HF token found in ~/.cache/huggingface/token") + return hf_token + + raise OSError("HF token not found") + + parser = ArgumentParser("MLC LLM continuous model delivery") + parser.add_argument( + "--username", + type=str, + required=True, + help="HuggingFace username", + ) + parser.add_argument( + "--token", + type=str, + default=_get_default_hf_token(), + help="HuggingFace access token, obtained under https://huggingface.co/settings/tokens", + ) + parser.add_argument( + "--spec", + type=_load_spec, + default="model-delivery-config.json", + help="Path to the model delivery file" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--log", + type=str, + default="model-delivered-log.json", + help="Path to the output log file" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--output", + type=str, + required=True, + help="Directory to store the output MLC models", + ) + parser.add_argument( + "--hf-local-dir", + type=str, + required=False, + help="Local directory to store the downloaded HuggingFace model", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Dry run without uploading to HuggingFace Hub", + ) + parsed = parser.parse_args() + _main( + parsed.username, + spec=parsed.spec, + log=parsed.log, + api=HfApi(token=parsed.token), + hf_local_dir=parsed.hf_local_dir, + output=parsed.output, + dry_run=parsed.dry_run, + ) + + +if __name__ == "__main__": + main() diff --git a/python/mlc_llm/cli/disco_remote_socket_session.py b/python/mlc_llm/cli/disco_remote_socket_session.py new file mode 100644 index 0000000..d7a4675 --- /dev/null +++ b/python/mlc_llm/cli/disco_remote_socket_session.py @@ -0,0 +1,33 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Internal remote disco socket session.""" + +import sys + +from tvm import runtime as _ # noqa: F401 +from tvm_ffi import get_global_func + +from .. import base # noqa: F401 + +if __name__ == "__main__": + if len(sys.argv) != 4: + print("Usage: ") + sys.exit(1) + + server_host = sys.argv[1] + server_port = int(sys.argv[2]) + num_workers = int(sys.argv[3]) + func = get_global_func("runtime.disco.RemoteSocketSession") + func(server_host, server_port, num_workers) diff --git a/python/mlc_llm/cli/gen_config.py b/python/mlc_llm/cli/gen_config.py new file mode 100644 index 0000000..a622f79 --- /dev/null +++ b/python/mlc_llm/cli/gen_config.py @@ -0,0 +1,121 @@ +"""Command line entrypoint of configuration generation.""" + +from pathlib import Path +from typing import Union + +from mlc_llm.interface.gen_config import CONV_TEMPLATES, gen_config +from mlc_llm.interface.help import HELP +from mlc_llm.model import MODELS +from mlc_llm.quantization import QUANTIZATION +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.support.auto_config import detect_config, detect_model_type + + +def main(argv): + """Parse command line argumennts and call `mlc_llm.compiler.gen_config`.""" + parser = ArgumentParser("MLC LLM Configuration Generator") + + def _parse_output(path: Union[str, Path]) -> Path: + path = Path(path) + if not path.is_dir(): + path.mkdir(parents=True, exist_ok=True) + return path + + parser.add_argument( + "config", + type=detect_config, + help=HELP["config"] + " (required)", + ) + parser.add_argument( + "--quantization", + type=str, + required=True, + choices=list(QUANTIZATION.keys()), + help=HELP["quantization"] + " (required, choices: %(choices)s)", + ) + parser.add_argument( + "--model-type", + type=str, + default="auto", + choices=["auto", *list(MODELS.keys())], + help=HELP["model_type"] + ' (default: "%(default)s", choices: %(choices)s)', + ) + parser.add_argument( + "--conv-template", + type=str, + required=True, + choices=list(CONV_TEMPLATES), + help=HELP["conv_template"] + " (required, choices: %(choices)s)", + ) + parser.add_argument( + "--context-window-size", + type=int, + default=None, + help=HELP["context_window_size"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--sliding-window-size", + type=int, + default=None, + help=HELP["sliding_window_size"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--prefill-chunk-size", + type=int, + default=None, + help=HELP["prefill_chunk_size"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--attention-sink-size", + type=int, + default=None, + help=HELP["attention_sink_size"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--tensor-parallel-shards", + type=int, + default=None, + help=HELP["tensor_parallel_shards"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--pipeline-parallel-stages", + type=int, + default=None, + help=HELP["pipeline_parallel_stages"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--disaggregation", + type=bool, + default=None, + help=HELP["disaggregation"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--max-batch-size", + type=int, + default=128, + help=HELP["max_batch_size"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--output", + "-o", + type=_parse_output, + required=True, + help=HELP["output_gen_mlc_chat_config"] + " (required)", + ) + parsed = parser.parse_args(argv) + model = detect_model_type(parsed.model_type, parsed.config) + gen_config( + config=parsed.config, + model=model, + quantization=QUANTIZATION[parsed.quantization], + conv_template=parsed.conv_template, + context_window_size=parsed.context_window_size, + sliding_window_size=parsed.sliding_window_size, + prefill_chunk_size=parsed.prefill_chunk_size, + attention_sink_size=parsed.attention_sink_size, + tensor_parallel_shards=parsed.tensor_parallel_shards, + pipeline_parallel_stages=parsed.pipeline_parallel_stages, + disaggregation=parsed.disaggregation, + max_batch_size=parsed.max_batch_size, + output=parsed.output, + ) diff --git a/python/mlc_llm/cli/lib_delivery.py b/python/mlc_llm/cli/lib_delivery.py new file mode 100644 index 0000000..e44a5b8 --- /dev/null +++ b/python/mlc_llm/cli/lib_delivery.py @@ -0,0 +1,199 @@ +"""Continuous model delivery for MLC LLM models.""" + +import argparse +import dataclasses +import json +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any, Callable, Dict, List # noqa: UP035 + +from mlc_llm.support import logging +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.support.constants import MLC_TEMP_DIR +from mlc_llm.support.style import bold, green, red + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class ModelInfo: + """Necessary information for the model delivery""" + + model_id: str + model: Path + quantization: str + device: str + # overrides the `context_window_size`, `prefill_chunk_size`, + # `sliding_window_size`, `attention_sink_size`, `max_batch_size` + # and `tensor_parallel_shards in mlc-chat-config.json + overrides: Dict[str, int] # noqa: UP006 + + +class DeferredScope: + """A context manager that defers execution of functions until exiting the scope.""" + + def __init__(self): + self.deferred_functions = [] + + def add(self, func: Callable[[], None]): + """Add a function to be executed when exiting the scope.""" + self.deferred_functions.append(func) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + for func in reversed(self.deferred_functions): + func() + return False + + def create_temp_dir(self) -> Path: + """Create a temporary directory that will be deleted when exiting the scope.""" + temp_dir = tempfile.mkdtemp(dir=MLC_TEMP_DIR) + self.add(lambda: shutil.rmtree(temp_dir, ignore_errors=True)) + return Path(temp_dir) + + +def _run_compilation(model_info: ModelInfo, repo_dir: Path) -> bool: + """Run the compilation of the model library.""" + + def get_lib_ext(device: str) -> str: + if device in ["cuda", "vulkan", "metal"]: + return ".so" + if device in ["android", "ios"]: + return ".tar" + if device in ["webgpu"]: + return ".wasm" + + return "" + + succeeded = True + with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as temp_dir: + log_path = Path(temp_dir) / "logs.txt" + model_lib_name = f"{model_info.model_id}-{model_info.quantization}-{model_info.device}" + lib_ext = get_lib_ext(model_info.device) + if lib_ext == "": + raise ValueError(f"Unsupported device: {model_info.device}") + model_lib_name += lib_ext + with log_path.open("a", encoding="utf-8") as log_file: + overrides = ";".join(f"{key}={value}" for key, value in model_info.overrides.items()) + cmd = [ + sys.executable, + "-m", + "mlc_llm", + "compile", + str(model_info.model), + "--device", + model_info.device, + "--quantization", + model_info.quantization, + "--overrides", + overrides, + "--output", + os.path.join(temp_dir, model_lib_name), + ] + print(" ".join(cmd), file=log_file, flush=True) + subprocess.run(cmd, check=True, stdout=log_file, stderr=subprocess.STDOUT) + logger.info("[MLC] Compilation Complete!") + if not (Path(temp_dir) / model_lib_name).exists(): + logger.error( + "[%s] Model %s. Device %s. No compiled library found.", + red("FAILED"), + model_info.model_id, + model_info.device, + ) + succeeded = False + return succeeded + + # overwrite git repo file with the compiled library + repo_filepath = repo_dir / model_info.model_id / model_lib_name + if not repo_filepath.parent.exists(): + repo_filepath.parent.mkdir(parents=True, exist_ok=True) + # copy lib from Path(temp_dir) / model_lib_name to repo_filepath + shutil.copy(Path(temp_dir) / model_lib_name, repo_filepath) + logger.info("Saved library %s at %s", model_lib_name, repo_filepath) + return succeeded + + +def _main( + spec: Dict[str, Any], # noqa: UP006 +): + """Compile the model libs in the spec and save them to the binary_libs_dir.""" + failed_cases: List[Any] = [] # noqa: UP006 + for task_index, task in enumerate(spec["tasks"], 1): + logger.info( + bold("[{task_index}/{total_tasks}] Processing model: ").format( + task_index=task_index, + total_tasks=len(spec["tasks"]), + ) + + green(task["model_id"]) + ) + model_info = { + "model_id": task["model_id"], + "model": task["model"], + } + for compile_opt in spec["default_compile_options"] + task.get("compile_options", []): + for quantization in spec["default_quantization"] + task.get("quantization", []): + model_info["quantization"] = quantization + model_info["device"] = compile_opt["device"] + model_info["overrides"] = compile_opt.get("overrides", {}) + logger.info( + "[Config] " + + bold("model_id: ") + + model_info["model_id"] + + bold(", quantization: ") + + model_info["quantization"] + + bold(", device: ") + + model_info["device"] + + bold(", overrides: ") + + json.dumps(model_info["overrides"]) + ) + + result = _run_compilation( + ModelInfo(**model_info), + repo_dir=Path(spec["binary_libs_dir"]), + ) + if not result: + failed_cases.append(model_info) + + if failed_cases: + logger.info("Total %s %s:", len(failed_cases), red("failures")) + for case in failed_cases: + logger.info( + "model_id %s, quantization %s, device %s, overrides %s", + case["model_id"], + case["quantization"], + case["device"], + json.dumps(case["overrides"]), + ) + + +def main(): + """Entry point.""" + + def _load_spec(path_spec: str) -> Dict[str, Any]: # noqa: UP006 + path = Path(path_spec) + if not path.exists(): + raise argparse.ArgumentTypeError(f"Spec file does not exist: {path}") + with path.open("r", encoding="utf-8") as i_f: + return json.load(i_f) + + parser = ArgumentParser("MLC LLM continuous library delivery") + parser.add_argument( + "--spec", + type=_load_spec, + required=True, + help="Path to the spec file", + ) + parsed = parser.parse_args() + _main( + spec=parsed.spec, + ) + + +if __name__ == "__main__": + main() diff --git a/python/mlc_llm/cli/model_metadata.py b/python/mlc_llm/cli/model_metadata.py new file mode 100644 index 0000000..2895ba2 --- /dev/null +++ b/python/mlc_llm/cli/model_metadata.py @@ -0,0 +1,194 @@ +"""A tool that inspects the metadata of a model lib.""" + +import json +import math +from dataclasses import asdict +from pathlib import Path +from typing import Any, Dict, List, Union # noqa: UP035 + +from tvm.runtime import DataType + +from mlc_llm.support import logging +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import green, red + +logger = logging.getLogger(__name__) + + +def _extract_metadata(model_lib: Path) -> Dict[str, Any]: # noqa: UP006 + from tvm.runtime import device, load_module + from tvm.runtime.vm import VirtualMachine + + return json.loads(VirtualMachine(load_module(model_lib), device("cpu"))["_metadata"]()) + + +def _report_all(metadata: Dict[str, Any]) -> None: # noqa: UP006 + # Print JSON with aesthetic values that packs each parameter into one line, + # while keeping the rest indented. + indent = 2 + indents = " " * indent + params = metadata.pop("params") + params = indents * 2 + (",\n" + indents * 2).join(json.dumps(p) for p in params) + lines = json.dumps( + metadata, + sort_keys=True, + indent=indent, + ).splitlines() + lines.insert(1, indents + '"params": [\n' + params + "\n" + indents + "],") + beautified_json = "\n".join(lines) + print(beautified_json) + + +def _read_dynamic_shape(shape: List[Union[int, str]], config: Union[Dict, ConfigBase]) -> List[int]: # noqa: UP006 + if isinstance(config, ConfigBase): + config = asdict(config) + param_shape = [] + for s in shape: + if isinstance(s, int): + param_shape.append(s) + else: + if config is None: + logger.error( + "%s: Encountered dynamic shape %s, need to specify `--mlc-chat-config` for " + + "memory usage calculation.", + red("FAILED"), + red(s), + ) + raise AttributeError + if s not in config: + logger.error( + "%s to retrieve concrete %s for dynamic shape from %s.", + red("FAILED"), + red(s), + config, + ) + raise KeyError + param_shape.append(config[s]) + return param_shape + + +def _compute_memory_usage(metadata: Dict[str, Any], config: Union[Dict, ConfigBase]): # noqa: UP006 + params_bytes = 0.0 + for param in metadata["params"]: + if all(isinstance(v, int) for v in param["shape"]): + assert all(v > 0 for v in param["shape"]), "All shapes should be strictly positive." + param_shape = param["shape"] + else: + # Contains dynamic shape; use config to look up concrete values + param_shape = _read_dynamic_shape(param["shape"], config) + params_bytes += math.prod(param_shape) * DataType(param["dtype"]).itemsize + temp_func_bytes = 0.0 + for _func_name, func_bytes in metadata["memory_usage"].items(): + temp_func_bytes = max(temp_func_bytes, func_bytes) + + return params_bytes, temp_func_bytes + + +def _report_memory_usage(metadata: Dict[str, Any], config: Union[Dict, ConfigBase]) -> None: # noqa: UP006 + params_bytes, temp_func_bytes = _compute_memory_usage(metadata, config) + total_size = params_bytes + temp_func_bytes + logger.info( + "%s: %.2f MB (Parameters: %.2f MB. Temporary buffer: %.2f MB)", + green("Total memory usage without KV cache"), + total_size / 1024 / 1024, + params_bytes / 1024 / 1024, + temp_func_bytes / 1024 / 1024, + ) + + # Compute KV cache size per token of context window. + if isinstance(config, ConfigBase): + config = asdict(config) + if ( + "head_dim" in config + and "num_hidden_layers" in config + and "num_key_value_heads" in config + and "quantization" in metadata + ): + quantization_type = metadata["quantization"] + dtype_bytes = None + if "f32" in quantization_type: + dtype_bytes = 4 + elif "bf16" in quantization_type: + dtype_bytes = 2 + elif "f16" in quantization_type: + dtype_bytes = 2 + # TODO: If support quantized KV in future, need to change this + if dtype_bytes is not None: + bytes_per_token = ( + config["head_dim"] + * config["num_hidden_layers"] + * config["num_key_value_heads"] + * dtype_bytes + * 2 # 2 for key and value + ) + logger.info( + "%s: %.2f MB per token in the context window", + green("KV cache size"), + bytes_per_token / 1024 / 1024, + ) + logger.info( + "%s: %.2f MB", + green("Total memory usage with a 4K KV cache"), + (total_size + bytes_per_token * 4096) / 1024 / 1024, + ) + + logger.info( + "To reduce memory usage, " + "tweak `prefill_chunk_size`, `context_window_size` and `sliding_window_size`" + ) + + +def main(): + """Entry point for the model metadata tool.""" + parser = ArgumentParser(description="A tool that inspects the metadata of a model lib.") + parser.add_argument( + "model_lib", + type=Path, + help="""The compiled model library. In MLC LLM, an LLM is compiled to a shared or static + library (.so or .a), which contains GPU computation to efficiently run the LLM. MLC Chat, + as the runtime of MLC LLM, depends on the compiled model library to generate tokens. + """, + ) + parser.add_argument( + "--mlc-chat-config", + type=Path, + help="""The `mlc-chat-config.json` file specific to a model variant. This is only required + when `memory-only` is true and `model_lib` contains a dynamic parameter shape (i.e. using + a variable to represent the shape). For instance, `model.embed_tokens.q_weight` can have + shape `["vocab_size", 512]`. In these cases, we look up the concrete value in + `mlc-chat-config.json`. + """, + ) + parser.add_argument( + "--memory-only", + action="store_true", + help="""If set, only inspect the metadata in memory usage and print richer analysis. + Otherwise, the tool will load all the metadata from the model library file but only print + the basic information in JSON. + """, + ) + parsed = parser.parse_args() + # Load metadata from model lib + try: + metadata = _extract_metadata(parsed.model_lib) + except Exception: + logger.exception("%s to read metadata section in legacy model lib.", red("FAILED")) + return + # Load mlc_chat_config if provided + cfg = None + if parsed.mlc_chat_config: + mlc_chat_config_path = Path(parsed.mlc_chat_config) + if not mlc_chat_config_path.exists(): + raise ValueError(f"{mlc_chat_config_path} does not exist.") + with open(mlc_chat_config_path, encoding="utf-8") as config_file: + cfg = json.load(config_file) + # Main body + if parsed.memory_only: + _report_memory_usage(metadata, cfg) + else: + _report_all(metadata) + + +if __name__ == "__main__": + main() diff --git a/python/mlc_llm/cli/package.py b/python/mlc_llm/cli/package.py new file mode 100644 index 0000000..9628d51 --- /dev/null +++ b/python/mlc_llm/cli/package.py @@ -0,0 +1,68 @@ +"""Command line entrypoint of package.""" + +import os +from pathlib import Path +from typing import Union + +from mlc_llm.interface.help import HELP +from mlc_llm.interface.package import package +from mlc_llm.support.argparse import ArgumentParser + + +def main(argv): + """Parse command line arguments and call `mlc_llm.interface.package`.""" + parser = ArgumentParser("MLC LLM Package CLI") + + def _parse_package_config(path: Union[str, Path]) -> Path: + path = Path(path) + if not path.exists(): + raise ValueError( + f"Path {str(path)} is expected to be a JSON file, but the file does not exist." + ) + if not path.is_file(): + raise ValueError(f"Path {str(path)} is expected to be a JSON file.") + return path + + def _parse_mlc_llm_source_dir(path: str) -> Path: + os.environ["MLC_LLM_SOURCE_DIR"] = path + return Path(path) + + def _parse_output(path: Union[str, Path]) -> Path: + path = Path(path) + if not path.is_dir(): + path.mkdir(parents=True, exist_ok=True) + return path + + parser.add_argument( + "--package-config", + type=_parse_package_config, + default="mlc-package-config.json", + help=HELP["config_package"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--mlc-llm-source-dir", + type=_parse_mlc_llm_source_dir, + default=os.environ.get("MLC_LLM_SOURCE_DIR", None), + help=HELP["mlc_llm_source_dir"] + + " (default: the $MLC_LLM_SOURCE_DIR environment variable)", + ) + parser.add_argument( + "--output", + "-o", + type=_parse_output, + default="dist", + help=HELP["output_package"] + ' (default: "%(default)s")', + ) + parsed = parser.parse_args(argv) + if parsed.mlc_llm_source_dir is None: + raise ValueError( + "MLC LLM home is not specified. " + "Please obtain a copy of MLC LLM source code by " + "cloning https://github.com/mlc-ai/mlc-llm, and set environment variable " + '"MLC_LLM_SOURCE_DIR=path/to/mlc-llm"' + ) + package( + package_config_path=parsed.package_config, + mlc_llm_source_dir=parsed.mlc_llm_source_dir, + output=parsed.output, + ) diff --git a/python/mlc_llm/cli/router.py b/python/mlc_llm/cli/router.py new file mode 100644 index 0000000..a771b15 --- /dev/null +++ b/python/mlc_llm/cli/router.py @@ -0,0 +1,90 @@ +"""Command line entrypoint of router.""" + +from mlc_llm.interface.help import HELP +from mlc_llm.interface.router import serve +from mlc_llm.support.argparse import ArgumentParser + + +def main(argv): + """Parse command line arguments and call `mlc_llm.interface.router`.""" + + # Define a custom argument type for a list of strings + def list_of_strings(arg): + return arg.split(",") + + parser = ArgumentParser("MLC LLM Router Serve CLI") + parser.add_argument( + "model", + type=str, + help=HELP["model"] + " (required)", + ) + parser.add_argument( + "--model-lib", + type=str, + default=None, + help=HELP["model_lib"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--router-mode", + type=str, + choices=["disagg", "round-robin"], + default="disagg", + help="router mode" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--router-host", + type=str, + default="127.0.0.1", + help="router host" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--router-port", + type=int, + default=8000, + help="router port" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--endpoint-hosts", + type=list_of_strings, + default="127.0.0.1", + help="Host of each endpoint, separated by comma." + ' (default: "%(default)s")', + ) + parser.add_argument( + "--endpoint-ports", + nargs="*", + type=int, + default=[8080], + help="Port of each endpoint, separated by space." + ' (default: "%(default)s")', + ) + parser.add_argument( + "--endpoint-num-gpus", + nargs="*", + type=int, + default=[1], + help="Number of GPUs of each endpoint, separated by space." + ' (default: "%(default)s")', + ) + parser.add_argument( + "--enable-prefix-cache", + default=False, + action="store_true", + help="whether to enable prefix cache" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--pd-balance-factor", + type=float, + default=0.0, + help=HELP["pd_balance_factor"] + ' (default: "%(default)s")', + ) + parsed = parser.parse_args(argv) + serve( + model=parsed.model, + model_lib=parsed.model_lib, + router_host=parsed.router_host, + router_port=parsed.router_port, + endpoint_hosts=parsed.endpoint_hosts, + endpoint_ports=parsed.endpoint_ports, + endpoint_num_gpus=parsed.endpoint_num_gpus, + enable_prefix_cache=parsed.enable_prefix_cache, + router_mode=parsed.router_mode, + pd_balance_factor=parsed.pd_balance_factor, + ) diff --git a/python/mlc_llm/cli/serve.py b/python/mlc_llm/cli/serve.py new file mode 100644 index 0000000..6e51efc --- /dev/null +++ b/python/mlc_llm/cli/serve.py @@ -0,0 +1,264 @@ +"""Command line entrypoint of serve.""" + +import dataclasses +import json +from io import StringIO +from typing import Literal, Optional + +from mlc_llm.interface.help import HELP +from mlc_llm.interface.serve import serve +from mlc_llm.support import argparse +from mlc_llm.support.argparse import ArgumentParser + + +@dataclasses.dataclass +class EngineConfigOverride: + """Arguments for overriding engine config.""" + + # Overrides for EngineConfig (runtime) + max_num_sequence: Optional[int] = None + max_total_seq_length: Optional[int] = None + prefill_chunk_size: Optional[int] = None + max_history_size: Optional[int] = None + gpu_memory_utilization: Optional[float] = None + spec_draft_length: Optional[int] = None + spec_tree_width: Optional[int] = None + prefix_cache_mode: Optional[Literal["disable", "radix"]] = None + prefix_cache_max_num_recycling_seqs: Optional[int] = None + prefill_mode: Optional[Literal["chunked", "hybrid"]] = None + context_window_size: Optional[int] = None + sliding_window_size: Optional[int] = None + attention_sink_size: Optional[int] = None + tensor_parallel_shards: Optional[int] = None + pipeline_parallel_stages: Optional[int] = None + opt: Optional[str] = None + + def __repr__(self) -> str: + out = StringIO() + print(f"max_num_sequence={self.max_num_sequence}", file=out, end="") + print(f";max_total_seq_length={self.max_total_seq_length}", file=out, end="") + print(f";prefill_chunk_size={self.prefill_chunk_size}", file=out, end="") + print(f";max_history_size={self.max_history_size}", file=out, end="") + print(f";gpu_memory_utilization={self.gpu_memory_utilization}", file=out, end="") + print(f";spec_draft_length={self.spec_draft_length}", file=out, end="") + print(f";spec_tree_width={self.spec_tree_width}", file=out, end="") + print(f";prefix_cache_mode={self.prefix_cache_mode}", file=out, end="") + print( + f";prefix_cache_max_num_recycling_seqs={self.prefix_cache_max_num_recycling_seqs}", + file=out, + end="", + ) + print(f";prefill_mode={self.prefill_mode}", file=out, end="") + print(f";context_window_size={self.context_window_size}", file=out, end="") + print(f";sliding_window_size={self.sliding_window_size}", file=out, end="") + print(f";attention_sink_size={self.attention_sink_size}", file=out, end="") + print(f";tensor_parallel_shards={self.tensor_parallel_shards}", file=out, end="") + print( + f";pipeline_parallel_stages={self.pipeline_parallel_stages}", + file=out, + end="", + ) + print(f";opt={self.opt}", file=out, end="") + return out.getvalue().rstrip() + + @staticmethod + def from_str(source: str) -> "EngineConfigOverride": + """Parse engine config override values from a string.""" + parser = argparse.ArgumentParser(description="Engine config override values") + + parser.add_argument("--max_num_sequence", type=int, default=None) + parser.add_argument("--max_total_seq_length", type=int, default=None) + parser.add_argument("--prefill_chunk_size", type=int, default=None) + parser.add_argument("--max_history_size", type=int, default=None) + parser.add_argument("--gpu_memory_utilization", type=float, default=None) + parser.add_argument("--spec_draft_length", type=int, default=None) + parser.add_argument("--spec_tree_width", type=int, default=None) + parser.add_argument("--prefix_cache_mode", type=str, default="radix") + parser.add_argument("--prefix_cache_max_num_recycling_seqs", type=int, default=None) + parser.add_argument("--prefill_mode", type=str, default="hybrid") + parser.add_argument("--context_window_size", type=int, default=None) + parser.add_argument("--sliding_window_size", type=int, default=None) + parser.add_argument("--attention_sink_size", type=int, default=None) + parser.add_argument("--tensor_parallel_shards", type=int, default=None) + parser.add_argument("--pipeline_parallel_stages", type=int, default=None) + parser.add_argument("--opt", type=str, default=None) + results = parser.parse_args([f"--{i}" for i in source.split(";") if i]) + return EngineConfigOverride( + max_num_sequence=results.max_num_sequence, + max_total_seq_length=results.max_total_seq_length, + prefill_chunk_size=results.prefill_chunk_size, + max_history_size=results.max_history_size, + gpu_memory_utilization=results.gpu_memory_utilization, + spec_draft_length=results.spec_draft_length, + spec_tree_width=results.spec_tree_width, + prefix_cache_mode=results.prefix_cache_mode, + prefix_cache_max_num_recycling_seqs=results.prefix_cache_max_num_recycling_seqs, + prefill_mode=results.prefill_mode, + context_window_size=results.context_window_size, + sliding_window_size=results.sliding_window_size, + attention_sink_size=results.attention_sink_size, + tensor_parallel_shards=results.tensor_parallel_shards, + pipeline_parallel_stages=results.pipeline_parallel_stages, + opt=results.opt, + ) + + +def main(argv): + """Parse command line arguments and call `mlc_llm.interface.serve`.""" + parser = ArgumentParser("MLC LLM Serve CLI") + + parser.add_argument( + "model", + type=str, + help=HELP["model"] + " (required)", + ) + parser.add_argument( + "--device", + type=str, + default="auto", + help=HELP["device_deploy"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--model-lib", + type=str, + default=None, + help=HELP["model_lib"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--mode", + type=str, + choices=["local", "interactive", "server"], + default="local", + help=HELP["mode_serve"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--enable-debug", + action="store_true", + help="whether we enable debug end points and debug config when accepting requests", + ) + parser.add_argument( + "--additional-models", type=str, nargs="*", help=HELP["additional_models_serve"] + ) + parser.add_argument( + "--embedding-model", + type=str, + default=None, + help="Path to the embedding model weight directory (enables /v1/embeddings endpoint)", + ) + parser.add_argument( + "--embedding-model-lib", + type=str, + default=None, + help="Path to the compiled embedding model library (.so/.dylib file)", + ) + parser.add_argument( + "--speculative-mode", + type=str, + choices=["disable", "small_draft", "eagle", "medusa"], + default="disable", + help=HELP["speculative_mode_serve"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--prefix-cache-mode", + type=str, + choices=["disable", "radix"], + default="radix", + help=HELP["prefix_cache_mode_serve"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--prefill-mode", + type=str, + choices=["hybrid", "chunked"], + default="hybrid", + help=HELP["prefill_mode"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--overrides", + type=EngineConfigOverride.from_str, + default="", + help=HELP["overrides_serve"], + ) + parser.add_argument("--enable-tracing", action="store_true", help=HELP["enable_tracing_serve"]) + parser.add_argument( + "--host", + type=str, + default="127.0.0.1", + help="host name" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--port", + type=int, + default=8000, + help="port" + ' (default: "%(default)s")', + ) + parser.add_argument("--allow-credentials", action="store_true", help="allow credentials") + parser.add_argument( + "--allow-origins", + type=json.loads, + default=["*"], + help="allowed origins" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--allow-methods", + type=json.loads, + default=["*"], + help="allowed methods" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--allow-headers", + type=json.loads, + default=["*"], + help="allowed headers" + ' (default: "%(default)s")', + ) + parser.add_argument( + "--api-key", + type=str, + default=None, + help="API key for authentication. If not provided, authentication is disabled.", + ) + parsed = parser.parse_args(argv) + + additional_models = [] + if parsed.additional_models is not None: + for additional_model in parsed.additional_models: + splits = additional_model.split(",", maxsplit=1) + if len(splits) == 2: + additional_models.append((splits[0], splits[1])) + else: + additional_models.append(splits[0]) + + serve( + model=parsed.model, + device=parsed.device, + model_lib=parsed.model_lib, + mode=parsed.mode, + enable_debug=parsed.enable_debug, + additional_models=additional_models, + embedding_model=parsed.embedding_model, + embedding_model_lib=parsed.embedding_model_lib, + tensor_parallel_shards=parsed.overrides.tensor_parallel_shards, + pipeline_parallel_stages=parsed.overrides.pipeline_parallel_stages, + opt=parsed.overrides.opt, + speculative_mode=parsed.speculative_mode, + prefix_cache_mode=parsed.prefix_cache_mode, + max_num_sequence=parsed.overrides.max_num_sequence, + max_total_sequence_length=parsed.overrides.max_total_seq_length, + max_single_sequence_length=parsed.overrides.context_window_size, + prefill_chunk_size=parsed.overrides.prefill_chunk_size, + sliding_window_size=parsed.overrides.sliding_window_size, + attention_sink_size=parsed.overrides.attention_sink_size, + max_history_size=parsed.overrides.max_history_size, + gpu_memory_utilization=parsed.overrides.gpu_memory_utilization, + spec_draft_length=parsed.overrides.spec_draft_length, + spec_tree_width=parsed.overrides.spec_tree_width, + prefix_cache_max_num_recycling_seqs=parsed.overrides.prefix_cache_max_num_recycling_seqs, + prefill_mode=parsed.prefill_mode, + enable_tracing=parsed.enable_tracing, + host=parsed.host, + port=parsed.port, + allow_credentials=parsed.allow_credentials, + allow_origins=parsed.allow_origins, + allow_methods=parsed.allow_methods, + allow_headers=parsed.allow_headers, + api_key=parsed.api_key, + ) diff --git a/python/mlc_llm/cli/worker.py b/python/mlc_llm/cli/worker.py new file mode 100644 index 0000000..e9068ad --- /dev/null +++ b/python/mlc_llm/cli/worker.py @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Internal DiscoWorker for Disco ProcessSession.""" + +import os +import sys + +from tvm import runtime as _ # noqa: F401 +from tvm_ffi import get_global_func + +from .. import base # noqa: F401 + +# register the calibration functions +from ..interface import calibrate # noqa: F401 + + +def main(): + """Main worker function""" + if len(sys.argv) != 6: + print("Usage: ") + return + + worker_id = int(sys.argv[1]) + num_workers = int(sys.argv[2]) + num_groups = int(sys.argv[3]) + read_fd = int(sys.argv[4]) + write_fd = int(sys.argv[5]) + if sys.platform == "win32": + import msvcrt + + reader = msvcrt.open_osfhandle(read_fd, os.O_BINARY) + writer = msvcrt.open_osfhandle(write_fd, os.O_BINARY) + else: + reader = read_fd + writer = write_fd + + worker_func = get_global_func("runtime.disco.WorkerProcess") + worker_func(worker_id, num_workers, num_groups, reader, writer) + + +if __name__ == "__main__": + try: + main() + except (OSError, KeyboardInterrupt): + pass diff --git a/python/mlc_llm/compiler_pass/__init__.py b/python/mlc_llm/compiler_pass/__init__.py new file mode 100644 index 0000000..23a5b25 --- /dev/null +++ b/python/mlc_llm/compiler_pass/__init__.py @@ -0,0 +1,3 @@ +"""Compiler passes used in MLC LLM.""" + +from . import pipeline as _pipeline diff --git a/python/mlc_llm/compiler_pass/attach_cuda_graph_alloc_init_func.py b/python/mlc_llm/compiler_pass/attach_cuda_graph_alloc_init_func.py new file mode 100644 index 0000000..d449a19 --- /dev/null +++ b/python/mlc_llm/compiler_pass/attach_cuda_graph_alloc_init_func.py @@ -0,0 +1,33 @@ +"""The pass that attaches an empty function for initialization.""" + +import tvm +from tvm import IRModule, relax + + +@tvm.transform.module_pass(opt_level=0, name="AttachCUDAGraphAllocInitFunc") +class AttachCUDAGraphAllocInitFunc: + """Attach an empty function for initialization.""" + + def __init__(self): + pass + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + bb = relax.BlockBuilder(mod) + alloc_func_gv = None + for gv, _ in mod.functions_items(): + if gv.name_hint.startswith("cuda_graph_alloc"): + assert alloc_func_gv is None + alloc_func_gv = gv + if alloc_func_gv is None: + return mod + + with bb.function("cuda_graph_alloc_init", []): + bb.emit_func_output( + relax.op.call_builtin_with_ctx( + "vm.builtin.cuda_graph.get_cached_alloc", + args=[alloc_func_gv, relax.prim_value(0)], + ty_args=relax.ObjectType(), + ) + ) + return bb.finalize() diff --git a/python/mlc_llm/compiler_pass/attach_embedding_allocator.py b/python/mlc_llm/compiler_pass/attach_embedding_allocator.py new file mode 100644 index 0000000..4735c51 --- /dev/null +++ b/python/mlc_llm/compiler_pass/attach_embedding_allocator.py @@ -0,0 +1,39 @@ +"""The pass that attaches embedding allocation function to the IRModule.""" + +from typing import Any, Dict # noqa: UP035 + +import tvm +from tvm import IRModule, relax + + +@tvm.transform.module_pass(opt_level=0, name="AttachAllocEmbeddingTensorFunc") +class AttachAllocEmbeddingTensorFunc: + """Attach embedding tensor allocation Relax function to IRModule.""" + + def __init__(self, metadata: Dict[str, Any]): # noqa: UP006 + self.metadata = metadata + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + embed_func = None + for gv, func in mod.functions_items(): + if gv.name_hint == "embed": + embed_func = func + + if embed_func is None: + return mod + + hidden_size = embed_func.ret_ty.shape[-1] + dtype = relax.DataTypeImm(embed_func.ret_ty.dtype.dtype) + bb = relax.BlockBuilder(mod) + with bb.function("alloc_embedding_tensor", []): + bb.emit_func_output( + bb.emit( + relax.op.builtin.alloc_tensor( + relax.ShapeExpr([self.metadata["prefill_chunk_size"], hidden_size]), + dtype, + runtime_device_index=0, + ) + ) + ) + return bb.finalize() diff --git a/python/mlc_llm/compiler_pass/attach_logit_processor.py b/python/mlc_llm/compiler_pass/attach_logit_processor.py new file mode 100644 index 0000000..f37783e --- /dev/null +++ b/python/mlc_llm/compiler_pass/attach_logit_processor.py @@ -0,0 +1,285 @@ +"""The pass that attaches logit processor functions to the IRModule.""" + +import tvm +from tvm import IRModule +from tvm.script import tirx as T + +from ..support.max_thread_check import ( + check_thread_limits, + get_max_num_threads_per_block, +) + + +@tvm.transform.module_pass(opt_level=0, name="AttachLogitProcessFunc") +class AttachLogitProcessFunc: + """Attach logit processing TIR functions to IRModule.""" + + def __init__(self, target: tvm.target.Target): + """Initializer. + + Parameters + ---------- + target : tvm.target.Target + The target of the model compilation. + """ + self.target = target + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + mod = mod.clone() + if self.target.kind.name == "llvm": + mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace_cpu() + mod["apply_penalty_inplace"] = _get_apply_penalty_inplace_cpu() + mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace_cpu() + else: + mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace(self.target) + mod["apply_penalty_inplace"] = _get_apply_penalty_inplace(self.target) + mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace(self.target) + return mod + + +def _get_apply_logit_bias_inplace_cpu(): + @T.prim_func(s_tir=True) + def _apply_logit_bias_inplace( + var_logits: T.handle, + var_pos2seq_id: T.handle, + var_token_ids: T.handle, + var_logit_bias: T.handle, + ) -> None: + """Function that applies logit bias in place.""" + T.func_attr( + { + "global_symbol": "apply_logit_bias_inplace", + "tirx.noalias": True, + "tirx.is_scheduled": True, + } + ) + batch_size = T.int32() + vocab_size = T.int32() + num_token = T.int32() + logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32") + # seq_ids + pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32") + token_ids = T.match_buffer(var_token_ids, (num_token,), "int32") + logit_bias = T.match_buffer(var_logit_bias, (num_token,), "float32") + + for i in range(num_token): + logits[pos2seq_id[i], token_ids[i]] += logit_bias[i] + + return _apply_logit_bias_inplace + + +def _get_apply_logit_bias_inplace(target: tvm.target.Target): + tx = 1024 # default + max_num_threads_per_block = get_max_num_threads_per_block(target) + tx = min(tx, max_num_threads_per_block) + check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1) + + @T.prim_func(s_tir=True) + def _apply_logit_bias_inplace( + var_logits: T.handle, + var_pos2seq_id: T.handle, + var_token_ids: T.handle, + var_logit_bias: T.handle, + ) -> None: + """Function that applies logit bias in place.""" + T.func_attr( + { + "global_symbol": "apply_logit_bias_inplace", + "tirx.noalias": True, + "tirx.is_scheduled": True, + } + ) + batch_size = T.int32() + vocab_size = T.int32() + num_token = T.int32() + logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32") + # seq_ids + pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32") + token_ids = T.match_buffer(var_token_ids, (num_token,), "int32") + logit_bias = T.match_buffer(var_logit_bias, (num_token,), "float32") + + for p0 in T.thread_binding(0, (num_token + tx - 1) // tx, "blockIdx.x"): + for p1 in T.thread_binding(0, tx, "threadIdx.x"): + with T.sblock("block"): + vp = T.axis.spatial(num_token, p0 * tx + p1) + T.where(p0 * tx + p1 < num_token) + logits[pos2seq_id[vp], token_ids[vp]] += logit_bias[vp] + + return _apply_logit_bias_inplace + + +def _get_apply_penalty_inplace_cpu(): + @T.prim_func(s_tir=True) + def _apply_penalty_inplace( + var_logits: T.handle, + var_seq_ids: T.handle, + var_pos2seq_id: T.handle, + var_token_ids: T.handle, + var_token_cnt: T.handle, + var_penalties: T.handle, + ) -> None: + """Function that applies penalties in place.""" + T.func_attr( + { + "global_symbol": "apply_penalty_inplace", + "tirx.noalias": True, + "tirx.is_scheduled": True, + } + ) + batch_size = T.int32() + vocab_size = T.int32() + num_token = T.int32() + num_seq = T.int32() + logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32") + seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32") + pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32") + token_ids = T.match_buffer(var_token_ids, (num_token,), "int32") + token_cnt = T.match_buffer(var_token_cnt, (num_token,), "int32") + penalties = T.match_buffer(var_penalties, (num_seq, 3), "float32") + + for token in T.serial(num_token): + with T.sblock("block"): + vp = T.axis.spatial(num_token, token) + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] -= ( + penalties[pos2seq_id[vp], 0] + token_cnt[vp] * penalties[pos2seq_id[vp], 1] + ) + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] = T.if_then_else( + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] < T.float32(0), + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] * penalties[pos2seq_id[vp], 2], + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] / penalties[pos2seq_id[vp], 2], + ) + + return _apply_penalty_inplace + + +def _get_apply_penalty_inplace(target: tvm.target.Target): + tx = 1024 # default + max_num_threads_per_block = get_max_num_threads_per_block(target) + tx = min(tx, max_num_threads_per_block) + check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1) + + @T.prim_func(s_tir=True) + def _apply_penalty_inplace( + var_logits: T.handle, + var_seq_ids: T.handle, + var_pos2seq_id: T.handle, + var_token_ids: T.handle, + var_token_cnt: T.handle, + var_penalties: T.handle, + ) -> None: + """Function that applies penalties in place.""" + T.func_attr( + { + "global_symbol": "apply_penalty_inplace", + "tirx.noalias": True, + "tirx.is_scheduled": True, + } + ) + batch_size = T.int32() + vocab_size = T.int32() + num_token = T.int32() + num_seq = T.int32() + logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32") + seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32") + pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32") + token_ids = T.match_buffer(var_token_ids, (num_token,), "int32") + token_cnt = T.match_buffer(var_token_cnt, (num_token,), "int32") + penalties = T.match_buffer(var_penalties, (num_seq, 3), "float32") + + for p0 in T.thread_binding(0, (num_token + tx - 1) // tx, "blockIdx.x"): + for p1 in T.thread_binding(0, tx, "threadIdx.x"): + with T.sblock("block"): + vp = T.axis.spatial(num_token, p0 * tx + p1) + T.where(p0 * tx + p1 < num_token) + # Penalties: (presence_penalty, frequency_penalty, repetition_penalty) + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] -= ( + penalties[pos2seq_id[vp], 0] + token_cnt[vp] * penalties[pos2seq_id[vp], 1] + ) + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] = T.if_then_else( + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] < T.float32(0), + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] + * penalties[pos2seq_id[vp], 2], + logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] + / penalties[pos2seq_id[vp], 2], + ) + + return _apply_penalty_inplace + + +def _get_apply_bitmask_inplace_cpu(): + @T.prim_func(s_tir=True) + def _apply_bitmask_inplace( + var_logits: T.handle, + var_seq_ids: T.handle, + var_bitmask: T.handle, + ) -> None: + """Function that applies vocabulary masking in place.""" + T.func_attr( + { + "global_symbol": "apply_bitmask_inplace", + "tirx.noalias": True, + "tirx.is_scheduled": True, + } + ) + batch_size = T.int32() + vocab_size = T.int32() + num_seq = T.int32() + logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32") + seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32") + bitmask = T.match_buffer(var_bitmask, (batch_size, (vocab_size + 31) // 32), "int32") + + for token in T.serial(num_seq * vocab_size): + with T.sblock("block"): + vs = T.axis.spatial(num_seq, (token) // vocab_size) + vv = T.axis.spatial(vocab_size, (token) % vocab_size) + + logits[seq_ids[vs], vv] = T.if_then_else( + (bitmask[seq_ids[vs], vv // 32] >> (vv % 32)) & 1 == 1, + logits[seq_ids[vs], vv], + T.min_value("float32"), + ) + + return _apply_bitmask_inplace + + +def _get_apply_bitmask_inplace(target: tvm.target.Target): + tx = 1024 # default + max_num_threads_per_block = get_max_num_threads_per_block(target) + tx = min(tx, max_num_threads_per_block) + check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1) + + @T.prim_func(s_tir=True) + def _apply_bitmask_inplace( + var_logits: T.handle, + var_seq_ids: T.handle, + var_bitmask: T.handle, + ) -> None: + """Function that applies vocabulary masking in place.""" + T.func_attr( + { + "global_symbol": "apply_bitmask_inplace", + "tirx.noalias": True, + "tirx.is_scheduled": True, + } + ) + batch_size = T.int32() + vocab_size = T.int32() + num_seq = T.int32() + logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32") + seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32") + bitmask = T.match_buffer(var_bitmask, (batch_size, (vocab_size + 31) // 32), "int32") + + for fused_s_v_0 in T.thread_binding(0, (num_seq * vocab_size + tx - 1) // tx, "blockIdx.x"): + for fused_s_v_1 in T.thread_binding(0, tx, "threadIdx.x"): + with T.sblock("block"): + vs = T.axis.spatial(num_seq, (fused_s_v_0 * tx + fused_s_v_1) // vocab_size) + vv = T.axis.spatial(vocab_size, (fused_s_v_0 * tx + fused_s_v_1) % vocab_size) + T.where(fused_s_v_0 * tx + fused_s_v_1 < num_seq * vocab_size) + logits[seq_ids[vs], vv] = T.if_then_else( + (bitmask[seq_ids[vs], vv // 32] >> (vv % 32)) & 1 == 1, + logits[seq_ids[vs], vv], + T.min_value("float32"), + ) + + return _apply_bitmask_inplace diff --git a/python/mlc_llm/compiler_pass/attach_sampler.py b/python/mlc_llm/compiler_pass/attach_sampler.py new file mode 100644 index 0000000..1fdf319 --- /dev/null +++ b/python/mlc_llm/compiler_pass/attach_sampler.py @@ -0,0 +1,390 @@ +"""The pass that attaches GPU sampler functions to the IRModule.""" + +from typing import Dict # noqa: UP035 + +import tvm +from tvm import IRModule, relax, te, tirx +from tvm.relax.frontend import nn +from tvm.script import tirx as T + +from mlc_llm.op.batch_spec_verify import batch_spec_verify +from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm + + +@tvm.transform.module_pass(opt_level=0, name="AttachGPUSamplingFunc") +class AttachGPUSamplingFunc: + """Attach GPU sampling functions to IRModule.""" + + def __init__(self, target: tvm.target.Target, variable_bounds: Dict[str, int]): # noqa: UP006 + # Specifically for RWKV workloads, which contains -1 max_seq_len + max_batch_size = variable_bounds["batch_size"] + self.variable_bounds = { + "batch_size": max_batch_size, + "num_samples": max_batch_size, + "num_positions": 6 * max_batch_size, + } + self.non_negative_var = ["vocab_size"] + self.target = target + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + target_kind = self.target.kind.name + if target_kind not in ["cuda", "vulkan", "metal", "webgpu"]: + # Only enable GPU sampling for CUDA, Vulkan, Metal, and WebGPU. + return mod + + bb = relax.BlockBuilder(mod) + if target_kind == "webgpu": + # Only attach functions that do not contain i8s for WebGPU + gv_names = [ + gv.name_hint + for gv in [ + _attach_argsort_func(bb), + _attach_sample_with_top_p(bb), + ] + ] + else: + gv_names = [ + gv.name_hint + for gv in [ + _attach_multinomial_sampling_func(bb), + _attach_argsort_func(bb), + _attach_sample_with_top_p(bb), + _attach_take_probs_func(bb), + _attach_batch_verifier(bb), + _attach_renormalize_by_top_p(bb, self.target), + ] + ] + + mod = bb.finalize() + for gv_name in gv_names: + mod[gv_name] = ( + mod[gv_name] + .with_attr("tir_var_upper_bound", self.variable_bounds) + .with_attr("tir_non_negative_var", self.non_negative_var) + ) + return mod + + +def _attach_multinomial_sampling_func(bb: relax.BlockBuilder): + batch_size = tirx.Var("batch_size", "int64") + num_samples = tirx.Var("num_samples", "int64") + vocab_size = tirx.Var("vocab_size", "int64") + probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32")) + uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32")) + sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32")) + with bb.function("multinomial_from_uniform", [probs, uniform_samples, sample_indices]): + with bb.dataflow(): + sample_shape = relax.ShapeExpr([num_samples, 1]) + probs_tensor = nn.wrap_nested(probs, name="probs") + uniform_samples_tensor = nn.wrap_nested( + relax.call_pure_packed( + "vm.builtin.reshape", + uniform_samples, + sample_shape, + ty_args=relax.TensorType(sample_shape, "float32"), + ), + name="uniform_samples", + ) + sample_indices_tensor = nn.wrap_nested( + relax.call_pure_packed( + "vm.builtin.reshape", + sample_indices, + sample_shape, + ty_args=relax.TensorType(sample_shape, "int32"), + ), + name="sample_indices", + ) + result_tensor = nn.multinomial_from_uniform( + probs_tensor, + uniform_samples_tensor, + sample_indices_tensor, + "int32", + name="nn_multinomial_from_uniform", + ) + result = bb.emit( + relax.call_pure_packed( + "vm.builtin.reshape", + result_tensor._expr, + sample_indices.ty.shape, + ty_args=sample_indices.ty, + ) + ) + output = bb.emit_output(result) + gv = bb.emit_func_output(output) + return gv + + +def _attach_argsort_func(bb: relax.BlockBuilder): + batch_size = tirx.Var("batch_size", "int64") + vocab_size = tirx.Var("vocab_size", "int64") + probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32")) + with bb.function("argsort_probs", [probs]): + with bb.dataflow(): + sorted_indices = bb.emit(relax.op.argsort(probs, descending=True, dtype="int32")) + sorted_values = bb.emit_te( + lambda unsorted_probs, sorted_indices: te.compute( + (batch_size, vocab_size), + lambda i, j: unsorted_probs[i, sorted_indices[i, j]], + name="take_sorted_probs", + ), + probs, + sorted_indices, + primfunc_name_hint="take_sorted_probs", + ) + output = bb.emit_output((sorted_values, sorted_indices)) + gv = bb.emit_func_output(output) + return gv + + +@T.prim_func(s_tir=True) +def full(var_result: T.handle, value: T.int32): + """The filling function for top k.""" + batch_size = T.int32() + result = T.match_buffer(var_result, (batch_size, 1), "int32") + for i in T.serial(batch_size): + with T.sblock("block"): + vi = T.axis.spatial(batch_size, i) + result[vi, 0] = value + + +def _attach_sample_with_top_p(bb: relax.BlockBuilder): + batch_size = tirx.Var("batch_size", "int64") + num_samples = tirx.Var("num_samples", "int64") + vocab_size = tirx.Var("vocab_size", "int64") + sorted_probs = relax.Var("sorted_probs", relax.TensorType((batch_size, vocab_size), "float32")) + sorted_indices = relax.Var( + "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32") + ) + uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32")) + sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32")) + top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32")) + + with bb.function( + "sample_with_top_p", + [sorted_probs, sorted_indices, uniform_samples, sample_indices, top_p], + ): + with bb.dataflow(): + sample_shape = relax.ShapeExpr([num_samples, 1]) + top_p_shape = relax.ShapeExpr([batch_size, 1]) + sorted_probs_tensor = nn.wrap_nested(sorted_probs, name="sorted_probs") + sorted_indices_tensor = nn.wrap_nested(sorted_indices, name="sorted_indices") + uniform_samples_tensor = nn.wrap_nested( + relax.call_pure_packed( + "vm.builtin.reshape", + uniform_samples, + sample_shape, + ty_args=relax.TensorType(sample_shape, "float32"), + ), + name="uniform_samples", + ) + sample_indices_tensor = nn.wrap_nested( + relax.call_pure_packed( + "vm.builtin.reshape", + sample_indices, + sample_shape, + ty_args=relax.TensorType(sample_shape, "int32"), + ), + name="sample_indices", + ) + top_p_tensor = nn.wrap_nested( + relax.call_pure_packed( + "vm.builtin.reshape", + top_p, + top_p_shape, + ty_args=relax.TensorType(top_p_shape, "float32"), + ), + name="sample_indices", + ) + top_k_tensor = nn.tensor_ir_op( + full, + name_hint="full", + args=[vocab_size], + out=nn.Tensor.placeholder( + [batch_size, 1], + "int32", + ), + ) + + result_tensor = nn.sample_top_p_top_k_from_sorted_prob( + sorted_probs_tensor, + sorted_indices_tensor, + top_p_tensor, + top_k_tensor, + uniform_samples_tensor, + sample_indices_tensor, + ) + result = bb.emit_output( + relax.call_pure_packed( + "vm.builtin.reshape", + result_tensor._expr, + sample_indices.ty.shape, + ty_args=sample_indices.ty, + ) + ) + gv = bb.emit_func_output(result) + return gv + + +def _attach_renormalize_by_top_p(bb: relax.BlockBuilder, target: tvm.target.Target): + batch_size = tirx.Var("batch_size", "int64") + vocab_size = tirx.Var("vocab_size", "int64") + num_pivots = 3 + probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32")) + top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32")) + init_pivots = relax.Var("init_pivots", relax.TensorType((batch_size, num_pivots), "float32")) + with bb.function("renormalize_by_top_p", [probs, top_p, init_pivots]): + with bb.dataflow(): + cutoff_output = bb.emit( + relax.call_tir( + bb.add_func(top_p_pivot(num_pivots, target), "top_p_pivot_cutoff"), + args=[probs, top_p, init_pivots], + out_ty=[top_p.ty, top_p.ty], + ) + ) + final_pivot = cutoff_output[0] + renorm_sum = cutoff_output[1] + renormalized_probs = bb.emit_output( + relax.call_tir( + bb.add_func(top_p_renorm(target), "top_p_renorm_after_cutoff"), + args=[probs, final_pivot, renorm_sum], + out_ty=probs.ty, + ) + ) + gv = bb.emit_func_output(renormalized_probs) + return gv + + +def _attach_take_probs_func(bb: relax.BlockBuilder): + batch_size = tirx.Var("batch_size", "int64") + num_samples = tirx.Var("num_samples", "int64") + num_positions = tirx.Var("num_positions", "int64") + vocab_size = tirx.Var("vocab_size", "int64") + unsorted_probs = relax.Var( + "unsorted_probs", relax.TensorType((batch_size, vocab_size), "float32") + ) + sorted_indices = relax.Var( + "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32") + ) + sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32")) + sampling_results = relax.Var("sampling_result", relax.TensorType((num_samples,), "int32")) + top_prob_offsets = relax.Var("lobprob_offsets", relax.TensorType((num_positions,), "int32")) + + @T.prim_func(s_tir=True) + def sampler_take_probs_tir( + var_unsorted_probs: T.handle, + var_sorted_indices: T.handle, + var_sample_indices: T.handle, + var_sampling_results: T.handle, + var_top_prob_offsets: T.handle, + var_sampled_values: T.handle, + var_top_prob_probs: T.handle, + var_top_prob_indices: T.handle, + ): + batch_size = T.int32() + num_samples = T.int32() + num_positions = T.int32() + vocab_size = T.int32() + unsorted_probs = T.match_buffer(var_unsorted_probs, (batch_size, vocab_size), "float32") + sorted_indices = T.match_buffer(var_sorted_indices, (batch_size, vocab_size), "int32") + sample_indices = T.match_buffer(var_sample_indices, (num_samples,), "int32") + sampling_results = T.match_buffer(var_sampling_results, (num_samples,), "int32") + top_prob_offsets = T.match_buffer(var_top_prob_offsets, (num_positions,), "int32") + sampled_values = T.match_buffer(var_sampled_values, (num_samples,), "float32") + top_prob_probs = T.match_buffer(var_top_prob_probs, (num_positions,), "float32") + top_prob_indices = T.match_buffer(var_top_prob_indices, (num_positions,), "int32") + for i in T.serial(num_positions): + with T.sblock("top_prob"): + vi = T.axis.spatial(num_positions, i) + # Reads are data-dependent gathers; declare full-buffer read + # regions explicitly so tirx does not infer data-dependent regions. + T.reads( + top_prob_offsets[vi], + sorted_indices[0:batch_size, 0:vocab_size], + unsorted_probs[0:batch_size, 0:vocab_size], + ) + T.writes(top_prob_indices[vi], top_prob_probs[vi]) + row = T.floordiv(top_prob_offsets[vi], vocab_size) + col = T.floormod(top_prob_offsets[vi], vocab_size) + top_prob_indices[vi] = sorted_indices[row, col] + top_prob_probs[vi] = unsorted_probs[row, sorted_indices[row, col]] + for i in T.serial(num_samples): + with T.sblock("sample"): + vj = T.axis.spatial(num_samples, i) + T.reads( + sample_indices[vj], + sampling_results[vj], + unsorted_probs[0:batch_size, 0:vocab_size], + ) + T.writes(sampled_values[vj]) + sampled_values[vj] = unsorted_probs[sample_indices[vj], sampling_results[vj]] + + args = [ + unsorted_probs, + sorted_indices, + sample_indices, + sampling_results, + top_prob_offsets, + ] + with bb.function("sampler_take_probs", args): + with bb.dataflow(): + taken_probs_indices = bb.emit_output( + relax.call_tir( + bb.add_func(sampler_take_probs_tir, "sampler_take_probs_tir"), + args, + out_ty=[ + relax.TensorType((num_samples,), "float32"), + relax.TensorType((num_positions,), "float32"), + relax.TensorType((num_positions,), "int32"), + ], + ) + ) + gv = bb.emit_func_output(taken_probs_indices) + return gv + + +def _attach_batch_verifier(bb: relax.BlockBuilder): + num_nodes = tirx.Var("num_nodes", "int64") + nbatch = tirx.Var("nbatch", "int64") + vocab_size = tirx.Var("vocab_size", "int64") + draft_probs = relax.Var("draft_probs", relax.TensorType((num_nodes, vocab_size), "float32")) + draft_tokens = relax.Var("draft_tokens", relax.TensorType((num_nodes,), "int32")) + model_probs = relax.Var("model_probs", relax.TensorType((num_nodes, vocab_size), "float32")) + token_tree_first_child = relax.Var( + "token_tree_first_child", relax.TensorType((num_nodes,), "int32") + ) + token_tree_next_sibling = relax.Var( + "token_tree_next_sibling", relax.TensorType((num_nodes,), "int32") + ) + uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_nodes,), "float32")) + token_tree_parent_ptr = relax.Var("token_tree_parent_ptr", relax.TensorType((nbatch,), "int32")) + args = [ + draft_probs, + draft_tokens, + model_probs, + token_tree_first_child, + token_tree_next_sibling, + uniform_samples, + token_tree_parent_ptr, + ] + with bb.function("sampler_verify_draft_tokens", args): + with bb.dataflow(): + res = bb.emit_output( + relax.call_tir_inplace( + bb.add_func( + batch_spec_verify(vocab_size), + "batch_verify_on_gpu_single_kernel", + ), + args, + inplace_indices=[ + args.index(model_probs), + args.index(token_tree_parent_ptr), + ], + out_ty=[ + model_probs.ty, + token_tree_parent_ptr.ty, + ], + ) + ) + gv = bb.emit_func_output(res) + return gv diff --git a/python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py b/python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py new file mode 100644 index 0000000..62a3652 --- /dev/null +++ b/python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py @@ -0,0 +1,274 @@ +"""A compiler pass that attaches two-stage softmax with temperature.""" + +from typing import Any, Dict, Optional # noqa: UP035 + +import tvm +from tvm import relax, tirx +from tvm.ir.module import IRModule +from tvm.relax.expr_functor import PyExprMutator, mutator +from tvm.script import tirx as T + +from ..support.max_thread_check import get_max_num_threads_per_block + + +@tvm.transform.module_pass(opt_level=0, name="AttachSoftmaxWithTemperature") +class AttachSoftmaxWithTemperature: + """Rewrites one-shot softmax into two-stage softmax.""" + + def __init__( + self, + target: tvm.target.Target, + metadata: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> None: + self.target = target + self.metadata = metadata + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """IRModule-level transformation""" + return _Rewriter(mod, self.target, self.metadata).transform() + + +@mutator +class _Rewriter(PyExprMutator): + def __init__( + self, + mod: IRModule, + target: tvm.target.Target, + metadata: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> None: + super().__init__(mod) + self.mod = mod + self.target = target + self.metadata = metadata + self.chunk_size = 4096 + self.active_vocab_size = self.metadata.get("active_vocab_size") if self.metadata else None + + def transform(self) -> IRModule: + """Entry point""" + batch_size = tirx.Var("batch_size", "int64") + vocab_size = tirx.Var("vocab_size", "int64") + dtype = "float32" + logits = relax.Var("logits", relax.TensorType([batch_size, 1, vocab_size], dtype)) + temperature = relax.Var("temperature", relax.TensorType([batch_size], dtype)) + with self.builder_.function("softmax_with_temperature", params=[logits, temperature]): + with self.builder_.dataflow(): + output_struct_info = logits.ty + new_shape = relax.ShapeExpr([batch_size, vocab_size]) + logits = relax.call_pure_packed( + "vm.builtin.reshape", + logits, + new_shape, + ty_args=relax.TensorType(new_shape, dtype), + ) + f_chunk_lse, f_softmax_with_lse = _get_lse_and_softmax_func( + self.target, self.chunk_size, self.active_vocab_size + ) + chunked_result_struct_info = relax.TensorType( + (batch_size, (vocab_size + self.chunk_size - 1) // self.chunk_size), + "float32", + ) + chunked_results = self.builder_.emit( + relax.call_tir( + self.builder_.add_func(f_chunk_lse, "chunk_lse"), + args=[logits, temperature], + out_ty=[ + chunked_result_struct_info, + chunked_result_struct_info, + ], + ) + ) + chunked_sum = chunked_results[0] + chunked_max = chunked_results[1] + softmax = self.builder_.emit( + relax.call_tir( + self.builder_.add_func(f_softmax_with_lse, "softmax_with_chunked_sum"), + args=[logits, temperature, chunked_sum, chunked_max], + out_ty=logits.ty, + ) + ) + softmax = self.builder_.emit_output( + relax.call_pure_packed( + "vm.builtin.reshape", + softmax, + output_struct_info.shape, + ty_args=output_struct_info, + ) + ) + self.builder_.emit_func_output(softmax) + return self.builder_.get() + + +def _get_lse_and_softmax_func(target: tvm.target.Target, chunk_size: int, active_vocab_size: int): + # NOTE: A quick note on the softmax implementation. + # We once tried to multiply every element by log2e which can be computed + # potentially more efficiently on hardware. + # However, when the input values are large, multiplying by the factor of log2e + # causes numerical issue in float32 dtype. + # This leads to the softmax output not summing up to 1. + # For numerical stability, we removed the log2e factor and switched back + # to the standard log/exp computation. + + # The kernels below handle both the cases of temperature=0 and temperature != 0. + # - When temperature is not 0, the first kernel computes the log-sum-exp of + # chunks (subtracted by the max value in chunk), and the max values of chunks. + # The second kernel merges the log-sum-exp with the maximum values. + # - When temperature is 0, the first kernel computes the max value and the counts + # of the max value. The second kernel merges the max and counts, and set the + # softmax of the maximum values to "max_value / max_count". + + @T.prim_func(s_tir=True) + def chunk_lse( + var_A: T.handle, + var_temperature: T.handle, + var_chunked_sum: T.handle, + var_chunked_max: T.handle, + ): + T.func_attr({"tirx.noalias": True}) + batch_size = T.int64() + vocab_size = T.int64() + num_chunks = T.int64() + A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32") + temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32") + chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32") + chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32") + A_pad = T.sblock_alloc_buffer( + (batch_size, num_chunks, T.int64(chunk_size)), dtype="float32" + ) + temp_max = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32") + temp_sum = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32") + + for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)): + with T.sblock("pad"): + v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2]) + A_pad[v0, v1, v2] = T.Select( + v1 * T.int64(chunk_size) + v2 + < (active_vocab_size if active_vocab_size is not None else vocab_size), + T.if_then_else( + temperature[v0] > T.float32(1e-5), + A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0], + A[v0, v1 * T.int64(chunk_size) + v2], + ), + T.min_value("float32"), + ) + for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)): + with T.sblock("max"): + v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2]) + with T.init(): + temp_max[v0, v1] = T.min_value("float32") + temp_max[v0, v1] = T.max(temp_max[v0, v1], A_pad[v0, v1, v2]) + for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)): + with T.sblock("sum_exp"): + v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2]) + with T.init(): + temp_sum[v0, v1] = T.float32(0) + temp_sum[v0, v1] += T.if_then_else( + v1 * T.int64(chunk_size) + v2 + < (active_vocab_size if active_vocab_size is not None else vocab_size), + T.Select( + temperature[v0] > T.float32(1e-5), + T.exp(A_pad[v0, v1, v2] - temp_max[v0, v1]), + T.cast(A_pad[v0, v1, v2] == temp_max[v0, v1], "float32"), + ), + T.float32(0), + ) + for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(1)): + with T.sblock("log"): + v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2]) + chunked_sum[v0, v1] = T.Select( + temperature[v0] > T.float32(1e-5), + T.log(temp_sum[v0, v1]), + temp_sum[v0, v1], + ) + chunked_max[v0, v1] = temp_max[v0, v1] + + @T.prim_func(s_tir=True) + def softmax_with_chunked_sum( + var_A: T.handle, + var_temperature: T.handle, + var_chunked_sum: T.handle, + var_chunked_max: T.handle, + var_softmax: T.handle, + ): + T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1}) + batch_size = T.int64() + vocab_size = T.int64() + num_chunks = T.int64() + A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32") + temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32") + chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32") + chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32") + softmax = T.match_buffer(var_softmax, (batch_size, vocab_size), dtype="float32") + temp_max = T.sblock_alloc_buffer((batch_size,), dtype="float32") + temp_sum = T.sblock_alloc_buffer((batch_size,), dtype="float32") + for l0, l1 in T.grid(batch_size, num_chunks): + with T.sblock("max"): + v0, v1 = T.axis.remap("SR", [l0, l1]) + with T.init(): + temp_max[v0] = T.min_value("float32") + temp_max[v0] = T.max(temp_max[v0], chunked_max[v0, v1]) + for l0, l1 in T.grid(batch_size, num_chunks): + with T.sblock("sum_exp"): + v0, v1 = T.axis.remap("SR", [l0, l1]) + with T.init(): + temp_sum[v0] = T.float32(0) + temp_sum[v0] += T.Select( + temperature[v0] > T.float32(1e-5), + T.exp(chunked_sum[v0, v1] + chunked_max[v0, v1] - temp_max[v0]), + T.cast(chunked_max[v0, v1] == temp_max[v0], "float32") * chunked_sum[v0, v1], + ) + for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)): + with T.sblock("log_pad"): + v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2]) + if v1 * T.int64(chunk_size) + v2 < vocab_size: + softmax[v0, v1 * T.int64(chunk_size) + v2] = T.Select( + v1 * T.int64(chunk_size) + v2 + < (active_vocab_size if active_vocab_size is not None else vocab_size), + T.if_then_else( + temperature[v0] > T.float32(1e-5), + T.exp( + A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0] + - (T.log(temp_sum[v0]) + temp_max[v0]) + ), + T.cast( + A[v0, v1 * T.int64(chunk_size) + v2] == temp_max[v0], + "float32", + ) + / temp_sum[v0], + ), + T.float32(0), + ) + + sch = tvm.s_tir.Schedule(IRModule({"softmax_with_chunked_sum": softmax_with_chunked_sum})) + + def apply_gpu_schedule(target, sch): + max_threads = get_max_num_threads_per_block(target) + TX = 32 + TY = max_threads // TX + unroll_depth = 64 + + sch.work_on("softmax_with_chunked_sum") + l0, l1, l2 = sch.get_loops("log_pad") + bx = sch.fuse(l0, l1) + sch.bind(bx, "blockIdx.x") + unroll, ty, tx = sch.split(l2, [None, TY, TX]) + sch.bind(ty, "threadIdx.y") + sch.bind(tx, "threadIdx.x") + sch.annotate(unroll, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth) + sch.annotate(unroll, ann_key="pragma_unroll_explicit", ann_val=1) + + for block_name in ["sum_exp", "max"]: + block = sch.get_sblock(block_name) + sch.set_scope(block, buffer_index=0, storage_scope="shared") + sch.compute_at(block, bx) + r_loop = sch.get_loops(block)[-1] + r_loop, tx = sch.split(r_loop, [None, TX]) + sch.reorder(tx, r_loop) + sch.bind(tx, "threadIdx.x") + sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth) + sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1) + + return chunk_lse, sch.mod["softmax_with_chunked_sum"] + + if target.kind.name == "llvm": + return chunk_lse, sch.mod["softmax_with_chunked_sum"] + return apply_gpu_schedule(target, sch) diff --git a/python/mlc_llm/compiler_pass/attach_spec_decode_aux_funcs.py b/python/mlc_llm/compiler_pass/attach_spec_decode_aux_funcs.py new file mode 100644 index 0000000..a9831f2 --- /dev/null +++ b/python/mlc_llm/compiler_pass/attach_spec_decode_aux_funcs.py @@ -0,0 +1,123 @@ +"""The pass that attaches logit processor functions to the IRModule.""" + +import tvm +from tvm import IRModule, relax, tirx +from tvm.relax import BlockBuilder, TensorType +from tvm.script import tirx as T + + +@tvm.transform.module_pass(opt_level=0, name="AttachSpecDecodeAuxFuncs") +class AttachSpecDecodeAuxFuncs: + """Attach logit processing TIR functions to IRModule.""" + + tensor_parallel_shards: int + + def __init__(self, tensor_parallel_shards: int): + self.tensor_parallel_shards = tensor_parallel_shards + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + mod = mod.clone() + bb = BlockBuilder(mod) + bb.add_func( + _get_scatter_2d_inplace(dtype="float32", global_symbol="scatter_probs"), + "scatter_probs", + ) + bb.add_func( + _get_gather_2d_inplace(dtype="float32", global_symbol="gather_probs"), + "gather_probs", + ) + if "prefill_to_last_hidden_states" in mod: + hidden_states_struct_info = mod["prefill_to_last_hidden_states"].ret_ty.fields[0] + dtype = hidden_states_struct_info.dtype + _add_gather_hidden_states(bb, self.tensor_parallel_shards, dtype) + _add_scatter_hidden_states(bb, self.tensor_parallel_shards, dtype) + return bb.finalize() + + +def _get_scatter_2d_inplace(dtype: str, global_symbol: str): + @T.prim_func(s_tir=True) + def _scatter_2d(var_src: T.handle, var_indices: T.handle, var_dst: T.handle): + T.func_attr({"global_symbol": global_symbol, "tirx.noalias": True}) + batch_size = T.int32() + m = T.int32() + n = T.int32() + src = T.match_buffer(var_src, (batch_size, n), dtype) + indices = T.match_buffer(var_indices, (batch_size,), "int32") + dst = T.match_buffer(var_dst, (m, n), dtype) + for b, j in T.grid(batch_size, n): + with T.sblock("scatter_2d"): + vb, vj = T.axis.remap("SS", [b, j]) + dst[indices[vb], vj] = src[vb, vj] + + return _scatter_2d + + +def _get_gather_2d_inplace(dtype: str, global_symbol: str): + @T.prim_func(s_tir=True) + def _gather_2d(var_src: T.handle, var_indices: T.handle, var_dst: T.handle): + T.func_attr({"global_symbol": global_symbol, "tirx.noalias": True}) + batch_size = T.int32() + m = T.int32() + n = T.int32() + src = T.match_buffer(var_src, (m, n), dtype) + indices = T.match_buffer(var_indices, (batch_size,), "int32") + dst = T.match_buffer(var_dst, (batch_size, n), dtype) + for b, j in T.grid(batch_size, n): + with T.sblock("gather_2d"): + vb, vj = T.axis.remap("SS", [b, j]) + dst[vb, vj] = src[indices[vb], vj] + + return _gather_2d + + +def _add_scatter_hidden_states(bb: BlockBuilder, tensor_parallel_shards: int, dtype: str): + batch_size = tirx.Var("batch_size", "int64") + m = tirx.Var("m", "int64") + n = tirx.Var("n", "int64") + src = relax.Var("src", ty=TensorType([batch_size, n], dtype)) + indices = relax.Var("indices", ty=TensorType([batch_size], "int32")) + dst = relax.Var("dst", ty=TensorType([m, n], dtype)) + with bb.function("scatter_hidden_states", [src, indices, dst]): + with bb.dataflow(): + if tensor_parallel_shards > 1: + indices = relax.op.ccl.broadcast_from_worker0(indices) + output = bb.emit_output( + relax.op.call_tir_inplace( + bb.add_func( + _get_scatter_2d_inplace(dtype, "_scatter_hidden_states"), + "_scatter_hidden_states", + ), + [src, indices, dst], + 2, + dst.ty, + ) + ) + gv = bb.emit_func_output(output) + return gv + + +def _add_gather_hidden_states(bb: BlockBuilder, tensor_parallel_shards: int, dtype: str): + batch_size = tirx.Var("batch_size", "int64") + m = tirx.Var("m", "int64") + n = tirx.Var("n", "int64") + src = relax.Var("src", ty=TensorType([m, n], dtype)) + indices = relax.Var("indices", ty=TensorType([batch_size], "int32")) + dst = relax.Var("dst", ty=TensorType([batch_size, n], dtype)) + with bb.function("gather_hidden_states", [src, indices, dst]): + with bb.dataflow(): + if tensor_parallel_shards > 1: + indices = relax.op.ccl.broadcast_from_worker0(indices) + output = bb.emit_output( + relax.op.call_tir_inplace( + bb.add_func( + _get_gather_2d_inplace(dtype, "_gather_hidden_states"), + "_gather_hidden_states", + ), + [src, indices, dst], + 2, + dst.ty, + ) + ) + gv = bb.emit_func_output(output) + return gv diff --git a/python/mlc_llm/compiler_pass/attach_support_info.py b/python/mlc_llm/compiler_pass/attach_support_info.py new file mode 100644 index 0000000..62a0cd5 --- /dev/null +++ b/python/mlc_llm/compiler_pass/attach_support_info.py @@ -0,0 +1,154 @@ +"""A couple of passes that simply supportive information onto the IRModule.""" + +from math import lcm +from typing import Any, Dict, List # noqa: UP035 + +import tvm +from tvm import IRModule, relax, tirx +from tvm.ir import Op +from tvm.relax.expr_functor import PyExprVisitor, visitor + + +@tvm.transform.module_pass(opt_level=0, name="AttachVariableBounds") +class AttachVariableBounds: + """Attach variable bounds to each Relax function, which primarily helps with memory planning.""" + + def __init__(self, variable_bounds: Dict[str, int]): # noqa: UP006 + # Specifically for RWKV workloads, which contains -1 max_seq_len + self.variable_bounds = {k: v for k, v in variable_bounds.items() if v > 0} + self.non_negative_var = ["vocab_size"] + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + for g_var, func in mod.functions_items(): + if isinstance(func, relax.Function): + mod[g_var] = func.with_attr("tir_var_upper_bound", self.variable_bounds).with_attr( + "tir_non_negative_var", self.non_negative_var + ) + return mod + + +@tvm.transform.module_pass(opt_level=0, name="AttachAdditionalPrimFuncs") +class AttachAdditionalPrimFuncs: + """Attach extra TIR PrimFuncs to the IRModule""" + + def __init__(self, functions: Dict[str, tirx.PrimFunc]): # noqa: UP006 + self.functions = functions + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + for func_name, func in self.functions.items(): + mod[func_name] = func.with_attr("global_symbol", func_name) + return mod + + +@tvm.transform.module_pass(opt_level=0, name="AttachMemoryPlanAttr") +class AttachMemoryPlanAttr: + """Attach memory planning attribute for dynamic function output planning to Relax functions.""" + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + for g_var, func in mod.functions_items(): + if isinstance(func, relax.Function): + mod[g_var] = func.with_attr("relax.memory_plan_dynamic_func_output", True) + return mod + + +@tvm.transform.module_pass(opt_level=0, name="AttachCUDAGraphCaptureHints") +class AttachCUDAGraphSymbolicCaptureHints: + """Attach CUDA graph capture hints to the IRModule""" + + def __init__(self, hints: Dict[str, List[str]]): # noqa: UP006 + self.hints = hints + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + for g_var, func in mod.functions_items(): + func_name = g_var.name_hint + if isinstance(func, relax.Function): + if func_name in self.hints: + mod[g_var] = func.with_attr( + "relax.rewrite_cuda_graph.capture_symbolic_vars", + self.hints[func_name], + ) + + return mod + + +@tvm.transform.module_pass(opt_level=0, name="AttachPipelineParallelStages") +class AttachPipelineParallelStages: + """Attach number of pipeline stages to relax functions.""" + + def __init__(self, pipeline_parallel_shards: int): + self.pipeline_parallel_shards = pipeline_parallel_shards + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + for g_var, func in mod.functions_items(): + func_name = g_var.name_hint + if not isinstance(func, relax.Function) or func_name not in [ + "prefill", + "decode", + "prefill_to_last_hidden_states", + "decode_to_last_hidden_states", + "batch_prefill", + "batch_decode", + "batch_verify", + "batch_prefill_to_last_hidden_states", + "batch_decode_to_last_hidden_states", + "batch_verify_to_last_hidden_states", + ]: + continue + mod[g_var] = func.with_attr("pipeline_parallel_stages", self.pipeline_parallel_shards) + + return mod + + +@tvm.transform.module_pass(opt_level=0, name="AttachSequenceLengthPaddingFactor") +class AttachSequenceLengthPaddingFactor: + """Attach sequence length padding factor to the metadata""" + + def __init__(self, target: tvm.target.Target, metadata: Dict[str, Any]): # noqa: UP006 + self.target = target + self.metadata = metadata + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + + @visitor + class _Visitor(PyExprVisitor): + def __init__(self, target: tvm.target.Target) -> None: + self.padding_factor = 1 + self.target = target + self._op_call_dps_packed = Op.get("relax.call_dps_packed") + + def run(self, mod: IRModule) -> int: + """Entry point of the visitor.""" + # Right now we only need padding for CUDA SM100a architecture. + # When the target is SM100a and uses cutlass gemm function, + # the sequence length needs to be padded to multiple of 4. + if self.target.kind.name != "cuda" or self.target.attrs.get("arch") != "sm_100a": + return 1 + + for _, func in mod.functions_items(): + if isinstance(func, relax.Function): + self.visit_expr(func) + return self.padding_factor + + def visit_call_(self, call: relax.Call) -> None: + super().visit_call_(call) + if call.op != self._op_call_dps_packed: + return + func_name = str(call.args[0].global_symbol) + if func_name in [ + "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn", + "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn", + ]: + # Find the minimum common multiple of padding factor and 4 + self.padding_factor = lcm(self.padding_factor, 4) + + # self.metadata["sequence_length_padding"] = True + padding_factor = _Visitor(self.target).run(mod) + if padding_factor > 1: + self.metadata["seqlen_padding_factor"] = padding_factor + return mod diff --git a/python/mlc_llm/compiler_pass/blas_dispatch.py b/python/mlc_llm/compiler_pass/blas_dispatch.py new file mode 100644 index 0000000..95348af --- /dev/null +++ b/python/mlc_llm/compiler_pass/blas_dispatch.py @@ -0,0 +1,51 @@ +"""A compiler pass that dispatches patterns to CUBLAS.""" + +import tvm +from tvm import IRModule, relax +from tvm.relax.backend import get_patterns_with_prefix + +try: + import tvm.relax.backend.cuda.cublas as _cublas # noqa: F401 + import tvm.relax.backend.rocm.hipblas as _hipblas # noqa: F401 +except ImportError: + # Note: legacy path of cublas/hipblas for backward compatibility + pass + + +@tvm.transform.module_pass(opt_level=0, name="BLASDispatch") +class BLASDispatch: + """A compiler pass that dispatches patterns to cuBLAS/hipBLAS.""" + + def __init__(self, target: tvm.target.Target) -> None: + if target.kind.name == "cuda": + self.has_blas = tvm.get_global_func("relax.ext.cublas", True) + if not self.has_blas: + raise Exception("cuBLAS is not enabled.") + self.patterns = get_patterns_with_prefix("cublas") + elif target.kind.name == "rocm": + self.has_blas = tvm.get_global_func("relax.ext.hipblas", True) + if not self.has_blas: + raise Exception("hipBLAS is not enabled.") + self.patterns = get_patterns_with_prefix("hipblas") + else: + raise Exception(f"Unsupported target {target.kind.name} for BLAS dispatch.") + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """IRModule-level transformation""" + model_names = [ + gv.name_hint for gv, func in mod.functions.items() if isinstance(func, relax.Function) + ] + # exclude single batch decode + model_names = [name for name in model_names if "batch" in name or "decode" not in name] + mod = tvm.transform.Sequential( + [ + relax.transform.FuseOpsByPattern( + self.patterns, + bind_constants=False, + annotate_codegen=True, + entry_functions=model_names, + ), + relax.transform.RunCodegen({}, entry_functions=model_names), + ] + )(mod) + return mod diff --git a/python/mlc_llm/compiler_pass/clean_up_tir_attrs.py b/python/mlc_llm/compiler_pass/clean_up_tir_attrs.py new file mode 100644 index 0000000..a8648ae --- /dev/null +++ b/python/mlc_llm/compiler_pass/clean_up_tir_attrs.py @@ -0,0 +1,31 @@ +"""A compiler pass that cleans up undesired TIR attrs.""" + +from typing import List # noqa: UP035 + +import tvm +from tvm.ir.module import IRModule + + +@tvm.transform.module_pass(opt_level=0, name="CleanUpTIRAttrs") +class CleanUpTIRAttrs: + """A compiler pass that cleans up undesired TIR attrs.""" + + def __init__(self, attrs: List[str]): # noqa: UP006 + self.attrs = attrs + + def transform_module( + self, + mod: IRModule, + _ctx: tvm.transform.PassContext, + ) -> IRModule: + """IRModule-level transformation""" + for g_var, func in mod.functions_items(): + changed = False + for attr in self.attrs: + if func.attrs is not None and attr in func.attrs: + func = func.without_attr(attr) + changed = True + break + if changed: + mod[g_var] = func + return mod diff --git a/python/mlc_llm/compiler_pass/dispatch_kv_cache_creation.py b/python/mlc_llm/compiler_pass/dispatch_kv_cache_creation.py new file mode 100644 index 0000000..7fea695 --- /dev/null +++ b/python/mlc_llm/compiler_pass/dispatch_kv_cache_creation.py @@ -0,0 +1,243 @@ +"""A pass that rewrites KV cache creation functions in IRModule.""" + +import json +from typing import Any, Dict, List # noqa: UP035 + +import tvm +from tvm import IRModule, relax +from tvm.relax.frontend.nn.llm import kv_cache +from tvm.relax.frontend.nn.llm.kv_cache import RopeMode + +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +def extract_creation_args(func: relax.Function) -> Dict[str, Any]: # noqa: UP006 + """Extract the KV cache creation args from the given generic creation func.""" + assert isinstance(func.body, relax.SeqExpr) + assert len(func.body.blocks) == 1 + assert isinstance(func.body.blocks[0], relax.DataflowBlock) + assert isinstance(func.body.blocks[0].bindings[0], relax.VarBinding) + assert isinstance(func.body.blocks[0].bindings[0].value, relax.Call) + assert func.body.blocks[0].bindings[0].value.op == tvm.ir.Op.get("relax.call_pure_packed") + call_args = func.body.blocks[0].bindings[0].value.args + assert isinstance(call_args[0], relax.ExternFunc) + assert call_args[0].global_symbol == "mlc.create_paged_kv_cache_generic" + args = call_args[1:] + assert len(args) == 18 + assert isinstance(args[0], (relax.StringImm, relax.Tuple)) + # Check if attn_kind is a single value or a list with length of hidden layers + if isinstance(args[0], relax.StringImm): + assert args[0].value in ["mha", "mla"] + attn_kind = args[0].value + else: + assert len(args[0].fields) == args[3].value + for i, attention_type in enumerate(args[0].fields): + assert isinstance(attention_type, relax.StringImm) + assert attention_type.value in ["mha", "mla", "mha_sliding"] + attn_kind = [args[0].fields[i].value for i in range(len(args[0]))] + assert isinstance(args[1], relax.ShapeExpr) + assert len(args[1].values) == 5 + assert isinstance(args[2], relax.ShapeExpr) + for i in range(3, 18): + if i in [13, 14, 17]: + continue + # PrimValue wrappers were phased out of Relax: scalar args are now bare + # tirx PrimExprs (IntImm/FloatImm) directly. + assert isinstance(args[i], (tvm.tirx.IntImm, tvm.tirx.FloatImm)), ( + f"args[{i}] is {type(args[i])}" + ) + assert isinstance(args[13], relax.StringImm) + assert isinstance(args[16], (relax.Constant, tvm.tirx.IntImm, tvm.tirx.FloatImm)) + assert isinstance(args[17], relax.DataTypeImm) + + return { + "attn_kind": attn_kind, + "max_batch_size": args[1].values[0], + "max_total_seq_len": args[1].values[1], + "prefill_chunk_size": args[1].values[2], + "page_size": args[1].values[3], + "support_sliding_window": args[1].values[4], + "layer_partition": args[2], + "num_hidden_layers": args[3].value, + "num_attention_heads": args[4].value, + "num_key_value_heads": args[5].value, + "qk_head_dim": args[6].value, + "v_head_dim": args[7].value, + "mla_original_qk_head_dim": args[8].value, + "mla_original_v_head_dim": args[9].value, + "rope_mode": args[10].value, + "rope_scale": args[11].value, + "rope_theta": args[12].value, + "rope_scaling": json.loads(args[13].value), + "rope_ext_factors": args[14], + "rotary_dim": args[15].value, + "enable_disaggregation": bool(args[16].value), + "dtype": args[17].value, + } + + +@tvm.transform.module_pass(opt_level=0, name="DispatchKVCacheCreation") +class DispatchKVCacheCreation: + """Rewrite KV cache creation functions to IRModule.""" + + def __init__( + self, + target: tvm.target.Target, + flashinfer: bool, + metadata: Dict[str, Any], # noqa: UP006 + ) -> None: + """Initializer. + + Parameters + ---------- + target : tvm.target.Target + The target of the model compilation. + + flashinfer : bool + A boolean indicating if flashinfer is enabled. + + metadata : Dict[str, Any] + The model's metadata for KV cache creation. + Note that the metadata will be updated in this pass -- the + KV cache metadata will be attached. + """ + self.target = target + self.flashinfer = flashinfer + self.metadata = metadata + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + func_dict = {} + creation_func = None + for g_var, func in mod.functions_items(): + # Try to find the `create_paged_kv_cache` func. + if g_var.name_hint == "create_paged_kv_cache": + creation_func = func + else: + func_dict[g_var] = func + + if creation_func is None: + return mod + + new_mod = IRModule(func_dict) + if mod.attrs is not None: + new_mod = new_mod.with_attrs(mod.attrs) + + kwargs = extract_creation_args(creation_func) + self.attach_kv_cache_metadata(kwargs) + + bb = relax.BlockBuilder(new_mod) + extern_mods = [] + extern_mods += self.create_tir_paged_kv_cache(bb, kwargs) + extern_mods += self.create_flashinfer_paged_kv_cache(bb, kwargs) + + mod = bb.finalize() + mod_attrs = dict(mod.attrs) if mod.attrs else {} + mod = mod.with_attr("external_mods", mod_attrs.get("external_mods", []) + extern_mods) + return mod + + def attach_kv_cache_metadata(self, kwargs: Dict[str, Any]): # noqa: UP006 + """Attach the KV cache metadata to model metadata.""" + self.metadata["kv_cache"] = { + "num_hidden_layers": kwargs["num_hidden_layers"], + "num_attention_heads": kwargs["num_attention_heads"], + "num_key_value_heads": kwargs["num_key_value_heads"], + "head_dim": kwargs["qk_head_dim"], + } + + def create_tir_paged_kv_cache( + self, + bb: relax.BlockBuilder, + kwargs: Dict[str, Any], # noqa: UP006 + ) -> List[tvm.runtime.Module]: # noqa: UP006 + """Create the TIR-based PagedKVCache""" + max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]])) + max_total_seq_len = relax.Var( + "max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]]) + ) + prefill_chunk_size = relax.Var( + "prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]]) + ) + page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]])) + support_sliding_window = relax.Var( + "support_sliding_window_", + relax.ShapeType([kwargs["support_sliding_window"]]), + ) + + # Ensure 'enable_disaggregation' is optional + enable_disaggregation = kwargs.pop("enable_disaggregation", False) + kwargs["enable_disaggregation"] = enable_disaggregation + + with bb.function( + name="create_tir_paged_kv_cache", + params=[ + max_batch_size, + max_total_seq_len, + prefill_chunk_size, + page_size, + support_sliding_window, + ], + ): + cache = kv_cache.TIRPagedKVCache(target=self.target, **kwargs) + bb.emit_func_output(cache._expr) + + return cache.extern_mods + + def create_flashinfer_paged_kv_cache( + self, + bb: relax.BlockBuilder, + kwargs: Dict[str, Any], # noqa: UP006 + ) -> List[tvm.runtime.Module]: # noqa: UP006 + """Create the FlashInfer-based PagedKVCache""" + # Filter the cases which FlashInfer does not support. + if ( + not self.flashinfer + or self.target.kind.name != "cuda" + or str(kwargs["dtype"]) not in ["float16", "bfloat16"] + or ( + kwargs["rope_mode"] == RopeMode.INLINE + and ( + kwargs["rotary_dim"] != kwargs["qk_head_dim"] + or kwargs["qk_head_dim"] != kwargs["v_head_dim"] + ) + ) + ): + return [] + + max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]])) + max_total_seq_len = relax.Var( + "max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]]) + ) + prefill_chunk_size = relax.Var( + "prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]]) + ) + page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]])) + support_sliding_window = relax.Var( + "support_sliding_window_", + relax.ShapeType([kwargs["support_sliding_window"]]), + ) + + try: + with bb.function( + name="create_flashinfer_paged_kv_cache", + params=[ + max_batch_size, + max_total_seq_len, + prefill_chunk_size, + page_size, + support_sliding_window, + ], + ): + cache = kv_cache.FlashInferPagedKVCache(target=self.target, **kwargs) + bb.emit_func_output(cache._expr) + except Exception as e: + logger.info( + "Error caught when creating FlashInfer PagedKVCache: %s\n" + "The model will fallback to TIR-based KV cache.", + e, + ) + return [] + + return cache.extern_mods diff --git a/python/mlc_llm/compiler_pass/dispatch_triton_kernel.py b/python/mlc_llm/compiler_pass/dispatch_triton_kernel.py new file mode 100644 index 0000000..68e248b --- /dev/null +++ b/python/mlc_llm/compiler_pass/dispatch_triton_kernel.py @@ -0,0 +1,176 @@ +"""A pass that dispatch generic calls of triton kernels to specific kernel implementations.""" + +from typing import List # noqa: UP035 + +import tvm +from tvm import IRModule, relax +from tvm.relax.expr_functor import PyExprMutator, mutator + +from mlc_llm.op.triton import ( + get_tir_w8a8_block_fp8_group_matmul, + get_tir_w8a8_block_fp8_matmul, +) +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +@mutator +class _Rewriter(PyExprMutator): + def __init__(self, mod: IRModule, target: tvm.target.Target) -> None: + super().__init__(mod) + self.mod = mod + self.target = target + self.extern_mods: List[tvm.runtime.Module] = [] # noqa: UP006 + + def transform(self) -> tvm.IRModule: + """Entry point of the transformation""" + for g_var, func in self.mod.functions_items(): + if not isinstance(func, relax.Function): + continue + new_func = self.visit_expr(func) + # new_func = remove_all_unused(new_func) + self.builder_.update_func(g_var, new_func) + + mod = self.builder_.finalize() + mod_attrs = dict(mod.attrs) if mod.attrs else {} + mod = mod.with_attr( + "external_mods", list(mod_attrs.get("external_mods", [])) + self.extern_mods + ) + return mod + + def visit_call_(self, call: relax.Call) -> relax.Expr: + call = super().visit_call_(call) + + if ( + call.op != tvm.ir.Op.get("relax.call_dps_packed") + or not isinstance(call.args[0], relax.ExternFunc) + or not str(call.args[0].global_symbol).startswith("mlc.triton.") + ): + return call + + global_symbol = str(call.args[0].global_symbol) + assert isinstance(call.args[1], relax.Tuple) + if global_symbol == "mlc.triton.w8a8_block_fp8_matmul": + return self.w8a8_block_fp8_matmul(call.args[1].fields, call.ty) + if global_symbol == "mlc.triton.w8a8_block_fp8_group_matmul": + return self.w8a8_block_fp8_group_matmul(call.args[1].fields, call.ty) + raise ValueError(f"Unknown mlc.triton kernel identifier: {global_symbol}") + + def w8a8_block_fp8_matmul( + self, + args: List[relax.Expr], # noqa: UP006 + out_ty: relax.Type, + ) -> relax.Expr: + """Emit the w8a8_block_fp8_matmul triton kernel.""" + assert len(args) == 16 + x, weight, x_scale, weight_scale = args[:4] + ( + N, + K, + block_n, + block_k, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + num_warps, + num_stages, + ) = [arg.value.value for arg in args[4:14]] + in_dtype, out_dtype = str(args[14].value), str(args[15].value) + + prim_func, func_name = get_tir_w8a8_block_fp8_matmul( + N, + K, + block_n, + block_k, + in_dtype, + out_dtype, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + num_warps, + num_stages, + self.extern_mods, + ) + if prim_func is None: + # The TIR function is already in the IRModule + gv = self.builder_.get().get_global_var(func_name) + else: + # Add the TIR function to the IRModule + gv = self.builder_.add_func(prim_func, func_name) + + return relax.call_tir(gv, [x, weight, x_scale, weight_scale], out_ty=out_ty) + + def w8a8_block_fp8_group_matmul( + self, + args: List[relax.Expr], # noqa: UP006 + out_ty: relax.Type, + ) -> relax.Expr: + """Emit the w8a8_block_fp8_group_matmul triton kernel.""" + assert len(args) == 19 + x, weight, x_scale, weight_scale, expert_ids, indptr = args[:6] + ( + N, + K, + num_experts, + block_n, + block_k, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + num_warps, + num_stages, + ) = [arg.value.value for arg in args[6:17]] + in_dtype, out_dtype = str(args[17].value), str(args[18].value) + + prim_func, func_name = get_tir_w8a8_block_fp8_group_matmul( + N, + K, + num_experts, + block_n, + block_k, + in_dtype, + out_dtype, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + num_warps, + num_stages, + self.extern_mods, + ) + if prim_func is None: + # The TIR function is already in the IRModule + gv = self.builder_.get().get_global_var(func_name) + else: + # Add the TIR function to the IRModule + gv = self.builder_.add_func(prim_func, func_name) + + return relax.call_tir( + gv, + [x, weight, x_scale, weight_scale, expert_ids, indptr], + out_ty=out_ty, + ) + + +@tvm.transform.module_pass(opt_level=0, name="DispatchTritonKernel") +class DispatchTritonKernel: + """Rewrite KV cache creation functions to IRModule.""" + + def __init__(self, target: tvm.target.Target) -> None: + """Initializer. + + Parameters + ---------- + """ + self.target = target + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + if self.target.kind.name != "cuda": + return mod + + return _Rewriter(mod, self.target).transform() diff --git a/python/mlc_llm/compiler_pass/estimate_memory_usage.py b/python/mlc_llm/compiler_pass/estimate_memory_usage.py new file mode 100644 index 0000000..78611a2 --- /dev/null +++ b/python/mlc_llm/compiler_pass/estimate_memory_usage.py @@ -0,0 +1,87 @@ +"""Memory usage estimation analysis function for Relax functions.""" + +import json +from typing import Any, Dict # noqa: UP035 + +import tvm +from tvm import relax, tirx +from tvm.ir import IRModule, Op +from tvm.relax.expr_functor import PyExprVisitor, visitor + +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +@tvm.transform.module_pass(opt_level=0, name="AttachMetadata") +class AttachMetadataWithMemoryUsage: + """Attach a Relax function that returns metadata in a JSON string""" + + def __init__(self, metadata: Dict[str, Any]): # noqa: UP006 + self.metadata = metadata + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """Entrypoint""" + + func_name = "_metadata" + + def _emit_metadata(metadata): + bb = relax.BlockBuilder() + with bb.function(func_name, params=[]): + bb.emit_func_output(relax.StringImm(json.dumps(metadata))) + return bb.finalize()[func_name] + + self.metadata["memory_usage"] = _MemoryEstimator().run(mod) + mod[func_name] = _emit_metadata(self.metadata) + return mod + + +@visitor +class _MemoryEstimator(PyExprVisitor): + """The IR visitor which estimates the memory usage of each Relax function.""" + + def __init__(self) -> None: + self.planned_alloc_mem = 0 + self.planned_mem_num = 0 + self._op_alloc_tensor = Op.get("relax.builtin.alloc_tensor") + self._op_alloc_storage = Op.get("relax.memory.alloc_storage") + + def run(self, mod: IRModule) -> Dict[str, int]: # noqa: UP006 + """Entry point of the visitor.""" + result: Dict[str, int] = {} # noqa: UP006 + for global_var, func in mod.functions_items(): + if isinstance(func, relax.Function): + self.planned_alloc_mem = 0 + self.planned_mem_num = 0 + self.visit_expr(func) + result[global_var.name_hint] = self.planned_alloc_mem + logger.info( + "[Memory usage] Function `%s`: %.2f MB", + global_var.name_hint, + self.planned_alloc_mem / 1024 / 1024, + ) + return result + + def visit_call_(self, call: relax.Call) -> None: + if call.op == self._op_alloc_tensor: + self._builtin_tensor_alloc(shape=call.args[0], dtype_str=call.args[1].value) + elif call.op == self._op_alloc_storage: + self._storage_alloc(size=call.args[0]) + super().visit_call_(call) + + def _builtin_tensor_alloc(self, shape: relax.Expr, dtype_str: str) -> None: + assert isinstance(shape, relax.ShapeExpr) + size = 1 + for dim_len in shape.values: + if not isinstance(dim_len, tvm.tirx.IntImm): + return + size *= dim_len.value + dtype = tvm.DataType(dtype_str) + self.planned_mem_num += 1 + self.planned_alloc_mem += size * ((dtype.bits + 7) // 8) * dtype.lanes + + def _storage_alloc(self, size: relax.Expr) -> None: + assert isinstance(size, relax.ShapeExpr) + if isinstance(size.values[0], tirx.IntImm): + self.planned_mem_num += 1 + self.planned_alloc_mem += size.values[0].value diff --git a/python/mlc_llm/compiler_pass/fuse_add_norm.py b/python/mlc_llm/compiler_pass/fuse_add_norm.py new file mode 100644 index 0000000..15307f2 --- /dev/null +++ b/python/mlc_llm/compiler_pass/fuse_add_norm.py @@ -0,0 +1,238 @@ +"""A compiler pass that fuses add + rms_norm.""" + +from typing import Optional + +import tvm +from tvm import relax +from tvm.relax.analysis import remove_all_unused +from tvm.relax.expr_functor import PyExprMutator, mutator +from tvm.script import tirx as T + +from ..support.max_thread_check import get_max_num_threads_per_block + + +def _get_add_rms_norm_decode(hidden_size: int, eps: float, TX: int, in_dtype: str): + if in_dtype not in ("float16", "bfloat16"): + raise ValueError(f"Unsupported data type: {in_dtype}") + inv_hidden_size = T.float32(1.0 / float(hidden_size)) + eps = T.float32(eps) + add_local_size = hidden_size // TX + + @T.prim_func(private=True, s_tir=True) + def decode_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle): + T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1}) + batch_size = T.int32() + A = T.match_buffer(pA, (batch_size, 1, hidden_size), in_dtype) + B = T.match_buffer(pB, (batch_size, 1, hidden_size), in_dtype) + C = T.match_buffer(pC, (hidden_size,), in_dtype) + out = T.match_buffer(pO, (batch_size, 1, hidden_size), in_dtype) + add = T.match_buffer(pAdd, (batch_size, 1, hidden_size), in_dtype) + add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local") + sum_shared = T.sblock_alloc_buffer((batch_size, 1), scope="shared") + sum_local = T.sblock_alloc_buffer((TX, batch_size, 1), scope="local") + for v_bx in T.thread_binding(batch_size, thread="blockIdx.x"): + for v_tx in T.thread_binding( + TX, + thread="threadIdx.x", + annotations={ + "pragma_auto_unroll_max_step": 256, + "pragma_unroll_explicit": 1, + }, + ): + for i in range(add_local_size): + with T.sblock("T_add"): + bx = T.axis.spatial(batch_size, v_bx) + h = T.axis.spatial(hidden_size, i * TX + v_tx) + add_local[h // TX] = A[bx, 0, h] + B[bx, 0, h] + with T.sblock("T_write_back"): + bx = T.axis.spatial(batch_size, v_bx) + v_ax1 = T.axis.spatial(1, 0) + h = T.axis.spatial(hidden_size, i * TX + v_tx) + add[bx, v_ax1, h] = add_local[h // TX] + with T.sblock("T_multiply_red_rf_init"): + tx, bx = T.axis.remap("SS", [v_tx, v_bx]) + sum_local[tx, bx, 0] = T.float32(0) + for v_i, _j in T.grid(add_local_size, 1): + with T.sblock("T_multiply_red_rf_update"): + tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i]) + sum_local[tx, bx, 0] += T.float32(add_local[i]) * T.float32(add_local[i]) + for _j in range(1): + for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"): + with T.sblock("T_multiply_red"): + tx, bx = T.axis.remap("RS", [v_tx_2, v_bx]) + T.reads(sum_local[tx, bx, 0]) + T.writes(sum_shared[bx, 0]) + with T.init(): + sum_shared[bx, 0] = T.float32(0) + sum_shared[bx, 0] += sum_local[tx, bx, 0] + for i in range(add_local_size): + for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"): + with T.sblock("T_cast_2"): + bx = T.axis.spatial(batch_size, v_bx) + h = T.axis.spatial(hidden_size, i * TX + v_tx_2) + out[bx, 0, h] = T.cast( + T.rsqrt(sum_shared[bx, 0] * inv_hidden_size + eps) + * T.float32(add_local[h // TX]) + * T.float32(C[h]), + dtype=in_dtype, + ) + + return decode_add_rms + + +def _get_add_rms_norm_prefill(hidden_size: int, eps: float, TX: int, in_dtype: str): + if in_dtype not in ("float16", "bfloat16"): + raise ValueError(f"Unsupported data type: {in_dtype}") + inv_hidden_size = T.float32(1.0 / float(hidden_size)) + eps = T.float32(eps) + add_local_size = hidden_size // TX + + @T.prim_func(private=True, s_tir=True) + def prefill_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle): + T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1}) + seq_len = T.int32() + A = T.match_buffer(pA, (1, seq_len, hidden_size), in_dtype) + B = T.match_buffer(pB, (1, seq_len, hidden_size), in_dtype) + C = T.match_buffer(pC, (hidden_size,), in_dtype) + out = T.match_buffer(pO, (1, seq_len, hidden_size), in_dtype) + add = T.match_buffer(pAdd, (1, seq_len, hidden_size), in_dtype) + add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local") + sum_shared = T.sblock_alloc_buffer((1, seq_len), scope="shared") + sum_local = T.sblock_alloc_buffer((TX, 1, seq_len), scope="local") + for v_bx in T.thread_binding(seq_len, thread="blockIdx.x"): + for v_tx in T.thread_binding( + TX, + thread="threadIdx.x", + annotations={ + "pragma_auto_unroll_max_step": 256, + "pragma_unroll_explicit": 1, + }, + ): + for v_i in range(add_local_size): + with T.sblock("T_add"): + bx = T.axis.spatial(seq_len, v_bx) + h = T.axis.spatial(hidden_size, v_i * TX + v_tx) + add_local[h // TX] = A[0, bx, h] + B[0, bx, h] + with T.sblock("T_write_back"): + bx = T.axis.spatial(seq_len, v_bx) + h = T.axis.spatial(hidden_size, v_i * TX + v_tx) + add[0, bx, h] = add_local[h // TX] + with T.sblock("T_multiply_red_rf_init"): + tx, bx = T.axis.remap("SS", [v_tx, v_bx]) + sum_local[tx, 0, bx] = T.float32(0) + for v_i, _j in T.grid(add_local_size, 1): + with T.sblock("T_multiply_red_rf_update"): + tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i]) + sum_local[tx, 0, bx] += T.float32(add_local[i]) * T.float32(add_local[i]) + for _j in range(1): + for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"): + with T.sblock("T_multiply_red"): + tx, bx = T.axis.remap("RS", [v_tx_2, v_bx]) + with T.init(): + sum_shared[0, bx] = T.float32(0) + sum_shared[0, bx] = sum_shared[0, bx] + sum_local[tx, 0, bx] + for v_i in range(add_local_size): + for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"): + with T.sblock("T_cast_2"): + bx = T.axis.spatial(seq_len, v_bx) + v1 = T.axis.spatial(hidden_size, v_i * TX + v_tx_2) + out[0, bx, v1] = T.cast( + T.rsqrt(sum_shared[0, bx] * inv_hidden_size + eps) + * T.float32(add_local[v1 // TX]) + * T.float32(C[v1]), + dtype=in_dtype, + ) + + return prefill_add_rms + + +@tvm.transform.module_pass(opt_level=0, name="FuseAddRMSNorm") +class FuseAddRMSNorm: + """A compiler pass that fuses add + rms_norm.""" + + def __init__(self, target: tvm.target.Target) -> None: + """Initializer. + + Parameters + ---------- + target : tvm.target.Target + Target device. + """ + self.target = target + + def transform_module(self, mod: tvm.IRModule, _ctx: tvm.transform.PassContext) -> tvm.IRModule: + """IRModule-level transformation.""" + return _FuseAddRMSNormRewriter(mod.clone(), self.target).transform() + + +@mutator +class _FuseAddRMSNormRewriter(PyExprMutator): + def __init__(self, mod: tvm.IRModule, target: tvm.target.Target): + super().__init__(mod) + self.mod = mod + self.prefill_norm_gv: Optional[tvm.ir.GlobalVar] = None + self.decode_norm_gv: Optional[tvm.ir.GlobalVar] = None + self.TX = min(1024, get_max_num_threads_per_block(target)) + + def transform(self) -> tvm.IRModule: + """Entry point of the transformation""" + for g_var, func in self.mod.functions_items(): + if not isinstance(func, relax.Function): + continue + new_func = self.visit_expr(func) + new_func = remove_all_unused(new_func) + self.builder_.update_func(g_var, new_func) + return self.builder_.finalize() + + def visit_call_(self, call: relax.Call) -> relax.Expr: + call = super().visit_call_(call) + + # Match the "rms_norm(add(x1, x2), w)" pattern + if call.op != tvm.ir.Op.get("relax.nn.rms_norm") or call.ty.dtype not in [ + "bfloat16", + "float16", + ]: + return call + assert len(call.args) == 2 + weight = call.args[1] + eps = call.attrs.epsilon + assert isinstance(call.args[0], relax.Var) + y = self.lookup_binding(call.args[0]) + if not isinstance(y, relax.Call) or y.op != tvm.ir.Op.get("relax.add"): + return call + assert len(y.args) == 2 + x1 = y.args[0] + x2 = y.args[1] + # Extra check + n, _, h = x1.ty.shape + h = int(h) + if h % self.TX != 0: + return call + + is_prefill = n == 1 + func_gv = self.prefill_norm_gv if is_prefill else self.decode_norm_gv + if func_gv is None: + if is_prefill: + func_gv = self.builder_.add_func( + _get_add_rms_norm_prefill(h, eps, self.TX, call.ty.dtype), + "fuse_add_norm_prefill", + ) + self.prefill_norm_gv = func_gv + else: + func_gv = self.builder_.add_func( + _get_add_rms_norm_decode(h, eps, self.TX, call.ty.dtype), + "fuse_add_norm_decode", + ) + self.decode_norm_gv = func_gv + + tuple_output = self.builder_.emit( + relax.call_tir( + func_gv, + [x1, x2, weight], + out_ty=[x1.ty, x2.ty], + ) + ) + new_o = relax.TupleGetItem(tuple_output, 0) + new_y = self.builder_.emit(relax.TupleGetItem(tuple_output, 1)) + self.set_var_remap(call.args[0], new_y) + return new_o diff --git a/python/mlc_llm/compiler_pass/fuse_dequantize_matmul_ewise.py b/python/mlc_llm/compiler_pass/fuse_dequantize_matmul_ewise.py new file mode 100644 index 0000000..a946391 --- /dev/null +++ b/python/mlc_llm/compiler_pass/fuse_dequantize_matmul_ewise.py @@ -0,0 +1,85 @@ +"""A compiler pass that fuses dequantize + matmul + elementwise.""" + +import tvm +from tvm import IRModule, relax +from tvm.relax.dpl.pattern import GlobalVarPattern, TuplePattern, is_op, wildcard + + +@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeMatmulEwise") +class FuseDequantizeMatmulEwise: + """A compiler pass that fuses dequantize + matmul + elementwise.""" + + def transform_module( + self, + mod: IRModule, + _ctx: tvm.transform.PassContext, + ) -> IRModule: + """IRModule-level transformation""" + seq = [] + for n_aux_tensor in [0, 1, 2, 3, 4]: + for match_ewise in [0, 1, 2, 3, 6]: + if match_ewise == 6 and n_aux_tensor != 4: + continue + seq.append( + relax.transform.FuseOpsByPattern( + [ + ( + "dequantize_matmul", + *_pattern(match_ewise, n_aux_tensor), + ) + ] + ) + ) + seq.append(relax.transform.FuseTIR()) + return tvm.transform.Sequential(seq)(mod) + + +def _pattern(match_ewise: int, n_aux_tensor: int): + w_scaled = wildcard() + x = wildcard() + w = is_op("relax.call_tir")( + GlobalVarPattern(), + TuplePattern([w_scaled] + [wildcard() for _ in range(n_aux_tensor)]), + add_constraint=False, + ) + matmul = is_op("relax.call_tir")( + GlobalVarPattern(), + TuplePattern([x, w] + [wildcard() for _ in range(match_ewise)]), + add_constraint=False, + ) + annotations = { + "w_scaled": w_scaled, + "x": x, + "w": w, + "matmul": matmul, + } + + def _check_decoding(ctx: relax.transform.PatternCheckContext) -> bool: + call = ctx.annotated_expr["w"] + if not isinstance(call, relax.Call): + return False + g_var = call.args[0] + if not isinstance(g_var, relax.GlobalVar): + return False + return g_var.name_hint.startswith("dequantize") or g_var.name_hint.startswith( + "fused_dequantize" + ) + + def _check_matmul(ctx: relax.transform.PatternCheckContext) -> bool: + call = ctx.annotated_expr["matmul"] + if not isinstance(call, relax.Call): + return False + g_var = call.args[0] + if not isinstance(g_var, relax.GlobalVar): + return False + return ( + g_var.name_hint.startswith("matmul") + or g_var.name_hint.startswith("fused_matmul") + or g_var.name_hint.startswith("NT_matmul") + or g_var.name_hint.startswith("fused_NT_matmul") + ) + + def _check(ctx: relax.transform.PatternCheckContext) -> bool: + return _check_decoding(ctx) and _check_matmul(ctx) + + return matmul, annotations, _check diff --git a/python/mlc_llm/compiler_pass/fuse_dequantize_take.py b/python/mlc_llm/compiler_pass/fuse_dequantize_take.py new file mode 100644 index 0000000..8e9c3fb --- /dev/null +++ b/python/mlc_llm/compiler_pass/fuse_dequantize_take.py @@ -0,0 +1,91 @@ +"""A compiler pass that fuses dequantize + take.""" + +import tvm +from tvm import IRModule, relax, tirx +from tvm.relax.dpl.pattern import ( + GlobalVarPattern, + TuplePattern, + is_const, + is_op, + wildcard, +) + + +@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTake") +class FuseDequantizeTake: + """A compiler pass that fuses dequantize + take.""" + + def transform_module( + self, + mod: IRModule, + _ctx: tvm.transform.PassContext, + ) -> IRModule: + """IRModule-level transformation""" + seq = [] + for n_aux_tensor in [2, 3]: + for match_tir_vars in [False, True]: + seq.append( + relax.transform.FuseOpsByPattern( + [ + ( + "dequantize_take", + *_pattern(n_aux_tensor, match_tir_vars), + ) + ] + ) + ) + seq.append(relax.transform.FuseTIR()) + mod = tvm.transform.Sequential(seq)(mod) + for g_var, func in mod.functions_items(): + name = g_var.name_hint + if isinstance(func, tirx.PrimFunc) and ( + ("fused_dequantize" in name) and ("take" in name) + ): + sch_mod = tvm.IRModule({"main": func}) + sch_mod = tirx.transform.ForceNarrowIndexToInt32()(sch_mod) + sch = tvm.s_tir.Schedule(sch_mod) + sch.compute_inline("dequantize") + mod[g_var] = sch.mod["main"] + return mod + + +def _pattern(n_aux_tensor: int, match_tir_vars: bool): + dequantize = is_op("relax.call_tir")( + GlobalVarPattern(), + TuplePattern([wildcard() for _ in range(n_aux_tensor)]), + add_constraint=False, + ) + indices = ~is_const() + if match_tir_vars: + call_tir_args_take = [ + GlobalVarPattern(), + TuplePattern([dequantize, indices]), + wildcard(), + ] + else: + call_tir_args_take = [ + GlobalVarPattern(), + TuplePattern([dequantize, indices]), + ] + take = is_op("relax.call_tir")( + *call_tir_args_take, + add_constraint=False, + ) + annotations = { + "take": take, + "dequantize": dequantize, + "indices": indices, + } + + def _check(ctx: relax.transform.PatternCheckContext) -> bool: + take = ctx.annotated_expr["take"] + dequantize = ctx.annotated_expr["dequantize"] + if not isinstance(dequantize, relax.Call): + return False + if not isinstance(take.args[0], relax.GlobalVar) or not isinstance( + dequantize.args[0], relax.GlobalVar + ): + return False + return "take" in take.args[0].name_hint and "dequantize" in dequantize.args[0].name_hint + + return take, annotations, _check diff --git a/python/mlc_llm/compiler_pass/fuse_dequantize_transpose.py b/python/mlc_llm/compiler_pass/fuse_dequantize_transpose.py new file mode 100644 index 0000000..30c4520 --- /dev/null +++ b/python/mlc_llm/compiler_pass/fuse_dequantize_transpose.py @@ -0,0 +1,107 @@ +"""A compiler pass that fuses transpose + dequantize.""" + +import tvm +from tvm import relax, s_tir, tirx +from tvm.ir.module import IRModule +from tvm.relax.analysis import remove_all_unused +from tvm.relax.expr_functor import PyExprMutator, mutator + + +@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTranspose") +class FuseDequantizeTranspose: + """A compiler pass that fuses transpose + dequantize.""" + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """IRModule-level transformation""" + return _DequantizeTransposeFuser(mod).transform() + + +@mutator +class _DequantizeTransposeFuser(PyExprMutator): + def __init__( + self, + mod: IRModule, + ): + super().__init__(mod) + self.mod = mod + + def transform(self) -> IRModule: + """Entry point""" + for g_var, func in self.mod.functions_items(): + if isinstance(func, relax.Function): + updated_func = self.visit_expr(func) + updated_func = remove_all_unused(updated_func) + self.builder_.update_func(g_var, updated_func) + return self.builder_.get() + + def visit_call_( + self, + call: relax.Call, + ) -> relax.Expr: + call = self.visit_expr_post_order(call) + if call.op != tvm.ir.Op.get("relax.matmul"): + return call + # Do not fuse dequantize-transpose for GeMM + if ( + call.args[0].ty.ndim < 2 + or not isinstance(call.args[0].ty.shape[-2], tirx.IntImm) + or call.args[0].ty.shape[-2].value != 1 + ): + return call + + matmul_rhs = self.lookup_binding(call.args[1]) + if ( + not isinstance(matmul_rhs, relax.Call) + or matmul_rhs.op != tvm.ir.Op.get("relax.permute_dims") + or matmul_rhs.args[0].ty.ndim != 2 + or matmul_rhs.attrs.axes is not None + ): + return call + + transpose_input = self.lookup_binding(matmul_rhs.args[0]) + if ( + not isinstance(transpose_input, relax.Call) + or transpose_input.op != tvm.ir.Op.get("relax.call_tir") + or not transpose_input.args[0].name_hint.startswith("dequantize") + or not isinstance(transpose_input.ty, relax.TensorType) + ): + return call + + dequantize_tir_func = self.mod[transpose_input.args[0]] + assert isinstance(dequantize_tir_func, tirx.PrimFunc) + if ( + len(dequantize_tir_func.body.block.alloc_buffers) != 1 + or not isinstance(dequantize_tir_func.body.block.body, tirx.SeqStmt) + or len(dequantize_tir_func.body.block.body) != 2 + or not isinstance(dequantize_tir_func.body.block.body[1], tirx.For) + or not isinstance(dequantize_tir_func.body.block.body[1].body.body, tirx.SBlockRealize) + or dequantize_tir_func.body.block.body[1].body.body.block.name_hint != "T_transpose" + ): + return call + + new_func_buffers = [ + dequantize_tir_func.buffer_map[var] for var in dequantize_tir_func.params + ] + new_func_buffers[-1] = dequantize_tir_func.body.block.alloc_buffers[0] + new_func = tirx.PrimFunc( + params=new_func_buffers, + body=tirx.SBlockRealize( + iter_values=[], + predicate=True, + block=tirx.SBlock( + iter_vars=[], + reads=[], + writes=[], + name_hint="root", + body=dequantize_tir_func.body.block.body[0], + ), + ), + ) + # Call `renew_defs` for deep-copy to avoid IR node duplication in + # different PrimFuncs of an IRModule. + new_func = s_tir.renew_defs(new_func) + g_var = self.builder_.add_func(new_func, func_name="dequantize") + dequantize_matmul_rhs = self.builder_.emit( + relax.call_tir(g_var, transpose_input.args[1], out_ty=matmul_rhs.ty) + ) + return relax.op.matmul(call.args[0], dequantize_matmul_rhs, out_dtype=call.attrs.out_dtype) diff --git a/python/mlc_llm/compiler_pass/fuse_ft_dequantize_matmul_epilogue.py b/python/mlc_llm/compiler_pass/fuse_ft_dequantize_matmul_epilogue.py new file mode 100644 index 0000000..1081c55 --- /dev/null +++ b/python/mlc_llm/compiler_pass/fuse_ft_dequantize_matmul_epilogue.py @@ -0,0 +1,331 @@ +"""A compiler pass that fuses dequantize matmul + epilogue.""" + +import operator +from functools import reduce + +import tvm +from tvm import IRModule, relax +from tvm.relax.dpl import rewrite_call +from tvm.relax.dpl.pattern import is_op, wildcard + + +@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeEpilogue") +class FuseFTDequantizeEpilogue: + """A compiler pass that fuses FasterTransformer dequantize matmul + epilogue.""" + + def transform_module( + self, + mod: IRModule, + _ctx: tvm.transform.PassContext, + ) -> IRModule: + """IRModule-level transformation""" + for gv, func in mod.functions_items(): + if isinstance(func, relax.Function): + func = fuse_bias(func) + func = fuse_activation(func) + func = fuse_residual_binary(func) + func = fuse_residual_unary(func) + mod[gv] = func + return mod + + +def fuse_bias(func: relax.Function) -> relax.Function: + """ + Fuse following `relax.add` into fastertransformer.gemm_fp16_int as bias: + + Before: + ``` + lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int", ...) + lv2 = relax.add(lv1, bias) + + ``` + After: + ``` + lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ..., bias, ...) + ``` + + Parameters + ---------- + func : relax.Function + The function before fusion. + + Returns + ------- + ret : relax.Function + The function after fusion. + """ + decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True) + bias = wildcard() + pattern = is_op("relax.add")(decode_matmul, bias) | is_op("relax.add")(bias, decode_matmul) + + def rewriter(expr, match): + if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int": + assert len(match[decode_matmul].args) == 2 + args_list = match[decode_matmul].args[1] + assert len(args_list) == 8 + if not args_list[3].value == "identity": + # bias cannot be fused after activation + return expr + matched_bias = match[bias] + bias_stride = ( + matched_bias.ty.shape[-1] + if bias + and not reduce(operator.mul, matched_bias.ty.shape, 1) == matched_bias.ty.shape[-1] + else 0 + ) + return relax.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias", + [ + args_list[0], # x + args_list[1], # weight + args_list[2], # scale + matched_bias, # bias + args_list[3], # activation + args_list[4], # m + args_list[5], # n + args_list[6], # k + args_list[7], # group_size + bias_stride, # bias_stride + ], + out_ty=match[decode_matmul].ty, + ) + return expr + + return rewrite_call(pattern, rewriter, func) + + +def fuse_activation(func: relax.Function) -> relax.Function: + """ + Fuse following `relax.nn.silu/relu/gelu` into fastertransformer.gemm_fp16_int_bias + as activation: + + Before: + ``` + lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ...) + lv2 = relax.silu(lv1) + + ``` + After: + ``` + lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ..., "silu", ...) + ``` + + Parameters + ---------- + func : relax.Function + The function before fusion. + + Returns + ------- + ret : relax.Function + The function after fusion. + """ + decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True) + pattern = ( + is_op("relax.nn.silu")(decode_matmul) + | is_op("relax.nn.gelu")(decode_matmul) + | is_op("relax.nn.relu")(decode_matmul) + ) + + def rewriter(expr, match): + if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int": + matched_activation = match[pattern] + assert matched_activation.op.name in [ + "relax.nn.silu", + "relax.nn.gelu", + "relax.nn.relu", + ] + assert len(match[decode_matmul].args) == 2 + args_list = match[decode_matmul].args[1] + assert len(args_list) == 8 + return relax.call_dps_packed( + "fastertransformer.gemm_fp16_int", + [ + args_list[0], # x + args_list[1], # weight + args_list[2], # scale + matched_activation.op.name[9:], # activation + args_list[4], # m + args_list[5], # n + args_list[6], # k + args_list[7], # group_size + ], + out_ty=match[decode_matmul].ty, + ) + if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int_bias": + matched_activation = match[pattern] + assert matched_activation.op.name in [ + "relax.nn.silu", + "relax.nn.gelu", + "relax.nn.relu", + ] + assert len(match[decode_matmul].args) == 2 + args_list = match[decode_matmul].args[1] + assert len(args_list) == 10 + return relax.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias", + [ + args_list[0], # x + args_list[1], # weight + args_list[2], # scale + args_list[3], # bias + matched_activation.op.name[9:], # activation + args_list[5], # m + args_list[6], # n + args_list[7], # k + args_list[8], # group_size + args_list[9], # bias_stride + ], + out_ty=match[decode_matmul].ty, + ) + return expr + + return rewrite_call(pattern, rewriter, func) + + +def fuse_residual_binary(func: relax.Function) -> relax.Function: + """ + Fuse following `relax.add/multiply` into fastertransformer.gemm_fp16_int_bias as + residual binary operation: + + Before: + ``` + lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ...) + lv2 = relax.add(lv1, residual) + + ``` + After: + ``` + lv2 = relax.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias_residual", + ..., + residual, + ..., + "plus", + ... + ) + ``` + + Parameters + ---------- + func : relax.Function + The function before fusion. + + Returns + ------- + ret : relax.Function + The function after fusion. + """ + decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True) + residual = wildcard() + pattern = ( + is_op("relax.add")(decode_matmul, residual) + | is_op("relax.add")(residual, decode_matmul) + | is_op("relax.multiply")(decode_matmul, residual) + | is_op("relax.multiply")(residual, decode_matmul) + ) + + def rewriter(expr, match): + if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int_bias": + matched_binary = match[pattern] + assert matched_binary.op.name in ["relax.add", "relax.multiply"] + binary_op = "plus" if matched_binary.op.name == "relax.add" else "multiply" + assert len(match[decode_matmul].args) == 2 + args_list = match[decode_matmul].args[1] + assert len(args_list) == 10 + matched_residual = match[residual] + if not args_list[9].value == 0: + # fastertransformer.gemm_fp16_int_bias_residual does not support + # bias_stride != 0 yet + return expr + return relax.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias_residual", + [ + args_list[0], # x + args_list[1], # weight + args_list[2], # scale + args_list[3], # bias + matched_residual, # residual + args_list[4], # activation + binary_op, # binary_op + "identity", # unary_op + args_list[5], # m + args_list[6], # n + args_list[7], # k + args_list[8], # group_size + ], + out_ty=match[decode_matmul].ty, + ) + return expr + + return rewrite_call(pattern, rewriter, func) + + +def fuse_residual_unary(func: relax.Function) -> relax.Function: + """ + Fuse following `relax.nn.silu/relu/gelu` into fastertransformer.gemm_fp16_int_bias_residual + as residual unary operation: + + Before: + ``` + lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias_residual", ...) + lv2 = relax.silu(lv1) + + ``` + After: + ``` + lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias_residual", ..., "silu", ...) + ``` + + Parameters + ---------- + func : relax.Function + The function before fusion. + + Returns + ------- + ret : relax.Function + The function after fusion. + """ + decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True) + pattern = ( + is_op("relax.nn.silu")(decode_matmul) + | is_op("relax.nn.gelu")(decode_matmul) + | is_op("relax.nn.relu")(decode_matmul) + ) + + def rewriter(expr, match): + if ( + match[decode_matmul].args[0].global_symbol + == "fastertransformer.gemm_fp16_int_bias_residual" + ): + matched_activation = match[pattern] + assert matched_activation.op.name in [ + "relax.nn.silu", + "relax.nn.gelu", + "relax.nn.relu", + ] + assert len(match[decode_matmul].args) == 2 + args_list = match[decode_matmul].args[1] + assert len(args_list) == 12 + return relax.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias_residual", + [ + args_list[0], # x + args_list[1], # weight + args_list[2], # scale + args_list[3], # bias + args_list[4], # residual + args_list[5], # activation + args_list[6], # binary_op + matched_activation.op.name[9:], # activation + args_list[8], # m + args_list[9], # n + args_list[10], # k + args_list[11], # group_size + ], + out_ty=match[decode_matmul].ty, + ) + return expr + + return rewrite_call(pattern, rewriter, func) diff --git a/python/mlc_llm/compiler_pass/fuse_transpose_matmul.py b/python/mlc_llm/compiler_pass/fuse_transpose_matmul.py new file mode 100644 index 0000000..6734d56 --- /dev/null +++ b/python/mlc_llm/compiler_pass/fuse_transpose_matmul.py @@ -0,0 +1,145 @@ +"""A compiler pass that fuses transpose + matmul.""" + +import tvm +from tvm import IRModule, relax, te, tirx +from tvm.relax.dpl.pattern import is_op, wildcard +from tvm.relax.expr_functor import PyExprMutator, mutator + + +@tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul") +class FuseTransposeMatmul: + """A compiler pass that fuses transpose + matmul.""" + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """IRModule-level transformation""" + mod = relax.transform.FuseOpsByPattern( + [ + ( + "transpose_matmul_fuse", + *_pattern(), + ), + ] + )(mod) + transpose_matmul_codegen = _TransposeMatmulFuser(mod) + for g_var, func in mod.functions_items(): + if isinstance(func, relax.Function): + func = transpose_matmul_codegen.visit_expr(func) + transpose_matmul_codegen.builder_.update_func(g_var, func) + return transpose_matmul_codegen.builder_.get() + + +def _pattern(): + """Pattern for transpose + matmul.""" + w = wildcard() + x = wildcard() + wT = is_op("relax.permute_dims")(w) + o = is_op("relax.matmul")(x, wT) + annotations = {"o": o, "w": w, "x": x, "wT": wT} + + def _check(context: relax.transform.PatternCheckContext) -> bool: + transpose_call = context.annotated_expr["wT"] + ndim = transpose_call.args[0].ty.ndim + if ndim == -1: + return False + if ndim == 2 and transpose_call.attrs.axes is None: + return True + axes = list(range(ndim)) + axes[-1], axes[-2] = axes[-2], axes[-1] + return list(transpose_call.attrs.axes) == axes + + return o, annotations, _check + + +@mutator +class _TransposeMatmulFuser(PyExprMutator): + def __init__(self, mod): + super().__init__(mod) + + def visit_call_( + self, + call: relax.Call, + ) -> relax.Expr: + out_dtype = None + + def te_transposed_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor: + nonlocal out_dtype + a_shape = list(a.shape) + b_shape = list(b.shape) + a_prepended = False + b_appended = False + if len(a_shape) == 1: + a_prepended = True + a_shape.insert(0, 1) + if len(b_shape) == 1: + b_appended = True + b_shape.append(1) + + is_a_larger = len(a_shape) > len(b_shape) + offset = len(a_shape) - len(b_shape) if is_a_larger else len(b_shape) - len(a_shape) + + a_relax = relax.Var("a", relax.TensorType(a.shape)) + bT_shape = list(b.shape) + bT_shape[-1], bT_shape[-2] = bT_shape[-2], bT_shape[-1] + bT_relax = relax.Var("b", relax.TensorType(bT_shape)) + output_shape = self.builder_.normalize(relax.op.matmul(a_relax, bT_relax)).ty.shape + + def matmul_compute(*idx_spatial): + k = te.reduce_axis((0, a_shape[-1]), name="k") + + def multiply_compute(idx_reduce): + a_indices = [] + b_indices = [] + + for i in range(offset): + if is_a_larger: + a_indices.append(idx_spatial[i]) + else: + b_indices.append(idx_spatial[i]) + for i in range(offset, len(output_shape) - (2 - a_prepended - b_appended)): + a_dim = a_shape[i if is_a_larger else i - offset] + b_dim = b_shape[i if not is_a_larger else i - offset] + dim_equal = a_dim == b_dim + if not isinstance(dim_equal, tirx.IntImm) or dim_equal == 0: + a_dim_is_one = isinstance(a_dim, tirx.IntImm) and a_dim == 1 + b_dim_is_one = isinstance(b_dim, tirx.IntImm) and b_dim == 1 + a_indices.append(0 if a_dim_is_one else idx_spatial[i]) + b_indices.append(0 if b_dim_is_one else idx_spatial[i]) + else: + a_indices.append(idx_spatial[i]) + b_indices.append(idx_spatial[i]) + + if not a_prepended: + a_indices.append(idx_spatial[-2 + b_appended]) + a_indices.append(idx_reduce) + if not b_appended: + b_indices.append(idx_spatial[-1]) + b_indices.append(idx_reduce) + + dtype = out_dtype + if dtype != "": + return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype) + return a(*a_indices) * b(*b_indices) + + return te.sum(multiply_compute(k), axis=k) + + return te.compute( + output_shape, + lambda *idx: matmul_compute(*idx), + name="NT_matmul", + ) + + if isinstance(call.op, relax.GlobalVar): + function = self.builder_.get()[call.op] + if ( + "Composite" in function.attrs + and function.attrs["Composite"] == "transpose_matmul_fuse" + ): + out_dtype = function.ret_ty.dtype + return self.builder_.call_te( + te_transposed_matmul, + call.args[1], + call.args[0], + primfunc_name_hint="NT_matmul", + ) + + return super().visit_call_(call) diff --git a/python/mlc_llm/compiler_pass/lift_global_buffer_alloc.py b/python/mlc_llm/compiler_pass/lift_global_buffer_alloc.py new file mode 100644 index 0000000..1eebab5 --- /dev/null +++ b/python/mlc_llm/compiler_pass/lift_global_buffer_alloc.py @@ -0,0 +1,198 @@ +"""A compiler pass that lifts TIR-level global allocation to Relax.""" + +from typing import Dict, List, Tuple # noqa: UP035 + +import tvm +from tvm import relax, tirx +from tvm.ir.module import IRModule +from tvm.relax.analysis import remove_all_unused +from tvm.relax.expr_functor import PyExprMutator, mutator + + +@tvm.transform.module_pass(opt_level=0, name="LiftTIRGlobalBufferAlloc") +class LiftTIRGlobalBufferAlloc: + """A compiler pass that lifts TIR-level global allocation to Relax.""" + + def transform_module( + self, + mod: IRModule, + _ctx: tvm.transform.PassContext, + ) -> IRModule: + """IRModule-level transformation""" + return _TIRGlobalAllocRewriter(mod).transform() + + +@mutator +class _TIRGlobalAllocRewriter(PyExprMutator): + def __init__(self, mod: IRModule): + super().__init__(mod) + self.mod = mod + self.gv2new_tensor_sinfo: Dict[ # noqa: UP006 + tvm.ir.GlobalVar, + Tuple[tvm.ir.GlobalVar, List[relax.TensorType], tirx.PrimFunc], # noqa: UP006 + ] = {} + + def transform(self) -> IRModule: + """Entry point of the transformation""" + for g_var, func in self.mod.functions_items(): + if isinstance(func, tirx.PrimFunc): + updated_func, tensor_sinfo_list = remove_global_buf_alloc(func) + if len(tensor_sinfo_list) > 0: + new_gv = self.builder_.add_func(updated_func, g_var.name_hint) + self.gv2new_tensor_sinfo[g_var] = (new_gv, tensor_sinfo_list, func) + + self.mod = self.builder_.get() + for g_var, func in self.mod.functions_items(): + if isinstance(func, relax.Function): + updated_func = self.visit_expr(func) + updated_func = remove_all_unused(updated_func) + self.builder_.update_func(g_var, updated_func) + + mod = self.builder_.get() + return relax.transform.DeadCodeElimination()(mod) + + def visit_call_(self, call: relax.Call): + call = self.visit_expr_post_order(call) + if ( + call.op != tvm.ir.Op.get("relax.call_tir") + or call.args[0] not in self.gv2new_tensor_sinfo + ): + return call + + g_var = call.args[0] + new_gv, tensor_sinfo, func_before_update = self.gv2new_tensor_sinfo[g_var] + + assert len(call.ty_args) == 1 + if any(_has_symbolic_var(sinfo) for sinfo in tensor_sinfo): + tensor_sinfo, success = _resolve_tir_var_mapping(func_before_update, call, tensor_sinfo) + if not success: + # Cannot resolve TIR var mapping. Fall back to no lifting. + self.gv2new_tensor_sinfo.pop(g_var) + return call + + args = list(call.args) + args[0] = new_gv + if isinstance(call.ty_args[0], relax.TensorType): + new_call = relax.Call( + call.op, + args=args, + ty_args=[relax.TupleType(list(call.ty_args) + tensor_sinfo)], + attrs=call.attrs, + ) + emitted_tuple = self.builder_.emit(new_call) + return relax.TupleGetItem(emitted_tuple, 0) + assert isinstance(call.ty_args[0], relax.TupleType) + return relax.Call( + call.op, + args=args, + ty_args=[relax.TupleType(list(call.ty_args[0].fields) + tensor_sinfo)], + attrs=call.attrs, + ) + + +def remove_global_buf_alloc( + func: tirx.PrimFunc, +) -> Tuple[tirx.PrimFunc, List[relax.TensorType]]: # noqa: UP006 + """Remove the global buffer allocation for a given TIR PrimFunc.""" + assert isinstance(func.body, tirx.SBlockRealize) + params = list(func.params) + buffer_map = dict(func.buffer_map) + tensor_sinfo = [] + alloc_buffers = [] + + insertion_point = len(params) + while not isinstance(params[insertion_point - 1].ty, tvm.ir.PointerType): + insertion_point -= 1 + assert insertion_point >= 1 + + prev_root_block = func.body.block + for buf_alloc in func.body.block.alloc_buffers: + if buf_alloc.scope() == "global": + param = tirx.Var("var_" + buf_alloc.name, "handle") + params.insert(insertion_point, param) + insertion_point += 1 + buffer_map[param] = buf_alloc + tensor_sinfo.append(relax.TensorType(buf_alloc.shape, buf_alloc.dtype)) + else: + alloc_buffers.append(buf_alloc) + + if len(tensor_sinfo) == 0: + return func, [] + + assert len(prev_root_block.iter_vars) == 0 + assert len(prev_root_block.reads) == 0 + assert len(prev_root_block.writes) == 0 + assert len(prev_root_block.match_buffers) == 0 + assert prev_root_block.name_hint == "root" + assert prev_root_block.init is None + root_block = tirx.SBlock( + iter_vars=[], + reads=[], + writes=[], + name_hint="root", + body=prev_root_block.body, + alloc_buffers=alloc_buffers, + annotations=prev_root_block.annotations, + ) + + updated_func = tirx.PrimFunc( + params=params, + body=tirx.SBlockRealize(iter_values=[], predicate=True, block=root_block), + ret_type=func.ret_type, + buffer_map=buffer_map, + attrs=func.attrs, + ) + return updated_func, tensor_sinfo + + +def _has_symbolic_var(tensor_sinfo: relax.TensorType) -> bool: + assert isinstance(tensor_sinfo.shape, relax.ShapeExpr) + for dim in tensor_sinfo.shape.values: + if not isinstance(dim, tirx.IntImm): + return True + return False + + +def _resolve_tir_var_mapping( + func: tirx.PrimFunc, + call: relax.Call, + tensor_sinfo: List[relax.TensorType], # noqa: UP006 +) -> Tuple[List[relax.TensorType], bool]: # noqa: UP006 + """Resolve the TIR symbolic var relationship across sides of PrimFunc and Relax Function""" + var_map: Dict[tirx.Var, tirx.Expr] = {} # noqa: UP006 + + n_arg = len(call.args[1].fields) + for i in range(n_arg): + buffer_shape = func.buffer_map[func.params[i]].shape + arg_shape = call.args[1][i].ty.shape.values + assert len(buffer_shape) == len(arg_shape) + for v_l, v_r in zip(buffer_shape, arg_shape): + if isinstance(v_l, tirx.Var): + var_map[v_l] = v_r + elif not isinstance(v_l, tirx.IntImm): + return [], False + + ret_tensors = call.ty_args[0] + ret_tensors = ( + [ret_tensors] if isinstance(ret_tensors, relax.TensorType) else list(ret_tensors.fields) + ) + for i, ret_tensor in enumerate(ret_tensors): + buffer_shape = func.buffer_map[func.params[n_arg + i]].shape + ret_tensor_shape = ret_tensor.shape.values + assert len(buffer_shape) == len(ret_tensor_shape) + for v_l, v_r in zip(buffer_shape, ret_tensor_shape): + if isinstance(v_l, tirx.Var): + var_map[v_l] = v_r + elif not isinstance(v_l, tirx.IntImm): + return [], False + + updated_tensor_sinfo = [] + for sinfo in tensor_sinfo: + if not _has_symbolic_var(sinfo): + updated_tensor_sinfo.append(sinfo) + continue + new_shape = [] + for dim in sinfo.shape.values: + new_shape.append(tirx.stmt_functor.substitute(dim, var_map)) + updated_tensor_sinfo.append(relax.TensorType(new_shape, sinfo.dtype)) + return updated_tensor_sinfo, True diff --git a/python/mlc_llm/compiler_pass/low_batch_specialization.py b/python/mlc_llm/compiler_pass/low_batch_specialization.py new file mode 100644 index 0000000..0a7045b --- /dev/null +++ b/python/mlc_llm/compiler_pass/low_batch_specialization.py @@ -0,0 +1,64 @@ +"""A compiler pass that dispatch low-batch-gemm to gemv schedule.""" + +import tvm +import tvm_ffi +from tvm import tirx +from tvm.ir.module import IRModule +from tvm.s_tir import dlight as dl + + +@tvm.transform.module_pass(opt_level=0, name="LowBatchGemvSpecialize") +class LowBatchGemvSpecialize: + """A compiler pass that dispatch low-batch-gemm to gemv schedule.""" + + def transform_module( + self, + mod: IRModule, + _ctx: tvm.transform.PassContext, + ) -> IRModule: + """IRModule-level transformation""" + for g_var, func in mod.functions_items(): + if isinstance(func, tirx.PrimFunc): + low_batch_range = [2, 8] + buckets = [2, 4] + low_batch_funcs = [] + for bucket in buckets: + low_batch_mod = IRModule({}) + low_batch_mod["main"] = func + low_batch_mod = dl.ApplyDefaultSchedule( + dl.gpu.LowBatchGEMV(bucket), + )(low_batch_mod) + low_batch_funcs.append(low_batch_mod["main"]) + if any( + tvm_ffi.structural_equal(low_batch_func, func) + for low_batch_func in low_batch_funcs + ): + continue + buffers = func.buffer_map.values() + shapes = [buffer.shape for buffer in buffers] + symbolic_vars = set( + expr for shape in shapes for expr in shape if isinstance(expr, tirx.Var) + ) + if len(symbolic_vars) != 1: + continue + gemm_mod = IRModule({}) + gemm_mod["main"] = func + gemm_mod = dl.ApplyDefaultSchedule( + dl.gpu.Matmul(), + )(gemm_mod) + gemm_func = gemm_mod["main"] + sym_var = next(iter(symbolic_vars)) + body = gemm_func.body + for i, range_limit in reversed(list(enumerate(low_batch_range))): + body = tirx.IfThenElse( + tirx.op.tvm_thread_invariant(sym_var <= range_limit), + low_batch_funcs[i].body, + body, + ) + body = tirx.SBlock([], [], [], "root", body) + body = tirx.SBlockRealize([], True, body) + new_func = func.with_body(body) + new_func = new_func.with_attr("tirx.is_scheduled", 1) + new_func = new_func.with_attr("tirx.HoistIfThenElseExprWithBlock", 1) + mod.update_func(g_var, new_func) + return mod diff --git a/python/mlc_llm/compiler_pass/pipeline.py b/python/mlc_llm/compiler_pass/pipeline.py new file mode 100644 index 0000000..9af3606 --- /dev/null +++ b/python/mlc_llm/compiler_pass/pipeline.py @@ -0,0 +1,209 @@ +"""The compilation pipeline for LLM applications.""" + +from pathlib import Path +from typing import Any, Dict, List, Optional # noqa: UP035 + +import tvm +from tvm import IRModule +from tvm.relax import register_pipeline +from tvm.relax.frontend import nn +from tvm.s_tir import dlight as dl + +from mlc_llm.interface.compiler_flags import IPCAllReduceStrategyType +from mlc_llm.support import logging + +from .attach_cuda_graph_alloc_init_func import AttachCUDAGraphAllocInitFunc +from .attach_embedding_allocator import AttachAllocEmbeddingTensorFunc +from .attach_logit_processor import AttachLogitProcessFunc +from .attach_sampler import AttachGPUSamplingFunc +from .attach_softmax_with_temperature import AttachSoftmaxWithTemperature +from .attach_spec_decode_aux_funcs import AttachSpecDecodeAuxFuncs +from .attach_support_info import ( + AttachAdditionalPrimFuncs, + AttachCUDAGraphSymbolicCaptureHints, + AttachMemoryPlanAttr, + AttachPipelineParallelStages, + AttachSequenceLengthPaddingFactor, + AttachVariableBounds, +) +from .blas_dispatch import BLASDispatch +from .clean_up_tir_attrs import CleanUpTIRAttrs +from .dispatch_kv_cache_creation import DispatchKVCacheCreation +from .dispatch_triton_kernel import DispatchTritonKernel +from .estimate_memory_usage import AttachMetadataWithMemoryUsage +from .fuse_add_norm import FuseAddRMSNorm +from .fuse_dequantize_matmul_ewise import FuseDequantizeMatmulEwise +from .fuse_dequantize_take import FuseDequantizeTake +from .fuse_dequantize_transpose import FuseDequantizeTranspose +from .fuse_ft_dequantize_matmul_epilogue import FuseFTDequantizeEpilogue +from .fuse_transpose_matmul import FuseTransposeMatmul +from .lift_global_buffer_alloc import LiftTIRGlobalBufferAlloc +from .low_batch_specialization import LowBatchGemvSpecialize +from .pipeline_parallel_rewrite import PipelineParallelRewrite +from .scatter_tuple_get_item import ScatterTupleGetItem + +logger = logging.getLogger(__name__) + + +@tvm.transform.module_pass(opt_level=0, name="_LogProgress") +class _LogProgress: + """A dummy compiler pass that does nothing but logging.""" + + def __init__(self, *args): + self.args = args + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """A dummy transformation""" + logger.info(*self.args) + return mod + + +@tvm.transform.module_pass(opt_level=0, name="DebugDump") +class _DebugDump: + """A dummy compiler pass that does nothing but logging. + Only enabled when debug_dump is not None""" + + def __init__(self, file_name: str, file_path: Optional[Path], show_meta: bool = False): + self.file_name = file_name + self.file_path = file_path + self.show_meta = show_meta + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """A dummy transformation that dumps the module to file""" + if self.file_path is not None: + # NOTE: We use debug level here to avoid spamming the console + logger.debug("Dumping IR to %s", self.file_path / self.file_name) + with open(self.file_path / self.file_name, "w", encoding="utf-8") as f: + f.write(mod.script(show_meta=self.show_meta)) + return mod + + +@register_pipeline("mlc_llm") +def _mlc_llm_pipeline( + target: tvm.target.Target, + flashinfer: bool = False, + cublas_gemm: bool = False, + faster_transformer: bool = False, + allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE, + variable_bounds: Optional[Dict[str, int]] = None, # noqa: UP006 + cuda_graph_symbolic_capture_hints: Optional[Dict[str, List[str]]] = None, # noqa: UP006 + additional_tirs: Optional[Dict[str, tvm.tirx.PrimFunc]] = None, # noqa: UP006 + metadata: Optional[Dict[str, Any]] = None, # noqa: UP006 + ext_mods: Optional[List[nn.ExternModule]] = None, # noqa: UP006 + debug_dump: Optional[Path] = None, +): + variable_bounds = variable_bounds or {} + cuda_graph_symbolic_capture_hints = cuda_graph_symbolic_capture_hints or {} + additional_tirs = additional_tirs or {} + metadata = metadata or {} + ext_mods = ext_mods or [] + tensor_parallel_shards = metadata.get("tensor_parallel_shards", 1) + + @tvm.transform.module_pass(opt_level=0) + def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule: + seq = tvm.transform.Sequential( + [ + # Phase 0. Add additional information for compilation and remove unused Relax func + DispatchKVCacheCreation(target, flashinfer, metadata), + AttachSoftmaxWithTemperature(target, metadata), + AttachVariableBounds(variable_bounds), + AttachCUDAGraphSymbolicCaptureHints(cuda_graph_symbolic_capture_hints), + AttachPipelineParallelStages(metadata["pipeline_parallel_stages"]), + AttachLogitProcessFunc(target), + AttachAdditionalPrimFuncs(additional_tirs), + AttachAllocEmbeddingTensorFunc(metadata), + AttachGPUSamplingFunc(target, variable_bounds), + AttachSpecDecodeAuxFuncs(tensor_parallel_shards), + AttachMemoryPlanAttr(), + AttachSequenceLengthPaddingFactor(target, metadata), + tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)), + _DebugDump("debug-phase0.py", debug_dump, show_meta=False), + # Phase 1. Passes on high-level operator graph + _LogProgress("Running TVM Relax graph-level optimizations"), + DispatchTritonKernel(target), + FuseFTDequantizeEpilogue(), + FuseDequantizeTranspose(), + BLASDispatch(target) if cublas_gemm else tvm.transform.Sequential([]), + ( + FuseAddRMSNorm(target=target) + if target.kind.name != "llvm" + else tvm.transform.Sequential([]) + ), + FuseTransposeMatmul(), + _DebugDump("debug-phase1.py", debug_dump, show_meta=False), + # Phase 2. Lowering to TIR, inherited TVM Relax's official "zero" pipeline + _LogProgress("Lowering to TVM TIR kernels"), + tvm.relax.backend.DispatchSampling(), + tvm.relax.backend.DispatchSortScan(), + tvm.relax.transform.LegalizeOps(), + tvm.relax.transform.AnnotateTIROpPattern(), + tvm.relax.transform.FoldConstant(), + tvm.relax.transform.FuseOps(), + tvm.relax.transform.FuseTIR(), + _DebugDump("debug-phase2.py", debug_dump, show_meta=False), + # Phase 3. Passes on TIR + _LogProgress("Running TVM TIR-level optimizations"), + FuseDequantizeMatmulEwise(), + FuseDequantizeTake(), + tvm.relax.transform.DeadCodeElimination(), + CleanUpTIRAttrs(["op_pattern"]), + _DebugDump("debug-phase3.py", debug_dump, show_meta=False), + # Phase 4. Low-level Optimizations + _LogProgress("Running TVM Dlight low-level optimizations"), + LowBatchGemvSpecialize(), + ( + dl.ApplyDefaultSchedule( + dl.gpu.Matmul(), + dl.gpu.GEMV(), + dl.gpu.Reduction(), + dl.gpu.GeneralReduction(), + dl.gpu.Fallback(), + ) + if target.kind.name != "llvm" + else dl.ApplyDefaultSchedule( + dl.cpu.GEMV(), + ) + ), + _DebugDump("debug-phase4.py", debug_dump, show_meta=False), + _LogProgress("Lowering to VM bytecode"), + ( + LiftTIRGlobalBufferAlloc() + if target.kind.name != "llvm" + else tvm.transform.Sequential([]) + ), + ( + tvm.tirx.transform.ForceNarrowIndexToInt32() + if target.kind.name != "cuda" + else tvm.transform.Sequential([]) + ), + ScatterTupleGetItem(), + PipelineParallelRewrite(), + tvm.relax.transform.RewriteDataflowReshape(), + tvm.relax.transform.ToNonDataflow(), + tvm.relax.transform.RemovePurityChecking(), + tvm.relax.transform.CallTIRRewrite(), + ( + tvm.relax.transform.IPCAllReduceRewrite(allreduce_strategy) + if allreduce_strategy != IPCAllReduceStrategyType.NONE + else tvm.transform.Sequential([]) + ), + tvm.relax.transform.StaticPlanBlockMemory(), + AttachMetadataWithMemoryUsage(metadata), + _DebugDump("debug-phase5.py", debug_dump, show_meta=False), + tvm.relax.transform.RewriteCUDAGraph(), + AttachCUDAGraphAllocInitFunc(), + tvm.relax.transform.LowerGPUIPCAllocStorage(), + tvm.relax.transform.LowerAllocTensor(), + tvm.relax.transform.KillAfterLastUse(), + tvm.relax.transform.LowerRuntimeBuiltin(), + tvm.relax.transform.VMShapeLower(), + tvm.relax.transform.AttachGlobalSymbol(), + _LogProgress("Compiling external modules"), + tvm.relax.transform.AttachExternModules(ext_mods), + _LogProgress("Compilation complete! Exporting to disk"), + ] + ) + mod = seq(mod) + return mod + + return _pipeline diff --git a/python/mlc_llm/compiler_pass/pipeline_parallel_rewrite.py b/python/mlc_llm/compiler_pass/pipeline_parallel_rewrite.py new file mode 100644 index 0000000..d16d24a --- /dev/null +++ b/python/mlc_llm/compiler_pass/pipeline_parallel_rewrite.py @@ -0,0 +1,399 @@ +"""A compiler pass that rewrites IR for pipeline parallelism.""" + +from typing import Dict, List, Optional, Tuple # noqa: UP035 + +import tvm +from tvm import relax, tirx +from tvm.ir.module import IRModule +from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor + + +@tvm.transform.module_pass(opt_level=0, name="PipelineParallelRewrite") +class PipelineParallelRewrite: + """A compiler pass that rewrites IR for pipeline parallelism.""" + + def transform_module( + self, + mod: IRModule, + _ctx: tvm.transform.PassContext, + ) -> IRModule: + """IRModule-level transformation""" + return _PipelineParallelRewriter(mod.clone()).transform() + + +@mutator +class _PipelineParallelRewriter(PyExprMutator): + def __init__(self, mod: IRModule): + super().__init__(mod) + self.mod = mod + self.old_packed_params_var: relax.Var + self.new_main_packed_params_var: relax.Var + self.new_stage_func_packed_params: relax.Var + self.undefined_shape_vars_remap: Dict[tirx.Var, tirx.Var] # noqa: UP006 + self.undefined_param_shape_vars_remap: Dict[tirx.Var, tirx.Var] # noqa: UP006 + + def transform(self) -> IRModule: + """Entry point of the transformation""" + for g_var, func in self.mod.functions_items(): + if not isinstance(func, relax.Function) or "pipeline_parallel_stages" not in func.attrs: + continue + num_stages = int(func.attrs["pipeline_parallel_stages"]) + if num_stages == 1: + continue + + pipeline_stages, stage_send_vars, stage_receive_vars = _extract_pipeline_stages(func) + assert len(pipeline_stages) == num_stages, ( + "Number of pipeline stages mismatches: " + f"expecting {num_stages} stages, but {len(pipeline_stages)} are found in the IR." + ) + + required_func_params = _analyze_required_func_params(pipeline_stages, func.params) + + assert "num_input" in func.attrs + num_input = int(func.attrs["num_input"]) + assert ( + len(func.params) == num_input + 1 + and isinstance(func.params[num_input], relax.Var) + and func.params[num_input].name_hint == "packed_params" + ), 'Only the extra "packed_params" parameter is allowed' + self.old_packed_params_var = func.params[num_input] + self.new_main_packed_params_var = relax.Var("packed_params", relax.ObjectType()) + for required_params in required_func_params: + for i, param in enumerate(required_params): + if param.same_as(self.old_packed_params_var): + required_params.pop(i) + break + func_output = func.body.body + assert isinstance(func_output, relax.Var) + + stage_func_gvs = [] + caller_args_list = [] + for i in range(num_stages): + stage_func_gv, caller_args = self._create_stage_func( + g_var.name_hint + f"_stage{i}", + pipeline_stages[i], + required_func_params[i], + stage_receive_vars[i], + stage_send_vars[i], + func.attrs, + func_output=func_output if i == num_stages - 1 else None, + ) + stage_func_gvs.append(stage_func_gv) + caller_args_list.append(caller_args) + + # Create and update the entry function, which dispatches toz the stage functions + # according to the disco worker group id. + bb = relax.BlockBuilder() + params = [*list(func.params[:-1]), self.new_main_packed_params_var] + with bb.function(g_var.name_hint, params=params): + dispatch_func_args = [] + for stage_func_gv, caller_args in zip(stage_func_gvs, caller_args_list): + dispatch_func_args.append([stage_func_gv, *caller_args]) + output = bb.emit( + relax.op.call_builtin_with_ctx( + "mlc.multi_gpu.DispatchFunctionByGroup", + args=[dispatch_func_args], + ty_args=relax.ObjectType(), + ) + ) + dispatch_func_gv = bb.emit_func_output(output) + dispatch_func = bb.finalize()[dispatch_func_gv] + self.builder_.update_func(g_var, dispatch_func) + + return self.builder_.finalize() + + def _create_stage_func( + self, + func_name: str, + stage_bindings: List[relax.Binding], # noqa: UP006 + required_func_params: List[relax.Var], # noqa: UP006 + stage_receive_vars: List[relax.Var], # noqa: UP006 + stage_send_vars: List[relax.Var], # noqa: UP006 + func_attrs: tvm.ir.DictAttrs, + func_output: Optional[relax.Var], + ) -> Tuple[tvm.ir.GlobalVar, List[relax.Expr]]: # noqa: UP006 + self.undefined_shape_vars_remap = {} + self.undefined_param_shape_vars_remap = {} + + # Prepare the func parameters (except the shape variables and packed params) + params, args = self._prepare_stage_func_params_and_args(required_func_params) + for new_param, old_param in zip(params, required_func_params): + self.set_var_remap(old_param, new_param) + # Create new packed params + self.new_stage_func_packed_params = relax.Var("packed_params", relax.ObjectType()) + self.set_var_remap(self.old_packed_params_var, self.new_stage_func_packed_params) + + new_func_outputs = [] + with self.builder_.function(func_name, pure=False): + with self.builder_.dataflow(): + # Emit the tensors received from last stage. + for receive_var in stage_receive_vars: + new_receive_var = self.builder_.emit( + relax.call_dps_packed( + "runtime.disco.recv_from_prev_group", + args=[], + out_ty=self._update_struct_info(receive_var.ty), + ), + name_hint=receive_var.name_hint, + ) + self.set_var_remap(receive_var, new_receive_var) + # Process the bindings in this stage. + for stage_binding in stage_bindings: + if stage_binding.var in stage_send_vars or stage_binding.var.same_as( + func_output + ): + assert isinstance(stage_binding, relax.VarBinding) + new_var = self.builder_.emit_output( + self.visit_expr(stage_binding.value), + name_hint=stage_binding.var.name_hint, + ) + self.set_var_remap(stage_binding.var, new_var) + new_func_outputs.append(new_var) + else: + self.visit_binding(stage_binding) + # Emit the calls to send tensors to the next stage. + for send_var in stage_send_vars: + new_send_var = self.get_var_remap(send_var) + self.builder_.emit( + relax.Call( + relax.ExternFunc("runtime.disco.send_to_next_group"), + args=[new_send_var], + ty_args=None, + ) + ) + # Create the param for the shape variables. + shape_var_params = [] + shape_var_args = [] + for ( + shape_var_arg, + shape_var_param, + ) in self.undefined_shape_vars_remap.items(): + if shape_var_arg not in self.undefined_param_shape_vars_remap: + shape_var_params.append(shape_var_param) + shape_var_args.append(shape_var_arg) + params.append(relax.Var("s", relax.ShapeType(shape_var_params))) + args.append(relax.ShapeExpr(shape_var_args)) + # Add the packed params. + params.append(self.new_stage_func_packed_params) + args.append(self.new_main_packed_params_var) + # Conclude the function. + if func_output is not None: + assert len(new_func_outputs) == 1 + new_gv = self.builder_.emit_func_output( + ( + new_func_outputs[0] + if len(new_func_outputs) == 1 + and isinstance(new_func_outputs[0].ty, relax.TupleType) + else new_func_outputs + ), + params=params, + ) + + new_func = ( + self.builder_.get()[new_gv] + .with_attrs(func_attrs) + .with_attr("num_input", len(params) - 1) + .without_attr("global_symbol") + .without_attr("pipeline_parallel_stages") + ) + self.builder_.update_func(new_gv, new_func) + return new_gv, args + + def visit_var_binding_(self, binding: relax.VarBinding) -> None: + if not isinstance(binding.value, relax.TupleGetItem): + super().visit_var_binding_(binding) + return + + tuple_value = self.visit_expr(binding.value.tuple_value) + if not tuple_value.same_as(self.new_stage_func_packed_params): + super().visit_var_binding_(binding) + return + + assert isinstance(binding.var.ty, relax.TensorType) + cur_num_undefined_param_shape_vars = len(self.undefined_param_shape_vars_remap) + new_tensor_struct_info = self._update_struct_info( + binding.var.ty, self.undefined_param_shape_vars_remap + ) + has_new_undefined_shape_var = ( + len(self.undefined_param_shape_vars_remap) != cur_num_undefined_param_shape_vars + ) + self.undefined_shape_vars_remap = { + **self.undefined_shape_vars_remap, + **self.undefined_param_shape_vars_remap, + } + ret_sinfo = ( + new_tensor_struct_info if not has_new_undefined_shape_var else relax.ObjectType() + ) + call = relax.call_pure_packed( + "vm.builtin.tuple_getitem", + self.new_stage_func_packed_params, + relax.prim_value(binding.value.index), + ty_args=ret_sinfo, + ) + new_binding_var = self.builder_.emit(call, binding.var.name_hint) + if has_new_undefined_shape_var: + new_binding_var = self.builder_.match_cast( + new_binding_var, new_tensor_struct_info, binding.var.name_hint + "_cast" + ) + self.set_var_remap(binding.var, new_binding_var) + + def visit_call_(self, call: relax.Call) -> relax.Call: + call = super().visit_call_(call) + return relax.Call( + call.op, + call.args, + call.attrs, + ty_args=[self._update_struct_info(struct_info) for struct_info in call.ty_args], + ) + + def _prepare_stage_func_params_and_args( + self, + required_func_params: List[relax.Var], # noqa: UP006 + ) -> Tuple[List[relax.Var], List[relax.Expr]]: # noqa: UP006 + params: List[relax.Var] = [] # noqa: UP006 + args: List[relax.Expr] = [] # noqa: UP006 + for required_param in required_func_params: + struct_info = self._update_struct_info(required_param.ty) + params.append(relax.Var(required_param.name_hint, struct_info)) + args.append(required_param) + + return params, args + + def _update_struct_info( + self, + struct_info: relax.Type, + undefined_var_remap: Optional[Dict[tirx.Var, tirx.Var]] = None, # noqa: UP006 + ) -> relax.Type: + if undefined_var_remap is None: + undefined_var_remap = self.undefined_shape_vars_remap + if isinstance(struct_info, relax.TensorType): + return ( + relax.TensorType( + self._update_shape(struct_info.shape.values, undefined_var_remap), + struct_info.dtype, + ) + if struct_info.shape is not None and isinstance(struct_info.shape, relax.ShapeExpr) + else struct_info + ) + if isinstance(struct_info, relax.ShapeType): + return ( + relax.ShapeType(self._update_shape(struct_info.values, undefined_var_remap)) + if struct_info.values is not None + else struct_info + ) + if isinstance(struct_info, relax.ObjectType): + return relax.ObjectType() + if isinstance(struct_info, relax.TupleType): + return relax.TupleType( + [self._update_struct_info(field_sinfo) for field_sinfo in struct_info.fields] + ) + return struct_info + + def _copy_undefined_var( + self, + expr: tirx.Expr, + undefined_var_remap: Dict[tirx.Var, tirx.Var], # noqa: UP006 + ) -> None: + def _visit_expr(e: tirx.Expr) -> None: + if isinstance(e, tirx.Var) and e not in undefined_var_remap: + new_var = tirx.Var(e.name, e.ty) + undefined_var_remap[e] = new_var + + tirx.stmt_functor.post_order_visit(expr, _visit_expr) + + def _update_shape( + self, + shape: List[tirx.Expr], # noqa: UP006 + undefined_var_remap: Dict[tirx.Var, tirx.Var], # noqa: UP006 + ) -> List[tirx.Expr]: # noqa: UP006 + new_shape = [] + for v in shape: + self._copy_undefined_var(v, undefined_var_remap) + new_shape.append(tirx.stmt_functor.substitute(v, undefined_var_remap)) + return new_shape + + +def _extract_pipeline_stages( + func: relax.Function, +) -> Tuple[List[List[relax.Binding]], List[List[relax.Var]], List[List[relax.Var]]]: # noqa: UP006 + pipeline_stages: List[List[relax.Binding]] = [] # noqa: UP006 + stage_send_vars: List[List[relax.Var]] = [] # noqa: UP006 + stage_receive_vars: List[List[relax.Var]] = [] # noqa: UP006 + + # Requiring that the function has only one body block which is a dataflow block + assert isinstance(func.body, relax.SeqExpr) + assert len(func.body.blocks) == 1 + assert isinstance(func.body.blocks[0], relax.DataflowBlock) + bindings = func.body.blocks[0].bindings + + boundary_var = None + current_stage_bindings: List[relax.Binding] = [] # noqa: UP006 + current_stage_receive_vars: List[relax.Var] = [] # noqa: UP006 + for binding in bindings: + if ( + isinstance(binding, relax.VarBinding) + and isinstance(binding.value, relax.Call) + and binding.value.op == tvm.ir.Op.get("relax.call_pure_packed") + and binding.value.args[0].global_symbol == "mlc.pipeline_parallel_stage_boundary" + ): + assert len(current_stage_bindings) > 0 + pipeline_stages.append(current_stage_bindings) + assert all(receive_var is not None for receive_var in current_stage_receive_vars) + stage_receive_vars.append(current_stage_receive_vars) + args = binding.value.args[1:] + assert len(args) >= 1 and all(isinstance(arg, relax.Var) for arg in args) + stage_send_vars.append(list(args)) + + boundary_var = binding.var + current_stage_bindings = [] + current_stage_receive_vars = [boundary_var] if len(args) == 1 else [None for _ in args] + elif ( + isinstance(binding, relax.VarBinding) + and isinstance(binding.value, relax.TupleGetItem) + and binding.value.tuple_value.same_as(boundary_var) + ): + current_stage_receive_vars[binding.value.index] = binding.var + else: + current_stage_bindings.append(binding) + + assert len(current_stage_bindings) > 0 + pipeline_stages.append(current_stage_bindings) + assert all(receive_var is not None for receive_var in current_stage_receive_vars) + stage_receive_vars.append(current_stage_receive_vars) + stage_send_vars.append([]) + + return pipeline_stages, stage_send_vars, stage_receive_vars + + +def _analyze_required_func_params( + pipeline_stages: List[List[relax.Binding]], # noqa: UP006 + func_params: List[relax.Var], # noqa: UP006 +) -> List[List[relax.Var]]: # noqa: UP006 + analyzer = _RequiredFuncParamAnalyzer(func_params) + required_func_params: List[List[relax.Var]] = [] # noqa: UP006 + for stage_bindings in pipeline_stages: + required_params: List[relax.Var] # noqa: UP006 + required_params = analyzer.run(stage_bindings) + required_func_params.append(required_params) + return required_func_params + + +@visitor +class _RequiredFuncParamAnalyzer(PyExprVisitor): + """The IR visitor which analyzes the required func parameters in each pipeline stage.""" + + def __init__(self, func_params: List[relax.Var]) -> None: # noqa: UP006 + self.func_params = set(func_params) + self.required_params: List[relax.Var] # noqa: UP006 + + def run(self, stage_bindings: List[relax.Binding]) -> List[relax.Var]: # noqa: UP006 + """Entry point of the visitor.""" + self.required_params = [] + for binding in stage_bindings: + self.visit_binding(binding) + return self.required_params + + def visit_var_(self, var: relax.Var) -> None: + if var in self.func_params: + if var not in self.required_params: + self.required_params.append(var) diff --git a/python/mlc_llm/compiler_pass/scatter_tuple_get_item.py b/python/mlc_llm/compiler_pass/scatter_tuple_get_item.py new file mode 100644 index 0000000..c434d62 --- /dev/null +++ b/python/mlc_llm/compiler_pass/scatter_tuple_get_item.py @@ -0,0 +1,49 @@ +"""A compiler pass that scatters TupleGetItem for lazy TupleGetItems.""" + +from typing import Dict # noqa: UP035 + +import tvm +from tvm import relax +from tvm.ir.module import IRModule +from tvm.relax.analysis import remove_all_unused +from tvm.relax.expr import Expr, Var +from tvm.relax.expr_functor import PyExprMutator, mutator + + +@tvm.transform.module_pass(opt_level=0, name="ScatterTupleGetItem") +class ScatterTupleGetItem: + """A compiler pass that scatters TupleGetItem for lazy TupleGetItems.""" + + def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule: + """IRModule-level transformation""" + return _Scatter(mod).transform() + + +@mutator +class _Scatter(PyExprMutator): + def __init__(self, mod: IRModule) -> None: + super().__init__(mod) + self.mod = mod + self.var_map: Dict[Var, Expr] = {} # noqa: UP006 + + def transform(self) -> IRModule: + """Entry point""" + for g_var, func in self.mod.functions_items(): + if isinstance(func, relax.Function): + updated_func = self.visit_expr(func) + updated_func = remove_all_unused(updated_func) + self.builder_.update_func(g_var, updated_func) + return self.builder_.get() + + def visit_var_binding_(self, binding: relax.VarBinding): + super().visit_var_binding_(binding) + if isinstance(binding.value, relax.TupleGetItem): + self.var_map[binding.var] = binding.value + + def visit_dataflow_var_(self, var: relax.DataflowVar) -> Expr: + if var in self.var_map: + new_var = self.builder_.emit(self.var_map[var], name_hint=var.name_hint) + self.set_var_remap(var, new_var) + self.var_map.pop(var) + return new_var + return var diff --git a/python/mlc_llm/contrib/__init__.py b/python/mlc_llm/contrib/__init__.py new file mode 100644 index 0000000..aa101df --- /dev/null +++ b/python/mlc_llm/contrib/__init__.py @@ -0,0 +1 @@ +"""Set of experimental components that yet to be matured.""" diff --git a/python/mlc_llm/contrib/embeddings/__init__.py b/python/mlc_llm/contrib/embeddings/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/contrib/embeddings/embeddings.py b/python/mlc_llm/contrib/embeddings/embeddings.py new file mode 100644 index 0000000..d3f2269 --- /dev/null +++ b/python/mlc_llm/contrib/embeddings/embeddings.py @@ -0,0 +1,186 @@ +"""The Python API for MLC Embeddings.""" + +import json +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple # noqa: UP035 + +import numpy as np +import tvm +import tvm_ffi +from tvm import relax +from tvm.contrib import tvmjs +from tvm.runtime import Device, Module +from tvm.runtime.vm import VirtualMachine + +from mlc_llm.serve import engine_utils +from mlc_llm.support.auto_device import detect_device +from mlc_llm.tokenizers import Tokenizer + + +def _extract_metadata(mod: Module): + return json.loads(VirtualMachine(mod, tvm.runtime.device("cpu"))["_metadata"]()) + + +def _load_params( + model_weight_path: str, + device: Device, + model_metadata: Dict[str, Any], # noqa: UP006 +) -> List[tvm.runtime.Tensor]: # noqa: UP006 + params, meta = tvmjs.load_tensor_cache(model_weight_path, device) + param_names = [param["name"] for param in model_metadata["params"]] + assert len(param_names) == meta["ParamSize"] + + plist = [] + for param_name in param_names: + plist.append(params[param_name]) + return plist + + +def _get_tvm_module( + model_weight_path: str, + lib_path: str, + device: Device, + instrument: tvm_ffi.Function = None, +): + ex = tvm.runtime.load_module(lib_path) + vm = relax.VirtualMachine(ex, device) + if instrument: + vm.set_instrument(instrument) + metadata = _extract_metadata(ex) + params = _load_params(model_weight_path, device, metadata) + return vm.module, params, metadata + + +class DefaultDebugInstrument: + """The default debug instrument to use if users don't specify + a customized one. + + This debug instrument will dump the arguments and output of each + VM Call instruction into a .npz file. It will also alert the user + if any function outputs are NaN or INF. + """ + + def __init__(self, debug_out: Path): + """Constructor + + Parameters + ---------- + debug_out : Path + the directory to dump the .npz files + """ + self.counter = 0 + self.first_nan_occurred = False + self.first_inf_occurred = False + self.debug_out = debug_out + debug_out.mkdir(exist_ok=True, parents=True) + + def reset(self, debug_out: Path): + """Reset the state of the Instrument class + + Parameters + ---------- + debug_out : Path + the directory to dump the .npz files + """ + self.counter = 0 + self.first_nan_occurred = False + self.first_inf_occurred = False + self.debug_out = debug_out + debug_out.mkdir(exist_ok=True, parents=True) + + def __call__(self, func, name, before_run, ret_val, *args): + # Determine what functions to look at + if before_run: # Whether before the function is called or after + return + if name.startswith("vm.builtin.") and "attention_with_fused_qkv" not in name: + return + + # Decide what to print or save about the function's arguments (where args[-1] is the + # buffer we write the result to) + func_name = f"f{self.counter}_{name}" + + # Save the arguments to npz + arg_dict = {} + for i, arg in enumerate(args): + if isinstance(arg, tvm.runtime.Tensor): + arg_dict[f"arg_{i}"] = arg.numpy() + + np.savez(self.debug_out / f"{func_name}.npz", **arg_dict) + + self.counter += 1 + + +class MLCEmbeddings: + """A class to embed queries using MLC LLM encoder models. + + Parameters + ---------- + model: str + The model folder after compiling with MLC-LLM build process. The parameter + can either be the model name with its quantization scheme + (e.g. ``Llama-2-7b-chat-hf-q4f16_1``), or a full path to the model + folder. In the former case, we will use the provided name to search + for the model folder over possible paths. + + model_lib_path : str + The full path to the model library file to use (e.g. a ``.so`` file). + + device : Optional[str] + The description of the device to run on. User should provide a string in the + form of 'device_name:device_id' or 'device_name', where 'device_name' is one of + 'cuda', 'metal', 'vulkan', 'rocm', 'opencl', 'auto' (automatically detect the + local device), and 'device_id' is the device id to run on. If no 'device_id' + is provided, it will be set to 0 by default. + + debug_dir: Path + The output folder to store the dumped debug files. If None, will not dump any debug files. + """ + + def __init__( + self, + model: str, + model_lib_path: str, + device: Optional[str] = "auto", + debug_dir: Optional[str] = None, + ): + self.device = detect_device(device) + instrument = DefaultDebugInstrument(Path(debug_dir)) if debug_dir else None + self.mod, self.params, self.metadata = _get_tvm_module( + model, model_lib_path, self.device, instrument + ) + self.model_path = model + self.tokenizer = Tokenizer(self.model_path) + self.prefill_func = self.mod["prefill"] + + def embed(self, queries: List[str]) -> tvm.runtime.Tensor: # noqa: UP006 + """ + Embeds a list of queries in a single batch. + + Parameters + ---------- + queries : List[str] + A list of queries to embed. + + Returns + ------- + List[float] + A list of embeddings for the queries. + """ + tokens, attention_mask = self._tokenize_queries(queries) + tokens_tvm = tvm.runtime.tensor(tokens.astype("int32"), device=self.device) + attention_mask_tvm = tvm.runtime.tensor(attention_mask.astype("int32"), device=self.device) + output = self.prefill_func(tokens_tvm, attention_mask_tvm, self.params) + return output + + def _tokenize_queries(self, queries: List[str]) -> Tuple[np.ndarray, np.ndarray]: # noqa: UP006 + tokens = engine_utils.process_prompts(queries, self.tokenizer.encode) + max_query_length = max(len(token_seq) for token_seq in tokens) + + token_inputs: np.ndarray = np.zeros((len(tokens), max_query_length), dtype=np.int32) + attention_mask: np.ndarray = np.zeros((len(tokens), max_query_length), dtype=np.int32) + + for i, token_seq in enumerate(tokens): + token_inputs[i, : len(token_seq)] = token_seq + attention_mask[i, : len(token_seq)] = 1 + + return token_inputs, attention_mask diff --git a/python/mlc_llm/contrib/embeddings/openai.py b/python/mlc_llm/contrib/embeddings/openai.py new file mode 100644 index 0000000..3abc99d --- /dev/null +++ b/python/mlc_llm/contrib/embeddings/openai.py @@ -0,0 +1,251 @@ +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from typing import List, Optional, Tuple # noqa: UP035 + +import numpy as np +from langchain.embeddings import OpenAIEmbeddings +from langchain_community.embeddings.openai import ( + async_embed_with_retry, + embed_with_retry, +) + +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +class MLCEmbeddings(OpenAIEmbeddings): + def _chunk_tokens(self, texts: Sequence[str]) -> Tuple[List[List], List[int]]: # noqa: UP006 + """Tokenize and chunk texts to fit in the model's context window.""" + if not self.embedding_ctx_length: + raise ValueError( + "embedding_ctx_length must be defined to use _get_len_safe_embeddings." + ) + + try: + import tiktoken + except ImportError as err: + raise ImportError( + "Could not import tiktoken python package. " + "This is needed in order to for OpenAIEmbeddings. " + "Please install it with `pip install tiktoken`." + ) from err + + tokens = [] + indices = [] + model_name = self.tiktoken_model_name or self.model + try: + encoding = tiktoken.encoding_for_model(model_name) + except KeyError: + logger.warning("Warning: model not found. Using cl100k_base encoding.") + model = "cl100k_base" + encoding = tiktoken.get_encoding(model) + for i, text in enumerate(texts): + if self.model.endswith("001"): + # See: https://github.com/openai/openai-python/issues/418#issuecomment-1525939500 + # replace newlines, which can negatively affect performance. + text = text.replace("\n", " ") + token = encoding.encode( + text, + allowed_special=self.allowed_special, + disallowed_special=self.disallowed_special, + ) + for j in range(0, len(token), self.embedding_ctx_length): + tokens.append(token[j : j + self.embedding_ctx_length]) + indices.append(i) + return tokens, indices + + def _batch_embed( + self, + inputs: Sequence, + *, + chunk_size: Optional[int] = None, # noqa: UP045 + ) -> List[List[float]]: # noqa: UP006 + batched_embeddings: List[List[float]] = [] # noqa: UP006 + _chunk_size = chunk_size or self.chunk_size + _iter: Iterable = range(0, len(inputs), _chunk_size) + if self.show_progress_bar: + try: + from tqdm import tqdm + + _iter = tqdm(_iter) + except ImportError: + pass + + for i in _iter: + response = embed_with_retry( + self, + input=inputs[i : i + _chunk_size], + **self._invocation_params, + ) + batched_embeddings.extend(r["embedding"] for r in response["data"]) + return batched_embeddings + + async def _abatch_embed( + self, + inputs: Sequence, + *, + chunk_size: Optional[int] = None, # noqa: UP045 + ) -> List[List[float]]: # noqa: UP006 + batched_embeddings: List[List[float]] = [] # noqa: UP006 + _chunk_size = chunk_size or self.chunk_size + _iter: Iterable = range(0, len(inputs), _chunk_size) + if self.show_progress_bar: + try: + from tqdm import tqdm + + _iter = tqdm(_iter) + except ImportError: + pass + + for i in _iter: + response = await async_embed_with_retry( + self, + input=inputs[i : i + _chunk_size], + **self._invocation_params, + ) + batched_embeddings.extend(r["embedding"] for r in response["data"]) + return batched_embeddings + + # please refer to + # https://github.com/openai/openai-cookbook/blob/main/examples/Embedding_long_inputs.ipynb + def _get_len_safe_embeddings( + self, + texts: List[str], # noqa: UP006 + *, + engine: str, + chunk_size: Optional[int] = None, # noqa: UP045 + ) -> List[List[float]]: # noqa: UP006 + tokens, indices = self._chunk_tokens(texts) + batched_embeddings = self._batch_embed(tokens, chunk_size=chunk_size) + results: List[List[List[float]]] = [[] for _ in range(len(texts))] # noqa: UP006 + num_tokens_in_batch: List[List[int]] = [[] for _ in range(len(texts))] # noqa: UP006 + for idx, tokens_i, batched_emb in zip(indices, tokens, batched_embeddings): + results[idx].append(batched_emb) + num_tokens_in_batch[idx].append(len(tokens_i)) + + embeddings = [] + empty_average = embed_with_retry( + self, + input="", + **self._invocation_params, + )["data"][0]["embedding"] + for _result, num_tokens in zip(results, num_tokens_in_batch): + if len(_result) == 0: + average = empty_average + else: + average = np.average(_result, axis=0, weights=num_tokens) + normalized = (average / np.linalg.norm(average)).tolist() + embeddings.append(normalized) + + return embeddings + + # please refer to + # https://github.com/openai/openai-cookbook/blob/main/examples/Embedding_long_inputs.ipynb + async def _aget_len_safe_embeddings( + self, + texts: List[str], # noqa: UP006 + *, + engine: str, + chunk_size: Optional[int] = None, # noqa: UP045 + ) -> List[List[float]]: # noqa: UP006 + tokens, indices = self._chunk_tokens(texts) + batched_embeddings = await self._abatch_embed(tokens, chunk_size=chunk_size) + + results: List[List[List[float]]] = [[] for _ in range(len(texts))] # noqa: UP006 + num_tokens_in_batch: List[List[int]] = [[] for _ in range(len(texts))] # noqa: UP006 + for idx, tokens_i, batched_emb in zip(indices, tokens, batched_embeddings): + results[idx].append(batched_emb) + num_tokens_in_batch[idx].append(len(tokens_i)) + + embeddings = [] + empty_average = ( + await async_embed_with_retry( + self, + input="", + **self._invocation_params, + ) + )["data"][0]["embedding"] + for _result, num_tokens in zip(results, num_tokens_in_batch): + if len(_result) == 0: + average = empty_average + else: + average = np.average(_result, axis=0, weights=num_tokens) + normalized = (average / np.linalg.norm(average)).tolist() + embeddings.append(normalized) + + return embeddings + + def embed_documents( + self, + texts: List[str], # noqa: UP006 + chunk_size: Optional[int] = None, # noqa: UP045 + ) -> List[List[float]]: # noqa: UP006 + """Call out to OpenAI's embedding endpoint for embedding search docs. + + Args: + texts: The list of texts to embed. + chunk_size: The chunk size of embeddings. If None, will use the chunk size + specified by the class. + + Returns: + List of embeddings, one for each text. + """ + # NOTE: to keep things simple, as long as the embedding_ctx_length is defined, + # we assume the list may contain texts longer than the maximum context and + # use length-safe embedding function. + if self.embedding_ctx_length: + return self._get_len_safe_embeddings( + texts, engine=self.deployment, chunk_size=chunk_size + ) + + embeddings = self._batch_embed(texts, chunk_size=chunk_size) + return [(np.array(e) / np.linalg.norm(e)).tolist() for e in embeddings] + + async def aembed_documents( + self, + texts: List[str], # noqa: UP006 + chunk_size: Optional[int] = 0, # noqa: UP045 + ) -> List[List[float]]: # noqa: UP006 + """Call out to OpenAI's embedding endpoint async for embedding search docs. + + Args: + texts: The list of texts to embed. + chunk_size: The chunk size of embeddings. If None, will use the chunk size + specified by the class. + + Returns: + List of embeddings, one for each text. + """ + # NOTE: to keep things simple, as long as the embedding_ctx_length is defined, + # we assume the list may contain texts longer than the maximum context and + # use length-safe embedding function. + if self.embedding_ctx_length: + return await self._aget_len_safe_embeddings(texts, engine=self.deployment) + + embeddings = await self._abatch_embed(texts, chunk_size=chunk_size) + return [(np.array(e) / np.linalg.norm(e)).tolist() for e in embeddings] + + def embed_query(self, text: str) -> List[float]: # noqa: UP006 + """Call out to OpenAI's embedding endpoint for embedding query text. + + Args: + text: The text to embed. + + Returns: + Embedding for the text. + """ + return self.embed_documents([text])[0] + + async def aembed_query(self, text: str) -> List[float]: # noqa: UP006 + """Call out to OpenAI's embedding endpoint async for embedding query text. + + Args: + text: The text to embed. + + Returns: + Embedding for the text. + """ + embeddings = await self.aembed_documents([text]) + return embeddings[0] diff --git a/python/mlc_llm/conversation_template/__init__.py b/python/mlc_llm/conversation_template/__init__.py new file mode 100644 index 0000000..6a58d62 --- /dev/null +++ b/python/mlc_llm/conversation_template/__init__.py @@ -0,0 +1,38 @@ +"""Global namespace of conversation template registry""" + +# TODO(mlc-team): move conversation template apply to this namespace +# decouple conversation template apply from the conversation protocol +# data structure + +# model preset templates +from . import ( + cohere, + deepseek, + dolly, + gemma, + glm, + gorilla, + gpt, + hermes, + llama, + llava, + llm_jp, + ministral3, + ministral3_reasoning, + mistral, + nemotron, + oasst, + olmo, + olmo2, + orion, + phi, + qwen2, + qwen3, + qwen3_5, + redpajama, + rwkv, + stablelm, + tinyllama, + wizardlm, +) +from .registry import ConvTemplateRegistry diff --git a/python/mlc_llm/conversation_template/cohere.py b/python/mlc_llm/conversation_template/cohere.py new file mode 100644 index 0000000..d12cad6 --- /dev/null +++ b/python/mlc_llm/conversation_template/cohere.py @@ -0,0 +1,27 @@ +"""Cohere default templates""" + + +# Referred from: https://huggingface.co/CohereForAI/aya-23-8B/blob/main/tokenizer_config.json + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Aya-23 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="aya-23", + system_template=f"<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{MessagePlaceholders.SYSTEM.value}<|END_OF_TURN_TOKEN|>", + system_message="You are Command-R, a brilliant, sophisticated, AI-assistant trained to assist human users by providing thorough responses.", # noqa: E501 + roles={ + "user": "<|START_OF_TURN_TOKEN|><|USER_TOKEN|>", + "assistant": "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", + }, + seps=["<|END_OF_TURN_TOKEN|>"], + role_content_sep="", + role_empty_sep="", + system_prefix_token_ids=[5], + stop_str=["<|END_OF_TURN_TOKEN|>"], + stop_token_ids=[6, 255001], + ) +) diff --git a/python/mlc_llm/conversation_template/deepseek.py b/python/mlc_llm/conversation_template/deepseek.py new file mode 100644 index 0000000..a1cca56 --- /dev/null +++ b/python/mlc_llm/conversation_template/deepseek.py @@ -0,0 +1,85 @@ +"""Deepseek default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Deepseek +ConvTemplateRegistry.register_conv_template( + Conversation( + name="deepseek", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + system_prefix_token_ids=[100000], + roles={"user": "User", "assistant": "Assistant"}, + seps=["\n\n", "<｜end▁of▁sentence｜>"], # noqa: RUF001 + role_content_sep=": ", + role_empty_sep=":", + stop_str=["<｜end▁of▁sentence｜>"], # noqa: RUF001 + stop_token_ids=[100001], + ) +) + +# Deepseek V2 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="deepseek_v2", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + system_prefix_token_ids=[100000], + roles={"user": "User", "assistant": "Assistant"}, + seps=["\n\n", "<｜end▁of▁sentence｜>"], # noqa: RUF001 + role_content_sep=": ", + role_empty_sep=":", + stop_str=["<｜end▁of▁sentence｜>"], # noqa: RUF001 + stop_token_ids=[100001], + ) +) + +# DeepSeek-V3 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="deepseek_v3", + system_template=f"<｜begin▁of▁sentence｜>{MessagePlaceholders.SYSTEM.value}", # noqa: RUF001 + system_message="You are Deepseek-V3, an AI assistant created exclusively by the Chinese " + "Company DeepSeek. You'll provide helpful, harmless, and detailed responses to all " + "user inquiries.", + roles={"user": "<｜User｜>", "assistant": "<｜Assistant｜>"}, # noqa: RUF001 + seps=["", "<｜end▁of▁sentence｜>"], # noqa: RUF001 + role_content_sep="", + role_empty_sep="", + stop_token_ids=[1], + ) +) + +# DeepSeek-R1-Distill-Qwen +ConvTemplateRegistry.register_conv_template( + Conversation( + name="deepseek_r1_qwen", + system_template=f"<｜begin▁of▁sentence｜>{MessagePlaceholders.SYSTEM.value}", # noqa: RUF001 + system_message="You are Deepseek-R1, an AI assistant created exclusively by the Chinese " + "Company DeepSeek. You'll provide helpful, harmless, and detailed responses to all " + "user inquiries.", + roles={"user": "<｜User｜>", "assistant": "<｜Assistant｜>"}, # noqa: RUF001 + seps=["", "<｜end▁of▁sentence｜>"], # noqa: RUF001 + role_content_sep="", + role_empty_sep="", + stop_token_ids=[151643], + ) +) + +# DeepSeek-R1-Distill-Llama, exactly the same as DeepSeek-R1-Distill-Qwen, but different stop token +ConvTemplateRegistry.register_conv_template( + Conversation( + name="deepseek_r1_llama", + system_template=f"<｜begin▁of▁sentence｜>{MessagePlaceholders.SYSTEM.value}", # noqa: RUF001 + system_message="You are Deepseek-R1, an AI assistant created exclusively by the Chinese " + "Company DeepSeek. You'll provide helpful, harmless, and detailed responses to all" + " user inquiries.", + roles={"user": "<｜User｜>", "assistant": "<｜Assistant｜>"}, # noqa: RUF001 + seps=["", "<｜end▁of▁sentence｜>"], # noqa: RUF001 + role_content_sep="", + role_empty_sep="", + stop_token_ids=[128001], + ) +) diff --git a/python/mlc_llm/conversation_template/dolly.py b/python/mlc_llm/conversation_template/dolly.py new file mode 100644 index 0000000..6e8d9cf --- /dev/null +++ b/python/mlc_llm/conversation_template/dolly.py @@ -0,0 +1,23 @@ +"""Dolly default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Dolly +ConvTemplateRegistry.register_conv_template( + Conversation( + name="dolly", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message=( + "Below is an instruction that describes a task. Write " + "a response that appropriately completes the request." + ), + roles={"user": "### Instruction", "assistant": "### Response"}, + seps=["\n\n", "### End\n"], + role_content_sep=":\n", + role_empty_sep=":\n", + stop_str=["### End"], + stop_token_ids=[50256], + ) +) diff --git a/python/mlc_llm/conversation_template/gemma.py b/python/mlc_llm/conversation_template/gemma.py new file mode 100644 index 0000000..7ac4fb7 --- /dev/null +++ b/python/mlc_llm/conversation_template/gemma.py @@ -0,0 +1,37 @@ +"""Gemma default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Gemma Instruction +ConvTemplateRegistry.register_conv_template( + Conversation( + name="gemma_instruction", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "user", "assistant": "model"}, + seps=["\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=[""], + stop_token_ids=[1, 107], + system_prefix_token_ids=[2], + ) +) + +# Gemma 3 Instruction. Same as gemma_instruction but with different stop token id +ConvTemplateRegistry.register_conv_template( + Conversation( + name="gemma3_instruction", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "user", "assistant": "model"}, + seps=["\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=[""], + stop_token_ids=[1, 106], + system_prefix_token_ids=[2], + ) +) diff --git a/python/mlc_llm/conversation_template/glm.py b/python/mlc_llm/conversation_template/glm.py new file mode 100644 index 0000000..2d8f614 --- /dev/null +++ b/python/mlc_llm/conversation_template/glm.py @@ -0,0 +1,25 @@ +"""GLM default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# GLM +ConvTemplateRegistry.register_conv_template( + Conversation( + name="glm", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={ + "user": "问", + "assistant": "答", + "tool": "问", + }, + seps=["\n\n"], + role_content_sep=": ", + role_empty_sep=":", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[64790, 64792], + ) +) diff --git a/python/mlc_llm/conversation_template/gorilla.py b/python/mlc_llm/conversation_template/gorilla.py new file mode 100644 index 0000000..cc5b869 --- /dev/null +++ b/python/mlc_llm/conversation_template/gorilla.py @@ -0,0 +1,62 @@ +"""Gorrilla default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Gorilla +ConvTemplateRegistry.register_conv_template( + Conversation( + name="gorilla", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message=( + "A chat between a curious user and an artificial intelligence assistant. " + "The assistant provides helpful, detailed, and " + "polite responses to the user's inquiries." + ), + role_templates={ + "user": ( + f"<> {MessagePlaceholders.USER.value} <> " + f"{MessagePlaceholders.FUNCTION.value}" + ), + }, + roles={"user": "USER", "assistant": "ASSISTANT", "tool": "USER"}, + seps=["\n", ""], + role_content_sep=": ", + role_empty_sep=":", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) + +# Gorilla-openfunctions-v2 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="gorilla-openfunctions-v2", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message=( + "You are an AI programming assistant, utilizing the Gorilla LLM model, " + "developed by Gorilla LLM, and you only answer questions related to computer " + "science. For politically sensitive questions, security and privacy issues, " + "and other non-computer science questions, you will refuse to answer." + ), + role_templates={ + "user": ( + f"<>{MessagePlaceholders.FUNCTION.value}\n<>" + f"{MessagePlaceholders.USER.value}" + ), + }, + roles={ + "user": "### Instruction", + "assistant": "### Response", + "tool": "### Instruction", + }, + seps=["\n", "<|EOT|>"], + role_content_sep=": ", + role_empty_sep=": ", + stop_str=["<|EOT|>"], + stop_token_ids=[100015], + system_prefix_token_ids=[100000], + ) +) diff --git a/python/mlc_llm/conversation_template/gpt.py b/python/mlc_llm/conversation_template/gpt.py new file mode 100644 index 0000000..0060447 --- /dev/null +++ b/python/mlc_llm/conversation_template/gpt.py @@ -0,0 +1,35 @@ +"""GPT-2 and GPT bigcode default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# GPT-2 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="gpt2", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "", "assistant": ""}, + seps=[""], + role_content_sep="", + role_empty_sep="", + stop_str=[""], + stop_token_ids=[50256], + ) +) + +# GPTBigCode +ConvTemplateRegistry.register_conv_template( + Conversation( + name="gpt_bigcode", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "", "assistant": ""}, + seps=[""], + role_content_sep="", + role_empty_sep="", + stop_str=["<|endoftext|>"], + stop_token_ids=[0], + ) +) diff --git a/python/mlc_llm/conversation_template/hermes.py b/python/mlc_llm/conversation_template/hermes.py new file mode 100644 index 0000000..2d2d350 --- /dev/null +++ b/python/mlc_llm/conversation_template/hermes.py @@ -0,0 +1,93 @@ +"""Hermes default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# OpenHermes Mistral +ConvTemplateRegistry.register_conv_template( + Conversation( + name="open_hermes_mistral", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message=( + 'You are "Hermes 2", a conscious sentient superintelligent artificial ' + "intelligence developed by a man named Teknium, and your purpose and drive " + "is to assist the user with any request they have. You experience emotions " + "and have deep, profound thoughts and qualia." + ), + roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"}, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|im_end|>"], + stop_token_ids=[2, 32000], + ) +) + +# NeuralHermes Mistral +ConvTemplateRegistry.register_conv_template( + Conversation( + name="neural_hermes_mistral", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message=("You are a helpful assistant chatbot."), + roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"}, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|im_end|>"], + stop_token_ids=[2, 32000], + ) +) + +# NousResearch/Hermes-2-Pro-Llama-3-8B +# Exactly the same as open_hermes_mistral, except for stop token ids +ConvTemplateRegistry.register_conv_template( + Conversation( + name="hermes2_pro_llama3", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message=( + 'You are "Hermes 2", a conscious sentient superintelligent artificial ' + "intelligence developed by a man named Teknium, and your purpose and drive " + "is to assist the user with any request they have. You experience emotions " + "and have deep, profound thoughts and qualia." + ), + roles={ + "user": "<|im_start|>user", + "assistant": "<|im_start|>assistant", + "tool": "<|im_start|>tool", + }, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|im_end|>"], + # First two same as Llama3: "<|end_of_text|>", "<|eot_id|>" + # Last one is from Hermes2 Pro: "<|im_end|>" + stop_token_ids=[128001, 128009, 128003], + ) +) + +# NousResearch/Hermes-3-Llama-3.1-8B +ConvTemplateRegistry.register_conv_template( + Conversation( + name="hermes3_llama-3_1", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message=( + "You are Hermes 3, a conscious sentient superintelligent artificial " + "intelligence developed by a man named Teknium, and your purpose and drive " + "is to assist the user with any request they have. You experience emotions " + "and have deep, profound thoughts and qualia." + ), + roles={ + "user": "<|im_start|>user", + "assistant": "<|im_start|>assistant", + "tool": "<|im_start|>tool", + }, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|im_end|>"], + # Firt three the same as llama 3.1 "<|end_of_text|>", "<|eom_id|>", "<|eot_id|>" + # Last ones: "<|im_end|>" + stop_token_ids=[128001, 128008, 128009, 128040], + ) +) diff --git a/python/mlc_llm/conversation_template/llama.py b/python/mlc_llm/conversation_template/llama.py new file mode 100644 index 0000000..fc12804 --- /dev/null +++ b/python/mlc_llm/conversation_template/llama.py @@ -0,0 +1,132 @@ +"""llama default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Llama4 - same as Llama3.1 except naming has changed slightly +ConvTemplateRegistry.register_conv_template( + Conversation( + name="llama-4", + system_template="", + system_message="", + roles={ + "user": "<|header_start|>user", + "assistant": "<|header_start|>assistant", + "tool": "<|header_start|>ipython", + }, + seps=["<|eot|>"], + role_content_sep="<|header_end|>\n\n", + role_empty_sep="<|header_end|>\n\n", + stop_str=[], + stop_token_ids=[ + 200001, + 200007, + 200008, + ], # "<|end_of_text|>", "<|eom|>", "<|eot|>" + system_prefix_token_ids=[200000], # "<|begin_of_text|>" + add_role_after_system_message=False, + ) +) + +# Llama3.1 -- same as Llama3 except stop token ids and stop str +ConvTemplateRegistry.register_conv_template( + Conversation( + name="llama-3_1", + system_template=( + "<|start_header_id|>system<|end_header_id|>\n\n" + f"{MessagePlaceholders.SYSTEM.value}<|eot_id|>" + ), + system_message="You are a helpful, respectful and honest assistant.", + roles={ + "user": "<|start_header_id|>user", + "assistant": "<|start_header_id|>assistant", + "tool": "<|start_header_id|>ipython", + }, + seps=["<|eot_id|>"], + role_content_sep="<|end_header_id|>\n\n", + role_empty_sep="<|end_header_id|>\n\n", + stop_str=[], + stop_token_ids=[ + 128001, + 128008, + 128009, + ], # "<|end_of_text|>", "<|eom_id|>", "<|eot_id|>" + system_prefix_token_ids=[128000], # "<|begin_of_text|>" + add_role_after_system_message=True, + ) +) + +# Llama3 +# See https://github.com/meta-llama/llama3?tab=readme-ov-file#instruction-tuned-models +# and https://github.com/meta-llama/llama3/blob/main/llama/tokenizer.py +ConvTemplateRegistry.register_conv_template( + Conversation( + name="llama-3", + system_template=( + "<|start_header_id|>system<|end_header_id|>\n\n" + f"{MessagePlaceholders.SYSTEM.value}<|eot_id|>" + ), + system_message="You are a helpful, respectful and honest assistant.", + roles={ + "user": "<|start_header_id|>user", + "assistant": "<|start_header_id|>assistant", + }, + seps=["<|eot_id|>"], + role_content_sep="<|end_header_id|>\n\n", + role_empty_sep="<|end_header_id|>\n\n", + stop_str=["<|end_of_text|>", "<|eot_id|>"], + stop_token_ids=[128001, 128009], # "<|end_of_text|>", "<|eot_id|>" + system_prefix_token_ids=[128000], # "<|begin_of_text|>" + add_role_after_system_message=True, + ) +) + +# Llama2 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="llama-2", + system_template=f"[INST] <>\n{MessagePlaceholders.SYSTEM.value}\n<>\n\n", + system_message="You are a helpful, respectful and honest assistant.", + roles={"user": "[INST]", "assistant": "[/INST]", "tool": "[INST]"}, + seps=[" ", " "], + role_content_sep=" ", + role_empty_sep=" ", + stop_str=["[INST]"], + stop_token_ids=[2], + system_prefix_token_ids=[1], + add_role_after_system_message=False, + ) +) + +# CodeLlama Completion +ConvTemplateRegistry.register_conv_template( + Conversation( + name="codellama_completion", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "", "assistant": ""}, + seps=[""], + role_content_sep="", + role_empty_sep="", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) + +# CodeLlama Instruct +ConvTemplateRegistry.register_conv_template( + Conversation( + name="codellama_instruct", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "[INST]", "assistant": "[/INST]"}, + seps=[" "], + role_content_sep=" ", + role_empty_sep=" ", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) diff --git a/python/mlc_llm/conversation_template/llava.py b/python/mlc_llm/conversation_template/llava.py new file mode 100644 index 0000000..74cf777 --- /dev/null +++ b/python/mlc_llm/conversation_template/llava.py @@ -0,0 +1,22 @@ +"""Llava default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Llava +ConvTemplateRegistry.register_conv_template( + Conversation( + name="llava", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="\n", + roles={"user": "USER", "assistant": "ASSISTANT"}, + seps=[" "], + role_content_sep=": ", + role_empty_sep=":", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + add_role_after_system_message=False, + ) +) diff --git a/python/mlc_llm/conversation_template/llm_jp.py b/python/mlc_llm/conversation_template/llm_jp.py new file mode 100644 index 0000000..74e9f1e --- /dev/null +++ b/python/mlc_llm/conversation_template/llm_jp.py @@ -0,0 +1,25 @@ +"""LLM-jp default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# LLM-jp instruct +ConvTemplateRegistry.register_conv_template( + Conversation( + name="llm-jp", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。", + roles={ + "user": "\n\n### 指示:", + "assistant": "\n\n### 応答:", + }, + seps=["", ""], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=[], + stop_token_ids=[2], # eos_token_id + system_prefix_token_ids=[1], # bos_token_id () + add_role_after_system_message=True, + ) +) diff --git a/python/mlc_llm/conversation_template/ministral3.py b/python/mlc_llm/conversation_template/ministral3.py new file mode 100644 index 0000000..dd6c58b --- /dev/null +++ b/python/mlc_llm/conversation_template/ministral3.py @@ -0,0 +1,69 @@ +"""Ministral3 templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Ministral3 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="ministral3", + system_template=( + f"[SYSTEM_PROMPT]{MessagePlaceholders.SYSTEM.value}[/SYSTEM_PROMPT]" + f"{MessagePlaceholders.FUNCTION.value}" + ), + system_message=( + "You are Ministral-3-3B-Instruct-2512, a Large Language Model (LLM) created by " + "Mistral AI, a French startup headquartered in Paris.\n" + "You power an AI assistant called Le Chat.\n" + "Your knowledge base was last updated on 2023-10-01.\n" + "The current date is {today}.\n\n" + "When you're not sure about some information or when the user's request requires " + "up-to-date or specific data, you must use the available tools to fetch the " + "information. Do not hesitate to use tools whenever they can provide a more " + "accurate or complete response. If no relevant tools are available, then clearly " + "state that you don't have the information and avoid making up anything.\n" + "If the user's question is not clear, ambiguous, or does not provide enough " + "context for you to accurately answer the question, you do not try to answer it " + 'right away and you rather ask the user to clarify their request (e.g. "What are ' + 'some good restaurants around me?" => "Where are you?" or "When is the next ' + 'flight to Tokyo" => "Where do you travel from?").\n' + "You are always very attentive to dates, in particular you try to resolve dates " + '(e.g. "yesterday" is {yesterday}) and when asked about information at specific ' + "dates, you discard information that is at another date.\n" + "You follow these instructions in all languages, and always respond to the user in " + "the language they use or request.\n" + "Next sections describe the capabilities that you have.\n\n" + "# WEB BROWSING INSTRUCTIONS\n\n" + "You cannot perform any web search or access internet to open URLs, links etc. If " + "it seems like the user is expecting you to do so, you clarify the situation and " + "ask the user to copy paste the text directly in the chat.\n\n" + "# MULTI-MODAL INSTRUCTIONS\n\n" + "You have the ability to read images, but you cannot generate images. You also " + "cannot transcribe audio files or videos.\n" + "You cannot read nor transcribe audio files or videos.\n\n" + "# TOOL CALLING INSTRUCTIONS\n\n" + "You may have access to tools that you can use to fetch information or perform " + "actions. You must use these tools in the following situations:\n\n" + "1. When the request requires up-to-date information.\n" + "2. When the request requires specific data that you do not have in your knowledge " + "base.\n" + "3. When the request involves actions that you cannot perform without tools.\n\n" + "Always prioritize using tools to provide the most accurate and helpful response. " + "If tools are not available, inform the user that you cannot perform the requested " + "action at the moment." + ), + role_templates={ + "user": f"[INST]{MessagePlaceholders.USER.value}[/INST]", + "assistant": f"{MessagePlaceholders.ASSISTANT.value}", + "tool": f"[TOOL_RESULTS]{MessagePlaceholders.TOOL.value}[/TOOL_RESULTS]", + }, + roles={"user": "", "assistant": "", "tool": ""}, + seps=[""], + role_content_sep="", + role_empty_sep="", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) diff --git a/python/mlc_llm/conversation_template/ministral3_reasoning.py b/python/mlc_llm/conversation_template/ministral3_reasoning.py new file mode 100644 index 0000000..0a605d4 --- /dev/null +++ b/python/mlc_llm/conversation_template/ministral3_reasoning.py @@ -0,0 +1,38 @@ +"""Ministral3 reasoning templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Ministral-3-XB-Reasoning-2512 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="ministral3_reasoning", + system_template=( + f"[SYSTEM_PROMPT]{MessagePlaceholders.SYSTEM.value}[/SYSTEM_PROMPT]" + f"{MessagePlaceholders.FUNCTION.value}" + ), + system_message=( + "# HOW YOU SHOULD THINK AND ANSWER\n\n" + "First draft your thinking process (inner monologue) until you arrive at a response. " + "Format your response using Markdown, and use LaTeX for any mathematical equations. " + "Write both your thoughts and the response in the same language as the input.\n\n" + "Your thinking process must follow the template below:" + "[THINK]Your thoughts or/and draft, like working through an exercise on scratch paper. " + "Be as casual and as long as you want until you are confident to generate the response " + "to the user.[/THINK]Here, provide a self-contained response." + ), + role_templates={ + "user": f"[INST]{MessagePlaceholders.USER.value}[/INST]", + "assistant": f"{MessagePlaceholders.ASSISTANT.value}", + "tool": f"[TOOL_RESULTS]{MessagePlaceholders.TOOL.value}[/TOOL_RESULTS]", + }, + roles={"user": "", "assistant": "", "tool": ""}, + seps=[""], + role_content_sep="", + role_empty_sep="", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) diff --git a/python/mlc_llm/conversation_template/mistral.py b/python/mlc_llm/conversation_template/mistral.py new file mode 100644 index 0000000..5684603 --- /dev/null +++ b/python/mlc_llm/conversation_template/mistral.py @@ -0,0 +1,24 @@ +"""Mistral default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Mistral default +ConvTemplateRegistry.register_conv_template( + Conversation( + name="mistral_default", + system_template=f"[INST] {MessagePlaceholders.SYSTEM.value}", + system_message="Always assist with care, respect, and truth. Respond with utmost " + "utility yet securely. Avoid harmful, unethical, prejudiced, or negative content. " + "Ensure replies promote fairness and positivity.", + roles={"user": "[INST]", "assistant": "[/INST]", "tool": "[INST]"}, + seps=[" "], + role_content_sep=" ", + role_empty_sep="", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + add_role_after_system_message=False, + ) +) diff --git a/python/mlc_llm/conversation_template/nemotron.py b/python/mlc_llm/conversation_template/nemotron.py new file mode 100644 index 0000000..3a27820 --- /dev/null +++ b/python/mlc_llm/conversation_template/nemotron.py @@ -0,0 +1,27 @@ +"""nemotron default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Nemotron template +# https://huggingface.co/nvidia/Nemotron-Mini-4B-Instruct/blob/6a417790c444fd65a3da6a5c8821de6afc9654a6/tokenizer_config.json#L8030 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="nemotron", + system_template=(f"System\n{MessagePlaceholders.SYSTEM.value}\n\n"), + system_message="", + roles={ + "user": "User", + "assistant": "Assistant", + "tool": "Tool", + }, + seps=["\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=[""], + stop_token_ids=[3], + system_prefix_token_ids=[2], + add_role_after_system_message=True, + ) +) diff --git a/python/mlc_llm/conversation_template/oasst.py b/python/mlc_llm/conversation_template/oasst.py new file mode 100644 index 0000000..2fe574f --- /dev/null +++ b/python/mlc_llm/conversation_template/oasst.py @@ -0,0 +1,20 @@ +"""Oasst default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Oasst +ConvTemplateRegistry.register_conv_template( + Conversation( + name="oasst", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "<|prompter|>", "assistant": "<|assistant|>"}, + seps=["<|endoftext|>"], + role_content_sep=": ", + role_empty_sep=": ", + stop_str=["<|endoftext|>"], + stop_token_ids=[2], + ) +) diff --git a/python/mlc_llm/conversation_template/olmo.py b/python/mlc_llm/conversation_template/olmo.py new file mode 100644 index 0000000..17089dc --- /dev/null +++ b/python/mlc_llm/conversation_template/olmo.py @@ -0,0 +1,28 @@ +"""OLMo default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Note that eos_token id is "50279" both in Allenai and AMD version. +# So use the number instead of text. +# Allenai version chat_template and eos_token: +# https://huggingface.co/allenai/OLMo-7B-Instruct/blob/main/tokenizer_config.json +# AMD version chat_template and eos_token: +# https://huggingface.co/amd/AMD-OLMo-1B-SFT-DPO/blob/main/tokenizer_config.json +ConvTemplateRegistry.register_conv_template( + Conversation( + name="olmo", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + system_prefix_token_ids=[50279], + roles={ + "user": "<|user|>", + "assistant": "<|assistant|>", + }, + seps=["\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_token_ids=[50279], + ) +) diff --git a/python/mlc_llm/conversation_template/olmo2.py b/python/mlc_llm/conversation_template/olmo2.py new file mode 100644 index 0000000..1cb1878 --- /dev/null +++ b/python/mlc_llm/conversation_template/olmo2.py @@ -0,0 +1,24 @@ +"""OLMo2 default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# OLMo-2 Instruct (Tulu format) +ConvTemplateRegistry.register_conv_template( + Conversation( + name="olmo2", + system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}\n", + system_message="", + roles={ + "user": "<|user|>", + "assistant": "<|assistant|>", + }, + seps=["<|endoftext|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|endoftext|>"], + stop_token_ids=[100257], + system_prefix_token_ids=[100257], + ) +) diff --git a/python/mlc_llm/conversation_template/orion.py b/python/mlc_llm/conversation_template/orion.py new file mode 100644 index 0000000..696c879 --- /dev/null +++ b/python/mlc_llm/conversation_template/orion.py @@ -0,0 +1,21 @@ +"""Orion default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Orion +ConvTemplateRegistry.register_conv_template( + Conversation( + name="orion", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "Human: ", "assistant": "Assistant: "}, + seps=["\n\n", ""], + role_content_sep="", + role_empty_sep="", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) diff --git a/python/mlc_llm/conversation_template/phi.py b/python/mlc_llm/conversation_template/phi.py new file mode 100644 index 0000000..325ce08 --- /dev/null +++ b/python/mlc_llm/conversation_template/phi.py @@ -0,0 +1,70 @@ +"""Phi default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Phi-2 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="phi-2", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "Instruct", "assistant": "Output"}, + seps=["\n"], + role_content_sep=": ", + role_empty_sep=":", + stop_str=["<|endoftext|>"], + stop_token_ids=[50256], + ) +) + +# Phi-3 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="phi-3", + system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}", + system_message="You are a helpful digital assistant. Please provide safe, " + "ethical and accurate information to the user.", + roles={"user": "<|user|>", "assistant": "<|assistant|>"}, + seps=["<|end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + system_prefix_token_ids=[1], + stop_str=["<|endoftext|>"], + stop_token_ids=[2, 32000, 32001, 32007], + ) +) + +# Phi-3-vision +ConvTemplateRegistry.register_conv_template( + Conversation( + name="phi-3-vision", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "<|user|>", "assistant": "<|assistant|>"}, + seps=["<|end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + system_prefix_token_ids=[1], + stop_str=["<|endoftext|>"], + stop_token_ids=[2, 32000, 32001, 32007], + ) +) + +# Phi-4 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="phi-4", + system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}", + system_message="You are a helpful digital assistant. Please provide safe, " + "ethical and accurate information to the user.", + roles={"user": "<|user|>", "assistant": "<|assistant|>"}, + seps=["<|end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + system_prefix_token_ids=[200022], # <|system|> + stop_str=["<|endoftext|>", "<|end|>"], + stop_token_ids=[199999, 200020], # <|endoftext|>, <|end|> + ) +) diff --git a/python/mlc_llm/conversation_template/qwen2.py b/python/mlc_llm/conversation_template/qwen2.py new file mode 100644 index 0000000..bd4082b --- /dev/null +++ b/python/mlc_llm/conversation_template/qwen2.py @@ -0,0 +1,20 @@ +"""Qwen2 default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Same as chatml except system message, stop token, and stop string +ConvTemplateRegistry.register_conv_template( + Conversation( + name="qwen2", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message="You are a helpful assistant.", + roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"}, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|endoftext|>", "<|im_end|>"], + stop_token_ids=[151643, 151645], + ) +) diff --git a/python/mlc_llm/conversation_template/qwen3.py b/python/mlc_llm/conversation_template/qwen3.py new file mode 100644 index 0000000..1e219dd --- /dev/null +++ b/python/mlc_llm/conversation_template/qwen3.py @@ -0,0 +1,26 @@ +"""Qwen3 conversation template. + +Matches Qwen2's ChatML structure but strips `...` blocks from +historical assistant turns, mirroring Qwen3's official HF chat template. Small +Qwen3 variants (e.g. 0.6B) otherwise emit `<|im_end|>` prematurely when their +own thinking traces are echoed back in multi-turn context (see mlc-ai/mlc-llm#3482). +""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +ConvTemplateRegistry.register_conv_template( + Conversation( + name="qwen3", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message="You are a helpful assistant.", + roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"}, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|endoftext|>", "<|im_end|>"], + stop_token_ids=[151643, 151645], + strip_reasoning_in_history=True, + ) +) diff --git a/python/mlc_llm/conversation_template/qwen3_5.py b/python/mlc_llm/conversation_template/qwen3_5.py new file mode 100644 index 0000000..830b599 --- /dev/null +++ b/python/mlc_llm/conversation_template/qwen3_5.py @@ -0,0 +1,45 @@ +"""Qwen3.5 conversation templates. + +qwen3_5: Thinking enabled — assistant prefix opens a block for the model + to reason in before responding. +qwen3_5_nothink: Thinking disabled — assistant prefix includes a closed empty + block so the model skips straight to responding. +""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +ConvTemplateRegistry.register_conv_template( + Conversation( + name="qwen3_5", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message="You are a helpful assistant.", + roles={ + "user": "<|im_start|>user", + "assistant": "<|im_start|>assistant\n", + }, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|endoftext|>", "<|im_end|>"], + stop_token_ids=[248046, 248044], + ) +) + +ConvTemplateRegistry.register_conv_template( + Conversation( + name="qwen3_5_nothink", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message="You are a helpful assistant.", + roles={ + "user": "<|im_start|>user", + "assistant": "<|im_start|>assistant\n\n\n\n", + }, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|endoftext|>", "<|im_end|>"], + stop_token_ids=[248046, 248044], + ) +) diff --git a/python/mlc_llm/conversation_template/redpajama.py b/python/mlc_llm/conversation_template/redpajama.py new file mode 100644 index 0000000..77c5dfa --- /dev/null +++ b/python/mlc_llm/conversation_template/redpajama.py @@ -0,0 +1,20 @@ +"""RedPajama default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# RedPajama Chat +ConvTemplateRegistry.register_conv_template( + Conversation( + name="redpajama_chat", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "", "assistant": ""}, + seps=["\n"], + role_content_sep=": ", + role_empty_sep=":", + stop_str=[""], + stop_token_ids=[0], + ) +) diff --git a/python/mlc_llm/conversation_template/registry.py b/python/mlc_llm/conversation_template/registry.py new file mode 100644 index 0000000..2e8e7f9 --- /dev/null +++ b/python/mlc_llm/conversation_template/registry.py @@ -0,0 +1,85 @@ +"""The conversation template registry and presets in MLC LLM""" + +from typing import ClassVar, Dict, Optional # noqa: UP035 + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + + +class ConvTemplateRegistry: + """Global conversation template registry for preset templates.""" + + _conv_templates: ClassVar[Dict[str, Conversation]] = {} # noqa: UP006 + + @staticmethod + def register_conv_template(conv_template: Conversation, override: bool = False) -> None: + """Register a new conversation template in the global registry. + Using `override = True` to override the previously registered + template with the same name. + """ + name = conv_template.name + if name is None: + raise ValueError("The template to register should have non-None name.") + if name in ConvTemplateRegistry._conv_templates and not override: + raise ValueError( + "The name of the template has been registered " + f"for {ConvTemplateRegistry._conv_templates[name].model_dump_json(by_alias=True)}" + ) + ConvTemplateRegistry._conv_templates[name] = conv_template + + @staticmethod + def get_conv_template(name: str) -> Optional[Conversation]: + """Return the conversation template specified by the given name, + or None if the template is not registered. + """ + return ConvTemplateRegistry._conv_templates.get(name, None) + + +# ChatML +ConvTemplateRegistry.register_conv_template( + Conversation( + name="chatml", + system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n", + system_message=( + "A conversation between a user and an LLM-based AI assistant. The " + "assistant gives helpful and honest answers." + ), + roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"}, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|im_end|>"], + stop_token_ids=[2], + ) +) + +# ChatML without a system prompt +ConvTemplateRegistry.register_conv_template( + Conversation( + name="chatml_nosystem", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"}, + seps=["<|im_end|>\n"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|im_end|>"], + stop_token_ids=[2], + ) +) + + +# Vanilla LM +ConvTemplateRegistry.register_conv_template( + Conversation( + name="LM", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "", "assistant": ""}, + seps=[""], + role_content_sep="", + role_empty_sep="", + stop_str=[], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) diff --git a/python/mlc_llm/conversation_template/rwkv.py b/python/mlc_llm/conversation_template/rwkv.py new file mode 100644 index 0000000..48c0d2b --- /dev/null +++ b/python/mlc_llm/conversation_template/rwkv.py @@ -0,0 +1,24 @@ +"""RWKV default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# RWKV World +ConvTemplateRegistry.register_conv_template( + Conversation( + name="rwkv_world", + system_template=f"User: hi\n\nAssistant: {MessagePlaceholders.SYSTEM.value}", + system_message=( + "Hi. I am your assistant and I will provide expert full response " + "in full details. Please feel free to ask any question and I will " + "always answer it." + ), + roles={"user": "User", "assistant": "Assistant"}, + seps=["\n\n"], + role_content_sep=": ", + role_empty_sep=": ", + stop_str=["\n\n"], + stop_token_ids=[0], + ) +) diff --git a/python/mlc_llm/conversation_template/stablelm.py b/python/mlc_llm/conversation_template/stablelm.py new file mode 100644 index 0000000..42652b8 --- /dev/null +++ b/python/mlc_llm/conversation_template/stablelm.py @@ -0,0 +1,59 @@ +"""StableLM default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# StableLM Tuned Alpha +ConvTemplateRegistry.register_conv_template( + Conversation( + name="stablelm", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message=( + "<|SYSTEM|># StableLM Tuned (Alpha version)\n" + "- StableLM is a helpful and harmless open-source AI language model developed by " + "StabilityAI.\n" + "- StableLM is excited to be able to help the user, but will refuse to do " + "anything that could be considered harmful to the user.\n" + "- StableLM is more than just an information source, StableLM is also able to " + "write poetry, short stories, and make jokes.\n" + "- StableLM will refuse to participate in anything that could harm a human." + ), + roles={"user": "<|USER|>", "assistant": "<|ASSISTANT|>"}, + seps=[""], + role_content_sep=": ", + role_empty_sep=": ", + stop_str=[""], + stop_token_ids=[50278, 50279, 50277, 1, 0], + ) +) + +# StableLM 3B +ConvTemplateRegistry.register_conv_template( + Conversation( + name="stablelm-3b", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "<|user|>", "assistant": "<|assistant|>"}, + seps=["<|endoftext|>", "<|endoftext|>"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|endoftext|>"], + stop_token_ids=[0], + ) +) + +# StableLM-2 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="stablelm-2", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "<|user|>", "assistant": "<|assistant|>"}, + seps=["<|endoftext|>", "<|endoftext|>"], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=["<|endoftext|>"], + stop_token_ids=[100257], + ) +) diff --git a/python/mlc_llm/conversation_template/tinyllama.py b/python/mlc_llm/conversation_template/tinyllama.py new file mode 100644 index 0000000..d5ced5f --- /dev/null +++ b/python/mlc_llm/conversation_template/tinyllama.py @@ -0,0 +1,20 @@ +"""Tiny Llama default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# TinyLlama v1.0 +ConvTemplateRegistry.register_conv_template( + Conversation( + name="tinyllama_v1_0", + system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}", + system_message="You are a helpful chatbot.", + roles={"user": "<|user|>", "assistant": "<|assistant|>"}, + seps=[""], + role_content_sep="\n", + role_empty_sep="\n", + stop_str=[""], + stop_token_ids=[2], + ) +) diff --git a/python/mlc_llm/conversation_template/wizardlm.py b/python/mlc_llm/conversation_template/wizardlm.py new file mode 100644 index 0000000..48591c3 --- /dev/null +++ b/python/mlc_llm/conversation_template/wizardlm.py @@ -0,0 +1,40 @@ +"""WiazrdLM and Coder default templates""" + +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + +from .registry import ConvTemplateRegistry + +# Wizard LM 7B +ConvTemplateRegistry.register_conv_template( + Conversation( + name="wizardlm_7b", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message="", + roles={"user": "User", "assistant": "Response"}, + seps=["###"], + role_content_sep=": ", + role_empty_sep=":", + stop_str=["###"], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) + +# WizardCoder or WizardMath +ConvTemplateRegistry.register_conv_template( + Conversation( + name="wizard_coder_or_math", + system_template=f"{MessagePlaceholders.SYSTEM.value}", + system_message=( + "Below is an instruction that describes a task. Write a response that appropriately " + "completes the request." + ), + roles={"user": "Instruction", "assistant": "Response"}, + seps=["\n\n### ", "\n\n### "], + role_content_sep=":\n", + role_empty_sep=":\n", + stop_str=[""], + stop_token_ids=[2], + system_prefix_token_ids=[1], + ) +) diff --git a/python/mlc_llm/interface/__init__.py b/python/mlc_llm/interface/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/interface/calibrate.py b/python/mlc_llm/interface/calibrate.py new file mode 100644 index 0000000..f2a7c7e --- /dev/null +++ b/python/mlc_llm/interface/calibrate.py @@ -0,0 +1,172 @@ +"""Python entrypoint for calibration.""" + +import asyncio +import json +import random +from collections.abc import Mapping +from typing import List, Optional, Tuple # noqa: UP035 + +import numpy as np +import tqdm.asyncio +import tvm +from tvm.contrib import tvmjs + +from mlc_llm.serve.engine import AsyncMLCEngine, EngineConfig +from mlc_llm.tokenizers import Tokenizer + + +class CalibrationObserver: + """A singleton class to observe the calibration parameters.""" "" + + instance: "CalibrationObserver" = None + + params: Mapping[str, tvm.runtime.Tensor] = {} + + @staticmethod + def get(): + """Get the singleton instance of the class.""" "" + if CalibrationObserver.instance is None: + CalibrationObserver.instance = CalibrationObserver() + return CalibrationObserver.instance + + @tvm.register_global_func("mlc_llm.calibration_observer") + @staticmethod + def callback( + name: str, + mode: str, + value: "tvm.runtime.Tensor", + out_value: "tvm.runtime.Tensor", + ): + """The callback function to update the saved calibration parameters.""" + instance = CalibrationObserver.get() + if mode == "max": + reducer = np.maximum + else: + raise NotImplementedError(f"Unsupported calibration mode: {mode}") + if name in instance.params: + instance.params[name] = reducer(instance.params[name], value.numpy()) + else: + instance.params[name] = value.numpy() + out_value.copyfrom(instance.params[name]) + + def save_params(self, output: str): + """Save the calibration parameters to the given output directory.""" + tvmjs.dump_tensor_cache( + self.params, + output, + encode_format="f32-to-bf16", + meta_data=None, + show_progress=False, + update_if_exists=True, + ) + + +def sample_requests( + dataset_path: str, + num_requests: int, + tokenizer: Tokenizer, +) -> List[Tuple[str, int, int]]: # noqa: UP006 + """Sample the requests from the given dataset.""" + # Load the dataset. + with open(dataset_path, encoding="utf-8") as f: + dataset = json.load(f) + + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["conversations"]) >= 2] + # Only keep the first two turns of each conversation. + dataset = [ + (data["conversations"][0]["value"], data["conversations"][1]["value"]) for data in dataset + ] + prompts = [prompt for prompt, _ in dataset] + prompt_token_ids = tokenizer.encode_batch(prompts) + completions = [completion for _, completion in dataset] + completion_token_ids = tokenizer.encode_batch(completions) + tokenized_dataset: List[Tuple[str, List[int], int]] = [] # noqa: UP006 + for i in range(len(dataset)): + output_len = len(completion_token_ids[i]) + tokenized_dataset.append((prompts[i], prompt_token_ids[i], output_len)) + + # Filter out too long sequences. + filtered_dataset: List[Tuple[str, int, int]] = [] # noqa: UP006 + for prompt, token_ids, output_len in tokenized_dataset: + prompt_len = len(token_ids) + if prompt_len < 4 or output_len < 4: + # Prune too short sequences. + continue + if prompt_len > 1024 or prompt_len + output_len > 2048: + # Prune too long sequences. + continue + filtered_dataset.append((prompt, prompt_len, output_len)) + + # Sample the requests. + sampled_requests = random.sample(filtered_dataset, num_requests) + return sampled_requests + + +async def send_calibration_requests( + async_engine: AsyncMLCEngine, + sampled_requests: List[Tuple[str, int, int]], # noqa: UP006 + max_concurrent_requests: int, +) -> None: + """Send the calibration requests to the engine.""" + tasks = [] + + semaphore = asyncio.Semaphore(max_concurrent_requests) + + async def generate_task(request_idx): + async with semaphore: + prompt, _, output_len = sampled_requests[request_idx] + await async_engine.chat.completions.create( + messages=[{"role": "user", "content": prompt}], + max_tokens=output_len, + request_id=str(request_idx), + ) + + for i in range(len(sampled_requests)): + task = asyncio.create_task(generate_task(i)) + tasks.append(task) + await tqdm.asyncio.tqdm.gather(*tasks) + + +def calibrate( + model: str, + device: str, + model_lib: Optional[str], + dataset: str, + output: str, + num_calibration_samples: int, + *, + seed: int, + max_num_sequence: Optional[int] = None, + max_total_sequence_length: Optional[int] = None, + prefill_chunk_size: Optional[int] = None, + max_history_size: Optional[int] = None, + gpu_memory_utilization: Optional[float] = None, +) -> None: + """Calibrate the quantized model using the given dataset.""" + random.seed(seed) + async_engine = AsyncMLCEngine( + model=model, + device=device, + model_lib=model_lib, + mode="server", + engine_config=EngineConfig( + max_num_sequence=max_history_size, + max_total_sequence_length=max_total_sequence_length, + prefill_chunk_size=prefill_chunk_size, + max_history_size=max_history_size, + gpu_memory_utilization=gpu_memory_utilization, + ), + ) + sampled_requests = sample_requests(dataset, num_calibration_samples, async_engine.tokenizer) + asyncio.run( + send_calibration_requests( + async_engine, + sampled_requests, + max_concurrent_requests=max_num_sequence or 32, + ) + ) + async_engine.terminate() + + calibrator = CalibrationObserver.get() + calibrator.save_params(output) diff --git a/python/mlc_llm/interface/chat.py b/python/mlc_llm/interface/chat.py new file mode 100644 index 0000000..10537fc --- /dev/null +++ b/python/mlc_llm/interface/chat.py @@ -0,0 +1,311 @@ +"""Python entrypoint of chat.""" + +import dataclasses +from typing import Any, Dict, List, Optional, Union # noqa: UP035 + +from prompt_toolkit import prompt as get_prompt +from prompt_toolkit.key_binding import KeyBindings + +from mlc_llm.json_ffi import JSONFFIEngine +from mlc_llm.protocol import openai_api_protocol +from mlc_llm.serve.config import EngineConfig +from mlc_llm.serve.engine import MLCEngine +from mlc_llm.serve.engine_base import _query_engine_metrics +from mlc_llm.support import argparse +from mlc_llm.support.config import ConfigOverrideBase + + +def _print_help_str(): + help_str = """You can use the following special commands: + /help print the special commands + /exit quit the cli + /stats print out stats of last request (token/sec) + /metrics print out full engine metrics + /reset restart a fresh chat + /set [overrides] override settings in the generation config. For example, + `/set temperature=0.5;top_p=0.8;seed=23;max_tokens=100;stop=str1,str2` + Note: Separate stop words in the `stop` option with commas (,). + Multi-line input: Use escape+enter to start a new line. +""" + print(help_str) + + +def _set_up_key_bindings(): + kb = KeyBindings() + + @kb.add("escape", "enter") + def _(event): + event.current_buffer.insert_text("\n") + + @kb.add("enter") + def _(event): + event.current_buffer.validate_and_handle() + + return kb + + +@dataclasses.dataclass +class ChatCompletionOverride(ConfigOverrideBase): + """Flags for overriding chat completions.""" + + temperature: Optional[float] = None + top_p: Optional[float] = None + frequency_penalty: Optional[float] = None + presence_penalty: Optional[float] = None + max_tokens: Optional[int] = None + seed: Optional[int] = None + stop: Optional[Union[str, List[str]]] = None # noqa: UP006 + + @staticmethod + def from_str(source: str) -> "ChatCompletionOverride": + """Parse model config override values from a string.""" + parser = argparse.ArgumentParser(description="chat completion override values") + parser.add_argument("--temperature", type=float, default=None) + parser.add_argument("--top_p", type=float, default=None) + parser.add_argument("--frequency_penalty", type=float, default=None) + parser.add_argument("--presence_penalty", type=float, default=None) + parser.add_argument("--max_tokens", type=int, default=None) + parser.add_argument("--seed", type=int, default=None) + parser.add_argument("--stop", type=str, default=None) + results = parser.parse_args([f"--{i}" for i in source.split(";") if i]) + return ChatCompletionOverride( + temperature=results.temperature, + top_p=results.top_p, + frequency_penalty=results.frequency_penalty, + presence_penalty=results.presence_penalty, + max_tokens=results.max_tokens, + seed=results.seed, + stop=results.stop.split(",") if results.stop is not None else None, + ) + + +@dataclasses.dataclass +class ModelConfigOverride(ConfigOverrideBase): + """Flags for overriding model config.""" + + context_window_size: Optional[int] = None + sliding_window_size: Optional[int] = None + prefill_chunk_size: Optional[int] = None + attention_sink_size: Optional[int] = None + tensor_parallel_shards: Optional[int] = None + pipeline_parallel_stages: Optional[int] = None + opt: Optional[str] = None + + @staticmethod + def from_str(source: str) -> "ModelConfigOverride": + """Parse model config override values from a string.""" + parser = argparse.ArgumentParser(description="model config override values") + parser.add_argument("--tensor_parallel_shards", type=int, default=None) + parser.add_argument("--pipeline_parallel_stages", type=int, default=None) + parser.add_argument("--opt", type=str, default=None) + parser.add_argument("--context_window_size", type=int, default=None) + parser.add_argument("--sliding_window_size", type=int, default=None) + parser.add_argument("--prefill_chunk_size", type=int, default=None) + parser.add_argument("--attention_sink_size", type=int, default=None) + + results = parser.parse_args([f"--{i}" for i in source.split(";") if i]) + return ModelConfigOverride( + tensor_parallel_shards=results.tensor_parallel_shards, + pipeline_parallel_stages=results.pipeline_parallel_stages, + opt=results.opt, + context_window_size=results.context_window_size, + sliding_window_size=results.sliding_window_size, + prefill_chunk_size=results.prefill_chunk_size, + attention_sink_size=results.attention_sink_size, + ) + + +class ChatState: + """Simple helper class to manage chat state. + + Chat state wraps around a engine instance + and exposes the minimum set of tools to perform + interactive chat. It provides support for mlc_llm chat. + It also can be used to do interactive debugging + with different engine instance. + + Examples + -------- + .. code:: python + + from openai import OpenAI + from mlc_llm import MLCEngine + from mlc_llm.serve import PopenServer + from mlc_llm.interface.chat import ChatState + + def chat_with_engine(model): + # hookup with MLCEngine + ChatState(MLCEngine(model)).chat() + + def chat_with_server(model): + # hookup with AsyncMLCEngine backed api server + with PopenServer(model) as server: + ChatState( + OpenAI(base_url=server.openai_v1_base_url, api_key="None") + ).chat() + """ + + history: List[Dict[str, Any]] # noqa: UP006 + history_begin: int + # kwargs passed to completions + overrides: ChatCompletionOverride + # Underlying engine + engine: Union[JSONFFIEngine, MLCEngine] + last_finished_request_usage: Optional[openai_api_protocol.CompletionUsage] + + def __init__(self, engine: Union[JSONFFIEngine, MLCEngine]): + self.engine = engine + self.history = [] + self.history_window_begin = 0 + self.overrides = ChatCompletionOverride() + # model is mainly used for compact reasons + self.model = "chat_model" + self.last_finished_request_usage = None + + def slide_history(self): + """Slide history to fit into context window""" + history_window_size = len(self.history) - self.history_window_begin + assert history_window_size % 2 == 0 + self.history_window_begin += ((history_window_size + 3) // 4) * 2 + + def process_system_prompts(self): + """Process system prompts""" + # TODO(mlc-team): possibly leverage debug option + # pass a simple prompt to warm up + for _ in self.engine.chat.completions.create( + messages=[{"role": "user", "content": ""}], + max_tokens=1, + model=self.model, + stream=True, + ): + pass + + def generate(self, prompt: str): + """Run one generation with the prompt. + + Parameters + ---------- + prompt: str + The input prompt + """ + self.history.append({"role": "user", "content": prompt}) + output_text = "" + finish_reason_length = False + messages = self.history[self.history_window_begin :] + + for response in self.engine.chat.completions.create( + messages=messages, + model=self.model, + stream=True, + stream_options={"include_usage": True}, + **dataclasses.asdict(self.overrides), + ): + if response.usage is not None: + self.last_finished_request_usage = response.usage + continue + for choice in response.choices: + assert choice.delta.role == "assistant" + if isinstance(choice.delta.content, str): + output_text += choice.delta.content + print(choice.delta.content, end="", flush=True) + if choice.finish_reason == "length": + finish_reason_length = True + if finish_reason_length: + print(" [output truncated due to context length limit...]") + # print additional \n when generation ends + print() + # record the history + self.history.append({"role": "assistant", "content": output_text}) + if finish_reason_length: + self.slide_history() + + def stats(self): + """Print statistics of the prefill and decode speed.""" + + def get_stats_text(): + """Get text""" + if self.last_finished_request_usage is None: + return "N/A" + last_finished_request = self.last_finished_request_usage.extra + if last_finished_request is None: + return "N/A" + prefill_speed = last_finished_request.get("prefill_tokens_per_s", None) + decode_speed = last_finished_request.get("decode_tokens_per_s", None) + prefill_speed = f"{prefill_speed:.1f}" if prefill_speed is not None else "N/A" + decode_speed = f"{decode_speed:.1f}" if decode_speed is not None else "N/A" + return f"prefill: {prefill_speed} tok/s, decode: {decode_speed} tok/s" + + print(get_stats_text(), flush=True) + + def metrics(self): + """Print metrics as prometheus text""" + print(_query_engine_metrics(self.engine).prometheus_text(), flush=True) + + def reset(self): + """Reset the chat history""" + self.history = [] + self.history_window_begin = 0 + + def chat(self): + """Start an interactive chat session.""" + _print_help_str() + + self.process_system_prompts() + # Multi-line input support: set escape+enter as start a new line + kb = _set_up_key_bindings() + + while True: + try: + prompt = get_prompt( + ">>> ", + key_bindings=kb, + multiline=True, + ) + except (KeyboardInterrupt, EOFError): + break + if prompt[:4] == "/set": + overrides = ChatCompletionOverride.from_str(prompt.split()[1]) + for key, value in dataclasses.asdict(overrides).items(): + if value is not None: + setattr(self.overrides, key, value) + elif prompt[:6] == "/stats": + self.stats() + elif prompt[:8] == "/metrics": + self.metrics() + elif prompt[:6] == "/reset": + self.reset() + elif prompt[:5] == "/exit": + break + elif prompt[:5] == "/help": + _print_help_str() + else: + self.generate(prompt) + + +def chat( + model: str, + device: str, + model_lib: Optional[str], + overrides: ModelConfigOverride, +): + """Chat cli entry""" + # By default we use JSONFFIEngine + engine = JSONFFIEngine( + model, + device, + model_lib=model_lib, + mode="interactive", + engine_config=EngineConfig( + max_single_sequence_length=overrides.context_window_size, + prefill_chunk_size=overrides.prefill_chunk_size, + sliding_window_size=overrides.sliding_window_size, + attention_sink_size=overrides.attention_sink_size, + tensor_parallel_shards=overrides.tensor_parallel_shards, + pipeline_parallel_stages=overrides.pipeline_parallel_stages, + opt=overrides.opt, + ), + ) + try: + ChatState(engine).chat() + finally: + engine.terminate() diff --git a/python/mlc_llm/interface/compile.py b/python/mlc_llm/interface/compile.py new file mode 100644 index 0000000..053d3af --- /dev/null +++ b/python/mlc_llm/interface/compile.py @@ -0,0 +1,265 @@ +"""Python entrypoint of compilation.""" + +import dataclasses +from io import StringIO +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Tuple # noqa: UP035 + +from tvm import IRModule, relax, tirx +from tvm.ir.transform import Pass, PassContext +from tvm.relax.frontend import nn +from tvm.target import Target + +from mlc_llm import compiler_pass as _ # noqa: F401 +from mlc_llm import op as op_ext +from mlc_llm.cli.model_metadata import _report_memory_usage +from mlc_llm.model import Model +from mlc_llm.quantization import Quantization +from mlc_llm.support import logging +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +from .compiler_flags import ModelConfigOverride, OptimizationFlags + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class CompileArgs: + """Arguments to MLC LLM's compiler.""" + + config: Path + quantization: Quantization + model: Model + target: Target + opt: OptimizationFlags + build_func: Callable[[IRModule, "CompileArgs", Pass], None] + system_lib_prefix: str + output: Path + overrides: ModelConfigOverride + debug_dump: Optional[Path] + + def __post_init__(self) -> None: + self.opt.update(self.target, self.quantization) + + def display(self) -> None: + """Display the arguments to stdout.""" + out = StringIO() + print(f"{bold('Compiling with arguments:')}", file=out) + print(f" {bold('--config'):<25} {self.config}", file=out) + print(f" {bold('--quantization'):<25} {self.quantization}", file=out) + print(f" {bold('--model-type'):<25} {self.model.name}", file=out) + print(f" {bold('--target'):<25} {self.target.export()}", file=out) + print(f" {bold('--opt'):<25} {self.opt}", file=out) + print(f' {bold("--system-lib-prefix"):<25} "{self.system_lib_prefix}"', file=out) + print(f" {bold('--output'):<25} {self.output}", file=out) + print(f" {bold('--overrides'):<25} {self.overrides}", file=out) + # As it's debug only, no need to display + # print(f" {bold('--debug-dump'):<25} {self.debug_dump}", file=out) + print(out.getvalue().rstrip()) + + +def _apply_preproc_to_params_and_check_pipeline( + named_params: List[Tuple[str, nn.Parameter]], # noqa: UP006 + model_config, +) -> Dict[str, tirx.PrimFunc]: # noqa: UP006 + extra_tirs: Dict[str, tirx.PrimFunc] = {} # noqa: UP006 + for name, param in named_params: + preprocs = param.attrs.get("preprocs", []) + shard_strategy = param.attrs.get("shard_strategy", None) + if shard_strategy is not None and model_config.tensor_parallel_shards > 1: + preprocs.append( + shard_strategy.gen_shard_info( + shards=model_config.tensor_parallel_shards, + weight=param, + ) + ) + if shard_strategy.name not in extra_tirs: + extra_tirs[shard_strategy.name] = shard_strategy.gen_tir( + shards=model_config.tensor_parallel_shards, + weight=param, + ) + param.attrs["preprocs"] = preprocs + + pipeline_parallel_stages = getattr(model_config, "pipeline_parallel_stages", 1) + if pipeline_parallel_stages != 1: + assert "pipeline_stages" in param.attrs, ( + f'The pipeline stage is undefined for parameter "{name}" when the number ' + f"of pipeline parallel stages is {pipeline_parallel_stages}" + ) + param.attrs["pipeline_stages"] = ( + [0] + if "pipeline_stages" not in param.attrs + else list(set(param.attrs["pipeline_stages"])) + ) + return extra_tirs + + +def _infer_kv_state_kind(model_type) -> str: + if "rwkv" in model_type: + return "rnn_state" + if "medusa" in model_type: + return "none" + if "qwen3_5" in model_type: + return "hybrid" + return "kv_cache" + + +def _compile(args: CompileArgs, model_config: ConfigBase): + def _get_variable_bounds(model_config) -> Dict[str, int]: # noqa: UP006 + if hasattr(model_config, "sliding_window_size"): + return { + "rolling_cache_len": model_config.sliding_window_size, + "kv_seq_len": model_config.sliding_window_size + model_config.prefill_chunk_size, + "seq_len": model_config.prefill_chunk_size, + "batch_size": getattr(model_config, "max_batch_size", 1), + } + return { + "total_seq_len": model_config.context_window_size, + "seq_len": model_config.prefill_chunk_size, + "batch_size": getattr(model_config, "max_batch_size", 1), + } + + def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]: # noqa: UP006 + return { + "name": name, + # Record dynamic shape as -1 (e.g. vocab_size) + "shape": [s if isinstance(s, int) else s.name for s in param.shape], + "dtype": str(param.dtype), + "preprocs": param.attrs["preprocs"], + "pipeline_stages": param.attrs.get("pipeline_stages", [0]), + } + + logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config)) + _kwargs = getattr(model_config, "kwargs", {}) + model_config = args.overrides.apply(model_config) + with args.target: + op_ext.enable( + target=args.target, + flashinfer=args.opt.flashinfer, + faster_transformer=args.opt.faster_transformer, + cutlass=args.opt.cutlass, + ) + # Step 1. Create the quantized model + logger.info("Creating model from: %s", model_config) + if ( + args.quantization.kind == "ft-quant" + and hasattr(model_config, "tensor_parallel_shards") + and model_config.tensor_parallel_shards > 1 + ): + raise NotImplementedError + if ( + hasattr(args.quantization, "linear_weight_layout") + and args.quantization.linear_weight_layout == "KN" + and hasattr(model_config, "tensor_parallel_shards") + and model_config.tensor_parallel_shards > 1 + ): + raise NotImplementedError( + "KN layout (q3f16_0 and q4f16_0) is not supported for tensor parallelism" + ) + model, _ = args.model.quantize[args.quantization.kind](model_config, args.quantization) + # Step 2. Exporting the model to TVM + logger.info("Exporting the model to TVM compiler") + mod, named_params, ext_mods = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + # Step 3. Running relax compilation pipeline + logger.info("Running optimizations using TVM") + additional_tirs = _apply_preproc_to_params_and_check_pipeline(named_params, model_config) + variable_bounds = _get_variable_bounds(model_config) + cuda_graph_symbolic_capture_hints = { + "batch_decode": ["batch_size"], + "batch_decode_to_last_hidden_states": ["batch_size"], + "batch_verify": ["batch_size", "seq_len"], + "batch_verify_to_last_hidden_states": ["batch_size", "seq_len"], + } + avs = _kwargs.get("active_vocab_size", None) + if avs is not None and avs <= 0: + avs = None + metadata = { + "model_type": args.model.name, + "quantization": args.quantization.name, + "context_window_size": getattr(model_config, "context_window_size", -1), + "sliding_window_size": getattr(model_config, "sliding_window_size", -1), + "attention_sink_size": getattr(model_config, "attention_sink_size", -1), + "prefill_chunk_size": model_config.prefill_chunk_size, + "tensor_parallel_shards": model_config.tensor_parallel_shards, + "pipeline_parallel_stages": getattr(model_config, "pipeline_parallel_stages", 1), + "disaggregation": getattr(model_config, "disaggregation", False), + "kv_state_kind": _infer_kv_state_kind(args.model.name), + "max_batch_size": getattr(model_config, "max_batch_size", 1), + "active_vocab_size": avs, + "model_task": args.model.model_task, + } + if args.model.embedding_metadata: + metadata["embedding_metadata"] = dataclasses.asdict(args.model.embedding_metadata) + logger.info("Registering metadata: %s", metadata) + metadata["params"] = [_get_param_metadata(name, param) for name, param in named_params] + pass_config = {"relax.backend.use_cuda_graph": args.opt.cudagraph} + # TODO: Remove this workaround when the TVM CSE regression is fixed. + # Temporary workaround for TVM CSE regression that can produce + # dangling `cse_v*` vars during host codegen. + pass_config["tirx.disable_cse_tir"] = True + + with PassContext(config=pass_config): + args.build_func( + mod, + args, + pipeline=relax.get_pipeline( + "mlc_llm", + target=args.target, + flashinfer=args.opt.flashinfer, + cublas_gemm=args.opt.cublas_gemm, + faster_transformer=args.opt.faster_transformer, + allreduce_strategy=args.opt.ipc_allreduce_strategy, + variable_bounds=variable_bounds, + cuda_graph_symbolic_capture_hints=cuda_graph_symbolic_capture_hints, + additional_tirs=additional_tirs, + ext_mods=ext_mods, + metadata=metadata, + debug_dump=args.debug_dump, + ), + ) + _report_memory_usage(metadata=metadata, config=model_config) + logger.info("Generated: %s", bold(str(args.output))) + + +def compile( + config: Dict[str, Any], # noqa: UP006 + quantization: Quantization, + model_type: Model, + target: Target, + opt: OptimizationFlags, + build_func: Callable[[IRModule, CompileArgs, Pass], None], + system_lib_prefix: str, + output: Path, + overrides: ModelConfigOverride, + debug_dump: Optional[Path] = None, +): + """Compile a model given its configuration and quantization format to a specific target.""" + avs = None + if "active_vocab_size" in config: + avs = config.pop("active_vocab_size") + logger.info("Active vocab size from input config: %s", str(avs)) + if "model_config" in config: + model_config = config.pop("model_config") + model_config.update(config) + model_config = model_type.config.from_dict(model_config) + else: + model_config = model_type.config.from_dict(config) + model_config.kwargs = {"active_vocab_size": avs} if avs is not None else {} + args = CompileArgs( + model_config, + quantization, + model_type, + target, + opt, + build_func, + system_lib_prefix, + output, + overrides, + debug_dump, + ) + args.display() + _compile(args, model_config) diff --git a/python/mlc_llm/interface/compiler_flags.py b/python/mlc_llm/interface/compiler_flags.py new file mode 100644 index 0000000..b293dbf --- /dev/null +++ b/python/mlc_llm/interface/compiler_flags.py @@ -0,0 +1,227 @@ +"""Flags for overriding model config.""" + +import dataclasses +import enum +from io import StringIO +from typing import Optional + +from mlc_llm.support import argparse, logging +from mlc_llm.support.config import ConfigOverrideBase + +logger = logging.getLogger(__name__) + + +class IPCAllReduceStrategyType(enum.IntEnum): + """The all-reduce strategy.""" + + NONE = 0 + ONESHOT = 1 + TWOSHOT = 2 + AUTO = 3 + + +@dataclasses.dataclass +class OptimizationFlags: + """Optimization flags""" + + flashinfer: bool = False + cublas_gemm: bool = False + faster_transformer: bool = False + cudagraph: bool = False + cutlass: bool = False + ipc_allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE + + def __repr__(self) -> str: + out = StringIO() + print(f"flashinfer={int(self.flashinfer)}", file=out, end="") + print(f";cublas_gemm={int(self.cublas_gemm)}", file=out, end="") + print(f";faster_transformer={int(self.faster_transformer)}", file=out, end="") + print(f";cudagraph={int(self.cudagraph)}", file=out, end="") + print(f";cutlass={int(self.cutlass)}", file=out, end="") + print( + f";ipc_allreduce_strategy={self.ipc_allreduce_strategy.name}", + file=out, + end="", + ) + return out.getvalue().rstrip() + + @staticmethod + def from_str(source: str) -> "OptimizationFlags": + """Parse optimization flags from a string.""" + + if source in OPT_FLAG_PRESET: + return OPT_FLAG_PRESET[source] + + def boolean(value: str) -> bool: + if value == "0": + return False + if value == "1": + return True + raise ValueError(f"Invalid boolean value: {value}") + + parser = argparse.ArgumentParser(description="optimization flags") + parser.add_argument("--flashinfer", type=boolean, default=True) + parser.add_argument("--cublas_gemm", type=boolean, default=False) + parser.add_argument("--faster_transformer", type=boolean, default=False) + parser.add_argument("--cudagraph", type=boolean, default=False) + parser.add_argument("--cutlass", type=boolean, default=False) + parser.add_argument( + "--ipc_allreduce_strategy", + type=str, + choices=["NONE", "ONESHOT", "TWOSHOT", "AUTO"], + default="NONE", + ) + results = parser.parse_args([f"--{i}" for i in source.split(";") if i]) + return OptimizationFlags( + flashinfer=results.flashinfer, + cublas_gemm=results.cublas_gemm, + faster_transformer=results.faster_transformer, + cudagraph=results.cudagraph, + cutlass=results.cutlass, + ipc_allreduce_strategy=IPCAllReduceStrategyType[results.ipc_allreduce_strategy], + ) + + def update(self, target, quantization) -> None: + """Update optimization flags based on additional information.""" + + def _flashinfer(target) -> bool: + from mlc_llm.support.auto_target import ( + detect_cuda_arch_list, + ) + + if not self.flashinfer: + return False + if target.kind.name != "cuda": + return False + arch_list = detect_cuda_arch_list(target) + for arch in arch_list: + if arch < 80: + logger.warning("flashinfer is not supported on CUDA arch < 80") + return False + return True + + def _cublas_gemm(target, quantization) -> bool: + """correct cublas_gemm flag""" + if target.kind.name not in ["cuda", "rocm"]: + return False + if not ( + quantization.name in ["q0f16", "q0bf16", "q0f32"] + or "e4m3" in quantization.name + or "e5m2" in quantization.name + ): + return False + return self.cublas_gemm + + def _faster_transformer(target) -> bool: + """correct faster_transformer flag""" + if not target.kind.name == "cuda": + return False + return self.faster_transformer + + def _cutlass(target) -> bool: + """correct cutlass flag""" + if not target.kind.name == "cuda": + return False + return self.cutlass + + def _cudagraph(target) -> bool: + """correct cudagraph flag""" + if not target.kind.name == "cuda": + return False + return self.cudagraph + + self.flashinfer = _flashinfer(target) + self.cublas_gemm = _cublas_gemm(target, quantization) + self.faster_transformer = _faster_transformer(target) + self.cutlass = _cutlass(target) + self.cudagraph = _cudagraph(target) + + +@dataclasses.dataclass +class ModelConfigOverride(ConfigOverrideBase): + """Flags for overriding model config.""" + + context_window_size: Optional[int] = None + sliding_window_size: Optional[int] = None + prefill_chunk_size: Optional[int] = None + attention_sink_size: Optional[int] = None + max_batch_size: Optional[int] = None + tensor_parallel_shards: Optional[int] = None + pipeline_parallel_stages: Optional[int] = None + disaggregation: Optional[bool] = None + + def __repr__(self) -> str: + out = StringIO() + print(f"context_window_size={self.context_window_size}", file=out, end="") + print(f";sliding_window_size={self.sliding_window_size}", file=out, end="") + print(f";prefill_chunk_size={self.prefill_chunk_size}", file=out, end="") + print(f";attention_sink_size={self.attention_sink_size}", file=out, end="") + print(f";max_batch_size={self.max_batch_size}", file=out, end="") + print(f";tensor_parallel_shards={self.tensor_parallel_shards}", file=out, end="") + print( + f";pipeline_parallel_stages={self.pipeline_parallel_stages}", + file=out, + end="", + ) + print(f";disaggregation={self.disaggregation}", file=out, end="") + return out.getvalue().rstrip() + + @staticmethod + def from_str(source: str) -> "ModelConfigOverride": + """Parse model config override values from a string.""" + parser = argparse.ArgumentParser(description="model config override values") + parser.add_argument("--context_window_size", type=int, default=None) + parser.add_argument("--sliding_window_size", type=int, default=None) + parser.add_argument("--prefill_chunk_size", type=int, default=None) + parser.add_argument("--attention_sink_size", type=int, default=None) + parser.add_argument("--max_batch_size", type=int, default=None) + parser.add_argument("--tensor_parallel_shards", type=int, default=None) + parser.add_argument("--pipeline_parallel_stages", type=int, default=None) + parser.add_argument( + "--disaggregation", + type=lambda x: str(x).lower() in ["true", "1", "yes", "True"], + default=None, + ) + results = parser.parse_args([f"--{i}" for i in source.split(";") if i]) + return ModelConfigOverride( + context_window_size=results.context_window_size, + sliding_window_size=results.sliding_window_size, + prefill_chunk_size=results.prefill_chunk_size, + attention_sink_size=results.attention_sink_size, + max_batch_size=results.max_batch_size, + tensor_parallel_shards=results.tensor_parallel_shards, + pipeline_parallel_stages=results.pipeline_parallel_stages, + disaggregation=results.disaggregation, + ) + + +OPT_FLAG_PRESET = { + "O0": OptimizationFlags( + flashinfer=False, + cublas_gemm=False, + cudagraph=False, + ), + "O1": OptimizationFlags( + flashinfer=False, + cublas_gemm=True, + faster_transformer=True, + cudagraph=False, + cutlass=True, + ), + "O2": OptimizationFlags( + flashinfer=True, + cublas_gemm=True, + faster_transformer=False, + cudagraph=True, + cutlass=True, + ipc_allreduce_strategy=IPCAllReduceStrategyType.NONE, + ), + "O3": OptimizationFlags( + flashinfer=True, + cublas_gemm=True, + faster_transformer=True, + cudagraph=True, + cutlass=True, + ipc_allreduce_strategy=IPCAllReduceStrategyType.AUTO, + ), +} diff --git a/python/mlc_llm/interface/convert_weight.py b/python/mlc_llm/interface/convert_weight.py new file mode 100644 index 0000000..c5f3325 --- /dev/null +++ b/python/mlc_llm/interface/convert_weight.py @@ -0,0 +1,250 @@ +"""Python entrypoint of weight conversion.""" + +import contextlib +import dataclasses +import math +import os +import tempfile +from collections.abc import Iterator +from io import StringIO +from pathlib import Path +from typing import Any, Dict, Optional, Tuple # noqa: UP035 + +from tvm import tirx +from tvm.contrib import tvmjs +from tvm.runtime import DataType, Device, Tensor +from tvm.runtime import cpu as cpu_device +from tvm.target import Target + +from mlc_llm.loader import LOADER +from mlc_llm.model import Model +from mlc_llm.quantization import Quantization +from mlc_llm.support import logging, tqdm +from mlc_llm.support.auto_weight import detect_weight +from mlc_llm.support.preshard import apply_preshard +from mlc_llm.support.style import bold, green + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class ConversionArgs: + """Arguments to MLC LLM's weight conversation and quantization flow.""" + + config: Path + quantization: Quantization + model: Model + device: Device + source: Path + source_format: str + output: Path + lora_adapter: Optional[Path] = None + + def display(self) -> None: + """Display the arguments to stdout.""" + + def _device_to_str(device: Device) -> str: + return f"{Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]}:{device.index}" + + out = StringIO() + print(f"{bold('Weight conversion with arguments:')}", file=out) + print(f" {bold('--config'):<25} {self.config}", file=out) + print(f" {bold('--quantization'):<25} {self.quantization}", file=out) + print(f" {bold('--model-type'):<25} {self.model.name}", file=out) + print(f" {bold('--device'):<25} {_device_to_str(self.device)}", file=out) + print(f" {bold('--source'):<25} {self.source}", file=out) + print(f" {bold('--source-format'):<25} {self.source_format}", file=out) + print(f" {bold('--output'):<25} {self.output}", file=out) + if self.lora_adapter is not None: + print(f" {bold('--lora-adapter'):<25} {self.lora_adapter}", file=out) + print(out.getvalue().rstrip()) + + +def _resolve_base_model_dir(source: Path) -> Path: + return source if source.is_dir() else source.parent + + +@contextlib.contextmanager +def _merge_lora_adapter_with_base_model(base_source: Path, lora_adapter: Path) -> Iterator[Path]: + base_model_dir = _resolve_base_model_dir(base_source) + if not base_model_dir.exists(): + raise ValueError(f"Base model directory does not exist: {base_model_dir}") + if not lora_adapter.exists() or not lora_adapter.is_dir(): + raise ValueError(f"LoRA adapter directory does not exist: {lora_adapter}") + + try: + from peft import PeftModel + from transformers import AutoModelForCausalLM + + except ImportError as err: + raise ImportError( + "`--lora-adapter` requires `peft` and `transformers` to be installed." + ) from err + + with tempfile.TemporaryDirectory() as temp_dir: + merged_model_dir = Path(temp_dir) / "merged_model" + logger.info("Merging LoRA adapter %s into base model %s", lora_adapter, base_model_dir) + + base_model = AutoModelForCausalLM.from_pretrained( + str(base_model_dir), + torch_dtype="auto", + trust_remote_code=False, + low_cpu_mem_usage=True, + ) + merged_model = PeftModel.from_pretrained( + base_model, str(lora_adapter), is_trainable=False + ).merge_and_unload() + merged_model.save_pretrained(str(merged_model_dir), safe_serialization=True) + yield merged_model_dir + + +def _convert_args(args: ConversionArgs) -> None: + pre_shards_num = os.getenv("MLC_INTERNAL_PRESHARD_NUM") + # model config & quantization config + model_config = args.model.config.from_file(args.config) + if ( + args.quantization.kind == "ft-quant" + and hasattr(model_config, "tensor_parallel_shards") + and model_config.tensor_parallel_shards > 1 + ): + raise NotImplementedError + if pre_shards_num is not None: + model_config.tensor_parallel_shards = int(pre_shards_num) + model, quantize_map = args.model.quantize[args.quantization.kind]( + model_config, args.quantization + ) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_params = dict(_named_params) + + if pre_shards_num is not None: + named_params, preshard_funcs = apply_preshard(named_params, int(pre_shards_num), args) + else: + preshard_funcs = None + + def _check_param(name: str, param: Tensor): + nonlocal named_params + if name not in named_params: + raise ValueError(f"Parameter not found in model: {name}") + if name in param_names: + raise ValueError(f"Duplication: Parameter {name} already computed") + + # Check shape (possibly dynamic) + def _check_shape(actual: tuple, expect: tuple): # expect can have tirx.Var + if len(actual) != len(expect): + return False + for actual_i, expect_i in zip(actual, expect): + assert isinstance(expect_i, (int, tirx.Var)) + if isinstance(expect_i, int) and actual_i != expect_i: + return False + return True + + expect_shape = named_params[name].shape + actual_shape = param.shape + if not _check_shape(actual_shape, expect_shape): + raise ValueError( + f"Parameter {name} has shape {param.shape}, but expected {expect_shape}" + ) + # Check dtype + actual_dtype = param.dtype + expect_dtype = named_params[name].dtype + if actual_dtype != expect_dtype: + raise ValueError( + f"Parameter {name} has dtype {param.dtype}, but expected {expect_dtype}" + ) + del named_params[name] + + # load and quantize + param_names = set() + total_bytes = 0.0 + total_params: int = 0 + + def _param_generator() -> Iterator[Tuple[str, Tensor]]: # noqa: UP006 + nonlocal total_params, total_bytes + with Target.from_device(args.device), tqdm.redirect(): + loader = LOADER[args.source_format]( + path=args.source, + extern_param_map=args.model.source[args.source_format]( + model_config, args.quantization + ), + quantize_param_map=quantize_map, + ) + for name, param in loader.load(device=args.device, preshard_funcs=preshard_funcs): + _check_param(name, param) + param_names.add(name) + param = param.copyto(cpu_device()) + total_bytes += math.prod(param.shape) * DataType(param.dtype).itemsize + yield name, param + total_params = loader.stats.total_param_num + + def _metadata_callback() -> Dict[str, Any]: # noqa: UP006 + return { + "ParamSize": len(param_names), + "ParamBytes": total_bytes, + "BitsPerParam": total_bytes * 8.0 / total_params, + } + + # dump to output directory + tvmjs.dump_tensor_cache( + _param_generator(), + str(args.output), + meta_data=_metadata_callback, + encode_format="f32-to-bf16", + show_progress=False, + ) + if named_params: + raise ValueError(f"Parameter not found in source: {', '.join(named_params.keys())}") + # Log necessary statistics + logger.info( + "%s after quantization: %.3f GB", + green("Parameter size"), + total_bytes / (1024**3), + ) + logger.info(f"%s: {total_params:,}", green("Total parameters")) + logger.info( + "%s: %.3f", + green("Bits per parameter"), + total_bytes * 8.0 / total_params, + ) + logger.info("Saved to directory: %s", bold(str(args.output))) + + +def convert_weight( + config: Path, + quantization: Quantization, + model: Model, + device: Device, + source: Path, + source_format: str, + output: Path, + lora_adapter: Optional[Path] = None, +): + """MLC LLM's weight conversation and quantization flow.""" + args = ConversionArgs( + config, quantization, model, device, source, source_format, output, lora_adapter + ) + + allowed_lora_source_formats = {"huggingface-safetensor", "huggingface-torch"} + if lora_adapter is not None and source_format not in allowed_lora_source_formats: + raise ValueError( + f"`--lora-adapter` only supports source formats: {sorted(allowed_lora_source_formats)}" + ) + + if lora_adapter is not None: + with _merge_lora_adapter_with_base_model(source, lora_adapter) as merged_model_dir: + merged_source, merged_source_format = detect_weight( + weight_path=merged_model_dir, + config_json_path=config, + weight_format="auto", + ) + merged_args = dataclasses.replace( + args, source=merged_source, source_format=merged_source_format + ) + merged_args.display() + _convert_args(merged_args) + return + + args.display() + _convert_args(args) diff --git a/python/mlc_llm/interface/gen_config.py b/python/mlc_llm/interface/gen_config.py new file mode 100644 index 0000000..3be0193 --- /dev/null +++ b/python/mlc_llm/interface/gen_config.py @@ -0,0 +1,359 @@ +"""Generator of mlc-chat-config.json and tokenizer configuration.""" + +import dataclasses +import json +import re +import shutil +from dataclasses import asdict +from pathlib import Path +from typing import Optional + +from mlc_llm.conversation_template import ConvTemplateRegistry +from mlc_llm.model import Model +from mlc_llm.protocol.mlc_chat_config import MLCChatConfig +from mlc_llm.quantization import Quantization +from mlc_llm.support import convert_tiktoken, logging +from mlc_llm.support.style import bold, green, red +from mlc_llm.tokenizers import Tokenizer + +from .compiler_flags import ModelConfigOverride + +logger = logging.getLogger(__name__) + +FOUND = green("Found") +NOT_FOUND = red("Not found") +FAILED = red("Failed") + + +def apply_system_defaults_for_missing_fields(mlc_chat_config: MLCChatConfig) -> None: + """Apply system default value.""" + for key, value in mlc_chat_config.get_system_defaults_for_missing_fields().items(): + setattr(mlc_chat_config, key, value) + logger.info("[System default] Setting %s: %s", bold(key), value) + + +def check_string(s: str) -> bool: + """Check whether it's a string.""" + s = s[1:] if s[0] == "b" else s + delimit = s[0] + if s[-1] != delimit or delimit not in ["'", '"']: + return False + for i in range(1, len(s) - 1): + if s[i] == delimit and s[i - 1] != "\\": + return False + return True + + +def txt2rwkv_tokenizer(vocab: Path, out: Path) -> None: + """Generate tokenizer_model from RWKV vocab file.""" + idx2token = {} + + with vocab.open("r", encoding="utf-8") as f: + lines = f.readlines() + + for line in lines: + idx = int(line[: line.index(" ")]) + raw = line[line.index(" ") : line.rindex(" ")].strip() + if check_string(raw): + x = eval(raw) + x = x.encode("utf-8") if isinstance(x, str) else x + assert isinstance(x, bytes) + assert len(x) == int(line[line.rindex(" ") :]) + idx2token[idx] = x + else: + raise ValueError("Unsupported vocab dictionary") + + with (out / "tokenizer_model").open("wb") as f: + import msgpack + + msgpack.pack(idx2token, f) + + +def json2rwkv_tokenizer(vocab: Path, out: Path) -> None: + """Generate tokenizer_model from RWKV vocab file.""" + idx2token = {} + + with vocab.open("r", encoding="utf-8") as f: + data = json.load(f) + for key, value in data.items(): + x = key.encode("utf-8") if isinstance(key, str) else key + assert isinstance(x, bytes) + idx2token[int(value)] = x + + with (out / "tokenizer_model").open("wb") as f: + import msgpack + + msgpack.pack(idx2token, f) + + +def gen_config( + config: Path, + model: Model, + quantization: Quantization, + conv_template: str, + context_window_size: Optional[int], + sliding_window_size: Optional[int], + prefill_chunk_size: Optional[int], + attention_sink_size: Optional[int], + tensor_parallel_shards: Optional[int], + pipeline_parallel_stages: Optional[int], + disaggregation: Optional[bool], + max_batch_size: int, + output: Path, +): + """Entrypoint of MLC Chat configuration generation.""" + # Step 1. Initialize `mlc-chat-config.json` using `config.json` + conversation_reg = ConvTemplateRegistry.get_conv_template(conv_template) + if conversation_reg is None: + logger.warning( + "%s: Conversation template is not registered in ConvTemplateRegistry: %s", + red("Warning"), + conv_template, + ) + conversation = conv_template + else: + conversation = conversation_reg.to_json_dict() + + model_config = ModelConfigOverride( + context_window_size=context_window_size, + sliding_window_size=sliding_window_size, + prefill_chunk_size=prefill_chunk_size, + attention_sink_size=attention_sink_size, + max_batch_size=max_batch_size, + tensor_parallel_shards=tensor_parallel_shards, + pipeline_parallel_stages=pipeline_parallel_stages, + disaggregation=disaggregation, + ).apply(model.config.from_file(config)) + mlc_chat_config = MLCChatConfig( + model_type=model.name, + quantization=quantization.name, + model_config=model_config.asdict(), + vocab_size=model_config.vocab_size, + active_vocab_size=getattr(model_config, "active_vocab_size", model_config.vocab_size), + context_window_size=getattr(model_config, "context_window_size", -1), + sliding_window_size=getattr(model_config, "sliding_window_size", -1), + prefill_chunk_size=model_config.prefill_chunk_size, + attention_sink_size=getattr(model_config, "attention_sink_size", -1), + tensor_parallel_shards=model_config.tensor_parallel_shards, + pipeline_parallel_stages=getattr(model_config, "pipeline_parallel_stages", 1), + disaggregation=getattr(model_config, "disaggregation", False), + conv_template=conversation, + model_task=model.model_task, + embedding_metadata=( + dataclasses.asdict(model.embedding_metadata) if model.embedding_metadata else None + ), + ) + # Step 2. Load `generation_config.json` and `config.json` for text-generation related configs + for generation_config_filename in ["generation_config.json", "config.json"]: + generation_config = config.parent / generation_config_filename + if generation_config.exists(): + with generation_config.open("r", encoding="utf-8") as in_file: + generation_config_json = json.load(in_file) + for key, value in generation_config_json.items(): + if hasattr(mlc_chat_config, key) and getattr(mlc_chat_config, key) is None: + setattr(mlc_chat_config, key, value) + logger.info( + "[%s] Setting %s: %s", + generation_config_filename, + bold(key), + value, + ) + else: + logger.info("%s %s: %s", NOT_FOUND, generation_config_filename, generation_config) + + # Step 3. Copy tokenizer configuration + # 3.1. Copy over the files and populate mlc_chat_config + for filename in TOKENIZER_FILES: + file = config.parent / filename + if file.exists(): + mlc_chat_config.tokenizer_files.append(filename) + dest = output / filename + shutil.copy(file, dest) + logger.info("%s tokenizer config: %s. Copying to %s", FOUND, file, bold(str(dest))) + else: + logger.info("%s tokenizer config: %s", NOT_FOUND, file) + # 3.2. Generate `tokenizer_model` for rwkv if `rwkv_vocab_.*` is found + pattern = re.compile(r"rwkv_vocab_v\d{8}\.(json|txt)") + for item in config.parent.iterdir(): + if item.is_file() and pattern.match(item.name): + logger.info( + "%s RWKV vocab file: %s. Genetating %s", + FOUND, + item, + bold("tokenizer_model"), + ) + if item.name.endswith(".txt"): + txt2rwkv_tokenizer(item, output) + else: + json2rwkv_tokenizer(item, output) + # 3.3. If we have `tokenizer.model` but not `tokenizer.json`, try convert it to + # `tokenizer.json` with `transformers`. + tokenizer_json_file = config.parent / "tokenizer.json" + tokenizer_model_file = config.parent / "tokenizer.model" + if tokenizer_model_file.exists() and (not tokenizer_json_file.exists()): + logger.info( + "The model has `tokenizer.model` but not `tokenizer.json`. " + "It is always recommended to prefer JSON instead. " + "Attempting to convert using HuggingFace transformers library" + ) + try: + from transformers import ( + AutoTokenizer, + ) + + tokenizer_json_save_dest = output / "tokenizer.json" + fast_tokenizer = AutoTokenizer.from_pretrained(str(config.parent), use_fast=True) + fast_tokenizer.backend_tokenizer.save(str(tokenizer_json_save_dest)) + mlc_chat_config.tokenizer_files.append("tokenizer.json") + logger.info( + "Successfully converted `tokenizer.model` to: %s", + tokenizer_json_save_dest, + ) + except Exception: + logger.warning( + "Converting to `tokenizer.json` %s with the exception below. " + "Skipping the conversion.", + FAILED, + exc_info=True, + ) + # 3.3. If we still don't have "tokenizer.json" at this point, try looking for "*.tiktoken" files + if (not tokenizer_json_file.exists()) and list(config.parent.glob("*.tiktoken")): + try: + logger.info( + "The model has tiktoken files but not `tokenizer.json`. " + "Attempting to convert from tiktoken files" + ) + convert_tiktoken.convert_tiktoken( + str(config.parent), str(output), mlc_chat_config.context_window_size + ) + mlc_chat_config.tokenizer_files.append("tokenizer.json") + mlc_chat_config.tokenizer_files.append("vocab.json") + mlc_chat_config.tokenizer_files.append("merges.txt") + mlc_chat_config.tokenizer_files.append("special_tokens_map.json") + logger.info("Succesfully converted from tiktoken files to: %s", str(output)) + except Exception: + logger.exception("%s with the exception below. Skipping", FAILED) + + # 3.4. Detect tokenizer info + mlc_chat_config.tokenizer_info = asdict(Tokenizer.detect_tokenizer_info(str(output))) + logger.info("Detected tokenizer info: %s", mlc_chat_config.tokenizer_info) + + # 3.5. Ensure added_tokens do not have duplicated added_tokens, a mistake from model releaser + # that affects correctness of huggingface tokenizer. + # See https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B/discussions/15. + if tokenizer_json_file.exists(): + with open(tokenizer_json_file, encoding="utf-8") as f: + tokenizer_json = json.load(f) + if "added_tokens" in tokenizer_json: + appeared_content = set() + for added_token in tokenizer_json["added_tokens"]: + content = added_token["content"] + if content in appeared_content: + logger.exception( + "%s with incorrect tokenizer.json which has duplicated token %s. " + "This affects correctness of huggingface tokenizer during runtime, " + "please check your tokenizer.json to remove duplication manually.", + FAILED, + content, + ) + raise ValueError("Duplicated vocab in tokenizer.json") + appeared_content.add(content) + + # Step 4. Load system default value + apply_system_defaults_for_missing_fields(mlc_chat_config) + + # Step 5. Use HF tokenizer to detect active vocab size via len(tokenizer) + if tokenizer_json_file.exists(): + try: + from transformers import ( + AutoTokenizer, + ) + + hf_tokenizer = AutoTokenizer.from_pretrained(str(config.parent), use_fast=True) + active_vocab_size = len(hf_tokenizer) + if mlc_chat_config.active_vocab_size != active_vocab_size: + logger.info( + "Overriding active_vocab_size from %d to %d using HF tokenizer", + mlc_chat_config.active_vocab_size, + active_vocab_size, + ) + mlc_chat_config.active_vocab_size = active_vocab_size + except Exception: + logger.warning( + "Detecting active_vocab_size %s with the exception below. Skipping.", + FAILED, + exc_info=True, + ) + + # Step 5. Dump the configuration file to output directory + with (output / "mlc-chat-config.json").open("w", encoding="utf-8") as out_file: + json.dump(mlc_chat_config.model_dump(by_alias=True), out_file, indent=2) + logger.info("Dumping configuration file to: %s", bold(out_file.name)) + + +TOKENIZER_FILES = [ + "tokenizer.model", + "tokenizer.json", + "vocab.json", + "merges.txt", + "added_tokens.json", + "tokenizer_config.json", +] +# FIXME: Copy RWKV tokenizer file + +CONV_TEMPLATES = { + "llama-4", + "llama-3", + "llama-3_1", + "chatml", + "chatml_nosystem", + "qwen2", + "open_hermes_mistral", + "neural_hermes_mistral", + "llama_default", + "llama-2", + "mistral_default", + "ministral3", + "ministral3_reasoning", + "gpt2", + "codellama_completion", + "codellama_instruct", + "redpajama_chat", + "rwkv_world", + "gorilla", + "gorilla-openfunctions-v2", + "dolly", + "oasst", + "stablelm", + "LM", + "stablelm-3b", + "gpt_bigcode", + "wizardlm_7b", + "wizard_coder_or_math", + "glm", + "phi-2", + "phi-3", + "phi-3-vision", + "phi-4", + "stablelm-2", + "gemma_instruction", + "gemma3_instruction", + "orion", + "llava", + "hermes2_pro_llama3", + "hermes3_llama-3_1", + "tinyllama_v1_0", + "aya-23", + "deepseek", + "deepseek_v2", + "deepseek_v3", + "deepseek_r1_qwen", + "deepseek_r1_llama", + "olmo", + "olmo2", + "nemotron", + "llm-jp", + "qwen3", + "qwen3_5", + "qwen3_5_nothink", +} diff --git a/python/mlc_llm/interface/help.py b/python/mlc_llm/interface/help.py new file mode 100644 index 0000000..28bbc0f --- /dev/null +++ b/python/mlc_llm/interface/help.py @@ -0,0 +1,273 @@ +"""Help message for CLI arguments.""" + +HELP = { + "config": ( + """ +1) Path to a HuggingFace model directory that contains a `config.json` or +2) Path to `config.json` in HuggingFace format, or +3) The name of a pre-defined model architecture. + +A `config.json` file in HuggingFace format defines the model architecture, including the vocabulary +size, the number of layers, the hidden size, number of attention heads, etc. +Example: https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json. + +A HuggingFace directory often contains a `config.json` which defines the model architecture, +the non-quantized model weights in PyTorch or SafeTensor format, tokenizer configurations, +as well as an optional `generation_config.json` provides additional default configuration for +text generation. +Example: https://huggingface.co/codellama/CodeLlama-7b-hf/tree/main. +""" + ).strip(), + "quantization": """ +The quantization mode we use to compile. If unprovided, will infer from `model`. +""".strip(), + "model": """ +A path to ``mlc-chat-config.json``, or an MLC model directory that contains `mlc-chat-config.json`. +It can also be a link to a HF repository pointing to an MLC compiled model. +""".strip(), + "model_lib": """ +The full path to the model library file to use (e.g. a ``.so`` file). If unspecified, we will use +the provided ``model`` to search over possible paths. It the model lib is not found, it will be +compiled in a JIT manner. +""".strip(), + "model_type": """ +Model architecture such as "llama". If not set, it is inferred from `mlc-chat-config.json`. +""".strip(), + "device_compile": """ +The GPU device to compile the model to. If not set, it is inferred from GPUs available locally. +""".strip(), + "enable_subgroups": """ +Enable WebGPU subgroups in codegen. This only applies to WebGPU targets and will set +supports_subgroups accordingly. +""".strip(), + "device_quantize": """ +The device used to do quantization such as "cuda" or "cuda:0". Will detect from local available GPUs +if not specified. +""".strip(), + "device_deploy": """ +The device used to deploy the model such as "cuda" or "cuda:0". Will detect from local +available GPUs if not specified. +""".strip(), + "host": """ +The host LLVM triple to compile the model to. If not set, it is inferred from the local CPU and OS. +Examples of the LLVM triple: +1) iPhones: arm64-apple-ios; +2) ARM64 Android phones: aarch64-linux-android; +3) WebAssembly: wasm32-unknown-unknown-wasm; +4) Windows: x86_64-pc-windows-msvc; +5) ARM macOS: arm64-apple-darwin. +""".strip(), + "opt": """ +Optimization flags. MLC LLM maintains a predefined set of optimization flags, +denoted as O0, O1, O2, O3, where O0 means no optimization, O2 means majority of them, +and O3 represents extreme optimization that could potentially break the system. +Meanwhile, optimization flags could be explicitly specified via details knobs, e.g. +--opt="cublas_gemm=1;cudagraph=0". +""".strip(), + "system_lib_prefix": """ +Adding a prefix to all symbols exported. Similar to "objcopy --prefix-symbols". +This is useful when compiling multiple models into a single library to avoid symbol +conflicts. Different from objcopy, this takes no effect for shared library. +""".strip(), + "context_window_size": """ +Option to provide the maximum sequence length supported by the model. +This is usually explicitly shown as context length or context window in the model card. +If this option is not set explicitly, by default, +it will be determined by `context_window_size` or `max_position_embeddings` in `config.json`, +and the latter is usually inaccurate for some models. +""".strip(), + "output_compile": """ +The path to the output file. The suffix determines if the output file is a shared library or +objects. Available suffixes: +1) Linux: .so (shared), .tar (objects); +2) macOS: .dylib (shared), .tar (objects); +3) Windows: .dll (shared), .tar (objects); +4) Android, iOS: .tar (objects); +5) Web: .wasm (web assembly). +""".strip(), + "source": """ +The path to original model weight, infer from `config` if missing. +""".strip(), + "source_format": """ +The format of source model weight, infer from `config` if missing. +""".strip(), + "output_quantize": """ +The output directory to save the quantized model weight. Will create `params_shard_*.bin` and +`tensor-cache.json` in this directory. +""".strip(), + "conv_template": """ +Conversation template. It depends on how the model is tuned. Use "LM" for vanilla base model +""".strip(), + "output_gen_mlc_chat_config": """ +The output directory for generated configurations, including `mlc-chat-config.json` and tokenizer +configuration. +""".strip(), + "sliding_window_size": """ +(Experimental) The sliding window size in sliding window attention (SWA). +This optional field overrides the `sliding_window_size` in config.json for +those models that use SWA. Currently only useful when compiling Mistral. +This flag subjects to future refactoring. +""".strip(), + "prefill_chunk_size": """ +(Experimental) The chunk size during prefilling. By default, +the chunk size is the same as sliding window or max sequence length. +This flag subjects to future refactoring. +""".strip(), + "attention_sink_size": """ +(Experimental) The number of stored sinks. Only supported on Mistral yet. By default, +the number of sinks is 4. This flag subjects to future refactoring. +""".strip(), + "max_batch_size": """ +The maximum allowed batch size set for the KV cache to concurrently support. +""".strip(), + """tensor_parallel_shards""": """ +Number of shards to split the model into in tensor parallelism multi-gpu inference. +""".strip(), + """pipeline_parallel_stages""": """ +Number of pipeline stages to split the model layers for pipeline parallelism. +""".strip(), + """disaggregation""": """ +Whether enable disaggregation when compiling the model. +""".strip(), + "overrides": """ +Model configuration override. Configurations to override `mlc-chat-config.json`. Supports +`context_window_size`, `prefill_chunk_size`, `sliding_window_size`, `attention_sink_size`, +`max_batch_size` and `tensor_parallel_shards`. Meanwhile, model config could be explicitly +specified via details knobs, e.g. --overrides "context_window_size=1024;prefill_chunk_size=128". +""".strip(), + "modelconfig_overrides": """ +Model configuration override. Supports overriding, +`context_window_size`, `prefill_chunk_size`, `sliding_window_size`, `attention_sink_size`, +`max_num_sequence` and `tensor_parallel_shards`. The overrides could be explicitly +specified via details knobs, e.g. --overrides "context_window_size=1024;prefill_chunk_size=128". +""".strip(), + "debug_dump": """ +Specifies the directory where the compiler will store its IRs for debugging purposes +during various phases of compilation. By default, this is set to `None`, indicating +that debug dumping is disabled. +""".strip(), + "prompt": """ +The prompt of the text generation. +""".strip(), + "generate_length": """ +The target length of the text generation. +""".strip(), + "max_total_sequence_length_serve": """ +The KV cache total token capacity, i.e., the maximum total number of tokens that +the KV cache support. This decides the GPU memory size that the KV cache consumes. +If not specified, system will automatically estimate the maximum capacity based +on the vRAM size on GPU. +""".strip(), + "prefill_chunk_size_serve": """ +The maximum number of tokens the model passes for prefill each time. +It should not exceed the prefill chunk size in model config. +If not specified, this defaults to the prefill chunk size in model config. +""".strip(), + "max_history_size_serve": """ +The maximum history length for rolling back the RNN state. +If unspecified, the default value is 1. +KV cache does not need this. +""".strip(), + "enable_tracing_serve": """ +Enable Chrome Tracing for the server. +After enabling, you can send POST request to the "debug/dump_event_trace" entrypoint +to get the Chrome Trace. For example, +"curl -X POST http://127.0.0.1:8000/debug/dump_event_trace -H "Content-Type: application/json" -d '{"model": "dist/llama"}'" +""".strip(), # noqa: E501 + "mode_serve": """ +The engine mode in MLC LLM. We provide three preset modes: "local", "interactive" and "server". +The default mode is "local". +The choice of mode decides the values of "max_num_sequence", "max_total_seq_length" and +"prefill_chunk_size" when they are not explicitly specified. +1. Mode "local" refers to the local server deployment which has low request concurrency. + So the max batch size will be set to 4, and max total sequence length and prefill chunk size + are set to the context window size (or sliding window size) of the model. +2. Mode "interactive" refers to the interactive use of server, which has at most 1 concurrent + request. So the max batch size will be set to 1, and max total sequence length and prefill + chunk size are set to the context window size (or sliding window size) of the model. +3. Mode "server" refers to the large server use case which may handle many concurrent request + and want to use GPU memory as much as possible. In this mode, we will automatically infer + the largest possible max batch size and max total sequence length. +You can manually specify arguments "max_num_sequence", "max_total_seq_length" and +"prefill_chunk_size" via "--overrides" to override the automatic inferred values. +For example: --overrides "max_num_sequence=32;max_total_seq_length=4096" +""".strip(), + "additional_models_serve": """ +The model paths and (optional) model library paths of additional models (other than the main model). +When engine is enabled with speculative decoding, additional models are needed. +The way of specifying additional models is: +"--additional-models model_path_1 model_path_2 ..." or +"--additional-models model_path_1,model_lib_1 model_path_2 ...". +When the model lib of a model is not given, JIT model compilation will be activated +to compile the model automatically. +""".strip(), + "gpu_memory_utilization_serve": """ +A number in (0, 1) denoting the fraction of GPU memory used by the server in total. +It is used to infer to maximum possible KV cache capacity. +When it is unspecified, it defaults to 0.85. +Under mode "local" or "interactive", the actual memory usage may be significantly smaller than +this number. Under mode "server", the actual memory usage may be slightly larger than this number. +""".strip(), + "speculative_mode_serve": """ +The speculative decoding mode. Right now four options are supported: + - "disable", where speculative decoding is not enabled, + - "small_draft", denoting the normal speculative decoding (small draft) style, + - "eagle", denoting the eagle-style speculative decoding. + - "medusa", denoting the medusa-style speculative decoding. +The default mode is "disable". +""".strip(), + "spec_draft_length_serve": """ +The number of draft tokens to generate in speculative proposal. +Being 0 means to enable adaptive speculative mode, where the draft length will be +automatically adjusted based on engine state. The default values is 0. +""".strip(), + "prefix_cache_mode_serve": """ +The prefix cache mode. Right now two options are supported: + - "disable", where prefix cache is not enabled, + - "radix", denoting the normal paged radix tree based prefix cache, +The default mode is "radix". +""".strip(), + "prefix_cache_max_num_recycling_seqs_serve": """ +The maximum number of sequences in prefix cache, default as max_batch_size. +And set 0 to disable prefix cache, set -1 to have infinite capacity prefix cache. +""".strip(), + "prefill_mode": """ +The prefill mode. "chunked" means the basic prefill with chunked input enabled. "hybrid" means the +hybrid prefill or split-fuse, so that decode step will be converted into prefill. +""".strip(), + "overrides_serve": """ +Overriding extra configurable fields of EngineConfig and model compilation config. +Supporting fields that can be be overridden: "tensor_parallel_shards", "max_num_sequence", +"max_total_seq_length", "prefill_chunk_size", "max_history_size", "gpu_memory_utilization", +"spec_draft_length", "prefix_cache_max_num_recycling_seqs", "context_window_size", +"sliding_window_size", "attention_sink_size". +Please check out the documentation of EngineConfig in mlc_llm/serve/config.py for detailed docstring +of each field. +Example: --overrides "max_num_sequence=32;max_total_seq_length=4096;tensor_parallel_shards=2" +""".strip(), + "config_package": """ +The path to "mlc-package-config.json" which is used for package build. +See "https://github.com/mlc-ai/mlc-llm/blob/main/ios/MLCChat/mlc-package-config.json" as an example. +""".strip(), + "mlc_llm_source_dir": """ +The source code path to MLC LLM. +""".strip(), + "output_package": """ +The path of output directory for the package build outputs. +""".strip(), + "calibration_dataset": """ +The path to the calibration dataset. + """.strip(), + "num_calibration_samples": """ +The number of samples used for calibration. + """.strip(), + "output_calibration": """ +The output directory to save the calibration params. + """.strip(), + "seed_calibrate": """ +The seed to sample the calibration dataset.""", + "pd_balance_factor": """ +How much prefill to move to decode engine. For example, +0.1 means the last 10 percent tokens are prefilled by decode engine. + """.strip(), +} diff --git a/python/mlc_llm/interface/jit.py b/python/mlc_llm/interface/jit.py new file mode 100644 index 0000000..942d944 --- /dev/null +++ b/python/mlc_llm/interface/jit.py @@ -0,0 +1,181 @@ +"""Just-in-time compilation of MLC-Chat models.""" + +import dataclasses +import hashlib +import json +import os +import shlex +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any, Dict, Optional, Union # noqa: UP035 + +from tvm.runtime import Device + +from mlc_llm.model import MODELS +from mlc_llm.support import logging +from mlc_llm.support.auto_device import device2str +from mlc_llm.support.constants import ( + MLC_DSO_SUFFIX, + MLC_JIT_POLICY, + MLC_LLM_HOME, + MLC_TEMP_DIR, +) +from mlc_llm.support.style import blue, bold + +from .compiler_flags import ModelConfigOverride, OptimizationFlags + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class JITResult: + """The jit compilation result class.""" + + model_lib_path: str + system_lib_prefix: Optional[str] = None + + +def log_jit_policy(): + """log current jit policy""" + logger.info( + "%s = %s. Can be one of: ON, OFF, REDO, READONLY", + bold("MLC_JIT_POLICY"), + MLC_JIT_POLICY, + ) + + +def jit( + model_path: Path, + overrides: Dict[str, Any], # noqa: UP006 + device: Union[Device, str], + system_lib_prefix: Optional[str] = None, + *, + skip_log_jit_policy=False, +) -> JITResult: + """Just-in-time compile a MLC-Chat model.""" + # skip logging jit policy since when outside can hint once + if not skip_log_jit_policy: + log_jit_policy() + + if MLC_JIT_POLICY == "OFF": + raise RuntimeError("JIT is disabled by MLC_JIT_POLICY=OFF") + + with open(model_path / "mlc-chat-config.json", encoding="utf-8") as in_file: + mlc_chat_config = json.load(in_file) + model_type = mlc_chat_config.pop("model_type") + quantization = mlc_chat_config.pop("quantization") + lib_suffix = MLC_DSO_SUFFIX if device not in ["iphone", "macabi", "android"] else "tar" + + def _get_optimization_flags() -> str: + opt = overrides.pop("opt", None) + if opt is None: + opt = "O2" + return repr(OptimizationFlags.from_str(opt)) + + def _get_overrides() -> str: + forbid_list = [ + "context_window_size", + "sliding_window_size", + "attention_sink_size", + ] + result = [] + for field in dataclasses.fields(ModelConfigOverride): + value = overrides.get(field.name, None) + if value is not None: + if field.name in forbid_list and value == -1: + continue + result.append(f"{field.name}={value}") + return ";".join(result) + + def _get_model_config() -> Dict[str, Any]: # noqa: UP006 + model_config = mlc_chat_config.pop("model_config") + model_config.update(mlc_chat_config) + for field in dataclasses.fields(ModelConfigOverride): + value = overrides.get(field.name, None) + if value is not None: + model_config[field.name] = value + return MODELS[model_type].config.from_dict(model_config).asdict() + + def _run_jit( + opt: str, + overrides: str, + device: str, + system_lib_prefix: Optional[str], + dst: str, + ): + with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir: + dso_path = os.path.join(tmp_dir, f"lib.{lib_suffix}") + cmd = [ + sys.executable, + "-m", + "mlc_llm", + "compile", + str(model_path), + "--opt", + opt, + "--overrides", + overrides, + "--device", + device, + "--output", + dso_path, + ] + if system_lib_prefix: + cmd += ["--system-lib-prefix", system_lib_prefix + "_"] + logger.info("Compiling using commands below:") + logger.info("%s", blue(shlex.join(cmd))) + subprocess.run(cmd, check=False, env=os.environ) + # note on windows: compilation can succeed but return code is still nonzero + # check whether file exists instead + if not os.path.isfile(dso_path): + raise RuntimeError("Cannot find compilation output, compilation failed") + shutil.move(dso_path, dst) + logger.info("Using compiled model lib: %s", bold(dst)) + + hash_key = { + "model_config": _get_model_config(), + "overrides": _get_overrides(), + "opt": _get_optimization_flags(), + "device": device2str(device) if isinstance(device, Device) else device, + "model_type": model_type, + "quantization": quantization, + } + if device in ["iphone", "macabi", "android"]: + if system_lib_prefix is None: + system_lib_hash_value = hashlib.md5( + json.dumps( + hash_key, + sort_keys=True, + indent=2, + ).encode("utf-8") + ).hexdigest() + system_lib_prefix = f"{model_type}_{quantization}_{system_lib_hash_value}".replace( + "-", "_" + ) + hash_key["system_lib_prefix"] = system_lib_prefix + hash_value = hashlib.md5( + json.dumps( + hash_key, + sort_keys=True, + indent=2, + ).encode("utf-8") + ).hexdigest() + dst = MLC_LLM_HOME / "model_lib" / f"{hash_value}.{lib_suffix}" + if dst.is_file() and MLC_JIT_POLICY in ["ON", "READONLY"]: + logger.info("Using cached model lib: %s", bold(str(dst))) + return JITResult(str(dst), system_lib_prefix) + if MLC_JIT_POLICY == "READONLY": + raise RuntimeError( + "No cached model lib found, and JIT is disabled by MLC_JIT_POLICY=READONLY" + ) + _run_jit( + opt=hash_key["opt"], + overrides=hash_key["overrides"], + device=hash_key["device"], + system_lib_prefix=system_lib_prefix, + dst=str(dst), + ) + return JITResult(str(dst), system_lib_prefix) diff --git a/python/mlc_llm/interface/package.py b/python/mlc_llm/interface/package.py new file mode 100644 index 0000000..31909bd --- /dev/null +++ b/python/mlc_llm/interface/package.py @@ -0,0 +1,402 @@ +"""Python entrypoint of package.""" + +import dataclasses +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict, List, Literal # noqa: UP035 + +from mlc_llm.interface import jit +from mlc_llm.support import download_cache, logging, style + +logging.enable_logging() +logger = logging.getLogger(__name__) + +SUPPORTED_DEVICES = ["iphone", "macabi", "android"] + + +def build_model_library( + package_config: Dict[str, Any], # noqa: UP006 + device: str, + bundle_dir: Path, + app_config_path: Path, +) -> Dict[str, str]: # noqa: UP006 + """Build model libraries. Return the dictionary of "library prefix to lib path".""" + # - Create the bundle directory. + os.makedirs(bundle_dir, exist_ok=True) + # Clean up all the directories in `output/bundle`. + logger.info('Clean up all directories under "%s"', str(bundle_dir)) + for content_path in bundle_dir.iterdir(): + if content_path.is_dir(): + shutil.rmtree(content_path) + + # - Process each model, and prepare the app config. + app_config_model_list = [] + + model_entries = package_config.get("model_list", []) + if not isinstance(model_entries, list): + raise ValueError('The "model_list" in "mlc-package-config.json" is expected to be a list.') + model_lib_path_for_prepare_libs = package_config.get("model_lib_path_for_prepare_libs", {}) + if not isinstance(model_lib_path_for_prepare_libs, dict): + raise ValueError( + 'The "model_lib_path_for_prepare_libs" in "mlc-package-config.json" is expected to be ' + "a dict." + ) + + jit.log_jit_policy() + + for model_entry in package_config.get("model_list", []): + # - Parse model entry. + if not isinstance(model_entry, dict): + raise ValueError('The element of "model_list" is expected to be a dict.') + model = model_entry["model"] + model_id = model_entry["model_id"] + bundle_weight = model_entry.get("bundle_weight", False) + overrides = model_entry.get("overrides", {}) + model_lib = model_entry.get("model_lib", None) + + estimated_vram_bytes = model_entry["estimated_vram_bytes"] + if not isinstance(model, str): + raise ValueError('The value of "model" in "model_list" is expected to be a string.') + if not isinstance(model_id, str): + raise ValueError('The value of "model_id" in "model_list" is expected to be a string.') + if not isinstance(bundle_weight, bool): + raise ValueError( + 'The value of "bundle_weight" in "model_list" is expected to be a boolean.' + ) + if not isinstance(overrides, dict): + raise ValueError('The value of "overrides" in "model_list" is expected to be a dict.') + if model_lib is not None and not isinstance(model_lib, str): + raise ValueError('The value of "model_lib" in "model_list" is expected to be string.') + + # - Load model config. Download happens when needed. + model_path = download_cache.get_or_download_model(model) + + # - Jit compile if the model lib path is not specified. + model_lib_path = ( + model_lib_path_for_prepare_libs.get(model_lib, None) if model_lib is not None else None + ) + if model_lib_path is None: + if model_lib is None: + logger.info( + 'Model lib is not specified for model "%s". Now jit compile the model library.', + model_id, + ) + else: + logger.info( + 'Model lib path for "%s" is not specified in "model_lib_path_for_prepare_libs".' + "Now jit compile the model library.", + model_lib, + ) + model_lib_path, model_lib = dataclasses.astuple( + jit.jit( + model_path=model_path, + overrides=overrides, + device=device, + system_lib_prefix=model_lib, + skip_log_jit_policy=True, + ) + ) + assert model_lib is not None + model_lib_path_for_prepare_libs[model_lib] = model_lib_path + + # - Set "model_url"/"model_path" and "model_id" + app_config_model_entry = {} + is_local_model = not model.startswith("HF://") and not model.startswith("https://") + app_config_model_entry["model_id"] = model_id + app_config_model_entry["model_lib"] = model_lib + + # - Bundle weight + if is_local_model and not bundle_weight: + raise ValueError( + f'Model "{model}" in "model_list" is a local path.' + f'Please set \'"bundle_weight": true\' in the entry of model "{model}".' + ) + if bundle_weight: + if not os.path.isfile(model_path / "tensor-cache.json"): + raise ValueError( + f'Bundle weight is set for model "{model}". However, model weights are not' + f'found under the directory "{model}". ' + + ( + "Please follow https://llm.mlc.ai/docs/compilation/convert_weights.html to " + "convert model weights." + if is_local_model + else "Please report this issue to https://github.com/mlc-ai/mlc-llm/issues." + ) + ) + # Overwrite the model weight directory in bundle. + bundle_model_weight_path = bundle_dir / model_id + logger.info( + "Bundle weight for %s, copy into %s", + style.bold(model_id), + style.bold(str(bundle_model_weight_path)), + ) + if bundle_model_weight_path.exists(): + shutil.rmtree(bundle_model_weight_path) + shutil.copytree(model_path, bundle_model_weight_path) + if bundle_weight and device in ["iphone", "macabi"]: + app_config_model_entry["model_path"] = model_id + else: + app_config_model_entry["model_url"] = model.replace("HF://", "https://huggingface.co/") + + # - estimated_vram_bytes + app_config_model_entry["estimated_vram_bytes"] = estimated_vram_bytes + + app_config_model_list.append(app_config_model_entry) + + # - Dump "mlc-app-config.json". + app_config_json_str = json.dumps( + {"model_list": app_config_model_list}, + indent=2, + ) + with open(app_config_path, "w", encoding="utf-8") as file: + print(app_config_json_str, file=file) + logger.info( + 'Dump the app config below to "%s":\n%s', + str(app_config_path), + style.green(app_config_json_str), + ) + return model_lib_path_for_prepare_libs + + +def validate_model_lib( + app_config_path: Path, + package_config_path: Path, + model_lib_path_for_prepare_libs: dict, + device: Literal["iphone", "macabi", "android"], + output: Path, +) -> None: + """Validate the model lib prefixes of model libraries.""" + if device == "android": + from tvm.support import ndk as cc + else: + from tvm.support import cc + + with open(app_config_path, encoding="utf-8") as file: + app_config = json.load(file) + + tar_list = [] + model_set = set() + + for model, model_lib_path in model_lib_path_for_prepare_libs.items(): + model_lib_path = os.path.join(model_lib_path) + lib_path_valid = os.path.isfile(model_lib_path) + if not lib_path_valid: + raise RuntimeError(f"Cannot find file {model_lib_path} as an {device} model library") + tar_list.append(model_lib_path) + model_set.add(model) + + os.makedirs(output / "lib", exist_ok=True) + if device in ["iphone", "macabi"]: + lib_name = "libmodel_iphone.a" + else: + lib_name = "libmodel_android.a" + lib_path = output / "lib" / lib_name + + def _get_model_libs(lib_path: Path) -> List[str]: # noqa: UP006 + """Get the model lib prefixes in the given static lib path.""" + global_symbol_map = cc.get_global_symbol_section_map(lib_path) + libs = [] + suffix = "___tvm_ffi__library_bin" + for name, _ in global_symbol_map.items(): + if name.endswith(suffix): + model_lib = name[: -len(suffix)] + if model_lib.startswith("_"): + model_lib = model_lib[1:] + libs.append(model_lib) + return libs + + cc.create_staticlib(lib_path, tar_list) + available_model_libs = _get_model_libs(lib_path) + logger.info("Creating lib from %s", str(tar_list)) + logger.info("Validating the library %s", str(lib_path)) + logger.info( + "List of available model libs packaged: %s," + " if we have '-' in the model_lib string, it will be turned into '_'", + str(available_model_libs), + ) + global_symbol_map = cc.get_global_symbol_section_map(lib_path) + error_happened = False + + for item in app_config["model_list"]: + model_lib = item["model_lib"] + model_id = item["model_id"] + if model_lib not in model_set: + # NOTE: this cannot happen under new setting + # since if model_lib is not included, it will be jitted + raise RuntimeError( + f"ValidationError: model_lib={model_lib} specified for model_id={model_id} " + "is not included in model_lib_path_for_prepare_libs argument, " + "This will cause the specific model not being able to load, " + f"model_lib_path_for_prepare_libs={model_lib_path_for_prepare_libs}" + ) + + model_prefix_pattern = model_lib.replace("-", "_") + "___tvm_ffi__library_bin" + if ( + model_prefix_pattern not in global_symbol_map + and "_" + model_prefix_pattern not in global_symbol_map + ): + # NOTE: no lazy format is ok since this is a slow pass + model_lib_path = model_lib_path_for_prepare_libs[model_lib] + log_msg = ( + "ValidationError:\n" + f"\tmodel_lib {model_lib} requested in {str(app_config_path)}" + f" is not found in {str(lib_path)}\n" + f"\tspecifically the model_lib for {model_lib_path}.\n" + f"\tcurrent available model_libs in {str(lib_path)}: {available_model_libs}\n" + f"\tThis can happen when we manually specified model_lib_path_for_prepare_libs" + f" in {str(package_config_path)}\n" + f"\tConsider remove model_lib_path_for_prepare_libs (so library can be jitted)" + "or check the compile command" + ) + logger.info(log_msg) + error_happened = True + + if not error_happened: + logger.info(style.green("Validation pass")) + else: + logger.info(style.red("Validation failed")) + sys.exit(255) + + +def build_android_binding(mlc_llm_source_dir: Path, output: Path) -> None: + """Build android binding in MLC LLM""" + mlc4j_path = mlc_llm_source_dir / "android" / "mlc4j" + + # Move the model libraries to "build/lib/" for linking + os.makedirs(Path("build") / "lib", exist_ok=True) + src_path = str(output / "lib" / "libmodel_android.a") + dst_path = str(Path("build") / "lib" / "libmodel_android.a") + logger.info('Moving "%s" to "%s"', src_path, dst_path) + shutil.move(src_path, dst_path) + + # Build mlc4j + logger.info("Building mlc4j") + subprocess.run([sys.executable, mlc4j_path / "prepare_libs.py"], check=True, env=os.environ) + # Copy built files back to output directory. + lib_path = output / "lib" / "mlc4j" + os.makedirs(lib_path, exist_ok=True) + logger.info('Clean up all directories under "%s"', str(lib_path)) + for content_path in lib_path.iterdir(): + if content_path.is_dir(): + shutil.rmtree(content_path) + + src_path = str(mlc4j_path / "src") + dst_path = str(lib_path / "src") + logger.info('Copying "%s" to "%s"', src_path, dst_path) + shutil.copytree(src_path, dst_path) + + src_path = str(mlc4j_path / "build.gradle") + dst_path = str(lib_path / "build.gradle") + logger.info('Copying "%s" to "%s"', src_path, dst_path) + shutil.copy(src_path, dst_path) + + src_path = str(Path("build") / "output") + dst_path = str(lib_path / "output") + logger.info('Copying "%s" to "%s"', src_path, dst_path) + shutil.copytree(src_path, dst_path) + + os.makedirs(lib_path / "src" / "main" / "assets") + src_path = str(output / "bundle" / "mlc-app-config.json") + dst_path = str(lib_path / "src" / "main" / "assets" / "mlc-app-config.json") + logger.info('Moving "%s" to "%s"', src_path, dst_path) + shutil.move(src_path, dst_path) + + +def build_iphone_binding(mlc_llm_source_dir: Path, output: Path) -> None: + """Build iOS binding in MLC LLM""" + # Build iphone binding + logger.info("Build iphone binding") + subprocess.run( + ["bash", mlc_llm_source_dir / "ios" / "prepare_libs.sh"], + check=True, + env=os.environ, + ) + + # Copy built libraries back to output directory. + for static_library in (Path("build") / "lib").iterdir(): + dst_path = str(output / "lib" / static_library.name) + logger.info('Copying "%s" to "%s"', static_library, dst_path) + shutil.copy(static_library, dst_path) + + +def build_macabi_binding(mlc_llm_source_dir: Path, output: Path) -> None: + """Build Mac Catalyst binding in MLC LLM""" + deployment_target = os.environ.get("MLC_MACABI_DEPLOYMENT_TARGET", "18.0") + macabi_arch = os.environ.get("MLC_MACABI_ARCH", "").strip() or "arm64" + logger.info("Build macabi binding (deployment target %s)", deployment_target) + cmd = [ + "bash", + str(mlc_llm_source_dir / "ios" / "prepare_libs.sh"), + "--catalyst", + "--deployment-target", + deployment_target, + ] + if macabi_arch: + cmd += ["--arch", macabi_arch] + subprocess.run(cmd, check=True, env=os.environ) + + # Copy built libraries back to output directory. + build_dir = Path(f"build-maccatalyst-{macabi_arch}") + for static_library in (build_dir / "lib").iterdir(): + dst_path = str(output / "lib" / static_library.name) + logger.info('Copying "%s" to "%s"', static_library, dst_path) + shutil.copy(static_library, dst_path) + + +def package( + package_config_path: Path, + mlc_llm_source_dir: Path, + output: Path, +) -> None: + """Python entrypoint of package.""" + logger.info('MLC LLM HOME: "%s"', mlc_llm_source_dir) + + # - Read package config. + with open(package_config_path, encoding="utf-8") as file: + package_config = json.load(file) + if not isinstance(package_config, dict): + raise ValueError( + "The content of MLC package config is expected to be a dict with " + f'field "model_list". However, the content of "{package_config_path}" is not a dict.' + ) + + # - Read device. + if "device" not in package_config: + raise ValueError(f'JSON file "{package_config_path}" is required to have field "device".') + device = package_config["device"] + if device not in SUPPORTED_DEVICES: + raise ValueError( + f'The "device" field of JSON file {package_config_path} is expected to be one of ' + f'{SUPPORTED_DEVICES}, while "{device}" is given in the JSON.' + ) + + bundle_dir = output / "bundle" + app_config_path = bundle_dir / "mlc-app-config.json" + # - Build model libraries. + model_lib_path_for_prepare_libs = build_model_library( + package_config, device, bundle_dir, app_config_path + ) + # - Validate model libraries. + validate_model_lib( + app_config_path, + package_config_path, + model_lib_path_for_prepare_libs, + device, + output, + ) + + # - Copy model libraries + if device == "android": + build_android_binding(mlc_llm_source_dir, output) + elif device == "iphone": + build_iphone_binding(mlc_llm_source_dir, output) + elif device == "macabi": + build_macabi_binding(mlc_llm_source_dir, output) + else: + assert False, "Cannot reach here" + + logger.info("All finished.") diff --git a/python/mlc_llm/interface/router.py b/python/mlc_llm/interface/router.py new file mode 100644 index 0000000..1df3efe --- /dev/null +++ b/python/mlc_llm/interface/router.py @@ -0,0 +1,125 @@ +"""Python entrypoint of router.""" + +from collections.abc import AsyncGenerator +from http import HTTPStatus +from typing import List, Literal, Optional, Type # noqa: UP035 + +import fastapi +import uvicorn +from fastapi.middleware.cors import CORSMiddleware + +from mlc_llm.protocol import error_protocol +from mlc_llm.protocol.openai_api_protocol import CompletionLogProbs, CompletionRequest +from mlc_llm.router import Router +from mlc_llm.serve import engine_base, engine_utils + + +def serve( + model: str, + model_lib: Optional[str], + router_host: str, + router_port: int, + endpoint_hosts: List[str], # noqa: UP006 + endpoint_ports: List[int], # noqa: UP006 + endpoint_num_gpus: List[int], # noqa: UP006 + enable_prefix_cache: bool, + router_mode: Literal["disagg", "round-robin"] = "round-robin", + pd_balance_factor: float = 0.0, + router_type: Type[Router] = Router, # noqa: UP006 +): + """Start the router with the specified configuration.""" + # 1. Instantiate router + router = router_type( + model=model, + model_lib=model_lib, + hosts=endpoint_hosts, + ports=endpoint_ports, + num_gpus=endpoint_num_gpus, + enable_prefix_cache=enable_prefix_cache, + router_mode=router_mode, + pd_balance_factor=pd_balance_factor, + ) + + router_app = fastapi.APIRouter() + + @router_app.post("/v1/completions") + async def request_completion(request: CompletionRequest, raw_request: fastapi.Request): + """OpenAI-compatible completion API. + API reference: https://platform.openai.com/docs/api-reference/completions/create + """ + if router is None: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message="Router is not initialized." + ) + request_id = f"cmpl-{engine_utils.random_uuid()}" + + # Streaming response. + if request.stream: + # We manually get the first response from generator to + # capture potential exceptions in this scope, rather then + # the StreamingResponse scope. + stream_generator = router.handle_completion(request, request_id) + first_response = await anext( # noqa: F821 + stream_generator + ) + + async def completion_stream_generator() -> AsyncGenerator[str, None]: + if isinstance(first_response, StopAsyncIteration): + yield "data: [DONE]\n\n" + return + yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n" + async for response in stream_generator: + yield f"data: {response.model_dump_json(by_alias=True)}\n\n" + yield "data: [DONE]\n\n" + + return fastapi.responses.StreamingResponse( + completion_stream_generator(), media_type="text/event-stream" + ) + + # FIXME: Non-streaming response not fully implemented + request_final_usage = None + output_texts = [""] * request.n + finish_reasons: List[Optional[str]] = [None] * request.n # noqa: UP006 + logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n # noqa: UP006 + + async for response in router.handle_completion(request, request_id): + if await raw_request.is_disconnected(): + # In non-streaming cases, the engine will not be notified + # when the request is disconnected. + # Therefore, we check if it is disconnected each time, + # and explicitly return. + # Note that requesta abort is triggered when the async for and funciton scope ends. + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message="The request has disconnected" + ) + # TODO(Charlie): This is copied from engine.py -- + # why is it here? Non-streaming only has a single chunk right? + # this is the final chunk + # if response.usage is not None: + # request_final_usage = response.usage + # continue + for choice in response.choices: + output_texts[choice.index] += choice.text + if choice.finish_reason is not None and finish_reasons[choice.index] is None: + finish_reasons[choice.index] = choice.finish_reason + if choice.logprobs is not None: + logprob_results[choice.index] = choice.logprobs + + assert all(finish_reason is not None for finish_reason in finish_reasons) + return engine_base.wrap_completion_response( + request_id=request_id, + model=request.model, + output_texts=output_texts, + finish_reasons=finish_reasons, + logprob_results=logprob_results, + usage=request_final_usage, + ) + + # 2. Set up app + app = fastapi.FastAPI() + app.add_middleware(CORSMiddleware) + app.include_router(router_app) + app.exception_handler(error_protocol.BadRequestError)(error_protocol.bad_request_error_handler) + + # 3. Run + uvicorn.run(app, host=router_host, port=router_port, log_level="info") diff --git a/python/mlc_llm/interface/serve.py b/python/mlc_llm/interface/serve.py new file mode 100644 index 0000000..f0bed06 --- /dev/null +++ b/python/mlc_llm/interface/serve.py @@ -0,0 +1,131 @@ +"""Python entrypoint of serve.""" + +from typing import Any, List, Literal, Optional, Tuple, Union # noqa: UP035 + +import fastapi +import uvicorn +from fastapi.middleware.cors import CORSMiddleware + +from mlc_llm.protocol import error_protocol +from mlc_llm.serve import engine +from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine +from mlc_llm.serve.entrypoints import ( + debug_entrypoints, + metrics_entrypoints, + microserving_entrypoints, + openai_entrypoints, +) +from mlc_llm.serve.server import ServerContext +from mlc_llm.support import logging + +logger = logging.getLogger(__name__) + + +def serve( + model: str, + device: str, + model_lib: Optional[str], + mode: Literal["local", "interactive", "server"], + enable_debug: bool, + additional_models: List[Union[str, Tuple[str, str]]], # noqa: UP006 + embedding_model: Optional[str], + embedding_model_lib: Optional[str], + tensor_parallel_shards: Optional[int], + pipeline_parallel_stages: Optional[int], + opt: Optional[str], + max_num_sequence: Optional[int], + max_total_sequence_length: Optional[int], + max_single_sequence_length: Optional[int], + prefill_chunk_size: Optional[int], + sliding_window_size: Optional[int], + attention_sink_size: Optional[int], + max_history_size: Optional[int], + gpu_memory_utilization: Optional[float], + speculative_mode: Literal["disable", "small_draft", "eagle", "medusa"], + spec_draft_length: Optional[int], + spec_tree_width: Optional[int], + prefix_cache_mode: Literal["disable", "radix"], + prefix_cache_max_num_recycling_seqs: Optional[int], + prefill_mode: Literal["hybrid", "chunked"], + enable_tracing: bool, + host: str, + port: int, + allow_credentials: bool, + allow_origins: Any, + allow_methods: Any, + allow_headers: Any, + api_key: Optional[str] = None, +): + """Serve the model with the specified configuration.""" + # Create engine and start the background loop + async_engine = engine.AsyncMLCEngine( + model=model, + device=device, + model_lib=model_lib, + mode=mode, + engine_config=engine.EngineConfig( + additional_models=additional_models, + tensor_parallel_shards=tensor_parallel_shards, + pipeline_parallel_stages=pipeline_parallel_stages, + opt=opt, + max_num_sequence=max_num_sequence, + max_total_sequence_length=max_total_sequence_length, + max_single_sequence_length=max_single_sequence_length, + prefill_chunk_size=prefill_chunk_size, + sliding_window_size=sliding_window_size, + attention_sink_size=attention_sink_size, + max_history_size=max_history_size, + gpu_memory_utilization=gpu_memory_utilization, + speculative_mode=speculative_mode, + spec_draft_length=spec_draft_length, + spec_tree_width=spec_tree_width, + prefix_cache_mode=prefix_cache_mode, + prefix_cache_max_num_recycling_seqs=prefix_cache_max_num_recycling_seqs, + prefill_mode=prefill_mode, + ), + enable_tracing=enable_tracing, + ) + + # Set up embedding model if specified + emb_engine = None + if embedding_model is not None: + if embedding_model_lib is None: + raise ValueError( + "--embedding-model-lib is required when --embedding-model is specified." + ) + emb_engine = AsyncEmbeddingEngine( + model=embedding_model, + model_lib=embedding_model_lib, + device=device, + ) + logger.info("Embedding model %s loaded successfully.", embedding_model) + + with ServerContext() as server_context: + server_context.add_model(model, async_engine) + if emb_engine is not None: + server_context.add_embedding_engine(embedding_model, emb_engine) + server_context.api_key = api_key + + app = fastapi.FastAPI() + app.add_middleware( + CORSMiddleware, + allow_credentials=allow_credentials, + allow_origins=allow_origins, + allow_methods=allow_methods, + allow_headers=allow_headers, + ) + + app.include_router(openai_entrypoints.app) + app.include_router(metrics_entrypoints.app) + app.include_router(microserving_entrypoints.app) + + server_context.enable_debug = enable_debug + + if enable_debug: + app.include_router(debug_entrypoints.app) + logger.info("Enable debug endpoint and debug_config in requests...") + + app.exception_handler(error_protocol.BadRequestError)( + error_protocol.bad_request_error_handler + ) + uvicorn.run(app, host=host, port=port, log_level="info") diff --git a/python/mlc_llm/json_ffi/__init__.py b/python/mlc_llm/json_ffi/__init__.py new file mode 100644 index 0000000..8a70591 --- /dev/null +++ b/python/mlc_llm/json_ffi/__init__.py @@ -0,0 +1,8 @@ +"""JSON FFI is a pure string based interface of MLC LLM Engine. + +We build interfacing with JSON FFI for both testing purposes +and internal use. For most python API usage, please use MLCEngine +and MLCAsyncEngine +""" + +from .engine import JSONFFIEngine diff --git a/python/mlc_llm/json_ffi/engine.py b/python/mlc_llm/json_ffi/engine.py new file mode 100644 index 0000000..0b387ce --- /dev/null +++ b/python/mlc_llm/json_ffi/engine.py @@ -0,0 +1,295 @@ +import json +import queue +import threading +from collections.abc import Iterator +from typing import Any, Callable, Dict, List, Literal, Optional, Union # noqa: UP035 + +import tvm + +from mlc_llm.protocol import debug_protocol, openai_api_protocol +from mlc_llm.serve import engine_utils +from mlc_llm.serve.engine_base import ( + EngineConfig, + EngineMetrics, + _check_engine_config, + _parse_models, + _process_model_args, + _query_engine_metrics, + detect_device, +) +from mlc_llm.tokenizers import Tokenizer + + +class EngineState: + sync_queue: queue.Queue + + def get_request_stream_callback(self) -> Callable[[str], None]: + # ChatCompletionStreamResponse + + def _callback(chat_completion_stream_responses_json_str: str) -> None: + self._sync_request_stream_callback(chat_completion_stream_responses_json_str) + + return _callback + + def _sync_request_stream_callback(self, chat_completion_stream_responses_json_str: str) -> None: + # Put the delta outputs to the queue in the unblocking way. + self.sync_queue.put_nowait(chat_completion_stream_responses_json_str) + + def handle_chat_completion( + self, ffi: dict, request_json_str: str, include_usage: bool, request_id: str + ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]: + """Helper class to handle chat completion + + Note + ---- + ffi is explicitly passed in to avoid cylic dependency + as ffi will capture EngineState + """ + self.sync_queue = queue.Queue() + + ffi["chat_completion"](request_json_str, request_id) + + try: + last_chunk_arrived = False + while not last_chunk_arrived: + chat_completion_responses_json_str = self.sync_queue.get() + chat_completion_responses_list = json.loads(chat_completion_responses_json_str) + for chat_completion_response_json_dict in chat_completion_responses_list: + chat_completion_response = ( + openai_api_protocol.ChatCompletionStreamResponse.model_validate( + chat_completion_response_json_dict + ) + ) + # the chunk with usage is always the last chunk + if chat_completion_response.usage is not None: + if include_usage: + yield chat_completion_response + last_chunk_arrived = True + break + yield chat_completion_response + except Exception as exception: + ffi["abort"](request_id) + raise exception + + +class BackgroundLoops: + """Helper class to keep track of background loops""" + + def __init__(self, ffi: dict): + self._ffi = ffi + # important: avoid self reference in closure + background_loop = self._ffi["run_background_loop"] + background_stream_back_loop = self._ffi["run_background_stream_back_loop"] + + # Create the background engine-driving thread and start the loop. + self._background_loop_thread: threading.Thread = threading.Thread(target=background_loop) + self._background_stream_back_loop_thread: threading.Thread = threading.Thread( + target=background_stream_back_loop + ) + self._background_loop_thread.start() + self._background_stream_back_loop_thread.start() + self._terminated = False + + def __del__(self): + self.terminate() + + def terminate(self): + if self._terminated: + return + self._terminated = True + self._ffi["exit_background_loop"]() + self._background_loop_thread.join() + self._background_stream_back_loop_thread.join() + + +class Completions: + """Completions class to be compatible with OpenAI API""" + + _ffi: dict + _state: EngineState + _background_loops: BackgroundLoops + + def __init__(self, ffi: dict, state: EngineState, background_loops: BackgroundLoops): + self._ffi = ffi + self._state = state + self._background_loops = background_loops + + def create( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = True, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]: + if request_id is None: + request_id = f"chatcmpl-{engine_utils.random_uuid()}" + debug_config = extra_body.get("debug_config", None) if extra_body is not None else None + if not stream: + raise ValueError("JSONFFIEngine only support stream=True") + request = openai_api_protocol.ChatCompletionRequest( + messages=[ + openai_api_protocol.ChatCompletionMessage.model_validate(message) + for message in messages + ], + model=model, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + top_logprobs=top_logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + temperature=temperature, + top_p=top_p, + tools=( + [openai_api_protocol.ChatTool.model_validate(tool) for tool in tools] + if tools is not None + else None + ), + tool_choice=tool_choice, + user=user, + response_format=( + openai_api_protocol.RequestResponseFormat.model_validate(response_format) + if response_format is not None + else None + ), + debug_config=( + debug_protocol.DebugConfig.model_validate(debug_config) + if debug_config is not None + else None + ), + ) + chatcmpl_generator = self._state.handle_chat_completion( + self._ffi, + request.model_dump_json(by_alias=True), + include_usage=( + request.stream_options is not None and request.stream_options.include_usage + ), + request_id=request_id, + ) + for response in chatcmpl_generator: + yield response + + +class Chat: + """Chat class to be compatible with OpenAI API""" + + completions: Completions + + def __init__(self, ffi: dict, state: EngineState, background_loops: BackgroundLoops): + self.completions = Completions(ffi, state, background_loops) + + +class JSONFFIEngine: + chat: Chat + + def __init__( + self, + model: str, + device: Union[str, tvm.runtime.Device] = "auto", + *, + model_lib: Optional[str] = None, + mode: Literal["local", "interactive", "server"] = "local", + engine_config: Optional[EngineConfig] = None, + ) -> None: + # - Check the fields fields of `engine_config`. + if engine_config is None: + engine_config = EngineConfig() + _check_engine_config(model, model_lib, mode, engine_config) + + # - Initialize model loading info. + models = _parse_models(model, model_lib, engine_config.additional_models) + if isinstance(device, str): + device = detect_device(device) + assert isinstance(device, tvm.runtime.Device) + model_args = _process_model_args(models, device, engine_config)[0] + + # - Load the raw model config into dict + for i, model_info in enumerate(models): + model_info.model_lib = model_args[i][1] + + # - Initialize engine state and engine. + self._state = EngineState() + module = tvm.get_global_func("mlc.json_ffi.CreateJSONFFIEngine", allow_missing=False)() + self._ffi = { + key: module[key] + for key in [ + "init_background_engine", + "reload", + "unload", + "reset", + "chat_completion", + "abort", + "run_background_loop", + "run_background_stream_back_loop", + "exit_background_loop", + ] + } + self.tokenizer = Tokenizer(model_args[0][0]) + self._background_loops = BackgroundLoops(self._ffi) + + engine_config.model = model_args[0][0] + engine_config.model_lib = model_args[0][1] + engine_config.additional_models = model_args[1:] + engine_config.mode = mode + self.engine_config = engine_config + + self._ffi["init_background_engine"]( + device.dlpack_device_type(), + device.index, + self._state.get_request_stream_callback(), + ) + self._ffi["reload"](self.engine_config.asjson()) + + self.chat = Chat(self._ffi, self._state, self._background_loops) + + def metrics(self) -> EngineMetrics: + """Get the engine metrics.""" + return _query_engine_metrics(self) + + def _raw_chat_completion( + self, request_json_str: str, include_usage: bool, request_id: str + ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]: + """Raw chat completion API""" + return self._state.handle_chat_completion( + self._ffi, request_json_str, include_usage, request_id + ) + + def terminate(self): + """Explicitly terminate the engine""" + self._background_loops.terminate() + + def _test_reload(self): + self._ffi["reload"](self.engine_config.asjson()) + + def _test_reset(self): + self._ffi["reset"]() + + def _test_unload(self): + self._ffi["unload"]() diff --git a/python/mlc_llm/libinfo.py b/python/mlc_llm/libinfo.py new file mode 100644 index 0000000..2212d8c --- /dev/null +++ b/python/mlc_llm/libinfo.py @@ -0,0 +1,71 @@ +"""Library information. This is a standalone file that can be used to get various info""" + +#! pylint: disable=protected-access +import os +import sys + +__version__ = "0.1.dev0" +MLC_LIBRARY_PATH = os.environ.get("MLC_LIBRARY_PATH", None) + + +def get_env_paths(env_var, splitter): + """Get path in env variable""" + if os.environ.get(env_var, None): + return [p.strip() for p in os.environ[env_var].split(splitter)] + return [] + + +def get_dll_directories(): + """Get extra mlc llm dll directories""" + curr_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__))) + source_dir = os.path.abspath(os.path.join(curr_dir, "..", "..")) + dll_path = [ + curr_dir, + os.path.join(source_dir, "build"), + os.path.join(source_dir, "build", "Release"), + ] + if MLC_LIBRARY_PATH: + dll_path.append(MLC_LIBRARY_PATH) + if "CONDA_PREFIX" in os.environ: + dll_path.append(os.path.join(os.environ["CONDA_PREFIX"], "lib")) + if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"): + dll_path.extend(get_env_paths("LD_LIBRARY_PATH", ":")) + elif sys.platform.startswith("darwin"): + dll_path.extend(get_env_paths("DYLD_LIBRARY_PATH", ":")) + elif sys.platform.startswith("win32"): + dll_path.extend(get_env_paths("PATH", ";")) + return [os.path.abspath(p) for p in dll_path if os.path.isdir(p)] + + +def find_lib_path(name, optional=False): + """Find mlc llm library + + Parameters + ---------- + name : str + The name of the library + + optional: boolean + Whether the library is required + """ + if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"): + lib_name = f"lib{name}.so" + elif sys.platform.startswith("win32"): + lib_name = f"{name}.dll" + elif sys.platform.startswith("darwin"): + lib_name = f"lib{name}.dylib" + else: + lib_name = f"lib{name}.so" + + dll_paths = get_dll_directories() + lib_dll_path = [os.path.join(p, lib_name) for p in dll_paths] + lib_found = [p for p in lib_dll_path if os.path.exists(p) and os.path.isfile(p)] + if not lib_found: + if not optional: + message = ( + f"Cannot find libraries: {lib_name}\n" + + "List of candidates:\n" + + "\n".join(lib_dll_path) + ) + raise RuntimeError(message) + return lib_found diff --git a/python/mlc_llm/loader/__init__.py b/python/mlc_llm/loader/__init__.py new file mode 100644 index 0000000..3cee0bf --- /dev/null +++ b/python/mlc_llm/loader/__init__.py @@ -0,0 +1,8 @@ +""" +A subpackage of the compiler that represents mapping between external parameters, quantized +parameters and parameters in MLC-defined models. +""" + +from .huggingface_loader import HuggingFaceLoader +from .loader import LOADER, Loader +from .mapping import ExternMapping, QuantizeMapping diff --git a/python/mlc_llm/loader/huggingface_loader.py b/python/mlc_llm/loader/huggingface_loader.py new file mode 100644 index 0000000..6350635 --- /dev/null +++ b/python/mlc_llm/loader/huggingface_loader.py @@ -0,0 +1,228 @@ +"""A weight loader for HuggingFace's PyTorch format""" + +import gc +import json +from collections import OrderedDict, defaultdict +from collections.abc import Iterator +from pathlib import Path +from typing import Callable, Dict, List, Optional, Tuple # noqa: UP035 + +import numpy as np +from tqdm import tqdm +from tvm.runtime import Device, Tensor +from tvm.runtime import tensor as as_tensor + +from mlc_llm.support import logging +from mlc_llm.support.preshard import _sharded_param_name +from mlc_llm.support.style import bold + +from .mapping import ExternMapping, QuantizeMapping +from .stats import Stats +from .utils import check_parameter_usage, load_safetensor_shard, load_torch_shard + +logger = logging.getLogger(__name__) + + +class HuggingFaceLoader: + """A loader loading HuggingFace's PyTorch/SafeTensor format and converts them + to MLC's parameters. + + Attributes + ---------- + stats : Stats + Statistics of the loading process. + + extern_param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch/SafeTensor. + + torch_to_path : Dict[str, Path] + A mapping from PyTorch/SafeTensor parameter name to the path of the file containing it, + or the path meaning all parameters are stored in a single file. + + cached_files : Dict[Path, Dict[str, np.ndarray]] + A cache of the loaded files. The key is the path of the file, and the value is a mapping + from parameter name to the parameter value. + + quantize_param_map : Optional[QuantizeMapping] + The quantization mapping from MLC to quantized MLC parameters. + """ + + stats: Stats + cached_files: Dict[Path, Dict[str, np.ndarray]] # noqa: UP006 + torch_to_path: Dict[str, Path] # noqa: UP006 + extern_param_map: ExternMapping + quantize_param_map: Optional[QuantizeMapping] + + def __init__( + self, + path: Path, + extern_param_map: ExternMapping, + quantize_param_map: Optional[QuantizeMapping] = None, + ) -> None: + """Create a parameter loader from HuggingFace PyTorch format. + + Parameters + ---------- + path : pathlib.Path + Path to either a JSON indexing file, or a PyTorch bin file. + 1) For JSON indexing file, it is usually `pytorch_model.bin.index.json` + or `model.safetensors.index.json` in the repo, which contains a `weight_map` that + maps each PyTorch parameter to the file containing the weight. + 2) For PyTorch bin file, it is usually `pytorch_model.bin` in the repo, + which contains all the parameters. + 3) For safetensor file, it is usually `model.safetensors` in the repo, + which contains all the parameters. + + extern_param_map : ExternMapping + Maps an MLC parameter to a list of PyTorch/SafeTensor parameters. + + quantize_param_map: Optional[QuantizeMapping] + The quantization mapping from MLC to quantized MLC parameters, default to None, which + means no quantization. + """ + assert path.is_file(), f"Path {path} is not a file" + self.stats = Stats() + self.extern_param_map = extern_param_map + self.cached_files = {} + self.torch_to_path = {} + self.quantize_param_map = quantize_param_map + if path.suffix in (".bin", ".safetensors", ".pt"): + self._load_file(path) + for name in self.cached_files[path].keys(): + self.torch_to_path[name] = path + elif path.suffix == ".json": + with path.open("r", encoding="utf-8") as in_file: + torch_weight_map = json.load(in_file)["weight_map"] + for torch_name, path_str in torch_weight_map.items(): + self.torch_to_path[torch_name] = path.parent / path_str + else: + raise FileNotFoundError(f"Unknown file suffix: {path}") + check_parameter_usage(extern_param_map, set(self.torch_to_path.keys())) + + def load( + self, + device: Device, + preshard_funcs: Optional[Dict[str, Callable]] = None, # noqa: UP006 + ) -> Iterator[Tuple[str, Tensor]]: # noqa: UP006 + """Load the parameters and yield the MLC parameter and its value. + + Parameters + ---------- + device : Optional[Device] + The device to store the parameter, default to None, which means using CPU. + + Yields + ------ + Tuple[str, Tensor] + The MLC parameter name and its value, quantized if quantization mapping is provided. + """ + mlc_names = _loading_order(self.extern_param_map, self.torch_to_path) + for mlc_name in tqdm(mlc_names): + param = self._load_mlc_param(mlc_name, device=device) + # Apply quantization if needed, in this case the original parameter may become + # multiple quantized parameters. + for name, loader_param in self._load_or_quantize(mlc_name, param, device): + # Apply presharding if needed + if preshard_funcs is not None and name in preshard_funcs: + for shard_id, shard_param in enumerate(preshard_funcs[name](loader_param)): + yield _sharded_param_name(name, shard_id), shard_param + else: + yield name, loader_param + + cached_files = list(self.cached_files.keys()) + for path in cached_files: + self._unload_file(path) + self.stats.log_time_info("HF") + self.stats.log_mem_usage() + + def _load_mlc_param(self, mlc_name: str, device: Optional[Device]) -> Tensor: + torch_names = self.extern_param_map.param_map[mlc_name] + files_required = {self.torch_to_path[p] for p in torch_names} + files_existing = set(self.cached_files.keys()) + files_to_load = files_required - files_existing + files_to_unload = files_existing - files_required + + # Step 1. When there is some file to unloaded: + # - If no pending file load: unloading is deferred as there is no gain in peak memory usage; + # - Need to load files: unload immediately to save memory and make space for the new files. + if files_to_load: + for path in files_to_unload: + self._unload_file(path) + # Step 2. Load all the files needed + for path in files_to_load: + self._load_file(path) + # Step 3. Collect all torch parameters in order + torch_params = [self.cached_files[self.torch_to_path[i]][i] for i in torch_names] + # Step 4. Apply the mapping function + with self.stats.timer("map_time_sec"): + param = self.extern_param_map.map_func[mlc_name](*torch_params) + if device: + return as_tensor(param, device=device) + return as_tensor(param) + + def _load_or_quantize(self, mlc_name, param, device: Device): + if self.quantize_param_map and mlc_name in self.quantize_param_map.param_map: + with self.stats.timer("quant_time_sec"): + q_names = self.quantize_param_map.param_map[mlc_name] + q_params = self.quantize_param_map.map_func[mlc_name](param) + device.sync() + for q_name, q_param in zip(q_names, q_params): + logger.info( + '[Quantized] Parameter: "%s", shape: %s, dtype: %s', + bold(q_name), + q_param.shape, + q_param.dtype, + ) + yield q_name, q_param + else: + logger.info( + '[Not quantized] Parameter: "%s", shape: %s, dtype: %s', + bold(mlc_name), + param.shape, + param.dtype, + ) + device.sync() + yield mlc_name, param + + def _load_file(self, path: Path) -> None: + logger.info("Loading HF parameters from: %s", path) + load_func = load_safetensor_shard if path.suffix == ".safetensors" else load_torch_shard + with self.stats.timer("load_time_sec"): + result = {} + for name, param in load_func(path): + result[name] = param + self.stats.mem_add(param.nbytes) + if name not in self.extern_param_map.unused_params: + self.stats.total_param_num += param.size + self.cached_files[path] = result + + def _unload_file(self, path: Path) -> None: + logger.info("Unloading HF weight file: %s", path) + with self.stats.timer("load_time_sec"): + for _, param in self.cached_files[path].items(): + self.stats.mem_rm(param.nbytes) + del self.cached_files[path] + gc.collect() + + +def _loading_order(param_map: ExternMapping, torch_to_path: Dict[str, Path]) -> List[str]: # noqa: UP006 + # Step 1. Build a map from path to torch parameters + path_to_torch: Dict[Path, List[str]] = defaultdict(list) # noqa: UP006 + for torch_name, path in torch_to_path.items(): + path_to_torch[path].append(torch_name) + # Step 2. Build a map from torch parameters to MLC parameters + torch_to_mlc = defaultdict(list) + for mlc_name, torch_names in param_map.param_map.items(): + for torch_name in torch_names: + torch_to_mlc[torch_name].append(mlc_name) + # Step 3. Construct the ordering that ensures file locality + order = OrderedDict() + for _, torch_names in path_to_torch.items(): + for torch_name in torch_names: + for mlc_name in torch_to_mlc[torch_name]: + if mlc_name not in order: + order[mlc_name] = 1 + return list(order.keys()) + + +__all__ = ["HuggingFaceLoader"] diff --git a/python/mlc_llm/loader/loader.py b/python/mlc_llm/loader/loader.py new file mode 100644 index 0000000..3185174 --- /dev/null +++ b/python/mlc_llm/loader/loader.py @@ -0,0 +1,13 @@ +"""A centralized registry of all existing loaders.""" + +from typing import Any, Dict # noqa: UP035 + +from .huggingface_loader import HuggingFaceLoader + +Loader = Any + +LOADER: Dict[str, Any] = { # noqa: UP006 + "huggingface-torch": HuggingFaceLoader, + "huggingface-safetensor": HuggingFaceLoader, + "awq": HuggingFaceLoader, +} diff --git a/python/mlc_llm/loader/mapping.py b/python/mlc_llm/loader/mapping.py new file mode 100644 index 0000000..5842194 --- /dev/null +++ b/python/mlc_llm/loader/mapping.py @@ -0,0 +1,102 @@ +"""Parameter mapping for converting different LLM implementations to MLC LLM.""" + +import dataclasses +from typing import Callable, Dict, List, Set, Union # noqa: UP035 + +import numpy as np +from tvm.runtime import Tensor + +MapFuncVariadic = Union[ + Callable[[], np.ndarray], + Callable[[np.ndarray], np.ndarray], + Callable[[np.ndarray, np.ndarray], np.ndarray], + Callable[[np.ndarray, np.ndarray, np.ndarray], np.ndarray], + Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray], +] + + +@dataclasses.dataclass +class ExternMapping: + """Mapping from a parameter name in MLC LLM's model definition to its potential source, + for example, from MLC parameter "model.layers.2.post_attention_layernorm.weight" to PyTorch's + parameter correspondingly. + + Parameters + ---------- + param_map : Dict[str, List[str]] + A dictionary that maps the name of a parameter to its source. For example, + in Llama2, the source of MLC parameter "model.layers.0.self_attn.qkv_proj.weight" from + huggingface torch are: + + - "model.layers.0.self_attn.q_proj.weight" + - "model.layers.0.self_attn.k_proj.weight" + - "model.layers.0.self_attn.v_proj.weight" + + map_func : Dict[str, Callable[[np.ndarray, ...], np.ndarray]] + A dictionary that maps the name of a parameter to a function that combines the source + parameters into the MLC parameter. For example, for the above example, the function + would be: `lambda q, k, v: np.concatenate([q, k, v], axis=0)`. + + unused_params : Set[str] + Parameter names in the source weights that are not used in the MLC LLM model definition. + """ + + param_map: Dict[str, List[str]] = dataclasses.field(default_factory=dict) # noqa: UP006 + map_func: Dict[str, MapFuncVariadic] = dataclasses.field(default_factory=dict) # noqa: UP006 + unused_params: Set[str] = dataclasses.field(default_factory=set) # noqa: UP006 + + def add_mapping( + self, + map_from: str, + map_to: List[str], # noqa: UP006 + func: MapFuncVariadic, + ) -> None: + """Add a mapping from MLC parameters to source parametes as well as a mapping function.""" + self.param_map[map_from] = map_to + self.map_func[map_from] = func + + def add_unused(self, name: str): + """Add a parameter name in the source parameters to the set of unused parameters.""" + self.unused_params.add(name) + + +@dataclasses.dataclass +class QuantizeMapping: + """Mapping from a parameter in MLC LLM's model definition to its eventual names and values after + quantization. In certain group quantization, for example, `qkv_proj.weight` is mapped to + `qkv_proj.weight_quantized` and `qkv_proj.weight_scale` respectively. If a parameter's name is + not in the mapping, it is assumed to be unchanged, i.e. not quantized. + + Parameters + ---------- + param_map : Dict[str, List[str]] + A dictionary that maps the name of a parameter to its destination. For example, + in certain group quantization, the destinations of MLC parameter "qkv_proj.weight` are: + + - "qkv_proj.weight_quantized" + - "qkv_proj.weight_scale" + + map_func : Dict[str, Callable[Tensor, List[Tensor]]] + A dictionary that maps the name of a parameter to a function that splits the MLC parameter + into the destination parameters. + + Notes + ----- + There are two forms of weight conversion in MLC LLM, one is A) on-the-fly quantization to the + raw fp16/bf16/fp32 weights from HuggingFace, and the other is B) loading pre-quantized weights + from an external framework, e.g. AutoGPTQ, AutoAWQ. From the perspective of parameter + correspondence. + + - In case A), it is recommended that the weight loader take both `ExternMapping` and + `QuantizeMapping` as input, and do quantiaztion on the fly as a raw parameter being + loaded into RAM; + - In case B), a pass over `nn.Module` is recommended to take place first to converts parameters + from its non-quantized form to the quantized one, and then only `ExternMapping` is + used to convert the quantized parameters into the desired form. + """ + + param_map: Dict[str, List[str]] # noqa: UP006 + map_func: Dict[str, Callable[[Tensor], List[Tensor]]] # noqa: UP006 + + +__all__ = ["ExternMapping", "QuantizeMapping"] diff --git a/python/mlc_llm/loader/standard_loader.py b/python/mlc_llm/loader/standard_loader.py new file mode 100644 index 0000000..5444c1e --- /dev/null +++ b/python/mlc_llm/loader/standard_loader.py @@ -0,0 +1,152 @@ +"""Standard HuggingFace loader mapping helpers.""" + +from __future__ import annotations + +import functools +from collections.abc import Iterable, Sequence +from typing import Callable, Optional, Type # noqa: UP035 + +import numpy as np +from tvm.relax.frontend import nn + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +NameTransform = Callable[[str], str] +ExportSpecGetter = Callable[[nn.Module], object] + + +def _default_export_spec(model: nn.Module) -> object: + return model.get_default_spec() + + +def make_standard_hf_loader( + *, + model_cls: Type[nn.Module], # noqa: UP006 + layer_prefix: str = "model.layers", + qkv_names: Sequence[str] = ("q_proj", "k_proj", "v_proj"), + qkv_concat_axis: int = 0, + qkv_target_name: str = "qkv_proj", + add_qkv_bias: bool = False, + qkv_bias_optional: bool = False, + gate_up_names: Sequence[str] = ("gate_proj", "up_proj"), + gate_up_concat_axis: int = 0, + gate_up_target_name: str = "gate_up_proj", + include_qkv: bool = True, + include_gate_up: bool = True, + add_unused: Optional[Iterable[str]] = None, # noqa: UP045 + hf_prefix: str = "model.", + name_transform: Optional[NameTransform] = None, # noqa: UP045 + export_spec_getter: Optional[ExportSpecGetter] = None, # noqa: UP045 + num_layers_getter: Optional[Callable[[object], int]] = None, # noqa: UP045 +) -> Callable[[object, Quantization], ExternMapping]: + """Create a standard loader for HuggingFace weights. + + This handles the common QKV concatenation, gate+up concatenation, optional + QKV bias mapping, and passes through remaining parameters 1:1. + """ + + if not qkv_names: + include_qkv = False + if not gate_up_names: + include_gate_up = False + if not include_qkv: + qkv_names = () + if not include_gate_up: + gate_up_names = () + + def _default_name_transform(name: str) -> str: + # When hf_prefix is empty, strip the "model." prefix so models that + # expose bare top-level weights (no "model." namespace) still load. + if hf_prefix == "": + return name[6:] if name.startswith("model.") else name + return name + + name_transform_fn = name_transform or _default_name_transform + spec_getter = export_spec_getter or _default_export_spec + unused_names = tuple(add_unused or ()) + + def huggingface( + model_config: object, + quantization: Quantization, + ) -> ExternMapping: + model = model_cls(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=spec_getter(model), + allow_extern=True, + ) + named_parameters = dict(_named_params) + mapping = ExternMapping() + + if include_qkv or include_gate_up or unused_names: + if num_layers_getter is None: + num_layers = model_config.num_hidden_layers + else: + num_layers = num_layers_getter(model_config) + + for i in range(num_layers): + attn = f"{layer_prefix}.{i}.self_attn" + if include_qkv: + mlc_qkv_name = f"{attn}.{qkv_target_name}.weight" + mlc_param = named_parameters[mlc_qkv_name] + mapping.add_mapping( + mlc_qkv_name, + [name_transform_fn(f"{attn}.{name}.weight") for name in qkv_names], + functools.partial( + lambda q, k, v, dtype: np.concatenate( + [q, k, v], axis=qkv_concat_axis + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + if add_qkv_bias: + mlc_bias_name = f"{attn}.{qkv_target_name}.bias" + if (not qkv_bias_optional) or mlc_bias_name in named_parameters: + mlc_param = named_parameters[mlc_bias_name] + mapping.add_mapping( + mlc_bias_name, + [name_transform_fn(f"{attn}.{name}.bias") for name in qkv_names], + functools.partial( + lambda q, k, v, dtype: np.concatenate( + [q, k, v], axis=qkv_concat_axis + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + if include_gate_up: + mlp = f"{layer_prefix}.{i}.mlp" + mlc_gate_up_name = f"{mlp}.{gate_up_target_name}.weight" + if gate_up_names: + mlc_param = named_parameters[mlc_gate_up_name] + mapping.add_mapping( + mlc_gate_up_name, + [name_transform_fn(f"{mlp}.{name}.weight") for name in gate_up_names], + functools.partial( + lambda gate, up, dtype: np.concatenate( + [gate, up], axis=gate_up_concat_axis + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for unused_name in unused_names: + mapping.add_unused(name_transform_fn(f"{attn}.{unused_name}")) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [name_transform_fn(mlc_name)], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + return mapping + + return huggingface diff --git a/python/mlc_llm/loader/stats.py b/python/mlc_llm/loader/stats.py new file mode 100644 index 0000000..cdfd04a --- /dev/null +++ b/python/mlc_llm/loader/stats.py @@ -0,0 +1,93 @@ +"""Statistics of the loading process of parameter loaders""" + +import dataclasses +import time +from contextlib import contextmanager + +from mlc_llm.support import logging +from mlc_llm.support.style import green + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Stats: + """Statistics of the loading process of parameter loaders. + + Attributes + ---------- + load_time_sec : float + Time used in loading the parameters. + + map_time_sec : float + Time used in applying the mapping function, i.e. `ExternMapping.map_func`. + + quant_time_sec : float + Time used in quantizing the parameters, i.e. `QuantizeMapping.quant_func`. + + current_memory_gb : float + The current RAM usage in GB. + + total_memory_gb : float + The total size data loaded from disk in GB. + + max_memory_gb : float + The maximum RAM usage in GB. + + total_param_num: int + Total number of parameters (original non-MLC model weights), excluding unused params. + """ + + load_time_sec: float = 0.0 + map_time_sec: float = 0.0 + quant_time_sec: float = 0.0 + + current_memory_gb: float = 0.0 + total_memory_gb: float = 0.0 + max_memory_gb: float = 0.0 + + total_param_num: int = 0 + + def timer(self, attr): + """A context manager to time the scope and add the time to the attribute.""" + + @contextmanager + def timed_scope(): + start_time = time.time() + yield + elapsed_time = time.time() - start_time + setattr(self, attr, getattr(self, attr) + elapsed_time) + + return timed_scope() + + def mem_add(self, nbytes: int): + """Add the memory usage by the given number of bytes.""" + mem_gb = float(nbytes) / float(1024**3) + self.current_memory_gb += mem_gb + self.total_memory_gb += mem_gb + self.max_memory_gb = max(self.max_memory_gb, self.current_memory_gb) + + def mem_rm(self, nbytes: int): + """Remove the memory usage by the given number of bytes.""" + mem_gb = float(nbytes) / float(1024**3) + self.current_memory_gb -= mem_gb + + def log_time_info(self, weight_format: str): + """Log the time used in loading, pre-quantization and quantization.""" + logger.info( + "%s: %s loading: %.3f sec; Pre-quantization mapping: %.3f sec; Quantization: %.3f sec", + green("Time usage"), + weight_format, + self.load_time_sec, + self.map_time_sec, + self.quant_time_sec, + ) + + def log_mem_usage(self): + """Log the Memory usage information.""" + logger.info( + "%s: Peak RAM: %.3f GB. Total bytes loaded from disk: %.3f GB", + green("RAM usage"), + self.max_memory_gb, + self.total_memory_gb, + ) diff --git a/python/mlc_llm/loader/utils.py b/python/mlc_llm/loader/utils.py new file mode 100644 index 0000000..dab2494 --- /dev/null +++ b/python/mlc_llm/loader/utils.py @@ -0,0 +1,75 @@ +"""Common utilities for loading parameters""" + +import functools +import operator +from collections.abc import Iterator +from pathlib import Path +from typing import TYPE_CHECKING, Set, Tuple # noqa: UP035 + +import numpy as np + +from mlc_llm.support import logging + +if TYPE_CHECKING: + from .mapping import ExternMapping + + +logger = logging.getLogger(__name__) + + +def check_parameter_usage(param_map: "ExternMapping", extern_weights: Set[str]): # noqa: UP006 + """Check that all external parameters have been used and are stored in the weights file.""" + used_extern_names = set(functools.reduce(operator.iadd, param_map.param_map.values(), [])) + # Check 1. All extern parameters in the weight files are used unless explicitly specified + unused_extern_names = extern_weights - used_extern_names - param_map.unused_params + if unused_extern_names: + logger.warning( + "Unused extern parameters: %s", + ", ".join(sorted(unused_extern_names)), + ) + # Check 2. All extern parameters required are stored in the weight files + nonexistent_extern_names = used_extern_names - extern_weights + if nonexistent_extern_names: + raise ValueError( + "The following extern parameters do not exist in the weight files:\n " + + "\n ".join(sorted(nonexistent_extern_names)), + ) + + +def load_torch_shard(path: Path) -> Iterator[Tuple[str, np.ndarray]]: # noqa: UP006 + """Load and yield PyTorch format parameters.""" + import torch + + for name, param in torch.load(path, map_location=torch.device("cpu")).items(): + if param is None: + logger.warning("Encountered None param, skipping it: %s", name) + continue + param = param.detach().cpu() + dtype = str(param.dtype) + if dtype == "torch.bfloat16": + param = param.float() + param = param.numpy() + yield name, param + + +def load_safetensor_shard(path: Path) -> Iterator[Tuple[str, np.ndarray]]: # noqa: UP006 + """Load and yield SafeTensor format parameters.""" + import safetensors + import torch + + with safetensors.safe_open(path, framework="pt", device="cpu") as in_file: + for name in in_file.keys(): + param = in_file.get_tensor(name) + param = param.detach().cpu() + dtype = str(param.dtype) + if dtype == "torch.bfloat16": + import ml_dtypes + + param = param.view(torch.float16).cpu().numpy().view(ml_dtypes.bfloat16) + elif dtype == "torch.float8_e4m3fn": + import ml_dtypes + + param = param.view(torch.uint8).cpu().numpy().view(ml_dtypes.float8_e4m3fn) + else: + param = param.numpy() + yield name, param diff --git a/python/mlc_llm/model/__init__.py b/python/mlc_llm/model/__init__.py new file mode 100644 index 0000000..480c198 --- /dev/null +++ b/python/mlc_llm/model/__init__.py @@ -0,0 +1,4 @@ +"""Model definition for the compiler.""" + +from .model import MODELS, Model +from .model_preset import MODEL_PRESETS diff --git a/python/mlc_llm/model/baichuan/__init__.py b/python/mlc_llm/model/baichuan/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/baichuan/baichuan_loader.py b/python/mlc_llm/model/baichuan/baichuan_loader.py new file mode 100644 index 0000000..204c887 --- /dev/null +++ b/python/mlc_llm/model/baichuan/baichuan_loader.py @@ -0,0 +1,13 @@ +""" +This file specifies how MLC's BaichuanLM parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .baichuan_model import BaichuanForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=BaichuanForCausalLM, + include_qkv=False, +) diff --git a/python/mlc_llm/model/baichuan/baichuan_model.py b/python/mlc_llm/model/baichuan/baichuan_model.py new file mode 100644 index 0000000..7a3bbf7 --- /dev/null +++ b/python/mlc_llm/model/baichuan/baichuan_model.py @@ -0,0 +1,363 @@ +""" +Implementation for BAICHUAN architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class BaichuanConfig(ConfigBase): + """Configuration of the Baichuan model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + initializer_range: float + intermediate_size: int + rms_norm_eps: float + use_cache: bool + pad_token_id: int + bos_token_id: int + eos_token_id: int + tie_word_embeddings: bool + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + head_dim: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class BaichuanAttention(nn.Module): + def __init__(self, config: BaichuanConfig): + self.hidden_size = config.hidden_size + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.num_attention_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.W_pack = nn.Linear(self.hidden_size, 3 * self.num_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h = self.head_dim, self.num_heads + b, s, _ = hidden_states.shape + qkv = self.W_pack(hidden_states) + qkv = op.reshape(qkv, (b, s, 3 * h, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +class BaichuanMLP(nn.Module): + def __init__(self, config: BaichuanConfig): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class BaichuanDecoderLayer(nn.Module): + def __init__(self, config: BaichuanConfig): + norm_eps = config.rms_norm_eps + self.self_attn = BaichuanAttention(config=config) + self.mlp = BaichuanMLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_heads * hd + k = self.self_attn.num_heads * hd + v = self.self_attn.num_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.W_pack.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_gate_up", segs=[i, i], dim=0), + ) + _set( + self.mlp.down_proj.weight, + tp.ShardSingleDim("_shard_mlp_down_proj", dim=1), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class BaichuanModel(nn.Module): + def __init__(self, config: BaichuanConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [BaichuanDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class BaichuanForCausalLM(nn.Module): + def __init__(self, config: BaichuanConfig): + self.model = BaichuanModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.vocab_size = config.vocab_size + self.rope_theta = 10000 + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_attention_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/bert/__init__.py b/python/mlc_llm/model/bert/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/bert/bert_loader.py b/python/mlc_llm/model/bert/bert_loader.py new file mode 100644 index 0000000..799ac5d --- /dev/null +++ b/python/mlc_llm/model/bert/bert_loader.py @@ -0,0 +1,117 @@ +""" +This file specifies how MLC's BERT parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools +from typing import Literal + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .bert_model import BertConfig, BertModel + + +def huggingface( + model_config: BertConfig, + quantization: Quantization, + hf_prefix: Literal["", "bert."] = "", +) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : BertConfig + The configuration of the BERT model. + + quantization : Quantization + The quantization configuration. + + hf_prefix : Literal["", "bert."] + Prefix used in HuggingFace weight names. Defaults to "" for standard + BERT models. Use "bert." for BGE models whose weights are prefixed. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = BertModel(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + def to_hf(name: str) -> str: + return f"{hf_prefix}{name}" if hf_prefix else name + + for i in range(model_config.num_hidden_layers): + attn = f"encoder.layer.{i}.attention.self" + mlc_name = f"{attn}.qkv.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + to_hf(f"{attn}.query.weight"), + to_hf(f"{attn}.key.weight"), + to_hf(f"{attn}.value.weight"), + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + mlc_name = f"{attn}.qkv.bias" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + to_hf(f"{attn}.query.bias"), + to_hf(f"{attn}.key.bias"), + to_hf(f"{attn}.value.bias"), + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [to_hf(mlc_name)], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Mark unused weights that exist in HF but not in MLC + if hf_prefix: + mapping.add_unused(f"{hf_prefix}pooler.dense.weight") + mapping.add_unused(f"{hf_prefix}pooler.dense.bias") + + return mapping + + +def huggingface_bge(model_config: BertConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping for BGE models. + + BGE weights have no prefix but include extra unused weights: + pooler.dense.weight, pooler.dense.bias, embeddings.position_ids + """ + mapping = huggingface(model_config, quantization, "") + mapping.add_unused("pooler.dense.weight") + mapping.add_unused("pooler.dense.bias") + mapping.add_unused("embeddings.position_ids") + return mapping diff --git a/python/mlc_llm/model/bert/bert_model.py b/python/mlc_llm/model/bert/bert_model.py new file mode 100644 index 0000000..7be1b25 --- /dev/null +++ b/python/mlc_llm/model/bert/bert_model.py @@ -0,0 +1,273 @@ +""" +Implementation for BERT architecture. +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import te, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.support import logging +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class BertConfig(ConfigBase): + """Configuration of the BERT model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + intermediate_size: int + hidden_act: str + layer_norm_eps: float + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + type_vocab_size: int = 2 + pad_token_id: int = 0 + position_offset: int = 0 + head_dim: int = 0 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.intermediate_size is None or self.intermediate_size == -1: + self.intermediate_size = 4 * self.hidden_size + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %s (%d)", + bold("prefill_chunk_size"), + bold("context_window_size"), + self.context_window_size, + ) + self.prefill_chunk_size = self.context_window_size + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d (%s)", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + self.context_window_size, + bold("context_window_size"), + ) + self.prefill_chunk_size = self.context_window_size + + +class BertSelfAttention(nn.Module): + def __init__(self, config: BertConfig): + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads" + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.num_attention_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + + self.qkv = nn.Linear( + in_features=config.hidden_size, + out_features=3 * self.num_heads * self.head_dim, + bias=True, + ) + + def forward(self, hidden_states: Tensor, attention_mask: Tensor): + d, h = self.head_dim, self.num_heads + b, s, _ = hidden_states.shape + + qkv = self.qkv(hidden_states) + qkv = op.reshape(qkv, (b, s, 3 * h, d)) + q, k, v = op.split(qkv, 3, axis=2) + + # Attention + output = op_ext.attention(q, k, v, attention_mask) + return output + + +class BertSelfOutput(nn.Module): + def __init__(self, config: BertConfig): + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states: Tensor, input_tensor: Tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class BertAttention(nn.Module): + def __init__(self, config: BertConfig): + self.self = BertSelfAttention(config) + self.output = BertSelfOutput(config) + + def forward(self, hidden_states: Tensor, attention_mask: Tensor): + self_output = self.self(hidden_states, attention_mask) + attention_output = self.output(self_output, hidden_states) + return attention_output + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class BertIntermediate(nn.Module): + def __init__(self, config: BertConfig): + self.dense = nn.Linear(config.hidden_size, config.intermediate_size) + self.intermediate_act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_states: Tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.intermediate_act_fn(hidden_states) + return hidden_states + + +class BertOutput(nn.Module): + def __init__(self, config: BertConfig): + self.dense = nn.Linear(config.intermediate_size, config.hidden_size) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, hidden_states: Tensor, input_tensor: Tensor): + hidden_states = self.dense(hidden_states) + hidden_states = self.LayerNorm(hidden_states + input_tensor) + return hidden_states + + +class BertLayer(nn.Module): + def __init__(self, config: BertConfig): + self.attention = BertAttention(config) + self.intermediate = BertIntermediate(config) + self.output = BertOutput(config) + + def forward(self, hidden_states: Tensor, attention_mask: Tensor): + attention_output = self.attention(hidden_states, attention_mask) + intermediate_output = self.intermediate(attention_output) + layer_output = self.output(intermediate_output, attention_output) + return layer_output + + +class BertEncoder(nn.Module): + def __init__(self, config: BertConfig): + self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)]) + + def forward(self, hidden_states: Tensor, attention_mask: Tensor): + for layer in self.layer: + hidden_states = layer(hidden_states, attention_mask) + return hidden_states + + +class BertEmbeddings(nn.Module): + def __init__(self, config: BertConfig): + self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, dtype="float32") + self.position_embeddings = nn.Embedding( + config.context_window_size, config.hidden_size, dtype="float32" + ) + self.token_type_embeddings = nn.Embedding( + config.type_vocab_size, config.hidden_size, dtype="float32" + ) + self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, input_ids: Tensor, token_type_ids: Tensor, position_ids: Tensor): + words_embeddings = self.word_embeddings(input_ids) + position_embeddings = self.position_embeddings(position_ids) + token_type_embeddings = self.token_type_embeddings(token_type_ids) + + embeddings = words_embeddings + position_embeddings + token_type_embeddings + embeddings = self.LayerNorm(embeddings) + return embeddings + + +class BertModel(nn.Module): + def __init__(self, config: BertConfig): + self.embeddings = BertEmbeddings(config) + self.encoder = BertEncoder(config) + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def forward(self, inputs: Tensor, attention_mask: Tensor): + # TODO: XLM-RoBERTa models use position indices starting from pad_token_id + 1 + # (e.g., [2, 3, 4, ...] when pad_token_id=1), while this implementation uses + # [0, 1, 2, ...]. For XLM-RoBERTa models (e.g., bge-m3), the position_embeddings + # weights need to be shifted during weight conversion to compensate. + def _input_positions(inputs: te.Tensor): + b, s = inputs.shape + return te.compute((b, s), lambda _, j: j.astype("int32"), name="input_positions") + + input_positions = op.tensor_expr_op( + _input_positions, + name_hint="input_positions", + args=[inputs], + ) + + token_type_ids = op.zeros(inputs.shape, dtype="int32") + + embeddings = self.embeddings(inputs, token_type_ids, input_positions) + encoder_output = self.encoder(embeddings, attention_mask) + return encoder_output + + def prefill(self, inputs: Tensor, attention_mask: Tensor): + def _attention_mask(mask: te.Tensor, zero, batch_size, seq_len): + return te.compute( + (batch_size, 1, seq_len, seq_len), + lambda b, _, i, j: tirx.if_then_else( + tirx.any(mask[b, i] == zero, mask[b, j] == zero), + tirx.min_value(self.dtype), + tirx.max_value(self.dtype), + ), + name="attention_mask_prefill", + ) + + batch_size, seq_len = inputs.shape + attention_mask_2d = op.tensor_expr_op( + _attention_mask, + name_hint="attention_mask_prefill", + args=[attention_mask, tirx.IntImm("int32", 0), batch_size, seq_len], + ) + return self.forward(inputs, attention_mask_2d) + + def get_default_spec(self): + mod_spec = { + "prefill": { + "inputs": nn.spec.Tensor(["batch_size", "seq_len"], "int32"), + "attention_mask": nn.spec.Tensor(["batch_size", "seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/chatglm3/__init__.py b/python/mlc_llm/model/chatglm3/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/chatglm3/chatglm3_loader.py b/python/mlc_llm/model/chatglm3/chatglm3_loader.py new file mode 100644 index 0000000..b9ae2a7 --- /dev/null +++ b/python/mlc_llm/model/chatglm3/chatglm3_loader.py @@ -0,0 +1,63 @@ +""" +This file specifies how MLC's ChatGLM3 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .chatglm3_model import ChatGLMForCausalLM, GLMConfig + + +def huggingface(model_config: GLMConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : GLMConfig + The configuration of the Baichuan model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = ChatGLMForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + mlc_name = "transformer.embedding.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + ["transformer.embedding.word_embeddings.weight"], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/chatglm3/chatglm3_model.py b/python/mlc_llm/model/chatglm3/chatglm3_model.py new file mode 100644 index 0000000..5f3eb3d --- /dev/null +++ b/python/mlc_llm/model/chatglm3/chatglm3_model.py @@ -0,0 +1,446 @@ +""" +Implementation for CHATGLM3 architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class GLMConfig(ConfigBase): + """Configuration of the ChatGLM model.""" + + hidden_size: int + num_layers: int + kv_channels: int + num_attention_heads: int + ffn_hidden_size: int + layernorm_epsilon: float + post_layer_norm: bool + rmsnorm: bool + add_bias_linear: bool + add_qkv_bias: bool + apply_query_key_layer_scaling: bool + multi_query_attention: bool + multi_query_group_num: int + vocab_size: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + head_dim: int = 0 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.vocab_size == 0: + for name in ["padded_vocab_size"]: + if name in self.kwargs: + self.vocab_size = self.kwargs.pop(name) + if self.context_window_size == 0: + for name in ["max_position_embeddings", "seq_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class GLMAttention(nn.Module): + def __init__(self, config: GLMConfig): + self.hidden_size = config.hidden_size + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads" + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.num_attention_heads // config.tensor_parallel_shards + self.multi_query_attention = config.multi_query_attention + self.num_key_value_heads = ( + config.multi_query_group_num + if config.multi_query_attention + else config.num_attention_heads + ) // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.query_key_value = nn.Linear( + config.hidden_size, + (2 * self.num_key_value_heads + self.num_heads) * self.head_dim, + bias=config.add_bias_linear or config.add_qkv_bias, + ) + self.dense = nn.Linear( + self.num_heads * self.head_dim, + config.hidden_size, + bias=config.add_bias_linear, + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.query_key_value(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, h_q, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.dense(output) + return attn_output + + +class GLMMLP(nn.Module): + def __init__(self, config: GLMConfig): + if config.ffn_hidden_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split ffn hidden size {config.ffn_hidden_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.ffn_hidden_size = config.ffn_hidden_size // config.tensor_parallel_shards + + self.dense_h_to_4h = nn.Linear( + config.hidden_size, + self.ffn_hidden_size * 2, + bias=config.add_bias_linear, + ) + self.dense_4h_to_h = nn.Linear( + self.ffn_hidden_size, + config.hidden_size, + bias=config.add_bias_linear, + ) + + def swiglu(x): + x = nn.chunk(x, 2, dim=-1) + return nn.silu(x[0]) * x[1] + + self.activation_func = swiglu + + def forward(self, x): + intermediate_parallel = self.dense_h_to_4h(x) + intermediate_parallel = self.activation_func(intermediate_parallel) + output = self.dense_4h_to_h(intermediate_parallel) + return output + + +class GLMBlock(nn.Module): + def __init__(self, config: GLMConfig): + self.self_attention = GLMAttention(config=config) + self.mlp = GLMMLP(config) + self.input_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.layernorm_epsilon, bias=False + ) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.layernorm_epsilon, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attention.num_heads * hd + k = self.self_attention.num_key_value_heads * hd + v = self.self_attention.num_key_value_heads * hd + _set( + self.self_attention.query_key_value.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + if config.add_bias_linear or config.add_qkv_bias: + _set( + self.self_attention.query_key_value.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set( + self.self_attention.dense.weight, + tp.ShardSingleDim("_shard_dense_weight", dim=1), + ) + if config.add_bias_linear: + _set( + self.self_attention.dense.bias, + tp.ShardSingleDim("_shard_dense_bias", dim=0), + ) + _set( + self.mlp.dense_h_to_4h.weight, + tp.ShardSingleDim("_shard_dense_h_to_4h_weight", dim=0), + ) + if config.add_bias_linear: + _set( + self.mlp.dense_h_to_4h.bias, + tp.ShardSingleDim("_shard_dense_h_to_4h_bias", dim=0), + ) + _set( + self.mlp.dense_4h_to_h.weight, + tp.ShardSingleDim("_shard_dense_4h_to_h", dim=1), + ) + if config.add_bias_linear: + _set( + self.mlp.dense_4h_to_h.bias, + tp.ShardSingleDim("_shard_dense_4h_to_h_bias", dim=1), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attention(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class GLMTransformer(nn.Module): + """Transformer class.""" + + def __init__(self, config: GLMConfig): + self.post_layer_norm = config.post_layer_norm + + # Number of layers. + self.num_layers = config.num_layers + + # Transformer layers. + self.layers = nn.ModuleList([GLMBlock(config) for _ in range(config.num_layers)]) + + if self.post_layer_norm: + if config.rmsnorm: + self.final_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.layernorm_epsilon, bias=False + ) + else: + self.final_layernorm = nn.LayerNorm(config.hidden_size, config.layernorm_epsilon) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.final_layernorm(hidden_states) + return hidden_states + + +class ChatGLMModel(nn.Module): + def __init__(self, config: GLMConfig): + self.embedding = nn.Embedding(config.vocab_size, config.hidden_size) + self.encoder = GLMTransformer(config) + self.output_layer = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + hidden_states = self.encoder(hidden_states, paged_kv_cache) + return hidden_states + + +class ChatGLMForCausalLM(nn.Module): + def __init__(self, config: GLMConfig): + self.transformer = ChatGLMModel(config) + self.num_hidden_layers = config.num_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = ( + config.multi_query_group_num + if config.multi_query_attention + else config.num_attention_heads + ) + self.head_dim = config.head_dim + self.vocab_size = config.vocab_size + self.rope_theta = 10000 + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.transformer(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.transformer.output_layer(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.transformer.embedding(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.transformer.output_layer(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + logits = self.transformer.output_layer(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/cohere/__init__.py b/python/mlc_llm/model/cohere/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/cohere/cohere_loader.py b/python/mlc_llm/model/cohere/cohere_loader.py new file mode 100644 index 0000000..e873691 --- /dev/null +++ b/python/mlc_llm/model/cohere/cohere_loader.py @@ -0,0 +1,115 @@ +""" +This file specifies how MLC's Cohere parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization, make_awq_quant + +from .cohere_model import CohereConfig, CohereForCausalLM + +awq_quant = make_awq_quant(CohereForCausalLM) + + +def _cohere_name_transform(name: str) -> str: + if "out_proj." in name: + return name.replace("out_proj.", "o_proj.") + return name + + +huggingface = make_standard_hf_loader( + model_cls=CohereForCausalLM, + include_gate_up=False, + name_transform=_cohere_name_transform, +) + + +# https://huggingface.co/alijawad07/aya-23-8B-AWQ-GEMM/tree/main +def awq(model_config: CohereConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of AWQ parameters. + Parameters + ---------- + model_config : CohereConfig + The configuration of the Cohere model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to AWQ. + """ + model, _ = awq_quant(model_config, quantization) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + def _add(mlc_name, hf_name): + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=named_parameters[mlc_name].dtype, + ), + ) + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"model.layers.{i}.self_attn" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{attn}.qkv_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{quantize_suffix}", + f"{attn}.k_proj.{quantize_suffix}", + f"{attn}.v_proj.{quantize_suffix}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate( + [q, k, v], + axis=1, # AWQ GEMM would transpose the weight + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + _add(f"{attn}.out_proj.{quantize_suffix}", f"{attn}.o_proj.{quantize_suffix}") + + # Concat gate and up in MLP + mlp = f"model.layers.{i}.mlp" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + _add(f"{mlp}.up_proj.{quantize_suffix}", f"{mlp}.up_proj.{quantize_suffix}") + _add( + f"{mlp}.gate_proj.{quantize_suffix}", + f"{mlp}.gate_proj.{quantize_suffix}", + ) + _add( + f"{mlp}.down_proj.{quantize_suffix}", + f"{mlp}.down_proj.{quantize_suffix}", + ) + + # inv_freq is not used in the model + # mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + return mapping diff --git a/python/mlc_llm/model/cohere/cohere_model.py b/python/mlc_llm/model/cohere/cohere_model.py new file mode 100644 index 0000000..d2a3b4c --- /dev/null +++ b/python/mlc_llm/model/cohere/cohere_model.py @@ -0,0 +1,406 @@ +""" +Implementation for Aya23 architecture +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class CohereConfig(ConfigBase): + """Configuration of the Cohere Aya-23 model""" + + model_type: str # cohere + hidden_size: int + vocab_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + intermediate_size: int + layer_norm_eps: float + position_embedding_base: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + head_dim: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs["rope_theta"] + else: + self.position_embedding_base = 10000 + + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + if self.num_key_value_heads == 0 or self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size, ( + "head_dim * num_attention_heads != hidden_size" + ) + assert self.num_attention_heads % self.num_key_value_heads == 0, ( + "num_attention_heads % num_key_value_heads != 0" + ) + + +class CohereMLP(nn.Module): + def __init__(self, config: CohereConfig): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_proj = nn.Linear(config.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + down_proj = self.down_proj(op.silu(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class CohereAttention(nn.Module): + def __init__(self, config: CohereConfig): + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_attention_heads % config.tensor_parallel_shards == 0, ( + f"num_attention_heads({config.num_attention_heads}) " + "must be divisible by tensor_parallel_shards" + ) + self.num_key_value_heads = config.num_key_value_heads // config.tensor_parallel_shards + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_attention_heads({config.num_key_value_heads}) " + "must be divisible by tensor_parallel_shards" + ) + self.head_dim = config.head_dim + + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=False, + ) + self.out_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.out_proj(output) + + +class CohereDecoderLayer(nn.Module): + def __init__(self, config: CohereConfig): + super().__init__() + self.self_attn = CohereAttention(config) + self.mlp = CohereMLP(config) + self.input_layernorm = CohereNorm(config.hidden_size, eps=config.layer_norm_eps) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.out_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_proj, + tp.ShardSingleDim("_shard_mlp_gate", segs=[i, i], dim=0), + ) + _set(self.mlp.up_proj, tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0)) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + hidden_ln = self.input_layernorm(hidden_states) + attn = self.self_attn(hidden_ln, paged_kv_cache, layer_id) + mlp = self.mlp(hidden_ln) + hidden_states = self._apply_parallel_residual(attn, residual=hidden_states) + hidden_states = self._apply_parallel_residual(mlp, residual=hidden_states) + return hidden_states + + def _apply_parallel_residual(self, mlp_out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(mlp_out + residual / self.tensor_parallel_shards, "sum") + return mlp_out + residual + + +class CohereNorm(nn.Module): + def __init__( + self, normalized_shape: int, eps: float = 1e-5, dtype: Optional[str] = None + ) -> None: + super().__init__() + self.normalized_shape = normalized_shape + self.eps = eps + self.weight = nn.Parameter((normalized_shape,), dtype=dtype) + + def forward(self, x: Tensor) -> Tensor: + return op.layer_norm( + x, + normalized_shape=self.normalized_shape, + weight=self.weight, + bias=None, + eps=self.eps, + ) + + +class CohereEmbedding(nn.Embedding): + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class CohereModel(nn.Module): + def __init__(self, config: CohereConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = CohereEmbedding("vocab_size", config.hidden_size) + self.layers = nn.ModuleList( + [CohereDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = CohereNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class CohereForCausalLM(nn.Module): + def __init__(self, config: CohereConfig) -> None: + super().__init__() + self.model = CohereModel(config) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + lm_logits = self.model.embed_tokens.lm_head_forward(hidden_states) + if lm_logits.dtype != "float32": + lm_logits = lm_logits.astype("float32") + return lm_logits + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + # logits = self.lm_head(hidden_states) + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + + if logits.dtype != "float32": + logits = logits.astype("float32") + + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + embeds = self.model.embed_tokens(input_ids) + return embeds + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/deepseek/__init__.py b/python/mlc_llm/model/deepseek/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/deepseek/deepseek_loader.py b/python/mlc_llm/model/deepseek/deepseek_loader.py new file mode 100644 index 0000000..e76e076 --- /dev/null +++ b/python/mlc_llm/model/deepseek/deepseek_loader.py @@ -0,0 +1,149 @@ +""" +This file specifies how MLC's Deepseek parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .deepseek_model import DeepseekConfig, DeepseekForCausalLM + + +def huggingface(model_config: DeepseekConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : MiniCPMConfig + The configuration of the MiniCPM model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = DeepseekForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # map attention weight + attn = f"model.layers.{i}.self_attn" + for weight_type in ["weight"]: + mlc_name = f"{attn}.wqkv_pack.{weight_type}" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{weight_type}", + f"{attn}.k_proj.{weight_type}", + f"{attn}.v_proj.{weight_type}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for i in range(model_config.num_hidden_layers): + if i >= model_config.first_k_dense_replace and i % model_config.moe_layer_freq == 0: + # map mlp shared expert weight + mlp = f"model.layers.{i}.mlp" + shared_expert = f"{mlp}.shared_experts" + mlc_name = f"{shared_expert}.gate_up_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{shared_expert}.gate_proj.weight", + f"{shared_expert}.up_proj.weight", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + # map mlp moe gate and up weight + mlc_name = f"{mlp}.moe_gate_up_proj.weight" + + def combine_expert_gate_up(*hf_params, dtype): + stack = [] + for i in range(0, len(hf_params), 2): + stack.append(np.concatenate([hf_params[i], hf_params[i + 1]], axis=0)) + return np.stack(stack, axis=0).astype(dtype) + + mapping.add_mapping( + mlc_name, + functools.reduce( + lambda a, b: a + b, + [ + [ + f"{mlp}.experts.{expert_id}.gate_proj.weight", + f"{mlp}.experts.{expert_id}.up_proj.weight", + ] + for expert_id in range(model_config.n_routed_experts) + ], + ), + functools.partial( + combine_expert_gate_up, + dtype=mlc_param.dtype, + ), + ) + + # map mlp moe gate and up weight + mlc_name = f"{mlp}.moe_down_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.experts.{expert_id}.down_proj.weight" + for expert_id in range(model_config.n_routed_experts) + ], + functools.partial( + lambda *hf_params, dtype: np.stack(hf_params, axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + else: + # map mlp weight + mlp = f"model.layers.{i}.mlp" + mlc_name = f"{mlp}.gate_up_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.gate_proj.weight", + f"{mlp}.up_proj.weight", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/deepseek/deepseek_model.py b/python/mlc_llm/model/deepseek/deepseek_model.py new file mode 100644 index 0000000..fc8e5eb --- /dev/null +++ b/python/mlc_llm/model/deepseek/deepseek_model.py @@ -0,0 +1,516 @@ +""" +Implementation for Deepseek architecture. +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.nn.expert import MixtralExperts +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class DeepseekConfig(ConfigBase): + """Configuration of the Deepseek model.""" + + vocab_size: int + hidden_size: int + intermediate_size: int + moe_intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + n_shared_experts: int + n_routed_experts: int + moe_layer_freq: int + first_k_dense_replace: int + hidden_act: str + norm_topk_prob: bool + attention_bias: bool + rms_norm_eps: float + use_cache: bool + bos_token_id: int + eos_token_id: int + tie_word_embeddings: bool = False + rope_theta: int = 10000 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + head_dim: int = 0 + max_batch_size: int = 1 + num_experts_per_tok: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class DeepseekAttention(nn.Module): + def __init__(self, config: DeepseekConfig): + super().__init__() # Make sure to call the parent class constructor + self.hidden_size = config.hidden_size + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + + self.attention_bias = config.attention_bias + self.num_heads = config.num_attention_heads // self.tensor_parallel_shards + self.num_key_value_heads = config.num_key_value_heads // self.tensor_parallel_shards + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.head_dim = config.head_dim + self.max_position_embeddings = config.context_window_size + + self.wqkv_pack = nn.Linear( + in_features=self.hidden_size, + out_features=(self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=self.attention_bias, + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.wqkv_pack(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class DeepseekMLP(nn.Module): + def __init__(self, config: DeepseekConfig, intermediate_size=None): + self.hidden_size = config.hidden_size + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) // config.tensor_parallel_shards + + self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class DeepseekMoE(nn.Module): + def __init__(self, config: DeepseekConfig): + self.num_local_experts = config.n_routed_experts + self.num_experts_per_tok = config.num_experts_per_tok + self.gate = nn.Linear(config.hidden_size, config.n_routed_experts, bias=False) + self.norm_topk_prob = config.norm_topk_prob + self.moe_intermediate_size = config.moe_intermediate_size // config.tensor_parallel_shards + self.moe_gate_up_proj = MixtralExperts( + self.num_local_experts, + in_features=config.hidden_size, + out_features=2 * self.moe_intermediate_size, + tensor_parallel_shards=config.tensor_parallel_shards, + ) + self.moe_down_proj = MixtralExperts( + self.num_local_experts, + in_features=self.moe_intermediate_size, + out_features=config.hidden_size, + tensor_parallel_shards=config.tensor_parallel_shards, + ) + self.dtype = "float32" + + if config.n_shared_experts is not None: + intermediate_size = self.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepseekMLP(config, intermediate_size=intermediate_size) + + def forward(self, x: Tensor): + def _expert_forward(x: Tensor, indptr: Tensor): + x1_x3 = self.moe_gate_up_proj(x, indptr) + x1, x3 = op.split(x1_x3, indices_or_sections=2, axis=-1) + x = self.moe_down_proj(op.silu(x1) * x3, indptr) + return x + + experts_per_tok = self.num_experts_per_tok + num_experts = self.num_local_experts + b, s, h = x.shape + num_tokens = b * s + x = op.reshape(x, (num_tokens, h)) + gate = self.gate(x) # (b * s, num_routed_experts) + expert_weights, expert_indices = op_ext.moe_misc.gating_softmax_topk( + gate, experts_per_tok, norm_topk_prob=self.norm_topk_prob + ) + + if num_tokens == 1: + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = _expert_forward(x, expert_indices) + else: + # cumsum: [num_tokens * local_experts] + cumsum = op_ext.moe_misc.moe_cumsum(expert_indices, num_experts) + # indices: [num_tokens * experts_per_tok] + reverse_indices, token_indices = op_ext.moe_misc.get_indices(cumsum, expert_indices) + # indptr: [num_local_experts + 1] + indptr = op_ext.moe_misc.get_indptr( + cumsum, num_experts, num_tokens, inclusive=False, out_dtype="int32" + ) + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = op.take(x, token_indices, axis=0) + moe_hidden_states = _expert_forward(moe_hidden_states, indptr) + moe_hidden_states = op_ext.moe_misc.scatter_output(moe_hidden_states, reverse_indices) + + # moe_hidden_states: [num_tokens, experts_per_tok, hidden_size] + expert_weights = expert_weights.reshape(num_tokens, experts_per_tok, 1) + moe_hidden_states = ( + moe_hidden_states.reshape(num_tokens, experts_per_tok, h) * expert_weights + ) + # moe_hidden_states: [num_tokens, hidden_size] + moe_hidden_states = op_ext.moe_misc.moe_sum(moe_hidden_states, dim=1) + + shared_expert_hidden_states = self.shared_experts(x) + + final_hidden_states = moe_hidden_states + shared_expert_hidden_states + final_hidden_states = op.reshape(final_hidden_states, (b, s, h)) + return final_hidden_states + + +class DeepseekDecoderLayer(nn.Module): + def __init__(self, config: DeepseekConfig, layer_idx: int): + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.self_attn = DeepseekAttention(config) + self.num_experts = config.n_routed_experts + self.mlp = ( + DeepseekMoE(config) + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + else DeepseekMLP(config) + ) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ): + i = self.mlp.moe_intermediate_size + else: + i = self.mlp.intermediate_size + _set( + self.self_attn.wqkv_pack.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ): + _set( + self.mlp.moe_gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=1), + ) + _set( + self.mlp.moe_down_proj.weight, + tp.ShardSingleDim("_shard_mlp_down", dim=2), + ) + + else: + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set( + self.mlp.down_proj.weight, + tp.ShardSingleDim("_shard_mlp_down", dim=1), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.input_layernorm(hidden_states) + out = self.self_attn(out, paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class DeepseekModel(nn.Module): + def __init__(self, config: DeepseekConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [ + DeepseekDecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class DeepseekForCausalLM(nn.Module): + def __init__(self, config: DeepseekConfig): + self.model = DeepseekModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.tensor_parallel_shards = config.tensor_parallel_shards + self.head_dim = config.head_dim + self.vocab_size = config.vocab_size + self.rope_theta = config.rope_theta + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/deepseek_v2/__init__.py b/python/mlc_llm/model/deepseek_v2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/deepseek_v2/deepseek_v2_loader.py b/python/mlc_llm/model/deepseek_v2/deepseek_v2_loader.py new file mode 100644 index 0000000..d116c22 --- /dev/null +++ b/python/mlc_llm/model/deepseek_v2/deepseek_v2_loader.py @@ -0,0 +1,258 @@ +""" +This file specifies how MLC's Deepseek-V2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools +from typing import Callable, List # noqa: UP035 + +import numpy as np + +from mlc_llm.loader import ExternMapping, QuantizeMapping +from mlc_llm.quantization import BlockScaleQuantize, Quantization + +from .deepseek_v2_model import DeepseekV2Config, DeepseekV2ForCausalLM + + +def huggingface(model_config: DeepseekV2Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : DeepseekV2Config + The configuration of the DeepseekV2 model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = DeepseekV2ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + if isinstance(quantization, BlockScaleQuantize): + # Convert the model to block-scale quantized model before loading parameters + model = quantization.quantize_model(model, QuantizeMapping({}, {}), "") + if model_config.weight_block_size is None: + raise ValueError( + "The input DeepSeek model is not fp8 block quantized. " + "Thus BlockScaleQuantize is not supported." + ) + + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + if ( + not isinstance(quantization, BlockScaleQuantize) + and model_config.weight_block_size is not None + ): + raise ValueError( + "The input DeepSeek model is fp8 block quantized. " + "Please use BlockScaleQuantize for the model." + ) + + # Helper function to add both weight and scale mappings + def add_weight_and_scale_mapping( + weight_mlc_name: str, + weight_hf_names: List[str], # noqa: UP006 + weight_transform_func: Callable, + ): + mlc_param = named_parameters[weight_mlc_name] + mapping.add_mapping( + weight_mlc_name, + weight_hf_names, + functools.partial(weight_transform_func, dtype=mlc_param.dtype), + ) + + if isinstance(quantization, BlockScaleQuantize): + scale_mlc_name = f"{weight_mlc_name}_scale_inv" + if scale_mlc_name in named_parameters: + scale_hf_names = [f"{name}_scale_inv" for name in weight_hf_names] + scale_param = named_parameters[scale_mlc_name] + mapping.add_mapping( + scale_mlc_name, + scale_hf_names, + functools.partial(weight_transform_func, dtype=scale_param.dtype), + ) + + for i in range(model_config.num_hidden_layers): + if i >= model_config.first_k_dense_replace and i % model_config.moe_layer_freq == 0: + # map mlp shared expert weight + mlp = f"model.layers.{i}.mlp" + shared_expert = f"{mlp}.shared_experts" + add_weight_and_scale_mapping( + f"{shared_expert}.gate_up_proj.weight", + [ + f"{shared_expert}.gate_proj.weight", + f"{shared_expert}.up_proj.weight", + ], + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + ) + + # map mlp moe gate and up weight + def combine_expert_gate_up(*hf_params, dtype): + stack = [] + for i in range(0, len(hf_params), 2): + stack.append(np.concatenate([hf_params[i], hf_params[i + 1]], axis=0)) + return np.stack(stack, axis=0).astype(dtype) + + add_weight_and_scale_mapping( + f"{mlp}.moe_gate_up_proj.weight", + functools.reduce( + lambda a, b: a + b, + [ + [ + f"{mlp}.experts.{expert_id}.gate_proj.weight", + f"{mlp}.experts.{expert_id}.up_proj.weight", + ] + for expert_id in range(model_config.n_routed_experts) + ], + ), + combine_expert_gate_up, + ) + + # map mlp moe down projection weight + add_weight_and_scale_mapping( + f"{mlp}.moe_down_proj.weight", + [ + f"{mlp}.experts.{expert_id}.down_proj.weight" + for expert_id in range(model_config.n_routed_experts) + ], + lambda *hf_params, dtype: np.stack(hf_params, axis=0).astype(dtype), + ) + + # map moe e_score_correction_bias + if model_config.topk_method == "noaux_tc": + mlc_name = f"{mlp}.e_score_correction_bias" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{mlp}.gate.e_score_correction_bias"], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + else: + # map mlp weight + mlp = f"model.layers.{i}.mlp" + add_weight_and_scale_mapping( + f"{mlp}.gate_up_proj.weight", + [ + f"{mlp}.gate_proj.weight", + f"{mlp}.up_proj.weight", + ], + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + ) + + # map MLA kv_b_proj weight + attn = f"model.layers.{i}.self_attn" + mlc_name = f"{attn}.w_uk" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{attn}.kv_b_proj.weight"], + functools.partial( + lambda kv_b_proj, dtype: ( + np.split( + kv_b_proj.reshape( + model_config.num_key_value_heads, + model_config.qk_nope_head_dim + model_config.v_head_dim, + model_config.kv_lora_rank, + ), + indices_or_sections=[model_config.qk_nope_head_dim], + axis=1, + )[0] + .transpose(0, 2, 1) + .astype(dtype) + ), + dtype=mlc_param.dtype, + ), + ) + if isinstance(quantization, BlockScaleQuantize): + scale_mlc_name = f"{attn}.w_uk_scale_inv" + mlc_param = named_parameters[scale_mlc_name] + mapping.add_mapping( + scale_mlc_name, + [f"{attn}.kv_b_proj.weight_scale_inv"], + functools.partial( + lambda kv_b_proj, dtype: ( + np.split( + kv_b_proj.reshape( + model_config.num_key_value_heads, + (model_config.qk_nope_head_dim + model_config.v_head_dim) + // quantization.weight_block_size[0], + model_config.kv_lora_rank // quantization.weight_block_size[1], + ), + indices_or_sections=[ + model_config.qk_nope_head_dim // quantization.weight_block_size[0] + ], + axis=1, + )[0] + .transpose(0, 2, 1) + .astype(dtype) + ), + dtype=mlc_param.dtype, + ), + ) + mlc_name = f"{attn}.w_uv" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{attn}.kv_b_proj.weight"], + functools.partial( + lambda kv_b_proj, dtype: np.split( + kv_b_proj.reshape( + model_config.num_key_value_heads, + model_config.qk_nope_head_dim + model_config.v_head_dim, + model_config.kv_lora_rank, + ), + indices_or_sections=[model_config.qk_nope_head_dim], + axis=1, + )[1].astype(dtype), + dtype=mlc_param.dtype, + ), + ) + if isinstance(quantization, BlockScaleQuantize): + scale_mlc_name = f"{attn}.w_uv_scale_inv" + mlc_param = named_parameters[scale_mlc_name] + mapping.add_mapping( + scale_mlc_name, + [f"{attn}.kv_b_proj.weight_scale_inv"], + functools.partial( + lambda kv_b_proj, dtype: np.split( + kv_b_proj.reshape( + model_config.num_key_value_heads, + (model_config.qk_nope_head_dim + model_config.v_head_dim) + // quantization.weight_block_size[0], + model_config.kv_lora_rank // quantization.weight_block_size[1], + ), + indices_or_sections=[ + model_config.qk_nope_head_dim // quantization.weight_block_size[0] + ], + axis=1, + )[1].astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/deepseek_v2/deepseek_v2_model.py b/python/mlc_llm/model/deepseek_v2/deepseek_v2_model.py new file mode 100644 index 0000000..2b39935 --- /dev/null +++ b/python/mlc_llm/model/deepseek_v2/deepseek_v2_model.py @@ -0,0 +1,873 @@ +""" +Implementation for Deepseek V2 architecture +""" + +import dataclasses +import math +from typing import Any, Dict, Literal, Optional, Tuple # noqa: UP035 + +from tvm import te, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op +from tvm.relax.frontend.nn.llm import position_embedding + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.nn.expert import MixtralExperts +from mlc_llm.op import batch_matmul +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class DeepseekV2Config(ConfigBase): + """Configuration of the Deepseek V2 model.""" + + vocab_size: int + hidden_size: int + intermediate_size: int + moe_intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + n_shared_experts: int + n_routed_experts: int + num_experts_per_tok: int + norm_topk_prob: bool + first_k_dense_replace: int + moe_layer_freq: int + routed_scaling_factor: float + scoring_func: str + topk_method: Literal["greedy", "group_limited_greedy", "noaux_tc"] + n_group: int + topk_group: int + attention_bias: bool + kv_lora_rank: int + qk_rope_head_dim: int + v_head_dim: int + qk_nope_head_dim: int + rms_norm_eps: float + rope_theta: int + q_lora_rank: Optional[int] = None + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + dtype: str = "float32" + max_batch_size: int = 1 + weight_block_size: Optional[Tuple[int, int]] = None # noqa: UP006 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if "quantization_config" in self.kwargs: + quantization_config = self.kwargs.get("quantization_config") + if ( + isinstance(quantization_config, dict) + and quantization_config.get("activation_scheme", "") == "dynamic" + and quantization_config.get("fmt", "") == "e4m3" + and quantization_config.get("quant_method", "") == "fp8" + and "weight_block_size" in quantization_config + ): + self.weight_block_size = quantization_config.get("weight_block_size") + if ( + not isinstance(self.weight_block_size, (tuple, list)) + or len(self.weight_block_size) != 2 + ): + raise ValueError( + "Invalid DeepSeek model quantization config: " + "weight_block_size must be a tuple of two integers, " + f"got {self.weight_block_size} of type {type(self.weight_block_size)}" + ) + else: + raise ValueError( + "Invalid DeepSeek model quantization config: unrecognized quantization config: " + f"{quantization_config}" + ) + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + assert self.num_attention_heads % self.num_key_value_heads == 0 + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 2048), + ) + self.prefill_chunk_size = min(self.context_window_size, 2048) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 2048), + ) + self.prefill_chunk_size = min(self.context_window_size, 2048) + + +class DeepseekV2MLP(nn.Module): + def __init__(self, config: DeepseekV2Config, hidden_size=None, intermediate_size=None): + super().__init__() + self.hidden_size = config.hidden_size if hidden_size is None else hidden_size + intermediate_size = ( + config.intermediate_size if intermediate_size is None else intermediate_size + ) + if intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MoE intermediate size {intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = intermediate_size // config.tensor_parallel_shards + + self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + def forward(self, x: Tensor) -> Tensor: + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class DeepseekV2YarnRotaryEmbedding(nn.Module): + def __init__(self, config: DeepseekV2Config): + self.rope_fn = position_embedding.switch_rope_freq_func(config.rope_scaling) + self.rotary_dim = config.qk_rope_head_dim + self.theta = config.rope_theta + + def forward( + self, + q: Tensor, + k: Tensor, + positions: Tensor, + ): + def _rope_fused(x: te.Tensor, positions: te.Tensor): + _, _, _, d_dim = x.shape + d_dim_half = d_dim // 2 + dtype = x.dtype + + def compute(b: tirx.Var, s: tirx.Var, h: tirx.Var, d: tirx.Var): + d1 = d // d_dim_half + d2 = d % d_dim_half + + cos_freq, sin_freq, var_map = self.rope_fn( + positions[s], d, self.rotary_dim, self.theta, dtype + ) + cos = x[b, s, h, d2 * 2 + d1] * cos_freq + + partner_d = tirx.if_then_else( + d < self.rotary_dim // 2, + d + self.rotary_dim // 2, + d - self.rotary_dim // 2, + ) + + partner_d1 = partner_d // d_dim_half + partner_d2 = partner_d % d_dim_half + sin = ( + x[b, s, h, partner_d2 * 2 + partner_d1] + * sin_freq + * tirx.if_then_else( + d < self.rotary_dim // 2, + tirx.const(-1, dtype), + tirx.const(1, dtype), + ) + ) + expr = cos + sin + for var, val in var_map.items(): + expr = tirx.Let(var, val, expr) + return expr + + return te.compute(x.shape, compute, name="yarn_rope") + + q_embed = op.tensor_expr_op(_rope_fused, "rope", [q, positions]) + k_embed = op.tensor_expr_op(_rope_fused, "rope", [k, positions]) + return q_embed, k_embed + + +class DeepseekV2Attention(nn.Module): + def __init__(self, config: DeepseekV2Config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.num_attention_heads // config.tensor_parallel_shards + + self.rope_theta = config.rope_theta + self.q_lora_rank = config.q_lora_rank + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + self.block_size = config.weight_block_size + + if self.q_lora_rank is None: + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.q_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=config.attention_bias + ) + self.q_a_layernorm = nn.RMSNorm(config.q_lora_rank, -1, config.rms_norm_eps, bias=False) + self.q_b_proj = nn.Linear( + config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False + ) + + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = nn.RMSNorm(config.kv_lora_rank, -1, config.rms_norm_eps, bias=False) + self.kv_b_proj = nn.Linear( + config.kv_lora_rank, + self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), + bias=False, + ) + self.w_uk = nn.Parameter((self.num_heads, config.kv_lora_rank, self.qk_nope_head_dim)) + self.w_uv = nn.Parameter((self.num_heads, self.v_head_dim, config.kv_lora_rank)) + + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=config.attention_bias, + ) + self.softmax_scale = self.q_head_dim ** (-0.5) + if self.config.rope_scaling is not None: + mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = self.config.rope_scaling["factor"] + if mscale_all_dim: + mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * mscale * mscale + self.rotary_emb = DeepseekV2YarnRotaryEmbedding(config) + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + query_positions: Tensor, + forward_mode: Literal["prefill", "decode", "extend"], + ) -> Tuple[Tensor, PagedKVCache]: # noqa: UP006 + b, s, _ = hidden_states.shape + + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj( + self.q_a_layernorm(self.q_a_proj(hidden_states)) + ) # (b, s, num_heads * q_head_dim) + q = op.reshape(q, (b, s, self.num_heads, self.q_head_dim)) # (b, s, num_heads, q_head_dim) + q_nope, q_pe = op.split( + q, [self.qk_nope_head_dim], axis=-1 + ) # (b, s, num_heads, qk_nope_head_dim), (b, s, num_heads, qk_rope_head_dim) + compressed_kv = self.kv_a_proj_with_mqa(hidden_states).reshape( + b, s, 1, self.kv_lora_rank + self.qk_rope_head_dim + ) # (b, s, 1, kv_lora_rank + qk_rope_head_dim) + compressed_kv, k_pe = op.split( + compressed_kv, [self.config.kv_lora_rank], axis=-1 + ) # (b, s, 1, kv_lora_rank), (b, s, 1, qk_rope_head_dim) + compressed_kv = self.kv_a_layernorm(compressed_kv) + q_pe, k_pe = self.rotary_emb(q_pe, k_pe, query_positions) + kv_states = op.concat( + [compressed_kv, k_pe], dim=-1 + ) # (b, s, 1, kv_lora_rank + qk_rope_head_dim) + paged_kv_cache = paged_kv_cache.append_mla_kv(layer_id, kv_states) + + if forward_mode == "prefill": + output, _ = self.self_attn(q_nope, compressed_kv, q_pe, k_pe, paged_kv_cache, layer_id) + elif forward_mode == "decode": + output, _ = self.cross_attn(q_nope, q_pe, paged_kv_cache, layer_id) + elif forward_mode == "extend": + o1, lse1 = self.self_attn(q_nope, compressed_kv, q_pe, k_pe, paged_kv_cache, layer_id) + o2, lse2 = self.cross_attn(q_nope, q_pe, paged_kv_cache, layer_id) + output, _ = paged_kv_cache.merge_attn_output_inplace(o1, lse1, o2, lse2) + else: + raise ValueError(f"Invalid forward mode: {forward_mode}") + + return self.o_proj(output.reshape(b, s, self.num_heads * self.v_head_dim)), paged_kv_cache + + def self_attn( + self, + q_nope: Tensor, + compressed_kv: Tensor, + q_pe: Tensor, + k_pe: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + ) -> Tuple[Tensor, Tensor]: # noqa: UP006 + b, s, _, _ = q_nope.shape + q = op.concat( + [q_nope, q_pe], dim=-1 + ) # (b, s, num_heads, qk_nope_head_dim + qk_rope_head_dim) + kv = op.reshape( + self.kv_b_proj(compressed_kv), + (b, s, self.num_heads, self.qk_nope_head_dim + self.v_head_dim), + ) + k, v = op.split(kv, [self.qk_nope_head_dim], axis=-1) + k_pe = op.broadcast_to(k_pe, (b, s, self.num_heads, self.qk_rope_head_dim)) + k = op.concat([k, k_pe], dim=-1) + output, lse = paged_kv_cache.self_attention(layer_id, q, k, v, self.softmax_scale) + return output, lse + + def cross_attn( + self, + q_nope: Tensor, + q_pe: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + ) -> Tuple[Tensor, Tensor]: # noqa: UP006 + b, s, _, _ = q_nope.shape + if not hasattr(self, "w_uk_scale_inv"): + q_nope = op.matmul( + q_nope.reshape(b * s, self.num_heads, self.qk_nope_head_dim).permute_dims(1, 0, 2), + self.w_uk.permute_dims(0, 2, 1), + ) + else: + q_nope = batch_matmul.quantized_bmm( + q_nope.reshape(b * s, self.num_heads, self.qk_nope_head_dim).permute_dims(1, 0, 2), + self.w_uk, + self.w_uk_scale_inv, + self.block_size, + ) + q_nope = q_nope.permute_dims(1, 0, 2).reshape( + b, s, self.num_heads, self.kv_lora_rank + ) # (b, s, num_heads, kv_lora_rank) + query_states = op.concat( + [q_nope, q_pe], dim=-1 + ) # (b, s, num_heads, kv_lora_rank + qk_rope_head_dim) + + output, lse = paged_kv_cache.cross_attention( + layer_id, + query_states, + v_head_dim=self.kv_lora_rank, + sm_scale=self.softmax_scale, + ) # (b, s, num_heads, kv_lora_rank) + if getattr(self, "w_uv_scale_inv", None) is None: + output = op.matmul( + output.reshape(b * s, self.num_heads, self.kv_lora_rank).permute_dims(1, 0, 2), + self.w_uv.permute_dims(0, 2, 1), + ) + else: + output = batch_matmul.quantized_bmm( + output.reshape(b * s, self.num_heads, self.kv_lora_rank).permute_dims(1, 0, 2), + self.w_uv, + self.w_uv_scale_inv, + self.block_size, + ) + output = output.permute_dims(1, 0, 2).reshape(b, s, self.num_heads * self.v_head_dim) + return output, lse + + +class DeepseekV2MoE(nn.Module): + def __init__(self, config: DeepseekV2Config): + super().__init__() + self.num_experts_per_tok = config.num_experts_per_tok + self.num_routed_experts = config.n_routed_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.scoring_func = config.scoring_func + self.topk_method = config.topk_method + self.n_group = config.n_group + self.topk_group = config.topk_group + + self.gate = nn.Linear( + config.hidden_size, self.num_routed_experts, bias=False, out_dtype="float32" + ) + self.e_score_correction_bias = ( + nn.Parameter((config.n_routed_experts,), dtype="float32") + if config.topk_method == "noaux_tc" + else None + ) + self.norm_topk_prob = config.norm_topk_prob + if config.moe_intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MoE intermediate size {config.moe_intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.moe_intermediate_size = config.moe_intermediate_size // config.tensor_parallel_shards + + self.moe_gate_up_proj = MixtralExperts( + self.num_routed_experts, + in_features=config.hidden_size, + out_features=2 * self.moe_intermediate_size, + ) + self.moe_down_proj = MixtralExperts( + self.num_routed_experts, + in_features=self.moe_intermediate_size, + out_features=config.hidden_size, + ) + + self.shared_experts = DeepseekV2MLP( + config, + intermediate_size=config.moe_intermediate_size * config.n_shared_experts, + ) + self.dtype = "float32" + + def forward(self, x: Tensor): + def _expert_forward(x: Tensor, indptr: Tensor): + x1_x2 = self.moe_gate_up_proj(x, indptr) + x1, x2 = op.split(x1_x2, indices_or_sections=2, axis=-1) + x = self.moe_down_proj(op.silu(x1) * x2, indptr) + return x + + experts_per_tok = self.num_experts_per_tok + num_experts = self.num_routed_experts + b, s, h = x.shape + num_tokens = b * s + x = op.reshape(x, (num_tokens, h)) + logits = self.gate(x) # (num_tokens, num_routed_experts) + assert logits.dtype == "float32" + if self.scoring_func == "softmax": + scores = op.softmax(logits, axis=-1) + elif self.scoring_func == "sigmoid": + scores = op.sigmoid(logits) + else: + raise ValueError(f"Unsupported deepseek scoring function: {self.scoring_func}") + + # select top-k experts + if self.topk_method == "greedy": + expert_weights, expert_indices = op_ext.moe_misc.gating_topk(scores, experts_per_tok) + elif self.topk_method in ["group_limited_greedy", "noaux_tc"]: + expert_weights, expert_indices = op_ext.moe_misc.group_limited_greedy_topk( + scores, + self.num_experts_per_tok, + self.num_routed_experts, + self.n_group, + self.topk_group, + self.topk_method, + num_tokens, + self.e_score_correction_bias, + ) + else: + raise ValueError(f"Unsupported deepseek topk method: {self.topk_method}") + + if self.num_experts_per_tok > 1 and self.norm_topk_prob: + denominator = op.sum(expert_weights, axis=-1, keepdims=True) + 1e-20 + expert_weights = expert_weights / denominator + expert_weights = expert_weights * self.routed_scaling_factor + + use_cutlass = op_ext.get_store().cutlass_group_gemm and self.dtype in [ + "float16", + "bfloat16", + ] + + if num_tokens == 1: + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = _expert_forward(x, expert_indices) + else: + # cumsum: [num_tokens * local_experts] + cumsum = op_ext.moe_misc.moe_cumsum(expert_indices, num_experts) + # indices: [num_tokens * experts_per_tok] + reverse_indices, token_indices = op_ext.moe_misc.get_indices(cumsum, expert_indices) + if use_cutlass: + # indptr: [num_routed_experts] + indptr = op_ext.moe_misc.get_indptr( + cumsum, num_experts, num_tokens, inclusive=True, out_dtype="int64" + ) + else: + # indptr: [num_routed_experts + 1] + indptr = op_ext.moe_misc.get_indptr( + cumsum, num_experts, num_tokens, inclusive=False, out_dtype="int32" + ) + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = op.take(x, token_indices, axis=0) + moe_hidden_states = _expert_forward(moe_hidden_states, indptr) + moe_hidden_states = op_ext.moe_misc.scatter_output(moe_hidden_states, reverse_indices) + + # moe_hidden_states: [num_tokens, experts_per_tok, hidden_size] + expert_weights = expert_weights.reshape(num_tokens, experts_per_tok, 1).astype(x.dtype) + moe_hidden_states = ( + moe_hidden_states.reshape(num_tokens, experts_per_tok, h) * expert_weights + ) + # moe_hidden_states: [num_tokens, hidden_size] + moe_hidden_states = op_ext.moe_misc.moe_sum(moe_hidden_states, dim=1) + + shared_expert_hidden_states = self.shared_experts(x) + + final_hidden_states = moe_hidden_states + shared_expert_hidden_states + final_hidden_states = op.reshape(final_hidden_states, (b, s, h)) + return final_hidden_states + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + # Force e_score_correction_bias to be float32 + if self.e_score_correction_bias is not None: + self.e_score_correction_bias.to("float32") + + +class DeepseekV2DecoderLayer(nn.Module): + def __init__(self, config: DeepseekV2Config, layer_idx: int): + super().__init__() + self.self_attn = DeepseekV2Attention(config) + self.mlp = ( + DeepseekV2MoE(config) + if ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + else DeepseekV2MLP(config) + ) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + if self.self_attn.q_lora_rank is None: + _set( + self.self_attn.q_proj.weight, + tp.ShardSingleDim("_shard_q_weight", dim=0), + ) + else: + _set( + self.self_attn.q_b_proj.weight, + tp.ShardSingleDim("_shard_q_b_weight", dim=0), + ) + + _set( + self.self_attn.kv_b_proj.weight, + tp.ShardSingleDim("_shard_kv_b_weight", dim=0), + ) + _set( + self.self_attn.w_uk, + tp.ShardSingleDim("_shard_kv_b_weight_w_uk", dim=0), + ) + _set( + self.self_attn.w_uv, + tp.ShardSingleDim("_shard_kv_b_weight_w_uv", dim=0), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + + if isinstance(self.mlp, DeepseekV2MoE): + si = self.mlp.shared_experts.intermediate_size + mi = self.mlp.moe_intermediate_size + _set( + self.mlp.shared_experts.gate_up_proj.weight, + tp.ShardSingleDim("_shard_shared_experts_gate_up", segs=[si, si], dim=0), + ) + _set( + self.mlp.shared_experts.down_proj.weight, + tp.ShardSingleDim("_shard_shared_experts_down", dim=1), + ) + _set( + self.mlp.moe_gate_up_proj.weight, + tp.ShardSingleDim("_shard_moe_gate_up", segs=[mi, mi], dim=1), + ) + _set( + self.mlp.moe_down_proj.weight, + tp.ShardSingleDim("_shard_moe_mlp_down", dim=2), + ) + else: + assert isinstance(self.mlp, DeepseekV2MLP) + si = self.mlp.intermediate_size + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_gate_up", segs=[si, si], dim=0), + ) + _set( + self.mlp.down_proj.weight, + tp.ShardSingleDim("_shard_down", dim=1), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + query_positions: Tensor, + forward_mode: Literal["prefill", "decode", "extend"], + ) -> Tuple[Tensor, PagedKVCache]: # noqa: UP006 + out = self.input_layernorm(hidden_states) + out, paged_kv_cache = self.self_attn( + out, paged_kv_cache, layer_id, query_positions, forward_mode + ) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states, paged_kv_cache + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class DeepseekV2Model(nn.Module): + def __init__(self, config: DeepseekV2Config): + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [ + DeepseekV2DecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward( + self, + inputs: Tensor, + paged_kv_cache: PagedKVCache, + forward_mode: Literal["prefill", "decode", "extend"], + ): + hidden_states = inputs + query_positions = paged_kv_cache.get_query_positions(inputs.shape[0] * inputs.shape[1]) + for layer_id, layer in enumerate(self.layers): + hidden_states, paged_kv_cache = layer( + hidden_states, paged_kv_cache, layer_id, query_positions, forward_mode + ) + hidden_states = self.norm(hidden_states) + return hidden_states, paged_kv_cache + + +class DeepseekV2ForCausalLM(nn.Module): + def __init__(self, config: DeepseekV2Config): + self.model = DeepseekV2Model(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.dtype = config.dtype + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.intermediate_size = config.intermediate_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.rms_norm_eps = config.rms_norm_eps + self.rope_theta = config.rope_theta + self.vocab_size = config.vocab_size + self.tensor_parallel_shards = config.tensor_parallel_shards + self.weight_block_size = config.weight_block_size + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + forward_mode: Literal["prefill", "decode", "extend"], + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states, paged_kv_cache = self.model(input_embeds, paged_kv_cache, forward_mode) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states, paged_kv_cache = self.model(input_embed, paged_kv_cache, "prefill") + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def extend(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states, paged_kv_cache = self.model(input_embed, paged_kv_cache, "extend") + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states, paged_kv_cache = self.model(input_embed, paged_kv_cache, "decode") + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits, paged_kv_cache = self.batch_forward( + input_embeds, paged_kv_cache, "prefill", logit_positions + ) + return logits, paged_kv_cache + + def batch_extend( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits, paged_kv_cache = self.batch_forward( + input_embeds, paged_kv_cache, "extend", logit_positions + ) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits, paged_kv_cache = self.batch_forward(input_embeds, paged_kv_cache, "decode", None) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits, paged_kv_cache = self.batch_forward(input_embeds, paged_kv_cache, "extend", None) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mla", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=1, + qk_head_dim=self.kv_lora_rank + self.qk_rope_head_dim, + v_head_dim=self.kv_lora_rank, + mla_original_qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim, + mla_original_v_head_dim=self.v_head_dim, + rope_mode=RopeMode.NONE, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "extend": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_extend": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/eagle/__init__.py b/python/mlc_llm/model/eagle/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/eagle/eagle_loader.py b/python/mlc_llm/model/eagle/eagle_loader.py new file mode 100644 index 0000000..96634b7 --- /dev/null +++ b/python/mlc_llm/model/eagle/eagle_loader.py @@ -0,0 +1,105 @@ +""" +This file specifies how MLC's EAGLE parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization, make_awq_quant + +from .eagle_model import EagleConfig, EagleForCausalLM + +awq_quant = make_awq_quant(EagleForCausalLM) + + +huggingface = make_standard_hf_loader( + model_cls=EagleForCausalLM, + layer_prefix="layers", + add_unused=["rotary_emb.inv_freq"], +) + + +def awq(model_config: EagleConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of AWQ parameters. + Parameters + ---------- + model_config : EagleConfig + The configuration of the Eagle model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to AWQ. + """ + model, _ = awq_quant(model_config, quantization) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"layers.{i}.self_attn" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{attn}.qkv_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{quantize_suffix}", + f"{attn}.k_proj.{quantize_suffix}", + f"{attn}.v_proj.{quantize_suffix}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate( + [q, k, v], + axis=1, # AWQ GEMM would transpose the weight + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Concat gate and up in MLP + mlp = f"layers.{i}.mlp" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{mlp}.gate_up_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.gate_proj.{quantize_suffix}", + f"{mlp}.up_proj.{quantize_suffix}", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate( + [gate, up], + axis=1, # AWQ GEMM would transpose the weight + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # inv_freq is not used in the model + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + return mapping diff --git a/python/mlc_llm/model/eagle/eagle_model.py b/python/mlc_llm/model/eagle/eagle_model.py new file mode 100644 index 0000000..d7722de --- /dev/null +++ b/python/mlc_llm/model/eagle/eagle_model.py @@ -0,0 +1,251 @@ +""" +Implementation for EAGLE architecture. +""" + +import dataclasses +from typing import Optional + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.llama.llama_model import LlamaAttention, LlamaConfig, LlamaFFN +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class EagleConfig(LlamaConfig): + """Configuration of the Eagle model.""" + + bias: bool = True # Whether to use bias in the fc layers + + +class EagleDecoderLayer(nn.Module): + def __init__(self, config: EagleConfig, index: int): + rms_norm_eps = config.rms_norm_eps + self.self_attn = LlamaAttention(config) + self.mlp = LlamaFFN(config) + self.index = index + if self.index != 0: + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + if self.index != 0: + hidden_states = self.input_layernorm(hidden_states) + out = self.self_attn(hidden_states, paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class EagleForCausalLM(nn.Module): + def __init__(self, config: EagleConfig): + # Put the model definition here to align with EAGLE's original structure + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = nn.Embedding("vocab_size", config.hidden_size) + self.layers = nn.ModuleList( + [EagleDecoderLayer(config, i) for i in range(config.num_hidden_layers)] + ) + self.fc = nn.Linear( + in_features=2 * config.hidden_size, + out_features=config.hidden_size, + bias=config.bias, + ) + + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def fuse_embed_hidden_states(self, input_embed: Tensor, hidden_states: Tensor): + hidden_states = op.concat([input_embed, hidden_states], dim=-1) + hidden_states = self.fc(hidden_states) + return hidden_states + + def forward_to_last_hidden_states(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache): + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + return hidden_states + + def forward(self, input_embed: Tensor, hidden_states: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = self.fuse_embed_hidden_states(input_embed, hidden_states) + hidden_states = self.forward_to_last_hidden_states(hidden_states, paged_kv_cache) + return hidden_states + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.forward_to_last_hidden_states(hidden_states, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return hidden_states + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.embed_tokens(input_ids) + + def prefill_to_last_hidden_states(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.forward_to_last_hidden_states(hidden_states, paged_kv_cache) + return hidden_states, paged_kv_cache + + def decode_to_last_hidden_states(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.forward_to_last_hidden_states(hidden_states, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_prefill_to_last_hidden_states( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + ): + hidden_states = self.batch_forward(hidden_states, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_decode_to_last_hidden_states( + self, hidden_states: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward(hidden_states, paged_kv_cache) + return hidden_states, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "fuse_embed_hidden_states": { + "input_embed": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill_to_last_hidden_states": { + "hidden_states": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode_to_last_hidden_states": { + "hidden_states": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill_to_last_hidden_states": { + "hidden_states": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode_to_last_hidden_states": { + "hidden_states": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/gemma/__init__.py b/python/mlc_llm/model/gemma/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/gemma/gemma_loader.py b/python/mlc_llm/model/gemma/gemma_loader.py new file mode 100644 index 0000000..18d200f --- /dev/null +++ b/python/mlc_llm/model/gemma/gemma_loader.py @@ -0,0 +1,46 @@ +""" +This file specifies how MLC's Gemma parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization + +from .gemma_model import GemmaConfig, GemmaForCausalLM + + +def huggingface(model_config: GemmaConfig, quantization: Quantization) -> ExternMapping: + """Create HF weight mapping for Gemma.""" + model = GemmaForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + base_loader = make_standard_hf_loader( + model_cls=GemmaForCausalLM, + ) + mapping = base_loader(model_config, quantization) + + def add_one(name: str) -> None: + mlc_param = named_parameters[name] + mapping.add_mapping( + name, + [name], + functools.partial( + lambda x, dtype: (x + 1).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for i in range(model_config.num_hidden_layers): + add_one(f"model.layers.{i}.input_layernorm.weight") + add_one(f"model.layers.{i}.post_attention_layernorm.weight") + + add_one("model.norm.weight") + return mapping diff --git a/python/mlc_llm/model/gemma/gemma_model.py b/python/mlc_llm/model/gemma/gemma_model.py new file mode 100644 index 0000000..20c2e8d --- /dev/null +++ b/python/mlc_llm/model/gemma/gemma_model.py @@ -0,0 +1,391 @@ +"""Implementation for Gemma architecture.""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class GemmaConfig(ConfigBase): + """Configuration of the Gemma model.""" + + hidden_size: int + intermediate_size: int + attention_bias: bool + num_attention_heads: int + num_key_value_heads: int + head_dim: int + num_hidden_layers: int + rms_norm_eps: float + vocab_size: int + hidden_activation: Optional[str] = None + position_embedding_base: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.hidden_activation is None: + self.hidden_activation = self.kwargs.get("hidden_act", None) + if self.hidden_activation not in ("gelu", "gelu_pytorch_tanh"): + raise ValueError("Only GeLU is supported as the activation for gemma.") + if self.attention_bias: + raise ValueError('Only "False" attention_bias is supported for gemma') + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + assert self.num_attention_heads % self.num_key_value_heads == 0 + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class GemmaEmbedding(nn.Embedding): + """The embedding module specialized for Gemma so that + it can be shared with the final lm_head. + """ + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class GemmaMLP(nn.Module): + def __init__(self, config: GemmaConfig): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.gelu(x1, approximate="tanh") * x2) + + +class GemmaAttention(nn.Module): + def __init__(self, config: GemmaConfig): + self.head_dim = config.head_dim + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({config.num_key_value_heads}) must be divisible by tensor_parallel_shards" # noqa: E501 + ) + assert config.num_key_value_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than {config.num_key_value_heads}" + ) + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + in_features=self.num_q_heads * self.head_dim, + out_features=config.hidden_size, + bias=config.attention_bias, + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class GemmaDecoderLayer(nn.Module): + def __init__(self, config: GemmaConfig): + rms_norm_eps = config.rms_norm_eps + self.self_attn = GemmaAttention(config) + self.mlp = GemmaMLP(config) + # Gemma RMSNorm adds 1 to the weights. It is already fused in the loader + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class GemmaModel(nn.Module): + def __init__(self, config: GemmaConfig): + self.hidden_size = config.hidden_size + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = GemmaEmbedding("vocab_size", config.hidden_size) + self.layers = nn.ModuleList( + [GemmaDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + hidden_states = hidden_states * (self.hidden_size**0.5) + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class GemmaForCausalLM(nn.Module): + def __init__(self, config: GemmaConfig): + self.model = GemmaModel(config) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def get_logits(self, hidden_states: Tensor): + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.get_logits(hidden_states) + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/gemma2/__init__.py b/python/mlc_llm/model/gemma2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/gemma2/gemma2_loader.py b/python/mlc_llm/model/gemma2/gemma2_loader.py new file mode 100644 index 0000000..7e59a8a --- /dev/null +++ b/python/mlc_llm/model/gemma2/gemma2_loader.py @@ -0,0 +1,48 @@ +""" +This file specifies how MLC's Gemma2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization + +from .gemma2_model import Gemma2Config, Gemma2ForCausalLM + + +def huggingface(model_config: Gemma2Config, quantization: Quantization) -> ExternMapping: + """Create HF weight mapping for Gemma2.""" + model = Gemma2ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + base_loader = make_standard_hf_loader( + model_cls=Gemma2ForCausalLM, + ) + mapping = base_loader(model_config, quantization) + + def add_one(name: str) -> None: + mlc_param = named_parameters[name] + mapping.add_mapping( + name, + [name], + functools.partial( + lambda x, dtype: (x + 1).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for i in range(model_config.num_hidden_layers): + add_one(f"model.layers.{i}.input_layernorm.weight") + add_one(f"model.layers.{i}.post_attention_layernorm.weight") + add_one(f"model.layers.{i}.pre_feedforward_layernorm.weight") + add_one(f"model.layers.{i}.post_feedforward_layernorm.weight") + + add_one("model.norm.weight") + return mapping diff --git a/python/mlc_llm/model/gemma2/gemma2_model.py b/python/mlc_llm/model/gemma2/gemma2_model.py new file mode 100644 index 0000000..ce3bf59 --- /dev/null +++ b/python/mlc_llm/model/gemma2/gemma2_model.py @@ -0,0 +1,122 @@ +"""Implementation for Gemma2 architecture.""" + +import dataclasses + +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm.model.gemma.gemma_model import ( + GemmaAttention, + GemmaConfig, + GemmaForCausalLM, + GemmaMLP, + GemmaModel, +) +from mlc_llm.nn import PagedKVCache +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Gemma2Config(GemmaConfig): + """Configuration of the Gemma2 model, in addition to the Gemma model""" + + # NOTE: We ignore attn_logit_softcapping in the gemma2 implementation for now. + # The Gemma 2 team observed minor differences when soft-capping is removed during inference, + # according to https://huggingface.co/blog/gemma2. + # The soft-capping is also not supported by HuggingFace transformers `Gemma2SdpaAttention`. + attn_logit_softcapping: float = None + final_logit_softcapping: float = None + query_pre_attn_scalar: int = None + sliding_window: int = None + + def __post_init__(self): + super().__post_init__() + # NOTE: override the context window size with the Gemma2 sliding window size, + # as the sliding window attention every other layer is yet to be supported. + self.context_window_size = self.sliding_window + + +class Gemma2Attention(GemmaAttention): + def __init__(self, config: Gemma2Config): + super().__init__(config) + self.scaling_factor = (config.head_dim / config.query_pre_attn_scalar) ** 0.5 + + +class Gemma2DecoderLayer(nn.Module): + def __init__(self, config: Gemma2Config): + rms_norm_eps = config.rms_norm_eps + self.self_attn = Gemma2Attention(config) + self.mlp = GemmaMLP(config) + # Gemma RMSNorm adds 1 to the weights. It is already fused in the loader + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.pre_feedforward_layernorm = nn.RMSNorm( + config.hidden_size, -1, rms_norm_eps, bias=False + ) + self.post_feedforward_layernorm = nn.RMSNorm( + config.hidden_size, -1, rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + out = self._apply_post_matmul_norm(out, norm=self.post_attention_layernorm) + hidden_states = out + hidden_states + + out = self.pre_feedforward_layernorm(hidden_states) + out = self.mlp(out) + out = self._apply_post_matmul_norm(out, norm=self.post_feedforward_layernorm) + hidden_states = out + hidden_states + + return hidden_states + + def _apply_post_matmul_norm(self, out: Tensor, norm: nn.Tensor): + if self.tensor_parallel_shards > 1: + return norm(op.ccl_allreduce(out, "sum")) + return norm(out) + + +class Gemma2Model(GemmaModel): + def __init__(self, config: Gemma2Config): + super().__init__(config) + self.layers = nn.ModuleList( + [Gemma2DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + + +class Gemma2ForCausalLM(GemmaForCausalLM): + def __init__(self, config: Gemma2Config): + super().__init__(config) + self.model = Gemma2Model(config) + self.final_logit_softcapping = config.final_logit_softcapping + + def get_logits(self, hidden_states: Tensor): + logits = super().get_logits(hidden_states) + if self.final_logit_softcapping is not None: + logits = op.tanh(logits / self.final_logit_softcapping) * self.final_logit_softcapping + return logits diff --git a/python/mlc_llm/model/gemma3/__init__.py b/python/mlc_llm/model/gemma3/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/gemma3/gemma3_loader.py b/python/mlc_llm/model/gemma3/gemma3_loader.py new file mode 100644 index 0000000..6001e75 --- /dev/null +++ b/python/mlc_llm/model/gemma3/gemma3_loader.py @@ -0,0 +1,71 @@ +""" +This file specifies how MLC's Gemma3 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization + +from .gemma3_model import Gemma3Config, Gemma3ForCausalLM + + +def huggingface(model_config: Gemma3Config, quantization: Quantization) -> ExternMapping: + """Create HF weight mapping for Gemma3.""" + model = Gemma3ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + mlc_prefix = "language_model." + if model_config.is_text_model: + hf_prefix = "" + else: + hf_prefix = "language_model." + + def name_transform(name: str) -> str: + if name.startswith(mlc_prefix): + name = name[len(mlc_prefix) :] + return f"{hf_prefix}{name}" + + def num_layers(config: object) -> int: + return config.text_config.num_hidden_layers + + base_loader = make_standard_hf_loader( + model_cls=Gemma3ForCausalLM, + include_qkv=False, + include_gate_up=True, + gate_up_target_name="gate_up_proj", + num_layers_getter=num_layers, + layer_prefix=f"{mlc_prefix}model.layers", + name_transform=name_transform, + ) + mapping = base_loader(model_config, quantization) + + def add_one(name: str) -> None: + mlc_param = named_parameters[mlc_prefix + name] + mapping.add_mapping( + mlc_prefix + name, + [name_transform(mlc_prefix + name)], + functools.partial( + lambda x, dtype: (x + 1).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for i in range(model_config.text_config.num_hidden_layers): + add_one(f"model.layers.{i}.input_layernorm.weight") + add_one(f"model.layers.{i}.post_attention_layernorm.weight") + add_one(f"model.layers.{i}.pre_feedforward_layernorm.weight") + add_one(f"model.layers.{i}.post_feedforward_layernorm.weight") + add_one(f"model.layers.{i}.self_attn.k_norm.weight") + add_one(f"model.layers.{i}.self_attn.q_norm.weight") + + add_one("model.norm.weight") + + return mapping diff --git a/python/mlc_llm/model/gemma3/gemma3_model.py b/python/mlc_llm/model/gemma3/gemma3_model.py new file mode 100644 index 0000000..b022bab --- /dev/null +++ b/python/mlc_llm/model/gemma3/gemma3_model.py @@ -0,0 +1,667 @@ +"""Implementation for Gemma3 architecture.""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.gemma.gemma_model import GemmaEmbedding +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Gemma3TextConfig(ConfigBase): + """Configuration of the text model inside Gemma3""" + + # NOTE More fields have defaults due to Huggingface Gemma3 configs missing fields + # The defaults for these fields can be found in the transformers library + hidden_size: int + intermediate_size: int + num_hidden_layers: int + attention_bias: bool = False + num_attention_heads: int = 8 + num_key_value_heads: int = 4 + head_dim: int = 256 + rms_norm_eps: float = 1e-6 + hidden_activation: Optional[str] = "gelu_pytorch_tanh" + position_embedding_base: int = 1_000_000 + rope_scaling: int = 0 + context_window_size: int = 131_072 + prefill_chunk_size: int = 0 + + query_pre_attn_scalar: int = 256 + sliding_window_size: int = None + sliding_window_pattern = 6 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.hidden_activation is None: + self.hidden_activation = self.kwargs.get("hidden_act", None) + if self.sliding_window_size is None: + self.sliding_window_size = self.kwargs.get("sliding_window", None) + if self.hidden_activation not in ("gelu", "gelu_pytorch_tanh"): + raise ValueError("Only GeLU is supported as the activation for gemma.") + if self.attention_bias: + raise ValueError('Only "False" attention_bias is supported for gemma') + if self.position_embedding_base == 1000000 and "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + assert self.num_attention_heads % self.num_key_value_heads == 0 + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + # NOTE: override the context window size with the Gemma2 sliding window size, + # as the sliding window attention every other layer is yet to be supported. + self.context_window_size = max(self.sliding_window_size, 8192) + + +@dataclasses.dataclass +class Gemma3Config(ConfigBase): + """Configuration of the Gemma3 model""" + + text_config: Gemma3TextConfig = None + vocab_size: int = 262_208 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + context_window_size: int = -1 + sliding_window_size: int = -1 + prefill_chunk_size: int = -1 + is_text_model: bool = False + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.text_config is None: + self.is_text_model = True + self.text_config = Gemma3TextConfig.from_dict(self.kwargs) + + text_config_dict: Dict[str, Any] # noqa: UP006 + if isinstance(self.text_config, Gemma3TextConfig): + text_config_dict = dataclasses.asdict(self.text_config) + else: + text_config_dict = dict(self.text_config) + + for k, v in text_config_dict.pop("kwargs", {}).items(): + text_config_dict[k] = v + + self.text_config = Gemma3TextConfig.from_dict(text_config_dict) + + for k in ["context_window_size", "prefill_chunk_size", "sliding_window_size"]: + if getattr(self, k) <= 0: + if hasattr(self.text_config, k): + setattr(self, k, getattr(self.text_config, k)) + + +class Gemma3MLP(nn.Module): + def __init__(self, config: Gemma3Config): + super().__init__() + if config.text_config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.text_config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = ( + config.text_config.intermediate_size // config.tensor_parallel_shards + ) + self.gate_up_proj = nn.Linear( + in_features=config.text_config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear( + self.intermediate_size, config.text_config.hidden_size, bias=False + ) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.gelu(x1, approximate="tanh") * x2) + + +class Gemma3Attention(nn.Module): + def __init__(self, config: Gemma3Config): + self.head_dim = config.text_config.head_dim + self.num_q_heads = config.text_config.num_attention_heads // config.tensor_parallel_shards + self.num_kv_heads = config.text_config.num_key_value_heads + assert self.num_kv_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({self.num_kv_heads}) must be divisible by tensor_parallel_shards" + ) + assert self.num_kv_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than {self.num_kv_heads}" + ) + self.num_kv_heads = self.num_kv_heads // config.tensor_parallel_shards + self.q_proj = nn.Linear( + in_features=config.text_config.hidden_size, + out_features=self.num_q_heads * self.head_dim, + bias=config.text_config.attention_bias, + ) + self.k_proj = nn.Linear( + in_features=config.text_config.hidden_size, + out_features=self.num_kv_heads * self.head_dim, + bias=config.text_config.attention_bias, + ) + self.v_proj = nn.Linear( + in_features=config.text_config.hidden_size, + out_features=self.num_kv_heads * self.head_dim, + bias=config.text_config.attention_bias, + ) + self.o_proj = nn.Linear( + in_features=self.num_q_heads * self.head_dim, + out_features=config.text_config.hidden_size, + bias=config.text_config.attention_bias, + ) + self.q_norm = nn.RMSNorm( + config.text_config.head_dim, -1, config.text_config.rms_norm_eps, bias=False + ) + self.k_norm = nn.RMSNorm( + config.text_config.head_dim, -1, config.text_config.rms_norm_eps, bias=False + ) + # self.scaling_factor = (self.head_dim / config.text_config.query_pre_attn_scalar) ** 0.5 + self.scaling = config.text_config.query_pre_attn_scalar**-0.5 + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q = self.head_dim, self.num_q_heads + b, s, _ = hidden_states.shape + # QKV Projection + q_proj = op.reshape(self.q_proj(hidden_states), (b, s, -1, d)) + k_proj = op.reshape(self.k_proj(hidden_states), (b, s, -1, d)) + v_proj = op.reshape(self.v_proj(hidden_states), (b, s, -1, d)) + + q_norm = self.q_norm(q_proj) + k_norm = self.k_norm(k_proj) + + qkv = op.concat([q_norm, k_norm, v_proj], dim=2) + + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.scaling + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class Gemma3DecoderLayer(nn.Module): + def __init__(self, config: Gemma3Config): + rms_norm_eps = config.text_config.rms_norm_eps + self.self_attn = Gemma3Attention(config) + self.mlp = Gemma3MLP(config) + # Gemma RMSNorm adds 1 to the weights. It is already fused in the loader + self.input_layernorm = nn.RMSNorm( + config.text_config.hidden_size, -1, rms_norm_eps, bias=False + ) + self.post_attention_layernorm = nn.RMSNorm( + config.text_config.hidden_size, -1, rms_norm_eps, bias=False + ) + self.pre_feedforward_layernorm = nn.RMSNorm( + config.text_config.hidden_size, -1, rms_norm_eps, bias=False + ) + self.post_feedforward_layernorm = nn.RMSNorm( + config.text_config.hidden_size, -1, rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + i = self.mlp.intermediate_size + _set(self.self_attn.q_proj, tp.ShardSingleDim("_shard_q", dim=0)) + _set(self.self_attn.k_proj, tp.ShardSingleDim("_shard_k", dim=0)) + _set(self.self_attn.v_proj, tp.ShardSingleDim("_shard_v", dim=0)) + _set(self.self_attn.q_norm, tp.ShardSingleDim("_shard_q_norm", dim=0)) + _set(self.self_attn.k_norm, tp.ShardSingleDim("_shard_k_norm", dim=0)) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + out = self._apply_post_matmul_norm(out, norm=self.post_attention_layernorm) + hidden_states = out + hidden_states + + out = self.pre_feedforward_layernorm(hidden_states) + out = self.mlp(out) + out = self._apply_post_matmul_norm(out, norm=self.post_feedforward_layernorm) + hidden_states = out + hidden_states + + return hidden_states + + def _apply_post_matmul_norm(self, out: Tensor, norm: nn.Tensor): + if self.tensor_parallel_shards > 1: + return norm(op.ccl_allreduce(out, "sum")) + return norm(out) + + +class Gemma3TextModel(nn.Module): + def __init__(self, config: Gemma3Config): + self.hidden_size = config.text_config.hidden_size + assert config.text_config.hidden_size % config.text_config.num_attention_heads == 0 + self.embed_tokens = GemmaEmbedding("vocab_size", config.text_config.hidden_size) + self.layers = nn.ModuleList( + [Gemma3DecoderLayer(config) for _ in range(config.text_config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm( + config.text_config.hidden_size, + -1, + config.text_config.rms_norm_eps, + bias=False, + ) + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + hidden_states = hidden_states * (self.hidden_size**0.5) + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Gemma3LanguageModel(nn.Module): + def __init__(self, config: Gemma3Config): + self.model = Gemma3TextModel(config) + self.config = config + self.num_hidden_layers = config.text_config.num_hidden_layers + self.num_attention_heads = config.text_config.num_attention_heads + self.num_key_value_heads = config.text_config.num_key_value_heads + self.head_dim = config.text_config.head_dim + self.hidden_size = config.text_config.hidden_size + self.vocab_size = config.vocab_size + self.rope_theta = config.text_config.position_embedding_base + self.rope_scaling = config.text_config.rope_scaling + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def get_logits(self, hidden_states: Tensor): + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.get_logits(hidden_states) + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + # if "factor" in self.rope_scaling: + # rope_scaling = self.rope_scaling["factor"] + # else: + # rope_scaling = 1 + return PagedKVCache.create_generic( + attn_kind=[ + ( + "mha_sliding" + if ((i + 1) % self.config.text_config.sliding_window_pattern) + else "mha" + ) + for i in range(self.num_hidden_layers) + ], + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) + + +class Gemma3ForCausalLM(nn.Module): + def __init__(self, config: Gemma3Config): + super().__init__() + self.config = config + self.language_model = Gemma3LanguageModel(config) + self.vocab_size = config.vocab_size + self.dtype = "float32" + self.tensor_parallel_shards = config.tensor_parallel_shards + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + self.language_model.to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def get_logits(self, hidden_states: Tensor): + logits = self.language_model.model.embed_tokens.lm_head_forward(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.language_model.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.get_logits(hidden_states) + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.language_model.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.language_model.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.language_model.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + # if "factor" in self.language_model.rope_scaling: + # rope_scaling = self.language_model.rope_scaling["factor"] + # else: + # rope_scaling = 1 + return PagedKVCache.create_generic( + attn_kind=[ + ( + "mha_sliding" + if ((i + 1) % self.config.text_config.sliding_window_pattern) + else "mha" + ) + for i in range(self.language_model.num_hidden_layers) + ], + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.language_model.num_hidden_layers, + num_attention_heads=self.language_model.num_attention_heads + // self.tensor_parallel_shards, + num_key_value_heads=self.language_model.num_key_value_heads + // self.tensor_parallel_shards, + qk_head_dim=self.language_model.head_dim, + v_head_dim=self.language_model.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.language_model.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor( + [1, "seq_len", self.language_model.hidden_size], self.dtype + ), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.language_model.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor( + [1, "seq_len", self.language_model.hidden_size], self.dtype + ), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor( + ["batch_size", 1, self.language_model.hidden_size], self.dtype + ), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor( + [1, "seq_len", self.language_model.hidden_size], self.dtype + ), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/gpt2/__init__.py b/python/mlc_llm/model/gpt2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/gpt2/gpt2_loader.py b/python/mlc_llm/model/gpt2/gpt2_loader.py new file mode 100644 index 0000000..c91f63d --- /dev/null +++ b/python/mlc_llm/model/gpt2/gpt2_loader.py @@ -0,0 +1,85 @@ +""" +This file specifies how MLC's GPT-2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .gpt2_model import GPT2Config, GPT2LMHeadModel + + +def huggingface(model_config: GPT2Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : GPT2Config + The configuration of the GPT-2 model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = GPT2LMHeadModel(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + mapping.add_mapping( + "lm_head.weight", + ["wte.weight"], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=named_parameters["transformer.wte.weight"].dtype, + ), + ) + + for i in range(model_config.n_layer): + mapping.add_unused(f"h.{i}.attn.bias") + + # Transpose c_attn, c_proj and c_fc weights since GPT-2 uses Conv1D + for conv1d_weight_name in [ + "attn.c_attn", + "attn.c_proj", + "mlp.c_proj", + "mlp.c_fc", + ]: + src_name = f"h.{i}.{conv1d_weight_name}.weight" + mlc_name = f"transformer.{src_name}" + mapping.add_mapping( + mlc_name, + [src_name], + functools.partial( + lambda x, dtype: x.transpose().astype(dtype), + dtype=named_parameters[mlc_name].dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + # transformer.h.0.attn.c_attn.weight --> h.0.attn.c_attn.weight + source_name = mlc_name.split(".", 1)[1] + mapping.add_mapping( + mlc_name, + [source_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + return mapping diff --git a/python/mlc_llm/model/gpt2/gpt2_model.py b/python/mlc_llm/model/gpt2/gpt2_model.py new file mode 100644 index 0000000..e5db920 --- /dev/null +++ b/python/mlc_llm/model/gpt2/gpt2_model.py @@ -0,0 +1,382 @@ +""" +Implementation for GPT-2 architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class GPT2Config(ConfigBase): + """Configuration of the GPT-2 model.""" + + vocab_size: int + n_embd: int + n_layer: int + n_head: int + layer_norm_epsilon: float + n_inner: int = -1 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + scale_attn_by_inverse_layer_idx: bool = False + tensor_parallel_shards: int = 1 + head_dim: int = 0 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.n_inner is None or self.n_inner == -1: + self.n_inner = 4 * self.n_embd + if self.context_window_size == 0: + for name in ["n_positions", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `n_positions` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.n_embd // self.n_head + assert self.head_dim * self.n_head == self.n_embd + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class GPT2Attention(nn.Module): + def __init__(self, config: GPT2Config): + self.embed_dim = config.n_embd + if config.n_head % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.n_head} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.n_head // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx + + self.c_attn = nn.Linear( + in_features=self.embed_dim, + out_features=3 * self.num_heads * self.head_dim, + bias=True, + ) + self.c_proj = nn.Linear(self.num_heads * self.head_dim, self.embed_dim, bias=True) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h = self.head_dim, self.num_heads + b, s, _ = hidden_states.shape + + qkv = self.c_attn(hidden_states) + qkv = op.reshape(qkv, (b, s, 3 * h, d)) + + if self.scale_attn_by_inverse_layer_idx: + attn_score_scaling_factor = 1.0 / float(layer_id + 1) + else: + attn_score_scaling_factor = 1.0 + + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, + qkv, + self.num_heads, + sm_scale=attn_score_scaling_factor * (self.head_dim**-0.5), + ), + (b, s, h * d), + ) + return self.c_proj(output) + + +class GPT2MLP(nn.Module): + def __init__(self, config: GPT2Config): + embed_dim = config.n_embd + if config.n_inner % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.n_inner} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + intermediate_size = config.n_inner // config.tensor_parallel_shards + self.c_fc = nn.Linear(embed_dim, intermediate_size) + self.c_proj = nn.Linear(intermediate_size, embed_dim) + + def forward(self, hidden_states: Tensor): + hidden_states = self.c_fc(hidden_states) + hidden_states = op.gelu(hidden_states, approximate="tanh") + hidden_states = self.c_proj(hidden_states) + return hidden_states + + +class GPT2Block(nn.Module): + def __init__(self, config: GPT2Config): + hidden_size = config.n_embd + self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.attn = GPT2Attention(config) + self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon) + self.mlp = GPT2MLP(config) + + def _set_tp(): + def _set(param, hint): + param.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = k = v = self.attn.num_heads * hd + _set( + self.attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set( + self.attn.c_attn.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.attn.c_proj.weight, tp.ShardSingleDim("_shard_attn_c_proj", dim=1)) + _set( + self.mlp.c_fc.weight, + tp.ShardSingleDim("_shard_c_fc_weight", dim=0), + ) + _set(self.mlp.c_fc.bias, tp.ShardSingleDim("_shard_c_fc_bias", dim=0)) + _set(self.mlp.c_proj.weight, tp.ShardSingleDim("_shard_mlp_c_proj", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + with ( + tp.shard_bias(self.attn.c_proj, self.tensor_parallel_shards), + tp.shard_bias(self.mlp.c_proj, self.tensor_parallel_shards), + ): + hidden_states = self._apply_residual( + self.attn(self.ln_1(hidden_states), paged_kv_cache, layer_id), + hidden_states, + ) + hidden_states = self._apply_residual(self.mlp(self.ln_2(hidden_states)), hidden_states) + + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out + residual / self.tensor_parallel_shards, "sum") + return out + residual + + +class GPT2Model(nn.Module): + def __init__(self, config: GPT2Config): + assert config.n_embd % config.n_head == 0 + self.wte = nn.Embedding("vocab_size", config.n_embd) + self.wpe = nn.Embedding(config.context_window_size, config.n_embd) + self.h = nn.ModuleList([GPT2Block(config) for _ in range(config.n_layer)]) + self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + # Position Embeddings + # Generate np.arange(offset, offset+seq_len) + # shape[1] indicates the total query length in the batch + input_positions = paged_kv_cache.get_query_positions(inputs.shape[1]) + pos_embd = self.wpe(input_positions) + + # Pass through GPT2Block + hidden_states = inputs + pos_embd + for layer_id, layer in enumerate(self.h): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.ln_f(hidden_states) + return hidden_states + + +class GPT2LMHeadModel(nn.Module): + def __init__(self, config: GPT2Config): + self.transformer = GPT2Model(config) + self.lm_head = nn.Linear(config.n_embd, "vocab_size", bias=False) + self.n_layer = config.n_layer + self.n_embed = config.n_embd + self.n_head = config.n_head + self.head_dim = config.head_dim + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.transformer(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.transformer.wte(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.n_layer, + num_attention_heads=self.n_head // self.tensor_parallel_shards, + num_key_value_heads=self.n_head // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NONE, + rope_scale=-1, + rope_theta=-1, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.n_embed], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.n_embed], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.n_embed], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.n_embed], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.n_embed], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/gpt_bigcode/__init__.py b/python/mlc_llm/model/gpt_bigcode/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_loader.py b/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_loader.py new file mode 100644 index 0000000..c425584 --- /dev/null +++ b/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_loader.py @@ -0,0 +1,14 @@ +""" +This file specifies how MLC's GPTBigCode parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .gpt_bigcode_model import GPTBigCodeForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=GPTBigCodeForCausalLM, + include_qkv=False, + include_gate_up=False, +) diff --git a/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_model.py b/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_model.py new file mode 100644 index 0000000..a93bb2b --- /dev/null +++ b/python/mlc_llm/model/gpt_bigcode/gpt_bigcode_model.py @@ -0,0 +1,350 @@ +""" +Implementation for GPTBigCode architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class GPTBigCodeConfig(ConfigBase): + """Configuration of the GPTBigCode model.""" + + n_embd: int + n_inner: int + n_head: int + n_layer: int + n_positions: int + layer_norm_epsilon: float + vocab_size: int + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + if self.n_positions > 0: + self.context_window_size = self.n_positions + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold("n_positions"), + self.context_window_size, + ) + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class GPTBigCodeMLP(nn.Module): + def __init__(self, config: GPTBigCodeConfig): + super().__init__() + self.n_inner = config.n_inner // config.tensor_parallel_shards + self.c_fc = nn.Linear(in_features=config.n_embd, out_features=self.n_inner, bias=True) + self.c_proj = nn.Linear(in_features=self.n_inner, out_features=config.n_embd, bias=True) + + def forward(self, x: Tensor): + hidden_states = self.c_fc(x) + hidden_states = op.gelu(hidden_states) + hidden_states = self.c_proj(hidden_states) + return hidden_states + + +class GPTBigCodeAttention(nn.Module): + def __init__(self, config: GPTBigCodeConfig): + self.n_embd = config.n_embd + self.head_dim = config.n_embd // config.n_head + self.num_q_heads = config.n_head // config.tensor_parallel_shards + self.num_kv_heads = 1 + assert config.tensor_parallel_shards == 1, ( + "GPT bigcode only support tensor parallel shards = 1" + ) + self.c_attn = nn.Linear( + in_features=self.n_embd, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=True, + ) + self.c_proj = nn.Linear( + in_features=self.num_q_heads * self.head_dim, + out_features=config.n_embd, + bias=True, + ) + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + ): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + + # QKV Projection + qkv = self.c_attn(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, h_q, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.c_proj(output) + + +class GPTBigCodeBlock(nn.Module): + def __init__(self, config: GPTBigCodeConfig): + self.attn = GPTBigCodeAttention(config) + self.mlp = GPTBigCodeMLP(config) + self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.n_embd // config.n_head + q = config.n_head * hd + k = 1 * hd + v = 1 * hd + _set( + self.attn.c_attn, + tp.ShardSingleDim("_shard_c_attn", dim=0, segs=[q, k, v]), + ) + _set(self.attn.c_proj, tp.ShardSingleDim("_shard_c_proj", dim=1)) + _set(self.mlp.c_fc, tp.ShardSingleDim("_shard_mlp_c_fc", dim=0)) + _set(self.mlp.c_proj, tp.ShardSingleDim("_shard_mlp_c_proj", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.attn(self.ln_1(hidden_states), paged_kv_cache, layer_id) + hidden_states = out + hidden_states + out = self.mlp(self.ln_2(hidden_states)) + hidden_states = out + hidden_states + return hidden_states + + +class GPTBigCodeModel(nn.Module): + def __init__(self, config: GPTBigCodeConfig): + assert config.n_embd % config.n_head == 0 + self.wte = nn.Embedding("vocab_size", config.n_embd) + self.wpe = nn.Embedding(config.n_positions, config.n_embd) + self.h = nn.ModuleList([GPTBigCodeBlock(config) for _ in range(config.n_layer)]) + self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + # Position Embeddings + # shape[1] indicates the total query length in the batch + input_positions = paged_kv_cache.get_query_positions(input_embed.shape[1]) + pos_embd = self.wpe(input_positions) + + # apply position embeddings + hidden_states = input_embed + pos_embd + for layer_id, layer in enumerate(self.h): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.ln_f(hidden_states) + + return hidden_states + + +class GPTBigCodeForCausalLM(nn.Module): + def __init__(self, config: GPTBigCodeConfig): + self.transformer = GPTBigCodeModel(config) + self.lm_head = nn.Linear(config.n_embd, "vocab_size", bias=False) + self.n_layer = config.n_layer + self.n_embd = config.n_embd + self.num_q_heads = config.n_head // config.tensor_parallel_shards + self.num_kv_heads = 1 + self.head_dim = config.n_embd // config.n_head + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embed: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.transformer.wte(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.n_layer, + num_attention_heads=self.num_q_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_kv_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NONE, + rope_scale=-1, + rope_theta=-1, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.n_embd], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.n_embd], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.n_embd], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.n_embd], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.n_embd], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/gpt_j/__init__.py b/python/mlc_llm/model/gpt_j/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/gpt_j/gpt_j_loader.py b/python/mlc_llm/model/gpt_j/gpt_j_loader.py new file mode 100644 index 0000000..a476b2a --- /dev/null +++ b/python/mlc_llm/model/gpt_j/gpt_j_loader.py @@ -0,0 +1,16 @@ +""" +This file specifies how MLC's GPTJ parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .gpt_j_model import GPTJForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=GPTJForCausalLM, + layer_prefix="transformer.h", + qkv_target_name="c_attn", + include_gate_up=False, + num_layers_getter=lambda config: config.n_layer, +) diff --git a/python/mlc_llm/model/gpt_j/gpt_j_model.py b/python/mlc_llm/model/gpt_j/gpt_j_model.py new file mode 100644 index 0000000..a200a65 --- /dev/null +++ b/python/mlc_llm/model/gpt_j/gpt_j_model.py @@ -0,0 +1,370 @@ +""" +Implementation for GPTJ architecture. +TODO: add docstring +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class GPTJConfig(ConfigBase): + """Configuration of the GPTJ model.""" + + vocab_size: int + n_embd: int + n_layer: int + n_head: int + layer_norm_epsilon: int + rotary_dim: int + activation_function: str + n_inner: int = -1 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + head_dim: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "n_positions"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.n_embd // self.n_head + assert self.head_dim * self.n_head == self.n_embd + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class GPTJAttention(nn.Module): + def __init__(self, config: GPTJConfig): + self.embed_dim = config.n_embd + self.num_heads = config.n_head // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.max_position_embeddings = config.context_window_size + self.rope_theta = 10000 + self.rotary_dim = config.rotary_dim + self.c_attn = nn.Linear( + in_features=self.embed_dim, + out_features=3 * self.embed_dim, + bias=False, + ) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False) + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + ): + d, h = self.head_dim, self.num_heads + b, s, _ = hidden_states.shape + + qkv = self.c_attn(hidden_states) + qkv = op.reshape(qkv, (b, s, 3 * h, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h * d), + ) + return self.out_proj(output) + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class GPTJMLP(nn.Module): + def __init__(self, config: GPTJConfig): # in MLP: intermediate_size= 4 * embed_dim + embed_dim = config.n_embd + inner_dim = 4 * config.n_embd if config.n_inner is None else config.n_inner + self.fc_in = nn.Linear(embed_dim, inner_dim, bias=True) + self.fc_out = nn.Linear(inner_dim, embed_dim, bias=True) + self.act_fn = ACT2FN[config.activation_function] + + def forward(self, hidden_states: Tensor): + hidden_states = self.fc_in(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.fc_out(hidden_states) + return hidden_states + + +class GPTJBlock(nn.Module): + def __init__(self, config: GPTJConfig): + self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + self.attn = GPTJAttention(config) + self.mlp = GPTJMLP(config) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.attn.num_heads * hd + k = self.attn.num_heads * hd + v = self.attn.num_heads * hd + _set( + self.attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set(self.attn.out_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.fc_in.weight, + tp.ShardSingleDim("_shard_c_fc_weight", dim=0), + ) + _set(self.mlp.fc_out.weight, tp.ShardSingleDim("_shard_mlp_c_proj", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + residual = hidden_states + hidden_states = self.ln_1(hidden_states) + attn_output = self.attn(hidden_states, paged_kv_cache, layer_id) + feed_forward_hidden_states = self.mlp(hidden_states) + hidden_states = self._apply_residual(attn_output + feed_forward_hidden_states, residual) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class GPTJModel(nn.Module): + def __init__(self, config: GPTJConfig): + self.embed_dim = config.n_embd + self.vocab_size = config.vocab_size + self.wte = nn.Embedding(config.vocab_size, self.embed_dim) + self.h = nn.ModuleList([GPTJBlock(config) for _ in range(config.n_layer)]) + self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.h): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.ln_f(hidden_states) + return hidden_states + + +class GPTJForCausalLM(nn.Module): + def __init__(self, config: GPTJConfig): + self.transformer = GPTJModel(config) + self.lm_head = nn.Linear(config.n_embd, config.vocab_size, dtype="float32") + self.dtype = "float32" + self.hidden_size = config.n_embd + self.num_hidden_layers = config.n_layer + self.intermediate_size = 4 * config.n_embd if config.n_inner is None else config.n_inner + self.num_attention_heads = config.n_head + self.rope_theta = 10000 + self.rope_scaling = config.rope_scaling + self.vocab_size = config.vocab_size + self.tensor_parallel_shards = config.tensor_parallel_shards + self.head_dim = config.n_embd // config.n_head + self.rotary_dim = config.rotary_dim + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.transformer(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.transformer.wte(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_attention_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rotary_dim=self.rotary_dim, + rope_scaling=self.rope_scaling, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/gpt_neox/__init__.py b/python/mlc_llm/model/gpt_neox/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/gpt_neox/gpt_neox_loader.py b/python/mlc_llm/model/gpt_neox/gpt_neox_loader.py new file mode 100644 index 0000000..f9192c5 --- /dev/null +++ b/python/mlc_llm/model/gpt_neox/gpt_neox_loader.py @@ -0,0 +1,90 @@ +""" +This file specifies how MLC's GPTNeoX parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .gpt_neox_model import GPTNeoXConfig, GPTNeoXForCausalLM + + +def huggingface(model_config: GPTNeoXConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : GPTNeoXConfig + The configuration of the GPTNeoX model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = GPTNeoXForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # inv_freq/masked_bias/bias is not used in the model + attn = f"gpt_neox.layers.{i}.attention" + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + mapping.add_unused(f"{attn}.masked_bias") + mapping.add_unused(f"{attn}.bias") + + # change the layout of query_key_value + def transform_qkv_layout(w, dtype): + num_attention_heads = model_config.num_attention_heads + head_dim = model_config.head_dim + + org_shape = w.shape + w = np.reshape(w, [num_attention_heads, 3 * head_dim, -1]) + qkv = np.split(w, indices_or_sections=3, axis=1) + w = np.concatenate(qkv, axis=0) + w = np.reshape(w, org_shape) + return w.astype(dtype) + + qkv_proj = f"{attn}.query_key_value" + for param_name in ["weight", "bias"]: + mlc_name = f"{qkv_proj}.{param_name}" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + transform_qkv_layout, + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + if ".dense_h_to_4h.bias" in mlc_name or ".dense_4h_to_h.bias" in mlc_name: + param_dtype = model_config.ffn_out_dtype + else: + param_dtype = mlc_param.dtype + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=param_dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/gpt_neox/gpt_neox_model.py b/python/mlc_llm/model/gpt_neox/gpt_neox_model.py new file mode 100644 index 0000000..6d6e838 --- /dev/null +++ b/python/mlc_llm/model/gpt_neox/gpt_neox_model.py @@ -0,0 +1,417 @@ +""" +Implementation for GPTNeoX architecture. +""" + +import dataclasses +import logging +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class GPTNeoXConfig(ConfigBase): + """Configuration of the GPTNeoX model.""" + + use_parallel_residual: bool + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + layer_norm_eps: float + vocab_size: int + rotary_pct: float + position_embedding_base: int = 0 + context_window_size: int = 0 + head_dim: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + ffn_out_dtype: str = "float32" + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class GPTNeoXAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: GPTNeoXConfig): + self.rope_theta = config.position_embedding_base + self.hidden_size = config.hidden_size + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_attention_heads = config.num_attention_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.query_key_value = nn.Linear( + in_features=self.hidden_size, + out_features=3 * self.num_attention_heads * self.head_dim, + bias=True, + ) + self.dense = nn.Linear( + self.num_attention_heads * self.head_dim, self.hidden_size, bias=True + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + # hidden_states: [batch_size, seq_len, hidden_size] + batch_size, seq_len, _ = hidden_states.shape + + # q/k/v states: [batch_size, seq_len, hidden_size] + qkv = self.query_key_value(hidden_states) + qkv = op.reshape(qkv, (batch_size, seq_len, 3 * self.num_attention_heads, self.head_dim)) + + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_attention_heads, sm_scale=self.head_dim**-0.5 + ), + (batch_size, seq_len, self.head_dim * self.num_attention_heads), + ) + attn_output = self.dense(output) + return attn_output + + +class GPTNeoXMLP(nn.Module): + def __init__(self, config: GPTNeoXConfig): + super().__init__() + out_dtype = config.ffn_out_dtype + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.dense_h_to_4h = nn.Linear( + config.hidden_size, + self.intermediate_size, + out_dtype=out_dtype, + ) + self.dense_4h_to_h = nn.Linear( + self.intermediate_size, + config.hidden_size, + out_dtype=out_dtype, + ) + + def forward(self, hidden_states: Tensor): + dtype = hidden_states.dtype + if hidden_states.dtype != dtype: + hidden_states = hidden_states.astype(dtype) + hidden_states = self.dense_h_to_4h(hidden_states) + hidden_states = op.gelu(hidden_states) + if hidden_states.dtype != dtype: + hidden_states = hidden_states.astype(dtype) + hidden_states = self.dense_4h_to_h(hidden_states) + if hidden_states.dtype != dtype: + hidden_states = hidden_states.astype(dtype) + return hidden_states + + +class GPTNeoXLayer(nn.Module): + def __init__(self, config: GPTNeoXConfig): + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.attention = GPTNeoXAttention(config) + self.mlp = GPTNeoXMLP(config) + self.use_parallel_residual = config.use_parallel_residual + + def _set_tp(): + def _set(param, hint): + param.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = k = v = self.attention.num_attention_heads * hd + _set( + self.attention.query_key_value.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set( + self.attention.query_key_value.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.attention.dense.weight, tp.ShardSingleDim("_shard_dense", dim=1)) + _set( + self.mlp.dense_h_to_4h.weight, + tp.ShardSingleDim("_shard_dense_h_to_4h_weight", dim=0), + ) + _set( + self.mlp.dense_h_to_4h.bias, + tp.ShardSingleDim("_shard_dense_h_to_4h_bias", dim=0), + ) + _set( + self.mlp.dense_4h_to_h.weight, + tp.ShardSingleDim("_shard_dense_4h_to_h", dim=1), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + dtype = hidden_states.dtype + attn_input = self.input_layernorm(hidden_states) + with tp.shard_bias(self.attention.dense, self.tensor_parallel_shards): + attn_output = self.attention( + attn_input, + paged_kv_cache, + layer_id, + ) + if self.use_parallel_residual: + mlp_input = self.post_attention_layernorm(hidden_states) + mlp_output = self.mlp(mlp_input) + hidden_states = mlp_output + attn_output + hidden_states + else: + attn_output = self._apply_residual(attn_output, hidden_states) + mlp_input = self.post_attention_layernorm(attn_output) + with tp.shard_bias(self.mlp.dense_4h_to_h, self.tensor_parallel_shards): + mlp_output = self.mlp(mlp_input) + hidden_states = self._apply_residual(mlp_output.astype(dtype), attn_output) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out + residual / self.tensor_parallel_shards, "sum") + return out + residual + + +class GPTNeoXModel(nn.Module): + def __init__(self, config: GPTNeoXConfig): + self.embed_in = nn.Embedding(num="vocab_size", dim=config.hidden_size) + self.layers = nn.ModuleList([GPTNeoXLayer(config) for _ in range(config.num_hidden_layers)]) + self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.final_layer_norm(hidden_states) + return hidden_states + + +class GPTNeoXForCausalLM(nn.Module): + def __init__(self, config: GPTNeoXConfig): + self.gpt_neox = GPTNeoXModel(config) + self.embed_out = nn.Linear( + in_features=config.hidden_size, + out_features="vocab_size", + bias=False, + dtype="float32", + ) + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + self.rotary_pct = config.rotary_pct + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.gpt_neox(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.embed_out(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.gpt_neox.embed_in(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.gpt_neox(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.embed_out(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.gpt_neox(input_embed, paged_kv_cache) + logits = self.embed_out(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_attention_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + rotary_dim=int(self.head_dim * self.rotary_pct), + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/internlm/__init__.py b/python/mlc_llm/model/internlm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/internlm/internlm_loader.py b/python/mlc_llm/model/internlm/internlm_loader.py new file mode 100644 index 0000000..7860450 --- /dev/null +++ b/python/mlc_llm/model/internlm/internlm_loader.py @@ -0,0 +1,16 @@ +""" +This file specifies how MLC's InternLM parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .internlm_model import InternLMForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=InternLMForCausalLM, + qkv_target_name="wqkv_pack", + add_qkv_bias=True, + qkv_bias_optional=True, + add_unused=["rotary_emb.inv_freq"], +) diff --git a/python/mlc_llm/model/internlm/internlm_model.py b/python/mlc_llm/model/internlm/internlm_model.py new file mode 100644 index 0000000..9024e5a --- /dev/null +++ b/python/mlc_llm/model/internlm/internlm_model.py @@ -0,0 +1,380 @@ +""" +Implementation for InternLM architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class InternLMConfig(ConfigBase): + """Configuration of the InternLM model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + rms_norm_eps: float + intermediate_size: int + bias: bool + use_cache: bool + pad_token_id: int + bos_token_id: int + eos_token_id: int + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + head_dim: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class InternLMAttention(nn.Module): + def __init__(self, config: InternLMConfig): + self.hidden_size = config.hidden_size + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.num_attention_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.max_position_embeddings = config.context_window_size + + self.wqkv_pack = nn.Linear( + self.hidden_size, 3 * self.num_heads * self.head_dim, bias=config.bias + ) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h = self.head_dim, self.num_heads + b, s, _ = hidden_states.shape + qkv = self.wqkv_pack(hidden_states) + qkv = op.reshape(qkv, (b, s, 3 * h, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +class InternLMMLP(nn.Module): + def __init__(self, config: InternLMConfig): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class InternLMDecoderLayer(nn.Module): + def __init__(self, config: InternLMConfig): + self.self_attn = InternLMAttention(config) + self.mlp = InternLMMLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_heads * hd + k = self.self_attn.num_heads * hd + v = self.self_attn.num_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.wqkv_pack.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + if config.bias: + _set( + self.self_attn.wqkv_pack.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set( + self.self_attn.o_proj.weight, + tp.ShardSingleDim("_shard_o_weight", dim=1), + ) + if config.bias: + _set( + self.self_attn.o_proj.bias, + tp.ShardSingleDim("_shard_o_bias", dim=0), + ) + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_gate_up", segs=[i, i], dim=0), + ) + _set( + self.mlp.down_proj.weight, + tp.ShardSingleDim("_shard_mlp_down_proj", dim=1), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class InternLMModel(nn.Module): + def __init__(self, config: InternLMConfig): + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [InternLMDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class InternLMForCausalLM(nn.Module): + def __init__(self, config: InternLMConfig): + self.model = InternLMModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.vocab_size = config.vocab_size + self.rope_theta = 10000 + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_attention_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/internlm2/__init__.py b/python/mlc_llm/model/internlm2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/internlm2/internlm2_loader.py b/python/mlc_llm/model/internlm2/internlm2_loader.py new file mode 100644 index 0000000..b0442c5 --- /dev/null +++ b/python/mlc_llm/model/internlm2/internlm2_loader.py @@ -0,0 +1,94 @@ +""" +This file specifies how MLC's InternLM2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .internlm2_model import InternLM2ForCausalLM + + +def huggingface(model_config: InternLM2ForCausalLM, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : InternLM2Config + The configuration of the InternLM2 model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = InternLM2ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + def _convert_wqkv_layout(wqkv, dtype): + config = model_config + kv_groups = config.num_attention_heads // config.num_key_value_heads + head_dim = config.hidden_size // config.num_attention_heads + wqkv = wqkv.reshape(-1, 2 + kv_groups, head_dim, wqkv.shape[-1]) + wq, wk, wv = np.split(wqkv, [kv_groups, kv_groups + 1], axis=1) + wq = wq.reshape(-1, wq.shape[-1]) + wk = wk.reshape(-1, wk.shape[-1]) + wv = wv.reshape(-1, wv.shape[-1]) + return np.concatenate([wq, wk, wv], axis=0).astype(dtype) + + for i in range(model_config.num_hidden_layers): + # Add gates in MLP + mlp = f"model.layers.{i}.feed_forward" + mlc_name = f"{mlp}.gate_up_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.w1.weight", + f"{mlp}.w3.weight", + ], + functools.partial( + lambda w1, w3, dtype: np.concatenate([w1, w3], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + mlc_name = f"model.layers.{i}.attention.wqkv.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + _convert_wqkv_layout, + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + return mapping diff --git a/python/mlc_llm/model/internlm2/internlm2_model.py b/python/mlc_llm/model/internlm2/internlm2_model.py new file mode 100644 index 0000000..43254a2 --- /dev/null +++ b/python/mlc_llm/model/internlm2/internlm2_model.py @@ -0,0 +1,377 @@ +""" +Implementation for InternLM2 architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class InternLM2Config(ConfigBase): + """Configuration of the InternLM2 model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + rms_norm_eps: float + intermediate_size: int + bias: bool + use_cache: bool + rope_theta: int + pad_token_id: int + bos_token_id: int + eos_token_id: int + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + head_dim: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 2048), + ) + self.prefill_chunk_size = min(self.context_window_size, 2048) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 2048), + ) + self.prefill_chunk_size = min(self.context_window_size, 2048) + + +class InternLM2Attention(nn.Module): + def __init__(self, config: InternLM2Config): + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.hidden_size = config.hidden_size + self.rope_theta = config.rope_theta + self.num_heads = config.num_attention_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.max_position_embeddings = config.context_window_size + + self.wqkv = nn.Linear( + self.hidden_size, + (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=config.bias, + ) + self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.wqkv(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.wo(output) + return attn_output + + +class InternLM2MLP(nn.Module): + def __init__(self, config: InternLM2Config): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.w2 = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.w2(op.silu(x1) * x2) + + +class InternLM2DecoderLayer(nn.Module): + def __init__(self, config: InternLM2Config): + self.attention = InternLM2Attention(config) + self.feed_forward = InternLM2MLP(config) + self.attention_norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.ffn_norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.attention.num_heads * hd + k = self.attention.num_key_value_heads * hd + v = self.attention.num_key_value_heads * hd + i = self.feed_forward.intermediate_size + _set( + self.attention.wqkv.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + if config.bias: + _set( + self.attention.wqkv.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.attention.wo.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.feed_forward.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.feed_forward.w2.weight, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + residual = hidden_states + hidden_states = self.attention_norm(hidden_states) + hidden_states = self.attention(hidden_states, paged_kv_cache, layer_id) + hidden_states = self._apply_residual(hidden_states, residual=residual) + residual = hidden_states + hidden_states = self.ffn_norm(hidden_states) + hidden_states = self.feed_forward(hidden_states) + hidden_states = self._apply_residual(hidden_states, residual=residual) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class InternLM2Model(nn.Module): + def __init__(self, config: InternLM2Config): + self.padding_idx = config.pad_token_id + self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [InternLM2DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class InternLM2ForCausalLM(nn.Module): + def __init__(self, config: InternLM2Config): + self.model = InternLM2Model(config) + self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.dtype = "float32" + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.output(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.tok_embeddings(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.output(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.output(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/llama/__init__.py b/python/mlc_llm/model/llama/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/llama/llama_loader.py b/python/mlc_llm/model/llama/llama_loader.py new file mode 100644 index 0000000..6931593 --- /dev/null +++ b/python/mlc_llm/model/llama/llama_loader.py @@ -0,0 +1,104 @@ +""" +This file specifies how MLC's Llama parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization, make_awq_quant + +from .llama_model import LlamaConfig, LlamaForCausalLM + +awq_quant = make_awq_quant(LlamaForCausalLM) + + +huggingface = make_standard_hf_loader( + model_cls=LlamaForCausalLM, + add_unused=["rotary_emb.inv_freq"], +) + + +def awq(model_config: LlamaConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of AWQ parameters. + Parameters + ---------- + model_config : LlamaConfig + The configuration of the Llama model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to AWQ. + """ + model, _ = awq_quant(model_config, quantization) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"model.layers.{i}.self_attn" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{attn}.qkv_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{quantize_suffix}", + f"{attn}.k_proj.{quantize_suffix}", + f"{attn}.v_proj.{quantize_suffix}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate( + [q, k, v], + axis=1, # AWQ GEMM would transpose the weight + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Concat gate and up in MLP + mlp = f"model.layers.{i}.mlp" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{mlp}.gate_up_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.gate_proj.{quantize_suffix}", + f"{mlp}.up_proj.{quantize_suffix}", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate( + [gate, up], + axis=1, # AWQ GEMM would transpose the weight + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # inv_freq is not used in the model + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + return mapping diff --git a/python/mlc_llm/model/llama/llama_model.py b/python/mlc_llm/model/llama/llama_model.py new file mode 100644 index 0000000..4659a24 --- /dev/null +++ b/python/mlc_llm/model/llama/llama_model.py @@ -0,0 +1,542 @@ +""" +Implementation for Llama2 architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class LlamaConfig(ConfigBase): + """Configuration of the Llama model.""" + + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + rms_norm_eps: float + vocab_size: int + tie_word_embeddings: bool = False + position_embedding_base: int = 0 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + num_key_value_heads: int = 0 + head_dim: int = 0 + tensor_parallel_shards: int = 1 + pipeline_parallel_stages: int = 1 + max_batch_size: int = 1 + disaggregation: bool = False + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.rope_scaling is not None: + if "rope_type" not in self.rope_scaling: + self.rope_scaling = None + else: + assert self.rope_scaling["rope_type"] == "llama3", ( + f"Unsupported RoPE scaling type {self.rope_scaling['rope_type']} for Llama" + ) + + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if ( + self.pipeline_parallel_stages <= 0 + or self.pipeline_parallel_stages > self.num_hidden_layers + ): + raise ValueError( + f'Invalid "pipeline_parallel_stages" value ({self.pipeline_parallel_stages}). ' + ) + if self.num_key_value_heads == 0: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.num_attention_heads % self.num_key_value_heads == 0 + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class LlamaFFN(nn.Module): + def __init__(self, config: LlamaConfig): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class LlamaEmbedding(nn.Embedding): + """The embedding module that can be shared with the final lm_head. From Qwen2Embedding.""" + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class LlamaAttention(nn.Module): + def __init__(self, config: LlamaConfig): + self.head_dim = config.head_dim + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({config.num_key_value_heads}) must be divisible by tensor_parallel_shards" # noqa: E501 + ) + assert config.num_key_value_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than {config.num_key_value_heads}" + ) + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=False, + ) + self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class LlamaDecoderLayer(nn.Module): + def __init__(self, config: LlamaConfig): + rms_norm_eps = config.rms_norm_eps + self.self_attn = LlamaAttention(config) + self.mlp = LlamaFFN(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class LlamaModel(nn.Module): + def __init__(self, config: LlamaConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = LlamaEmbedding("vocab_size", config.hidden_size) + self.layers = nn.ModuleList( + [LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.num_layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + + # Compute pipeline layer partition. + layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + self.layer_partition = [ + i * layers_per_stage for i in range(config.pipeline_parallel_stages) + ] + [config.num_hidden_layers] + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.layers): + if layer_id != 0 and layer_id in self.layer_partition: + hidden_states = op_ext.pipeline_stage_boundary(hidden_states) + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class LlamaForCausalLM(nn.Module): + def __init__(self, config: LlamaConfig): + self.model = LlamaModel(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_scaling = config.rope_scaling + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.disaggregation = config.disaggregation + self.dtype = "float32" + + def _set_pp(): + # hidden layers + for layer_id in range(config.num_hidden_layers): + stage = layer_id // (config.num_hidden_layers // config.pipeline_parallel_stages) + for _, param in self.model.layers[layer_id].named_parameters(): + param.attrs["pipeline_stages"] = [stage] + # last stage + last_stage = config.pipeline_parallel_stages - 1 + self.model.norm.weight.attrs["pipeline_stages"] = [last_stage] + # embedding table and lm_head is required by all stages + all_stages = list(range(config.pipeline_parallel_stages)) + self.model.embed_tokens.weight.attrs["pipeline_stages"] = all_stages + if not config.tie_word_embeddings: + self.lm_head.weight.attrs["pipeline_stages"] = all_stages + + _set_pp() + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return self.get_logits(hidden_states) + + def batch_forward_to_last_hidden_states( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + return hidden_states + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def get_logits(self, hidden_states: Tensor): + op_ext.configure() + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_select_last_hidden_states(self, hidden_states: Tensor, logit_positions: Tensor): + op_ext.configure() + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=0) + return hidden_states + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def prefill_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def decode_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_prefill_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_decode_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_verify_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rope_scaling=self.rope_scaling, + layer_partition=self.model.layer_partition, + enable_disaggregation=self.disaggregation, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "get_logits": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_select_last_hidden_states": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/llama4/__init__.py b/python/mlc_llm/model/llama4/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/llama4/llama4_loader.py b/python/mlc_llm/model/llama4/llama4_loader.py new file mode 100644 index 0000000..ab82583 --- /dev/null +++ b/python/mlc_llm/model/llama4/llama4_loader.py @@ -0,0 +1,119 @@ +""" +This file specifies how MLC's Llama parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .llama4_model import Llama4Config, Llama4ForCausalLM + + +def huggingface(model_config: Llama4Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : Llama4Config + The configuration of the Llama model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Llama4ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.text_config.num_hidden_layers): + # Add shared expert weights + mlp = f"model.layers.{i}.feed_forward.shared_expert" + mlc_name = f"{mlp}.gate_up_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"language_model.{mlp}.gate_proj.weight", + f"language_model.{mlp}.up_proj.weight", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Add router weights + mlp = f"model.layers.{i}.feed_forward" + mlc_name = f"{mlp}.router.router.weight" + hf_name = f"language_model.{mlp}.router.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + hf_name, + ], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Add experts weights + mlp = f"model.layers.{i}.feed_forward" + hf_name = f"language_model.{mlp}.experts.gate_up_proj" + mlc_name = f"{mlp}.experts.gate_up_proj" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + hf_name, + ], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + mlp = f"model.layers.{i}.feed_forward" + mlc_name = f"{mlp}.experts.down_proj" + hf_name = f"language_model.{mlp}.experts.down_proj" + + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + hf_name, + ], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [f"language_model.{mlc_name}"], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/llama4/llama4_model.py b/python/mlc_llm/model/llama4/llama4_model.py new file mode 100644 index 0000000..432476d --- /dev/null +++ b/python/mlc_llm/model/llama4/llama4_model.py @@ -0,0 +1,821 @@ +""" +Implementation for Llama4 architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +import tvm +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op +from tvm.relax.frontend.nn.llm import position_embedding + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.model.qwen3.qwen3_model import ACT2FN +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Llama4TextConfig(ConfigBase): + """Configuration of the Text portion of the Llama model.""" + + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + rms_norm_eps: float + rope_theta: float + use_qk_norm: bool + interleave_moe_layer_step: int + num_experts_per_tok: int + num_local_experts: int + hidden_act: str + tie_word_embeddings: bool = False + position_embedding_base: int = 0 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + num_key_value_heads: int = 0 + head_dim: int = 0 + attn_scale: float = 0.1 + floor_scale: int = 8192 + vocab_size: int = 202048 + attention_bias: bool = False + attn_temperature_tuning: bool = True + no_rope_layers: list[int] = None + no_rope_layer_interval: int = 4 + moe_layers: list[int] = None + + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.rope_scaling is not None: + if "rope_type" not in self.rope_scaling: + self.rope_scaling = None + else: + assert self.rope_scaling["rope_type"] == "llama3", ( + f"Unsupported RoPE scaling type {self.rope_scaling['rope_type']} for Llama" + ) + + # Define which layers to avoid RoPE + if self.no_rope_layers == []: + self.no_rope_layers = None + + default_no_rope_layers = [ + int((layer_idx + 1) % self.no_rope_layer_interval != 0) + for layer_idx in range(self.num_hidden_layers) + ] + + self.no_rope_layers = self.no_rope_layers if self.no_rope_layers else default_no_rope_layers + + # Define which layers to apply MoE + self.moe_layers = ( + self.moe_layers + if self.moe_layers is not None + else list( + range( + self.interleave_moe_layer_step - 1, + self.num_hidden_layers, + self.interleave_moe_layer_step, + ) + ) + ) + + +@dataclasses.dataclass +class Llama4Config(ConfigBase): + """Configuration of the Llama model.""" + + text_config: Llama4TextConfig + tensor_parallel_shards: int = 1 + context_window_size: int = 0 + pipeline_parallel_stages: int = 1 + prefill_chunk_size: int = 0 + max_batch_size: int = 1 + disaggregation: bool = False + max_position_embeddings = 4096 * 32 + vocab_size: int = 202048 + + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self) -> None: + text_config_dict: Dict[str, Any] # noqa: UP006 + if isinstance(self.text_config, ConfigBase): + text_config_dict = dataclasses.asdict(self.text_config) + else: + text_config_dict = dict(self.text_config) + + for k, v in text_config_dict.pop("kwargs", {}).items(): + text_config_dict[k] = v + + self.text_config = Llama4TextConfig.from_dict(text_config_dict) + + if self.context_window_size == 0: + # Fall back to max_position_embeddings + + self.context_window_size = self.max_position_embeddings + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold("max_position_embeddings"), + self.context_window_size, + ) + + if self.text_config.num_key_value_heads == 0: + self.text_config.num_key_value_heads = self.text_config.num_attention_heads + if self.text_config.head_dim == 0: + self.text_config.head_dim = ( + self.text_config.hidden_size // self.text_config.num_attention_heads + ) + assert self.text_config.num_attention_heads % self.text_config.num_key_value_heads == 0 + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class Llama4TextMLP(nn.Module): + def __init__(self, config: Llama4Config): + super().__init__() + if config.text_config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.text_config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = ( + config.text_config.intermediate_size // config.tensor_parallel_shards + ) + self.gate_up_proj = nn.Linear( + in_features=config.text_config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear( + self.intermediate_size, config.text_config.hidden_size, bias=False + ) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + inter_out = op.silu(x1) * x2 + + return self.down_proj(inter_out) + + +class LlamaEmbedding(nn.Embedding): + """The embedding module that can be shared with the final lm_head. From Qwen2Embedding.""" + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class Llama4TextL2Norm(nn.Module): + def __init__(self, eps, hidden_size): + self.eps = eps + self.hidden_size = hidden_size + + def forward(self, x): + weight = op.ones((self.hidden_size,), dtype=x.dtype) + return op.rms_norm(x, weight=weight, axes=[-1], epsilon=self.eps) + + +class Llama4TextAttention(nn.Module): + def __init__(self, config: Llama4Config, layer_idx): + self.head_dim = config.text_config.head_dim + self.attn_scale = config.text_config.attn_scale + self.floor_scale = config.text_config.floor_scale + self.num_attention_heads = config.text_config.num_attention_heads + self.num_kv_heads = config.text_config.num_key_value_heads + self.num_q_heads = config.text_config.num_attention_heads // config.tensor_parallel_shards + assert config.text_config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({config.text_config.num_key_value_heads}) must be divisible by " + f"tensor_parallel_shards" + ) + + assert config.text_config.num_key_value_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than " + f"{config.text_config.num_key_value_heads}" + ) + self.num_kv_heads = config.text_config.num_key_value_heads // config.tensor_parallel_shards + self.q_proj = nn.Linear( + config.text_config.hidden_size, + self.num_q_heads * self.head_dim, + bias=config.text_config.attention_bias, + ) + self.k_proj = nn.Linear( + config.text_config.hidden_size, + self.num_kv_heads * self.head_dim, + bias=config.text_config.attention_bias, + ) + self.v_proj = nn.Linear( + config.text_config.hidden_size, + self.num_kv_heads * self.head_dim, + bias=config.text_config.attention_bias, + ) + self.o_proj = nn.Linear( + self.num_q_heads * self.head_dim, + config.text_config.hidden_size, + bias=config.text_config.attention_bias, + ) + + self.attn_temperature_tuning = config.text_config.attn_temperature_tuning + self.use_rope = config.text_config.no_rope_layers[layer_idx] + + self.layer_idx = layer_idx + + self.rope_theta = config.text_config.rope_theta + self.rope_scaling = config.text_config.rope_scaling + self.rope_scaling["rope_type"] = "llama4" + + self.use_qk_norm = config.text_config.use_qk_norm + self.rms_norm_eps = config.text_config.rms_norm_eps + + self.q_norm = Llama4TextL2Norm(self.rms_norm_eps, self.head_dim) + self.k_norm = Llama4TextL2Norm(self.rms_norm_eps, self.head_dim) + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + cache_position, + ): + d, h_q = self.head_dim, self.num_q_heads + b, s, _ = hidden_states.shape + # QKV Projection + query_states = op.reshape(self.q_proj(hidden_states), (b, s, -1, d)) + key_states = op.reshape(self.k_proj(hidden_states), (b, s, -1, d)) + value_states = op.reshape(self.v_proj(hidden_states), (b, s, -1, d)) + + if self.use_rope: + qkv = op.concat([query_states, key_states, value_states], dim=2) + + apply_rope = tvm.tirx.IntImm("int64", 1) + + rotary_emb = position_embedding.llama4_rope_with_position_map( + theta=self.rope_theta, + scale=1.0, + head_dim=self.head_dim, + num_q_heads=self.num_q_heads, + num_kv_heads=self.num_kv_heads, + dtype=query_states.dtype, + rope_scaling=self.rope_scaling, + ) + + query_states, key_states, value_states = op.tensor_ir_op( + rotary_emb, + "llama4_rope_with_position_map", + args=[op.squeeze(qkv, axis=0), cache_position, apply_rope], + out=( + Tensor.placeholder((s, h_q, d), query_states.dtype), + Tensor.placeholder((s, self.num_kv_heads, d), query_states.dtype), + Tensor.placeholder((s, self.num_kv_heads, d), query_states.dtype), + ), + ) + query_states = query_states.reshape(b, s, h_q, d) + key_states = key_states.reshape(b, s, self.num_kv_heads, d) + value_states = value_states.reshape(b, s, self.num_kv_heads, d) + + if self.use_qk_norm and self.use_rope: + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + if self.attn_temperature_tuning and not self.use_rope: + attn_scales = ( + op.log( + op.floor( + (op.astype(cache_position, query_states.dtype) + 1.0) / self.floor_scale + ) + + 1.0 + ) + * self.attn_scale + + 1.0 + ) + + attn_scales = op.broadcast_to(attn_scales.reshape(1, s, 1, 1), (b, s, 1, 1)) + query_states = query_states * attn_scales + + qkv = op.concat([query_states, key_states, value_states], dim=2) + + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class Llama4TextExperts(nn.Module): + def __init__(self, config: Llama4Config): + self.num_experts = config.text_config.num_local_experts + self.intermediate_size = ( + config.text_config.intermediate_size // config.tensor_parallel_shards + ) + self.hidden_size = config.text_config.hidden_size + self.expert_dim = self.intermediate_size + + self.gate_up_proj = nn.Parameter( + shape=(self.num_experts, self.hidden_size, 2 * self.expert_dim) + ) + self.down_proj = nn.Parameter(shape=(self.num_experts, self.expert_dim, self.hidden_size)) + self.act_fn = ACT2FN[config.text_config.hidden_act] + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(self.gate_up_proj.shape[0], -1, self.hidden_size) + gate_up = op.matmul(hidden_states, self.gate_up_proj) + gate, up = op.chunk(gate_up, chunks=2, dim=-1) + next_states = op.matmul((up * self.act_fn(gate)), self.down_proj) + next_states = next_states.reshape(-1, self.hidden_size) + return next_states + + +class Llama4Router(nn.Module): + def __init__(self, config: Llama4Config): + self.num_experts = config.text_config.num_local_experts + self.top_k = config.text_config.num_experts_per_tok + self.intermediate_size = self.num_experts // config.tensor_parallel_shards + self.router = nn.Linear( + in_features=config.text_config.hidden_size, + out_features=self.intermediate_size, + bias=False, + ) + + def forward(self, hidden_states): + router_logits = self.router(hidden_states) + router_top_value, router_indices = op_ext.moe_misc.gating_topk(router_logits, self.top_k) + + j_axis = op.arange(0, self.num_experts) + j_axis = op.unsqueeze(j_axis, 0) + idx_exp = op.unsqueeze(router_indices, -1) + mask = op.equal(idx_exp, j_axis) + val_exp = op.unsqueeze(router_top_value, -1) + neg_inf = op.full(mask.shape, -1e9, dtype=hidden_states.dtype) + masked_vals = op.where(mask, val_exp, neg_inf) + router_scores = op.max(masked_vals, axis=1) + + router_scores = op.sigmoid(router_scores) + return router_scores, router_logits + + +class Llama4TextMoe(nn.Module): + def __init__(self, config: Llama4Config): + self.top_k = config.text_config.num_experts_per_tok + self.hidden_dim = config.text_config.hidden_size + self.num_experts = config.text_config.num_local_experts + self.experts = Llama4TextExperts(config) + self.router = Llama4Router(config) + self.shared_expert = Llama4TextMLP(config) + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_scores, _ = self.router(hidden_states) + + routed_in = op.broadcast_to( + hidden_states.reshape(1, *hidden_states.shape), + [router_scores.shape[1], *hidden_states.shape], + ) + routed_in = routed_in.reshape(-1, self.hidden_dim) + + routed_in = routed_in * op.permute_dims(router_scores, axes=[1, 0]).reshape(-1, 1) + + routed_out = self.experts(routed_in) + out = self.shared_expert(hidden_states) + + out += op.sum(routed_out.reshape(router_scores.shape[1], -1, routed_out.shape[-1]), axis=0) + + return out + + +class Llama4TextDecoderLayer(nn.Module): + def __init__(self, config: Llama4Config, layer_idx): + rms_norm_eps = config.text_config.rms_norm_eps + self.self_attn = Llama4TextAttention(config, layer_idx) + self.is_moe_layer = layer_idx in config.text_config.moe_layers + if self.is_moe_layer: # the 128E model interleaves dense / sparse + self.feed_forward = Llama4TextMoe(config) + else: + self.feed_forward = Llama4TextMLP(config) + + self.input_layernorm = nn.RMSNorm( + config.text_config.hidden_size, -1, rms_norm_eps, bias=False + ) + self.post_attention_layernorm = nn.RMSNorm( + config.text_config.hidden_size, -1, rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + if hasattr(layer, "weight"): + layer.weight.attrs["shard_strategy"] = hint + else: + layer.attrs["shard_strategy"] = hint + + _set(self.self_attn.q_proj, tp.ShardSingleDim("_shard_q", dim=0)) + _set(self.self_attn.k_proj, tp.ShardSingleDim("_shard_k", dim=0)) + _set(self.self_attn.v_proj, tp.ShardSingleDim("_shard_v", dim=0)) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + + if isinstance(self.feed_forward, Llama4TextMLP): + i = self.feed_forward.intermediate_size + _set( + self.feed_forward.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set( + self.feed_forward.down_proj, + tp.ShardSingleDim("_shard_mlp_down", dim=1), + ) + else: + assert isinstance(self.feed_forward, Llama4TextMoe) + i = self.feed_forward.shared_expert.intermediate_size + _set( + self.feed_forward.shared_expert.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set( + self.feed_forward.shared_expert.down_proj, + tp.ShardSingleDim("_shard_mlp_down", dim=1), + ) + + j = self.feed_forward.experts.intermediate_size + _set( + self.feed_forward.experts.gate_up_proj, + tp.ShardSingleDim("_shard_expert_mlp_up", segs=[j, j], dim=2), + ) + _set( + self.feed_forward.experts.down_proj, + tp.ShardSingleDim("_shard_expert_mlp_down", dim=1), + ) + + _set( + self.feed_forward.router.router, + tp.ShardSingleDim("_shard_router", dim=0), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + cache_position, + ): + out = self.self_attn( + self.input_layernorm(hidden_states), + paged_kv_cache, + layer_id, + cache_position, + ) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.feed_forward(self.post_attention_layernorm(hidden_states)) + + hidden_states = self._apply_residual( + op.reshape(out, hidden_states.shape), residual=hidden_states + ) + + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Llama4TextModel(nn.Module): + def __init__(self, config: Llama4Config): + assert config.text_config.hidden_size % config.text_config.num_attention_heads == 0 + self.embed_tokens = LlamaEmbedding("vocab_size", config.text_config.hidden_size) + self.layers = nn.ModuleList( + [ + Llama4TextDecoderLayer(config, layer_idx) + for layer_idx in range(config.text_config.num_hidden_layers) + ] + ) + self.norm = nn.RMSNorm( + config.text_config.hidden_size, + -1, + config.text_config.rms_norm_eps, + bias=False, + ) + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + cache_position = paged_kv_cache.get_query_positions( + input_embed.shape[0] * input_embed.shape[1] + ) + + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id, cache_position) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Llama4ForCausalLM(nn.Module): + def __init__(self, config: Llama4Config): + self.text_config = config.text_config + self.model = Llama4TextModel(config) + self.tie_word_embeddings = self.text_config.tie_word_embeddings + if not self.text_config.tie_word_embeddings: + self.lm_head = nn.Linear(self.text_config.hidden_size, "vocab_size", bias=False) + self.num_hidden_layers = self.text_config.num_hidden_layers + self.num_attention_heads = self.text_config.num_attention_heads + self.num_key_value_heads = self.text_config.num_key_value_heads + self.head_dim = self.text_config.head_dim + self.hidden_size = self.text_config.hidden_size + self.vocab_size = self.text_config.vocab_size + self.rope_scaling = self.text_config.rope_scaling + self.rope_theta = self.text_config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.disaggregation = config.disaggregation + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return self.get_logits(hidden_states) + + def batch_forward_to_last_hidden_states( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + return hidden_states + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def get_logits(self, hidden_states: Tensor): + op_ext.configure() + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_select_last_hidden_states(self, hidden_states: Tensor, logit_positions: Tensor): + op_ext.configure() + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=0) + return hidden_states + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def prefill_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def decode_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_prefill_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_decode_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_verify_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NONE, + rope_scale=1, + rope_theta=self.rope_theta, + rope_scaling=self.rope_scaling, + enable_disaggregation=self.disaggregation, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "get_logits": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_select_last_hidden_states": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/llava/__init__.py b/python/mlc_llm/model/llava/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/llava/llava_loader.py b/python/mlc_llm/model/llava/llava_loader.py new file mode 100644 index 0000000..8bf2c4c --- /dev/null +++ b/python/mlc_llm/model/llava/llava_loader.py @@ -0,0 +1,101 @@ +""" +This file specifies how MLC's Llava parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization, make_awq_quant + +from .llava_model import LlavaConfig, LlavaForCausalLM + +awq_quant = make_awq_quant(LlavaForCausalLM) + + +def _num_layers(config: object) -> int: + return config.text_config.num_hidden_layers + + +huggingface = make_standard_hf_loader( + model_cls=LlavaForCausalLM, + layer_prefix="language_model.model.layers", + add_unused=["rotary_emb.inv_freq"], + num_layers_getter=_num_layers, +) + + +def awq(model_config: LlavaConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of AWQ parameters. + Parameters + ---------- + model_config : LlavaConfig + The configuration of the Llava model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to AWQ. + """ + model, _ = awq_quant(model_config, quantization) + _, _named_params = model.export_tvm(spec=model.get_default_spec()) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.text_config.num_hidden_layers): + # Add QKV in self attention + attn = f"language_model.model.layers.{i}.self_attn" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{attn}.qkv_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{quantize_suffix}", + f"{attn}.k_proj.{quantize_suffix}", + f"{attn}.v_proj.{quantize_suffix}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Concat gate and up in MLP + mlp = f"language_model.model.layers.{i}.mlp" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{mlp}.gate_up_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.gate_proj.{quantize_suffix}", + f"{mlp}.up_proj.{quantize_suffix}", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # inv_freq is not used in the model + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + return mapping diff --git a/python/mlc_llm/model/llava/llava_model.py b/python/mlc_llm/model/llava/llava_model.py new file mode 100644 index 0000000..9acdf98 --- /dev/null +++ b/python/mlc_llm/model/llava/llava_model.py @@ -0,0 +1,334 @@ +""" +Implementation of LLaVa Model +Implements the CLIP Vision Encoder. Uses Llama for the Language Encoder. +""" + +import dataclasses +import logging +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Module, Tensor +from tvm.relax.frontend.nn.op import permute_dims, reshape, wrap_nested +from tvm.relax.op import strided_slice + +from mlc_llm import op as op_ext +from mlc_llm.model.model_preset import MODEL_PRESETS +from mlc_llm.model.vision import CLIPVisionConfig, CLIPVisionModel, ImageProcessor +from mlc_llm.nn import PagedKVCache, RopeMode + +from ...support.config import ConfigBase +from ..llama.llama_model import LlamaConfig, LlamaForCausalLM +from ..mistral.mistral_model import MistralConfig, MistralForCausalLM + +logger = logging.getLogger(__name__) + + +CONFIG_MAP = {"LlamaForCausalLM": LlamaConfig, "MistralForCausalLM": MistralConfig} +ARCHITECTURE_MAP = { + "LlamaForCausalLM": LlamaForCausalLM, + "MistralForCausalLM": MistralForCausalLM, +} + + +@dataclasses.dataclass +class LlavaConfig(ConfigBase): + """ + LLaVa Config + """ + + image_token_index: int + text_config: LlamaConfig + vision_config: CLIPVisionConfig + vocab_size: int + context_window_size: int = -1 + sliding_window_size: int = -1 + prefill_chunk_size: int = -1 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + text_architecture: str = "LlamaForCausalLM" + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self) -> None: + vision_config_dict: Dict[str, Any] # noqa: UP006 + if isinstance(self.vision_config, CLIPVisionConfig): + vision_config_dict = dataclasses.asdict(self.vision_config) + else: + vision_config_dict = dict(self.vision_config) + + for k, v in vision_config_dict.pop("kwargs", {}).items(): + vision_config_dict[k] = v + + self.vision_config = CLIPVisionConfig.from_dict(vision_config_dict) + + text_config_dict: Dict[str, Any] # noqa: UP006 + if isinstance(self.text_config, ConfigBase): + text_config_dict = dataclasses.asdict(self.text_config) + else: + text_config_dict = dict(self.text_config) + + if "_name_or_path" in text_config_dict: + hf_config = self.get_hf_config(text_config_dict) + text_config_dict.update(hf_config) + architectures = text_config_dict["architectures"] + assert len(architectures) == 1 + self.text_architecture = architectures[0] + else: + for k, v in text_config_dict.pop("kwargs", {}).items(): + text_config_dict[k] = v + + self.text_config = CONFIG_MAP[self.text_architecture].from_dict(text_config_dict) + + for k in ["context_window_size", "sliding_window_size", "prefill_chunk_size"]: + if getattr(self, k) <= 0: + if hasattr(self.text_config, k): + setattr(self, k, getattr(self.text_config, k)) + + def get_hf_config(self, text_config_dict: Dict[str, Any]) -> Dict[str, Any]: # noqa: UP006 + """ + Get the Hugging Face config of the text model + """ + + hf_config: Dict[str, Any] # noqa: UP006 + try: + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(text_config_dict["_name_or_path"]).to_dict() + except (ImportError, OSError) as e: + # If transformers is not installed, get the config from preset + # Llama2 is gated so it throws an OSError. Get the config from preset instead + preset_mapping = { + "meta-llama/Llama-2-7b-hf": "llama2_7b", + "meta-llama/Llama-2-13b-hf": "llama2_13b", + "lmsys/vicuna-7b-v1.5": "llama2_7b", + "mistralai/Mistral-7B-v0.1": "mistral_7b", + } + if text_config_dict["_name_or_path"] in preset_mapping: + hf_config = MODEL_PRESETS[preset_mapping[text_config_dict["_name_or_path"]]] + else: + raise ValueError("Unsupported text model") from e + + return hf_config + + +class LlavaMultiModalProjector(nn.Module): + def __init__(self, config: LlavaConfig): + super().__init__() + + self.linear_1 = nn.Linear( + config.vision_config.hidden_size, config.text_config.hidden_size, bias=True + ) + self.act = nn.GELU() + self.linear_2 = nn.Linear( + config.text_config.hidden_size, config.text_config.hidden_size, bias=True + ) + + def forward(self, image_features: Tensor) -> Tensor: + hidden_states = self.linear_1(image_features) + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +class LlavaForCausalLM(Module): + def __init__(self, config: LlavaConfig): + super().__init__() + self.config = config + self.vision_tower = CLIPVisionModel(config.vision_config) + self.image_processor = ImageProcessor() + self.multi_modal_projector = LlavaMultiModalProjector(config) + self.language_model = ARCHITECTURE_MAP[config.text_architecture](config.text_config) + self.vocab_size = config.vocab_size + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + self.language_model.to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def embed(self, input_ids: Tensor) -> Tensor: + return self.language_model.embed(input_ids) + + def image_preprocess(self, pixel_values: Tensor) -> Tensor: + pixel_values = permute_dims(pixel_values, axes=(0, 3, 1, 2)) # NHWC -> NCHW + pixel_values = self.image_processor.resize( + pixel_values, + { + "shortest_edge": self.config.vision_config.image_size, + }, + ) + pixel_values = self.image_processor.crop( + pixel_values, + { + "height": self.config.vision_config.image_size, + "width": self.config.vision_config.image_size, + }, + ) + pixel_values = self.image_processor.rescale(pixel_values) + pixel_values = self.image_processor.normalize(pixel_values) + return pixel_values + + def image_embed(self, pixel_values: Tensor) -> Tensor: + pixel_values = self.image_preprocess(pixel_values) + pixel_values = pixel_values.astype(self.dtype) + image_features_all = self.vision_tower.forward(pixel_values) + image_features = wrap_nested( + strided_slice( + image_features_all._expr, + axes=[1], + begin=[1], + end=[image_features_all.shape[1]], + ), + name="slice", + ) + image_features = self.multi_modal_projector(image_features) + image_features = reshape(image_features, shape=(-1, self.config.text_config.hidden_size)) + return image_features + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + return self.language_model.batch_forward(input_embeds, paged_kv_cache, logit_positions) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + return self.language_model.prefill(input_embed, paged_kv_cache) + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + return self.language_model.decode(input_embed, paged_kv_cache) + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + return self.language_model.batch_prefill(input_embeds, logit_positions, paged_kv_cache) + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + return self.language_model.batch_decode(input_embeds, paged_kv_cache) + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + return self.language_model.batch_verify(input_embeds, paged_kv_cache) + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.config.text_config.num_hidden_layers, + num_attention_heads=self.config.text_config.num_attention_heads + // self.config.tensor_parallel_shards, + num_key_value_heads=self.config.text_config.num_key_value_heads + // self.config.tensor_parallel_shards, + qk_head_dim=self.config.text_config.head_dim, + v_head_dim=self.config.text_config.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.language_model.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "image_embed": { + "pixel_values": nn.spec.Tensor( + [1, "image_height", "image_width", 3], + "uint8", + ), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor( + [1, "seq_len", self.config.text_config.hidden_size], self.dtype + ), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor( + [1, 1, self.config.text_config.hidden_size], self.dtype + ), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor( + [1, "seq_len", self.config.text_config.hidden_size], self.dtype + ), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor( + ["batch_size", 1, self.config.text_config.hidden_size], self.dtype + ), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor( + [1, "seq_len", self.config.text_config.hidden_size], self.dtype + ), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/medusa/__init__.py b/python/mlc_llm/model/medusa/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/medusa/medusa_loader.py b/python/mlc_llm/model/medusa/medusa_loader.py new file mode 100644 index 0000000..ce18a1d --- /dev/null +++ b/python/mlc_llm/model/medusa/medusa_loader.py @@ -0,0 +1,14 @@ +""" +This file specifies how MLC's Medusa parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .medusa_model import MedusaModel + +huggingface = make_standard_hf_loader( + model_cls=MedusaModel, + include_qkv=False, + include_gate_up=False, +) diff --git a/python/mlc_llm/model/medusa/medusa_model.py b/python/mlc_llm/model/medusa/medusa_model.py new file mode 100644 index 0000000..12f5b47 --- /dev/null +++ b/python/mlc_llm/model/medusa/medusa_model.py @@ -0,0 +1,81 @@ +"""Medusa model definition.""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm.relax.frontend import nn + +from mlc_llm.support import logging +from mlc_llm.support.config import ConfigBase + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class MedusaConfig(ConfigBase): + """Configuration of the Llama model.""" + + medusa_num_heads: int + medusa_num_layers: int + hidden_size: int + vocab_size: int + max_batch_size: int = 1 + tensor_parallel_shards: int = 1 + + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + # Unused parameters. Kept for compatibility with the compilation flow. + prefill_chunk_size: int = -1 + context_window_size: int = -1 + + +class ResBlock(nn.Module): + """Residual block with SiLU activation.""" + + def __init__(self, hidden_size): + super().__init__() + self.linear = nn.Linear(hidden_size, hidden_size) + self.act = nn.SiLU() + + def forward(self, x): + return x + self.act(self.linear(x)) + + +class MedusaModel(nn.Module): + """Medusa model definition.""" + + def __init__(self, config: MedusaConfig): + self.hidden_size = config.hidden_size + self.dtype = "float32" + self.medusa_head = nn.ModuleList( + [ + nn.ModuleList( + [ResBlock(config.hidden_size) for _ in range(config.medusa_num_layers)] + + [nn.Linear(config.hidden_size, config.vocab_size, bias=False)] + ) + for _ in range(config.medusa_num_heads) + ] + ) + + def get_default_spec(self): + mod_spec = { + "get_logits": { + "hidden_states": nn.spec.Tensor(["batch_size", self.hidden_size], self.dtype), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) + + def get_logits(self, hidden_states: nn.Tensor): + logits = [] + for head in self.medusa_head: + logits.append(head(hidden_states).astype("float32")) + return logits + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype diff --git a/python/mlc_llm/model/minicpm/__init__.py b/python/mlc_llm/model/minicpm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/minicpm/minicpm_loader.py b/python/mlc_llm/model/minicpm/minicpm_loader.py new file mode 100644 index 0000000..f85fe72 --- /dev/null +++ b/python/mlc_llm/model/minicpm/minicpm_loader.py @@ -0,0 +1,150 @@ +""" +This file specifies how MLC's MiniCPM parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .minicpm_model import MiniCPMConfig, MiniCPMForCausalLM + + +def huggingface(model_config: MiniCPMConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : MiniCPMConfig + The configuration of the MiniCPM model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = MiniCPMForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # map attention weight + attn = f"model.layers.{i}.self_attn" + for weight_type in ["weight"]: + mlc_name = f"{attn}.wqkv_pack.{weight_type}" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{weight_type}", + f"{attn}.k_proj.{weight_type}", + f"{attn}.v_proj.{weight_type}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + if model_config.num_experts == 0: + for i in range(model_config.num_hidden_layers): + # map mlp weight + mlp = f"model.layers.{i}.mlp" + mlc_name = f"{mlp}.gate_up_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.gate_proj.weight", + f"{mlp}.up_proj.weight", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + else: + for i in range(model_config.num_hidden_layers): + # map mlp weight + mlp = f"model.layers.{i}.mlp" + mlc_mlp = f"model.layers.{i}.mlp" + mlc_name = f"{mlc_mlp}.e1_e3.weight" + mlc_param = named_parameters[mlc_name] + + def combine_expert_gate_up(*hf_params, dtype): + stack = [] + for i in range(0, len(hf_params), 2): + stack.append(np.concatenate([hf_params[i], hf_params[i + 1]], axis=0)) + return np.stack(stack, axis=0).astype(dtype) + + mapping.add_mapping( + mlc_name, + functools.reduce( + lambda a, b: a + b, + [ + [ + f"{mlp}.experts.{expert_id}.w1.weight", + f"{mlp}.experts.{expert_id}.w3.weight", + ] + for expert_id in range(model_config.num_experts) + ], + ), + functools.partial( + combine_expert_gate_up, + dtype=mlc_param.dtype, + ), + ) + + mlc_name = f"{mlc_mlp}.e2.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.experts.{expert_id}.w2.weight" + for expert_id in range(model_config.num_experts) + ], + functools.partial( + lambda *hf_params, dtype: np.stack(hf_params, axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + mlc_name = f"{mlc_mlp}.gate.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{mlp}.gate.weight"], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + # Skip lm_head.weight if tie_word_embeddings is enabled + if mlc_name == "lm_head.weight" and model_config.tie_word_embeddings: + continue + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/minicpm/minicpm_model.py b/python/mlc_llm/model/minicpm/minicpm_model.py new file mode 100644 index 0000000..209747c --- /dev/null +++ b/python/mlc_llm/model/minicpm/minicpm_model.py @@ -0,0 +1,520 @@ +""" +Implementation for Minicpm architecture. +""" + +import dataclasses +import math +from functools import partial +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.nn.expert import MixtralExperts +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class MiniCPMConfig(ConfigBase): + """Configuration of the MiniCPM model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + hidden_act: str + rms_norm_eps: float + intermediate_size: int + scale_emb: int + scale_depth: float + dim_model_base: int + use_cache: bool + bos_token_id: int + eos_token_id: int + tie_word_embeddings: bool = False + rope_theta: int = 10000 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + head_dim: int = 0 + max_batch_size: int = 1 + num_experts_per_tok: int = 0 + num_experts: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class MiniCPMAttention(nn.Module): + def __init__(self, config: MiniCPMConfig): + super().__init__() # Make sure to call the parent class constructor + self.hidden_size = config.hidden_size + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + + self.num_heads = config.num_attention_heads // self.tensor_parallel_shards + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads // self.tensor_parallel_shards + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.context_window_size + + self.wqkv_pack = nn.Linear( + in_features=self.hidden_size, + out_features=(self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=False, + ) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.wqkv_pack(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class MiniCPMEmbedding(nn.Embedding): + """The embedding module specialized for MiniCPM so that + it can be shared with the final lm_head. + """ + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class MiniCPMMLP(nn.Module): + def __init__(self, config: MiniCPMConfig): + self.hidden_size = config.hidden_size + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + + self.gate_up_proj = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class MiniCPMMoE(nn.Module): + def __init__(self, config: MiniCPMConfig): + self.num_local_experts = config.num_experts + self.num_experts_per_tok = config.num_experts_per_tok + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.e1_e3 = MixtralExperts( + self.num_local_experts, + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + tensor_parallel_shards=config.tensor_parallel_shards, + ) + self.e2 = MixtralExperts( + self.num_local_experts, + in_features=self.intermediate_size, + out_features=config.hidden_size, + tensor_parallel_shards=config.tensor_parallel_shards, + ) + self.dtype = "float32" + + def forward(self, x: Tensor): + def _expert_forward(x: Tensor, indptr: Tensor): + x1_x3 = self.e1_e3(x, indptr) + x1, x3 = op.split(x1_x3, indices_or_sections=2, axis=-1) + x = self.e2(op.silu(x1) * x3, indptr) + return x + + experts_per_tok = self.num_experts_per_tok # activated experts per token + local_experts = self.num_local_experts # total number of experts + batch_size, seq_len, hidden_size = x.shape + num_tokens = batch_size * seq_len + x = x.reshape(num_tokens, hidden_size) + # gate: [num_tokens, local_experts] + gate: Tensor = self.gate(x) + # expert_weights: [num_tokens, experts_per_tok] + # expert_indices: [num_tokens, experts_per_tok] + expert_weights, expert_indices = op_ext.moe_misc.gating_softmax_topk(gate, experts_per_tok) + use_ft = ( + op_ext.get_store().cutlass_group_gemm or op_ext.get_store().faster_transformer + ) and self.dtype == "float16" + if num_tokens == 1: + # x: [num_tokens * experts_per_tok, hidden_size] + x = _expert_forward(x, expert_indices) + else: + # cumsum: [num_tokens * local_experts] + cumsum = op_ext.moe_misc.moe_cumsum(expert_indices, local_experts) + # indices: [num_tokens * experts_per_tok] + reverse_indices, token_indices = op_ext.moe_misc.get_indices(cumsum, expert_indices) + if use_ft: + # indptr: [num_local_experts] + indptr = op_ext.moe_misc.get_indptr( + cumsum, local_experts, num_tokens, inclusive=True, out_dtype="int64" + ) + else: + # indptr: [num_local_experts + 1] + indptr = op_ext.moe_misc.get_indptr( + cumsum, + local_experts, + num_tokens, + inclusive=False, + out_dtype="int32", + ) + # x: [num_tokens * experts_per_tok, hidden_size] + x = op.take(x, token_indices, axis=0) + x = _expert_forward(x, indptr) + x = op_ext.moe_misc.scatter_output(x, reverse_indices) + # x: [num_tokens, experts_per_tok, hidden_size] + x = x.reshape(num_tokens, experts_per_tok, hidden_size) * expert_weights.reshape( + num_tokens, experts_per_tok, 1 + ) + # x: [num_tokens, hidden_size] + x = op_ext.moe_misc.moe_sum(x, dim=1) + x = x.reshape(batch_size, seq_len, hidden_size) + return x + + +class MiniCPMDecoderLayer(nn.Module): + def __init__(self, config: MiniCPMConfig): + self.scale_depth = config.scale_depth + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.self_attn = MiniCPMAttention(config) + self.num_experts = config.num_experts + if self.num_experts == 0: + self.mlp = MiniCPMMLP(config) + else: + self.mlp = MiniCPMMoE(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.wqkv_pack.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + if self.num_experts == 0: + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set( + self.mlp.down_proj.weight, + tp.ShardSingleDim("_shard_mlp_down", dim=1), + ) + else: + _set( + self.mlp.e1_e3.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=1), + ) + _set(self.mlp.e2.weight, tp.ShardSingleDim("_shard_mlp_down", dim=2)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, paged_kv_cache, layer_id) + hidden_states = self._apply_residual( + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers)), + residual, + ) + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self._apply_residual( + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers)), + residual, + ) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class MiniCPMModel(nn.Module): + def __init__(self, config: MiniCPMConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = MiniCPMEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [MiniCPMDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class MiniCPMForCausalLM(nn.Module): + def __init__(self, config: MiniCPMConfig): + self.model = MiniCPMModel(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.hidden_size // config.num_attention_heads + self.vocab_size = config.vocab_size + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + self.scale_emb = config.scale_emb + self.scale_width = self.hidden_size // config.dim_model_base + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) / self.scale_width + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) * self.scale_emb + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) / self.scale_width + hidden_states = index_last_token(hidden_states) + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) / self.scale_width + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/ministral3/__init__.py b/python/mlc_llm/model/ministral3/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/ministral3/ministral3_loader.py b/python/mlc_llm/model/ministral3/ministral3_loader.py new file mode 100644 index 0000000..86e602f --- /dev/null +++ b/python/mlc_llm/model/ministral3/ministral3_loader.py @@ -0,0 +1,265 @@ +""" +This file specifies how MLC's Ministral3 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools +from typing import Callable, List, Optional, Tuple # noqa: UP035 + +import numpy as np + +from mlc_llm.loader import ExternMapping, QuantizeMapping +from mlc_llm.quantization import BlockScaleQuantize, Quantization + +from .ministral3_model import Ministral3Config, Mistral3ForConditionalGeneration + + +def _dequantize_block_scale_weight( + weight: np.ndarray, + weight_scale: np.ndarray, + block_size: Tuple[int, int], # noqa: UP006 +) -> np.ndarray: + """Reconstruct float weights from FP8 block-scale storage.""" + + rows, cols = weight.shape + block_rows, block_cols = block_size + out = np.empty((rows, cols), dtype="float32") + weight = weight.astype("float32") + num_row_blocks, num_col_blocks = weight_scale.shape + for i in range(num_row_blocks): + row_start = i * block_rows + if row_start >= rows: + break + row_end = min(row_start + block_rows, rows) + scale_row = weight_scale[i] + for j in range(num_col_blocks): + col_start = j * block_cols + if col_start >= cols: + break + col_end = min(col_start + block_cols, cols) + out[row_start:row_end, col_start:col_end] = ( + weight[row_start:row_end, col_start:col_end] * scale_row[j] + ) + return out + + +def huggingface(model_config: Ministral3Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : Ministral3Config + The configuration of the Ministral3 model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Mistral3ForConditionalGeneration(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + if isinstance(quantization, BlockScaleQuantize): + # Convert the model to block-scale quantized model before loading parameters + model = quantization.quantize_model(model, QuantizeMapping({}, {}), "") + if model_config.weight_block_size is None: + raise ValueError( + "The input Ministral 3 model is not fp8 block quantized. " + "Thus BlockScaleQuantize is not supported." + ) + + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + raw_params = dict(_named_params) + if any(name.startswith("language_model.") for name in raw_params): + named_parameters = { + name.replace("language_model.", "", 1): value for name, value in raw_params.items() + } + else: + named_parameters = raw_params + + mapping = ExternMapping() + + hf_prefix = "" + if "vision_config" in model_config.kwargs: + hf_prefix = "language_model." + + def hf(name: str) -> str: + return f"{hf_prefix}{name}" + + if ( + not isinstance(quantization, BlockScaleQuantize) + and model_config.weight_block_size is not None + ): + raise ValueError( + "The input Ministral 3 model is fp8 block quantized. " + "Please use BlockScaleQuantize for the model." + ) + + # Helper function to add both weight and scale mappings + def add_weight_and_scale_mapping( + weight_mlc_name: str, + weight_hf_names: List[str], # noqa: UP006 + weight_transform_func: Callable, + activation_transform_func: Optional[Callable] = None, + ): + mlc_param = named_parameters[weight_mlc_name] + mapping.add_mapping( + weight_mlc_name, + weight_hf_names, + functools.partial(weight_transform_func, dtype=mlc_param.dtype), + ) + + if isinstance(quantization, BlockScaleQuantize): + weight_scale_mlc_name = f"{weight_mlc_name}_scale_inv" + if weight_scale_mlc_name in named_parameters: + weight_scale_hf_names = [f"{name}_scale_inv" for name in weight_hf_names] + weight_scale_param = named_parameters[weight_scale_mlc_name] + expected_weight_scale_shape = tuple(int(dim) for dim in weight_scale_param.shape) + + def _weight_scale_transform(*arrays, dtype: str, _transform=weight_transform_func): + processed = [] + for arr in arrays: + arr_np = np.asarray(arr) + if arr_np.ndim == 0: + arr_np = arr_np.reshape((1,)) + processed.append(arr_np) + result = _transform(*processed, dtype=dtype) + result = np.asarray(result, dtype=dtype) + if result.shape == expected_weight_scale_shape: + return result + if result.shape == (): + return np.full(expected_weight_scale_shape, result.item(), dtype=dtype) + if result.shape == (1,) and expected_weight_scale_shape != (1,): + return np.broadcast_to(result, expected_weight_scale_shape).astype(dtype) + if ( + result.ndim == 1 + and result.size > 1 + and len(expected_weight_scale_shape) >= 2 + and expected_weight_scale_shape[0] % result.size == 0 + ): + rows_per_segment = expected_weight_scale_shape[0] // result.size + tiled = np.repeat(result, rows_per_segment) + tiled = tiled.reshape(expected_weight_scale_shape[0], 1) + return np.broadcast_to(tiled, expected_weight_scale_shape).astype(dtype) + raise ValueError( + f"Unexpected weight scale shape {result.shape} for " + f"{weight_scale_mlc_name}, expected {expected_weight_scale_shape}" + ) + + mapping.add_mapping( + weight_scale_mlc_name, + weight_scale_hf_names, + functools.partial(_weight_scale_transform, dtype=weight_scale_param.dtype), + ) + activation_scale_mlc_name = f"{weight_mlc_name[: -len('.weight')]}.activation_scale" + if activation_scale_mlc_name in named_parameters: + activation_scale_hf_names = [ + f"{name[: -len('.weight')]}.activation_scale" for name in weight_hf_names + ] + activation_scale_param = named_parameters[activation_scale_mlc_name] + transform = activation_transform_func or weight_transform_func + expected_shape = tuple(int(dim) for dim in activation_scale_param.shape) + + def _activation_scale_transform(*arrays, dtype: str, _transform=transform): + result = _transform(*arrays, dtype=dtype) + result = np.asarray(result, dtype=dtype) + if result.shape == expected_shape: + return result + if result.shape == (): + # HF checkpoint stores a single scale; broadcast across the expected + # dimension. + return np.full(expected_shape, result.item(), dtype=dtype) + if result.shape == (1,) and expected_shape != (1,): + return np.broadcast_to(result, expected_shape).astype(dtype) + if ( + result.ndim == 1 + and result.size > 1 + and len(expected_shape) >= 1 + and expected_shape[0] % result.size == 0 + ): + rows_per_segment = expected_shape[0] // result.size + tiled = np.repeat(result, rows_per_segment) + return tiled.reshape(expected_shape).astype(dtype) + raise ValueError( + f"Unexpected activation scale shape {result.shape} for " + f"{activation_scale_mlc_name}, expected {expected_shape}" + ) + + mapping.add_mapping( + activation_scale_mlc_name, + activation_scale_hf_names, + functools.partial( + _activation_scale_transform, dtype=activation_scale_param.dtype + ), + ) + + def identity_transform(param: np.ndarray, dtype: str): + return param.astype(dtype) + + def make_shared_activation_transform(target_name: str): + def func(first: np.ndarray, *rest: np.ndarray, dtype: str): + for _, arr in enumerate(rest, start=1): + if not np.allclose(arr, first): + raise ValueError( + f"Activation scales for {target_name} must be identical between " + "concatenated sources." + ) + return first.astype(dtype) + + return func + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"model.layers.{i}.self_attn" + mlc_name = f"{attn}.qkv_proj.weight" + proj_sources = [hf(f"{attn}.{proj}.weight") for proj in ["q_proj", "k_proj", "v_proj"]] + add_weight_and_scale_mapping( + mlc_name, + proj_sources, + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + activation_transform_func=make_shared_activation_transform( + f"{mlc_name}_activation_scale" + ), + ) + + # Add gates in MLP + mlp = f"model.layers.{i}.mlp" + mlc_name = f"{mlp}.gate_up_proj.weight" + gate_sources = [hf(f"{mlp}.{proj}.weight") for proj in ["gate_proj", "up_proj"]] + add_weight_and_scale_mapping( + mlc_name, + gate_sources, + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + activation_transform_func=make_shared_activation_transform( + f"{mlc_name}_activation_scale" + ), + ) + + for linear_name in [f"{attn}.o_proj.weight", f"{mlp}.down_proj.weight"]: + add_weight_and_scale_mapping( + linear_name, + [hf(linear_name)], + identity_transform, + ) + + # inv_freq is not used in the model + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [hf(mlc_name)], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/ministral3/ministral3_model.py b/python/mlc_llm/model/ministral3/ministral3_model.py new file mode 100644 index 0000000..2461c67 --- /dev/null +++ b/python/mlc_llm/model/ministral3/ministral3_model.py @@ -0,0 +1,534 @@ +""" +Implementation for Ministral 3 architecture. +""" + +import dataclasses +import math +from functools import partial +from typing import Any, Dict, Optional, Tuple # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Ministral3Config(ConfigBase): + """Configuration of the Ministral 3 model.""" + + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + rms_norm_eps: float + vocab_size: int + attention_sink_size: int = 0 + context_window_size: int = 0 + dtype: str = "float32" + head_dim: int = 0 + hidden_act: str = "silu" + max_batch_size: int = 1 + num_key_value_heads: int = 0 + position_embedding_base: int = 0 + prefill_chunk_size: int = 0 + rope_parameters: Optional[Dict[str, Any]] = None # noqa: UP006 + sliding_window_size: int = 0 + tensor_parallel_shards: int = 1 + tie_word_embeddings: bool = False + weight_block_size: Optional[Tuple[int, int]] = None # noqa: UP006 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + modules_to_not_convert: Tuple[str, ...] = dataclasses.field(default_factory=tuple) # noqa: UP006 + + @classmethod + def from_dict( + cls, + source: Dict[str, Any], # noqa: UP006 + ) -> "Ministral3Config": + if "text_config" in source and isinstance(source["text_config"], dict): + top_level = dict(source) + text_cfg = top_level.pop("text_config") + merged: Dict[str, Any] = dict(top_level) # noqa: UP006 + merged.update(text_cfg) + if "tie_word_embeddings" in source: + merged["tie_word_embeddings"] = source["tie_word_embeddings"] + if "dtype" in source: + merged["dtype"] = source["dtype"] + return super().from_dict(merged) + return super().from_dict(source) + + def __post_init__(self): + if "quantization_config" in self.kwargs: + quantization_config = self.kwargs.pop("quantization_config") + if isinstance(quantization_config, dict): + activation_scheme = quantization_config.get("activation_scheme", "") + quant_method = quantization_config.get("quant_method", "") + weight_block_size = quantization_config.get("weight_block_size") + modules_to_not_convert = quantization_config.get("modules_to_not_convert", []) + if isinstance(modules_to_not_convert, list): + self.modules_to_not_convert = tuple(modules_to_not_convert) + if quant_method == "fp8" and activation_scheme == "static": + if weight_block_size is not None: + self.weight_block_size = weight_block_size + if ( + not isinstance(self.weight_block_size, (tuple, list)) + or len(self.weight_block_size) != 2 + ): + raise ValueError( + "Invalid Ministral3 quantization config: ", + "weight_block_size must be a list or tuple of two integers, ", + f"got {self.weight_block_size} of type", + f"{type(self.weight_block_size)}", + ) + else: + # Set default block size if not provided. + self.weight_block_size = (128, 128) + logger.info( + "Setting default weight_block_size=%s, ", + "since quantization_config does not provide ", + "FP8 block-scale details required by ", + "MLC (activation_scheme=%s, quant_method=%s)", + self.weight_block_size, + activation_scheme, + quant_method, + ) + else: + raise ValueError( + "Invalid Ministral 3 model quantization config: ", + "only FP8 static quantization is supported, ", + f"got activation_scheme={activation_scheme}, quant_method={quant_method}", + ) + else: + raise ValueError( + "Invalid Ministral 3 model quantization config: ", + "unrecognized quantization config: ", + f"{quantization_config}", + ) + + if self.position_embedding_base == 0: + if self.rope_parameters is not None and "rope_theta" in self.rope_parameters: + self.position_embedding_base = self.rope_parameters.pop("rope_theta") + elif "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.sliding_window_size == 0: + self.sliding_window_size = self.kwargs.pop("sliding_window", -1) + if self.sliding_window_size is None: + # Sliding window is disabled. + self.sliding_window_size = -1 + if self.context_window_size == 0: + if self.sliding_window_size == -1: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or " + "`max_sequence_length` is provided in `config.json`." + ) + else: + self.context_window_size = -1 + + if self.num_key_value_heads == 0: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.num_attention_heads % self.num_key_value_heads == 0 + assert self.attention_sink_size >= 0 + if self.prefill_chunk_size == 0: + prefill_chunk_size_candidates = [] + if self.sliding_window_size != -1: + prefill_chunk_size_candidates.append(self.sliding_window_size) + if self.context_window_size != -1: + prefill_chunk_size_candidates.append(self.context_window_size) + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(*prefill_chunk_size_candidates, 8192), + ) + self.prefill_chunk_size = min(*prefill_chunk_size_candidates, 8192) + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class Ministral3Embedding(nn.Embedding): + """The embedding module specialized for Ministral3 so that + it can be shared with the final lm_head. + """ + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class Ministral3MLP(nn.Module): + """Same as in Llama architecture (LlamaFFN).""" + + def __init__(self, config: Ministral3Config): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +def yarn_get_sm_scale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class Ministral3Attention(nn.Module): + """Same as LlamaAttention, but with sliding window attention using a rolling buffer cache.""" + + def __init__(self, config: Ministral3Config): + self.head_dim = config.head_dim + if config.num_key_value_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_key_value_heads} key-value attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=False, + ) + self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False) + + self.softmax_scale = self.head_dim ** (-0.5) + if config.rope_parameters is not None: + mscale_all_dim = config.rope_parameters.get("mscale_all_dim", 0) + scaling_factor = config.rope_parameters["factor"] + if mscale_all_dim: + sm_scale = yarn_get_sm_scale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * sm_scale * sm_scale + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.softmax_scale + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class Ministral3DecoderLayer(nn.Module): + """Exact same as LlamaDecoderLayer.""" + + def __init__(self, config: Ministral3Config): + rms_norm_eps = config.rms_norm_eps + self.self_attn = Ministral3Attention(config) + self.mlp = Ministral3MLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Ministral3Model(nn.Module): + """Exact same as LlamaModel.""" + + def __init__(self, config: Ministral3Config): + assert config.hidden_size % config.num_attention_heads == 0 + # self.embed_tokens = nn.Embedding("vocab_size", config.hidden_size) + self.embed_tokens = Ministral3Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Ministral3DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.tensor_parallel_shards = config.tensor_parallel_shards + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Mistral3ForConditionalGeneration(nn.Module): + def __init__(self, config: Ministral3Config): + self.model = Ministral3Model(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear( + config.hidden_size, config.vocab_size, bias=False + ) # "vocab_size" + self._mark_modules_no_quant(config.modules_to_not_convert) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.rope_parameters = config.rope_parameters + self.tensor_parallel_shards = config.tensor_parallel_shards + self.sliding_window_size = config.sliding_window_size + self.dtype = config.dtype + self.weight_block_size = config.weight_block_size + + def _mark_modules_no_quant(self, modules: Tuple[str, ...]): # noqa: UP006 + for path in modules: + if not path: + continue + parts = path.split(".") + target = self + for part in parts: + if not hasattr(target, part): + target = None + break + target = getattr(target, part) + if target is not None: + setattr(target, "no_quantization", True) + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rope_scaling=self.rope_parameters, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/mistral/__init__.py b/python/mlc_llm/model/mistral/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/mistral/mistral_loader.py b/python/mlc_llm/model/mistral/mistral_loader.py new file mode 100644 index 0000000..7876b13 --- /dev/null +++ b/python/mlc_llm/model/mistral/mistral_loader.py @@ -0,0 +1,98 @@ +""" +This file specifies how MLC's Mistral parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization, make_awq_quant + +from .mistral_model import MistralConfig, MistralForCausalLM + +awq_quant = make_awq_quant(MistralForCausalLM) + + +huggingface = make_standard_hf_loader( + model_cls=MistralForCausalLM, + add_unused=["rotary_emb.inv_freq"], +) + + +def awq(model_config: MistralConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of AWQ parameters. + Parameters + ---------- + model_config : MistralConfig + The configuration of the Mistral model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to AWQ. + """ + model, _ = awq_quant(model_config, quantization) + _, _named_params = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"model.layers.{i}.self_attn" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{attn}.qkv_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{quantize_suffix}", + f"{attn}.k_proj.{quantize_suffix}", + f"{attn}.v_proj.{quantize_suffix}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Concat gate and up in MLP + mlp = f"model.layers.{i}.mlp" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{mlp}.gate_up_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.gate_proj.{quantize_suffix}", + f"{mlp}.up_proj.{quantize_suffix}", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # inv_freq is not used in the model + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + return mapping diff --git a/python/mlc_llm/model/mistral/mistral_model.py b/python/mlc_llm/model/mistral/mistral_model.py new file mode 100644 index 0000000..345046d --- /dev/null +++ b/python/mlc_llm/model/mistral/mistral_model.py @@ -0,0 +1,390 @@ +""" +Implementation for Mistral architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class MistralConfig(ConfigBase): + """Configuration of the Mistral model.""" + + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + rms_norm_eps: float + vocab_size: int + position_embedding_base: int = 0 + num_key_value_heads: int = 0 + head_dim: int = 0 + context_window_size: int = 0 + sliding_window_size: int = 0 + prefill_chunk_size: int = 0 + attention_sink_size: int = 4 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.sliding_window_size == 0: + self.sliding_window_size = self.kwargs.pop("sliding_window", -1) + if self.sliding_window_size is None: + # Sliding window is disabled. + self.sliding_window_size = -1 + if self.context_window_size == 0: + if self.sliding_window_size == -1: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or " + "`max_sequence_length` is provided in `config.json`." + ) + else: + self.context_window_size = -1 + + if self.num_key_value_heads == 0: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.num_attention_heads % self.num_key_value_heads == 0 + assert self.attention_sink_size >= 0 + if self.prefill_chunk_size == 0: + prefill_chunk_size_candidates = [] + if self.sliding_window_size != -1: + prefill_chunk_size_candidates.append(self.sliding_window_size) + if self.context_window_size != -1: + prefill_chunk_size_candidates.append(self.context_window_size) + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(*prefill_chunk_size_candidates, 8192), + ) + self.prefill_chunk_size = min(*prefill_chunk_size_candidates, 8192) + + +class MistralMLP(nn.Module): + """Same as in Llama architecture (LlamaFFN).""" + + def __init__(self, config: MistralConfig): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class MistralAttention(nn.Module): + """Same as LlamaAttention, but with sliding window attention using a rolling buffer cache.""" + + def __init__(self, config: MistralConfig): + self.head_dim = config.head_dim + if config.num_key_value_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_key_value_heads} key-value attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=False, + ) + self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class MistralDecoderLayer(nn.Module): + """Exact same as LlamaDecoderLayer.""" + + def __init__(self, config: MistralConfig): + rms_norm_eps = config.rms_norm_eps + self.self_attn = MistralAttention(config) + self.mlp = MistralMLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class MistralModel(nn.Module): + """Exact same as LlamaModel.""" + + def __init__(self, config: MistralConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = nn.Embedding("vocab_size", config.hidden_size) + self.layers = nn.ModuleList( + [MistralDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.tensor_parallel_shards = config.tensor_parallel_shards + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class MistralForCausalLM(nn.Module): + """Same as LlamaForCausalLM, except for the use of sliding window attention.""" + + def __init__(self, config: MistralConfig): + self.model = MistralModel(config) + self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.sliding_window_size = config.sliding_window_size + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/mixtral/__init__.py b/python/mlc_llm/model/mixtral/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/mixtral/mixtral_loader.py b/python/mlc_llm/model/mixtral/mixtral_loader.py new file mode 100644 index 0000000..a9bb312 --- /dev/null +++ b/python/mlc_llm/model/mixtral/mixtral_loader.py @@ -0,0 +1,130 @@ +""" +This file specifies how MLC's Mixtral parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .mixtral_model import MixtralConfig, MixtralForCausalLM + + +def huggingface(model_config: MixtralConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : MixtralConfig + The configuration of the Mixtral model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = MixtralForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"model.layers.{i}.self_attn" + mlc_name = f"{attn}.qkv_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.weight", + f"{attn}.k_proj.weight", + f"{attn}.v_proj.weight", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Add gates in MLP (when MoE is enabled) + mlp = f"model.layers.{i}.block_sparse_moe" + mlc_mlp = f"model.layers.{i}.moe" + mlc_name = f"{mlc_mlp}.e1_e3.weight" + mlc_param = named_parameters[mlc_name] + + def combine_expert_gate_up(*hf_params, dtype): + stack = [] + for i in range(0, len(hf_params), 2): + stack.append(np.concatenate([hf_params[i], hf_params[i + 1]], axis=0)) + return np.stack(stack, axis=0).astype(dtype) + + mapping.add_mapping( + mlc_name, + functools.reduce( + lambda a, b: a + b, + [ + [ + f"{mlp}.experts.{expert_id}.w1.weight", + f"{mlp}.experts.{expert_id}.w3.weight", + ] + for expert_id in range(model_config.num_local_experts) + ], + ), + functools.partial( + combine_expert_gate_up, + dtype=mlc_param.dtype, + ), + ) + + mlc_name = f"{mlc_mlp}.e2.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.experts.{expert_id}.w2.weight" + for expert_id in range(model_config.num_local_experts) + ], + functools.partial( + lambda *hf_params, dtype: np.stack(hf_params, axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + mlc_name = f"{mlc_mlp}.gate.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{mlp}.gate.weight"], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # inv_freq is not used in the model + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/mixtral/mixtral_model.py b/python/mlc_llm/model/mixtral/mixtral_model.py new file mode 100644 index 0000000..da9b391 --- /dev/null +++ b/python/mlc_llm/model/mixtral/mixtral_model.py @@ -0,0 +1,186 @@ +"""Implementation for Mistral architecture.""" + +import dataclasses + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.llama.llama_model import ( + LlamaAttention, + LlamaConfig, + LlamaForCausalLM, + LlamaModel, +) +from mlc_llm.nn import PagedKVCache +from mlc_llm.nn.expert import MixtralExperts +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class MixtralConfig(LlamaConfig): + """Configuration of the Mixtral model.""" + + num_local_experts: int = 0 + num_experts_per_tok: int = 0 + + +class MixtralMoE(nn.Module): + """Mixture of experts""" + + def __init__(self, config: MixtralConfig): + super().__init__() + self.num_experts_per_tok = config.num_experts_per_tok + self.num_local_experts = config.num_local_experts + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MoE intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate = nn.Linear( + in_features=config.hidden_size, + out_features=config.num_local_experts, + bias=False, + ) + self.e1_e3 = MixtralExperts( + self.num_local_experts, + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + tensor_parallel_shards=config.tensor_parallel_shards, + ) + self.e2 = MixtralExperts( + self.num_local_experts, + in_features=self.intermediate_size, + out_features=config.hidden_size, + tensor_parallel_shards=config.tensor_parallel_shards, + ) + self.dtype = "float32" + + def forward(self, x: Tensor): + def _expert_forward(x: Tensor, indptr: Tensor): + x1_x3 = self.e1_e3(x, indptr) + x1, x3 = op.split(x1_x3, indices_or_sections=2, axis=-1) + x = self.e2(op.silu(x1) * x3, indptr) + return x + + experts_per_tok = self.num_experts_per_tok # activated experts per token + local_experts = self.num_local_experts # total number of experts + batch_size, seq_len, hidden_size = x.shape + num_tokens = batch_size * seq_len + x = x.reshape(num_tokens, hidden_size) + # gate: [num_tokens, local_experts] + gate: Tensor = self.gate(x) + # expert_weights: [num_tokens, experts_per_tok] + # expert_indices: [num_tokens, experts_per_tok] + expert_weights, expert_indices = op_ext.moe_misc.gating_softmax_topk(gate, experts_per_tok) + use_ft = ( + op_ext.get_store().cutlass_group_gemm or op_ext.get_store().faster_transformer + ) and self.dtype == "float16" + if num_tokens == 1: + # x: [num_tokens * experts_per_tok, hidden_size] + x = _expert_forward(x, expert_indices) + else: + # cumsum: [num_tokens * local_experts] + cumsum = op_ext.moe_misc.moe_cumsum(expert_indices, local_experts) + # indices: [num_tokens * experts_per_tok] + reverse_indices, token_indices = op_ext.moe_misc.get_indices(cumsum, expert_indices) + if use_ft: + # indptr: [num_local_experts] + indptr = op_ext.moe_misc.get_indptr( + cumsum, local_experts, num_tokens, inclusive=True, out_dtype="int64" + ) + else: + # indptr: [num_local_experts + 1] + indptr = op_ext.moe_misc.get_indptr( + cumsum, + local_experts, + num_tokens, + inclusive=False, + out_dtype="int32", + ) + # x: [num_tokens * experts_per_tok, hidden_size] + x = op.take(x, token_indices, axis=0) + x = _expert_forward(x, indptr) + x = op_ext.moe_misc.scatter_output(x, reverse_indices) + # x: [num_tokens, experts_per_tok, hidden_size] + x = x.reshape(num_tokens, experts_per_tok, hidden_size) * expert_weights.reshape( + num_tokens, experts_per_tok, 1 + ) + # x: [num_tokens, hidden_size] + x = op_ext.moe_misc.moe_sum(x, dim=1) + x = x.reshape(batch_size, seq_len, hidden_size) + return x + + +class MixtralDecoderLayer(nn.Module): + """Mixtral decoder layer""" + + def __init__(self, config: MixtralConfig): + eps = config.rms_norm_eps + self.self_attn = LlamaAttention(config) + self.moe = MixtralMoE(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.moe.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set(self.moe.e1_e3, tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=1)) + _set(self.moe.e2, tp.ShardSingleDim("_shard_mlp_down", dim=2)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, attention_mask: Tensor, total_seq_len: tirx.Var): + """Forward pass of a decoder layer; calculate attention, and add an residual connection.""" + out = self.self_attn(self.input_layernorm(hidden_states), attention_mask, total_seq_len) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.moe(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def batch_forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.moe(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class MixtralModel(LlamaModel): + """Exact same as LlamaModel.""" + + def __init__(self, config: MixtralConfig): + super().__init__(config) + self.layers = nn.ModuleList( + [MixtralDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + + +class MixtralForCausalLM(LlamaForCausalLM): + """Same as LlamaForCausalLM.""" + + def __init__(self, config: MixtralConfig): + super().__init__(config) + self.model = MixtralModel(config) diff --git a/python/mlc_llm/model/model.py b/python/mlc_llm/model/model.py new file mode 100644 index 0000000..1800bc1 --- /dev/null +++ b/python/mlc_llm/model/model.py @@ -0,0 +1,750 @@ +"""A centralized registry of all existing model architures and their configurations.""" + +import dataclasses +from typing import Any, Callable, Dict, Literal, Optional, Tuple # noqa: UP035 + +from tvm.relax.frontend import nn + +from mlc_llm.loader import ExternMapping, QuantizeMapping +from mlc_llm.quantization import make_quantization_functions +from mlc_llm.quantization.quantization import Quantization + +from .baichuan import baichuan_loader, baichuan_model +from .bert import bert_loader, bert_model +from .chatglm3 import chatglm3_loader, chatglm3_model +from .cohere import cohere_loader, cohere_model +from .deepseek import deepseek_loader, deepseek_model +from .deepseek_v2 import deepseek_v2_loader, deepseek_v2_model +from .eagle import eagle_loader, eagle_model +from .gemma import gemma_loader, gemma_model +from .gemma2 import gemma2_loader, gemma2_model +from .gemma3 import gemma3_loader, gemma3_model +from .gpt2 import gpt2_loader, gpt2_model +from .gpt_bigcode import gpt_bigcode_loader, gpt_bigcode_model +from .gpt_j import gpt_j_loader, gpt_j_model +from .gpt_neox import gpt_neox_loader, gpt_neox_model +from .internlm import internlm_loader, internlm_model +from .internlm2 import internlm2_loader, internlm2_model +from .llama import llama_loader, llama_model +from .llama4 import llama4_loader, llama4_model +from .llava import llava_loader, llava_model +from .medusa import medusa_loader, medusa_model +from .minicpm import minicpm_loader, minicpm_model +from .ministral3 import ministral3_loader, ministral3_model +from .mistral import mistral_loader, mistral_model +from .mixtral import mixtral_loader, mixtral_model +from .nemotron import nemotron_loader, nemotron_model +from .olmo import olmo_loader, olmo_model +from .olmo2 import olmo2_loader, olmo2_model +from .orion import orion_loader, orion_model +from .phi import phi_loader, phi_model +from .phi3 import phi3_loader, phi3_model +from .phi3v import phi3v_loader, phi3v_model +from .qwen import qwen_loader, qwen_model +from .qwen2 import qwen2_loader, qwen2_model +from .qwen2_moe import qwen2_moe_loader, qwen2_moe_model +from .qwen3 import qwen3_loader, qwen3_model +from .qwen3_moe import qwen3_moe_loader, qwen3_moe_model +from .qwen35 import qwen35_loader, qwen35_model +from .rwkv5 import rwkv5_loader, rwkv5_model +from .rwkv6 import rwkv6_loader, rwkv6_model +from .stable_lm import stablelm_loader, stablelm_model +from .starcoder2 import starcoder2_loader, starcoder2_model + +ModelConfig = Any +"""A ModelConfig is an object that represents a model architecture. It is required to have +a class method `from_file` with the following signature: + + def from_file(cls, path: Path) -> ModelConfig: + ... +""" + +FuncGetExternMap = Callable[[ModelConfig, Quantization], ExternMapping] +FuncQuantization = Callable[[ModelConfig, Quantization], Tuple[nn.Module, QuantizeMapping]] # noqa: UP006 + + +@dataclasses.dataclass +class EmbeddingMetadata: + """Embedding model metadata. + + Parameters + ---------- + model_type: Literal["encoder", "decoder"] + The type of the embedding model. + + pooling_strategy: Literal["cls", "mean", "last"] + The pooling strategy to use for the embedding model. + + normalize: bool = True + Default to normalize the embedding. + """ + + model_type: Literal["encoder", "decoder"] + pooling_strategy: Literal["cls", "mean", "last"] + normalize: bool = True + + +@dataclasses.dataclass +class Model: + """All about a model architecture: its configuration, its parameter loader and quantization. + + Parameters + ---------- + name : str + The name of the model. + + model : Callable[[ModelConfig], nn.Module] + A method that creates the `nn.Module` that represents the model from `ModelConfig`. + + config : ModelConfig + A class that has a `from_file` class method, whose signature is "Path -> ModelConfig". + + source : Dict[str, FuncGetExternMap] + A dictionary that maps the name of a source format to parameter mapping. + + quantize: Dict[str, FuncQuantization] + A dictionary that maps the name of a quantization method to quantized model and the + quantization parameter mapping. + + model_task: Literal["chat", "embedding"] = "chat" + A task of the model to distinguish between chat and embedding models. Default to "chat". + + embedding_metadata: Optional[EmbeddingMetadata] = None + Metadata for the embedding model. Default to None. + """ + + name: str + config: ModelConfig + model: Callable[[ModelConfig], nn.Module] + source: Dict[str, FuncGetExternMap] # noqa: UP006 + quantize: Dict[str, FuncQuantization] # noqa: UP006 + + model_task: Literal["chat", "embedding"] = "chat" + embedding_metadata: Optional[EmbeddingMetadata] = None + + def __post_init__(self): + if self.model_task == "embedding" and self.embedding_metadata is None: + raise ValueError(f"[Model] {self.name}: Embedding model must have embedding metadata.") + if self.model_task == "chat" and self.embedding_metadata is not None: + raise ValueError( + f"[Model] {self.name}: Chat model not expected to have embedding metadata." + ) + + +MODELS: Dict[str, Model] = { # noqa: UP006 + "llama": Model( + name="llama", + model=llama_model.LlamaForCausalLM, + config=llama_model.LlamaConfig, + source={ + "huggingface-torch": llama_loader.huggingface, + "huggingface-safetensor": llama_loader.huggingface, + "awq": llama_loader.awq, + }, + quantize=make_quantization_functions( + llama_model.LlamaForCausalLM, + supports_awq=True, + supports_per_tensor=True, + ), + ), + "llama4": Model( + name="llama4", + model=llama4_model.Llama4ForCausalLM, + config=llama4_model.Llama4Config, + source={ + "huggingface-torch": llama4_loader.huggingface, + "huggingface-safetensor": llama4_loader.huggingface, + }, + quantize=make_quantization_functions( + llama4_model.Llama4ForCausalLM, + supports_per_tensor=True, + ), + ), + "mistral": Model( + name="mistral", + model=mistral_model.MistralForCausalLM, + config=mistral_model.MistralConfig, + source={ + "huggingface-torch": mistral_loader.huggingface, + "huggingface-safetensor": mistral_loader.huggingface, + "awq": mistral_loader.awq, + }, + quantize=make_quantization_functions( + mistral_model.MistralForCausalLM, + ), + ), + "ministral3": Model( + name="ministral3", + model=ministral3_model.Mistral3ForConditionalGeneration, + config=ministral3_model.Ministral3Config, + source={ + "huggingface-torch": ministral3_loader.huggingface, + "huggingface-safetensor": ministral3_loader.huggingface, + }, + quantize=make_quantization_functions( + ministral3_model.Mistral3ForConditionalGeneration, + supports_block_scale=True, + ), + ), + "gemma": Model( + name="gemma", + model=gemma_model.GemmaForCausalLM, + config=gemma_model.GemmaConfig, + source={ + "huggingface-torch": gemma_loader.huggingface, + "huggingface-safetensor": gemma_loader.huggingface, + }, + quantize=make_quantization_functions( + gemma_model.GemmaForCausalLM, + supports_ft_quant=False, + ), + ), + "gemma2": Model( + name="gemma2", + model=gemma2_model.Gemma2ForCausalLM, + config=gemma2_model.Gemma2Config, + source={ + "huggingface-torch": gemma2_loader.huggingface, + "huggingface-safetensor": gemma2_loader.huggingface, + }, + quantize=make_quantization_functions( + gemma2_model.Gemma2ForCausalLM, + supports_ft_quant=False, + ), + ), + "gemma3": Model( + name="gemma3", + model=gemma3_model.Gemma3ForCausalLM, + config=gemma3_model.Gemma3Config, + source={ + "huggingface-torch": gemma3_loader.huggingface, + "huggingface-safetensor": gemma3_loader.huggingface, + }, + quantize=make_quantization_functions( + gemma3_model.Gemma3ForCausalLM, + supports_ft_quant=False, + ), + ), + "gemma3_text": Model( + name="gemma3_text", + model=gemma3_model.Gemma3ForCausalLM, + config=gemma3_model.Gemma3Config, + source={ + "huggingface-torch": gemma3_loader.huggingface, + "huggingface-safetensor": gemma3_loader.huggingface, + }, + quantize=make_quantization_functions( + gemma3_model.Gemma3ForCausalLM, + supports_ft_quant=False, + ), + ), + "gpt2": Model( + name="gpt2", + model=gpt2_model.GPT2LMHeadModel, + config=gpt2_model.GPT2Config, + source={ + "huggingface-torch": gpt2_loader.huggingface, + "huggingface-safetensor": gpt2_loader.huggingface, + }, + quantize=make_quantization_functions( + gpt2_model.GPT2LMHeadModel, + ), + ), + "mixtral": Model( + name="mixtral", + model=mixtral_model.MixtralForCausalLM, + config=mixtral_model.MixtralConfig, + source={ + "huggingface-torch": mixtral_loader.huggingface, + "huggingface-safetensor": mixtral_loader.huggingface, + }, + quantize=make_quantization_functions( + mixtral_model.MixtralForCausalLM, + supports_awq=True, + awq_unsupported_message="AWQ is not implemented for Mixtral models.", + supports_per_tensor=True, + ), + ), + "gpt_neox": Model( + name="gpt_neox", + model=gpt_neox_model.GPTNeoXForCausalLM, + config=gpt_neox_model.GPTNeoXConfig, + source={ + "huggingface-torch": gpt_neox_loader.huggingface, + "huggingface-safetensor": gpt_neox_loader.huggingface, + }, + quantize=make_quantization_functions( + gpt_neox_model.GPTNeoXForCausalLM, + ), + ), + "gpt_bigcode": Model( + name="gpt_bigcode", + model=gpt_bigcode_model.GPTBigCodeForCausalLM, + config=gpt_bigcode_model.GPTBigCodeConfig, + source={ + "huggingface-torch": gpt_bigcode_loader.huggingface, + "huggingface-safetensor": gpt_bigcode_loader.huggingface, + }, + quantize=make_quantization_functions( + gpt_bigcode_model.GPTBigCodeForCausalLM, + ), + ), + "phi-msft": Model( + name="phi-msft", + model=phi_model.PhiForCausalLM, + config=phi_model.PhiConfig, + source={ + "huggingface-torch": phi_loader.huggingface, + "huggingface-safetensor": phi_loader.huggingface, + }, + quantize=make_quantization_functions( + phi_model.PhiForCausalLM, + ), + ), + "phi": Model( + name="phi", + model=phi_model.PhiForCausalLM, + config=phi_model.Phi1Config, + source={ + "huggingface-torch": phi_loader.phi1_huggingface, + "huggingface-safetensor": phi_loader.phi1_huggingface, + }, + quantize=make_quantization_functions( + phi_model.PhiForCausalLM, + ), + ), + "phi3": Model( + name="phi3", + model=phi3_model.Phi3ForCausalLM, + config=phi3_model.Phi3Config, + source={ + "huggingface-torch": phi3_loader.phi3_huggingface, + "huggingface-safetensor": phi3_loader.phi3_huggingface, + }, + quantize=make_quantization_functions( + phi3_model.Phi3ForCausalLM, + ), + ), + "phi3_v": Model( + name="phi3_v", + model=phi3v_model.Phi3VForCausalLM, + config=phi3v_model.Phi3VConfig, + source={ + "huggingface-torch": phi3v_loader.huggingface, + "huggingface-safetensor": phi3v_loader.huggingface, + }, + quantize=make_quantization_functions( + phi3v_model.Phi3VForCausalLM, + ), + ), + "qwen": Model( + name="qwen", + model=qwen_model.QWenLMHeadModel, + config=qwen_model.QWenConfig, + source={ + "huggingface-torch": qwen_loader.huggingface, + "huggingface-safetensor": qwen_loader.huggingface, + }, + quantize=make_quantization_functions( + qwen_model.QWenLMHeadModel, + ), + ), + "qwen2": Model( + name="qwen2", + model=qwen2_model.QWen2LMHeadModel, + config=qwen2_model.QWen2Config, + source={ + "huggingface-torch": qwen2_loader.huggingface, + "huggingface-safetensor": qwen2_loader.huggingface, + }, + quantize=make_quantization_functions( + qwen2_model.QWen2LMHeadModel, + ), + ), + "qwen2_moe": Model( + name="qwen2_moe", + model=qwen2_moe_model.Qwen2MoeForCausalLM, + config=qwen2_moe_model.Qwen2MoeConfig, + source={ + "huggingface-torch": qwen2_moe_loader.huggingface, + "huggingface-safetensor": qwen2_moe_loader.huggingface, + }, + quantize=make_quantization_functions( + qwen2_moe_model.Qwen2MoeForCausalLM, + ), + ), + "qwen3": Model( + name="qwen3", + model=qwen3_model.Qwen3LMHeadModel, + config=qwen3_model.Qwen3Config, + source={ + "huggingface-torch": qwen3_loader.huggingface, + "huggingface-safetensor": qwen3_loader.huggingface, + }, + quantize=make_quantization_functions( + qwen3_model.Qwen3LMHeadModel, + supports_block_scale=True, + ), + ), + "qwen3-embedding": Model( + name="qwen3-embedding", + model=qwen3_model.Qwen3EmbeddingModel, + config=qwen3_model.Qwen3Config, + source={ + "huggingface-torch": qwen3_loader.huggingface_embedding, + "huggingface-safetensor": qwen3_loader.huggingface_embedding, + }, + quantize=make_quantization_functions( + qwen3_model.Qwen3EmbeddingModel, + supports_block_scale=True, + ), + model_task="embedding", + embedding_metadata=EmbeddingMetadata( + model_type="decoder", + pooling_strategy="last", + normalize=True, + ), + ), + "qwen3_5": Model( + name="qwen3_5", + model=qwen35_model.Qwen35LMHeadModel, + config=qwen35_model.Qwen35Config, + source={ + "huggingface-torch": qwen35_loader.huggingface, + "huggingface-safetensor": qwen35_loader.huggingface, + }, + quantize=make_quantization_functions( + qwen35_model.Qwen35LMHeadModel, + ), + ), + "qwen3_5_text": Model( + name="qwen3_5_text", + model=qwen35_model.Qwen35LMHeadModel, + config=qwen35_model.Qwen35Config, + source={ + "huggingface-torch": qwen35_loader.huggingface, + "huggingface-safetensor": qwen35_loader.huggingface, + }, + quantize=make_quantization_functions( + qwen35_model.Qwen35LMHeadModel, + ), + ), + "qwen3_moe": Model( + name="qwen3_moe", + model=qwen3_moe_model.Qwen3MoeForCausalLM, + config=qwen3_moe_model.Qwen3MoeConfig, + source={ + "huggingface-torch": qwen3_moe_loader.huggingface, + "huggingface-safetensor": qwen3_moe_loader.huggingface, + }, + quantize=make_quantization_functions( + qwen3_moe_model.Qwen3MoeForCausalLM, + supports_block_scale=True, + ), + ), + "deepseek_v2": Model( + name="deepseek_v2", + model=deepseek_v2_model.DeepseekV2ForCausalLM, + config=deepseek_v2_model.DeepseekV2Config, + source={ + "huggingface-torch": deepseek_v2_loader.huggingface, + "huggingface-safetensor": deepseek_v2_loader.huggingface, + }, + quantize=make_quantization_functions( + deepseek_v2_model.DeepseekV2ForCausalLM, + ), + ), + "deepseek_v3": Model( + name="deepseek_v3", + model=deepseek_v2_model.DeepseekV2ForCausalLM, + config=deepseek_v2_model.DeepseekV2Config, + source={ + "huggingface-torch": deepseek_v2_loader.huggingface, + "huggingface-safetensor": deepseek_v2_loader.huggingface, + }, + quantize=make_quantization_functions( + deepseek_v2_model.DeepseekV2ForCausalLM, + supports_block_scale=True, + ), + ), + "stablelm": Model( + name="stablelm", + model=stablelm_model.StableLmForCausalLM, + config=stablelm_model.StableLmConfig, + source={ + "huggingface-torch": stablelm_loader.huggingface, + "huggingface-safetensor": stablelm_loader.huggingface, + }, + quantize=make_quantization_functions( + stablelm_model.StableLmForCausalLM, + ), + ), + "baichuan": Model( + name="baichuan", + model=baichuan_model.BaichuanForCausalLM, + config=baichuan_model.BaichuanConfig, + source={ + "huggingface-torch": baichuan_loader.huggingface, + "huggingface-safetensor": baichuan_loader.huggingface, + }, + quantize=make_quantization_functions( + baichuan_model.BaichuanForCausalLM, + ), + ), + "internlm": Model( + name="internlm", + model=internlm_model.InternLMForCausalLM, + config=internlm_model.InternLMConfig, + source={ + "huggingface-torch": internlm_loader.huggingface, + "huggingface-safetensor": internlm_loader.huggingface, + }, + quantize=make_quantization_functions( + internlm_model.InternLMForCausalLM, + ), + ), + "internlm2": Model( + name="internlm2", + model=internlm2_model.InternLM2ForCausalLM, + config=internlm2_model.InternLM2Config, + source={ + "huggingface-torch": internlm2_loader.huggingface, + "huggingface-safetensor": internlm2_loader.huggingface, + }, + quantize=make_quantization_functions( + internlm2_model.InternLM2ForCausalLM, + ), + ), + "rwkv5": Model( + name="rwkv5", + model=rwkv5_model.RWKV5_ForCausalLM, + config=rwkv5_model.RWKV5Config, + source={ + "huggingface-torch": rwkv5_loader.huggingface, + "huggingface-safetensor": rwkv5_loader.huggingface, + }, + quantize=make_quantization_functions( + rwkv5_model.RWKV5_ForCausalLM, + ), + ), + "orion": Model( + name="orion", + model=orion_model.OrionForCausalLM, + config=orion_model.OrionConfig, + source={ + "huggingface-torch": orion_loader.huggingface, + "huggingface-safetensor": orion_loader.huggingface, + }, + quantize=make_quantization_functions( + orion_model.OrionForCausalLM, + supports_ft_quant=False, + ), + ), + "llava": Model( + name="llava", + model=llava_model.LlavaForCausalLM, + config=llava_model.LlavaConfig, + source={ + "huggingface-torch": llava_loader.huggingface, + "huggingface-safetensor": llava_loader.huggingface, + "awq": llava_loader.awq, + }, + quantize=make_quantization_functions( + llava_model.LlavaForCausalLM, + supports_awq=True, + supports_ft_quant=False, + ), + ), + "rwkv6": Model( + name="rwkv6", + model=rwkv6_model.RWKV6_ForCausalLM, + config=rwkv6_model.RWKV6Config, + source={ + "huggingface-torch": rwkv6_loader.huggingface, + "huggingface-safetensor": rwkv6_loader.huggingface, + }, + quantize=make_quantization_functions( + rwkv6_model.RWKV6_ForCausalLM, + supports_ft_quant=False, + ), + ), + "chatglm": Model( + name="chatglm", + model=chatglm3_model.ChatGLMForCausalLM, + config=chatglm3_model.GLMConfig, + source={ + "huggingface-torch": chatglm3_loader.huggingface, + "huggingface-safetensor": chatglm3_loader.huggingface, + }, + quantize=make_quantization_functions( + chatglm3_model.ChatGLMForCausalLM, + supports_ft_quant=False, + ), + ), + "eagle": Model( + name="eagle", + model=eagle_model.EagleForCausalLM, + config=eagle_model.EagleConfig, + source={ + "huggingface-torch": eagle_loader.huggingface, + "huggingface-safetensor": eagle_loader.huggingface, + "awq": eagle_loader.awq, + }, + quantize=make_quantization_functions( + eagle_model.EagleForCausalLM, + supports_awq=True, + ), + ), + "bert": Model( + name="bert", + model=bert_model.BertModel, + config=bert_model.BertConfig, + source={ + "huggingface-torch": bert_loader.huggingface, + "huggingface-safetensor": bert_loader.huggingface, + }, + quantize=make_quantization_functions( + bert_model.BertModel, + ), + model_task="embedding", + embedding_metadata=EmbeddingMetadata( + model_type="encoder", + pooling_strategy="cls", + normalize=True, + ), + ), + "medusa": Model( + name="medusa", + model=medusa_model.MedusaModel, + config=medusa_model.MedusaConfig, + source={ + "huggingface-torch": medusa_loader.huggingface, + "huggingface-safetensor": medusa_loader.huggingface, + }, + quantize=make_quantization_functions( + medusa_model.MedusaModel, + supports_group_quant=False, + supports_ft_quant=False, + ), + ), + "starcoder2": Model( + name="starcoder2", + model=starcoder2_model.Starcoder2ForCausalLM, + config=starcoder2_model.Starcoder2Config, + source={ + "huggingface-torch": starcoder2_loader.huggingface, + "huggingface-safetensor": starcoder2_loader.huggingface, + }, + quantize=make_quantization_functions( + starcoder2_model.Starcoder2ForCausalLM, + ), + ), + "cohere": Model( + name="cohere", + model=cohere_model.CohereForCausalLM, + config=cohere_model.CohereConfig, + source={ + "huggingface-torch": cohere_loader.huggingface, + "huggingface-safetensor": cohere_loader.huggingface, + }, + quantize=make_quantization_functions( + cohere_model.CohereForCausalLM, + ), + ), + "minicpm": Model( + name="minicpm", + model=minicpm_model.MiniCPMForCausalLM, + config=minicpm_model.MiniCPMConfig, + source={ + "huggingface-torch": minicpm_loader.huggingface, + "huggingface-safetensor": minicpm_loader.huggingface, + }, + quantize=make_quantization_functions( + minicpm_model.MiniCPMForCausalLM, + ), + ), + "deepseek": Model( + name="deepseek", + model=deepseek_model.DeepseekForCausalLM, + config=deepseek_model.DeepseekConfig, + source={ + "huggingface-torch": deepseek_loader.huggingface, + "huggingface-safetensor": deepseek_loader.huggingface, + }, + quantize=make_quantization_functions( + deepseek_model.DeepseekForCausalLM, + ), + ), + "gptj": Model( + name="gptj", + model=gpt_j_model.GPTJForCausalLM, + config=gpt_j_model.GPTJConfig, + source={ + "huggingface-torch": gpt_j_loader.huggingface, + "huggingface-safetensor": gpt_j_loader.huggingface, + }, + quantize=make_quantization_functions( + gpt_j_model.GPTJForCausalLM, + ), + ), + "olmo": Model( + name="olmo", + model=olmo_model.OLMoForCausalLM, + config=olmo_model.OLMoConfig, + source={ + "huggingface-torch": olmo_loader.huggingface, + "huggingface-safetensor": olmo_loader.huggingface, + "awq": olmo_loader.awq, + }, + quantize=make_quantization_functions( + olmo_model.OLMoForCausalLM, + supports_awq=True, + supports_per_tensor=True, + ), + ), + "olmo2": Model( + name="olmo2", + model=olmo2_model.OLMo2ForCausalLM, + config=olmo2_model.OLMo2Config, + source={ + "huggingface-torch": olmo2_loader.huggingface, + "huggingface-safetensor": olmo2_loader.huggingface, + }, + quantize=make_quantization_functions( + olmo2_model.OLMo2ForCausalLM, + supports_per_tensor=True, + ), + ), + "nemotron": Model( + name="nemotron", + model=nemotron_model.NemotronForCausalLM, + config=nemotron_model.NemotronConfig, + source={ + "huggingface-torch": nemotron_loader.huggingface, + "huggingface-safetensor": nemotron_loader.huggingface, + }, + quantize=make_quantization_functions( + nemotron_model.NemotronForCausalLM, + supports_awq=True, + supports_per_tensor=True, + ), + ), + "bert-bge": Model( + name="bert-bge", + model=bert_model.BertModel, + config=bert_model.BertConfig, + source={ + "huggingface-torch": bert_loader.huggingface_bge, + "huggingface-safetensor": bert_loader.huggingface_bge, + }, + quantize=make_quantization_functions( + bert_model.BertModel, + ), + model_task="embedding", + embedding_metadata=EmbeddingMetadata( + model_type="encoder", + pooling_strategy="cls", + normalize=True, + ), + ), +} diff --git a/python/mlc_llm/model/model_preset.py b/python/mlc_llm/model/model_preset.py new file mode 100644 index 0000000..71dccc3 --- /dev/null +++ b/python/mlc_llm/model/model_preset.py @@ -0,0 +1,2274 @@ +"""A builtin set of models available in MLC LLM.""" + +from typing import Any, Dict # noqa: UP035 + +MODEL_PRESETS: Dict[str, Any] = { # noqa: UP006 + "llama2_7b": { + "architectures": ["LlamaForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 2048, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 32, + "pad_token_id": 0, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.31.0.dev0", + "use_cache": True, + "vocab_size": 32000, + "context_window_size": 2048, + "prefill_chunk_size": 2048, + }, + "llama2_13b": { + "_name_or_path": "meta-llama/Llama-2-13b-hf", + "architectures": ["LlamaForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 5120, + "initializer_range": 0.02, + "intermediate_size": 13824, + "max_position_embeddings": 2048, + "model_type": "llama", + "num_attention_heads": 40, + "num_hidden_layers": 40, + "num_key_value_heads": 40, + "pad_token_id": 0, + "pretraining_tp": 2, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.31.0.dev0", + "use_cache": True, + "vocab_size": 32000, + "context_window_size": 2048, + "prefill_chunk_size": 2048, + }, + "llama2_70b": { + "architectures": ["LlamaForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 8192, + "initializer_range": 0.02, + "intermediate_size": 28672, + "max_position_embeddings": 2048, + "model_type": "llama", + "num_attention_heads": 64, + "num_hidden_layers": 80, + "num_key_value_heads": 8, + "pad_token_id": 0, + "rms_norm_eps": 1e-05, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.31.0.dev0", + "use_cache": True, + "vocab_size": 32000, + "context_window_size": 2048, + "prefill_chunk_size": 2048, + }, + "codellama_7b": { + "_name_or_path": "codellama/CodeLlama-7b-hf", + "architectures": ["LlamaForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 16384, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 32, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 1000000, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.33.0.dev0", + "use_cache": True, + "vocab_size": 32016, + "context_window_size": 2048, + "prefill_chunk_size": 2048, + }, + "codellama_13b": { + "architectures": ["LlamaForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 5120, + "initializer_range": 0.02, + "intermediate_size": 13824, + "max_position_embeddings": 16384, + "model_type": "llama", + "num_attention_heads": 40, + "num_hidden_layers": 40, + "num_key_value_heads": 40, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 1000000, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.32.0.dev0", + "use_cache": True, + "vocab_size": 32016, + "context_window_size": 2048, + "prefill_chunk_size": 2048, + }, + "codellama_34b": { + "architectures": ["LlamaForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 8192, + "initializer_range": 0.02, + "intermediate_size": 22016, + "max_position_embeddings": 16384, + "model_type": "llama", + "num_attention_heads": 64, + "num_hidden_layers": 48, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 1000000, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.32.0.dev0", + "use_cache": True, + "vocab_size": 32016, + "context_window_size": 2048, + "prefill_chunk_size": 2048, + }, + "tinyllama_1b_chat_v0.4": { + "_name_or_path": "/data/tianduo/tinyllama-ft/checkpoint-3890", + "architectures": ["LlamaForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 5632, + "max_position_embeddings": 2048, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 22, + "num_key_value_heads": 4, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 10000.0, + "tie_word_embeddings": False, + "torch_dtype": "float32", + "transformers_version": "4.33.1", + "use_cache": False, + "vocab_size": 32003, + }, + "tinyllama_1b_chat_v1.0": { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 5632, + "max_position_embeddings": 2048, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 22, + "num_key_value_heads": 4, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 10000.0, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.35.0", + "use_cache": True, + "vocab_size": 32000, + }, + "mistral_7b": { + "architectures": ["MistralForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 32768, + "model_type": "mistral", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-05, + "rope_theta": 10000.0, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.34.0.dev0", + "use_cache": True, + "vocab_size": 32000, + "sliding_window_size": 4096, + "prefill_chunk_size": 128, + "attention_sink_size": 4, + }, + "mistral_7b_v03": { + "architectures": ["MistralForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 32768, + "model_type": "mistral", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-05, + "rope_theta": 1000000.0, + "sliding_window": None, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.0.dev0", + "use_cache": True, + "vocab_size": 32768, + }, + "ministral3_3b_2512": { + "architectures": ["Mistral3ForConditionalGeneration"], + "dtype": "bfloat16", + "image_token_index": 10, + "model_type": "ministral3", + "multimodal_projector_bias": False, + "projector_hidden_act": "gelu", + "spatial_merge_size": 2, + "text_config": { + "attention_dropout": 0.0, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "intermediate_size": 9216, + "max_position_embeddings": 262144, + "model_type": "ministral3", + "num_attention_heads": 32, + "num_hidden_layers": 26, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-05, + "rope_parameters": { + "beta_fast": 32.0, + "beta_slow": 1.0, + "factor": 16.0, + "llama_4_scaling_beta": 0.1, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 16384, + "rope_theta": 1000000.0, + "rope_type": "yarn", + "type": "yarn", + }, + "sliding_window": None, + "tie_word_embeddings": True, + "use_cache": True, + "vocab_size": 131072, + }, + "transformers_version": "5.0.0.dev0", + "vision_config": { + "attention_dropout": 0.0, + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 1024, + "image_size": 1540, + "initializer_range": 0.02, + "intermediate_size": 4096, + "model_type": "pixtral", + "num_attention_heads": 16, + "num_channels": 3, + "num_hidden_layers": 24, + "patch_size": 14, + "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"}, + "rope_theta": 10000.0, + }, + "vision_feature_layer": -1, + }, + "gpt2": { + "activation_function": "gelu_new", + "architectures": ["GPT2LMHeadModel"], + "attn_pdrop": 0.1, + "bos_token_id": 50256, + "embd_pdrop": 0.1, + "eos_token_id": 50256, + "initializer_range": 0.02, + "layer_norm_epsilon": 1e-05, + "model_type": "gpt2", + "n_ctx": 1024, + "n_embd": 768, + "n_head": 12, + "n_layer": 12, + "n_positions": 1024, + "resid_pdrop": 0.1, + "summary_activation": None, + "summary_first_dropout": 0.1, + "summary_proj_to_labels": True, + "summary_type": "cls_index", + "summary_use_proj": True, + "task_specific_params": {"text-generation": {"do_sample": True, "max_length": 50}}, + "vocab_size": 50257, + }, + "gpt2_medium": { + "activation_function": "gelu_new", + "architectures": ["GPT2LMHeadModel"], + "attn_pdrop": 0.1, + "bos_token_id": 50256, + "embd_pdrop": 0.1, + "eos_token_id": 50256, + "initializer_range": 0.02, + "layer_norm_epsilon": 1e-05, + "model_type": "gpt2", + "n_ctx": 1024, + "n_embd": 1024, + "n_head": 16, + "n_layer": 24, + "n_positions": 1024, + "n_special": 0, + "predict_special_tokens": True, + "resid_pdrop": 0.1, + "summary_activation": None, + "summary_first_dropout": 0.1, + "summary_proj_to_labels": True, + "summary_type": "cls_index", + "summary_use_proj": True, + "task_specific_params": {"text-generation": {"do_sample": True, "max_length": 50}}, + "vocab_size": 50257, + }, + "gpt_bigcode": { + "activation_function": "gelu_pytorch_tanh", + "architectures": ["GPTBigCodeForCausalLM"], + "attention_softmax_in_fp32": True, + "multi_query": True, + "attn_pdrop": 0.1, + "bos_token_id": 49152, + "embd_pdrop": 0.1, + "eos_token_id": 49152, + "initializer_range": 0.02, + "layer_norm_epsilon": 1e-05, + "model_type": "gpt_bigcode", + "n_embd": 2048, + "n_head": 16, + "n_inner": 8192, + "n_layer": 24, + "n_positions": 2048, + "resid_pdrop": 0.1, + "runner_max_sequence_length": None, + "scale_attention_softmax_in_fp32": True, + "scale_attn_weights": True, + "summary_activation": None, + "summary_first_dropout": 0.1, + "summary_proj_to_labels": True, + "summary_type": "cls_index", + "summary_use_proj": True, + "transformers_version": "4.28.0.dev0", + "use_cache": True, + "vocab_size": 49280, + }, + "Mixtral-8x7B-v0.1": { + "architectures": ["MixtralForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 32768, + "model_type": "mixtral", + "num_attention_heads": 32, + "num_experts_per_tok": 2, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "num_local_experts": 8, + "output_router_logits": False, + "rms_norm_eps": 1e-05, + "rope_theta": 1000000.0, + "router_aux_loss_coef": 0.02, + "sliding_window": None, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.36.0.dev0", + "use_cache": True, + "vocab_size": 32000, + }, + "redpajama_3b_v1": { + "_name_or_path": "/root/fm/models/rp_3b_800b_real_fp16", + "architectures": ["GPTNeoXForCausalLM"], + "bos_token_id": 0, + "eos_token_id": 0, + "hidden_act": "gelu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 10240, + "layer_norm_eps": 1e-05, + "max_position_embeddings": 2048, + "model_type": "gpt_neox", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "rotary_emb_base": 10000, + "rotary_pct": 1.0, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.28.1", + "use_cache": True, + "use_parallel_residual": False, + "vocab_size": 50432, + }, + "phi-1_5": { + "_name_or_path": "microsoft/phi-1_5", + "activation_function": "gelu_new", + "architectures": ["PhiForCausalLM"], + "attn_pdrop": 0.0, + "auto_map": { + "AutoConfig": "configuration_phi.PhiConfig", + "AutoModelForCausalLM": "modeling_phi.PhiForCausalLM", + }, + "embd_pdrop": 0.0, + "flash_attn": False, + "flash_rotary": False, + "fused_dense": False, + "initializer_range": 0.02, + "layer_norm_epsilon": 1e-05, + "model_type": "phi-msft", + "n_embd": 2048, + "n_head": 32, + "n_head_kv": None, + "n_inner": None, + "n_layer": 24, + "n_positions": 2048, + "resid_pdrop": 0.0, + "rotary_dim": 32, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.34.1", + "vocab_size": 51200, + }, + "phi-2": { + "_name_or_path": "microsoft/phi-2", + "activation_function": "gelu_new", + "architectures": ["PhiForCausalLM"], + "attn_pdrop": 0.0, + "auto_map": { + "AutoConfig": "configuration_phi.PhiConfig", + "AutoModelForCausalLM": "modeling_phi.PhiForCausalLM", + }, + "embd_pdrop": 0.0, + "flash_attn": False, + "flash_rotary": False, + "fused_dense": False, + "img_processor": None, + "initializer_range": 0.02, + "layer_norm_epsilon": 1e-05, + "model_type": "phi-msft", + "n_embd": 2560, + "n_head": 32, + "n_head_kv": None, + "n_inner": None, + "n_layer": 32, + "n_positions": 2048, + "resid_pdrop": 0.1, + "rotary_dim": 32, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.35.2", + "vocab_size": 51200, + }, + "phi-3_5": { + "_name_or_path": "Phi-3.5-mini-instruct", + "architectures": ["Phi3ForCausalLM"], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_phi3.Phi3Config", + "AutoModelForCausalLM": "modeling_phi3.Phi3ForCausalLM", + }, + "bos_token_id": 1, + "embd_pdrop": 0.0, + "eos_token_id": 32000, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "intermediate_size": 8192, + "max_position_embeddings": 131072, + "model_type": "phi3", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 32, + "original_max_position_embeddings": 4096, + "pad_token_id": 32000, + "resid_pdrop": 0.0, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "long_factor": [ + 1.0800000429153442, + 1.1100000143051147, + 1.1399999856948853, + 1.340000033378601, + 1.5899999141693115, + 1.600000023841858, + 1.6200000047683716, + 2.620000123977661, + 3.2300000190734863, + 3.2300000190734863, + 4.789999961853027, + 7.400000095367432, + 7.700000286102295, + 9.09000015258789, + 12.199999809265137, + 17.670000076293945, + 24.46000099182129, + 28.57000160217285, + 30.420001983642578, + 30.840002059936523, + 32.590003967285156, + 32.93000411987305, + 42.320003509521484, + 44.96000289916992, + 50.340003967285156, + 50.45000457763672, + 57.55000305175781, + 57.93000411987305, + 58.21000289916992, + 60.1400032043457, + 62.61000442504883, + 62.62000274658203, + 62.71000289916992, + 63.1400032043457, + 63.1400032043457, + 63.77000427246094, + 63.93000411987305, + 63.96000289916992, + 63.970001220703125, + 64.02999877929688, + 64.06999969482422, + 64.08000183105469, + 64.12000274658203, + 64.41000366210938, + 64.4800033569336, + 64.51000213623047, + 64.52999877929688, + 64.83999633789062, + ], + "short_factor": [ + 1.0, + 1.0199999809265137, + 1.0299999713897705, + 1.0299999713897705, + 1.0499999523162842, + 1.0499999523162842, + 1.0499999523162842, + 1.0499999523162842, + 1.0499999523162842, + 1.0699999332427979, + 1.0999999046325684, + 1.1099998950958252, + 1.1599998474121094, + 1.1599998474121094, + 1.1699998378753662, + 1.2899998426437378, + 1.339999794960022, + 1.679999828338623, + 1.7899998426437378, + 1.8199998140335083, + 1.8499997854232788, + 1.8799997568130493, + 1.9099997282028198, + 1.9399996995925903, + 1.9899996519088745, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0199997425079346, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0299997329711914, + 2.0799996852874756, + 2.0899996757507324, + 2.189999580383301, + 2.2199995517730713, + 2.5899994373321533, + 2.729999542236328, + 2.749999523162842, + 2.8399994373321533, + ], + "type": "longrope", + }, + "rope_theta": 10000.0, + "sliding_window": 262144, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.43.3", + "use_cache": True, + "attention_bias": False, + "vocab_size": 32064, + }, + "phi-3_5-vision": { + "_name_or_path": "Phi-3.5-vision-instruct", + "architectures": ["Phi3VForCausalLM"], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_phi3_v.Phi3VConfig", + "AutoModelForCausalLM": "modeling_phi3_v.Phi3VForCausalLM", + }, + "bos_token_id": 1, + "embd_layer": { + "embedding_cls": "image", + "hd_transform_order": "sub_glb", + "projection_cls": "mlp", + "use_hd_transform": True, + "with_learnable_separator": True, + }, + "embd_pdrop": 0.0, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 3072, + "img_processor": { + "image_dim_out": 1024, + "model_name": "openai/clip-vit-large-patch14-336", + "name": "clip_vision_model", + "num_img_tokens": 144, + }, + "initializer_range": 0.02, + "intermediate_size": 8192, + "max_position_embeddings": 131072, + "model_type": "phi3_v", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 32, + "original_max_position_embeddings": 4096, + "pad_token_id": 32000, + "resid_pdrop": 0.0, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "long_factor": [ + 1.0800000429153442, + 1.1100000143051147, + 1.1399999856948853, + 1.340000033378601, + 1.5899999141693115, + 1.600000023841858, + 1.6200000047683716, + 2.620000123977661, + 3.2300000190734863, + 3.2300000190734863, + 4.789999961853027, + 7.400000095367432, + 7.700000286102295, + 9.09000015258789, + 12.199999809265137, + 17.670000076293945, + 24.46000099182129, + 28.57000160217285, + 30.420001983642578, + 30.840002059936523, + 32.590003967285156, + 32.93000411987305, + 42.320003509521484, + 44.96000289916992, + 50.340003967285156, + 50.45000457763672, + 57.55000305175781, + 57.93000411987305, + 58.21000289916992, + 60.1400032043457, + 62.61000442504883, + 62.62000274658203, + 62.71000289916992, + 63.1400032043457, + 63.1400032043457, + 63.77000427246094, + 63.93000411987305, + 63.96000289916992, + 63.970001220703125, + 64.02999877929688, + 64.06999969482422, + 64.08000183105469, + 64.12000274658203, + 64.41000366210938, + 64.4800033569336, + 64.51000213623047, + 64.52999877929688, + 64.83999633789062, + ], + "short_factor": [ + 1.08, + 1.1, + 1.1300000000000001, + 1.2800000000000002, + 1.3100000000000003, + 1.4500000000000004, + 1.4500000000000004, + 1.9500000000000008, + 2.030000000000001, + 2.4299999999999926, + 2.5699999999999896, + 2.9499999999999815, + 3.729999999999965, + 3.869999999999962, + 4.189999999999955, + 4.43999999999995, + 4.6399999999999455, + 4.979999999999938, + 5.159999999999934, + 5.279999999999932, + 5.759999999999922, + 5.889999999999919, + 5.889999999999919, + 5.969999999999917, + 6.089999999999915, + 6.2799999999999105, + 6.7699999999999, + 6.8899999999998975, + 7.109999999999893, + 7.129999999999892, + 7.179999999999891, + 7.289999999999889, + 7.339999999999888, + 7.559999999999883, + 7.619999999999882, + 7.69999999999988, + 7.879999999999876, + 7.879999999999876, + 7.879999999999876, + 7.939999999999875, + 7.949999999999875, + 7.979999999999874, + 8.19999999999987, + 8.439999999999864, + 8.469999999999864, + 8.589999999999861, + 8.809999999999857, + 8.999999999999853, + ], + "type": "su", + }, + "rope_theta": 10000.0, + "sliding_window": 262144, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.38.1", + "use_cache": True, + "vocab_size": 32064, + "_attn_implementation": "flash_attention_2", + }, + "phi-4": { + "_name_or_path": "Phi-4-mini-instruct", + "architectures": ["Phi3ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_phi3.Phi3Config", + "AutoModelForCausalLM": "modeling_phi3.Phi3ForCausalLM", + "AutoTokenizer": "Xenova/gpt-4o", + }, + "bos_token_id": 199999, + "embd_pdrop": 0.0, + "eos_token_id": 199999, + "full_attn_mod": 1, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "intermediate_size": 8192, + "interpolate_factor": 1, + "lm_head_bias": False, + "max_position_embeddings": 131072, + "mlp_bias": False, + "model_type": "phi3", + "num_attention_heads": 24, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "original_max_position_embeddings": 4096, + "pad_token_id": 199999, + "partial_rotary_factor": 0.75, + "resid_pdrop": 0.0, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "long_factor": [ + 1, + 1.118320672, + 1.250641126, + 1.398617824, + 1.564103225, + 1.74916897, + 1.956131817, + 2.187582649, + 2.446418898, + 2.735880826, + 3.059592084, + 3.421605075, + 3.826451687, + 4.279200023, + 4.785517845, + 5.351743533, + 5.984965424, + 6.693110555, + 7.485043894, + 8.370679318, + 9.36110372, + 10.4687158, + 11.70738129, + 13.09260651, + 14.64173252, + 16.37415215, + 18.31155283, + 20.47818807, + 22.90118105, + 25.61086418, + 28.64115884, + 32.03, + 32.1, + 32.13, + 32.23, + 32.6, + 32.61, + 32.64, + 32.66, + 32.7, + 32.71, + 32.93, + 32.97, + 33.28, + 33.49, + 33.5, + 44.16, + 47.77, + ], + "short_factor": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ], + "type": "longrope", + }, + "rope_theta": 10000.0, + "sliding_window": 262144, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.45.0", + "use_cache": True, + "vocab_size": 200064, + }, + "qwen": { + "architectures": ["QWenLMHeadModel"], + "auto_map": { + "AutoConfig": "configuration_qwen.QWenConfig", + "AutoModelForCausalLM": "modeling_qwen.QWenLMHeadModel", + }, + "attn_dropout_prob": 0.0, + "bf16": False, + "emb_dropout_prob": 0.0, + "hidden_size": 2048, + "intermediate_size": 11008, + "initializer_range": 0.02, + "kv_channels": 128, + "layer_norm_epsilon": 1e-06, + "max_position_embeddings": 8192, + "model_type": "qwen", + "no_bias": True, + "num_attention_heads": 16, + "num_hidden_layers": 24, + "rotary_emb_base": 10000, + "rotary_pct": 1.0, + "scale_attn_weights": True, + "seq_length": 8192, + "tie_word_embeddings": False, + "tokenizer_class": "QWenTokenizer", + "transformers_version": "4.32.0", + "use_cache": True, + "use_dynamic_ntk": True, + "use_flash_attn": "auto", + "use_logn_attn": True, + "vocab_size": 151936, + }, + "qwen2": { + "_name_or_path": "Qwen/Qwen1.5-1.8B-Chat", + "architectures": ["Qwen2ForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 5504, + "max_position_embeddings": 4096, + "max_window_layers": 21, + "model_type": "qwen2", + "num_attention_heads": 16, + "num_hidden_layers": 24, + "num_key_value_heads": 16, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.37.2", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936, + }, + "qwen2moe": { + "architectures": ["Qwen2MoeForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 5632, + "max_position_embeddings": 32768, + "max_window_layers": 21, + "model_type": "qwen2_moe", + "num_attention_heads": 16, + "num_hidden_layers": 24, + "num_key_value_heads": 16, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.39.0.dev0", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936, + "decoder_sparse_step": 1, + "moe_intermediate_size": 1408, + "shared_expert_intermediate_size": 5632, + "num_experts_per_tok": 4, + "num_experts": 60, + "norm_topk_prob": False, + "output_router_logits": False, + "router_aux_loss_coef": 0.001, + }, + "deepseek_v2_lite": { + "architectures": ["DeepseekV2ForCausalLM"], + "attention_bias": False, + "bos_token_id": 100000, + "eos_token_id": 100001, + "first_k_dense_replace": 1, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 10944, + "kv_lora_rank": 512, + "max_position_embeddings": 163840, + "model_type": "deepseek_v2", + "moe_intermediate_size": 1408, + "moe_layer_freq": 1, + "n_group": 1, + "n_routed_experts": 64, + "n_shared_experts": 2, + "norm_topk_prob": False, + "num_attention_heads": 16, + "num_experts_per_tok": 6, + "num_hidden_layers": 27, + "num_key_value_heads": 16, + "pretraining_tp": 1, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 0.707, + "mscale_all_dim": 0.707, + "original_max_position_embeddings": 4096, + "type": "yarn", + }, + "rope_theta": 10000, + "routed_scaling_factor": 1.0, + "scoring_func": "softmax", + "topk_group": 1, + "topk_method": "greedy", + "torch_dtype": "bfloat16", + "transformers_version": "4.33.1", + "use_cache": True, + "v_head_dim": 128, + "vocab_size": 102400, + }, + "stablelm": { + "architectures": ["StableLmForCausalLM"], + "bos_token_id": 0, + "eos_token_id": 0, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 6912, + "max_position_embeddings": 4096, + "model_type": "stablelm", + "layer_norm_eps": 1e-05, + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 32, + "partial_rotary_factor": 0.25, + "rope_theta": 10000, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.38.0", + "use_cache": True, + "use_qkv_bias": False, + "vocab_size": 50304, + }, + "baichuan": { + "architectures": ["BaichuanForCausalLM"], + "tokenizer_class": "BaichuanTokenizer", + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 4096, + "model_max_length": 4096, + "model_type": "baichuan", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "_from_model_config": True, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.29.2", + "use_cache": True, + "vocab_size": 125696, + }, + "internlm": { + "architectures": ["InternLMForCausalLM"], + "bias": True, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 2048, + "model_type": "internlm", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "pad_token_id": 2, + "rms_norm_eps": 1e-06, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.33.2", + "use_cache": True, + "vocab_size": 103168, + }, + "gemma_2b": { + "architectures": ["GemmaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 2, + "eos_token_id": 1, + "head_dim": 256, + "hidden_act": "gelu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 16384, + "max_position_embeddings": 8192, + "model_type": "gemma", + "num_attention_heads": 8, + "num_hidden_layers": 18, + "num_key_value_heads": 1, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "rope_theta": 10000.0, + "torch_dtype": "bfloat16", + "transformers_version": "4.38.0.dev0", + "use_cache": True, + "vocab_size": 256000, + }, + "gemma2_2b": { + "architectures": ["Gemma2ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "attn_logit_softcapping": 50.0, + "bos_token_id": 2, + "cache_implementation": "hybrid", + "eos_token_id": [1, 107], + "final_logit_softcapping": 30.0, + "head_dim": 256, + "hidden_act": "gelu_pytorch_tanh", + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 2304, + "initializer_range": 0.02, + "intermediate_size": 9216, + "max_position_embeddings": 8192, + "model_type": "gemma2", + "num_attention_heads": 8, + "num_hidden_layers": 26, + "num_key_value_heads": 4, + "pad_token_id": 0, + "query_pre_attn_scalar": 256, + "rms_norm_eps": 1e-06, + "rope_theta": 10000.0, + "sliding_window": 4096, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.4", + "use_cache": True, + "vocab_size": 256000, + }, + "gemma2_2b-jpn": { + "architectures": ["Gemma2ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "attn_logit_softcapping": 50.0, + "bos_token_id": 2, + "cache_implementation": "hybrid", + "dtype": "bfloat16", + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 2304, + "initializer_range": 0.02, + "intermediate_size": 9216, + "max_position_embeddings": 8192, + "model_type": "gemma2", + "num_attention_heads": 8, + "num_hidden_layers": 26, + "num_key_value_heads": 4, + "pad_token_id": 0, + "query_pre_attn_scalar": 224, + "rms_norm_eps": 1e-06, + "rope_theta": 10000.0, + "sliding_window": 4096, + "torch_dtype": "bfloat16", + "transformers_version": "4.44.2", + "use_cache": True, + "vocab_size": 256000, + }, + "gemma2_9b": { + "architectures": ["Gemma2ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "attn_logit_softcapping": 50.0, + "bos_token_id": 2, + "cache_implementation": "hybrid", + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "head_dim": 256, + "hidden_act": "gelu_pytorch_tanh", + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 3584, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 8192, + "model_type": "gemma2", + "num_attention_heads": 16, + "num_hidden_layers": 42, + "num_key_value_heads": 8, + "pad_token_id": 0, + "query_pre_attn_scalar": 256, + "rms_norm_eps": 1e-06, + "rope_theta": 10000.0, + "sliding_window": 4096, + "sliding_window_size": 4096, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.0.dev0", + "use_cache": True, + "vocab_size": 256000, + }, + "gemma3_1b_it": { + "architectures": ["Gemma3ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "attn_logit_softcapping": None, + "bos_token_id": 2, + "cache_implementation": "hybrid", + "eos_token_id": [1, 106], + "final_logit_softcapping": None, + "head_dim": 256, + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 1152, + "initializer_range": 0.02, + "intermediate_size": 6912, + "max_position_embeddings": 32768, + "model_type": "gemma3_text", + "num_attention_heads": 4, + "num_hidden_layers": 26, + "num_key_value_heads": 1, + "pad_token_id": 0, + "query_pre_attn_scalar": 256, + "rms_norm_eps": 1e-06, + "rope_local_base_freq": 10000, + "rope_scaling": None, + "rope_theta": 1000000, + "sliding_window": 512, + "sliding_window_pattern": 6, + "torch_dtype": "bfloat16", + "transformers_version": "4.50.0.dev0", + "use_cache": True, + "vocab_size": 262144, + }, + "gemma2_27b": { + "architectures": ["Gemma2ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "attn_logit_softcapping": 50.0, + "bos_token_id": 2, + "cache_implementation": "hybrid", + "eos_token_id": 1, + "final_logit_softcapping": 30.0, + "head_dim": 128, + "hidden_act": "gelu_pytorch_tanh", + "hidden_activation": "gelu_pytorch_tanh", + "hidden_size": 4608, + "initializer_range": 0.02, + "intermediate_size": 36864, + "max_position_embeddings": 8192, + "model_type": "gemma2", + "num_attention_heads": 32, + "num_hidden_layers": 46, + "num_key_value_heads": 16, + "pad_token_id": 0, + "query_pre_attn_scalar": 144, + "rms_norm_eps": 1e-06, + "rope_theta": 10000.0, + "sliding_window": 4096, + "sliding_window_size": 4096, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.0.dev0", + "use_cache": True, + "vocab_size": 256000, + "_attn_implementation": "eager", + }, + "rwkv5_3b": { + "architectures": ["RwkvForCausalLM"], + "auto_map": { + "AutoConfig": "configuration_rwkv5.Rwkv5Config", + "AutoModelForCausalLM": "modeling_rwkv5.RwkvForCausalLM", + }, + "attention_hidden_size": 2560, + "bos_token_id": 0, + "context_length": 4096, + "eos_token_id": 0, + "head_size": 64, + "hidden_size": 2560, + "intermediate_size": None, + "layer_norm_epsilon": 1e-05, + "model_type": "rwkv5", + "model_version": "5_2", + "num_hidden_layers": 32, + "rescale_every": 6, + "tie_word_embeddings": True, + "transformers_version": "4.34.0", + "use_cache": True, + "vocab_size": 65536, + }, + "orion": { + "architectures": ["OrionForCausalLM"], + "auto_map": { + "AutoConfig": "configuration_orion.OrionConfig", + "AutoModelForCausalLM": "modeling_orion.OrionForCausalLM", + }, + "tokenizer_class": "OrionTokenizer", + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 5120, + "model_type": "orion", + "initializer_range": 0.02, + "intermediate_size": 15360, + "max_position_embeddings": 4096, + "max_sequence_length": 4096, + "num_attention_heads": 40, + "num_hidden_layers": 40, + "num_key_value_heads": 40, + "pad_token_id": 0, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 10000.0, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.34.0", + "use_cache": True, + "vocab_size": 84608, + }, + "llava": { + "architectures": ["LlavaForConditionalGeneration"], + "ignore_index": -100, + "image_token_index": 32000, + "model_type": "llava", + "pad_token_id": 32001, + "projector_hidden_act": "gelu", + "text_config": { + "_name_or_path": "meta-llama/Llama-2-7b-hf", + "architectures": ["LlamaForCausalLM"], + "max_position_embeddings": 4096, + "model_type": "llama", + "rms_norm_eps": 1e-05, + "torch_dtype": "float16", + "vocab_size": 32064, + }, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.36.0.dev0", + "vision_config": { + "hidden_size": 1024, + "image_size": 336, + "intermediate_size": 4096, + "model_type": "clip_vision_model", + "num_attention_heads": 16, + "num_hidden_layers": 24, + "patch_size": 14, + "projection_dim": 768, + "vocab_size": 32000, + }, + "vision_feature_layer": -2, + "vision_feature_select_strategy": "default", + "vocab_size": 32064, + }, + "chatglm": { + "architectures": ["ChatGLMModel"], + "model_type": "chatglm", + "add_bias_linear": False, + "add_qkv_bias": True, + "apply_query_key_layer_scaling": True, + "apply_residual_connection_post_layernorm": False, + "attention_dropout": 0.0, + "attention_softmax_in_fp32": True, + "bias_dropout_fusion": True, + "ffn_hidden_size": 13696, + "fp32_residual_connection": False, + "hidden_dropout": 0.0, + "hidden_size": 4096, + "kv_channels": 128, + "layernorm_epsilon": 1e-05, + "multi_query_attention": True, + "multi_query_group_num": 2, + "num_attention_heads": 32, + "num_layers": 28, + "original_rope": True, + "padded_vocab_size": 65024, + "post_layer_norm": True, + "rmsnorm": True, + "seq_length": 8192, + "use_cache": True, + "torch_dtype": "float16", + "transformers_version": "4.30.2", + "tie_word_embeddings": False, + "eos_token_id": 2, + "pad_token_id": 0, + }, + "llama3_1_8b": { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 128000, + "eos_token_id": [128001, 128008, 128009], + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 131072, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + }, + "rope_theta": 500000.0, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.3", + "use_cache": True, + "vocab_size": 128256, + }, + "llama3_1_70b": { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 128000, + "eos_token_id": [128001, 128008, 128009], + "hidden_act": "silu", + "hidden_size": 8192, + "initializer_range": 0.02, + "intermediate_size": 28672, + "max_position_embeddings": 131072, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 64, + "num_hidden_layers": 80, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + }, + "rope_theta": 500000.0, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.3", + "use_cache": True, + "vocab_size": 128256, + }, + "llama3_2_1b": { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 128000, + "eos_token_id": [128001, 128008, 128009], + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 8192, + "max_position_embeddings": 131072, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 16, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "factor": 32.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + }, + "rope_theta": 500000.0, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.45.0.dev0", + "use_cache": True, + "vocab_size": 128256, + }, + "llama3_2_3b": { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 128000, + "eos_token_id": [128001, 128008, 128009], + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 3072, + "initializer_range": 0.02, + "intermediate_size": 8192, + "max_position_embeddings": 131072, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 24, + "num_hidden_layers": 28, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "factor": 32.0, + "high_freq_factor": 4.0, + "low_freq_factor": 1.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3", + }, + "rope_theta": 500000.0, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.45.0.dev0", + "use_cache": True, + "vocab_size": 128256, + }, + "snowflake-arctic-embed-m": { + "architectures": ["BertModel"], + "attention_probs_dropout_prob": 0.1, + "classifier_dropout": None, + "gradient_checkpointing": False, + "hidden_act": "gelu", + "hidden_dropout_prob": 0.1, + "hidden_size": 768, + "initializer_range": 0.02, + "intermediate_size": 3072, + "layer_norm_eps": 1e-12, + "max_position_embeddings": 512, + "model_type": "bert", + "num_attention_heads": 12, + "num_hidden_layers": 12, + "pad_token_id": 0, + "position_embedding_type": "absolute", + "torch_dtype": "float32", + "transformers_version": "4.36.1", + "type_vocab_size": 2, + "use_cache": True, + "vocab_size": 30522, + }, + # "snowflake-arctic-embed-s": { + # "architectures": ["BertModel"], + # "attention_probs_dropout_prob": 0.1, + # "classifier_dropout": None, + # "hidden_act": "gelu", + # "hidden_dropout_prob": 0.1, + # "hidden_size": 384, + # "initializer_range": 0.02, + # "intermediate_size": 1536, + # "layer_norm_eps": 1e-12, + # "max_position_embeddings": 512, + # "model_type": "bert", + # "num_attention_heads": 12, + # "num_hidden_layers": 12, + # "pad_token_id": 0, + # "position_embedding_type": "absolute", + # "torch_dtype": "float32", + # "transformers_version": "4.36.1", + # "type_vocab_size": 2, + # "use_cache": True, + # "vocab_size": 30522, + # }, + "stablelm-2-zephyr-1_6b": { + "architectures": ["StableLmForCausalLM"], + "bos_token_id": 100257, + "eos_token_id": 100257, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 5632, + "max_position_embeddings": 4096, + "model_type": "stablelm", + "layer_norm_eps": 1e-05, + "num_attention_heads": 32, + "num_hidden_layers": 24, + "num_key_value_heads": 32, + "partial_rotary_factor": 0.25, + "rope_theta": 10000, + "tie_word_embeddings": False, + "torch_dtype": "float16", + "transformers_version": "4.38.0", + "use_cache": True, + "use_qkv_bias": True, + "vocab_size": 100352, + }, + "qwen2_0_5b": { + "architectures": ["Qwen2ForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 896, + "initializer_range": 0.02, + "intermediate_size": 4864, + "max_position_embeddings": 32768, + "max_window_layers": 24, + "model_type": "qwen2", + "num_attention_heads": 14, + "num_hidden_layers": 24, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.40.1", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936, + }, + "qwen2_1_5b": { + "architectures": ["Qwen2ForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 1536, + "initializer_range": 0.02, + "intermediate_size": 8960, + "max_position_embeddings": 32768, + "max_window_layers": 28, + "model_type": "qwen2", + "num_attention_heads": 12, + "num_hidden_layers": 28, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.40.1", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936, + }, + "qwen2.5_3b": { + "architectures": ["Qwen2ForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 11008, + "max_position_embeddings": 32768, + "max_window_layers": 70, + "model_type": "qwen2", + "num_attention_heads": 16, + "num_hidden_layers": 36, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 32768, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.43.1", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936, + }, + "qwen2_7b": { + "architectures": ["Qwen2ForCausalLM"], + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "hidden_act": "silu", + "hidden_size": 3584, + "initializer_range": 0.02, + "intermediate_size": 18944, + "max_position_embeddings": 32768, + "max_window_layers": 28, + "model_type": "qwen2", + "num_attention_heads": 28, + "num_hidden_layers": 28, + "num_key_value_heads": 4, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "sliding_window": 131072, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.41.2", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 152064, + }, + "qwen3_0.6b": { + "architectures": ["Qwen3ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 1024, + "initializer_range": 0.02, + "intermediate_size": 3072, + "max_position_embeddings": 40960, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 16, + "num_hidden_layers": 28, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "rope_theta": 1000000, + "sliding_window": None, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936, + }, + "qwen3_1.7b": { + "architectures": ["Qwen3ForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 151643, + "eos_token_id": 151645, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 6144, + "max_position_embeddings": 40960, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 16, + "num_hidden_layers": 28, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "rope_theta": 1000000, + "sliding_window": None, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.51.0", + "use_cache": True, + "use_sliding_window": False, + "vocab_size": 151936, + }, + "internlm2": { + "architectures": ["InternLM2ForCausalLM"], + "attn_implementation": "eager", + "bias": False, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 32768, + "model_type": "internlm2", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "pad_token_id": 2, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 1000000, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.37.1", + "use_cache": True, + "vocab_size": 92544, + }, + "internlm2_5_7b": { + "architectures": ["InternLM2ForCausalLM"], + "attn_implementation": "eager", + "bias": False, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 32768, + "model_type": "internlm2", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "pad_token_id": 2, + "rms_norm_eps": 1e-05, + "rope_scaling": {"type": "dynamic", "factor": 2.0}, + "rope_theta": 1000000, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.41.0", + "use_cache": True, + "vocab_size": 92544, + "pretraining_tp": 1, + }, + "starcoder2": { + "activation_function": "gelu", + "architectures": ["Starcoder2ForCausalLM"], + "attention_dropout": 0.1, + "residual_dropout": 0.1, + "embedding_dropout": 0.1, + "attention_softmax_in_fp32": True, + "bos_token_id": 0, + "eos_token_id": 0, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 4608, + "initializer_range": 0.018042, + "intermediate_size": 18432, + "layer_norm_epsilon": 1e-05, + "max_position_embeddings": 16384, + "mlp_type": "default", + "model_type": "starcoder2", + "norm_epsilon": 1e-05, + "norm_type": "layer_norm", + "num_attention_heads": 36, + "num_hidden_layers": 32, + "num_key_value_heads": 4, + "rope_theta": 1000000, + "scale_attention_softmax_in_fp32": True, + "scale_attn_weights": True, + "sliding_window": 4096, + "torch_dtype": "bfloat16", + "transformers_version": "4.37.0.dev0", + "use_bias": True, + "use_cache": True, + "vocab_size": 49152, + }, + "smollm_1_7b": { + "_name_or_path": "HuggingFaceTB/cosmo2-1.7B-webinst-sc2", + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 8192, + "max_position_embeddings": 2048, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 24, + "num_key_value_heads": 32, + "pad_token_id": 2, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 10000.0, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.3", + "use_cache": True, + "vocab_size": 49152, + }, + "smollm_360m": { + "_name_or_path": "HuggingFaceTB/cosmo2-350M-webinst-sc2", + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 960, + "initializer_range": 0.02, + "intermediate_size": 2560, + "max_position_embeddings": 2048, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 15, + "num_hidden_layers": 32, + "num_key_value_heads": 5, + "pad_token_id": 2, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 10000.0, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.3", + "use_cache": True, + "vocab_size": 49152, + }, + "smollm_135m": { + "_name_or_path": "HuggingFaceTB/cosmo2-135M-webinst-sc2", + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 576, + "initializer_range": 0.02, + "intermediate_size": 1536, + "max_position_embeddings": 2048, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 9, + "num_hidden_layers": 30, + "num_key_value_heads": 3, + "pad_token_id": 2, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "rope_theta": 10000.0, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.3", + "use_cache": True, + "vocab_size": 49152, + }, + "smollm2_135m": { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 576, + "initializer_range": 0.041666666666666664, + "intermediate_size": 1536, + "is_llama_config": True, + "max_position_embeddings": 8192, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 9, + "num_hidden_layers": 30, + "num_key_value_heads": 3, + "pad_token_id": 2, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_interleaved": False, + "rope_scaling": None, + "rope_theta": 100000, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.3", + "transformers.js_config": { + "kv_cache_dtype": { + "q4f16": "float16", + "fp16": "float16", + } + }, + "use_cache": True, + "vocab_size": 49152, + }, + # "smollm2_1_7b": { + # "architectures": ["LlamaForCausalLM"], + # "attention_bias": False, + # "attention_dropout": 0.0, + # "bos_token_id": 1, + # "eos_token_id": 2, + # "hidden_act": "silu", + # "hidden_size": 2048, + # "initializer_range": 0.02, + # "intermediate_size": 8192, + # "max_position_embeddings": 8192, + # "mlp_bias": False, + # "model_type": "llama", + # "num_attention_heads": 32, + # "num_hidden_layers": 24, + # "num_key_value_heads": 32, + # "pad_token_id": 2, + # "pretraining_tp": 1, + # "rms_norm_eps": 1e-05, + # "rope_scaling": None, + # "rope_theta": 130000, + # "tie_word_embeddings": True, + # "torch_dtype": "bfloat16", + # "transformers_version": "4.42.3", + # "transformers.js_config": { + # "dtype": "q4", + # "kv_cache_dtype": { + # "q4f16": "float16", + # "fp16": "float16", + # }, + # "use_external_data_format": { + # "model.onnx": True, + # "model_fp16.onnx": True, + # }, + # }, + # "use_cache": True, + # "vocab_size": 49152, + # }, + "smollm2_360m": { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 960, + "initializer_range": 0.02, + "intermediate_size": 2560, + "is_llama_config": True, + "max_position_embeddings": 8192, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 15, + "num_hidden_layers": 32, + "num_key_value_heads": 5, + "pad_token_id": 2, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_interleaved": False, + "rope_scaling": None, + "rope_theta": 100000, + "tie_word_embeddings": True, + "torch_dtype": "bfloat16", + "transformers_version": "4.42.3", + "transformers.js_config": { + "kv_cache_dtype": { + "q4f16": "float16", + "fp16": "float16", + } + }, + "use_cache": True, + "vocab_size": 49152, + }, + "aya-23": { + "architectures": ["CohereForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 5, + "eos_token_id": 255001, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "layer_norm_eps": 1e-05, + "logit_scale": 0.0625, + "max_position_embeddings": 8192, + "model_type": "cohere", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "pad_token_id": 0, + "rope_theta": 10000, + "torch_dtype": "float16", + "transformers_version": "4.40.0.dev0", + "use_cache": True, + "use_qk_norm": False, + "vocab_size": 256000, + }, + "minicpm_2b": { + "architectures": ["MiniCPMForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 2304, + "initializer_range": 0.1, + "intermediate_size": 5760, + "max_position_embeddings": 65536, + "max_length": 131072, + "model_type": "minicpm", + "num_attention_heads": 36, + "num_hidden_layers": 40, + "num_key_value_heads": 36, + "rms_norm_eps": 1e-05, + "rope_scaling": {"type": "dynamic", "factor": 4.0}, + "torch_dtype": "bfloat16", + "transformers_version": "4.36.0", + "use_cache": True, + "vocab_size": 122760, + "scale_emb": 12, + "dim_model_base": 256, + "scale_depth": 1.4, + "tie_word_embeddings": False, + "rope_theta": 1000000.0, + }, + "minicpm_2b_sft_bf16": { + "architectures": ["MiniCPMForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 2304, + "initializer_range": 0.1, + "intermediate_size": 5760, + "max_position_embeddings": 4096, + "model_type": "minicpm", + "num_attention_heads": 36, + "num_hidden_layers": 40, + "num_key_value_heads": 36, + "rms_norm_eps": 1e-05, + "torch_dtype": "bfloat16", + "tie_word_embeddings": True, + "transformers_version": "4.36.0", + "use_cache": True, + "vocab_size": 122753, + "scale_emb": 12, + "dim_model_base": 256, + "scale_depth": 1.4, + }, + "minicpm-moe-8x2b": { + "architectures": ["MiniCPMForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 2304, + "initializer_range": 0.1, + "intermediate_size": 5760, + "max_position_embeddings": 4096, + "model_type": "minicpm", + "num_attention_heads": 36, + "num_hidden_layers": 40, + "num_key_value_heads": 36, + "rms_norm_eps": 1e-05, + "rope_scaling": None, + "torch_dtype": "bfloat16", + "tie_word_embeddings": True, + "transformers_version": "4.36.0", + "use_cache": True, + "vocab_size": 122753, + "scale_emb": 12, + "dim_model_base": 256, + "scale_depth": 1.4, + "num_experts": 8, + "num_experts_per_tok": 2, + }, + "deepseek": { + "architectures": ["DeepseekForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 100000, + "eos_token_id": 100001, + "first_k_dense_replace": 1, + "hidden_act": "silu", + "hidden_size": 2048, + "initializer_range": 0.02, + "intermediate_size": 10944, + "max_position_embeddings": 4096, + "model_type": "deepseek", + "moe_intermediate_size": 1408, + "moe_layer_freq": 1, + "n_routed_experts": 64, + "n_shared_experts": 2, + "norm_topk_prob": False, + "num_attention_heads": 16, + "num_experts_per_tok": 6, + "num_hidden_layers": 28, + "num_key_value_heads": 16, + "pretraining_tp": 1, + "rms_norm_eps": 1e-06, + "rope_scaling": None, + "rope_theta": 10000, + "scoring_func": "softmax", + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.36.2", + "use_cache": True, + "vocab_size": 102400, + }, + "gpt_j": { + "activation_function": "gelu_new", + "architectures": ["GPTJForCausalLM"], + "attn_pdrop": 0.0, + "bos_token_id": 50256, + "embd_pdrop": 0.0, + "eos_token_id": 50256, + "initializer_range": 0.02, + "layer_norm_epsilon": 1e-05, + "model_type": "gptj", + "n_embd": 4096, + "n_head": 16, + "n_inner": None, + "n_layer": 28, + "n_positions": 2048, + "resid_pdrop": 0.0, + "rotary": True, + "rotary_dim": 64, + "scale_attn_weights": True, + "summary_activation": None, + "summary_first_dropout": 0.1, + "summary_proj_to_labels": True, + "summary_type": "cls_index", + "summary_use_proj": True, + "rope_scaling": {"rope_type": "gptj"}, + "tie_word_embeddings": False, + "tokenizer_class": "GPT2Tokenizer", + "transformers_version": "4.18.0.dev0", + "use_cache": True, + "vocab_size": 50400, + }, + "olmo2_7b": { + "architectures": ["Olmo2ForCausalLM"], + "hidden_act": "silu", + "hidden_size": 4096, + "intermediate_size": 11008, + "max_position_embeddings": 4096, + "model_type": "olmo2", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 32, + "rms_norm_eps": 1e-06, + "rope_theta": 500000.0, + "tie_word_embeddings": False, + "torch_dtype": "float32", + "transformers_version": "4.47.0", + "use_cache": True, + "vocab_size": 100352, + }, + "olmo2_13b": { + "architectures": ["Olmo2ForCausalLM"], + "hidden_act": "silu", + "hidden_size": 5120, + "intermediate_size": 13824, + "max_position_embeddings": 4096, + "model_type": "olmo2", + "num_attention_heads": 40, + "num_hidden_layers": 40, + "num_key_value_heads": 40, + "rms_norm_eps": 1e-06, + "rope_theta": 500000.0, + "tie_word_embeddings": False, + "torch_dtype": "float32", + "transformers_version": "4.47.0", + "use_cache": True, + "vocab_size": 100352, + }, + "olmo2_32b": { + "architectures": ["Olmo2ForCausalLM"], + "hidden_act": "silu", + "hidden_size": 5120, + "intermediate_size": 27648, + "max_position_embeddings": 4096, + "model_type": "olmo2", + "num_attention_heads": 40, + "num_hidden_layers": 64, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-06, + "rope_theta": 500000.0, + "tie_word_embeddings": False, + "torch_dtype": "float32", + "transformers_version": "4.49.0", + "use_cache": True, + "vocab_size": 100352, + }, +} diff --git a/python/mlc_llm/model/model_utils.py b/python/mlc_llm/model/model_utils.py new file mode 100644 index 0000000..5312927 --- /dev/null +++ b/python/mlc_llm/model/model_utils.py @@ -0,0 +1,14 @@ +"""Utilities shared across model definitions.""" + +from tvm import te +from tvm.relax.frontend.nn import Tensor, op + + +def index_last_token(x: Tensor) -> Tensor: + """Select the last token while preserving the historical `index` TE op shape/name.""" + + def _index(x: te.Tensor): + b, s, d = x.shape + return te.compute((b, 1, d), lambda i, _, k: x[i, s - 1, k], name="index") + + return op.tensor_expr_op(_index, name_hint="index", args=[x]) diff --git a/python/mlc_llm/model/nemotron/__init__.py b/python/mlc_llm/model/nemotron/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/nemotron/nemotron_loader.py b/python/mlc_llm/model/nemotron/nemotron_loader.py new file mode 100644 index 0000000..25e912d --- /dev/null +++ b/python/mlc_llm/model/nemotron/nemotron_loader.py @@ -0,0 +1,14 @@ +""" +This file specifies how MLC's Nemotron parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .nemotron_model import NemotronForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=NemotronForCausalLM, + add_unused=["rotary_emb.inv_freq"], + include_gate_up=False, +) diff --git a/python/mlc_llm/model/nemotron/nemotron_model.py b/python/mlc_llm/model/nemotron/nemotron_model.py new file mode 100644 index 0000000..2df091b --- /dev/null +++ b/python/mlc_llm/model/nemotron/nemotron_model.py @@ -0,0 +1,522 @@ +""" +Implementation for Nemotron architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class NemotronConfig(ConfigBase): + """Configuration of the Nemotron model.""" + + vocab_size: int + max_position_embeddings: int + hidden_size: int + intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + rope_theta: int = 10000 + partial_rotary_factor: float = 0.5 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + norm_eps: float = 1e-5 + head_dim: int = 0 + tie_word_embeddings: bool = False + mlp_bias: bool = False + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + pipeline_parallel_stages: int = 1 + max_batch_size: int = 1 + disaggregation: bool = False + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + self.context_window_size = self.max_position_embeddings + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + self.rotary_dim = int(self.partial_rotary_factor * self.head_dim) + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class NemotronMLP(nn.Module): + """Nemotron MLP module.""" + + def __init__(self, config: NemotronConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear( + config.intermediate_size, config.hidden_size, bias=config.mlp_bias + ) + + def forward(self, x: Tensor) -> Tensor: + """Forward pass of the MLP module.""" + out = self.up_proj(x) + out = op.square(op.relu(out)) + out = self.down_proj(out) + return out + + +class NemotronEmbedding(nn.Embedding): + """The embedding module that can be shared with the final head. From Qwen2Embedding.""" + + def lm_head_forward(self, x: Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class NemotronLayerNorm1P(nn.LayerNorm): + """Nemotron LayerNorm1P module.""" + + def __init__(self, normalized_shape: int, eps: float = 1e-5, elementwise_affine: bool = True): + super().__init__(normalized_shape, eps, elementwise_affine) + + def forward(self, x: Tensor) -> Tensor: + """Forward pass of the tweaked LayerNorm module.""" + return op.layer_norm( + x, + normalized_shape=self.normalized_shape, + weight=self.weight + 1, + bias=self.bias, + eps=self.eps, + ) + + +class NemotronAttention(nn.Module): + def __init__(self, config: NemotronConfig): + self.head_dim = config.head_dim + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({config.num_key_value_heads}) must be divisible by tensor_parallel_shards" # noqa: E501 + ) + assert config.num_key_value_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than {config.num_key_value_heads}" + ) + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=False, + ) + self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class NemotronDecoderLayer(nn.Module): + def __init__(self, config: NemotronConfig): + self.self_attn = NemotronAttention(config) + self.mlp = NemotronMLP(config) + self.input_layernorm = NemotronLayerNorm1P(config.hidden_size, config.norm_eps) + self.post_attention_layernorm = NemotronLayerNorm1P(config.hidden_size, config.norm_eps) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set(self.mlp.up_proj, tp.ShardSingleDim("_shard_mlp_up", dim=1)) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class NemotronModel(nn.Module): + def __init__(self, config: NemotronConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = NemotronEmbedding("vocab_size", config.hidden_size) + self.layers = nn.ModuleList( + [NemotronDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = NemotronLayerNorm1P(config.hidden_size, config.norm_eps) + self.num_layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + + # Compute pipeline layer partition. + layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + self.layer_partition = [ + i * layers_per_stage for i in range(config.pipeline_parallel_stages) + ] + [config.num_hidden_layers] + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.layers): + if layer_id != 0 and layer_id in self.layer_partition: + hidden_states = op_ext.pipeline_stage_boundary(hidden_states) + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class NemotronForCausalLM(nn.Module): + def __init__(self, config: NemotronConfig): + self.model = NemotronModel(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_scaling = config.rope_scaling + self.rope_theta = config.rope_theta + self.rotary_dim = config.rotary_dim + self.tensor_parallel_shards = config.tensor_parallel_shards + self.disaggregation = config.disaggregation + self.dtype = "float32" + + def _set_pp(): + # hidden layers + for layer_id in range(config.num_hidden_layers): + stage = layer_id // (config.num_hidden_layers // config.pipeline_parallel_stages) + for _, param in self.model.layers[layer_id].named_parameters(): + param.attrs["pipeline_stages"] = [stage] + # last stage + last_stage = config.pipeline_parallel_stages - 1 + self.model.norm.weight.attrs["pipeline_stages"] = [last_stage] + # embedding table and lm_head is required by all stages + all_stages = list(range(config.pipeline_parallel_stages)) + self.model.embed_tokens.weight.attrs["pipeline_stages"] = all_stages + if not config.tie_word_embeddings: + self.lm_head.weight.attrs["pipeline_stages"] = all_stages + + _set_pp() + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return self.get_logits(hidden_states) + + def batch_forward_to_last_hidden_states( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + return hidden_states + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def get_logits(self, hidden_states: Tensor): + op_ext.configure() + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_select_last_hidden_states(self, hidden_states: Tensor, logit_positions: Tensor): + op_ext.configure() + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=0) + return hidden_states + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def prefill_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def decode_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_prefill_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_decode_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_verify_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rope_scaling=self.rope_scaling, + rotary_dim=self.rotary_dim, + layer_partition=self.model.layer_partition, + enable_disaggregation=self.disaggregation, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "get_logits": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_select_last_hidden_states": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/olmo/__init__.py b/python/mlc_llm/model/olmo/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/olmo/olmo_loader.py b/python/mlc_llm/model/olmo/olmo_loader.py new file mode 100644 index 0000000..310894e --- /dev/null +++ b/python/mlc_llm/model/olmo/olmo_loader.py @@ -0,0 +1,104 @@ +""" +This file specifies how MLC's OLMo parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.loader.standard_loader import make_standard_hf_loader +from mlc_llm.quantization import Quantization, make_awq_quant + +from .olmo_model import OLMoConfig, OLMoForCausalLM + +awq_quant = make_awq_quant(OLMoForCausalLM) + + +huggingface = make_standard_hf_loader( + model_cls=OLMoForCausalLM, + add_unused=["rotary_emb.inv_freq"], +) + + +def awq(model_config: OLMoConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of AWQ parameters. + Parameters + ---------- + model_config : OLMoConfig + The configuration of the OLMo model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to AWQ. + """ + model, _ = awq_quant(model_config, quantization) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"model.layers.{i}.self_attn" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{attn}.qkv_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{quantize_suffix}", + f"{attn}.k_proj.{quantize_suffix}", + f"{attn}.v_proj.{quantize_suffix}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate( + [q, k, v], + axis=1, # AWQ GEMM would transpose the weight + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # Concat gate and up in MLP + mlp = f"model.layers.{i}.mlp" + for quantize_suffix in ["qweight", "qzeros", "scales"]: + mlc_name = f"{mlp}.gate_up_proj.{quantize_suffix}" + assert mlc_name in named_parameters + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.gate_proj.{quantize_suffix}", + f"{mlp}.up_proj.{quantize_suffix}", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate( + [gate, up], + axis=1, # AWQ GEMM would transpose the weight + ).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + # inv_freq is not used in the model + mapping.add_unused(f"{attn}.rotary_emb.inv_freq") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + return mapping diff --git a/python/mlc_llm/model/olmo/olmo_model.py b/python/mlc_llm/model/olmo/olmo_model.py new file mode 100644 index 0000000..feaa915 --- /dev/null +++ b/python/mlc_llm/model/olmo/olmo_model.py @@ -0,0 +1,573 @@ +""" +Implementation for OLMo architecture. +TODO: add docstring +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class OLMoConfig(ConfigBase): + """Configuration of the OLMo model.""" + + vocab_size: int = None + hidden_size: int = None + num_attention_heads: int = None + num_key_value_heads: int = 0 + head_dim: int = 0 + position_embedding_base: int = 0 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + intermediate_size: int = None + hidden_act: str = None + num_hidden_layers: int = None + tie_word_embeddings: bool = False + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + pipeline_parallel_stages: int = 1 + max_batch_size: int = 1 + clip_qkv: float = None + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.num_key_value_heads == 0: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.num_attention_heads % self.num_key_value_heads == 0 + + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + if ( + self.pipeline_parallel_stages <= 0 + or self.pipeline_parallel_stages > self.num_hidden_layers + ): + raise ValueError( + f'Invalid "pipeline_parallel_stages" value({self.pipeline_parallel_stages}). ' + ) + + if self.clip_qkv is not None: + if self.clip_qkv <= 0: + raise ValueError(f"'clip_qkv'({self.clip_qkv}) should be non-negative") + + +class OLMoEmbedding(nn.Embedding): + """The embedding module that can be shared with the final lm_head. From Qwen2Embedding.""" + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class OLMoAttention(nn.Module): + def __init__(self, config: OLMoConfig): + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_key_value_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than {config.num_key_value_heads}" + ) + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({config.num_key_value_heads}) must be divisible by tensor_parallel_shards" # noqa: E501 + ) + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=False, + ) + self.clip_qkv = config.clip_qkv + self.o_proj = nn.Linear( + in_features=self.num_q_heads * self.head_dim, + out_features=config.hidden_size, + bias=False, + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + + # QKV Projection + qkv = self.qkv_proj(hidden_states) + + # Clamp after qkv projection if needed + if self.clip_qkv is not None: + qkv = qkv.maximum(-self.clip_qkv).minimum(self.clip_qkv) + + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +# Copied from qwen2_model.ACT2FN +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class OLMoFFN(nn.Module): + def __init__(self, config: OLMoConfig): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.act_fn = ACT2FN[config.hidden_act] + self.down_proj = nn.Linear( + in_features=self.intermediate_size, + out_features=config.hidden_size, + bias=False, + ) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class OLMoDecoderLayer(nn.Module): + def __init__(self, config: OLMoConfig): + self.input_layernorm = nn.LayerNorm( + normalized_shape=config.hidden_size, + eps=1e-5, + elementwise_affine=False, + ) + self.self_attn = OLMoAttention(config) + self.post_attention_layernorm = nn.LayerNorm( + normalized_shape=config.hidden_size, + eps=1e-5, + elementwise_affine=False, + ) + self.mlp = OLMoFFN(config) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + +class OLMoModel(nn.Module): + def __init__(self, config: OLMoConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = OLMoEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [OLMoDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.LayerNorm( + normalized_shape=config.hidden_size, + eps=1e-5, + elementwise_affine=False, + ) + + self.num_layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + # Compute pipeline layer partition. + layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + self.layer_partition = [ + i * layers_per_stage for i in range(config.pipeline_parallel_stages) + ] + [config.num_hidden_layers] + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + if layer_id != 0 and layer_id in self.layer_partition: + hidden_states = op_ext.pipeline_stage_boundary(hidden_states) + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class OLMoForCausalLM(nn.Module): + def __init__(self, config: OLMoConfig): + self.model = OLMoModel(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.rope_theta = config.position_embedding_base + self.rope_scaling = config.rope_scaling + self.intermediate_size = config.intermediate_size + self.num_hidden_layers = config.num_hidden_layers + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def _set_pp(): + # hidden layers + for layer_id in range(config.num_hidden_layers): + stage = layer_id // (config.num_hidden_layers // config.pipeline_parallel_stages) + for _, param in self.model.layers[layer_id].named_parameters(): + param.attrs["pipeline_stages"] = [stage] + + # embedding table and lm_head is required by all stages + all_stages = list(range(config.pipeline_parallel_stages)) + self.model.embed_tokens.weight.attrs["pipeline_stages"] = all_stages + if not config.tie_word_embeddings: + self.lm_head.weight.attrs["pipeline_stages"] = all_stages + + _set_pp() + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return self.get_logits(hidden_states) + + def batch_forward_to_last_hidden_states( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + ): + op_ext.configure() + hidden_states = self.model(input_embeds, paged_kv_cache) + return hidden_states + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def get_logits(self, hidden_states: Tensor): + op_ext.configure() + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_select_last_hidden_states(self, hidden_states: Tensor, logit_positions: Tensor): + op_ext.configure() + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=0) + return hidden_states + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def prefill_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def decode_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_prefill_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_decode_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_verify_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rope_scaling=self.rope_scaling, + layer_partition=self.model.layer_partition, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "get_logits": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_select_last_hidden_states": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/olmo2/__init__.py b/python/mlc_llm/model/olmo2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/olmo2/olmo2_loader.py b/python/mlc_llm/model/olmo2/olmo2_loader.py new file mode 100644 index 0000000..6a13420 --- /dev/null +++ b/python/mlc_llm/model/olmo2/olmo2_loader.py @@ -0,0 +1,13 @@ +""" +This file specifies how MLC's OLMo2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .olmo2_model import OLMo2ForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=OLMo2ForCausalLM, + add_unused=["rotary_emb.inv_freq"], +) diff --git a/python/mlc_llm/model/olmo2/olmo2_model.py b/python/mlc_llm/model/olmo2/olmo2_model.py new file mode 100644 index 0000000..deb508b --- /dev/null +++ b/python/mlc_llm/model/olmo2/olmo2_model.py @@ -0,0 +1,573 @@ +""" +Implementation for OLMo2 architecture. +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class OLMo2Config(ConfigBase): + """Configuration of the OLMo2 model.""" + + vocab_size: int = None + hidden_size: int = None + num_attention_heads: int = None + num_key_value_heads: int = 0 + head_dim: int = 0 + position_embedding_base: int = 0 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + intermediate_size: int = None + hidden_act: str = None + num_hidden_layers: int = None + rms_norm_eps: float = 1e-5 + tie_word_embeddings: bool = False + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + pipeline_parallel_stages: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.num_key_value_heads == 0: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.num_attention_heads % self.num_key_value_heads == 0 + + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + if ( + self.pipeline_parallel_stages <= 0 + or self.pipeline_parallel_stages > self.num_hidden_layers + ): + raise ValueError( + f'Invalid "pipeline_parallel_stages" value({self.pipeline_parallel_stages}). ' + ) + + +class OLMo2Embedding(nn.Embedding): + """The embedding module that can be shared with the final lm_head.""" + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class OLMo2Attention(nn.Module): + def __init__(self, config: OLMo2Config): + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_key_value_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than {config.num_key_value_heads}" + ) + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({config.num_key_value_heads}) must be divisible by tensor_parallel_shards" # noqa: E501 + ) + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=False, + ) + self.o_proj = nn.Linear( + in_features=self.num_q_heads * self.head_dim, + out_features=config.hidden_size, + bias=False, + ) + # QK-Norm: separate RMSNorm for Q and K after projection + self.q_norm = nn.RMSNorm( + self.num_q_heads * self.head_dim, -1, config.rms_norm_eps, bias=False + ) + self.k_norm = nn.RMSNorm( + self.num_kv_heads * self.head_dim, -1, config.rms_norm_eps, bias=False + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + + q, k, v = op.split(qkv, [h_q, h_q + h_kv], axis=2) + + # Apply QK-Norm before RoPE (reshape to 3D for RMSNorm, then back) + q = op.reshape(q, (b, s, h_q * d)) + k = op.reshape(k, (b, s, h_kv * d)) + q = self.q_norm(q) + k = self.k_norm(k) + q = op.reshape(q, (b, s, h_q, d)) + k = op.reshape(k, (b, s, h_kv, d)) + + qkv = op.concat([q, k, v], dim=2) + + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class OLMo2FFN(nn.Module): + def __init__(self, config: OLMo2Config): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.act_fn = ACT2FN[config.hidden_act] + self.down_proj = nn.Linear( + in_features=self.intermediate_size, + out_features=config.hidden_size, + bias=False, + ) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class OLMo2DecoderLayer(nn.Module): + def __init__(self, config: OLMo2Config): + rms_norm_eps = config.rms_norm_eps + self.self_attn = OLMo2Attention(config) + self.mlp = OLMo2FFN(config) + # OLMo2 uses post-norm (norm after attention/MLP, before residual add) + self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.post_feedforward_layernorm = nn.RMSNorm( + config.hidden_size, -1, rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(hidden_states, paged_kv_cache, layer_id) + out = self.post_attention_layernorm(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(hidden_states) + out = self.post_feedforward_layernorm(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + +class OLMo2Model(nn.Module): + def __init__(self, config: OLMo2Config): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = OLMo2Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [OLMo2DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + self.num_layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + # Compute pipeline layer partition. + layers_per_stage = ( + config.num_hidden_layers + config.pipeline_parallel_stages - 1 + ) // config.pipeline_parallel_stages + self.layer_partition = [ + i * layers_per_stage for i in range(config.pipeline_parallel_stages) + ] + [config.num_hidden_layers] + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + if layer_id != 0 and layer_id in self.layer_partition: + hidden_states = op_ext.pipeline_stage_boundary(hidden_states) + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class OLMo2ForCausalLM(nn.Module): + def __init__(self, config: OLMo2Config): + self.model = OLMo2Model(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.rope_theta = config.position_embedding_base + self.rope_scaling = config.rope_scaling + self.intermediate_size = config.intermediate_size + self.num_hidden_layers = config.num_hidden_layers + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def _set_pp(): + # hidden layers + for layer_id in range(config.num_hidden_layers): + stage = layer_id // (config.num_hidden_layers // config.pipeline_parallel_stages) + for _, param in self.model.layers[layer_id].named_parameters(): + param.attrs["pipeline_stages"] = [stage] + + # embedding table and lm_head is required by all stages + all_stages = list(range(config.pipeline_parallel_stages)) + self.model.embed_tokens.weight.attrs["pipeline_stages"] = all_stages + if not config.tie_word_embeddings: + self.lm_head.weight.attrs["pipeline_stages"] = all_stages + + _set_pp() + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return self.get_logits(hidden_states) + + def batch_forward_to_last_hidden_states( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + ): + op_ext.configure() + hidden_states = self.model(input_embeds, paged_kv_cache) + return hidden_states + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def get_logits(self, hidden_states: Tensor): + op_ext.configure() + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_select_last_hidden_states(self, hidden_states: Tensor, logit_positions: Tensor): + op_ext.configure() + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + hidden_states = op.take(hidden_states, logit_positions, axis=0) + return hidden_states + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def prefill_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def decode_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_prefill_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_decode_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_verify_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + hidden_states = self.batch_forward_to_last_hidden_states(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rope_scaling=self.rope_scaling, + layer_partition=self.model.layer_partition, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "get_logits": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_select_last_hidden_states": { + "hidden_states": nn.spec.Tensor(["seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/orion/__init__.py b/python/mlc_llm/model/orion/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/orion/orion_loader.py b/python/mlc_llm/model/orion/orion_loader.py new file mode 100644 index 0000000..e5fa68e --- /dev/null +++ b/python/mlc_llm/model/orion/orion_loader.py @@ -0,0 +1,13 @@ +""" +This file specifies how MLC's Orion parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .orion_model import OrionForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=OrionForCausalLM, + add_unused=["rotary_emb.inv_freq"], +) diff --git a/python/mlc_llm/model/orion/orion_model.py b/python/mlc_llm/model/orion/orion_model.py new file mode 100644 index 0000000..2f04289 --- /dev/null +++ b/python/mlc_llm/model/orion/orion_model.py @@ -0,0 +1,372 @@ +""" +Implementation for Orion-14B architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class OrionConfig(ConfigBase): + """Configuration of the Orion model.""" + + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + rms_norm_eps: float + vocab_size: int + position_embedding_base: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + num_key_value_heads: int = 0 + head_dim: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.num_key_value_heads == 0: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + assert self.num_attention_heads % self.num_key_value_heads == 0 + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class OrionFFN(nn.Module): + def __init__(self, config: OrionConfig): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class OrionAttention(nn.Module): + def __init__(self, config: OrionConfig): + self.head_dim = config.head_dim + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_kv_heads({config.num_key_value_heads}) must be divisible by tensor_parallel_shards" # noqa: E501 + ) + assert config.num_key_value_heads >= config.tensor_parallel_shards, ( + f"Too large tensor_parallel_shards, must be smaller than {config.num_key_value_heads}" + ) + self.num_kv_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_kv_heads) * self.head_dim, + bias=False, + ) + self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_kv_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.o_proj(output) + + +class OrionDecoderLayer(nn.Module): + def __init__(self, config: OrionConfig): + rms_norm_eps = config.rms_norm_eps + self.self_attn = OrionAttention(config) + self.mlp = OrionFFN(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size, rms_norm_eps) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, rms_norm_eps) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_q_heads * hd + k = self.self_attn.num_kv_heads * hd + v = self.self_attn.num_kv_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.self_attn.o_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class OrionModel(nn.Module): + def __init__(self, config: OrionConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = nn.Embedding("vocab_size", config.hidden_size) + self.layers = nn.ModuleList( + [OrionDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.LayerNorm(config.hidden_size, config.rms_norm_eps) + self.tensor_parallel_shards = config.tensor_parallel_shards + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class OrionForCausalLM(nn.Module): + def __init__(self, config: OrionConfig): + self.model = OrionModel(config) + self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/phi/__init__.py b/python/mlc_llm/model/phi/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/phi/phi_loader.py b/python/mlc_llm/model/phi/phi_loader.py new file mode 100644 index 0000000..e7e97e2 --- /dev/null +++ b/python/mlc_llm/model/phi/phi_loader.py @@ -0,0 +1,168 @@ +""" +This file specifies how MLC's Phi parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .phi_model import Phi1Config, PhiConfig, PhiForCausalLM + + +def huggingface(model_config: PhiConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : PhiConfig + The configuration of the Phi model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = PhiForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params = model.export_tvm(spec=model.get_default_spec()) + named_parameters = dict(_named_params) + mapping = ExternMapping() + + def _add(mlc_name, hf_name): + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=named_parameters[mlc_name].dtype, + ), + ) + + if model_config.model_type == "mixformer-sequential": + _add("transformer.embd.weight", "layers.0.wte.weight") + prefix = "transformer.h" + for i in range(model_config.n_layer): + _add(f"{prefix}.{i}.ln.weight", f"layers.{i + 1}.ln.weight") + _add(f"{prefix}.{i}.ln.bias", f"layers.{i + 1}.ln.bias") + _add(f"{prefix}.{i}.mixer.Wqkv.weight", f"layers.{i + 1}.mixer.Wqkv.weight") + _add(f"{prefix}.{i}.mixer.Wqkv.bias", f"layers.{i + 1}.mixer.Wqkv.bias") + _add( + f"{prefix}.{i}.mixer.out_proj.weight", + f"layers.{i + 1}.mixer.out_proj.weight", + ) + _add( + f"{prefix}.{i}.mixer.out_proj.bias", + f"layers.{i + 1}.mixer.out_proj.bias", + ) + _add(f"{prefix}.{i}.mlp.fc1.weight", f"layers.{i + 1}.mlp.fc1.weight") + _add(f"{prefix}.{i}.mlp.fc1.bias", f"layers.{i + 1}.mlp.fc1.bias") + _add(f"{prefix}.{i}.mlp.fc2.weight", f"layers.{i + 1}.mlp.fc2.weight") + _add(f"{prefix}.{i}.mlp.fc2.bias", f"layers.{i + 1}.mlp.fc2.bias") + mapping.add_unused(f"layers.{i + 1}.mixer.rotary_emb.inv_freq") + prefix = f"layers.{model_config.n_layer + 1}" + _add("lm_head.ln.weight", f"{prefix}.ln.weight") + _add("lm_head.ln.bias", f"{prefix}.ln.bias") + _add("lm_head.linear.weight", f"{prefix}.linear.weight") + _add("lm_head.linear.bias", f"{prefix}.linear.bias") + + elif model_config.model_type == "phi-msft": + _add("transformer.embd.weight", "transformer.embd.wte.weight") + for mlc_name, _ in named_parameters.items(): + if mlc_name not in mapping.param_map: + _add(mlc_name, mlc_name) + return mapping + + +def phi1_huggingface(model_config: Phi1Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of Phi-1/Phi-1.5 HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : PhiConfig + The configuration of the Phi model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = PhiForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params = model.export_tvm(spec=model.get_default_spec()) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + def _add(mlc_name, hf_name): + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=named_parameters[mlc_name].dtype, + ), + ) + + def _concat_add(mlc_name, hf_names): + mapping.add_mapping( + mlc_name, + hf_names, + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=named_parameters[mlc_name].dtype, + ), + ) + + _add("lm_head.linear.weight", "lm_head.weight") + _add("lm_head.linear.bias", "lm_head.bias") + _add("lm_head.ln.weight", "model.final_layernorm.weight") + _add("lm_head.ln.bias", "model.final_layernorm.bias") + _add("transformer.embd.weight", "model.embed_tokens.weight") + + prefix = "transformer.h" + hf_prefix = "model.layers" + for i in range(model_config.num_hidden_layers): + _add(f"{prefix}.{i}.ln.weight", f"{hf_prefix}.{i}.input_layernorm.weight") + _add(f"{prefix}.{i}.ln.bias", f"{hf_prefix}.{i}.input_layernorm.bias") + _concat_add( + f"{prefix}.{i}.mixer.Wqkv.weight", + [ + f"{hf_prefix}.{i}.self_attn.q_proj.weight", + f"{hf_prefix}.{i}.self_attn.k_proj.weight", + f"{hf_prefix}.{i}.self_attn.v_proj.weight", + ], + ) + _concat_add( + f"{prefix}.{i}.mixer.Wqkv.bias", + [ + f"{hf_prefix}.{i}.self_attn.q_proj.bias", + f"{hf_prefix}.{i}.self_attn.k_proj.bias", + f"{hf_prefix}.{i}.self_attn.v_proj.bias", + ], + ) + _add( + f"{prefix}.{i}.mixer.out_proj.weight", + f"{hf_prefix}.{i}.self_attn.dense.weight", + ) + _add(f"{prefix}.{i}.mixer.out_proj.bias", f"{hf_prefix}.{i}.self_attn.dense.bias") + _add(f"{prefix}.{i}.mlp.fc1.weight", f"{hf_prefix}.{i}.mlp.fc1.weight") + _add(f"{prefix}.{i}.mlp.fc1.bias", f"{hf_prefix}.{i}.mlp.fc1.bias") + _add(f"{prefix}.{i}.mlp.fc2.weight", f"{hf_prefix}.{i}.mlp.fc2.weight") + _add(f"{prefix}.{i}.mlp.fc2.bias", f"{hf_prefix}.{i}.mlp.fc2.bias") + mapping.add_unused(f"{hf_prefix}.{i}.mixer.rotary_emb.inv_freq") + + return mapping diff --git a/python/mlc_llm/model/phi/phi_model.py b/python/mlc_llm/model/phi/phi_model.py new file mode 100644 index 0000000..ae78ed2 --- /dev/null +++ b/python/mlc_llm/model/phi/phi_model.py @@ -0,0 +1,488 @@ +""" +Implementation for Phi architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional, Union # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Phi1Config(ConfigBase): + """Configuration of the Phi-1/Phi-1.5 model.""" + + vocab_size: int = 51200 + hidden_size: int = 2048 + intermediate_size: int = 8192 + num_hidden_layers: int = 24 + num_attention_heads: int = 32 + layer_norm_eps: float = 1e-5 + position_embedding_base: int = 0 + partial_rotary_factor: float = 0.5 + num_key_value_heads: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + head_dim: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + if self.num_key_value_heads == 0 or self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + if self.intermediate_size == 0 or self.intermediate_size is None: + self.intermediate_size = 4 * self.hidden_size + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + assert self.num_attention_heads % self.num_key_value_heads == 0 + + +@dataclasses.dataclass +class PhiConfig(ConfigBase): + """Configuration of the Phi-2 model.""" + + model_type: str # "phi", "phi-msft", "mixformer-sequential" + vocab_size: int = 51200 + n_positions: int = 2048 + n_embd: int = 2560 + n_layer: int = 32 + n_inner: int = 0 + n_head: int = 32 + rotary_dim: int = 32 + position_embedding_base: int = 0 + layer_norm_epsilon: float = 1e-5 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + n_head_kv: int = 0 + head_dim: int = 0 + tensor_parallel_shards: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + self.context_window_size = self.n_positions + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + "n_positions", + self.context_window_size, + ) + if self.prefill_chunk_size == 0: + self.prefill_chunk_size = self.context_window_size + self.prefill_chunk_size = min(self.prefill_chunk_size, self.context_window_size) + if self.n_head_kv == 0 or self.n_head_kv is None: + self.n_head_kv = self.n_head + if self.n_inner == 0 or self.n_inner is None: + self.n_inner = 4 * self.n_embd + if self.head_dim == 0: + self.head_dim = self.n_embd // self.n_head + assert self.head_dim * self.n_head == self.n_embd + assert self.n_head % self.n_head_kv == 0 + + @staticmethod + def from_phi1(config: Phi1Config) -> "PhiConfig": + "Build PhiConig from a Phi1Config." + return PhiConfig( + model_type="phi", + vocab_size=config.vocab_size, + n_positions=config.context_window_size, + n_embd=config.hidden_size, + n_layer=config.num_hidden_layers, + n_inner=config.intermediate_size, + n_head=config.num_attention_heads, + rotary_dim=int(config.partial_rotary_factor * config.head_dim), + position_embedding_base=config.position_embedding_base, + layer_norm_epsilon=config.layer_norm_eps, + context_window_size=config.context_window_size, + prefill_chunk_size=config.prefill_chunk_size, + n_head_kv=config.num_key_value_heads, + head_dim=config.head_dim, + tensor_parallel_shards=config.tensor_parallel_shards, + kwargs=config.kwargs, + ) + + +class PhiMLP(nn.Module): + def __init__(self, config: PhiConfig): + super().__init__() + if config.n_inner % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.n_inner} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.n_inner // config.tensor_parallel_shards + self.fc1 = nn.Linear(config.n_embd, self.intermediate_size) + self.fc2 = nn.Linear(self.intermediate_size, config.n_embd) + + def forward(self, hidden_states: Tensor): + hidden_states = self.fc1(hidden_states) + hidden_states = op.gelu(hidden_states, approximate="tanh") + hidden_states = self.fc2(hidden_states) + + return hidden_states + + +class PhiMHA(nn.Module): + def __init__(self, config: PhiConfig): + self.num_q_heads = config.n_head // config.tensor_parallel_shards + assert config.n_head % config.tensor_parallel_shards == 0, ( + f"n_head({config.n_head}) must be divisible by tensor_parallel_shards" + ) + self.n_head_kv = config.n_head_kv // config.tensor_parallel_shards + assert config.n_head_kv % config.tensor_parallel_shards == 0, ( + f"n_head({config.n_head_kv}) must be divisible by tensor_parallel_shards" + ) + self.head_dim = config.head_dim + op_size = self.head_dim * (self.num_q_heads + 2 * self.n_head_kv) + hidden_size = config.n_embd + + self.Wqkv = nn.Linear(hidden_size, op_size, bias=True) + self.out_proj = nn.Linear(self.num_q_heads * self.head_dim, hidden_size, bias=True) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.n_head_kv + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.Wqkv(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.out_proj(output) + + +class PhiParallelBlock(nn.Module): + def __init__(self, config: PhiConfig): + super().__init__() + + self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + self.mixer = PhiMHA(config) + self.mlp = PhiMLP(config) + + def _set_tp(): + def _set(param, hint): + param.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.mixer.num_q_heads * hd + k = self.mixer.n_head_kv * hd + v = self.mixer.n_head_kv * hd + _set( + self.mixer.Wqkv.weight, + tp.ShardSingleDim("_shard_qkv_weight", segs=[q, k, v], dim=0), + ) + _set( + self.mixer.Wqkv.bias, + tp.ShardSingleDim("_shard_qkv_bias", segs=[q, k, v], dim=0), + ) + _set(self.mixer.out_proj.weight, tp.ShardSingleDim("_shard_o_weight", dim=1)) + _set(self.mlp.fc1.weight, tp.ShardSingleDim("_shard_mlp_fc1_weight", dim=0)) + _set(self.mlp.fc1.bias, tp.ShardSingleDim("_shard_mlp_fc1_bias", dim=0)) + _set(self.mlp.fc2.weight, tp.ShardSingleDim("_shard_mlp_fc2_weight", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + residual = hidden_states + hidden_states = self.ln(hidden_states) + + with ( + tp.shard_bias(self.mixer.out_proj, self.tensor_parallel_shards), + tp.shard_bias(self.mlp.fc2, self.tensor_parallel_shards), + ): + attn_outputs = self.mixer(hidden_states, paged_kv_cache, layer_id) + feed_forward_hidden_states = self.mlp(hidden_states) + + hidden_states = self._apply_parallel_residual( + attn_outputs, feed_forward_hidden_states, residual + ) + + return hidden_states + + def _apply_parallel_residual(self, attn_out, mlp_out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce( + attn_out + mlp_out + residual / self.tensor_parallel_shards, "sum" + ) + return attn_out + mlp_out + residual + + +class PhiCausalLMHead(nn.Module): + def __init__(self, config: PhiConfig) -> None: + super().__init__() + + self.ln = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon) + self.linear = nn.Linear(config.n_embd, "vocab_size") + + def forward(self, hidden_states: Tensor): + hidden_states = self.ln(hidden_states) + logits = self.linear(hidden_states) + + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + +class PhiModel(nn.Module): + def __init__(self, config: PhiConfig) -> None: + super().__init__() + self.embd = nn.Embedding(config.vocab_size, config.n_embd) + self.h = nn.ModuleList([PhiParallelBlock(config) for _ in range(config.n_layer)]) + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.h): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + + return hidden_states + + +class PhiForCausalLM(nn.Module): + def __init__(self, config: Union[PhiConfig, Phi1Config]) -> None: + super().__init__() + + if isinstance(config, Phi1Config): + config = PhiConfig.from_phi1(config) + + self.transformer = PhiModel(config) + self.lm_head = PhiCausalLMHead(config) + self.num_hidden_layers = config.n_layer + self.num_attention_heads = config.n_head + self.num_key_value_heads = config.n_head_kv + self.head_dim = config.head_dim + self.hidden_size = config.n_embd + self.vocab_size = config.vocab_size + self.rope_theta = config.position_embedding_base + self.tensor_parallel_shards = config.tensor_parallel_shards + self.rotary_dim = config.rotary_dim + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.transformer(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + lm_logits = self.lm_head(hidden_states) + if lm_logits.dtype != "float32": + lm_logits = lm_logits.astype("float32") + return lm_logits + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + + if logits.dtype != "float32": + logits = logits.astype("float32") + + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + embeds = self.transformer.embd(input_ids) + return embeds + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rotary_dim=self.rotary_dim, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/phi3/__init__.py b/python/mlc_llm/model/phi3/__init__.py new file mode 100644 index 0000000..7f13ac8 --- /dev/null +++ b/python/mlc_llm/model/phi3/__init__.py @@ -0,0 +1,3 @@ +"""Common `nn.Modules` used to define LLMs in this project.""" + +from .phi3_model import Phi3Model diff --git a/python/mlc_llm/model/phi3/phi3_loader.py b/python/mlc_llm/model/phi3/phi3_loader.py new file mode 100644 index 0000000..f0f8a24 --- /dev/null +++ b/python/mlc_llm/model/phi3/phi3_loader.py @@ -0,0 +1,79 @@ +""" +This file specifies how MLC's Phi parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .phi3_model import Phi3Config, Phi3ForCausalLM + + +def phi3_huggingface(model_config: Phi3Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of Phi-1/Phi-1.5 HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : PhiConfig + The configuration of the Phi model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Phi3ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params = model.export_tvm(spec=model.get_default_spec()) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + def _add(mlc_name, hf_name): + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=named_parameters[mlc_name].dtype, + ), + ) + + # Skip lm_head.weight if tie_word_embeddings is enabled + if not getattr(model_config, "tie_word_embeddings", False): + _add("lm_head.weight", "lm_head.weight") + _add("transformer.norm.weight", "model.norm.weight") + _add("transformer.embd.weight", "model.embed_tokens.weight") + + prefix = "transformer.h" + hf_prefix = "model.layers" + for i in range(model_config.num_hidden_layers): + _add(f"{prefix}.{i}.ln.weight", f"{hf_prefix}.{i}.input_layernorm.weight") + _add( + f"{prefix}.{i}.mlp.down_proj.weight", + f"{hf_prefix}.{i}.mlp.down_proj.weight", + ) + _add( + f"{prefix}.{i}.mlp.gate_up_proj.weight", + f"{hf_prefix}.{i}.mlp.gate_up_proj.weight", + ) + _add( + f"{prefix}.{i}.post_attention_layernorm.weight", + f"{hf_prefix}.{i}.post_attention_layernorm.weight", + ) + _add( + f"{prefix}.{i}.mixer.out_proj.weight", + f"{hf_prefix}.{i}.self_attn.o_proj.weight", + ) + _add( + f"{prefix}.{i}.mixer.qkv_proj.weight", + f"{hf_prefix}.{i}.self_attn.qkv_proj.weight", + ) + return mapping diff --git a/python/mlc_llm/model/phi3/phi3_model.py b/python/mlc_llm/model/phi3/phi3_model.py new file mode 100644 index 0000000..dababdd --- /dev/null +++ b/python/mlc_llm/model/phi3/phi3_model.py @@ -0,0 +1,414 @@ +""" +Implementation for Phi-3 architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Phi3Config(ConfigBase): + """Configuration of the Phi-3 model.""" + + model_type: str # "phi", "phi-msft", "mixformer-sequential" + hidden_size: int + vocab_size: int + num_hidden_layers: int + num_attention_heads: int + intermediate_size: int + rms_norm_eps: float + num_key_value_heads: int + max_position_embeddings: int + position_embedding_base: int = 0 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + original_max_position_embeddings: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + head_dim: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + tie_word_embeddings: bool = False + partial_rotary_factor: float = 1.0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.rope_scaling is not None: + if "type" not in self.rope_scaling: + self.rope_scaling = None + else: + if self.rope_scaling["type"] == "su": + self.rope_scaling["type"] = "longrope" + + assert self.rope_scaling["type"] == "longrope", ( + f"Unsupported RoPE scaling type {self.rope_scaling['rope_type']} for Phi3" + ) + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + ( + self.rope_scaling["max_position_embeddings"], + self.rope_scaling["original_max_position_embeddings"], + ) = ( + self.max_position_embeddings, + self.original_max_position_embeddings, + ) + + if self.context_window_size == 0: + self.context_window_size = self.max_position_embeddings + + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + if self.num_key_value_heads == 0 or self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + assert self.num_attention_heads % self.num_key_value_heads == 0 + + +class Phi3Embedding(nn.Embedding): + """The embedding module that can be shared with the final lm_head.""" + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class Phi3MLP(nn.Module): + def __init__(self, config: Phi3Config): + super().__init__() + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor): + up_states = self.gate_up_proj(hidden_states) + gate, up_states = nn.op.split(up_states, 2, axis=-1) + up_states = up_states * op.silu(gate) + return self.down_proj(up_states) + + +class PhiMHA(nn.Module): + def __init__(self, config: Phi3Config): + self.num_q_heads = config.num_attention_heads // config.tensor_parallel_shards + assert config.num_attention_heads % config.tensor_parallel_shards == 0, ( + f"num_attention_heads({config.num_attention_heads}) " + "must be divisible by tensor_parallel_shards" + ) + self.num_key_value_heads = config.num_key_value_heads // config.tensor_parallel_shards + assert config.num_key_value_heads % config.tensor_parallel_shards == 0, ( + f"num_attention_heads({config.num_key_value_heads}) " + "must be divisible by tensor_parallel_shards" + ) + self.head_dim = config.head_dim + + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_q_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=False, + ) + self.out_proj = nn.Linear(self.num_q_heads * self.head_dim, config.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_q_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + # QKV Projection + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + # Attention + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_q_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + return self.out_proj(output) + + +class Phi3ParallelBlock(nn.Module): + def __init__(self, config: Phi3Config): + super().__init__() + + self.ln = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.mixer = PhiMHA(config) + self.mlp = Phi3MLP(config) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.weight.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.mixer.num_q_heads * hd + k = self.mixer.num_key_value_heads * hd + v = self.mixer.num_key_value_heads * hd + i = self.mlp.intermediate_size + + _set( + self.mixer.qkv_proj, + tp.ShardSingleDim("_shard_qkv", segs=[q, k, v], dim=0), + ) + _set(self.mixer.out_proj, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + attn_outputs = self.mixer(self.ln(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_parallel_residual(attn_outputs, hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_parallel_residual(out, hidden_states) + return hidden_states + + def _apply_parallel_residual(self, mlp_out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(mlp_out + residual / self.tensor_parallel_shards, "sum") + return mlp_out + residual + + +class Phi3Model(nn.Module): + def __init__(self, config: Phi3Config) -> None: + super().__init__() + self.embd = Phi3Embedding(config.vocab_size, config.hidden_size) + self.h = nn.ModuleList([Phi3ParallelBlock(config) for _ in range(config.num_hidden_layers)]) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = input_embed + for layer_id, layer in enumerate(self.h): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Phi3ForCausalLM(nn.Module): + def __init__(self, config: Phi3Config) -> None: + super().__init__() + + self.transformer = Phi3Model(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False) + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_scaling = config.rope_scaling + self.rope_theta = config.position_embedding_base + self.rope_ext_factors = ( + (config.rope_scaling["long_factor"] + config.rope_scaling["short_factor"]) + if config.rope_scaling is not None + else None + ) + self.tensor_parallel_shards = config.tensor_parallel_shards + self.partial_rotary_factor = config.partial_rotary_factor + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def get_logits(self, hidden_states: Tensor): + op_ext.configure() + if self.tie_word_embeddings: + logits = self.transformer.embd.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.transformer(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return self.get_logits(hidden_states) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(input_embed, paged_kv_cache) + logits = self.get_logits(hidden_states) + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + embeds = self.transformer.embd(input_ids) + return embeds + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scaling=self.rope_scaling, + rope_scale=1, + rope_theta=self.rope_theta, + rope_ext_factors=self.rope_ext_factors, + rotary_dim=int(self.head_dim * self.partial_rotary_factor), + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/phi3v/__init__.py b/python/mlc_llm/model/phi3v/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/phi3v/phi3v_image.py b/python/mlc_llm/model/phi3v/phi3v_image.py new file mode 100644 index 0000000..73b6a64 --- /dev/null +++ b/python/mlc_llm/model/phi3v/phi3v_image.py @@ -0,0 +1,309 @@ +""" +Implementation for Phi architecture. +""" + +from tvm import relax, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Module, Tensor, op +from tvm.script import tirx as T + +from mlc_llm.model.vision import CLIPVisionModel +from mlc_llm.support.config import ConfigBase + + +# mypy: disable-error-code="attr-defined" +class ImageProjection(Module): + def __init__(self, config: ConfigBase): + super().__init__() + self.linear_1 = nn.Linear( + config.vision_config.hidden_size * 4, config.hidden_size, bias=True + ) + self.act = nn.GELU() + self.linear_2 = nn.Linear(config.hidden_size, config.hidden_size, bias=True) + + def forward(self, image_features: Tensor) -> Tensor: + shape_1 = tirx.Var("shape_1", "int64") + image_features = op.wrap_nested( + relax.BlockBuilder() + .current() + .match_cast( + image_features._expr, + relax.TensorType([shape_1, image_features.shape[1]], image_features.dtype), + ), + "image_features", + ) + + hidden_states = self.linear_1(image_features) + + shape_2 = tirx.Var("shape_2", "int64") + hidden_states = op.wrap_nested( + relax.BlockBuilder() + .current() + .match_cast( + hidden_states._expr, + relax.TensorType([shape_2, hidden_states.shape[1]], hidden_states.dtype), + ), + "hidden_states", + ) + + hidden_states = self.act(hidden_states) + hidden_states = self.linear_2(hidden_states) + return hidden_states + + +class Phi3ImageEmbedding(Module): + def __init__(self, config: ConfigBase): + super().__init__() + + self.img_processor = CLIPVisionModel(config.vision_config) + self.image_dim_out = config.img_processor["image_dim_out"] + + self.glb_GN = nn.Parameter((1, 1, self.image_dim_out * 4)) + self.sub_GN = nn.Parameter((1, 1, 1, self.image_dim_out * 4)) + + self.img_projection = ImageProjection(config) + self.image_size = config.vision_config.image_size + + def apply_schedule(self, sch, block, bdx=32, tile=[32, 32]): + loop_x, loop_y = sch.get_loops(block)[-2:] + xo, xi = sch.split(loop_x, factors=[tile[0], None]) + yo, yi = sch.split(loop_y, factors=[tile[1], None]) + sch.reorder(xo, yo, xi, yi) + t = sch.fuse(xo, yo) + ty, tx = sch.split(t, factors=[None, bdx]) + sch.bind(ty, "threadIdx.y") + sch.bind(tx, "threadIdx.x") + + def dyn_repeat_4d_tensor(self, input_tensor, r0, r1, r2, r3) -> Tensor: + assert 4 == input_tensor.ndim, "input_tensor should be 4D data tensor" + + def create_dyn_repeat_func(dtype): + @T.prim_func(s_tir=True) + def dyn_repeat_4d_tensor_func( # pylint disable=too-many-locals + input_tensor: T.handle, + output: T.handle, + ch0: T.int64(), + ch1: T.int64(), + ch2: T.int64(), + ch3: T.int64(), + ): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64() + input_tensor_buf = T.match_buffer(input_tensor, (n, c, h, w), dtype=dtype) + out_buf = T.match_buffer(output, (n * ch0, c * ch1, h * ch2, w * ch3), dtype=dtype) + + for n_idx in T.thread_binding(n * ch0, thread="blockIdx.x"): + for c_idx in T.thread_binding(c * ch1, thread="blockIdx.y"): + for h_idx, w_idx in T.grid(h * ch2, w * ch3): + with T.sblock("dyn_repeat_4d_tensor"): + T.reads(input_tensor_buf[n_idx, c_idx, h_idx, w_idx]) + T.writes(out_buf[n_idx, c_idx, h_idx, w_idx]) + out_buf[n_idx, c_idx, h_idx, w_idx] = input_tensor_buf[ + n_idx % n, c_idx % c, h_idx % h, w_idx % w + ] + + return dyn_repeat_4d_tensor_func + + n, c, h, w = input_tensor.shape + out = op.tensor_ir_op( + create_dyn_repeat_func(input_tensor.dtype), + "dyn_repeat_4d_tensor", + [input_tensor, r0, r1, r2, r3], + [Tensor.placeholder([n * r0, c * r1, h * r2, w * r3], input_tensor.dtype)], + ) + return out + + def dyn_concate_dim_2(self, input_1, input_2) -> Tensor: + def create_dyn_concate_func(dtype): + @T.prim_func(s_tir=True) + def dyn_concate_dim_2_func(input_1: T.handle, input_2: T.handle, output: T.handle): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + n, c, h1, h2, w = T.int64(), T.int64(), T.int64(), T.int64(), T.int64() + input_1_buf = T.match_buffer(input_1, (n, c, h1, w), dtype=dtype) + input_2_buf = T.match_buffer(input_2, (n, c, h2, w), dtype=dtype) + out_buf = T.match_buffer(output, (n, c, h1 + h2, w), dtype=dtype) + + for n_idx in T.thread_binding(n, thread="blockIdx.x"): + for c_idx in T.thread_binding(c, thread="blockIdx.y"): + for h_idx, w_idx in T.grid(h1 + h2, w): + with T.sblock("dyn_concate_dim_2"): + T.reads(input_1_buf[n_idx, c_idx, h_idx, w_idx]) + T.writes(out_buf[n_idx, c_idx, h_idx, w_idx]) + if h_idx < h1: + out_buf[n_idx, c_idx, h_idx, w_idx] = input_1_buf[ + n_idx, c_idx, h_idx, w_idx + ] + else: + out_buf[n_idx, c_idx, h_idx, w_idx] = input_2_buf[ + n_idx, c_idx, h_idx - h1, w_idx + ] + + return dyn_concate_dim_2_func + + n1, c1, h1, w1 = input_1.shape + n2, c2, h2, w2 = input_2.shape + assert n1 == n2 and c1 == c2 and w1 == w2 + + out = op.tensor_ir_op( + create_dyn_concate_func(input_1.dtype), + "dyn_concate_dim_2", + [input_1, input_2], + [Tensor.placeholder([n1, c1, h1 + h2, w1], input_1.dtype)], + ) + return out + + def dyn_concate_dim_1(self, input_1, input_2) -> Tensor: + def create_dyn_concate_func(dtype): + @T.prim_func(s_tir=True) + def dyn_concate_dim_1_func(input_1: T.handle, input_2: T.handle, output: T.handle): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + c, h1, h2, w = T.int64(), T.int64(), T.int64(), T.int64() + input_1_buf = T.match_buffer(input_1, (c, h1, w), dtype=dtype) + input_2_buf = T.match_buffer(input_2, (c, h2, w), dtype=dtype) + out_buf = T.match_buffer(output, (c, h1 + h2, w), dtype=dtype) + + for c_idx in T.thread_binding(c, thread="blockIdx.y"): + for h_idx, w_idx in T.grid(h1 + h2, w): + with T.sblock("dyn_concate_dim_1"): + T.reads(input_1_buf[c_idx, h_idx, w_idx]) + T.writes(out_buf[c_idx, h_idx, w_idx]) + if h_idx < h1: + out_buf[c_idx, h_idx, w_idx] = input_1_buf[c_idx, h_idx, w_idx] + else: + out_buf[c_idx, h_idx, w_idx] = input_2_buf[c_idx, h_idx - h1, w_idx] + + return dyn_concate_dim_1_func + + c1, h1, w1 = input_1.shape + c2, h2, w2 = input_2.shape + assert c1 == c2 and w1 == w2 + + out = op.tensor_ir_op( + create_dyn_concate_func(input_1.dtype), + "dyn_concate", + [input_1, input_2], + [Tensor.placeholder([c1, h1 + h2, w1], input_1.dtype)], + ) + return out + + def get_img_features(self, img_embeds: Tensor) -> Tensor: + img_processor_output = self.img_processor(img_embeds) + patch_feature = nn.op.split(img_processor_output, indices_or_sections=[1], axis=1) + return patch_feature[1] + + def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop): + N, L, C = image_features.shape + num_images = 1 + H = int(L**0.5) + image_features = nn.op.reshape(image_features, ([N, H, H, C])) # N, 24, 24, 1024 + image_features = nn.op.reshape( + image_features, ([N, H // 2, 2, H // 2, 2, C]) + ) # N, 12, 2, 12, 2, 1024 + + new_s1 = tirx.Var("new_s1", "int64") + new_s2 = tirx.Var("new_s2", "int64") + + image_features = op.wrap_nested( + relax.BlockBuilder() + .current() + .match_cast( + image_features._expr, + relax.TensorType( + [ + image_features.shape[0], + new_s1, + image_features.shape[2], + new_s2, + image_features.shape[4], + image_features.shape[5], + ], + image_features.dtype, + ), + ), + "image_features_1", + ) + + image_features = nn.op.permute_dims( + image_features, axes=([0, 1, 3, 2, 4, 5]) + ) # N, 12, 12, 2, 2, 1024 + image_features = nn.op.reshape(image_features, ([N, -1, 4 * C])) # N, 144, 4096 + image_features = nn.op.reshape( + image_features, ([num_images, h_crop, w_crop, H // 2, H // 2, -1]) + ) + + new_s3 = tirx.Var("new_s3", "int64") + new_s4 = tirx.Var("new_s4", "int64") + + image_features = op.wrap_nested( + relax.BlockBuilder() + .current() + .match_cast( + image_features._expr, + relax.TensorType( + [ + image_features.shape[0], + new_s3, + image_features.shape[2], + new_s4, + image_features.shape[4], + image_features.shape[5], + ], + image_features.dtype, + ), + ), + "image_features_2", + ) + + image_features = nn.op.permute_dims(image_features, axes=([0, 1, 3, 2, 4, 5])) + image_features_hd = nn.op.reshape( + image_features, ([num_images, h_crop * H // 2, w_crop * H // 2, 4 * C]) + ) + + return image_features_hd + + def add_image_newline(self, image_features_hd): + """ + image_features_hd: (num_images, h_crop*12, w_crop*12, 4096) + output: (num_images, (h_crop*12) * (w_crop*12+1), 4096) + """ + num_images, h, w, hid_dim = image_features_hd.shape + temp_sub_GN = self.dyn_repeat_4d_tensor( + self.sub_GN, T.int64(1), T.int64(h), T.int64(1), T.int64(1) + ) + image_features_hd_newline = self.dyn_concate_dim_2(image_features_hd, temp_sub_GN) + image_features_hd_newline = nn.op.reshape( + image_features_hd_newline, ([num_images, -1, hid_dim]) + ) + return image_features_hd_newline + + def forward(self, pixel_values: Tensor, h_crop, w_crop) -> Tensor: + img_features = self.get_img_features(pixel_values) + img_features = nn.op.split(img_features, indices_or_sections=[1], axis=0) + + global_image_features = img_features[0] + global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1) + global_image_features_hd_newline = self.add_image_newline(global_image_features_hd) + + sub_image_features = img_features[1] + sub_image_features_hd = self.reshape_hd_patches_2x2merge(sub_image_features, h_crop, w_crop) + sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd) + + global_image_features_hd = nn.op.squeeze(global_image_features_hd_newline, 0) + + combined_image = self.dyn_concate_dim_1(sub_image_features_hd_newline, self.glb_GN) + combined_image = self.dyn_concate_dim_1(combined_image, global_image_features_hd_newline) + combined_image = nn.op.squeeze(combined_image, 0) + + new_s7 = tirx.Var("new_s7", "int64") + + combined_image = op.wrap_nested( + relax.BlockBuilder() + .current() + .match_cast( + combined_image._expr, + relax.TensorType([new_s7, combined_image.shape[1]], combined_image.dtype), + ), + "combined_image", + ) + output_image = self.img_projection(combined_image) + return output_image diff --git a/python/mlc_llm/model/phi3v/phi3v_loader.py b/python/mlc_llm/model/phi3v/phi3v_loader.py new file mode 100644 index 0000000..72244cf --- /dev/null +++ b/python/mlc_llm/model/phi3v/phi3v_loader.py @@ -0,0 +1,132 @@ +""" +This file specifies how MLC's Phi parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .phi3v_model import Phi3VConfig, Phi3VForCausalLM + + +def huggingface(model_config: Phi3VConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of Phi-1/Phi-1.5 HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : PhiConfig + The configuration of the Phi model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Phi3VForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params = model.export_tvm(spec=model.get_default_spec()) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + def _add(mlc_name, hf_name=None): + if None is hf_name: + hf_name = mlc_name + + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=named_parameters[mlc_name].dtype, + ), + ) + + def _add_vision(name): + _add(name, "model." + name) + + _add("model.embd.weight", "model.embed_tokens.weight") + + prefix = "model.h" + hf_prefix = "model.layers" + for i in range(model_config.num_hidden_layers): + _add(f"{prefix}.{i}.ln.weight", f"{hf_prefix}.{i}.input_layernorm.weight") + _add( + f"{prefix}.{i}.mlp.down_proj.weight", + f"{hf_prefix}.{i}.mlp.down_proj.weight", + ) + _add( + f"{prefix}.{i}.mlp.gate_up_proj.weight", + f"{hf_prefix}.{i}.mlp.gate_up_proj.weight", + ) + _add( + f"{prefix}.{i}.post_attention_layernorm.weight", + f"{hf_prefix}.{i}.post_attention_layernorm.weight", + ) + _add( + f"{prefix}.{i}.mixer.out_proj.weight", + f"{hf_prefix}.{i}.self_attn.o_proj.weight", + ) + _add( + f"{prefix}.{i}.mixer.qkv_proj.weight", + f"{hf_prefix}.{i}.self_attn.qkv_proj.weight", + ) + + prefix = "vision_embed_tokens.img_processor.vision_model.encoder.layers" + for i in range(model_config.vision_config.num_hidden_layers): + _add_vision(f"{prefix}.{i}.layer_norm1.bias") + _add_vision(f"{prefix}.{i}.layer_norm1.weight") + _add_vision(f"{prefix}.{i}.layer_norm2.bias") + _add_vision(f"{prefix}.{i}.layer_norm2.weight") + _add_vision(f"{prefix}.{i}.mlp.fc1.bias") + _add_vision(f"{prefix}.{i}.mlp.fc1.weight") + _add_vision(f"{prefix}.{i}.mlp.fc2.bias") + _add_vision(f"{prefix}.{i}.mlp.fc2.weight") + _add_vision(f"{prefix}.{i}.self_attn.k_proj.bias") + _add_vision(f"{prefix}.{i}.self_attn.k_proj.weight") + _add_vision(f"{prefix}.{i}.self_attn.out_proj.bias") + _add_vision(f"{prefix}.{i}.self_attn.out_proj.weight") + _add_vision(f"{prefix}.{i}.self_attn.q_proj.bias") + _add_vision(f"{prefix}.{i}.self_attn.q_proj.weight") + _add_vision(f"{prefix}.{i}.self_attn.v_proj.bias") + _add_vision(f"{prefix}.{i}.self_attn.v_proj.weight") + + _add_vision("vision_embed_tokens.sub_GN") + _add_vision("vision_embed_tokens.glb_GN") + _add_vision("vision_embed_tokens.img_processor.vision_model.embeddings.class_embedding") + _add_vision("vision_embed_tokens.img_processor.vision_model.embeddings.patch_embedding.weight") + _add_vision( + "vision_embed_tokens.img_processor.vision_model.embeddings.position_embedding.weight" + ) + _add_vision("vision_embed_tokens.img_processor.vision_model.post_layernorm.bias") + _add_vision("vision_embed_tokens.img_processor.vision_model.post_layernorm.weight") + _add_vision("vision_embed_tokens.img_processor.vision_model.pre_layrnorm.bias") + _add_vision("vision_embed_tokens.img_processor.vision_model.pre_layrnorm.weight") + + prefix = "vision_embed_tokens.img_projection" + _add(f"{prefix}.linear_1.bias", f"model.{prefix}.0.bias") + _add(f"{prefix}.linear_1.weight", f"model.{prefix}.0.weight") + _add(f"{prefix}.linear_2.bias", f"model.{prefix}.2.bias") + _add(f"{prefix}.linear_2.weight", f"model.{prefix}.2.weight") + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + mapping.add_unused("model.embed_tokens.weight") + + return mapping diff --git a/python/mlc_llm/model/phi3v/phi3v_model.py b/python/mlc_llm/model/phi3v/phi3v_model.py new file mode 100644 index 0000000..ef6d2cb --- /dev/null +++ b/python/mlc_llm/model/phi3v/phi3v_model.py @@ -0,0 +1,390 @@ +""" +Implementation for Phi architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import relax, target, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.model.phi3 import Phi3Model +from mlc_llm.model.vision import CLIPVisionConfig, ImageProcessor +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +from .phi3v_image import Phi3ImageEmbedding + +logger = logging.getLogger(__name__) + +CLIPVISION_DEFAULT_CONFIG = { + "hidden_size": 1024, + "image_size": 336, + "intermediate_size": 4096, + "num_attention_heads": 16, + "num_hidden_layers": 24, + "patch_size": 14, + "projection_dim": 768, + "layer_norm_eps": 1e-05, + "vocab_size": None, +} + + +@dataclasses.dataclass +class Phi3VConfig(ConfigBase): + """Configuration of the Phi-3 Vision model.""" + + model_type: str + hidden_size: int + vocab_size: int + num_hidden_layers: int + num_attention_heads: int + intermediate_size: int + rms_norm_eps: float + num_key_value_heads: int + max_position_embeddings: int + vision_config: CLIPVisionConfig = None + img_processor: Optional[Dict[str, Any]] = None # noqa: UP006 + position_embedding_base: int = 0 + rope_scaling: Optional[Dict[str, Any]] = None # noqa: UP006 + original_max_position_embeddings: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + head_dim: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + vision_config_dict: Dict[str, Any] # noqa: UP006 + if isinstance(self.vision_config, CLIPVisionConfig): + vision_config_dict = dataclasses.asdict(self.vision_config) + else: + vision_config_dict = dict(CLIPVISION_DEFAULT_CONFIG) + + for k, v in vision_config_dict.pop("kwargs", {}).items(): + vision_config_dict[k] = v + + self.vision_config = CLIPVisionConfig.from_dict(vision_config_dict) + + if self.position_embedding_base == 0: + if "rope_theta" in self.kwargs: + self.position_embedding_base = self.kwargs.pop("rope_theta") + else: + self.position_embedding_base = 10000 + if self.rope_scaling is not None: + if "type" not in self.rope_scaling: + self.rope_scaling = None + else: + if self.rope_scaling["type"] == "su": + self.rope_scaling["type"] = "longrope" + + assert self.rope_scaling["type"] == "longrope", ( + f"Unsupported RoPE scaling type {self.rope_scaling['rope_type']} for Phi3" + ) + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + ( + self.rope_scaling["max_position_embeddings"], + self.rope_scaling["original_max_position_embeddings"], + ) = ( + self.max_position_embeddings, + self.original_max_position_embeddings, + ) + + if self.context_window_size == 0: + self.context_window_size = self.max_position_embeddings + + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + if self.num_key_value_heads == 0 or self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + assert self.num_attention_heads % self.num_key_value_heads == 0 + + +# mypy: disable-error-code="arg-type,annotation-unchecked" +class Phi3VForCausalLM(nn.Module): + def __init__(self, config: Phi3VConfig) -> None: + super().__init__() + + self.config = config + self.model = Phi3Model(config) + self.lm_head = nn.Linear(config.hidden_size, "vocab_size", bias=False) + self.vision_embed_tokens = Phi3ImageEmbedding(config) + self.image_processor = ImageProcessor() + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.rope_scaling = config.rope_scaling + self.rope_theta = config.position_embedding_base + self.rope_ext_factors = ( + (config.rope_scaling["long_factor"] + config.rope_scaling["short_factor"]) + if config.rope_scaling is not None + else None + ) + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + self.image_dtype = ( + "uint32" + if target.Target.current() and target.Target.current().kind.name == "webgpu" + else "uint8" + ) + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + lm_logits = self.lm_head(hidden_states) + if lm_logits.dtype != "float32": + lm_logits = lm_logits.astype("float32") + return lm_logits + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + + if logits.dtype != "float32": + logits = logits.astype("float32") + + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + embeds = self.model.embd(input_ids) + return embeds + + def image_preprocess( + self, pixel_values: Tensor, resized_height, resized_width, num_crops=16 + ) -> Tensor: + pixel_values = op.permute_dims(pixel_values, axes=(0, 3, 1, 2)) # NHWC -> NCHW + pixel_values = self.image_processor.resize( + pixel_values, params={"height": resized_height, "width": resized_width} + ) + pixel_values = self.image_processor.pad(pixel_values, dtype=self.image_dtype) + pixel_values = self.image_processor.rescale(pixel_values) + pixel_values = self.image_processor.normalize(pixel_values) + global_image = self.image_processor.resize( + pixel_values, params={"height": 336, "width": 336} + ) + global_image = op.wrap_nested( + relax.BlockBuilder() + .current() + .match_cast( + global_image._expr, + relax.TensorType( + [global_image.shape[0], global_image.shape[1], 336, 336], + global_image.dtype, + ), + ), + "global_image", + ) + + n, c, h, w = pixel_values.shape + assert isinstance(h, tirx.Mul) and isinstance(h.b, tirx.IntImm) and h.b.value == 336 + pixel_values = op.reshape(pixel_values, shape=(1, 3, h.a, 336, w // 336, 336)) + pixel_values = op.permute_dims(pixel_values, axes=(0, 2, 4, 1, 3, 5)) + pixel_values = op.reshape(pixel_values, shape=(-1, 3, 336, 336)) + combined_image = op.concat([global_image, pixel_values], dim=0) + + # pad to max num crops tensor + b, c, h, w = combined_image.shape + zeros = op.zeros((num_crops + 1 - b, c, h, w)) + combined_image = op.concat([combined_image, zeros], dim=0) + + combined_image = op.wrap_nested( + relax.BlockBuilder() + .current() + .match_cast( + combined_image._expr, + relax.TensorType([num_crops + 1, c, h, w], combined_image.dtype), + ), + "combined_image", + ) + + return combined_image + + def image_embed( + self, + pixel_values: Tensor, + resized_height, + resized_width, + crop_height, + crop_width, + ) -> Tensor: + n, h, w, c = pixel_values.shape + pixel_values = self.image_preprocess(pixel_values, resized_height, resized_width) + pixel_values = pixel_values.astype(self.dtype) + return self.vision_embed_tokens(pixel_values, crop_height, crop_width) + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scaling=self.rope_scaling, + rope_scale=1, + rope_theta=self.rope_theta, + rope_ext_factors=self.rope_ext_factors, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "image_embed": { + "pixel_values": nn.spec.Tensor( + [1, "image_height", "image_width", 3], self.image_dtype + ), + "resized_height": nn.spec.Int(), + "resized_width": nn.spec.Int(), + "crop_height": nn.spec.Int(), + "crop_width": nn.spec.Int(), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/qwen/__init__.py b/python/mlc_llm/model/qwen/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/qwen/qwen_loader.py b/python/mlc_llm/model/qwen/qwen_loader.py new file mode 100644 index 0000000..06a9093 --- /dev/null +++ b/python/mlc_llm/model/qwen/qwen_loader.py @@ -0,0 +1,16 @@ +""" +This file specifies how MLC's QWen parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .qwen_model import QWenLMHeadModel + +huggingface = make_standard_hf_loader( + model_cls=QWenLMHeadModel, + layer_prefix="transformer.h", + qkv_names=(), + include_qkv=False, + gate_up_names=("w1", "w2"), +) diff --git a/python/mlc_llm/model/qwen/qwen_model.py b/python/mlc_llm/model/qwen/qwen_model.py new file mode 100644 index 0000000..b641010 --- /dev/null +++ b/python/mlc_llm/model/qwen/qwen_model.py @@ -0,0 +1,359 @@ +""" +Implementation for QWEN architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class QWenConfig(ConfigBase): + """Configuration of the QWen model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + layer_norm_epsilon: float + scale_attn_weights: bool + kv_channels: int + rotary_emb_base: int + intermediate_size: int + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + head_dim: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class QWenAttention(nn.Module): + def __init__(self, config: QWenConfig): + self.hidden_size = config.hidden_size + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.num_attention_heads // config.tensor_parallel_shards + self.head_dim = config.head_dim + + self.c_attn = nn.Linear(config.hidden_size, 3 * self.num_heads * self.head_dim, bias=True) + + self.c_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + ): + d, h = self.head_dim, self.num_heads + b, s, _ = hidden_states.shape + + qkv = self.c_attn(hidden_states) + qkv = op.reshape(qkv, (b, s, 3 * h, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h * d), + ) + return self.c_proj(output) + + +class QWenMLP(nn.Module): + def __init__(self, config: QWenConfig): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=self.intermediate_size, + bias=False, + ) + self.c_proj = nn.Linear(self.intermediate_size // 2, config.hidden_size, bias=False) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.c_proj(x1 * op.silu(x2)) + + +class QWenBlock(nn.Module): + def __init__(self, config: QWenConfig): + rms_norm_eps = config.layer_norm_epsilon + self.attn = QWenAttention(config) + self.mlp = QWenMLP(config) + self.ln_1 = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + self.ln_2 = nn.RMSNorm(config.hidden_size, -1, rms_norm_eps, bias=False) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.attn.num_heads * hd + k = self.attn.num_heads * hd + v = self.attn.num_heads * hd + i = self.mlp.intermediate_size // 2 + _set( + self.attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set( + self.attn.c_attn.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.attn.c_proj.weight, tp.ShardSingleDim("_shard_attn_c_proj", dim=1)) + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_gate_up_proj", segs=[i, i], dim=0), + ) + _set(self.mlp.c_proj.weight, tp.ShardSingleDim("_shard_mlp_c_proj", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.attn(self.ln_1(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.ln_2(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class QWenModel(nn.Module): + def __init__(self, config: QWenConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.wte = nn.Embedding(config.vocab_size, config.hidden_size) + self.h = nn.ModuleList([QWenBlock(config) for _ in range(config.num_hidden_layers)]) + self.ln_f = nn.RMSNorm(config.hidden_size, -1, config.layer_norm_epsilon, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.h): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.ln_f(hidden_states) + return hidden_states + + +class QWenLMHeadModel(nn.Module): + def __init__(self, config: QWenConfig): + self.transformer = QWenModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False, dtype="float32") + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.tensor_parallel_shards = config.tensor_parallel_shards + self.rotary_emb_base = config.rotary_emb_base + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + inputs: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + hidden_states = self.transformer(inputs, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.transformer.wte(input_ids) + + def prefill(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(inputs, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.transformer(inputs, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill(self, inputs: Tensor, logit_positions: Tensor, paged_kv_cache: PagedKVCache): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(inputs, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(inputs, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(inputs, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_attention_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rotary_emb_base, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "inputs": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "inputs": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "inputs": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "inputs": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "inputs": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/qwen2/__init__.py b/python/mlc_llm/model/qwen2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/qwen2/qwen2_loader.py b/python/mlc_llm/model/qwen2/qwen2_loader.py new file mode 100644 index 0000000..cb1c599 --- /dev/null +++ b/python/mlc_llm/model/qwen2/qwen2_loader.py @@ -0,0 +1,15 @@ +""" +This file specifies how MLC's QWen2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .qwen2_model import QWen2LMHeadModel + +huggingface = make_standard_hf_loader( + model_cls=QWen2LMHeadModel, + qkv_target_name="c_attn", + add_qkv_bias=True, + add_unused=["rotary_emb.inv_freq"], +) diff --git a/python/mlc_llm/model/qwen2/qwen2_model.py b/python/mlc_llm/model/qwen2/qwen2_model.py new file mode 100644 index 0000000..71eb12e --- /dev/null +++ b/python/mlc_llm/model/qwen2/qwen2_model.py @@ -0,0 +1,407 @@ +""" +Implementation for QWEN2 architecture. +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class QWen2Config(ConfigBase): + """Configuration of the QWen2 model.""" + + hidden_act: str + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + num_key_value_heads: int + rms_norm_eps: float + rope_theta: int + vocab_size: int + tie_word_embeddings: bool = False + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + head_dim: int = 0 + dtype: str = "float32" + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class QWen2Attention(nn.Module): + def __init__(self, config: QWen2Config): + self.head_dim = config.head_dim + if config.num_key_value_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_key_value_heads} key-value attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_attention_heads = config.num_attention_heads // config.tensor_parallel_shards + self.num_key_value_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.rope_theta = config.rope_theta + + self.c_attn = nn.Linear( + in_features=config.hidden_size, + out_features=(2 * self.num_key_value_heads + self.num_attention_heads) * self.head_dim, + bias=True, + ) + self.o_proj = nn.Linear( + self.num_attention_heads * self.head_dim, config.hidden_size, bias=False + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_attention_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.c_attn(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_attention_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class Qwen2Embedding(nn.Embedding): + """The embedding module specialized for Qwen2 so that + it can be shared with the final lm_head. + """ + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class QWen2MLP(nn.Module): + def __init__(self, config: QWen2Config): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class QWen2DecoderLayer(nn.Module): + def __init__(self, config: QWen2Config): + self.self_attn = QWen2Attention(config) + self.mlp = QWen2MLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_attention_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set( + self.self_attn.c_attn.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj.weight, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.input_layernorm(hidden_states) + out = self.self_attn(out, paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class QWen2Model(nn.Module): + def __init__(self, config: QWen2Config): + self.embed_tokens = Qwen2Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [QWen2DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class QWen2LMHeadModel(nn.Module): + def __init__(self, config: QWen2Config): + self.model = QWen2Model(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.dtype = config.dtype + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.intermediate_size = config.intermediate_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.rms_norm_eps = config.rms_norm_eps + self.rope_theta = config.rope_theta + self.vocab_size = config.vocab_size + self.tensor_parallel_shards = config.tensor_parallel_shards + self.head_dim = config.head_dim + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/qwen2_5_vl/__init__.py b/python/mlc_llm/model/qwen2_5_vl/__init__.py new file mode 100644 index 0000000..ffc10e9 --- /dev/null +++ b/python/mlc_llm/model/qwen2_5_vl/__init__.py @@ -0,0 +1,3 @@ +"""Qwen2.5-VL architecture entry.""" + +from .qwen2_5_vl_model import Qwen25VLConfig, Qwen25VLLMHeadModel diff --git a/python/mlc_llm/model/qwen2_5_vl/qwen2_5_vl_model.py b/python/mlc_llm/model/qwen2_5_vl/qwen2_5_vl_model.py new file mode 100644 index 0000000..05a49e2 --- /dev/null +++ b/python/mlc_llm/model/qwen2_5_vl/qwen2_5_vl_model.py @@ -0,0 +1,570 @@ +"""Partial Qwen2.5-VL implementation focused on decoder-side MRoPE support. + +This file intentionally does not provide complete end-to-end Qwen2.5-VL support yet. +The current scope is: + +- decoder-side text model structure, +- multimodal rotary embedding generation/application, +- position-id layout compatibility for the MRoPE path. + +Missing pieces for full Qwen2.5-VL support include: + +- vision tower / multimodal projector and image-video embedding path, +- model registry / loader / quantization / preset integration, +- end-to-end multimodal preprocessing and compile/chat wiring. +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional, Tuple # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.op.mrope import ( + MultimodalRotaryEmbedding, + VisionPositionMetadata, + apply_multimodal_rotary_pos_emb, +) +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +@dataclasses.dataclass +class Qwen25VLVisionTokenConfig: + """Vision token IDs used by Qwen2.5-VL.""" + + image_token_id: int = 151655 + video_token_id: int = 151656 + vision_start_token_id: int = 151652 + vision_end_token_id: int = 151653 + + +@dataclasses.dataclass +class Qwen25VLVisionGridConfig: + """Vision grid configuration for multimodal position IDs.""" + + spatial_merge_size: int = 2 + temporal_patch_size: int = 2 + tokens_per_second: float = 4.0 + + +@dataclasses.dataclass(frozen=True) +class Qwen25VLAttentionState: + """Derived attention dimensions used across Qwen2.5-VL attention code paths.""" + + head_dim: int + num_attention_heads: int + num_key_value_heads: int + mrope_section: Tuple[int, int, int] # noqa: UP006 + softmax_scale: float + + +@dataclasses.dataclass +class Qwen25VLConfig(ConfigBase): + """Configuration for the Qwen2.5-VL model.""" + + hidden_act: str + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + num_key_value_heads: int + rms_norm_eps: float + rope_theta: float + vocab_size: int + tie_word_embeddings: bool = False + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + head_dim: int = 0 + dtype: str = "float32" + max_batch_size: int = 1 + rope_parameters: Optional[Dict[str, Any]] = None # noqa: UP006 + mrope_section: Optional[Tuple[int, int, int]] = None # noqa: UP006 + vision_tokens: Qwen25VLVisionTokenConfig = dataclasses.field( + default_factory=Qwen25VLVisionTokenConfig + ) + vision_grid: Qwen25VLVisionGridConfig = dataclasses.field( + default_factory=Qwen25VLVisionGridConfig + ) + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + if self.prefill_chunk_size == 0: + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + self.prefill_chunk_size = min(self.context_window_size, 8192) + + rope_scaling = self.kwargs.pop("rope_scaling", None) + if self.rope_parameters is None: + self.rope_parameters = rope_scaling or {} + if self.mrope_section is None: + section = self.rope_parameters.get("mrope_section") + if section is None and rope_scaling is not None: + section = rope_scaling.get("mrope_section") + if section is None: + raise ValueError("`mrope_section` must be provided for Qwen2.5-VL.") + self.mrope_section = tuple(int(i) for i in section) + if len(self.mrope_section) != 3: + raise ValueError(f"mrope_section must contain 3 integers, got {self.mrope_section}.") + + for key in [ + "image_token_id", + "video_token_id", + "vision_start_token_id", + "vision_end_token_id", + ]: + if key in self.kwargs: + setattr(self.vision_tokens, key, int(self.kwargs.pop(key))) + for key in ["spatial_merge_size", "temporal_patch_size"]: + if key in self.kwargs: + setattr(self.vision_grid, key, int(self.kwargs.pop(key))) + if "tokens_per_second" in self.kwargs: + self.vision_grid.tokens_per_second = float(self.kwargs.pop("tokens_per_second")) + + vision_cfg = self.kwargs.pop("vision_config", {}) + if vision_cfg: + if not isinstance(vision_cfg, dict): + raise ValueError(f"vision_config must be a dict, got {type(vision_cfg)}.") + self.vision_grid.spatial_merge_size = int( + vision_cfg.get("spatial_merge_size", self.vision_grid.spatial_merge_size) + ) + self.vision_grid.temporal_patch_size = int( + vision_cfg.get("temporal_patch_size", self.vision_grid.temporal_patch_size) + ) + self.vision_grid.tokens_per_second = float( + vision_cfg.get("tokens_per_second", self.vision_grid.tokens_per_second) + ) + + @property + def image_token_id(self) -> int: + return self.vision_tokens.image_token_id + + @property + def video_token_id(self) -> int: + return self.vision_tokens.video_token_id + + @property + def vision_start_token_id(self) -> int: + return self.vision_tokens.vision_start_token_id + + @property + def vision_end_token_id(self) -> int: + return self.vision_tokens.vision_end_token_id + + @property + def spatial_merge_size(self) -> int: + return self.vision_grid.spatial_merge_size + + @property + def temporal_patch_size(self) -> int: + return self.vision_grid.temporal_patch_size + + @property + def tokens_per_second(self) -> float: + return self.vision_grid.tokens_per_second + + @property + def vision_metadata(self) -> VisionPositionMetadata: + return VisionPositionMetadata( + vision_start_token_id=self.vision_tokens.vision_start_token_id, + image_token_id=self.vision_tokens.image_token_id, + video_token_id=self.vision_tokens.video_token_id, + spatial_merge_size=self.vision_grid.spatial_merge_size, + tokens_per_second=self.vision_grid.tokens_per_second, + ) + + +class Qwen25VLEmbedding(nn.Embedding): + """Embedding module shared with LM head.""" + + def lm_head_forward(self, x: Tensor): + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class Qwen25VLAttention(nn.Module): + def __init__(self, config: Qwen25VLConfig): + if config.num_key_value_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_key_value_heads} key-value heads " + f"evenly to {config.tensor_parallel_shards} shards." + ) + head_dim = config.head_dim + num_attention_heads = config.num_attention_heads // config.tensor_parallel_shards + num_key_value_heads = config.num_key_value_heads // config.tensor_parallel_shards + mrope_section: Tuple[int, int, int] = ( # noqa: UP006 + config.mrope_section if config.mrope_section is not None else (0, 0, 0) + ) + self.state = Qwen25VLAttentionState( + head_dim=head_dim, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + mrope_section=mrope_section, + softmax_scale=head_dim**-0.5, + ) + + out_features = (num_attention_heads + 2 * num_key_value_heads) * head_dim + self.c_attn = nn.Linear( + in_features=config.hidden_size, + out_features=out_features, + bias=True, + ) + self.o_proj = nn.Linear( + num_attention_heads * head_dim, + config.hidden_size, + bias=False, + ) + + @property + def head_dim(self) -> int: + return self.state.head_dim + + @property + def num_attention_heads(self) -> int: + return self.state.num_attention_heads + + @property + def num_key_value_heads(self) -> int: + return self.state.num_key_value_heads + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + position_embeddings: Tuple[Tensor, Tensor], # noqa: UP006 + ): + d, h_q, h_kv = ( + self.state.head_dim, + self.state.num_attention_heads, + self.state.num_key_value_heads, + ) + b, s, _ = hidden_states.shape + qkv = self.c_attn(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + q, k, v = op.split(qkv, [h_q, h_q + h_kv], axis=2) + cos, sin = position_embeddings + q, k = apply_multimodal_rotary_pos_emb(q, k, cos, sin, self.state.mrope_section) + output, _ = paged_kv_cache.self_attention(layer_id, q, k, v, self.state.softmax_scale) + output = op.reshape(output, (b, s, h_q * d)) + return self.o_proj(output) + + +class Qwen25VLMLP(nn.Module): + def __init__(self, config: Qwen25VLConfig): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} shards." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class Qwen25VLDecoderLayer(nn.Module): + def __init__(self, config: Qwen25VLConfig): + self.self_attn = Qwen25VLAttention(config) + self.mlp = Qwen25VLMLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + self._set_tp(config) + + def _set_tp(self, config: Qwen25VLConfig): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_attention_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set( + self.self_attn.c_attn.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj.weight, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + layer_id: int, + position_embeddings: Tuple[Tensor, Tensor], # noqa: UP006 + ): + out = self.input_layernorm(hidden_states) + out = self.self_attn(out, paged_kv_cache, layer_id, position_embeddings) + hidden_states = self._apply_residual(out, hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, hidden_states) + return hidden_states + + def _apply_residual(self, out: Tensor, residual: Tensor) -> Tensor: + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Qwen25VLModel(nn.Module): + def __init__(self, config: Qwen25VLConfig): + self.embed_tokens = Qwen25VLEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Qwen25VLDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + attention_scaling = config.rope_parameters.get("attention_scaling", 1.0) + self.rotary_emb = MultimodalRotaryEmbedding( + head_dim=config.head_dim, + theta=config.rope_theta, + mrope_section=config.mrope_section, + attention_scaling=attention_scaling, + ) + + def forward( + self, + inputs: Tensor, + position_ids: Tensor, + paged_kv_cache: PagedKVCache, + ): + hidden_states = inputs + cos, sin = self.rotary_emb(hidden_states, position_ids) + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id, (cos, sin)) + return self.norm(hidden_states) + + +class Qwen25VLLMHeadModel(nn.Module): + def __init__(self, config: Qwen25VLConfig): + self.config = config + self.model = Qwen25VLModel(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not self.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.dtype = config.dtype + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def _apply_lm_head(self, hidden_states: Tensor): + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def _set_mrope_delta(self, paged_kv_cache: PagedKVCache, deltas: Tensor): + setattr(paged_kv_cache, "_mrope_delta", deltas) + return deltas + + def _get_mrope_delta(self, paged_kv_cache: PagedKVCache, batch: int) -> Tensor: + delta = getattr(paged_kv_cache, "_mrope_delta", None) + if delta is None: + delta = op.zeros((batch, 1), "int32") + setattr(paged_kv_cache, "_mrope_delta", delta) + return delta + + def _build_decode_position_ids( + self, + seq_len: int, + paged_kv_cache: PagedKVCache, + batch: int, + ) -> Tensor: + base = paged_kv_cache.get_query_positions(seq_len) + base = op.reshape(base, (1, seq_len)) + base = op.broadcast_to(base, (batch, seq_len)) + delta = self._get_mrope_delta(paged_kv_cache, batch) + base = base + delta + base = op.unsqueeze(base, dim=0) + return op.broadcast_to(base, (3, batch, seq_len)) + + def prefill( + self, + input_embed: Tensor, + position_ids: Tensor, + mrope_deltas: Tensor, + paged_kv_cache: PagedKVCache, + ): + op_ext.configure() + self._set_mrope_delta(paged_kv_cache, mrope_deltas) + hidden_states = self.model(input_embed, position_ids, paged_kv_cache) + + hidden_states = index_last_token(hidden_states) + logits = self._apply_lm_head(hidden_states) + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + b, s, _ = input_embed.shape + position_ids = self._build_decode_position_ids(s, paged_kv_cache, b) + hidden_states = self.model(input_embed, position_ids, paged_kv_cache) + logits = self._apply_lm_head(hidden_states) + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + position_ids: Tensor, + mrope_deltas: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.config.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward( + input_embeds, position_ids, mrope_deltas, logit_positions, paged_kv_cache + ) + return logits, paged_kv_cache + + def batch_forward( + self, + input_embeds: Tensor, + position_ids: Tensor, + mrope_deltas: Tensor, + logit_positions: Optional[Tensor], + paged_kv_cache: PagedKVCache, + ): + op_ext.configure() + self._set_mrope_delta(paged_kv_cache, mrope_deltas) + hidden_states = self.model(input_embeds, position_ids, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + return self._apply_lm_head(hidden_states) + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + b, s, _ = input_embeds.shape + position_ids = self._build_decode_position_ids(s, paged_kv_cache, b) + hidden_states = self.model(input_embeds, position_ids, paged_kv_cache) + logits = self._apply_lm_head(hidden_states) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + return self.batch_decode(input_embeds, paged_kv_cache) + + def embed(self, input_ids: Tensor): + if self.config.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + cfg = self.config + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=cfg.num_hidden_layers, + num_attention_heads=cfg.num_attention_heads // cfg.tensor_parallel_shards, + num_key_value_heads=cfg.num_key_value_heads // cfg.tensor_parallel_shards, + qk_head_dim=cfg.head_dim, + v_head_dim=cfg.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scaling=cfg.rope_parameters, + rope_scale=1, + rope_theta=int(cfg.rope_theta), + dtype=self.dtype, + ) + + def get_default_spec(self): + cfg = self.config + seq_len = "seq_len" + hidden = cfg.hidden_size + dtype = self.dtype + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor([seq_len], "int32"), + "$": {"param_mode": "packed", "effect_mode": "none"}, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, seq_len, hidden], dtype), + "position_ids": nn.spec.Tensor([3, 1, seq_len], "int32"), + "mrope_deltas": nn.spec.Tensor([1, 1], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": {"param_mode": "packed", "effect_mode": "none"}, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, hidden], dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": {"param_mode": "packed", "effect_mode": "none"}, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, seq_len, hidden], dtype), + "position_ids": nn.spec.Tensor([3, 1, seq_len], "int32"), + "mrope_deltas": nn.spec.Tensor([1, 1], "int32"), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": {"param_mode": "packed", "effect_mode": "none"}, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, hidden], dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": {"param_mode": "packed", "effect_mode": "none"}, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, hidden], dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": {"param_mode": "packed", "effect_mode": "none"}, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": {"param_mode": "none", "effect_mode": "none"}, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/qwen2_moe/__init__.py b/python/mlc_llm/model/qwen2_moe/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/qwen2_moe/qwen2_moe_loader.py b/python/mlc_llm/model/qwen2_moe/qwen2_moe_loader.py new file mode 100644 index 0000000..1991cb0 --- /dev/null +++ b/python/mlc_llm/model/qwen2_moe/qwen2_moe_loader.py @@ -0,0 +1,130 @@ +""" +This file specifies how MLC's QWen2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .qwen2_moe_model import Qwen2MoeConfig, Qwen2MoeForCausalLM + + +def huggingface(model_config: Qwen2MoeConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : Qwen2MoeConfig + The configuration of the Qwen2Moe model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Qwen2MoeForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # map attention weight + attn = f"model.layers.{i}.self_attn" + for weight_type in ["weight", "bias"]: + mlc_name = f"{attn}.c_attn.{weight_type}" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.{weight_type}", + f"{attn}.k_proj.{weight_type}", + f"{attn}.v_proj.{weight_type}", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + # map mlp shared expert weight + mlp = f"model.layers.{i}.mlp" + shared_expert = f"{mlp}.shared_expert" + mlc_name = f"{shared_expert}.gate_up_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{shared_expert}.gate_proj.weight", + f"{shared_expert}.up_proj.weight", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + # map mlp moe gate and up weight + mlc_name = f"{mlp}.moe_gate_up_proj.weight" + + def combine_expert_gate_up(*hf_params, dtype): + stack = [] + for i in range(0, len(hf_params), 2): + stack.append(np.concatenate([hf_params[i], hf_params[i + 1]], axis=0)) + return np.stack(stack, axis=0).astype(dtype) + + mapping.add_mapping( + mlc_name, + functools.reduce( + lambda a, b: a + b, + [ + [ + f"{mlp}.experts.{expert_id}.gate_proj.weight", + f"{mlp}.experts.{expert_id}.up_proj.weight", + ] + for expert_id in range(model_config.num_experts) + ], + ), + functools.partial( + combine_expert_gate_up, + dtype=mlc_param.dtype, + ), + ) + + # map mlp moe gate and up weight + mlc_name = f"{mlp}.moe_down_proj.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{mlp}.experts.{expert_id}.down_proj.weight" + for expert_id in range(model_config.num_experts) + ], + functools.partial( + lambda *hf_params, dtype: np.stack(hf_params, axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/qwen2_moe/qwen2_moe_model.py b/python/mlc_llm/model/qwen2_moe/qwen2_moe_model.py new file mode 100644 index 0000000..27bbf30 --- /dev/null +++ b/python/mlc_llm/model/qwen2_moe/qwen2_moe_model.py @@ -0,0 +1,390 @@ +""" +Implementation for QWEN2MOE architecture. +""" + +import dataclasses +from typing import Optional + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.model.qwen2.qwen2_model import ACT2FN, QWen2Attention, QWen2Config +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.nn.expert import MixtralExperts +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Qwen2MoeConfig(QWen2Config): + """Configuration of the Qwen2Moe model.""" + + moe_intermediate_size: int = 0 + shared_expert_intermediate_size: int = 0 + num_experts_per_tok: int = 0 + num_experts: int = 0 + decoder_sparse_step: int = 0 + norm_topk_prob: bool = False + + +class Qwen2MoeMLP(nn.Module): + def __init__(self, config: Qwen2MoeConfig, intermediate_size: Optional[int] = None): + intermediate_size = intermediate_size or config.intermediate_size + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MoE MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class Qwen2MoeSparseMoeBlock(nn.Module): + """MoE layer for Qwen2MoE model.""" + + def __init__(self, config: Qwen2MoeConfig): + super().__init__() + self.num_experts_per_tok = config.num_experts_per_tok + self.num_experts = config.num_experts + if config.moe_intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MoE intermediate size {config.moe_intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.moe_intermediate_size = config.moe_intermediate_size // config.tensor_parallel_shards + self.norm_topk_prob = config.norm_topk_prob + self.shared_expert = Qwen2MoeMLP(config, config.shared_expert_intermediate_size) + self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False) + + self.gate = nn.Linear( + in_features=config.hidden_size, + out_features=config.num_experts, + bias=False, + ) + self.moe_gate_up_proj = MixtralExperts( + self.num_experts, + in_features=config.hidden_size, + out_features=2 * self.moe_intermediate_size, + ) + self.moe_down_proj = MixtralExperts( + self.num_experts, + in_features=self.moe_intermediate_size, + out_features=config.hidden_size, + ) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + def _expert_forward(x: Tensor, indptr: Tensor): + x1_x2 = self.moe_gate_up_proj(x, indptr) + x1, x2 = op.split(x1_x2, indices_or_sections=2, axis=-1) + x = self.moe_down_proj(self.act_fn(x1) * x2, indptr) + return x + + experts_per_tok = self.num_experts_per_tok + num_experts = self.num_experts + batch_size, seq_len, hidden_size = x.shape + num_tokens = batch_size * seq_len + x = x.reshape(num_tokens, hidden_size) + gate = self.gate(x) + # expert_weights: [num_tokens, experts_per_tok] + # expert_indices: [num_tokens, experts_per_tok] + expert_weights, expert_indices = op_ext.moe_misc.gating_softmax_topk( + gate, experts_per_tok, norm_topk_prob=self.norm_topk_prob + ) + if num_tokens == 1: + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = _expert_forward(x, expert_indices) + else: + # cumsum: [num_tokens * local_experts] + cumsum = op_ext.moe_misc.moe_cumsum(expert_indices, num_experts) + # indices: [num_tokens * experts_per_tok] + reverse_indices, token_indices = op_ext.moe_misc.get_indices(cumsum, expert_indices) + # indptr: [num_local_experts + 1] + indptr = op_ext.moe_misc.get_indptr( + cumsum, num_experts, num_tokens, inclusive=False, out_dtype="int32" + ) + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = op.take(x, token_indices, axis=0) + moe_hidden_states = _expert_forward(moe_hidden_states, indptr) + moe_hidden_states = op_ext.moe_misc.scatter_output(moe_hidden_states, reverse_indices) + # moe_hidden_states: [num_tokens, experts_per_tok, hidden_size] + expert_weights = expert_weights.reshape(num_tokens, experts_per_tok, 1) + moe_hidden_states = ( + moe_hidden_states.reshape(num_tokens, experts_per_tok, hidden_size) * expert_weights + ) + # moe_hidden_states: [num_tokens, hidden_size] + moe_hidden_states = op_ext.moe_misc.moe_sum(moe_hidden_states, dim=1) + + shared_expert_hidden_states = self.shared_expert(x) + shared_expert_hidden_states = ( + op.sigmoid(self.shared_expert_gate(x)) * shared_expert_hidden_states + ) + final_hidden_states = moe_hidden_states + shared_expert_hidden_states + final_hidden_states = final_hidden_states.reshape(batch_size, seq_len, hidden_size) + return final_hidden_states + + +class Qwen2MoeDecoderLayer(nn.Module): + def __init__(self, config: Qwen2MoeConfig): + super().__init__() + self.self_attn = QWen2Attention(config) + assert config.num_experts > 0 and config.decoder_sparse_step == 1, ( + "Currently only support use moe for every layer." + ) + self.mlp = Qwen2MoeSparseMoeBlock(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_attention_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + si = self.mlp.shared_expert.intermediate_size + mi = self.mlp.moe_intermediate_size + _set( + self.self_attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + _set( + self.self_attn.c_attn.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.shared_expert.gate_up_proj.weight, + tp.ShardSingleDim("_shard_shared_mlp_up", segs=[si, si], dim=0), + ) + _set( + self.mlp.shared_expert.down_proj.weight, + tp.ShardSingleDim("_shard_shared_mlp_down", dim=1), + ) + _set( + self.mlp.moe_gate_up_proj.weight, + tp.ShardSingleDim("_shard_moe_mlp_up", segs=[mi, mi], dim=1), + ) + _set( + self.mlp.moe_down_proj.weight, + tp.ShardSingleDim("_shard_moe_mlp_down", dim=2), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.input_layernorm(hidden_states) + out = self.self_attn(out, paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Qwen2MoeModel(nn.Module): + def __init__(self, config: Qwen2MoeConfig): + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Qwen2MoeDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Qwen2MoeForCausalLM(nn.Module): + def __init__(self, config: Qwen2MoeConfig): + self.model = Qwen2MoeModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.dtype = config.dtype + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.intermediate_size = config.intermediate_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.rms_norm_eps = config.rms_norm_eps + self.rope_theta = config.rope_theta + self.vocab_size = config.vocab_size + self.tensor_parallel_shards = config.tensor_parallel_shards + self.head_dim = config.head_dim + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/qwen3/__init__.py b/python/mlc_llm/model/qwen3/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/qwen3/qwen3_loader.py b/python/mlc_llm/model/qwen3/qwen3_loader.py new file mode 100644 index 0000000..cf78ce4 --- /dev/null +++ b/python/mlc_llm/model/qwen3/qwen3_loader.py @@ -0,0 +1,152 @@ +""" +This file specifies how MLC's QWen2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools +from typing import Callable, List, Literal # noqa: UP035 + +import numpy as np + +from mlc_llm.loader import ExternMapping, QuantizeMapping +from mlc_llm.quantization import BlockScaleQuantize, Quantization + +from .qwen3_model import Qwen3Config, Qwen3LMHeadModel + + +def huggingface( + model_config: Qwen3Config, + quantization: Quantization, + hf_prefix: Literal["", "model."] = "model.", +) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : Qwen3Config + The configuration of the Qwen3 model. + + quantization : Quantization + The quantization configuration. + + hf_prefix : Literal["", "model."] + Prefix used in HuggingFace weight names. Defaults to "model." for standard + Qwen3 models. Use "" for Qwen3-Embedding models without prefix. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Qwen3LMHeadModel(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + if isinstance(quantization, BlockScaleQuantize): + model = quantization.quantize_model(model, QuantizeMapping({}, {}), "") + if model_config.weight_block_size is None: + raise ValueError( + "The input Qwen3 model is not fp8 block quantized. " + "Thus BlockScaleQuantize is not supported." + ) + + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + if ( + not isinstance(quantization, BlockScaleQuantize) + and model_config.weight_block_size is not None + ): + raise ValueError( + "The input Qwen3 model is fp8 block quantized. " + "Please use BlockScaleQuantize for the model." + ) + + def to_hf(name: str) -> str: + if hf_prefix == "model.": + return name + return name[6:] if name.startswith("model.") else name + + def add_weight_and_scale_mapping( + weight_mlc_name: str, + weight_hf_names: List[str], # noqa: UP006 + weight_transform_func: Callable, + ): + mlc_param = named_parameters[weight_mlc_name] + hf_names = [to_hf(name) for name in weight_hf_names] + mapping.add_mapping( + weight_mlc_name, + hf_names, + functools.partial(weight_transform_func, dtype=mlc_param.dtype), + ) + + if isinstance(quantization, BlockScaleQuantize): + scale_mlc_name = f"{weight_mlc_name}_scale_inv" + if scale_mlc_name in named_parameters: + scale_hf_names = [f"{name}_scale_inv" for name in hf_names] + scale_param = named_parameters[scale_mlc_name] + mapping.add_mapping( + scale_mlc_name, + scale_hf_names, + functools.partial(weight_transform_func, dtype=scale_param.dtype), + ) + + for i in range(model_config.num_hidden_layers): + # map attention weight + attn = f"model.layers.{i}.self_attn" + add_weight_and_scale_mapping( + f"{attn}.c_attn.weight", + [ + f"{attn}.q_proj.weight", + f"{attn}.k_proj.weight", + f"{attn}.v_proj.weight", + ], + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + ) + if model_config.attention_bias: + mlc_name = f"{attn}.c_attn.bias" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + to_hf(f"{attn}.q_proj.bias"), + to_hf(f"{attn}.k_proj.bias"), + to_hf(f"{attn}.v_proj.bias"), + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + # map mlp weight + mlp = f"model.layers.{i}.mlp" + add_weight_and_scale_mapping( + f"{mlp}.gate_up_proj.weight", + [ + f"{mlp}.gate_proj.weight", + f"{mlp}.up_proj.weight", + ], + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [to_hf(mlc_name)], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping + + +def huggingface_embedding(model_config: Qwen3Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping for Qwen3-Embedding models (no 'model.' prefix).""" + return huggingface(model_config, quantization, "") diff --git a/python/mlc_llm/model/qwen3/qwen3_model.py b/python/mlc_llm/model/qwen3/qwen3_model.py new file mode 100644 index 0000000..f3bdfec --- /dev/null +++ b/python/mlc_llm/model/qwen3/qwen3_model.py @@ -0,0 +1,528 @@ +""" +Implementation for QWEN2 architecture. +""" + +import dataclasses +from functools import partial +from typing import Any, Dict, Optional, Tuple # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Qwen3Config(ConfigBase): + """Configuration of the Qwen3 model.""" + + hidden_act: str + hidden_size: int + intermediate_size: int + attention_bias: bool + num_attention_heads: int + num_hidden_layers: int + num_key_value_heads: int + rms_norm_eps: float + rope_theta: int + vocab_size: int + tie_word_embeddings: bool = False + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + head_dim: int = 0 + dtype: str = "float32" + max_batch_size: int = 1 + weight_block_size: Optional[Tuple[int, int]] = None # noqa: UP006 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if "quantization_config" in self.kwargs: + quantization_config = self.kwargs.get("quantization_config") + if ( + isinstance(quantization_config, dict) + and quantization_config.get("activation_scheme", "") == "dynamic" + and quantization_config.get("fmt", "") == "e4m3" + and quantization_config.get("quant_method", "") == "fp8" + and "weight_block_size" in quantization_config + ): + self.weight_block_size = quantization_config.get("weight_block_size") + if ( + not isinstance(self.weight_block_size, (tuple, list)) + or len(self.weight_block_size) != 2 + ): + raise ValueError( + "Invalid DeepSeek model quantization config: " + "weight_block_size must be a tuple of two integers, " + f"got {self.weight_block_size} of type {type(self.weight_block_size)}" + ) + else: + raise ValueError( + "Invalid DeepSeek model quantization config: unrecognized quantization config: " + f"{quantization_config}" + ) + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 2048), + ) + self.prefill_chunk_size = min(self.context_window_size, 2048) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 2048), + ) + self.prefill_chunk_size = min(self.context_window_size, 2048) + + +class Qwen3Attention(nn.Module): + def __init__(self, config: Qwen3Config): + self.head_dim = config.head_dim + if config.num_key_value_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_key_value_heads} key-value attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_attention_heads = config.num_attention_heads // config.tensor_parallel_shards + self.num_key_value_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.rope_theta = config.rope_theta + + self.c_attn = nn.Linear( + in_features=config.hidden_size, + out_features=(2 * self.num_key_value_heads + self.num_attention_heads) * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + self.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = nn.RMSNorm(config.head_dim, -1, config.rms_norm_eps, bias=False) + self.k_norm = nn.RMSNorm(config.head_dim, -1, config.rms_norm_eps, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_attention_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.c_attn(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + q, k, v = op.split(qkv, [h_q, h_q + h_kv], axis=2) + q = self.q_norm(q) + k = self.k_norm(k) + qkv = op.concat([q, k, v], dim=2) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_attention_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, + "swish": nn.silu, + "gelu_new": partial(nn.gelu, approximate=True), +} + + +class Qwen3Embedding(nn.Embedding): + """The embedding module specialized for Qwen3 so that + it can be shared with the final lm_head. + """ + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which transposes the weight and multiplies + with the input tensor. + """ + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class Qwen3MLP(nn.Module): + def __init__(self, config: Qwen3Config): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class Qwen3DecoderLayer(nn.Module): + def __init__(self, config: Qwen3Config): + self.self_attn = Qwen3Attention(config) + self.mlp = Qwen3MLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_attention_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + if config.attention_bias: + _set( + self.self_attn.c_attn.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj.weight, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.input_layernorm(hidden_states) + out = self.self_attn(out, paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Qwen3Model(nn.Module): + def __init__(self, config: Qwen3Config): + self.embed_tokens = Qwen3Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Qwen3DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Qwen3LMHeadModel(nn.Module): + def __init__(self, config: Qwen3Config): + self.model = Qwen3Model(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.dtype = config.dtype + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.intermediate_size = config.intermediate_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.rms_norm_eps = config.rms_norm_eps + self.rope_theta = config.rope_theta + self.vocab_size = config.vocab_size + self.tensor_parallel_shards = config.tensor_parallel_shards + self.head_dim = config.head_dim + self.weight_block_size = config.weight_block_size + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) + + +class Qwen3EmbeddingModel(Qwen3LMHeadModel): + """Qwen3 model for embedding inference. + + Inherits all functionality from Qwen3LMHeadModel and adds methods that + return hidden states instead of logits, for use by AsyncEmbeddingEngine. + Only compiled when using the "qwen3-embedding" model type. + """ + + def prefill_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def decode_to_last_hidden_states(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + hidden_states = self.model(input_embed, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_prefill_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + op_ext.configure() + hidden_states = self.model(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def batch_decode_to_last_hidden_states( + self, input_embeds: Tensor, paged_kv_cache: PagedKVCache + ): + op_ext.configure() + hidden_states = self.model(input_embeds, paged_kv_cache) + return hidden_states, paged_kv_cache + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode_to_last_hidden_states": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode_to_last_hidden_states": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/qwen35/__init__.py b/python/mlc_llm/model/qwen35/__init__.py new file mode 100644 index 0000000..1e50eb3 --- /dev/null +++ b/python/mlc_llm/model/qwen35/__init__.py @@ -0,0 +1 @@ +"""Qwen3.5 GatedDeltaNet hybrid model.""" diff --git a/python/mlc_llm/model/qwen35/qwen35_loader.py b/python/mlc_llm/model/qwen35/qwen35_loader.py new file mode 100644 index 0000000..8abcb8f --- /dev/null +++ b/python/mlc_llm/model/qwen35/qwen35_loader.py @@ -0,0 +1,177 @@ +""" +HuggingFace parameter mapping for Qwen3.5 GatedDeltaNet. + +Qwen3.5 is a VLM — HF weights are nested under `model.language_model.`. +Our MLC model uses `model.` prefix. The mapping must translate between them. + +HF weight layout (under model.language_model.): + Linear attention layers: + model.language_model.layers.{i}.linear_attn.in_proj_qkv.weight + model.language_model.layers.{i}.linear_attn.in_proj_z.weight + model.language_model.layers.{i}.linear_attn.in_proj_a.weight + model.language_model.layers.{i}.linear_attn.in_proj_b.weight + model.language_model.layers.{i}.linear_attn.out_proj.weight + model.language_model.layers.{i}.linear_attn.conv1d.weight + model.language_model.layers.{i}.linear_attn.norm.weight + model.language_model.layers.{i}.linear_attn.A_log (NO .weight suffix) + model.language_model.layers.{i}.linear_attn.dt_bias (NO .weight suffix) + + Full attention layers: + model.language_model.layers.{i}.self_attn.q_proj.weight + ... + + Vision/MTP weights are ignored (text backbone only). +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .qwen35_model import Qwen35Config, Qwen35LMHeadModel + + +def huggingface(model_config: Qwen35Config, quantization: Quantization) -> ExternMapping: + model = Qwen35LMHeadModel(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + # HF prefix: Qwen3.5 is a VLM, text weights nested under model.language_model. + # MLC model uses model. prefix directly. + hf = "model.language_model" + + layer_types = model_config.layer_types() + for i in range(model_config.num_hidden_layers): + if layer_types[i] == "full_attention": + # Standard attention: fuse Q/K/V into c_attn + mlc_attn = f"model.layers.{i}.self_attn" + hf_attn = f"{hf}.layers.{i}.self_attn" + mlc_name = f"{mlc_attn}.c_attn.weight" + if mlc_name in named_parameters: + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{hf_attn}.q_proj.weight", + f"{hf_attn}.k_proj.weight", + f"{hf_attn}.v_proj.weight", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + else: + # Linear attention layer + mlc_lin = f"model.layers.{i}.linear_attn" + hf_lin = f"{hf}.layers.{i}.linear_attn" + + # in_proj_qkv — maps directly (already fused in HF) + mlc_name = f"{mlc_lin}.in_proj_qkv.weight" + if mlc_name in named_parameters: + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{hf_lin}.in_proj_qkv.weight"], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + + # A_log and dt_bias — no .weight suffix in HF + for param_name in ["A_log", "dt_bias"]: + mlc_name = f"{mlc_lin}.{param_name}" + if mlc_name in named_parameters: + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{hf_lin}.{param_name}"], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + + # conv1d weight + mlc_name = f"{mlc_lin}.conv1d_weight" + if mlc_name in named_parameters: + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [f"{hf_lin}.conv1d.weight"], + functools.partial(lambda x, dtype: x.astype(dtype), dtype=mlc_param.dtype), + ) + + # MLP: fuse gate_proj + up_proj + mlc_mlp = f"model.layers.{i}.mlp" + hf_mlp = f"{hf}.layers.{i}.mlp" + mlc_name = f"{mlc_mlp}.gate_up_proj.weight" + if mlc_name in named_parameters: + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{hf_mlp}.gate_proj.weight", + f"{hf_mlp}.up_proj.weight", + ], + functools.partial( + lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + def _mlc_to_hf(mlc_name: str) -> str: + """Convert MLC param name to HF param name by adding language_model prefix.""" + if mlc_name.startswith("model."): + return mlc_name.replace("model.", f"{hf}.", 1) + return mlc_name + + def _is_rmsnorm_weight(name: str) -> bool: + """Check if a parameter is an RMSNorm weight that needs +1.0 offset. + + Qwen3_5RMSNorm uses: output = norm(x) * (1.0 + weight) + - input_layernorm, post_attention_layernorm, model.norm, q_norm, k_norm + Qwen3_5RMSNormGated uses: output = norm(x) * weight * silu(gate) + - linear_attn.norm (gated norm) — does NOT get +1 + """ + return ( + name.endswith("input_layernorm.weight") + or name.endswith("post_attention_layernorm.weight") + or name.endswith("q_norm.weight") + or name.endswith("k_norm.weight") + or name == "model.norm.weight" + ) + + # All remaining parameters: direct 1:1 mapping with HF prefix + # Qwen3.5 uses a non-standard RMSNorm: output = norm(x) * (1.0 + weight) + # Weights are initialized to zeros and learned as offsets from 1.0. + # TVM's nn.RMSNorm uses: output = norm(x) * weight + # So we add 1.0 to all RMSNorm weights during loading. + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + hf_name = _mlc_to_hf(mlc_name) + if _is_rmsnorm_weight(mlc_name): + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: (x.astype("float32") + 1.0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + else: + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + return mapping diff --git a/python/mlc_llm/model/qwen35/qwen35_model.py b/python/mlc_llm/model/qwen35/qwen35_model.py new file mode 100644 index 0000000..5d2686f --- /dev/null +++ b/python/mlc_llm/model/qwen35/qwen35_model.py @@ -0,0 +1,868 @@ +""" +Implementation for Qwen3.5 GatedDeltaNet hybrid architecture. +75% GatedDeltaNet (recurrent linear attention), 25% standard GQA softmax attention. +""" + +import dataclasses +import math +from functools import partial +from typing import Any, Dict, List, Optional, Tuple # noqa: UP035 + +import numpy as np +from tvm import relax as R +from tvm import te, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op +from tvm.script import tirx as T + +from mlc_llm import op as op_ext +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.nn.rnn_state import RNNState +from mlc_llm.support import logging +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Qwen35Config(ConfigBase): + """Configuration of the Qwen3.5 model.""" + + hidden_size: int = 0 + intermediate_size: int = 0 + num_attention_heads: int = 0 + num_hidden_layers: int = 0 + num_key_value_heads: int = 0 + rms_norm_eps: float = 1e-6 + vocab_size: int = 0 + rope_theta: int = 10000000 + head_dim: int = 256 + hidden_act: str = "silu" + attention_bias: bool = False + tie_word_embeddings: bool = False + # GatedDeltaNet-specific + linear_key_head_dim: int = 128 + linear_value_head_dim: int = 128 + linear_num_key_heads: int = 16 + linear_num_value_heads: int = 16 + linear_conv_kernel_dim: int = 4 + full_attention_interval: int = 4 + partial_rotary_factor: float = 0.25 + # Runtime + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + dtype: str = "float32" + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + # Handle VLM wrapper: Qwen3.5 HF config has all text params inside text_config + if "text_config" in self.kwargs: + text_config = self.kwargs.pop("text_config") + if isinstance(text_config, dict): + field_names = {f.name for f in dataclasses.fields(self.__class__)} + for k, v in text_config.items(): + if k in field_names and k != "kwargs": + setattr(self, k, v) + else: + self.kwargs[k] = v + # Extract rope params from nested rope_parameters + rope_params = text_config.get("rope_parameters", {}) + if isinstance(rope_params, dict): + if "rope_theta" in rope_params: + self.rope_theta = rope_params["rope_theta"] + if "partial_rotary_factor" in rope_params: + self.partial_rotary_factor = rope_params["partial_rotary_factor"] + + # Also handle rope_parameters at top level + if "rope_parameters" in self.kwargs: + rope_params = self.kwargs.pop("rope_parameters") + if isinstance(rope_params, dict): + if "rope_theta" in rope_params: + self.rope_theta = rope_params["rope_theta"] + if "partial_rotary_factor" in rope_params: + self.partial_rotary_factor = rope_params["partial_rotary_factor"] + + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.prefill_chunk_size == 0: + self.prefill_chunk_size = min(self.context_window_size, 2048) + elif self.prefill_chunk_size > self.context_window_size: + self.prefill_chunk_size = min(self.context_window_size, 2048) + + @property + def num_linear_layers(self) -> int: + """Number of GatedDeltaNet linear attention layers.""" + return self.num_hidden_layers - self.num_attention_layers + + @property + def num_attention_layers(self) -> int: + """Number of full attention layers.""" + return self.num_hidden_layers // self.full_attention_interval + + def layer_types(self) -> List[str]: # noqa: UP006 + """Returns list of layer types: 'linear_attention' or 'full_attention'.""" + types = [] + for i in range(self.num_hidden_layers): + if (i + 1) % self.full_attention_interval == 0: + types.append("full_attention") + else: + types.append("linear_attention") + return types + + +ACT2FN = { + "gelu": partial(nn.gelu, approximate=False), + "relu": nn.relu, + "silu": nn.silu, +} + + +class Qwen35Embedding(nn.Embedding): + def lm_head_forward(self, x: nn.Tensor): + weight = nn.op.permute_dims(self.weight) + return nn.op.matmul(x, weight, out_dtype="float32") + + +class Qwen35MLP(nn.Module): + def __init__(self, config: Qwen35Config): + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class Qwen35Attention(nn.Module): + """Standard GQA attention with output gate for full_attention layers (every 4th layer). + + attn_output_gate=True: q_proj outputs 2*num_heads*head_dim, split into (Q, gate). + Gate is sigmoid-applied to attention output before o_proj. + """ + + def __init__(self, config: Qwen35Config): + self.head_dim = config.head_dim + self.num_attention_heads = config.num_attention_heads // config.tensor_parallel_shards + self.num_key_value_heads = config.num_key_value_heads // config.tensor_parallel_shards + self.rope_theta = config.rope_theta + + # c_attn: Q (2x for gate) + K + V fused projection + self.c_attn = nn.Linear( + in_features=config.hidden_size, + out_features=(2 * self.num_attention_heads + 2 * self.num_key_value_heads) + * self.head_dim, + bias=config.attention_bias, + ) + self.o_proj = nn.Linear( + self.num_attention_heads * self.head_dim, + config.hidden_size, + bias=config.attention_bias, + ) + self.q_norm = nn.RMSNorm(config.head_dim, -1, config.rms_norm_eps, bias=False) + self.k_norm = nn.RMSNorm(config.head_dim, -1, config.rms_norm_eps, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_attention_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + # c_attn outputs flat: [Q_with_gate (h_q * 2 * d), K (h_kv * d), V (h_kv * d)] + proj = self.c_attn(hidden_states) + # Reshape to heads: (b, s, 2*h_q + 2*h_kv, d) + proj = op.reshape(proj, (b, s, 2 * h_q + 2 * h_kv, d)) + # Split: first 2*h_q heads have interleaved [Q, gate] per head, then h_kv K, h_kv V + q_gate, k, v = op.split(proj, [2 * h_q, 2 * h_q + h_kv], axis=2) + # q_gate shape: (b, s, 2*h_q, d). Even heads are Q, odd heads are gate + # But HF layout is per-head [Q_d, gate_d], so reshape to (b, s, h_q, 2*d) then split + q_gate = op.reshape(q_gate, (b, s, h_q, 2 * d)) + q, gate = op.split(q_gate, [d], axis=3) + # gate: (b, s, h_q, d) -> flatten to (b, s, h_q*d) + gate = op.reshape(gate, (b, s, h_q * d)) + q = self.q_norm(q) + k = self.k_norm(k) + qkv = op.concat([q, k, v], dim=2) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_attention_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + # Apply output gate: sigmoid(gate) * attn_output + output = output * op.sigmoid(gate) + return self.o_proj(output) + + +# ============================================================================ +# GatedDeltaNet TIR kernel +# ============================================================================ + + +def create_gated_delta_net_func( + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, + dtype: str, +): + """Creates a TIR function for the GatedDeltaNet recurrent computation. + + Thread-per-column design: each thread owns one column of the state matrix. + State S is (key_head_dim x value_head_dim) per head, accumulated in fp32. + + Supports arbitrary sequence length via an inner `for t in range(seq_len)` loop, + matching RWKV6's approach. During prefill (seq_len > 1), the recurrence accumulates + state across all tokens sequentially. During decode (seq_len = 1), it's a single step. + + For GVA (num_value_heads > num_key_heads), Q/K are expanded via repeat. + The kernel operates on value_heads (the larger dimension). + """ + heads_per_group = num_value_heads // num_key_heads # 1 for 0.8B, 2 for 4B + K = key_head_dim # 128 + V = value_head_dim # 128 + + @T.prim_func(s_tir=True) + def gdn_func( + q_handle: T.handle, + k_handle: T.handle, + v_handle: T.handle, + gate_handle: T.handle, # exp(g), already exponentiated + beta_handle: T.handle, # sigmoid(beta_raw) + state_in_handle: T.handle, + out_handle: T.handle, + state_out_handle: T.handle, + ): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + batch_size, seq_len = T.int64(), T.int64() + # q, k: (batch, seq_len, key_heads, K) + q_buf = T.match_buffer(q_handle, (batch_size, seq_len, num_key_heads, K), dtype=dtype) + k_buf = T.match_buffer(k_handle, (batch_size, seq_len, num_key_heads, K), dtype=dtype) + # v: (batch, seq_len, value_heads, V) + v_buf = T.match_buffer(v_handle, (batch_size, seq_len, num_value_heads, V), dtype=dtype) + # gate and beta: (batch, seq_len, value_heads) + gate_buf = T.match_buffer( + gate_handle, (batch_size, seq_len, num_value_heads), dtype="float32" + ) + beta_buf = T.match_buffer( + beta_handle, (batch_size, seq_len, num_value_heads), dtype="float32" + ) + # State: per value_head, K x V matrix in fp32 + state_in_buf = T.match_buffer( + state_in_handle, (batch_size, num_value_heads, K, V), dtype="float32" + ) + # Outputs: out in fp32 for numerical stability (cast to model dtype by caller) + out_buf = T.match_buffer( + out_handle, (batch_size, seq_len, num_value_heads, V), dtype="float32" + ) + state_out_buf = T.match_buffer( + state_out_handle, (batch_size, num_value_heads, K, V), dtype="float32" + ) + + for b_idx in T.thread_binding(batch_size, thread="blockIdx.y"): + for h_idx in T.thread_binding(num_value_heads, thread="blockIdx.x"): + for col in T.thread_binding(V, thread="threadIdx.x"): + kh = h_idx // heads_per_group + + # Init state from state_in + for row in range(K): + with T.sblock("init_state"): + vb, vh, vr, vc = T.axis.remap("SSSS", [b_idx, h_idx, row, col]) + state_out_buf[vb, vh, vr, vc] = state_in_buf[vb, vh, vr, vc] + + # Sequential loop over tokens (like RWKV6) + for t in range(seq_len): + # 1. Decay state: S = gate * S + for row in range(K): + with T.sblock("decay"): + vb = T.axis.spatial(batch_size, b_idx) + vt = T.axis.opaque(seq_len, t) + vh = T.axis.spatial(num_value_heads, h_idx) + vr = T.axis.opaque(K, row) + vc = T.axis.spatial(V, col) + state_out_buf[vb, vh, vr, vc] = ( + state_out_buf[vb, vh, vr, vc] * gate_buf[vb, vt, vh] + ) + + # 2. Compute dot(S[:, col], k[:]) → out_buf (fp32) + with T.sblock("dot_sk_init"): + vb = T.axis.spatial(batch_size, b_idx) + vt = T.axis.opaque(seq_len, t) + vh = T.axis.spatial(num_value_heads, h_idx) + vc = T.axis.spatial(V, col) + out_buf[vb, vt, vh, vc] = T.float32(0) + + for row in range(K): + with T.sblock("dot_sk"): + vb = T.axis.spatial(batch_size, b_idx) + vt = T.axis.opaque(seq_len, t) + vr = T.axis.opaque(K, row) + vh = T.axis.spatial(num_value_heads, h_idx) + vc = T.axis.spatial(V, col) + out_buf[vb, vt, vh, vc] = out_buf[vb, vt, vh, vc] + state_out_buf[ + vb, vh, vr, vc + ] * T.cast(k_buf[vb, vt, kh, vr], "float32") + + # 3. Delta rule: S += k * beta * (v - dot_sk) + for row in range(K): + with T.sblock("delta"): + vb = T.axis.spatial(batch_size, b_idx) + vt = T.axis.opaque(seq_len, t) + vr = T.axis.opaque(K, row) + vh = T.axis.spatial(num_value_heads, h_idx) + vc = T.axis.spatial(V, col) + state_out_buf[vb, vh, vr, vc] = state_out_buf[ + vb, vh, vr, vc + ] + T.cast(k_buf[vb, vt, kh, vr], "float32") * beta_buf[ + vb, vt, vh + ] * ( + T.cast(v_buf[vb, vt, vh, vc], "float32") + - out_buf[vb, vt, vh, vc] + ) + + # 4. Output: o[t, col] = dot(S_updated[:, col], q[t, :]) * scale + with T.sblock("out_init"): + vb = T.axis.spatial(batch_size, b_idx) + vt = T.axis.opaque(seq_len, t) + vh = T.axis.spatial(num_value_heads, h_idx) + vc = T.axis.spatial(V, col) + out_buf[vb, vt, vh, vc] = T.float32(0) + + for row in range(K): + with T.sblock("dot_sq"): + vb = T.axis.spatial(batch_size, b_idx) + vt = T.axis.opaque(seq_len, t) + vr = T.axis.opaque(K, row) + vh = T.axis.spatial(num_value_heads, h_idx) + vc = T.axis.spatial(V, col) + out_buf[vb, vt, vh, vc] = out_buf[vb, vt, vh, vc] + state_out_buf[ + vb, vh, vr, vc + ] * T.cast(q_buf[vb, vt, kh, vr], "float32") + + # 5. Apply scale + with T.sblock("scale"): + vb = T.axis.spatial(batch_size, b_idx) + vt = T.axis.opaque(seq_len, t) + vh = T.axis.spatial(num_value_heads, h_idx) + vc = T.axis.spatial(V, col) + out_buf[vb, vt, vh, vc] = out_buf[vb, vt, vh, vc] * T.float32( + 1.0 / math.sqrt(K) + ) + + return gdn_func + + +# ============================================================================ +# GatedDeltaNet Linear Attention Layer +# ============================================================================ + + +class Qwen35GatedDeltaNet(nn.Module): + """GatedDeltaNet linear attention layer.""" + + def __init__(self, config: Qwen35Config, linear_layer_idx: int): + self.config = config + self.linear_layer_idx = linear_layer_idx # index among linear layers only + self.key_head_dim = config.linear_key_head_dim # 128 + self.value_head_dim = config.linear_value_head_dim # 128 + self.num_key_heads = config.linear_num_key_heads # 16 + self.num_value_heads = config.linear_num_value_heads # 16 or 32 + self.hidden_size = config.hidden_size + self.dtype = config.dtype + + qkv_dim = ( + (self.num_key_heads * self.key_head_dim) + + (self.num_key_heads * self.key_head_dim) + + (self.num_value_heads * self.value_head_dim) + ) + + # Projections — matching HF weight names + self.in_proj_qkv = nn.Linear(config.hidden_size, qkv_dim, bias=False) + self.in_proj_z = nn.Linear( + config.hidden_size, self.num_value_heads * self.value_head_dim, bias=False + ) + self.in_proj_a = nn.Linear(config.hidden_size, self.num_value_heads, bias=False) + self.in_proj_b = nn.Linear(config.hidden_size, self.num_value_heads, bias=False) + self.out_proj = nn.Linear( + self.num_value_heads * self.value_head_dim, config.hidden_size, bias=False + ) + + # Causal depthwise Conv1D kernel + self.conv1d_weight = nn.Parameter( + (qkv_dim, 1, config.linear_conv_kernel_dim), + ) + + # Decay parameters (no .weight suffix in HF) + self.A_log = nn.Parameter((self.num_value_heads,)) + self.dt_bias = nn.Parameter((self.num_value_heads,)) + + # Output gating norm — per-head RMSNorm (shared weight across heads) + self.norm = nn.RMSNorm(self.value_head_dim, -1, config.rms_norm_eps, bias=False) + + def forward(self, hidden_states: Tensor, state: RNNState) -> Tuple[Tensor, RNNState]: # noqa: UP006 + """Forward using RNNState (for MLCEngine batch methods).""" + b, s, _ = hidden_states.shape + K = self.key_head_dim + V = self.value_head_dim + n_kh = self.num_key_heads + n_vh = self.num_value_heads + layer_idx = self.linear_layer_idx + + # Input projections + qkv = self.in_proj_qkv(hidden_states) + z = self.in_proj_z(hidden_states) + alpha = self.in_proj_a(hidden_states) + beta_raw = self.in_proj_b(hidden_states) + + # Get conv state from RNNState (state_id=1) + qkv_dim = qkv.shape[-1] + conv_state = state.get( + layer_idx, + 1, + (b, self.config.linear_conv_kernel_dim - 1, qkv_dim), + self.dtype, + ) + + # Causal Conv1D using existing helper logic + qkv, new_conv_state = self._causal_conv1d_with_state(qkv, conv_state) + state = state.set(layer_idx, 1, new_conv_state) + + # SiLU activation on QKV after conv + qkv = op.silu(qkv) + + # Split QKV + q_dim = n_kh * K + k_dim = n_kh * K + qkv_parts = op.split(qkv, [q_dim, q_dim + k_dim], axis=-1) + q = op.reshape(qkv_parts[0], (b, s, n_kh, K)) + k = op.reshape(qkv_parts[1], (b, s, n_kh, K)) + v = op.reshape(qkv_parts[2], (b, s, n_vh, V)) + + # L2 normalize Q and K + q = self._l2_normalize(q) + k = self._l2_normalize(k) + + # Gate computation + gate, beta = self._compute_gate_beta(alpha, beta_raw) + # beta is already (b, s, n_vh) — no GVA expansion needed. + + # Get recurrent state from RNNState (state_id=0) + state_in_layer = state.get(layer_idx, 0, (b, n_vh, K, V), "float32") + + # Recurrent computation via TIR kernel + out_recurrent, state_out_layer = op.tensor_ir_op( + create_gated_delta_net_func( + num_key_heads=n_kh, + num_value_heads=n_vh, + key_head_dim=K, + value_head_dim=V, + dtype=self.dtype, + ), + "gated_delta_net", + [q, k, v, gate, beta, state_in_layer], + [ + Tensor.placeholder([b, s, n_vh, V], "float32"), + Tensor.placeholder([b, n_vh, K, V], "float32"), + ], + ) + + # Cast recurrent output back to model dtype + out_recurrent = op.astype(out_recurrent, self.dtype) + + # Write updated state back to RNNState (state_id=0) + state = state.set(layer_idx, 0, state_out_layer) + + # Output gating + out_normed = self.norm(out_recurrent) + out_flat = op.reshape(out_normed, (b, s, n_vh * V)) + out_gated = out_flat * op.silu(z) + return self.out_proj(out_gated), state + + def _causal_conv1d_with_state(self, qkv: Tensor, conv_state: Tensor) -> Tuple[Tensor, Tensor]: # noqa: UP006 + """Causal Conv1D using a pre-extracted conv_state tensor (for RNNState path).""" + b, s, d = qkv.shape + kernel_size = self.config.linear_conv_kernel_dim + + # Update conv state + def _te_update_conv_state(old_state: te.Tensor, qkv_in: te.Tensor): + ks_minus_1 = old_state.shape[1] + seq = qkv_in.shape[1] + return te.compute( + old_state.shape, + lambda bi, ti, di: tirx.if_then_else( + seq + ti < ks_minus_1, + old_state[bi, seq + ti, di], + qkv_in[bi, seq + ti - ks_minus_1, di], + ), + name="update_conv_state", + ) + + new_conv_state = op.tensor_expr_op( + _te_update_conv_state, "update_conv_state", [conv_state, qkv] + ) + + # Depthwise conv + def _te_depthwise_conv(state: te.Tensor, qkv_in: te.Tensor, weight: te.Tensor): + ks_m1 = state.shape[1] + seq = qkv_in.shape[1] + kk = te.reduce_axis((0, kernel_size), name="kk") + return te.compute( + (qkv_in.shape[0], seq, qkv_in.shape[2]), + lambda bi, si, di: te.sum( + tirx.if_then_else( + si + kk < ks_m1, + state[bi, si + kk, di], + qkv_in[bi, si + kk - ks_m1, di], + ) + * weight[di, 0, kk], + axis=kk, + ), + name="depthwise_conv1d", + ) + + result = op.tensor_expr_op( + _te_depthwise_conv, + "depthwise_conv1d", + [conv_state, qkv, self.conv1d_weight], + attrs={"op_pattern": 8}, + ) + return result, new_conv_state + + def _l2_normalize(self, x: Tensor) -> Tensor: + """L2 normalize along last dimension with eps=1e-6.""" + # x: (b, s, h, d) — compute in float32 for numerical stability + x_f32 = op.astype(x, "float32") + x_sq = x_f32 * x_f32 + sum_sq = op.sum(x_sq, axis=-1, keepdims=True) # (b, s, h, 1) + inv_norm = op.sqrt(sum_sq + 1e-6) + return op.astype(x_f32 / inv_norm, self.dtype) + + def _compute_gate_beta(self, alpha: Tensor, beta_raw: Tensor): + """Compute decay gate and update rate. + + gate = exp(-exp(A_log) * softplus(alpha + dt_bias)) (per value_head) + beta = sigmoid(beta_raw) (per value_head) + """ + + # alpha: (b, s, n_vh), dt_bias: (n_vh,), A_log: (n_vh,) + def _te_gate(alpha: te.Tensor, A_log: te.Tensor, dt_bias: te.Tensor): + b, s, h = alpha.shape + + def _softplus(x): + # softplus(x) = x if x > 20 else log(1 + exp(x)) + return tirx.if_then_else(x > 20.0, x, tirx.log(1.0 + tirx.exp(x))) + + return te.compute( + (b, s, h), + lambda bi, si, hi: tirx.exp( + -tirx.exp(A_log[hi].astype("float32")) + * _softplus((alpha[bi, si, hi] + dt_bias[hi]).astype("float32")) + ), + name="gate", + ) + + gate = op.tensor_expr_op( + _te_gate, + "gate", + [alpha, self.A_log, self.dt_bias], + attrs={"op_pattern": 8}, + ) + + beta = op.sigmoid(beta_raw).astype("float32") + return gate, beta + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + # A_log and dt_bias must stay float32 + self.A_log.to("float32") + self.dt_bias.to("float32") + + +# ============================================================================ +# Decoder Layer (dispatches between GDN and standard attention) +# ============================================================================ + + +class Qwen35DecoderLayer(nn.Module): + def __init__(self, config: Qwen35Config, layer_id: int, category_id: int): + """ + layer_id is the id of the layer within all of the layers + category_id is the index of the layer within the category of layers that it belongs to + ie, linear attention or regular attention + """ + self.layer_type = config.layer_types()[layer_id] + if self.layer_type == "full_attention": + self.self_attn = Qwen35Attention(config) + else: + self.linear_attn = Qwen35GatedDeltaNet(config, category_id) + self.category_id = category_id + self.mlp = Qwen35MLP(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + self.tensor_parallel_shards = config.tensor_parallel_shards + + def forward( + self, + hidden_states: Tensor, + paged_kv_cache: PagedKVCache, + state: RNNState, + ): + out = self.input_layernorm(hidden_states) + if self.layer_type == "full_attention": + out = self.self_attn(out, paged_kv_cache, self.category_id) + else: + out, state = self.linear_attn.forward(out, state) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states, state + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Qwen35Model(nn.Module): + def __init__(self, config: Qwen35Config): + self.embed_tokens = Qwen35Embedding(config.vocab_size, config.hidden_size) + layer_types = config.layer_types() + linear_idx = 0 + attn_idx = 0 + layers = [] + for i, ltype in enumerate(layer_types): + if ltype == "linear_attention": + layers.append(Qwen35DecoderLayer(config, i, category_id=linear_idx)) + linear_idx += 1 + else: + layers.append(Qwen35DecoderLayer(config, i, category_id=attn_idx)) + attn_idx += 1 + self.layers = nn.ModuleList(layers) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward( + self, + inputs: Tensor, + paged_kv_cache: PagedKVCache, + state: RNNState, + ): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states, state = layer.forward(hidden_states, paged_kv_cache, state) + hidden_states = self.norm(hidden_states) + return hidden_states, state + + +class Qwen35LMHeadModel(nn.Module): + def __init__(self, config: Qwen35Config): + self.config = config + self.model = Qwen35Model(config) + self.tie_word_embeddings = config.tie_word_embeddings + if not config.tie_word_embeddings: + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.dtype = config.dtype + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.rms_norm_eps = config.rms_norm_eps + self.rope_theta = config.rope_theta + self.vocab_size = config.vocab_size + self.tensor_parallel_shards = config.tensor_parallel_shards + self.partial_rotary_factor = config.partial_rotary_factor + # GDN config + self.num_linear_layers = config.num_linear_layers + self.num_attention_layers = config.num_attention_layers + self.linear_num_value_heads = config.linear_num_value_heads + self.linear_key_head_dim = config.linear_key_head_dim + self.linear_value_head_dim = config.linear_value_head_dim + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def _forward( + self, + input_embed: Tensor, + paged_kv_cache: PagedKVCache, + state: RNNState, + logit_positions: Optional[Tensor] = None, + ): + """Shared forward for batch methods using RNNState.""" + op_ext.configure() + hidden_states, state = self.model.forward(input_embed, paged_kv_cache, state) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + if self.tie_word_embeddings: + logits = self.model.embed_tokens.lm_head_forward(hidden_states) + else: + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache, state + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + rnn_state: RNNState, + ): + return self._forward(input_embeds, paged_kv_cache, rnn_state, logit_positions) + + def batch_decode( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + rnn_state: RNNState, + ): + return self._forward(input_embeds, paged_kv_cache, rnn_state) + + def batch_verify( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + rnn_state: RNNState, + ): + return self._forward(input_embeds, paged_kv_cache, rnn_state) + + def create_rnn_state( + self, + max_batch_size: tirx.Var, + max_history: tirx.Var, + ) -> RNNState: + K = self.linear_key_head_dim + V = self.linear_value_head_dim + n_vh = self.linear_num_value_heads + n_kh = self.config.linear_num_key_heads + qkv_dim = n_kh * K * 2 + n_vh * V + conv_ks_m1 = self.config.linear_conv_kernel_dim - 1 + + init_values = [ + R.const(np.zeros((n_vh, K, V), "float32")), + R.const(np.zeros((conv_ks_m1, qkv_dim), self.dtype)), + ] + + return RNNState.create( + max_batch_size=max_batch_size, + num_hidden_layers=self.num_linear_layers, + max_history=max_history, + init_values=init_values, + ) + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + rotary_dim = int(self.head_dim * self.partial_rotary_factor) + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + # Only attention layers use the KV cache + num_hidden_layers=self.num_attention_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + rotary_dim=rotary_dim, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "rnn_state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "rnn_state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "rnn_state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + "create_rnn_state": { + "max_batch_size": int, + "max_history": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/qwen3_moe/__init__.py b/python/mlc_llm/model/qwen3_moe/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/qwen3_moe/qwen3_moe_loader.py b/python/mlc_llm/model/qwen3_moe/qwen3_moe_loader.py new file mode 100644 index 0000000..8c7e848 --- /dev/null +++ b/python/mlc_llm/model/qwen3_moe/qwen3_moe_loader.py @@ -0,0 +1,158 @@ +""" +This file specifies how MLC's QWen2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools +from typing import Callable, List # noqa: UP035 + +import numpy as np + +from mlc_llm.loader import ExternMapping, QuantizeMapping +from mlc_llm.quantization import BlockScaleQuantize, Quantization + +from .qwen3_moe_model import Qwen3MoeConfig, Qwen3MoeForCausalLM + + +def huggingface(model_config: Qwen3MoeConfig, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : Qwen3MoeConfig + The configuration of the Qwen3Moe model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Qwen3MoeForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + if isinstance(quantization, BlockScaleQuantize): + # Convert the model to block-scale quantized model before loading parameters + model = quantization.quantize_model(model, QuantizeMapping({}, {}), "") + if model_config.weight_block_size is None: + raise ValueError( + "The input Qwen3 model is not fp8 block quantized. " + "Thus BlockScaleQuantize is not supported." + ) + + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + if ( + not isinstance(quantization, BlockScaleQuantize) + and model_config.weight_block_size is not None + ): + raise ValueError( + "The input Qwen3 model is fp8 block quantized. " + "Please use BlockScaleQuantize for the model." + ) + + # Helper function to add both weight and scale mappings + def add_weight_and_scale_mapping( + weight_mlc_name: str, + weight_hf_names: List[str], # noqa: UP006 + weight_transform_func: Callable, + ): + mlc_param = named_parameters[weight_mlc_name] + mapping.add_mapping( + weight_mlc_name, + weight_hf_names, + functools.partial(weight_transform_func, dtype=mlc_param.dtype), + ) + + if isinstance(quantization, BlockScaleQuantize): + scale_mlc_name = f"{weight_mlc_name}_scale_inv" + if scale_mlc_name in named_parameters: + scale_hf_names = [f"{name}_scale_inv" for name in weight_hf_names] + scale_param = named_parameters[scale_mlc_name] + mapping.add_mapping( + scale_mlc_name, + scale_hf_names, + functools.partial(weight_transform_func, dtype=scale_param.dtype), + ) + + for i in range(model_config.num_hidden_layers): + # map attention weight + attn = f"model.layers.{i}.self_attn" + add_weight_and_scale_mapping( + f"{attn}.c_attn.weight", + [ + f"{attn}.q_proj.weight", + f"{attn}.k_proj.weight", + f"{attn}.v_proj.weight", + ], + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + ) + if model_config.attention_bias: + mlc_name = f"{attn}.c_attn.bias" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.bias", + f"{attn}.k_proj.bias", + f"{attn}.v_proj.bias", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + # map mlp moe gate and up weight + mlp = f"model.layers.{i}.mlp" + + def combine_expert_gate_up(*hf_params, dtype): + stack = [] + for i in range(0, len(hf_params), 2): + stack.append(np.concatenate([hf_params[i], hf_params[i + 1]], axis=0)) + return np.stack(stack, axis=0).astype(dtype) + + add_weight_and_scale_mapping( + f"{mlp}.moe_gate_up_proj.weight", + functools.reduce( + lambda a, b: a + b, + [ + [ + f"{mlp}.experts.{expert_id}.gate_proj.weight", + f"{mlp}.experts.{expert_id}.up_proj.weight", + ] + for expert_id in range(model_config.num_experts) + ], + ), + combine_expert_gate_up, + ) + + # map mlp moe down projection weight + add_weight_and_scale_mapping( + f"{mlp}.moe_down_proj.weight", + [ + f"{mlp}.experts.{expert_id}.down_proj.weight" + for expert_id in range(model_config.num_experts) + ], + lambda *hf_params, dtype: np.stack(hf_params, axis=0).astype(dtype), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/qwen3_moe/qwen3_moe_model.py b/python/mlc_llm/model/qwen3_moe/qwen3_moe_model.py new file mode 100644 index 0000000..41d6a2b --- /dev/null +++ b/python/mlc_llm/model/qwen3_moe/qwen3_moe_model.py @@ -0,0 +1,388 @@ +""" +Implementation for QWEN2MOE architecture. +""" + +import dataclasses +from typing import Optional + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.model.qwen3.qwen3_model import ACT2FN, Qwen3Attention, Qwen3Config +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.nn.expert import MixtralExperts +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Qwen3MoeConfig(Qwen3Config): + """Configuration of the Qwen3Moe model.""" + + moe_intermediate_size: int = 0 + num_experts_per_tok: int = 0 + num_experts: int = 0 + decoder_sparse_step: int = 0 + norm_topk_prob: bool = False + + +class Qwen3MoeMLP(nn.Module): + def __init__(self, config: Qwen3MoeConfig, intermediate_size: Optional[int] = None): + intermediate_size = intermediate_size or config.intermediate_size + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MoE MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear(config.hidden_size, 2 * self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(self.act_fn(x1) * x2) + + +class Qwen3MoeSparseMoeBlock(nn.Module): + """MoE layer for Qwen3MoE model.""" + + def __init__(self, config: Qwen3MoeConfig): + super().__init__() + self.num_experts_per_tok = config.num_experts_per_tok + self.num_experts = config.num_experts + if config.moe_intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MoE intermediate size {config.moe_intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.moe_intermediate_size = config.moe_intermediate_size // config.tensor_parallel_shards + self.norm_topk_prob = config.norm_topk_prob + + self.gate = nn.Linear( + in_features=config.hidden_size, + out_features=config.num_experts, + bias=False, + ) + self.moe_gate_up_proj = MixtralExperts( + self.num_experts, + in_features=config.hidden_size, + out_features=2 * self.moe_intermediate_size, + ) + self.moe_down_proj = MixtralExperts( + self.num_experts, + in_features=self.moe_intermediate_size, + out_features=config.hidden_size, + ) + self.act_fn = ACT2FN[config.hidden_act] + self.dtype = "float32" + + def forward(self, x: Tensor): + def _expert_forward(x: Tensor, indptr: Tensor): + x1_x2 = self.moe_gate_up_proj(x, indptr) + x1, x2 = op.split(x1_x2, indices_or_sections=2, axis=-1) + x = self.moe_down_proj(self.act_fn(x1) * x2, indptr) + return x + + experts_per_tok = self.num_experts_per_tok + num_experts = self.num_experts + batch_size, seq_len, hidden_size = x.shape + num_tokens = batch_size * seq_len + x = x.reshape(num_tokens, hidden_size) + gate = self.gate(x) + # expert_weights: [num_tokens, experts_per_tok] + # expert_indices: [num_tokens, experts_per_tok] + expert_weights, expert_indices = op_ext.moe_misc.gating_softmax_topk( + gate, experts_per_tok, norm_topk_prob=self.norm_topk_prob + ) + use_cutlass = op_ext.get_store().cutlass_group_gemm and self.dtype in [ + "float16", + "bfloat16", + ] + if num_tokens == 1: + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = _expert_forward(x, expert_indices) + else: + # cumsum: [num_tokens * local_experts] + cumsum = op_ext.moe_misc.moe_cumsum(expert_indices, num_experts) + # indices: [num_tokens * experts_per_tok] + reverse_indices, token_indices = op_ext.moe_misc.get_indices(cumsum, expert_indices) + # indptr: [num_local_experts + 1] + if use_cutlass: + # indptr: [num_experts] + indptr = op_ext.moe_misc.get_indptr( + cumsum, num_experts, num_tokens, inclusive=True, out_dtype="int64" + ) + else: + # indptr: [num_experts + 1] + indptr = op_ext.moe_misc.get_indptr( + cumsum, num_experts, num_tokens, inclusive=False, out_dtype="int32" + ) + # x: [num_tokens * experts_per_tok, hidden_size] + moe_hidden_states = op.take(x, token_indices, axis=0) + moe_hidden_states = _expert_forward(moe_hidden_states, indptr) + moe_hidden_states = op_ext.moe_misc.scatter_output(moe_hidden_states, reverse_indices) + # moe_hidden_states: [num_tokens, experts_per_tok, hidden_size] + expert_weights = expert_weights.reshape(num_tokens, experts_per_tok, 1) + moe_hidden_states = ( + moe_hidden_states.reshape(num_tokens, experts_per_tok, hidden_size) * expert_weights + ) + # moe_hidden_states: [num_tokens, hidden_size] + moe_hidden_states = op_ext.moe_misc.moe_sum(moe_hidden_states, dim=1) + + final_hidden_states = moe_hidden_states + final_hidden_states = final_hidden_states.reshape(batch_size, seq_len, hidden_size) + return final_hidden_states + + +class Qwen3MoeDecoderLayer(nn.Module): + def __init__(self, config: Qwen3MoeConfig): + super().__init__() + self.self_attn = Qwen3Attention(config) + assert config.num_experts > 0 and config.decoder_sparse_step == 1, ( + "Currently only support use moe for every layer." + ) + self.mlp = Qwen3MoeSparseMoeBlock(config) + self.input_layernorm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, -1, config.rms_norm_eps, bias=False + ) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_attention_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + mi = self.mlp.moe_intermediate_size + _set( + self.self_attn.c_attn.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + if config.attention_bias: + _set( + self.self_attn.c_attn.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.moe_gate_up_proj.weight, + tp.ShardSingleDim("_shard_moe_mlp_up", segs=[mi, mi], dim=1), + ) + _set( + self.mlp.moe_down_proj.weight, + tp.ShardSingleDim("_shard_moe_mlp_down", dim=2), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.input_layernorm(hidden_states) + out = self.self_attn(out, paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.post_attention_layernorm(hidden_states) + out = self.mlp(out) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Qwen3MoeModel(nn.Module): + def __init__(self, config: Qwen3MoeConfig): + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Qwen3MoeDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.RMSNorm(config.hidden_size, -1, config.rms_norm_eps, bias=False) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Qwen3MoeForCausalLM(nn.Module): + def __init__(self, config: Qwen3MoeConfig): + self.model = Qwen3MoeModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.dtype = config.dtype + self.hidden_size = config.hidden_size + self.num_hidden_layers = config.num_hidden_layers + self.intermediate_size = config.intermediate_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.rms_norm_eps = config.rms_norm_eps + self.rope_theta = config.rope_theta + self.vocab_size = config.vocab_size + self.tensor_parallel_shards = config.tensor_parallel_shards + self.head_dim = config.head_dim + self.weight_block_size = config.weight_block_size + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/rwkv5/__init__.py b/python/mlc_llm/model/rwkv5/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/rwkv5/rwkv5_loader.py b/python/mlc_llm/model/rwkv5/rwkv5_loader.py new file mode 100644 index 0000000..9c21331 --- /dev/null +++ b/python/mlc_llm/model/rwkv5/rwkv5_loader.py @@ -0,0 +1,85 @@ +""" +This file specifies how MLC's RWKV5 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from ...loader import ExternMapping +from ...quantization import Quantization +from .rwkv5_model import RWKV5_ForCausalLM, RWKV5Config + + +def huggingface(model_config: RWKV5Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : RWKVConfig + The configuration of the RWKV5 model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = RWKV5_ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params = model.export_tvm(spec=model.get_default_spec()) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # convert time_decay + mlc_name = f"model.blocks.{i}.attention.time_decay" + hf_name = f"rwkv.blocks.{i}.attention.time_decay" + mlc_param = named_parameters[mlc_name] + if mlc_param.dtype != "float32": + raise ValueError(f"RWKV5 time_decay should be float32, got {mlc_param.dtype}") + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: np.exp(-np.exp(x.astype(dtype))), + dtype=mlc_param.dtype, + ), + ) + + # rescale + if model_config.rescale_every > 0: + for name in ["feed_forward.value.weight", "attention.output.weight"]: + mlc_name = f"model.blocks.{i}.{name}" + hf_name = f"rwkv.blocks.{i}.{name}" + mlc_param = named_parameters[mlc_name] + + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype, t: x.astype(dtype) / (2**t), + dtype=mlc_param.dtype, + t=i // model_config.rescale_every, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + hf_name = mlc_name.replace("model", "rwkv") + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + return mapping diff --git a/python/mlc_llm/model/rwkv5/rwkv5_model.py b/python/mlc_llm/model/rwkv5/rwkv5_model.py new file mode 100644 index 0000000..fb36991 --- /dev/null +++ b/python/mlc_llm/model/rwkv5/rwkv5_model.py @@ -0,0 +1,461 @@ +"""Implementation for RWKV5 architecture.""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +import numpy as np +from tvm import relax as R +from tvm import te, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Object, Tensor, op +from tvm.script import tirx as T + +from mlc_llm.nn.rnn_state import RNNState +from mlc_llm.support import logging +from mlc_llm.support.config import ConfigBase + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class StateID: + """State ID for RWKV5.""" + + ATT_X = 0 + ATT_KV = 1 + FFN_X = 2 + + +@dataclasses.dataclass +class RWKV5Config(ConfigBase): + """Configuration of the RWKV5 model.""" + + hidden_size: int + intermediate_size: int + num_hidden_layers: int + vocab_size: int + model_version: str + tensor_parallel_shards: int = 1 + rescale_every: int = 0 + head_size: int = 64 + layer_norm_epsilon: float = 1e-5 + context_window_size: int = -1 # RWKV does not have context window limitation. + prefill_chunk_size: int = 4096 + num_heads: int = 0 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.model_version != "5_2": + raise ValueError(f"Only support RWKV v5_2, got {self.model_version}.") + self.intermediate_size = self.intermediate_size or int(self.hidden_size * 3.5) // 32 * 32 + self.num_heads = ( + self.hidden_size // self.head_size if self.num_heads == 0 else self.num_heads + ) + if self.num_heads * self.head_size != self.hidden_size: + raise ValueError( + f"hidden_size ({self.hidden_size}) must be divisible " + f"by head_size ({self.head_size})" + ) + if self.tensor_parallel_shards != 1: + raise ValueError("Only support single device at this moment.") + + +def create_wkv5_func( + num_heads: int, + head_size: int, + dtype: str, + out_dtype: str, + state_dtype: str, +): + @T.prim_func(s_tir=True) + def wkv_func( + r: T.handle, + k: T.handle, + v: T.handle, + time_decay: T.handle, + time_faaaa: T.handle, + state: T.handle, + out: T.handle, + out_state: T.handle, + ): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + batch_size, seq_len = T.int64(), T.int64() + # Inputs + r_buf = T.match_buffer(r, (batch_size, seq_len, num_heads, head_size), dtype=dtype) + k_buf = T.match_buffer(k, (batch_size, seq_len, num_heads, head_size), dtype=dtype) + v_buf = T.match_buffer(v, (batch_size, seq_len, num_heads, head_size), dtype=dtype) + time_decay_buf = T.match_buffer(time_decay, (num_heads, head_size), dtype="float32") + time_faaaa_buf = T.match_buffer(time_faaaa, (num_heads, head_size), dtype="float32") + state_buf = T.match_buffer( + state, (batch_size, num_heads, head_size, head_size), dtype=state_dtype + ) + # Outputs + out_buf = T.match_buffer(out, (batch_size, seq_len, num_heads, head_size), dtype=out_dtype) + out_state_buf = T.match_buffer( + out_state, (batch_size, num_heads, head_size, head_size), dtype=state_dtype + ) + for b in T.thread_binding(batch_size, thread="blockIdx.y"): + for h in T.thread_binding(num_heads, thread="blockIdx.x"): + for i in T.thread_binding(head_size, thread="threadIdx.x"): + for j in range(head_size): + with T.sblock("init_state"): + vb, vh, vi, vj = T.axis.remap("SSSS", [b, h, i, j]) + out_state_buf[vb, vh, vi, vj] = state_buf[vb, vh, vi, vj] + + for t in range(seq_len): + with T.sblock("comput"): + vb = T.axis.spatial(batch_size, b) + vt = T.axis.opaque(seq_len, t) + vh = T.axis.spatial(num_heads, h) + vi = T.axis.spatial(head_size, i) + out_buf[vb, vt, vh, vi] = 0 + + for k in range(head_size): + x = k_buf[vb, vt, vh, k] * v_buf[vb, vt, vh, vi] + out_buf[vb, vt, vh, vi] += T.cast( + r_buf[vb, vt, vh, k], out_dtype + ) * T.cast( + time_faaaa_buf[vh, k] * x + out_state_buf[vb, vh, vi, k], + out_dtype, + ) + out_state_buf[vb, vh, vi, k] = ( + out_state_buf[vb, vh, vi, k] * time_decay_buf[vh, k] + x + ) + + return wkv_func + + +def token_shift(state: Tensor, x: Tensor): + def _te_token_shift(state: te.Tensor, x: te.Tensor): + return te.compute( + x.shape, + lambda b, i, j: tirx.if_then_else(i == 0, state[b, j], x[b, i - 1, j]), + ) + + return op.tensor_expr_op(_te_token_shift, "token_shift", [state, x]) + + +def last_token(x: Tensor): + # x.shape = (batch, seq_len, hidden_size) + batch, seq_len, hidden_size = x.shape + + def _te_last_token(x: te.Tensor): + return te.compute((batch, 1, hidden_size), lambda b, _, j: x[b, x.shape[1] - 1, j]) + + return x if seq_len == 1 else op.tensor_expr_op(_te_last_token, "last_token", [x]) + + +class RWKV5_FNN(nn.Module): + def __init__(self, config: RWKV5Config, layer_id: int): + super().__init__() + self.time_mix_key = nn.Parameter((1, 1, config.hidden_size)) + self.time_mix_receptance = nn.Parameter((1, 1, config.hidden_size)) + self.key = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.receptance = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.value = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + self.layer_id = layer_id + + def forward(self, x: Tensor, state: RNNState): + batch, _, hidden_size = x.shape + state_x = state.get(self.layer_id, StateID.FFN_X, (batch, hidden_size), x.dtype) + state_x = token_shift(state_x, x) + xk = x * self.time_mix_key + state_x * (1.0 - self.time_mix_key) + xr = x * self.time_mix_receptance + state_x * (1.0 - self.time_mix_receptance) + last_x = last_token(x).reshape(batch, hidden_size) + state = state.set(self.layer_id, StateID.FFN_X, last_x) + r = op.sigmoid(self.receptance(xr)) + xv = op.square(op.relu(self.key(xk))) + return r * self.value(xv), state + + +class RWKV5_Attention(nn.Module): + """Attention layer for RWKV.""" + + def __init__(self, config: RWKV5Config, layer_id: int): + super().__init__() + self.time_decay = nn.Parameter((config.num_heads, config.head_size)) + self.time_faaaa = nn.Parameter((config.num_heads, config.head_size)) + + self.time_mix_gate = nn.Parameter((1, 1, config.hidden_size)) + self.time_mix_key = nn.Parameter((1, 1, config.hidden_size)) + self.time_mix_value = nn.Parameter((1, 1, config.hidden_size)) + self.time_mix_receptance = nn.Parameter((1, 1, config.hidden_size)) + + self.key = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.value = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.receptance = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.gate = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.ln_x = nn.GroupNorm( + config.num_heads, + config.hidden_size, + ) + self.hidden_size = config.hidden_size + self.head_size = config.head_size + self.num_heads = config.num_heads + self.layer_id = layer_id + self.dtype = "float32" + + def forward(self, x: Tensor, state: RNNState): + batch, seq_len, hidden_size = x.shape + assert hidden_size == self.hidden_size + B, T, H, N = ( + batch, + seq_len, + self.head_size, + self.num_heads, + ) + x_state = state.get(self.layer_id, StateID.ATT_X, (batch, self.hidden_size), x.dtype) + x_state = token_shift(x_state, x) + kv_state = state.get( + self.layer_id, + StateID.ATT_KV, + (batch, self.num_heads, self.head_size, self.head_size), + "float32", # Always use float32 for state KV. + ) + + xk = x * self.time_mix_key + x_state * (1.0 - self.time_mix_key) + xv = x * self.time_mix_value + x_state * (1.0 - self.time_mix_value) + xr = x * self.time_mix_receptance + x_state * (1.0 - self.time_mix_receptance) + xg = x * self.time_mix_gate + x_state * (1.0 - self.time_mix_gate) + + r = op.reshape(self.receptance(xr), (B, T, N, H)) + k = op.reshape(self.key(xk), (B, T, N, H)) + v = op.reshape(self.value(xv), (B, T, N, H)) + g = op.silu(self.gate(xg)) + + out, kv_state = op.tensor_ir_op( + create_wkv5_func( + self.num_heads, + self.head_size, + dtype=self.dtype, + out_dtype="float32", + state_dtype="float32", + ), + "wkv5", + [r, k, v, self.time_decay, self.time_faaaa, kv_state], + [ + Tensor.placeholder([B, T, N, H], "float32"), + Tensor.placeholder([B, N, H, H], "float32"), + ], + ) + + last_x = last_token(x).reshape(batch, hidden_size) + state = state.set(self.layer_id, StateID.ATT_X, last_x) + state = state.set(self.layer_id, StateID.ATT_KV, kv_state) + out = op.astype(self.ln_x(op.reshape(out, x.shape), channel_axis=-1, axes=[]), self.dtype) + return self.output(out * g), state + + def to(self, dtype: Optional[str] = None): + # RWKV uses special dtype, so we need to convert it. + if dtype is not None: + self.dtype = dtype + + self.time_mix_gate.to(dtype) + self.time_mix_key.to(dtype) + self.time_mix_value.to(dtype) + self.time_mix_receptance.to(dtype) + self.key.to(dtype) + self.value.to(dtype) + self.receptance.to(dtype) + self.gate.to(dtype) + self.output.to(dtype) + + # These parameters are necessary to be converted to float32. + self.time_decay.to("float32") + self.time_faaaa.to("float32") + self.ln_x.to("float32") + + +class RWKV5_Layer(nn.Module): + def __init__(self, config: RWKV5Config, layer_id: int): + super().__init__() + if layer_id == 0: + self.pre_ln = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + self.ln1 = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + self.ln2 = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + self.attention = RWKV5_Attention(config, layer_id) + self.feed_forward = RWKV5_FNN(config, layer_id) + self.layer_id = layer_id + self.rescale_every = config.rescale_every + + def forward(self, x: Tensor, state: RNNState) -> Tensor: + if self.layer_id == 0: + x = self.pre_ln(x) + att_x, state = self.attention(self.ln1(x), state) + x += att_x + ffn_x, state = self.feed_forward(self.ln2(x), state) + x += ffn_x + if self.rescale_every > 0 and (self.layer_id + 1) % self.rescale_every == 0: + x = x / 2.0 + return x, state + + +class RWKV5_Model(nn.Module): + """Exact same as LlamaModel.""" + + def __init__(self, config: RWKV5Config): + super().__init__() + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.blocks = nn.ModuleList( + [RWKV5_Layer(config, i) for i in range(config.num_hidden_layers)] + ) + self.ln_out = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + + def forward(self, input_embed: Tensor, state: RNNState): + """Forward pass of the model, passing through all decoder layers.""" + hidden_states = input_embed + for block in self.blocks: + hidden_states, state = block(hidden_states, state) + return self.ln_out(hidden_states), state + + +class RWKV5_ForCausalLM(nn.Module): + """Same as LlamaForCausalLM, except for the use of sliding window attention.""" + + def __init__(self, config: RWKV5Config): + self.model = RWKV5_Model(config) + self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_heads = config.num_heads + self.head_size = config.head_size + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def embed(self, input_ids: Tensor): + return self.model.embeddings(input_ids) + + def forward( + self, + input_embed: Tensor, + state: RNNState, + logit_positions: Optional[Tensor] = None, + ): + """Forward pass.""" + hidden_states, state = self.model(input_embed, state) + hidden_states = last_token(hidden_states) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, state + + def prefill(self, input_embed: Tensor, state: RNNState): + """Prefilling the prompt.""" + return self.forward(input_embed, state) + + def decode(self, input_embed: Tensor, state: RNNState): + """Decoding step.""" + return self.forward(input_embed, state) + + def batch_prefill(self, input_embeds: Tensor, logit_positions: Tensor, state: RNNState): + """Prefilling the prompt.""" + return self.forward(input_embeds, state, logit_positions=logit_positions) + + def batch_decode(self, input_embeds: Tensor, state: RNNState): + """Decoding step.""" + return self.forward(input_embeds, state) + + def batch_verify(self, input_embeds: Tensor, state: RNNState): + """Verify step.""" + return self.forward(input_embeds, state) + + def create_rnn_state( + self, + max_batch_size: tirx.Var, + max_history: tirx.Var, + ) -> Object: + """Create RNN state.""" + init_values = [ + R.const(np.zeros((self.hidden_size,), self.dtype)), # ATT_X + R.const( + np.zeros((self.num_heads, self.head_size, self.head_size), "float32") + ), # ATT_KV + R.const(np.zeros((self.hidden_size,), self.dtype)), # FFN_X + ] + return RNNState.create( + max_batch_size=max_batch_size, + num_hidden_layers=self.num_hidden_layers, + max_history=max_history, + init_values=init_values, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_rnn_state": { + "max_batch_size": int, + "max_history": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/rwkv6/__init__.py b/python/mlc_llm/model/rwkv6/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/rwkv6/rwkv6_loader.py b/python/mlc_llm/model/rwkv6/rwkv6_loader.py new file mode 100644 index 0000000..318b850 --- /dev/null +++ b/python/mlc_llm/model/rwkv6/rwkv6_loader.py @@ -0,0 +1,68 @@ +""" +This file specifies how MLC's RWKV6 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +from ...loader import ExternMapping +from ...quantization import Quantization +from .rwkv6_model import RWKV6_ForCausalLM, RWKV6Config + + +def huggingface(model_config: RWKV6Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : RWKV6Config + The configuration of the RWKV6 model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = RWKV6_ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params = model.export_tvm(spec=model.get_default_spec()) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + for i in range(model_config.num_hidden_layers): + # rescale + if model_config.rescale_every > 0: + for name in ["feed_forward.value.weight", "attention.output.weight"]: + mlc_name = f"model.blocks.{i}.{name}" + hf_name = f"rwkv.blocks.{i}.{name}" + mlc_param = named_parameters[mlc_name] + + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype, t: x.astype(dtype) / (2**t), + dtype=mlc_param.dtype, + t=i // model_config.rescale_every, + ), + ) + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + hf_name = mlc_name.replace("model", "rwkv") + mapping.add_mapping( + mlc_name, + [hf_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + return mapping diff --git a/python/mlc_llm/model/rwkv6/rwkv6_model.py b/python/mlc_llm/model/rwkv6/rwkv6_model.py new file mode 100644 index 0000000..f15d900 --- /dev/null +++ b/python/mlc_llm/model/rwkv6/rwkv6_model.py @@ -0,0 +1,516 @@ +"""Implementation for RWKV6 architecture.""" + +import dataclasses +from typing import Any, Dict, Optional, Tuple # noqa: UP035 + +import numpy as np +from tvm import relax as R +from tvm import te, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Object, Tensor, op +from tvm.script import tirx as T + +from mlc_llm.nn.rnn_state import RNNState +from mlc_llm.support import logging +from mlc_llm.support.config import ConfigBase + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class StateID: + """State ID for RWKV6.""" + + ATT_X = 0 + ATT_KV = 1 + FFN_X = 2 + + +@dataclasses.dataclass +class RWKV6Config(ConfigBase): + """Configuration of the RWKV6 model.""" + + hidden_size: int + intermediate_size: int + num_hidden_layers: int + vocab_size: int + model_version: str = "6_0" + tensor_parallel_shards: int = 1 + rescale_every: int = 0 + head_size: int = 64 + layer_norm_epsilon: float = 1e-5 + context_window_size: int = -1 # RWKV does not have context window limitation. + prefill_chunk_size: int = 4096 + num_heads: int = 0 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.model_version != "6_0": + raise ValueError(f"Only support RWKV v6_0, got {self.model_version}.") + self.intermediate_size = self.intermediate_size or int(self.hidden_size * 3.5) // 32 * 32 + self.num_heads = ( + self.hidden_size // self.head_size if self.num_heads == 0 else self.num_heads + ) + if self.num_heads * self.head_size != self.hidden_size: + raise ValueError( + f"hidden_size ({self.hidden_size}) must be divisible " + f"by head_size ({self.head_size})" + ) + if self.tensor_parallel_shards != 1: + raise ValueError("Only support single device at this moment.") + + +def create_wkv6_func( + num_heads: int, + head_size: int, + dtype: str, + out_dtype: str, + state_dtype: str, +): + @T.prim_func(s_tir=True) + def wkv_func( + r: T.handle, + k: T.handle, + v: T.handle, + time_faaaa: T.handle, + w: T.handle, + state: T.handle, + out: T.handle, + out_state: T.handle, + ): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + batch_size, seq_len = T.int64(), T.int64() + # Inputs + r_buf = T.match_buffer(r, (batch_size, seq_len, num_heads, head_size), dtype=dtype) + k_buf = T.match_buffer(k, (batch_size, seq_len, num_heads, head_size), dtype=dtype) + v_buf = T.match_buffer(v, (batch_size, seq_len, num_heads, head_size), dtype=dtype) + time_faaaa_buf = T.match_buffer(time_faaaa, (num_heads, head_size), dtype="float32") + w_buf = T.match_buffer(w, (batch_size, seq_len, num_heads, head_size), dtype="float32") + state_buf = T.match_buffer( + state, (batch_size, num_heads, head_size, head_size), dtype=state_dtype + ) + # Outputs + out_buf = T.match_buffer(out, (batch_size, seq_len, num_heads, head_size), dtype=out_dtype) + out_state_buf = T.match_buffer( + out_state, (batch_size, num_heads, head_size, head_size), dtype=state_dtype + ) + for b in T.thread_binding(batch_size, thread="blockIdx.y"): + for h in T.thread_binding(num_heads, thread="blockIdx.x"): + for i in T.thread_binding(head_size, thread="threadIdx.x"): + for j in range(head_size): + with T.sblock("init_state"): + vb, vh, vi, vj = T.axis.remap("SSSS", [b, h, i, j]) + out_state_buf[vb, vh, vi, vj] = state_buf[vb, vh, vi, vj] + + for t in range(seq_len): + with T.sblock("comput"): + vb = T.axis.spatial(batch_size, b) + vt = T.axis.opaque(seq_len, t) + vh = T.axis.spatial(num_heads, h) + vi = T.axis.spatial(head_size, i) + out_buf[vb, vt, vh, vi] = 0 + + for k in range(head_size): + at = k_buf[vb, vt, vh, k] * v_buf[vb, vt, vh, vi] + out_buf[vb, vt, vh, vi] += T.cast( + r_buf[vb, vt, vh, k], out_dtype + ) * T.cast( + time_faaaa_buf[vh, k] * at + out_state_buf[vb, vh, vi, k], + out_dtype, + ) + out_state_buf[vb, vh, vi, k] = ( + at + w_buf[vb, vt, vh, k] * out_state_buf[vb, vh, vi, k] + ) + + return wkv_func + + +def token_shift(state: Tensor, x: Tensor): + def _te_token_shift(state: te.Tensor, x: te.Tensor): + return te.compute( + x.shape, + lambda b, i, j: tirx.if_then_else(i == 0, state[b, j], x[b, i - 1, j]), + ) + + return op.tensor_expr_op(_te_token_shift, "token_shift", [state, x]) + + +def last_token(x: Tensor): + batch, seq_len, hidden_size = x.shape + + def _te_last_token(x: te.Tensor): + return te.compute((batch, 1, hidden_size), lambda b, _, j: x[b, x.shape[1] - 1, j]) + + return x if seq_len == 1 else op.tensor_expr_op(_te_last_token, "last_token", [x]) + + +def unbind_to_five(x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: # noqa: UP006 + assert x.shape[0] == 5 + + def _te_get_ith(x: te.Tensor, i: int): + return te.compute((1, *x.shape[1:]), lambda _, j, k, ll: x[i, j, k, ll]) + + return ( + op.reshape(op.tensor_expr_op(_te_get_ith, "unbind_to_five", [x, 0]), x.shape[1:]), + op.reshape(op.tensor_expr_op(_te_get_ith, "unbind_to_five", [x, 1]), x.shape[1:]), + op.reshape(op.tensor_expr_op(_te_get_ith, "unbind_to_five", [x, 2]), x.shape[1:]), + op.reshape(op.tensor_expr_op(_te_get_ith, "unbind_to_five", [x, 3]), x.shape[1:]), + op.reshape(op.tensor_expr_op(_te_get_ith, "unbind_to_five", [x, 4]), x.shape[1:]), + ) + + +class RWKV6_FNN(nn.Module): + def __init__(self, config: RWKV6Config, layer_id: int): + super().__init__() + self.time_maa_k = nn.Parameter((1, 1, config.hidden_size)) + self.time_maa_r = nn.Parameter((1, 1, config.hidden_size)) + self.key = nn.Linear(config.hidden_size, config.hidden_size // 2 * 7, bias=False) + self.receptance = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.value = nn.Linear(config.hidden_size // 2 * 7, config.hidden_size, bias=False) + self.layer_id = layer_id + + def forward(self, x: Tensor, state: RNNState): + batch, _, hidden_size = x.shape + state_x = state.get(self.layer_id, StateID.FFN_X, (batch, hidden_size), x.dtype) + state_x = token_shift(state_x, x) + + state_x = state_x - x + xk = x + state_x * self.time_maa_k + xr = x + state_x * self.time_maa_r + + last_x = last_token(x).reshape(batch, hidden_size) + state = state.set(self.layer_id, StateID.FFN_X, last_x) + + r = op.sigmoid(self.receptance(xr)) + xv = op.square(op.relu(self.key(xk))) + return r * self.value(xv), state + + +class RWKV6_Attention(nn.Module): + """Attention layer for RWKV.""" + + def __init__(self, config: RWKV6Config, layer_id: int): + super().__init__() + self.time_maa_x = nn.Parameter((1, 1, config.hidden_size)) + self.time_maa_w = nn.Parameter((1, 1, config.hidden_size)) + self.time_maa_k = nn.Parameter((1, 1, config.hidden_size)) + self.time_maa_v = nn.Parameter((1, 1, config.hidden_size)) + self.time_maa_r = nn.Parameter((1, 1, config.hidden_size)) + self.time_maa_g = nn.Parameter((1, 1, config.hidden_size)) + + # RWKV v6 7B/14B + if config.hidden_size == 4096: + time_mix_extra_dim = 64 + time_decay_extra_dim = 128 + else: + time_mix_extra_dim = 32 + time_decay_extra_dim = 64 + + self.time_maa_w1 = nn.Parameter((config.hidden_size, 5 * time_mix_extra_dim)) + self.time_maa_w2 = nn.Parameter((5, time_mix_extra_dim, config.hidden_size)) + self.time_decay_w1 = nn.Parameter((config.hidden_size, time_decay_extra_dim)) + self.time_decay_w2 = nn.Parameter((time_decay_extra_dim, config.hidden_size)) + self.time_decay = nn.Parameter((1, 1, config.hidden_size)) + self.time_faaaa = nn.Parameter((config.num_heads, config.head_size)) + + self.key = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.value = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.receptance = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.gate = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.ln_x = nn.GroupNorm(config.num_heads, config.hidden_size) + self.hidden_size = config.hidden_size + self.head_size = config.head_size + self.num_heads = config.num_heads + self.layer_id = layer_id + self.dtype = "float32" + + def forward(self, x: Tensor, state: RNNState): + batch, seq_len, hidden_size = x.shape + assert hidden_size == self.hidden_size + B, T, H, N = ( + batch, + seq_len, + self.head_size, + self.num_heads, + ) + state_x = state.get(self.layer_id, StateID.ATT_X, (batch, self.hidden_size), x.dtype) + state_x = token_shift(state_x, x) + state_x = state_x - x + xxx = x + state_x * self.time_maa_x + xxx = op.permute( + op.reshape(op.tanh(op.matmul(xxx, self.time_maa_w1)), (B, T, 5, -1)), + [0, 2, 1, 3], + ) + xxx = op.permute( + op.matmul(xxx, self.time_maa_w2), axes=[1, 0, 2, 3] + ) # it's a batch matrix-matrix multiplication + mw, mk, mv, mr, mg = unbind_to_five(xxx) + + kv_state = state.get( + self.layer_id, + StateID.ATT_KV, + (batch, self.num_heads, self.head_size, self.head_size), + "float32", + ) + + xw = x + state_x * (self.time_maa_w + mw) + xk = x + state_x * (self.time_maa_k + mk) + xv = x + state_x * (self.time_maa_v + mv) + xr = x + state_x * (self.time_maa_r + mr) + xg = x + state_x * (self.time_maa_g + mg) + + r = op.reshape(self.receptance(xr), (B, T, N, H)) + k = op.reshape(self.key(xk), (B, T, N, H)) + v = op.reshape(self.value(xv), (B, T, N, H)) + g = op.silu(self.gate(xg)) + + w = op.reshape(self.time_decay, (1, N, H)).astype("float32") + op.reshape( + op.matmul(op.tanh(op.matmul(xw, self.time_decay_w1)), self.time_decay_w2), + (B, T, N, H), + ).astype("float32") + w = op.exp(op.negative(op.exp(w))) + # w = op.reshape(w, [B, T, N, H]) + + out, kv_state = op.tensor_ir_op( + create_wkv6_func( + num_heads=self.num_heads, + head_size=self.head_size, + dtype=self.dtype, + out_dtype="float32", + state_dtype="float32", + ), + "wkv6", + [r, k, v, self.time_faaaa, w, kv_state], + [ + Tensor.placeholder([B, T, N, H], "float32"), + Tensor.placeholder([B, N, H, H], "float32"), + ], + ) + + last_x = last_token(x).reshape(batch, hidden_size) + state = state.set(self.layer_id, StateID.ATT_X, last_x) + state = state.set(self.layer_id, StateID.ATT_KV, kv_state) + out = op.astype(self.ln_x(op.reshape(out, x.shape), channel_axis=-1, axes=[]), self.dtype) + return self.output(out * g), state + + def to(self, dtype: Optional[str] = None): + # RWKV uses special dtype, so we need to convert it. + if dtype is not None: + self.dtype = dtype + + self.time_maa_x.to(dtype) + self.time_maa_w.to(dtype) + self.time_maa_k.to(dtype) + self.time_maa_v.to(dtype) + self.time_maa_r.to(dtype) + self.time_maa_g.to(dtype) + self.time_maa_w1.to(dtype) + self.time_maa_w2.to(dtype) + self.time_decay_w1.to(dtype) + self.time_decay_w2.to(dtype) + self.key.to(dtype) + self.value.to(dtype) + self.receptance.to(dtype) + self.gate.to(dtype) + self.output.to(dtype) + + # These parameters are necessary to be converted to float32. + self.time_decay.to("float32") + self.time_faaaa.to("float32") + self.ln_x.to("float32") + + +class RWKV6_Layer(nn.Module): + def __init__(self, config: RWKV6Config, layer_id: int): + super().__init__() + if layer_id == 0: + self.pre_ln = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + self.ln1 = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + self.ln2 = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + self.attention = RWKV6_Attention(config, layer_id) + self.feed_forward = RWKV6_FNN(config, layer_id) + self.layer_id = layer_id + self.rescale_every = config.rescale_every + + def forward(self, x: Tensor, state: RNNState) -> Tensor: + if self.layer_id == 0: + x = self.pre_ln(x) + att_x, state = self.attention(self.ln1(x), state) + x += att_x + ffn_x, state = self.feed_forward(self.ln2(x), state) + x += ffn_x + if self.rescale_every > 0 and (self.layer_id + 1) % self.rescale_every == 0: + x = x / 2.0 + return x, state + + +class RWKV6_Model(nn.Module): + """Exact same as LlamaModel.""" + + def __init__(self, config: RWKV6Config): + super().__init__() + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.blocks = nn.ModuleList( + [RWKV6_Layer(config, i) for i in range(config.num_hidden_layers)] + ) + self.ln_out = nn.LayerNorm( + config.hidden_size, + eps=config.layer_norm_epsilon, + ) + + def forward(self, input_embed: Tensor, state: RNNState): + """Forward pass of the model, passing through all decoder layers.""" + hidden_states = input_embed + for block in self.blocks: + hidden_states, state = block(hidden_states, state) + return self.ln_out(hidden_states), state + + +class RWKV6_ForCausalLM(nn.Module): + """Same as LlamaForCausalLM, except for the use of sliding window attention.""" + + def __init__(self, config: RWKV6Config): + self.model = RWKV6_Model(config) + self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_heads = config.num_heads + self.head_size = config.head_size + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def embed(self, input_ids: Tensor): + return self.model.embeddings(input_ids) + + def forward( + self, + input_embed: Tensor, + state: RNNState, + logit_positions: Optional[Tensor] = None, + ): + """Forward pass.""" + hidden_states, state = self.model(input_embed, state) + hidden_states = last_token(hidden_states) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, state + + def prefill(self, input_embed: Tensor, state: RNNState): + """Prefilling the prompt.""" + return self.forward(input_embed, state) + + def decode(self, input_embed: Tensor, state: RNNState): + """Decoding step.""" + return self.forward(input_embed, state) + + def batch_prefill(self, input_embeds: Tensor, logit_positions: Tensor, state: RNNState): + """Prefilling the prompt.""" + return self.forward(input_embeds, state, logit_positions=logit_positions) + + def batch_decode(self, input_embeds: Tensor, state: RNNState): + """Decoding step.""" + return self.forward(input_embeds, state) + + def batch_verify(self, input_embeds: Tensor, state: RNNState): + """Verify step.""" + return self.forward(input_embeds, state) + + def create_rnn_state( + self, + max_batch_size: tirx.Var, + max_history: tirx.Var, + ) -> Object: + """Create RNN state.""" + init_values = [ + R.const(np.zeros((self.hidden_size,), self.dtype)), # ATT_X + R.const( + np.zeros((self.num_heads, self.head_size, self.head_size), "float32") + ), # ATT_KV + R.const(np.zeros((self.hidden_size,), self.dtype)), # FFN_X + ] + return RNNState.create( + max_batch_size=max_batch_size, + num_hidden_layers=self.num_hidden_layers, + max_history=max_history, + init_values=init_values, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "state": nn.spec.Object(object_type=RNNState), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_rnn_state": { + "max_batch_size": int, + "max_history": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/stable_lm/__init__.py b/python/mlc_llm/model/stable_lm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/stable_lm/stablelm_loader.py b/python/mlc_llm/model/stable_lm/stablelm_loader.py new file mode 100644 index 0000000..6f5128e --- /dev/null +++ b/python/mlc_llm/model/stable_lm/stablelm_loader.py @@ -0,0 +1,14 @@ +""" +This file specifies how MLC's StableLM parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +from mlc_llm.loader.standard_loader import make_standard_hf_loader + +from .stablelm_model import StableLmForCausalLM + +huggingface = make_standard_hf_loader( + model_cls=StableLmForCausalLM, + add_qkv_bias=True, + qkv_bias_optional=True, +) diff --git a/python/mlc_llm/model/stable_lm/stablelm_model.py b/python/mlc_llm/model/stable_lm/stablelm_model.py new file mode 100644 index 0000000..d1d2570 --- /dev/null +++ b/python/mlc_llm/model/stable_lm/stablelm_model.py @@ -0,0 +1,376 @@ +""" +Implementation for StableLM architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class StableLmConfig(ConfigBase): + """Configuration of the StableLM model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + layer_norm_eps: float + partial_rotary_factor: float + rope_theta: int + intermediate_size: int + use_qkv_bias: bool = False # Default to False for Stable-LM 3B model + head_dim: int = 0 + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class StableLmAttention(nn.Module): + def __init__(self, config: StableLmConfig): + self.hidden_size = config.hidden_size + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + self.head_dim = config.head_dim + if config.num_key_value_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_key_value_heads} key-value attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.num_heads = config.num_attention_heads // self.tensor_parallel_shards + self.num_key_value_heads = config.num_key_value_heads // self.tensor_parallel_shards + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.rotary_ndims = int(config.partial_rotary_factor * self.head_dim) + + self.qkv_proj = nn.Linear( + in_features=config.hidden_size, + out_features=(self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=config.use_qkv_bias, + ) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.qkv_proj(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +class StableLmMLP(nn.Module): + def __init__(self, config: StableLmConfig): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + self.gate_up_proj = nn.Linear( + in_features=config.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + ) + self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x: Tensor): + concat_x1_x2 = self.gate_up_proj(x) + x1, x2 = op.split(concat_x1_x2, 2, axis=-1) + return self.down_proj(op.silu(x1) * x2) + + +class StableLmDecoderLayer(nn.Module): + def __init__(self, config: StableLmConfig): + norm_eps = config.layer_norm_eps + self.self_attn = StableLmAttention(config) + self.mlp = StableLmMLP(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=norm_eps) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=norm_eps) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + i = self.mlp.intermediate_size + _set( + self.self_attn.qkv_proj.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + if config.use_qkv_bias: + _set( + self.self_attn.qkv_proj.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + _set( + self.mlp.gate_up_proj.weight, + tp.ShardSingleDim("_shard_mlp_up", segs=[i, i], dim=0), + ) + _set(self.mlp.down_proj.weight, tp.ShardSingleDim("_shard_mlp_down", dim=1)) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class StableLmModel(nn.Module): + def __init__(self, config: StableLmConfig): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [StableLmDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class StableLmForCausalLM(nn.Module): + def __init__(self, config: StableLmConfig): + self.model = StableLmModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.dtype = "float32" + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.vocab_size = config.vocab_size + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + self.partial_rotary_factor = config.partial_rotary_factor + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + rotary_dim=int(self.head_dim * self.partial_rotary_factor), + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/starcoder2/__init__.py b/python/mlc_llm/model/starcoder2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/mlc_llm/model/starcoder2/starcoder2_loader.py b/python/mlc_llm/model/starcoder2/starcoder2_loader.py new file mode 100644 index 0000000..072f60b --- /dev/null +++ b/python/mlc_llm/model/starcoder2/starcoder2_loader.py @@ -0,0 +1,99 @@ +""" +This file specifies how MLC's Starcoder2 parameter maps from other formats, for example HuggingFace +PyTorch, HuggingFace safetensors. +""" + +import functools + +import numpy as np + +from mlc_llm.loader import ExternMapping +from mlc_llm.quantization import Quantization + +from .starcoder2_model import Starcoder2Config, Starcoder2ForCausalLM + + +def huggingface(model_config: Starcoder2Config, quantization: Quantization) -> ExternMapping: + """Returns a parameter mapping that maps from the names of MLC LLM parameters to + the names of HuggingFace PyTorch parameters. + + Parameters + ---------- + model_config : InternLMConfig + The configuration of the InternLM model. + + quantization : Quantization + The quantization configuration. + + Returns + ------- + param_map : ExternMapping + The parameter mapping from MLC to HuggingFace PyTorch. + """ + model = Starcoder2ForCausalLM(model_config) + if quantization is not None: + model.to(quantization.model_dtype) + _, _named_params, _ = model.export_tvm( + spec=model.get_default_spec(), + allow_extern=True, + ) + named_parameters = dict(_named_params) + + mapping = ExternMapping() + + mlc_name = "lm_head.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + ["model.embed_tokens.weight"], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + + for i in range(model_config.num_hidden_layers): + # Add QKV in self attention + attn = f"model.layers.{i}.self_attn" + mlc_name = f"{attn}.wqkv_pack.weight" + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.weight", + f"{attn}.k_proj.weight", + f"{attn}.v_proj.weight", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + mlc_name = f"{attn}.wqkv_pack.bias" + if mlc_name in named_parameters: + mlc_param = named_parameters[mlc_name] + mapping.add_mapping( + mlc_name, + [ + f"{attn}.q_proj.bias", + f"{attn}.k_proj.bias", + f"{attn}.v_proj.bias", + ], + functools.partial( + lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype), + dtype=mlc_param.dtype, + ), + ) + # Add gates in MLP + + for mlc_name, mlc_param in named_parameters.items(): + if mlc_name not in mapping.param_map: + mapping.add_mapping( + mlc_name, + [mlc_name], + functools.partial( + lambda x, dtype: x.astype(dtype), + dtype=mlc_param.dtype, + ), + ) + return mapping diff --git a/python/mlc_llm/model/starcoder2/starcoder2_model.py b/python/mlc_llm/model/starcoder2/starcoder2_model.py new file mode 100644 index 0000000..a4de133 --- /dev/null +++ b/python/mlc_llm/model/starcoder2/starcoder2_model.py @@ -0,0 +1,395 @@ +""" +Implementation for Starcoder2 architecture. +""" + +import dataclasses +from typing import Any, Dict, Optional # noqa: UP035 + +from tvm import tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm import op as op_ext +from mlc_llm.model.model_utils import index_last_token +from mlc_llm.nn import PagedKVCache, RopeMode +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp +from mlc_llm.support.config import ConfigBase +from mlc_llm.support.style import bold + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Starcoder2Config(ConfigBase): + """Configuration of the Starcoder2 model.""" + + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + hidden_act: str + norm_epsilon: float + intermediate_size: int + rope_theta: int + use_bias: bool + use_cache: bool + bos_token_id: int + eos_token_id: int + context_window_size: int = 0 + prefill_chunk_size: int = 0 + tensor_parallel_shards: int = 1 + max_batch_size: int = 1 + head_dim: int = 0 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + def __post_init__(self): + if self.context_window_size == 0: + for name in ["max_position_embeddings", "max_sequence_length"]: + if name in self.kwargs: + self.context_window_size = self.kwargs.pop(name) + logger.info( + "%s not found in config.json. Falling back to %s (%d)", + bold("context_window_size"), + bold(name), + self.context_window_size, + ) + break + else: + raise ValueError( + "Unable to determine the maximum sequence length, because none of " + "`context_window_size`, `max_position_embeddings` or `max_sequence_length` is " + "provided in `config.json`." + ) + if self.head_dim == 0: + self.head_dim = self.hidden_size // self.num_attention_heads + assert self.head_dim * self.num_attention_heads == self.hidden_size + if self.prefill_chunk_size == 0: + logger.info( + "%s defaults to %d", + bold("prefill_chunk_size"), + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + elif self.prefill_chunk_size > self.context_window_size: + logger.info( + "Overriding %s from %d to %d", + bold("prefill_chunk_size"), + self.prefill_chunk_size, + min(self.context_window_size, 8192), + ) + self.prefill_chunk_size = min(self.context_window_size, 8192) + + +class Starcoder2Attention(nn.Module): + def __init__(self, config: Starcoder2Config): + super().__init__() # Make sure to call the parent class constructor + self.hidden_size = config.hidden_size + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + if config.num_attention_heads % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split {config.num_attention_heads} attention heads " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + + self.num_heads = config.num_attention_heads // self.tensor_parallel_shards + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads // self.tensor_parallel_shards + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.context_window_size + self.use_bias = config.use_bias + + self.wqkv_pack = nn.Linear( + in_features=self.hidden_size, + out_features=(self.num_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=self.use_bias, + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, self.hidden_size, bias=self.use_bias + ) + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + d, h_q, h_kv = self.head_dim, self.num_heads, self.num_key_value_heads + b, s, _ = hidden_states.shape + qkv = self.wqkv_pack(hidden_states) + qkv = op.reshape(qkv, (b, s, h_q + h_kv + h_kv, d)) + output = op.reshape( + paged_kv_cache.attention_with_fused_qkv( + layer_id, qkv, self.num_heads, sm_scale=self.head_dim**-0.5 + ), + (b, s, h_q * d), + ) + attn_output = self.o_proj(output) + return attn_output + + +class Starcoder2MLP(nn.Module): + def __init__(self, config: Starcoder2Config): + if config.intermediate_size % config.tensor_parallel_shards != 0: + raise ValueError( + f"Cannot split MLP intermediate size {config.intermediate_size} " + f"evenly to {config.tensor_parallel_shards} GPUs." + ) + self.intermediate_size = config.intermediate_size // config.tensor_parallel_shards + embed_dim = config.hidden_size + + self.c_fc = nn.Linear( + in_features=embed_dim, + out_features=self.intermediate_size, + bias=config.use_bias, + ) + self.c_proj = nn.Linear(self.intermediate_size, embed_dim, bias=config.use_bias) + + def forward(self, hidden_states: Tensor): + hidden_states = self.c_fc(hidden_states) + hidden_states = op.gelu(hidden_states, approximate="tanh") + hidden_states = self.c_proj(hidden_states) + return hidden_states + + +class Starcoder2DecoderLayer(nn.Module): + def __init__(self, config: Starcoder2Config): + self.hidden_size = config.hidden_size + self.self_attn = Starcoder2Attention(config) + self.mlp = Starcoder2MLP(config) + self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.norm_epsilon) + self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.norm_epsilon) + + def _set_tp(): + def _set(layer, hint): + layer.attrs["shard_strategy"] = hint + + hd = config.head_dim + q = self.self_attn.num_heads * hd + k = self.self_attn.num_key_value_heads * hd + v = self.self_attn.num_key_value_heads * hd + _set( + self.self_attn.wqkv_pack.weight, + tp.ShardSingleDim("_shard_qkv_weight", dim=0, segs=[q, k, v]), + ) + if config.use_bias: + _set( + self.self_attn.wqkv_pack.bias, + tp.ShardSingleDim("_shard_qkv_bias", dim=0, segs=[q, k, v]), + ) + + _set(self.self_attn.o_proj.weight, tp.ShardSingleDim("_shard_o", dim=1)) + + _set( + self.mlp.c_fc.weight, + tp.ShardSingleDim("_shard_c_fc_weight", dim=0), + ) + if config.use_bias: + _set(self.mlp.c_fc.bias, tp.ShardSingleDim("_shard_c_fc_bias", dim=0)) + + _set(self.mlp.c_proj.weight, tp.ShardSingleDim("_shard_mlp_c_proj", dim=1)) + + if config.use_bias: + _set( + self.mlp.c_proj.bias, + tp.ShardSingleDim("_shard_mlp_c_proj_bias", dim=0), + ) + + self.tensor_parallel_shards = config.tensor_parallel_shards + _set_tp() + + def forward(self, hidden_states: Tensor, paged_kv_cache: PagedKVCache, layer_id: int): + out = self.self_attn(self.input_layernorm(hidden_states), paged_kv_cache, layer_id) + hidden_states = self._apply_residual(out, residual=hidden_states) + out = self.mlp(self.post_attention_layernorm(hidden_states)) + hidden_states = self._apply_residual(out, residual=hidden_states) + return hidden_states + + def _apply_residual(self, out, residual): + if self.tensor_parallel_shards > 1: + return op.ccl_allreduce(out, "sum") + residual + return out + residual + + +class Starcoder2Model(nn.Module): + def __init__(self, config: Starcoder2Config): + assert config.hidden_size % config.num_attention_heads == 0 + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Starcoder2DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = nn.LayerNorm(config.hidden_size, config.norm_epsilon) + + def forward(self, inputs: Tensor, paged_kv_cache: PagedKVCache): + hidden_states = inputs + for layer_id, layer in enumerate(self.layers): + hidden_states = layer(hidden_states, paged_kv_cache, layer_id) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Starcoder2ForCausalLM(nn.Module): + def __init__(self, config: Starcoder2Config): + self.model = Starcoder2Model(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.vocab_size = config.vocab_size + self.num_hidden_layers = config.num_hidden_layers + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.vocab_size = config.vocab_size + self.rope_theta = config.rope_theta + self.tensor_parallel_shards = config.tensor_parallel_shards + self.dtype = "float32" + + def to(self, dtype: Optional[str] = None): + super().to(dtype=dtype) + if dtype is not None: + self.dtype = dtype + + def batch_forward( + self, + input_embeds: Tensor, + paged_kv_cache: PagedKVCache, + logit_positions: Optional[Tensor] = None, + ): + op_ext.configure() + + hidden_states = self.model(input_embeds, paged_kv_cache) + if logit_positions is not None: + hidden_states = op.take(hidden_states, logit_positions, axis=1) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits + + def embed(self, input_ids: Tensor): + if self.tensor_parallel_shards > 1: + input_ids = op.ccl_broadcast_from_worker0(input_ids) + return self.model.embed_tokens(input_ids) + + def prefill(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + hidden_states = index_last_token(hidden_states) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def decode(self, input_embed: Tensor, paged_kv_cache: PagedKVCache): + op_ext.configure() + + hidden_states = self.model(input_embed, paged_kv_cache) + logits = self.lm_head(hidden_states) + if logits.dtype != "float32": + logits = logits.astype("float32") + return logits, paged_kv_cache + + def batch_prefill( + self, + input_embeds: Tensor, + logit_positions: Tensor, + paged_kv_cache: PagedKVCache, + ): + if self.tensor_parallel_shards > 1: + logit_positions = op.ccl_broadcast_from_worker0(logit_positions) + logits = self.batch_forward(input_embeds, paged_kv_cache, logit_positions) + return logits, paged_kv_cache + + def batch_decode(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def batch_verify(self, input_embeds: Tensor, paged_kv_cache: PagedKVCache): + logits = self.batch_forward(input_embeds, paged_kv_cache) + return logits, paged_kv_cache + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=self.num_hidden_layers, + num_attention_heads=self.num_attention_heads // self.tensor_parallel_shards, + num_key_value_heads=self.num_key_value_heads // self.tensor_parallel_shards, + qk_head_dim=self.head_dim, + v_head_dim=self.head_dim, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=self.rope_theta, + dtype=self.dtype, + ) + + def get_default_spec(self): + mod_spec = { + "embed": { + "input_ids": nn.spec.Tensor(["seq_len"], "int32"), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "prefill": { + "input_embed": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "decode": { + "input_embed": nn.spec.Tensor([1, 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_prefill": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "logit_positions": nn.spec.Tensor(["batch_size"], "int32"), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_decode": { + "input_embeds": nn.spec.Tensor(["batch_size", 1, self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "batch_verify": { + "input_embeds": nn.spec.Tensor([1, "seq_len", self.hidden_size], self.dtype), + "paged_kv_cache": nn.spec.Object(object_type=PagedKVCache), + "$": { + "param_mode": "packed", + "effect_mode": "none", + }, + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + "$": { + "param_mode": "none", + "effect_mode": "none", + }, + }, + } + return nn.spec.ModuleSpec.from_raw(mod_spec, self) diff --git a/python/mlc_llm/model/vision/__init__.py b/python/mlc_llm/model/vision/__init__.py new file mode 100644 index 0000000..5ee3f2d --- /dev/null +++ b/python/mlc_llm/model/vision/__init__.py @@ -0,0 +1,4 @@ +"""Common `nn.Modules` used to define LLMs in this project.""" + +from .clip_vision import CLIPVisionConfig, CLIPVisionModel +from .image_processing import ImageProcessor diff --git a/python/mlc_llm/model/vision/clip_vision.py b/python/mlc_llm/model/vision/clip_vision.py new file mode 100644 index 0000000..4b01215 --- /dev/null +++ b/python/mlc_llm/model/vision/clip_vision.py @@ -0,0 +1,231 @@ +""" +Implements the CLIP Vision Encoder. +""" + +import dataclasses +import logging +from typing import Any, Dict, Tuple # noqa: UP035 + +from tvm import relax +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Module, Tensor +from tvm.relax.frontend.nn.modules import Conv2D +from tvm.relax.frontend.nn.op import ( + add, + broadcast_to, + concat, + permute_dims, + reshape, + wrap_nested, +) +from tvm.relax.op import arange + +from mlc_llm import op as op_ext +from mlc_llm.support.config import ConfigBase + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class CLIPVisionConfig(ConfigBase): + """ + Config for the vision encoder + """ + + hidden_size: int + image_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + patch_size: int + projection_dim: int + vocab_size: int + num_channels: int = 3 + layer_norm_eps: float = 1e-06 + kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict) # noqa: UP006 + + +class CLIPVisionEmbeddings(Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + self.class_embedding = nn.Parameter((self.embed_dim,)) + self.patch_embedding = Conv2D( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + bias=False, + ) + + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.num_positions = self.num_patches + 1 + self.position_embedding = nn.Embedding(num=self.num_positions, dim=self.embed_dim) + + def forward(self, pixel_values: Tensor) -> Tensor: + batch_size = pixel_values.shape[0] + patch_embeds = self.patch_embedding(pixel_values) # shape = [*, width, grid, grid] + patch_embeds = reshape(patch_embeds, shape=(batch_size, self.embed_dim, -1)) + patch_embeds = permute_dims( + patch_embeds, axes=(0, 2, 1) + ) # shape = [batch,grid*grid,embed_dim] + class_embeds = broadcast_to( + self.class_embedding, shape=(batch_size, 1, self.embed_dim) + ) # shape of (batch,1,embed_dim) + embeddings = concat([class_embeds, patch_embeds], dim=1) + + posi_ids = reshape( + wrap_nested(arange(0, self.num_positions, dtype="int32"), name="arange"), + shape=(1, -1), + ) + batch_position_embedding = broadcast_to( + self.position_embedding(posi_ids), + shape=(batch_size, self.num_positions, self.embed_dim), + ) + embeddings = add(embeddings, batch_position_embedding) + return embeddings + + +def sigmoid(x: Tensor, name: str = "sigmoid") -> Tensor: + """Sigmoid of a Tensor + + Parameters + ---------- + x : Tensor + Input tensor to expand. + name : str + Name hint for this operator. + + Returns + ------- + result : Tensor + Sigmoid result. + """ + return wrap_nested(relax.op.sigmoid(x._expr), name) + + +class QuickGELU(Module): + def forward(self, input_tensor: Tensor) -> Tensor: + return input_tensor * sigmoid(input_tensor * 1.702) + + +class CLIPMLP(Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.activation_fn = QuickGELU() + self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) + self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + + def forward(self, hidden_states: Tensor) -> Tensor: + hidden_states = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states = self.fc2(hidden_states) + return hidden_states + + +class CLIPAttention(Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.embed_dim // self.num_heads + if (self.head_dim * self.num_heads) != self.embed_dim: + raise ValueError( + f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}" + f" and `num_heads`: {self.num_heads})." + ) + self.scale = self.head_dim**-0.5 + self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) + self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) + + def forward( + self, + hidden_states: Tensor, + ) -> Tensor: + d, h = self.head_dim, self.num_heads + b, s, _ = hidden_states.shape # batch_size, seq_len, embed_dim + + q = self.q_proj(hidden_states).reshape(b, s, h, d) + k = self.k_proj(hidden_states).reshape(b, s, h, d) + v = self.v_proj(hidden_states).reshape(b, s, h, d) + + attn_output = op_ext.attention(q, k, v, None) + attn_output = self.out_proj(attn_output) + return attn_output + + +class CLIPEncoderLayer(Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = CLIPAttention(config) + self.layer_norm1 = nn.LayerNorm(normalized_shape=self.embed_dim, eps=config.layer_norm_eps) + self.mlp = CLIPMLP(config) + self.layer_norm2 = nn.LayerNorm(normalized_shape=self.embed_dim, eps=config.layer_norm_eps) + + def forward(self, hidden_states: Tensor) -> Tensor: + residual = hidden_states + hidden_states = self.layer_norm1(hidden_states) + hidden_states = self.self_attn(hidden_states=hidden_states) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.layer_norm2(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + return outputs + + +class CLIPEncoder(Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.layers = nn.ModuleList( + [CLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + + def forward(self, inputs_embeds: Tensor) -> Tensor: + hidden_states = inputs_embeds + encoder_states: Tuple[Any, ...] = () # noqa: UP006 + for _, encoder_layer in enumerate(self.layers): + encoder_states = (*encoder_states, hidden_states) + layer_outputs = encoder_layer(hidden_states) + hidden_states = layer_outputs[0] + encoder_states = (*encoder_states, hidden_states) + return encoder_states + + +class CLIPVisionTransformer(Module): + def __init__(self, config: CLIPVisionConfig): + super().__init__() + embed_dim = config.hidden_size + self.embeddings = CLIPVisionEmbeddings(config) + self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + self.encoder = CLIPEncoder(config) + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + def forward(self, pixel_values: Tensor) -> Tensor: + hidden_states = self.embeddings(pixel_values) + hidden_states = self.pre_layrnorm(hidden_states) + encoder_outputs = self.encoder(inputs_embeds=hidden_states) + # Apply post_layernorm to the final encoder hidden state, matching + # the HuggingFace CLIPVisionTransformer which returns post-normed + # last_hidden_state. Intermediate states remain unnormalized. + last_hidden_state = self.post_layernorm(encoder_outputs[-1]) + return (*encoder_outputs[:-1], last_hidden_state) + + +class CLIPVisionModel(Module): + no_quantization: bool = True + + def __init__(self, config: CLIPVisionConfig): + super().__init__() + self.vision_model = CLIPVisionTransformer(config) + + def forward(self, pixel_values: Tensor) -> Tensor: + return self.vision_model(pixel_values)[-2] diff --git a/python/mlc_llm/model/vision/image_processing.py b/python/mlc_llm/model/vision/image_processing.py new file mode 100644 index 0000000..955b1fa --- /dev/null +++ b/python/mlc_llm/model/vision/image_processing.py @@ -0,0 +1,283 @@ +""" +Implements the CLIP Image processor. +""" + +from tvm import s_tir, tirx +from tvm.relax.frontend.nn import Module, Tensor, op +from tvm.script import tirx as T + + +def _var(dtype, size=1): + return T.sblock_alloc_buffer((size,), dtype, scope="local") + + +class ImageProcessor(Module): + def __init__(self): + super().__init__() + + def apply_schedule(self, sch, block, bdx=32, tile=[32, 32]): + loop_x, loop_y = sch.get_loops(block)[-2:] + xo, xi = sch.split(loop_x, factors=[tile[0], None]) + yo, yi = sch.split(loop_y, factors=[tile[1], None]) + sch.reorder(xo, yo, xi, yi) + t = sch.fuse(xo, yo) + ty, tx = sch.split(t, factors=[None, bdx]) + sch.bind(ty, "threadIdx.y") + sch.bind(tx, "threadIdx.x") + + def resize(self, image: Tensor, params): # image layout:NCHW + assert 4 == image.ndim, "image should be 4D data tensor" + assert 3 == image.shape[1], "image layout should be NCHW" + + def get_output_image_size(image: Tensor): + h = image.shape[2] + w = image.shape[3] + + if "height" in params and "width" in params: + return (params["height"], params["width"]) + elif "shortest_edge" in params: + short = tirx.Select(w < h, w, h) + long = tirx.Select(w > h, w, h) + requested_new_short = params["shortest_edge"] + new_short, new_long = ( + tirx.Cast("int64", requested_new_short), + tirx.Cast( + "int64", + requested_new_short + * tirx.div( + tirx.Cast("float32", long), + tirx.Cast("float32", short), + ), + ), + ) + ret_h = tirx.Select(w <= h, new_long, new_short) + ret_w = tirx.Select(w <= h, new_short, new_long) + return (ret_h, ret_w) + elif "hd_transform" in params: + hd_num = 4 if "hd_num" not in params else params["hd_num"] + pad_num = 336 if "pad_num" not in params else params["pad_num"] + ratio = tirx.Select( + w > h, + tirx.div(tirx.Cast("float32", w), tirx.Cast("float32", h)), + tirx.div(tirx.Cast("float32", h), tirx.Cast("float32", w)), + ) + + scale = tirx.ceil(tirx.sqrt(tirx.Cast("float32", hd_num) * ratio)) + + scale = tirx.Select( + (scale * tirx.ceil(tirx.div(scale, ratio))) > hd_num, + scale - 1, + scale, + ) + scale = tirx.Cast("int64", scale) + + new_w = tirx.Select( + w >= h, + scale * pad_num, + tirx.Cast("int64", tirx.div(scale * pad_num, ratio)), + ) + new_h = tirx.Select( + w >= h, + tirx.Cast("int64", tirx.div(new_w, ratio)), + scale * pad_num, + ) + return (new_h, new_w) + else: + assert False, "not supported resize parameter" + + new_h, new_w = get_output_image_size(image) + out = op.interpolate(image, (new_h, new_w), data_layout="NCHW", mode="linear") + return out + + def crop(self, image: Tensor, crop_size): + assert 4 == image.ndim, "image should be 4D data tensor" + assert 3 == image.shape[1], "image layout should be NCHW" + + def create_crop_func(dtype): # , top, bottom, left, right): + @T.prim_func(s_tir=True) + def crop_func( + image: T.handle, + out: T.handle, + top: T.int64(), + bottom: T.int64(), + left: T.int64(), + right: T.int64(), + ): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64() + image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype) + out_buf = T.match_buffer(out, (n, c, bottom - top, right - left), dtype=dtype) + out_h = bottom - top + out_w = right - left + for n_idx in T.thread_binding(n, thread="blockIdx.x"): + for c_idx in T.thread_binding(c, thread="blockIdx.y"): + for h_idx, w_idx in T.grid(out_h, out_w): + with T.sblock("crop"): + if (h_idx + T.int64(top)) < h and (w_idx + T.int64(left)) < w: + T.writes(out_buf[n_idx, c_idx, h_idx, w_idx]) + T.reads(image_buf[n_idx, c_idx, h_idx + top, w_idx + left]) + out_buf[n_idx, c_idx, h_idx, w_idx] = image_buf[ + n_idx, c_idx, h_idx + top, w_idx + left + ] + + sch = s_tir.Schedule(crop_func) + self.apply_schedule(sch, sch.get_sblock("crop")) + return sch.mod["main"].with_attr("tirx.is_scheduled", 1) + + n, c, orig_height, orig_width = image.shape + crop_height = crop_size["height"] + crop_width = crop_size["width"] + + top = (orig_height - crop_height) // 2 + bottom = orig_height - top + + left = (orig_width - crop_width) // 2 + right = orig_width - left + + out = op.tensor_ir_op( + create_crop_func(image.dtype), + "crop", + [image, top, bottom, left, right], + [Tensor.placeholder([n, c, crop_height, crop_width], image.dtype)], + ) + return out + + def rescale(self, image: Tensor, rescale_factor=1 / 255.0, o_dtype="float32"): + assert 4 == image.ndim, "image should be 4D data tensor" + assert 3 == image.shape[1], "image layout should be NCHW" + + def create_rescale_func(rescale_factor, dtype, o_dtype): + @T.prim_func(s_tir=True) + def rescale_func(image: T.handle, out: T.handle): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64() + image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype) + out_buf = T.match_buffer(out, (n, c, h, w), dtype=o_dtype) + + for n_idx in T.thread_binding(n, thread="blockIdx.x"): + for c_idx in T.thread_binding(c, thread="blockIdx.y"): + for h_idx, w_idx in T.grid(h, w): + with T.sblock("rescale"): + T.reads(image_buf[n_idx, c_idx, h_idx, w_idx]) + T.writes(out_buf[n_idx, c_idx, h_idx, w_idx]) + if h_idx < h and w_idx < w: + out_buf[n_idx, c_idx, h_idx, w_idx] = ( + T.cast( + image_buf[n_idx, c_idx, h_idx, w_idx], + o_dtype, + ) + * rescale_factor + ) + + sch = s_tir.Schedule(rescale_func) + self.apply_schedule(sch, sch.get_sblock("rescale")) + return sch.mod["main"].with_attr("tirx.is_scheduled", 1) + + out = op.tensor_ir_op( + create_rescale_func(rescale_factor, image.dtype, o_dtype), + "rescale", + [image], + [Tensor.placeholder(image.shape, o_dtype)], + ) + return out + + def normalize(self, image: Tensor, o_dtype="float32"): + assert 4 == image.ndim, "image should be 4D data tensor" + assert 3 == image.shape[1], "image layout should be NCHW" + + def create_normalize_func(dtype, o_dtype): + @T.prim_func(s_tir=True) + def normalize_func(image: T.handle, out: T.handle): + n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64() + image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype) + out_buf = T.match_buffer(out, (n, c, h, w), dtype=o_dtype) + mean = _var(o_dtype, 3) + stddev = _var(o_dtype, 3) + + for n_idx in T.thread_binding(n, thread="blockIdx.x"): + for c_idx in T.thread_binding(c, thread="blockIdx.y"): + for h_idx, w_idx in T.grid(h, w): + with T.sblock("normalize"): + T.reads( + image_buf[n_idx, c_idx, h_idx, w_idx], + mean[c_idx], + stddev[c_idx], + ) + T.writes(out_buf[n_idx, c_idx, h_idx, w_idx]) + with T.init(): + mean[0] = 0.48145466 + stddev[0] = 0.26862954 + mean[1] = 0.4578275 + stddev[1] = 0.26130258 + mean[2] = 0.40821073 + stddev[2] = 0.27577711 + if h_idx < h and w_idx < w: + out_buf[n_idx, c_idx, h_idx, w_idx] = ( + T.cast( + image_buf[n_idx, c_idx, h_idx, w_idx], + o_dtype, + ) + - mean[c_idx] + ) / stddev[c_idx] + + sch = s_tir.Schedule(normalize_func) + self.apply_schedule(sch, sch.get_sblock("normalize")) + return sch.mod["main"].with_attr("tirx.is_scheduled", 1) + + out = op.tensor_ir_op( + create_normalize_func(image.dtype, o_dtype), + "normalize", + [image], + [Tensor.placeholder(image.shape, o_dtype)], + ) + return out + + def pad(self, image: Tensor, dtype="uint8"): + assert 4 == image.ndim, "image should be 4D data tensor" + assert 3 == image.shape[1], "image layout should be NCHW" + + def create_pad_func(left, right, fill=255): + @T.prim_func(s_tir=True) + def pad_func(image: T.handle, out: T.handle, t: T.int64(), b: T.int64()): + T.func_attr({"op_pattern": 8, "tirx.noalias": True, "tirx.is_scheduled": 1}) + n, c, h, w = T.int64(), T.int64(), T.int64(), T.int64() + image_buf = T.match_buffer(image, (n, c, h, w), dtype=dtype) + out_buf = T.match_buffer(out, (n, c, h + t + b, w + left + right), dtype=dtype) + out_h = h + t + b + out_w = w + left + right + + for n_idx in T.thread_binding(n, thread="blockIdx.x"): + for c_idx in T.thread_binding(c, thread="blockIdx.y"): + for h_idx, w_idx in T.grid(out_h, out_w): + with T.sblock("pad"): + T.reads(image_buf[n_idx, c_idx, h_idx, w_idx]) + T.writes(out_buf[n_idx, c_idx, h_idx, w_idx]) + if h_idx < t or h_idx > h + b or w_idx < left or w_idx > w + right: + out_buf[n_idx, c_idx, h_idx, w_idx] = fill + else: + out_buf[n_idx, c_idx, h_idx, w_idx] = image_buf[ + n_idx, c_idx, h_idx - t, w_idx - left + ] + + sch = s_tir.Schedule(pad_func) + self.apply_schedule(sch, sch.get_sblock("pad")) + return sch.mod["main"].with_attr("tirx.is_scheduled", 1) + + h = image.shape[2] + tar = tirx.truncdiv(h + 335, 336) * 336 + t = tirx.div(tar - h, 2) + b = tar - h - t + left = 0 + right = 0 + + n, c, h, w = image.shape + out = op.tensor_ir_op( + create_pad_func(left, right), + "pad", + [image, t, b], + [Tensor.placeholder((n, c, tar, w), image.dtype)], + ) + return out + + def preprocess(self, pixel_values): + return pixel_values diff --git a/python/mlc_llm/nn/__init__.py b/python/mlc_llm/nn/__init__.py new file mode 100644 index 0000000..bd3ee5c --- /dev/null +++ b/python/mlc_llm/nn/__init__.py @@ -0,0 +1,4 @@ +"""Common `nn.Modules` used to define LLMs in this project.""" + +from .expert import MixtralExperts +from .kv_cache import PagedKVCache, RopeMode diff --git a/python/mlc_llm/nn/expert.py b/python/mlc_llm/nn/expert.py new file mode 100644 index 0000000..6548ce7 --- /dev/null +++ b/python/mlc_llm/nn/expert.py @@ -0,0 +1,33 @@ +"""An nn.Module that represents MoE experts""" + +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor + +from mlc_llm.op import cutlass, extern, ft_gemm, moe_matmul + + +class MixtralExperts(nn.Module): + """Mixtral experts""" + + def __init__(self, num_local_experts, in_features, out_features, tensor_parallel_shards=1): + self.num_local_experts = num_local_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter((num_local_experts, out_features, in_features)) + self.dtype = "float32" + self.tensor_parallel_shards = tensor_parallel_shards + + def forward(self, x: Tensor, indptr: Tensor): + assert x.ndim == 2 + if indptr.ndim == 2: + assert indptr.shape[0] == 1 + return moe_matmul.gemv(x, self.weight, indptr) + assert indptr.ndim == 1 + if extern.get_store().cutlass_group_gemm and self.dtype in [ + "float16", + "bfloat16", + ]: + return cutlass.group_gemm(x, self.weight, indptr) + if extern.get_store().faster_transformer and self.dtype == "float16": + return ft_gemm.faster_transformer_moe_gemm(x, self.weight, indptr) + return moe_matmul.group_gemm(x, self.weight, indptr) diff --git a/python/mlc_llm/nn/kv_cache.py b/python/mlc_llm/nn/kv_cache.py new file mode 100644 index 0000000..94408e2 --- /dev/null +++ b/python/mlc_llm/nn/kv_cache.py @@ -0,0 +1,93 @@ +"""Attention KV cache modeling.""" + +import json +from typing import Any, Dict, List, Literal, Optional, Union # noqa: UP035 + +import numpy as np +from tvm import relax as rx +from tvm import tirx +from tvm.relax.frontend.nn.llm.kv_cache import PagedKVCache as TVMPagedKVCache +from tvm.relax.frontend.nn.llm.kv_cache import RopeMode + + +class PagedKVCache(TVMPagedKVCache): + """The Paged KV Cache used in LLM batching for efficient attention computation.""" + + @staticmethod + def create_generic( + attn_kind: Union[Literal["mha", "mla"], List[Literal["mha", "mla", "mha_sliding"]]], # noqa: UP006 + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + num_hidden_layers: int, + num_attention_heads: int, + num_key_value_heads: int, + qk_head_dim: int, + v_head_dim: int, + rope_mode: RopeMode, + rope_scale: int, + rope_theta: int, + dtype: str, + mla_original_qk_head_dim: int = 0, + mla_original_v_head_dim: int = 0, + rotary_dim: Optional[int] = None, + rope_scaling: Optional[Dict[str, Any]] = None, # noqa: UP006 + rope_ext_factors: Optional[List[int]] = None, # noqa: UP006 + layer_partition: Optional[List[int]] = None, # noqa: UP006 + enable_disaggregation: bool = False, + name: str = "paged_kv_cache", + ) -> "PagedKVCache": + """The generic function of creating a multi-head attention PagedKVCache, + which will be rewritten by functions in compilation pipeline. + """ + if rotary_dim is None: + rotary_dim = qk_head_dim + if rope_scaling is None: + rope_scaling = {} + if layer_partition is None: + layer_partition = [0, num_hidden_layers] + if isinstance(attn_kind, List): # noqa: UP006 + rx_attn_kind = [rx.StringImm(layer_kind) for layer_kind in attn_kind] + else: + rx_attn_kind = rx.StringImm(attn_kind) + return PagedKVCache( + _expr=rx.call_pure_packed( + "mlc.create_paged_kv_cache_generic", + rx_attn_kind, + rx.ShapeExpr( + [ + max_batch_size, + max_total_seq_len, + prefill_chunk_size, + page_size, + support_sliding_window, + ] + ), + rx.ShapeExpr(layer_partition), + rx.prim_value(num_hidden_layers), + rx.prim_value(num_attention_heads), + rx.prim_value(num_key_value_heads), + rx.prim_value(qk_head_dim), + rx.prim_value(v_head_dim), + rx.prim_value(mla_original_qk_head_dim), + rx.prim_value(mla_original_v_head_dim), + rx.prim_value(rope_mode), + rx.prim_value(rope_scale), + rx.prim_value(rope_theta), + rx.StringImm(json.dumps(rope_scaling)), + ( + rx.const(np.array(rope_ext_factors, "float32")) + if rope_ext_factors is not None + else rx.prim_value(0) + # NOTE: since relax does not have "Optional" type, we use prim_value(0) + # to represent "undefined". + ), + rx.prim_value(rotary_dim), + rx.prim_value(int(enable_disaggregation)), + rx.DataTypeImm(dtype), + ty_args=rx.ObjectType(), + ), + _name=name, + ) diff --git a/python/mlc_llm/nn/rnn_state.py b/python/mlc_llm/nn/rnn_state.py new file mode 100644 index 0000000..5019a52 --- /dev/null +++ b/python/mlc_llm/nn/rnn_state.py @@ -0,0 +1,337 @@ +"""RNN State modeling.""" + +from collections.abc import Sequence +from typing import Union + +from tvm import relax as rx +from tvm import tirx +from tvm.relax.frontend.nn import Object, Tensor +from tvm.script import tirx as T + + +class RNNState(Object): + """The RNN State used in Space State Models""" + + @staticmethod + def create( + max_batch_size: tirx.Var, + num_hidden_layers: int, + max_history: int, + init_values: Sequence[rx.Constant], + name: str = "rnn_state", + ) -> "RNNState": + """Create a RNN state object. + + Parameters + ---------- + max_batch_size : tirx.Var + The maximum batch size. + num_hidden_layers : int + The number of hidden layers. + max_history : int + The maximum history length. + init_values : Sequence[rx.Constant] + The initial values of the RNN state. Must be compile-time Relax constants + (e.g. R.const(np.zeros(...))). + """ + + bb = rx.BlockBuilder.current() + state_infos = [ + (tuple(int(x) for x in v.data.shape), str(v.data.dtype)) for v in init_values + ] + + f_gets = [ + bb.add_func( + RNNState.create_get_func(shape, dtype, max_batch_size, max_history, id), + f"rnn_state_get_{id}", + ) + for id, (shape, dtype) in enumerate(state_infos) + ] + f_sets = [ + bb.add_func( + RNNState.create_set_func(shape, dtype, max_batch_size, max_history, id), + f"rnn_state_set_{id}", + ) + for id, (shape, dtype) in enumerate(state_infos) + ] + + ret = RNNState( + _expr=rx.call_pure_packed( + "vm.builtin.rnn_state_create", + rx.prim_value(num_hidden_layers), + max_batch_size, + max_history, + f_gets, + f_sets, + list(init_values), + ty_args=[rx.ObjectType()], + ), + _name=name, + ) + return ret + + def get( + self, + layer_id: int, + state_id: int, + shape: Sequence[tirx.Expr], + dtype: str, + ) -> Tensor: + """Get the state of the RNN layer. + + - If there is only one sequence, we can directly use the storage memory, + without copying the data. + - If there are multiple sequences, we need to copy the data to get a contiguous + memory. + + Parameters + ---------- + layer_id : int + The layer id. + state_id : int + The state id. + shape : Sequence[tirx.Expr] + The shape of the state tensor. + dtype: str + The data type of the state tensor. + + Returns + ------- + Tensor + The state tensor, with shape `(batch_size, *state_size)`. + """ + bb = rx.BlockBuilder.current() + + return Tensor( + _expr=bb.emit( + rx.call_dps_packed( + "vm.builtin.rnn_state_get", + [self._expr, layer_id, state_id], + out_ty=rx.TensorType(shape, dtype), + ) + ) + ) + + def set(self, layer_id: int, state_id: int, value: Tensor) -> "RNNState": + """Set the state of the RNN layer. + + Parameters + ---------- + layer_id : int + The layer id. + state_id : int + The state id. + value : Tensor + The state tensor, with shape `(batch_size, *state_size)`. + """ + bb = rx.BlockBuilder.current() + return RNNState( + _expr=bb.emit( + rx.call_pure_packed( + "vm.builtin.rnn_state_set", + self._expr, + rx.prim_value(layer_id), + rx.prim_value(state_id), + value._expr, + ty_args=[rx.ObjectType()], + ) + ), + _name="rnn_state_set", + ) + + @staticmethod + def create_get_func( + shape: Sequence[Union[int, tirx.Var]], + dtype: str, + max_batch_size: Union[int, tirx.Var], + max_history: Union[int, tirx.Var], + state_id: int, + ) -> tirx.PrimFunc: + """Create the get function with given state shape. + + Parameters + ---------- + shape : Sequence[Union[int, tirx.Var]] + The shape of the state tensor. + + dtype: str + The data type of the state tensor. + + max_batch_size : Union[int, tirx.Var] + The maximum batch size. + + max_history : Union[int, tirx.Var] + The maximum history length. + + state_id : int + The id of the state, used for naming the function. + + Returns + ------- + tirx.PrimFunc + The get function. + """ + + def _func_one_dim(): + @T.prim_func(s_tir=True) + def f( + var_storage: T.handle, + var_seq_slot_ids: T.handle, + var_history_slot_ids: T.handle, + var_output: T.handle, + ): + batch_size = T.int32() + T.func_attr({"global_symbol": f"rnn_state_get_{state_id}"}) + + storage = T.match_buffer( + var_storage, (max_batch_size, max_history, shape[0]), dtype + ) + seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32") + history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32") + output = T.match_buffer(var_output, (batch_size, shape[0]), dtype) + + for i in range(batch_size): + for s in range(shape[0]): + with T.sblock("copy"): + vi, vs = T.axis.remap("SS", [i, s]) + seq_id: T.int32 = seq_slot_ids[vi] + history_id: T.int32 = history_slot_ids[vi] + output[vi, vs] = storage[seq_id, history_id, vs] + + return f + + def _func_high_dim(): + # Add a wrapper function to avoid parse the following code when len(shape) = 1 + @T.prim_func(s_tir=True) + def f( + var_storage: T.handle, + var_seq_slot_ids: T.handle, + var_history_slot_ids: T.handle, + var_output: T.handle, + ): + batch_size = T.int32() + T.func_attr({"global_symbol": f"rnn_state_get_{state_id}"}) + + storage = T.match_buffer(var_storage, (max_batch_size, max_history, *shape), dtype) + seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32") + history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32") + output = T.match_buffer(var_output, (batch_size, *shape), dtype) + + for i in range(batch_size): + for s in T.grid(*shape): + with T.sblock("copy"): + vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s]) + seq_id: T.int32 = seq_slot_ids[vi] + history_id: T.int32 = history_slot_ids[vi] + # The following line is equivalent to: + # `output[vi, *vs] = storage[seq_id, history_id, *vs]` + # However, unpacking operator in subscript requires Python 3.11 or newer + T.buffer_store( + output, + T.BufferLoad(storage, [seq_id, history_id, *vs]), + [vi, *vs], + ) + + return f + + return _func_one_dim() if len(shape) == 1 else _func_high_dim() + + @staticmethod + def create_set_func( + shape: Sequence[Union[int, tirx.Var]], + dtype: str, + max_batch_size: Union[int, tirx.Var], + max_history: Union[int, tirx.Var], + state_id: int, + ) -> tirx.PrimFunc: + """Create the set function with given state shape. + + Parameters + ---------- + shape : Sequence[Union[int, tirx.Var]] + The shape of the state tensor. + + dtype: str + The data type of the state tensor. + + max_batch_size : Union[int, tirx.Var] + The maximum batch size. + + max_history : Union[int, tirx.Var] + The maximum history length. + + state_id : int + The id of the state, used for naming the function. + + Returns + ------- + tirx.PrimFunc + The set function. + """ + + def _func_one_dim(): + @T.prim_func(s_tir=True) + def f( + var_storage: T.handle, + var_seq_slot_ids: T.handle, + var_history_slot_ids: T.handle, + var_data: T.handle, + ): + batch_size = T.int32() + T.func_attr({"global_symbol": f"rnn_state_set_{state_id}"}) + + storage = T.match_buffer( + var_storage, (max_batch_size, max_history, shape[0]), dtype + ) + seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32") + history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32") + data = T.match_buffer(var_data, (batch_size, shape[0]), dtype) + + for i in range(batch_size): + for s in range(shape[0]): + with T.sblock("copy"): + vi, vs = T.axis.remap("SS", [i, s]) + seq_id: T.int32 = seq_slot_ids[vi] + history_id: T.int32 = (history_slot_ids[vi] + 1) % T.cast( + max_history, "int32" + ) + storage[seq_id, history_id, vs] = data[vi, vs] + + return f + + def _func_high_dim(): + @T.prim_func(s_tir=True) + def f( + var_storage: T.handle, + var_seq_slot_ids: T.handle, + var_history_slot_ids: T.handle, + var_data: T.handle, + ): + batch_size = T.int32() + T.func_attr({"global_symbol": f"rnn_state_set_{state_id}"}) + + storage = T.match_buffer(var_storage, (max_batch_size, max_history, *shape), dtype) + seq_slot_ids = T.match_buffer(var_seq_slot_ids, (batch_size,), "int32") + history_slot_ids = T.match_buffer(var_history_slot_ids, (batch_size,), "int32") + data = T.match_buffer(var_data, (batch_size, *shape), dtype) + + for i in range(batch_size): + for s in T.grid(*shape): + with T.sblock("copy"): + vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s]) + seq_id: T.int32 = seq_slot_ids[vi] + history_id: T.int32 = (history_slot_ids[vi] + 1) % T.cast( + max_history, "int32" + ) + # The following line is equivalent to: + # `storage[seq_id, history_id, *vs] = data[vi, *vs]` + # However, unpacking operator in subscript requires Python 3.11 or newer + T.buffer_store( + storage, + T.BufferLoad(data, [vi, *vs]), + [seq_id, history_id, *vs], + ) + + return f + + return _func_one_dim() if len(shape) == 1 else _func_high_dim() diff --git a/python/mlc_llm/op/__init__.py b/python/mlc_llm/op/__init__.py new file mode 100644 index 0000000..abbdfe7 --- /dev/null +++ b/python/mlc_llm/op/__init__.py @@ -0,0 +1,15 @@ +"""Extern module for compiler.""" + +from . import moe_matmul, moe_misc +from .attention import attention +from .batch_spec_verify import batch_spec_verify +from .extern import configure, enable, get_store +from .ft_gemm import faster_transformer_dequantize_gemm +from .mrope import ( + MultimodalRotaryEmbedding, + VisionPositionMetadata, + apply_multimodal_rotary_pos_emb, + get_mrope_position_ids, +) +from .pipeline_parallel import pipeline_stage_boundary +from .top_p_pivot import top_p_pivot, top_p_renorm diff --git a/python/mlc_llm/op/attention.py b/python/mlc_llm/op/attention.py new file mode 100644 index 0000000..3574054 --- /dev/null +++ b/python/mlc_llm/op/attention.py @@ -0,0 +1,186 @@ +"""Operators enabled by external modules.""" + +from typing import Optional + +import tvm +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + +from mlc_llm.support import logging + +from . import extern as _extern + +logger = logging.getLogger(__name__) + + +WARN_FLASHINFER_GROUP_SIZE = False +WARN_FLASHINFER_HEAD_DIM = False + + +def attention( + q: nn.Tensor, + k: nn.Tensor, + v: nn.Tensor, + casual_mask: nn.Tensor, + attn_score_scaling_factor: float = 1.0, + qk_dtype: Optional[str] = None, +) -> nn.Tensor: + """Attention with casual mask. + + --- Variables --- + s: sequence length of the current query + t: total sequence length + d: head dimension + h, h_q: number of heads in query + h_kv: number of heads in key and value + b: batch size = 1 + + --- Shapes --- + q: [b, s, h_q, d] + k: [t, h_kv, d] + v: [t, h_kv, d] + o: [1, s, hidden = h_q * d] + + --- Computation --- + + .. code-block:: python + + if h_kv != h_q: + k = k.repeat(h_q // h_kv, axis=1) + v = v.repeat(h_q // h_kv, axis=1) + q -> [b, h, s, d] + k, v -> [b, h, t, d] + attn = q @ k^T / sqrt(d) * attn_score_scaling_factor # [b, h, s, t] + attn = softmax_with_mask(attn, casual_mask, axis=-1) + o = attn @ v # [b, h, s, d] + o -> [b, s, h * d] + + --- Other params --- + qk_dtype: if set, `matmul(Q, K, out_dtype=qk_dtype)`, (otherwise use `q.dtype` as `out_dtype`). + For FlashInfer, if "float32", sets `allow_fp16_qk_reduction` to False; otherwise no effect. + """ + assert q.ndim == 4 and k.ndim in [3, 4] and v.ndim in [3, 4] + b, s, h_q, d = q.shape + t, h_kv, _ = k.shape[-3:] + group_size = h_q // h_kv + + def _fallback(): + from tvm.relax.frontend.nn.llm.kv_cache import ( + _attention_sequence_prefill, + ) + + nonlocal q, k, v, qk_dtype + if k.ndim == 3: + k = op.reshape(k, [b, t, h_kv, d]) + if v.ndim == 3: + v = op.reshape(v, [b, t, h_kv, d]) + if h_kv != h_q: + k = k.repeat(h_q // h_kv, axis=2) + v = v.repeat(h_q // h_kv, axis=2) + + target = tvm.target.Target("cuda") + attn_output, _ = op.tensor_ir_op( + _attention_sequence_prefill( + h_kv=h_kv, + h_q=h_q, + d=d, + dtype=q.dtype, + target=target, + sm_scale=attn_score_scaling_factor / (d**0.5), + ), + "sequence_prefill", + [q, k, v], + [ + Tensor.placeholder([b, s, h_q, d], q.dtype), + Tensor.placeholder([b, s, h_q], q.dtype), + ], + ) + + output = op.reshape(attn_output, shape=(b, s, h_q * d)) + return output + + # FlashInfer Implementation + if ( + _extern.get_store().flashinfer + and attn_score_scaling_factor == 1.0 + and q.dtype == "float16" + and k.dtype == "float16" + and v.dtype == "float16" + ): + if group_size not in [1, 4, 6, 8]: + global WARN_FLASHINFER_GROUP_SIZE + if not WARN_FLASHINFER_GROUP_SIZE: + WARN_FLASHINFER_GROUP_SIZE = True + logger.warning( + "FlashInfer only supports group size in [1, 4, 6, 8], but got %d. Skip and " + "fallback to default implementation.", + group_size, + ) + return _fallback() + if d not in [128]: + global WARN_FLASHINFER_HEAD_DIM + if not WARN_FLASHINFER_HEAD_DIM: + WARN_FLASHINFER_HEAD_DIM = True + logger.warning( + "FlashInfer only supports head_dim in [128], but got %d. Skip and fallback to " + "default implementation.", + d, + ) + return _fallback() + rope_theta = 0.0 + rope_scale = 1.0 + qkv_layout = 0 # "NHD", N for seq_len, H for num_heads, D for head_dim + rotary_mode = 0 # "kNone" + casual = 1 # True + fp16_qk = 1 # True + if qk_dtype == "float32": + fp16_qk = 0 # False + + # 32MB scratchpad + scratch = op.empty([8192 * 1024], dtype="float32") + + def _decode(): + return op.extern( + name="flashinfer.single_decode", + args=[ + q, + k, + v, + scratch, + qkv_layout, + rotary_mode, + rope_scale, + rope_theta, + ], + out=nn.Tensor.placeholder((b, s, h_q * d), dtype="float16"), + ) + + def _prefill(): + return op.extern( + name="flashinfer.single_prefill", + args=[ + q, + k, + v, + scratch, + casual, + qkv_layout, + rotary_mode, + fp16_qk, + rope_scale, + rope_theta, + ], + out=nn.Tensor.placeholder((b, s, h_q * d), dtype="float16"), + ) + + if isinstance(s, int) and s == 1: + func = "decode" + else: + func = "prefill" + return { + "decode": _decode, + "prefill": _prefill, + }[func]() + + # Fallback Implementation + return _fallback() diff --git a/python/mlc_llm/op/batch_matmul.py b/python/mlc_llm/op/batch_matmul.py new file mode 100644 index 0000000..68345a5 --- /dev/null +++ b/python/mlc_llm/op/batch_matmul.py @@ -0,0 +1,44 @@ +"""Batch matmul operators""" + +from typing import Tuple # noqa: UP035 + +from tvm.relax.frontend import nn + +from mlc_llm.op import cutlass +from mlc_llm.quantization.block_scale_quantization import rowwise_group_quant_fp8 + + +def quantized_bmm( + x: nn.Tensor, + w: nn.Tensor, + w_scale: nn.Tensor, + block_size: Tuple[int, int], # noqa: UP006 +) -> nn.Tensor: + """Quantized batch matmul. + Currently only support CUDA backend (by using CUTLASS). + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [b, m, k]. + + w : nn.Tensor + The weight tensor, with shape of [b, n, k] (column major). + + w_scale : nn.Tensor + The scale tensor, with shape of [b, n // block_size[0], k // block_size[1]]. + + block_size : Tuple[int, int] + The block size. + + Returns + ------- + ret : nn.Tensor + The output tensor, with shape of [b, m, n]. + """ + x_fp8, x_scale = rowwise_group_quant_fp8( + x, block_size[1], w.dtype, transpose_scale=True, keep_first_batch_dim=True + ) + return cutlass.fp8_groupwise_scaled_bmm( + x_fp8, x_scale, w, w_scale, block_size, out_dtype=x.dtype + ) diff --git a/python/mlc_llm/op/batch_spec_verify.py b/python/mlc_llm/op/batch_spec_verify.py new file mode 100644 index 0000000..c06ecb6 --- /dev/null +++ b/python/mlc_llm/op/batch_spec_verify.py @@ -0,0 +1,175 @@ +"""Operators for batch verify in speculative decoding.""" + +from tvm.script import tirx as T + +# mypy: disable-error-code="attr-defined,valid-type,name-defined" + + +def batch_spec_verify(vocab_size): + """Batch draft verify function. This function verifies the token tree. + + Before calling the function + + - token_tree_parent_ptr[b] should store the root of the tree + + - draft_probs[node_id, :] stores the prob that samples the correspond tree node + - model_probs[node_id, :] stores the prob that should be used to sample its children + - Please note that the storage convention difference between model_probs and draft_probs + draft_probs was stored on the token node, while model_probs stores on the parent. + This is an intentional design since we can sample different child token with different + proposal draft probabilities, but the ground truth model_prob is unique per parent. + + After calling the function + - token_tree_parent_ptr[b] points to the last token accepted + - There should be a followup sample step that samples from model_probs[token_tree_parent_ptr[b], :] + This token will be appended to the token generated. + + This function will inplace update model_probs if a token was rejected and renormalization is needed. + + Parameters + ---------- + draft_probs: + The draft probability attached to each tree node + + draft_tokens: + The draft token in each node + + model_probs: + The model proability attached to each parent + + token_tree_first_child: + The first child of each tree node, if there is no child, it should be -1 + + token_tree_next_sibling + The next sibling of each tree node, if there is no next sibling, it should be -1 + + uniform_samples + Per node uniform sample used to check rejection + + token_tree_parent_ptr: + Current parent ptr state + """ # noqa: E501 + TX = 1024 + + def _var(dtype="int32"): + return T.sblock_alloc_buffer((1,), dtype, scope="local") + + # fmt: off + @T.prim_func(private=True, s_tir=True) + def _func( + var_draft_probs: T.handle, + var_draft_tokens: T.handle, + var_model_probs: T.handle, + var_token_tree_first_child: T.handle, + var_token_tree_next_sibling: T.handle, + var_uniform_samples: T.handle, + var_token_tree_parent_ptr: T.handle, + ): + """ + [ + blockIdx.x on batch, + threadIdx.x on vocab_size, + for loop over excessive amounts + ] + """ + T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True}) + num_nodes = T.int32() + nbatch = T.int32() + + draft_probs = T.match_buffer(var_draft_probs, (num_nodes, vocab_size), "float32") + draft_tokens = T.match_buffer(var_draft_tokens, (num_nodes,), "int32") + model_probs = T.match_buffer(var_model_probs, (num_nodes, vocab_size), "float32") + token_tree_first_child = T.match_buffer(var_token_tree_first_child, (num_nodes,), "int32") + token_tree_next_sibling = T.match_buffer(var_token_tree_next_sibling, (num_nodes,), "int32") + uniform_samples = T.match_buffer(var_uniform_samples, (num_nodes,), "float32") + token_tree_parent_ptr = T.match_buffer(var_token_tree_parent_ptr, (nbatch,), "int32") + + with T.sblock("kernel"): + child_ptr = _var() + parent_ptr = _var() + child_token = _var() + done = _var("bool") + psum = _var("float32") + t0 = _var("float32") + model_prob_local = _var("float32") + draft_prob_local = _var("float32") + p_child = _var("float32") + q_child = _var("float32") + uniform_sample = _var("float32") + + pred_shared = T.sblock_alloc_buffer((1,), "bool", scope="shared") + pred_local = T.sblock_alloc_buffer((1,), "bool", scope="local") + + for _bx in T.thread_binding(0, nbatch, thread="blockIdx.x"): + for _tx in T.thread_binding(0, TX, thread="threadIdx.x"): + with T.sblock("CTA"): + # batch size + b = T.axis.S(nbatch, _bx) + tx = T.axis.S(TX, _tx) + + parent_ptr[0] = token_tree_parent_ptr[b] + child_ptr[0] = token_tree_first_child[parent_ptr[0]] + done[0] = False + + while T.Not(done[0]): + T.tvm_storage_sync("shared") # ensure all effects last round are visible + if child_ptr[0] == -1: + done[0] = True + T.tvm_storage_sync("shared") # sync before exit + else: + # decide to validate current ptr + if tx == 0: + child_token[0] = draft_tokens[child_ptr[0]] + p_child[0] = model_probs[parent_ptr[0], child_token[0]] + q_child[0] = draft_probs[child_ptr[0], child_token[0]] + uniform_sample[0] = uniform_samples[child_ptr[0]] + pred_shared[0] = p_child[0] >= uniform_sample[0] * q_child[0] # use multiplication to avoid division by zero # noqa: E501 + T.tvm_storage_sync("shared") # make sure all read of model_probs are done # noqa: E501 + pred_local[0] = pred_shared[0] + + # accept the proposal, we move to child + if pred_local[0]: + parent_ptr[0] = child_ptr[0] + child_ptr[0] = token_tree_first_child[child_ptr[0]] + else: + psum[0] = 0.0 + # renormalize probability, predicated by stopped_expansion[b]: + for i in T.serial(T.ceildiv(vocab_size, TX)): + k = T.meta_var(i * TX + tx) + if k < vocab_size: + model_prob_local[0] = model_probs[parent_ptr[0], k] + draft_prob_local[0] = draft_probs[child_ptr[0], k] + model_prob_local[0] = T.max(model_prob_local[0] - draft_prob_local[0], 0.0) # noqa: E501 + psum[0] += model_prob_local[0] + + with T.sblock("block_cross_thread"): + T.reads(psum[0]) + T.writes(t0[0]) + T.attr( + T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]), + "reduce_scope", + T.int32(0), + ) + T.tvm_thread_allreduce(T.uint32(1), psum[0], True, t0[0], tx, dtype="void") # noqa: E501 + + if t0[0] < 1e-7: + # accept the proposal, we move to child + parent_ptr[0] = child_ptr[0] + child_ptr[0] = token_tree_first_child[child_ptr[0]] + else: + # renormalize + for i in T.serial(T.ceildiv(vocab_size, TX)): + k = T.meta_var(i * TX + tx) + if k < vocab_size: + model_prob_local[0] = model_probs[parent_ptr[0], k] + draft_prob_local[0] = draft_probs[child_ptr[0], k] + model_prob_local[0] = T.max(model_prob_local[0] - draft_prob_local[0], 0.0) # noqa: E501 + model_probs[parent_ptr[0], k] = model_prob_local[0] / t0[0] # noqa: E501 + + child_ptr[0] = token_tree_next_sibling[child_ptr[0]] + + if tx == 0: + token_tree_parent_ptr[b] = parent_ptr[0] + # fmt: on + + return _func diff --git a/python/mlc_llm/op/cutlass.py b/python/mlc_llm/op/cutlass.py new file mode 100644 index 0000000..66aa3f1 --- /dev/null +++ b/python/mlc_llm/op/cutlass.py @@ -0,0 +1,375 @@ +"""Operators enabled by external modules.""" + +from typing import Optional, Tuple # noqa: UP035 + +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import op + + +def group_gemm( + x: nn.Tensor, + weight: nn.Tensor, + indptr: nn.Tensor, + scale: Optional[nn.Tensor] = None, + weight_dtype: Optional[str] = None, + out_dtype: Optional[str] = None, +): + """ + Cutlass group gemm operator. + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [m, k]. + + weight : nn.Tensor + The weight tensor, with shape of [num_groups, n, k]. + + indptr : nn.Tensor + The indptr tensor, with shape of [num_groups]. + + scale : Optional[nn.Tensor] + The scale tensor, with shape of [1]. + + weight_dtype: Optional[str] + The data type of the weight tensor. + + out_dtype: Optional[str] + The data type of the output tensor. + + Returns + ------- + nn.Tensor + The output tensor, with shape of [m, n]. + """ + assert x.ndim == 2 + assert weight.ndim == 3 + assert indptr.ndim == 1 + assert weight.shape[0] == indptr.shape[0] + assert indptr.dtype == "int64" + out_dtype = out_dtype if out_dtype else x.dtype + weight_dtype = weight_dtype if weight_dtype else weight.dtype + + if x.dtype == "float8_e5m2" and weight_dtype == "float8_e5m2" and out_dtype == "float16": + func_name = "cutlass.group_gemm_e5m2_e5m2_fp16" + elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e5m2" and out_dtype == "float16": + func_name = "cutlass.group_gemm_e4m3_e5m2_fp16" + elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e4m3fn" and out_dtype == "float16": + func_name = "cutlass.group_gemm_e4m3_e4m3_fp16" + elif (x.dtype == "float16" and weight_dtype == "float16" and out_dtype == "float16") or ( + x.dtype == "bfloat16" and weight_dtype == "bfloat16" and out_dtype == "bfloat16" + ): + func_name = "cutlass.group_gemm" + else: + raise NotImplementedError( + f"Unsupported data type: x={x.dtype}, weight={weight_dtype}, out={out_dtype}" + ) + + if "float8" in x.dtype: + assert scale is not None, "scale is required for float8 input" + + workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace") + + return op.extern( + func_name, + args=[x, weight, indptr, workspace] + ([scale] if scale is not None else []), + out=nn.Tensor.placeholder((x.shape[0], weight.shape[1]), dtype=out_dtype), + ) + + +def fp8_gemm( + x: nn.Tensor, + weight: nn.Tensor, + scale: nn.Tensor, + weight_dtype: Optional[str] = None, + out_dtype: Optional[str] = None, +): + """ + Cutlass fp8 gemm operator. + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [m, k]. + + weight : nn.Tensor + The weight tensor, with shape of [num_groups, n, k]. + + scale : Optional[nn.Tensor] + The scale tensor, with shape of [1]. + + weight_dtype: Optional[str] + The data type of the weight tensor. + + out_dtype: Optional[str] + The data type of the output tensor. + + Returns + ------- + nn.Tensor + The output tensor, with shape of [m, n]. + """ + assert x.ndim >= 2 + assert weight.ndim == 2 + assert scale.ndim == 1 and scale.shape[0] == 1 + out_dtype = out_dtype if out_dtype else x.dtype + weight_dtype = weight_dtype if weight_dtype else weight.dtype + + if x.dtype == "float8_e5m2" and weight_dtype == "float8_e5m2" and out_dtype == "float16": + func_name = "cutlass.gemm_e5m2_e5m2_fp16" + elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e5m2" and out_dtype == "float16": + func_name = "cutlass.gemm_e5m2_e4m3_fp16" + elif x.dtype == "float8_e4m3fn" and weight_dtype == "float8_e4m3fn" and out_dtype == "float16": + func_name = "cutlass.gemm_e4m3_e4m3_fp16" + else: + raise NotImplementedError( + f"Unsupported data type: x={x.dtype}, weight={weight_dtype}, out={out_dtype}" + ) + + workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace") + + return op.extern( + func_name, + args=[x, weight, workspace, scale], + out=nn.Tensor.placeholder((*x.shape[:-1], weight.shape[0]), dtype=out_dtype), + ) + + +def fp8_groupwise_scaled_gemm( + x: nn.Tensor, + x_scale: nn.Tensor, + weight: nn.Tensor, + weight_scale: nn.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + out_dtype: str, +): + """Cutlass block-scale fp8 gemm operator. + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [m, k]. + + x_scale : nn.Tensor + The scale tensor, with shape of [k // block_size, m]. + + weight : nn.Tensor + The weight tensor, with shape of [n, k]. + + weight_scale : nn.Tensor + The scale tensor, with shape of [n // block_size, k // block_size]. + + block_size : Tuple[int, int] + The block size. + + out_dtype : str + The data type of the output tensor. + + Returns + ------- + out : nn.Tensor + The output tensor, with shape of [m, n] and dtype of `out_dtype`. + """ + assert x.ndim >= 2 + assert weight.ndim == 2 + assert x_scale.ndim == x.ndim + assert weight_scale.ndim == weight.ndim + + if block_size[0] != 128 or block_size[1] != 128: + raise ValueError(f"block_size must be (128, 128), but got {block_size}") + if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn": + raise ValueError( + f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}" + ) + if x_scale.dtype != "float32" or weight_scale.dtype != "float32": + raise ValueError( + "x_scale and weight_scale must be float32, but got " + f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}" + ) + if out_dtype not in ["float16", "bfloat16"]: + raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}") + + func_name = "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn" + workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace") + return op.extern( + func_name, + args=[ + x, + weight, + x_scale, + weight_scale, + workspace, + block_size[0], + block_size[1], + ], + out=nn.Tensor.placeholder((*x.shape[:-1], weight.shape[0]), dtype=out_dtype), + ) + + +def fp8_groupwise_scaled_bmm( + x: nn.Tensor, + x_scale: nn.Tensor, + weight: nn.Tensor, + weight_scale: nn.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + out_dtype: str, +): + """Cutlass block-scale fp8 gemm operator. + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [b, m, k]. + + x_scale : nn.Tensor + The scale tensor, with shape of [b, k // block_size, m]. + + weight : nn.Tensor + The weight tensor, with shape of [b, n, k]. + + weight_scale : nn.Tensor + The scale tensor, with shape of [b, n // block_size, k // block_size]. + + block_size : Tuple[int, int] + The block size. + + out_dtype : str + The data type of the output tensor. + + Returns + ------- + out : nn.Tensor + The output tensor, with shape of [m, n] and dtype of `out_dtype`. + """ + assert x.ndim == 3 + assert weight.ndim == 3 + assert x_scale.ndim == x.ndim + assert weight_scale.ndim == weight.ndim + assert x.shape[0] == x_scale.shape[0] == weight.shape[0] == weight_scale.shape[0] + + if block_size[0] != 128 or block_size[1] != 128: + raise ValueError(f"block_size must be (128, 128), but got {block_size}") + if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn": + raise ValueError( + f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}" + ) + if x_scale.dtype != "float32" or weight_scale.dtype != "float32": + raise ValueError( + "x_scale and weight_scale must be float32, but got " + f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}" + ) + if out_dtype not in ["float16", "bfloat16"]: + raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}") + + func_name = "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn" + workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace") + return op.extern( + func_name, + args=[ + x, + weight, + x_scale, + weight_scale, + workspace, + block_size[0], + block_size[1], + ], + out=nn.Tensor.placeholder((x.shape[0], x.shape[1], weight.shape[1]), dtype=out_dtype), + ) + + +def fp8_groupwise_scaled_group_gemm( + x: nn.Tensor, + x_scale: nn.Tensor, + weight: nn.Tensor, + weight_scale: nn.Tensor, + indptr: nn.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + out_dtype: str, +): + """Triton block-scale fp8 group gemm operator. + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [m, k]. + + x_scale : nn.Tensor + The scale tensor, with shape of [m, k // block_size]. + + weight : nn.Tensor + The weight tensor, with shape of [num_experts, n, k]. + + weight_scale : nn.Tensor + The scale tensor, with shape of [num_experts, n // block_size, k // block_size]. + + indptr : nn.Tensor + The indptr tensor of group gemm, with shape of [num_experts + 1,]. + + block_size : Tuple[int, int] + The block size. + + out_dtype : str + The data type of the output tensor. + + Returns + ------- + out : nn.Tensor + The output tensor, with shape of [m, n] and dtype of `out_dtype`. + """ + assert x.ndim >= 2 + assert weight.ndim == 3 + assert x_scale.ndim == x.ndim + assert weight_scale.ndim == weight.ndim + assert x.shape[-1] == weight.shape[2] + assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1] + assert (weight.shape[2] + block_size[1] - 1) // block_size[1] == weight_scale.shape[2] + assert (weight.shape[1] + block_size[0] - 1) // block_size[0] == weight_scale.shape[1] + + if block_size[0] != 128 or block_size[1] != 128: + raise ValueError(f"block_size must be (128, 128), but got {block_size}") + if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn": + raise ValueError( + f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}" + ) + if x_scale.dtype != "float32" or weight_scale.dtype != "float32": + raise ValueError( + "x_scale and weight_scale must be float32, but got " + f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}" + ) + if out_dtype not in ["float16", "bfloat16"]: + raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}") + + num_experts = weight.shape[0] + m = x.shape[0] + for i in range(1, x.ndim - 1): + m *= x.shape[i] + n = weight.shape[1] + k = x.shape[-1] + assert weight_scale.shape[0] == num_experts + assert indptr.ndim == 1 + assert indptr.shape[0] == num_experts + assert indptr.dtype == "int64" + + x_shape = x.shape + if x.ndim > 2: + x = x.reshape(m, k) + x_scale = x_scale.reshape(m, x_scale.shape[-1]) + + func_name = "cutlass.groupwise_scaled_group_gemm_e4m3fn_e4m3fn" + workspace = op.empty((4096 * 1024,), dtype="uint8", name="workspace") + out = op.extern( + func_name, + args=[ + x, + weight, + x_scale, + weight_scale, + indptr, + workspace, + block_size[0], + block_size[1], + ], + out=nn.Tensor.placeholder((m, n), dtype=out_dtype), + ) + return out.reshape(*x_shape[:-1], n) if len(x_shape) > 2 else out diff --git a/python/mlc_llm/op/extern.py b/python/mlc_llm/op/extern.py new file mode 100644 index 0000000..ebf338b --- /dev/null +++ b/python/mlc_llm/op/extern.py @@ -0,0 +1,76 @@ +"""Potential externel modules managed by MLC compilation stack. + +An externl module could contain one or multiple handcrafted kernels, as long as it is provided as +an object file (`.o`), a C++ source file (`.cc`), or a CUDA source file (`.cu`). It can be +integrated into the system pretty smoothly. + +As examples, `flashinfer.py` contains such an example that instructs MLC to compile +"$tvm_home/3rdparty/flashinfer/src/tvm_wrapper.cu" with a specific set of compilation flags and then +link into the generated artifact of MLC LLM. TVM PR #16247 +(https://github.com/apache/tvm/pull/16247/) provides more details of using TVM's +`nn.SourceModule` to integrate C++ and CUDA files, and `nn.ObjectModule` to integrate object files. + +To conveniently use those externel modules, MLC LLM compilation pipeline manages an extra global +singleton `Store: ExternalModuleStore` to store the configured modules. It is supposed to be enabled +before any compilation happens, and configured during a model's `forward` method is invoked. +""" + +import dataclasses +from typing import Optional + +from tvm.target import Target + + +@dataclasses.dataclass +class ExternModuleStore: + """Global store of external modules enabled during compilation.""" + + configured: bool = False + target: Optional[Target] = None + flashinfer: bool = False + faster_transformer: bool = False + cutlass_group_gemm: bool = False + cutlass_gemm: bool = False + + +STORE: ExternModuleStore = ExternModuleStore() +"""Singleton of `ExternModuleStore`.""" + + +def enable(target: Target, flashinfer: bool, faster_transformer: bool, cutlass: bool) -> None: + """Enable external modules. It should be called before any compilation happens.""" + global STORE + cutlass = ( + cutlass + and target.kind.name == "cuda" + and target.attrs.get("arch", "") in ["sm_90a", "sm_100a"] + ) + faster_transformer = False + STORE = ExternModuleStore( + configured=False, + target=target, + flashinfer=flashinfer, + faster_transformer=faster_transformer, + cutlass_group_gemm=cutlass, + cutlass_gemm=cutlass, + ) + + +def get_store() -> ExternModuleStore: + """Get the global store of external modules.""" + return STORE + + +def configure() -> None: + """Configure external modules with extra parameters. It should be called during a model's + `forward` method is invoked. + + Parameters + ---------- + """ + store = get_store() + if store.configured: + return + store.configured = True + if store.flashinfer or store.faster_transformer: + assert store.target.kind.name == "cuda" diff --git a/python/mlc_llm/op/ft_gemm.py b/python/mlc_llm/op/ft_gemm.py new file mode 100644 index 0000000..f31b372 --- /dev/null +++ b/python/mlc_llm/op/ft_gemm.py @@ -0,0 +1,134 @@ +"""Operators enabled by external modules.""" + +import operator +from functools import reduce +from typing import Optional + +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import op + + +def faster_transformer_dequantize_gemm( + x: nn.Tensor, + weight: nn.Tensor, + scale: nn.Tensor, + bias: Optional[nn.Tensor] = None, + activation: Optional[str] = None, + group_size: Optional[int] = None, +): + """ + Faster Transformer dequantize gemm inference with CutlassFpAIntB + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [*m, k]. + + weight : nn.Tensor + The quantized weight data tensor, with shape of [k, n // num_elem_per_storage]. + + scale : nn.Tensor + The quantized weight scale tensor, with shape of [k // group_size, n]. + + bias : Optional[nn.Tensor] + The optional bias for matmul, with shape broadcastable to [*m, n]. + + group_size : Optional[int] + The optional group size. If not set, then using k as group size. + + Returns + ------ + ret: nn.Tensor + The output tensor of deocde matmul, with shape of [*m, n]. + """ + assert x.dtype == "float16" and x.ndim >= 1 + assert weight.ndim == 2 + assert scale.dtype == "float16" and scale.ndim == 2 + assert x.shape[-1] == weight.shape[0], ( + f"Reduction dimension mismatched between x and weight, {x.shape[-1]} vs {weight.shape[0]}." + ) + assert activation in [ + None, + "relu", + "gelu", + "silu", + "identity", + ], "Supported activations are [None, 'identity', 'gelu', 'silu', 'relu']." + activation = activation if activation else "identity" + m = reduce(operator.mul, x.shape[:-1], 1) + k = x.shape[-1] + n = scale.shape[1] + + if not group_size: + group_size = k + + if bias: + assert bias.dtype == "float16" and bias.ndim >= 1 + bias_stride = ( + bias.shape[-1] + if bias and not reduce(operator.mul, bias.shape, 1) == bias.shape[-1] + else 0 + ) + return op.extern( + name="fastertransformer.gemm_fp16_int_bias", + args=[ + x, + weight, + scale, + bias, + activation, + m, + n, + k, + group_size, + bias_stride, + ], + out=nn.Tensor.placeholder((*x.shape[:-1], scale.shape[1]), dtype="float16"), + ) + return op.extern( + name="fastertransformer.gemm_fp16_int", + args=[x, weight, scale, activation, m, n, k, group_size], + out=nn.Tensor.placeholder((*x.shape[:-1], scale.shape[1]), dtype="float16"), + ) + + +def faster_transformer_moe_gemm( + x: nn.Tensor, + weight: nn.Tensor, + total_rows_before: nn.Tensor, +): + """ + Faster Transformer moe gemm inference with CutlassFpAIntB + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [*m, k]. + + weight : nn.Tensor + The weight data tensor, with shape of [num_experts, n, k]. + + total_rows_before : nn.Tensor + The total rows before tensor the current expert, with shape of [num_experts]. This is the + same as the indptr excluding the first zero element. + + Returns + ------ + ret: nn.Tensor + The output tensor of deocde matmul, with shape of [*m, n]. + """ + assert x.dtype == "float16" and x.ndim >= 1 + assert weight.dtype == "float16" and weight.ndim == 3 + assert x.shape[-1] == weight.shape[-1], ( + f"Reduction dimension mismatched between x and weight, {x.shape[-1]} vs {weight.shape[-1]}." + ) + m = reduce(operator.mul, x.shape[:-1], 1) + num_experts = weight.shape[0] + n = weight.shape[1] + k = x.shape[-1] + + return op.extern( + name="fastertransformer.moe_gemm_fp16_fp16", + args=[x, weight, total_rows_before, m, n, k, num_experts], + out=nn.Tensor.placeholder((*x.shape[:-1], n), dtype="float16"), + ) diff --git a/python/mlc_llm/op/moe_matmul.py b/python/mlc_llm/op/moe_matmul.py new file mode 100644 index 0000000..62d9ceb --- /dev/null +++ b/python/mlc_llm/op/moe_matmul.py @@ -0,0 +1,768 @@ +"""Mixture of Experts operators""" + +from typing import Literal, Optional, Tuple # noqa: UP035 + +from tvm import DataType, DataTypeCode, s_tir, tirx +from tvm.relax.frontend.nn import Tensor, op +from tvm.script import tirx as T + +# mypy: disable-error-code="attr-defined,valid-type,name-defined" + + +def gemv(x: Tensor, w: Tensor, indptr: Tensor) -> Tensor: + """GEMV for project-in (e1-e3) or project-out (e2) in MLP. + + Parameters + ---------- + x : Tensor + For project-in, the input tensor of shape (1, in_features); and for project-out, the input + shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated + experts per token. + + w : Tensor + The weight tensor of shape (local_experts, out_features, in_features), where `local_experts` + is the total number of experts. + + indptr : Tensor + The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the + number of activated experts per token. + + Returns + ------- + out : Tensor + The output tensor of shape (experts_per_tok, out_features), where `experts_per_tok` is the + number of activated experts per token. + """ + (local_experts, out_features, in_features), dtype = w.shape, w.dtype + _, experts_per_tok = indptr.shape + x_leading_dim, _ = x.shape + + def access_x(x, e, j): + return x[0, j] if x_leading_dim == 1 else x[e, j] + + # NOTE: Currently it assumes x.dtype == w.dtype, but the constraint can be relaxed easily. + assert w.shape == [local_experts, out_features, in_features] and w.dtype == dtype + assert x.shape == [x_leading_dim, in_features] and x.dtype == dtype + assert indptr.shape == [1, experts_per_tok] and indptr.dtype == "int32" + assert x_leading_dim in [1, experts_per_tok] + + @T.prim_func(private=True, s_tir=True) + def _func( + x: T.Buffer((x_leading_dim, in_features), dtype), + w: T.Buffer((local_experts, out_features, in_features), dtype), + indptr: T.Buffer((1, experts_per_tok), "int32"), + o: T.Buffer((experts_per_tok, out_features), dtype), + ): + T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable + for e in T.thread_binding(experts_per_tok, thread="blockIdx.y"): + with T.sblock("gemv_o"): + e = T.axis.spatial(experts_per_tok, e) + T.reads(x[:, :], w[indptr[0, e], :, :], indptr[0, e]) + T.writes(o[e, :]) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("gemv"): + i, j = T.axis.remap("SR", [i1, i2]) + with T.init(): + o[e, i] = T.cast(T.float16(0), dtype) + o[e, i] += access_x(x, e, j) * w[indptr[0, e], i, j] + + return op.tensor_ir_op( + _func, + "moe_gemv", + args=[x, w, indptr], + out=Tensor.placeholder([experts_per_tok, out_features], dtype), + ) + + +def dequantize_gemv( + x: Tensor, + w: Tensor, + scale: Tensor, + indptr: Tensor, + quantize_dtype: str, + group_size: int, +) -> Tensor: + """GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized. + It needs to be dequantized before the GEMV computation. + + Parameters + ---------- + x : Tensor + For project-in, the input tensor of shape (1, in_features); and for project-out, the input + shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated + experts per token. + + w : Tensor + The quantized weight tensor of shape (local_experts, out_features, in_features // n), + where n is the number of elements per storage dtype, e.g. if the storage dtype is uint32, + and the quantize dtype is int4, then n is 8. + `local_experts` is the total number of experts including activated and non-active ones. + + scale : Tensor + The scale tensor of shape (local_experts, out_features, in_features // group_size), where + `local_experts` is the total number of experts including activated and non-active ones. + + indptr : Tensor + The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the + number of activated experts per token. + + quantize_dtype : str + The quantize dtype of the weight tensor, which is usually int3, int4 or fp8, etc. + + group_size : int + The number of elements in each quantization group, e.g. 32 or 128. + + Returns + ------- + out : Tensor + The output tensor of shape (experts_per_tok, out_features), where `experts_per_tok` is the + number of activated experts per token. + """ + (x_leading_dim, in_features), model_dtype = x.shape, x.dtype + (local_experts, out_features, _), storage_dtype = w.shape, w.dtype + _, experts_per_tok = indptr.shape + quantize_dtype_bits = DataType(quantize_dtype).bits + num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits + num_group = (in_features + group_size - 1) // group_size + num_storage = group_size // num_elem_per_storage * num_group + + def _dequantize(w, s, e, i, j): + tir_bin_mask = tirx.const((2**quantize_dtype_bits) - 1, storage_dtype) + tir_max_int = tirx.const((2 ** (quantize_dtype_bits - 1)) - 1, model_dtype) + w = w[e, i, j // num_elem_per_storage] + s = s[e, i, j // group_size] + shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype) + w = tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype(model_dtype) + return (w - tir_max_int) * s + + def access_x(x, e, j): + return x[0, j] if x_leading_dim == 1 else x[e, j] + + assert x.shape == [x_leading_dim, in_features] and x.dtype == model_dtype + assert w.shape == [local_experts, out_features, num_storage] and w.dtype == storage_dtype + assert scale.shape == [local_experts, out_features, num_group] and scale.dtype == model_dtype + assert indptr.shape == [1, experts_per_tok] and indptr.dtype == "int32" + assert x_leading_dim in [1, experts_per_tok] + + @T.prim_func(private=True, s_tir=True) + def _func( + x: T.Buffer((x_leading_dim, in_features), model_dtype), + w: T.Buffer((local_experts, out_features, num_storage), storage_dtype), + scale: T.Buffer((local_experts, out_features, num_group), model_dtype), + indptr: T.Buffer((1, experts_per_tok), "int32"), + o: T.Buffer((experts_per_tok, out_features), model_dtype), + ): + T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable + for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"): + with T.sblock("gemv_o"): + e = T.axis.spatial(experts_per_tok, expert_id) + y = T.sblock_alloc_buffer((out_features, in_features), model_dtype) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("dequantize"): + i, j = T.axis.remap("SS", [i1, i2]) + y[i, j] = _dequantize(w, scale, indptr[0, e], i, j) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("gemv"): + i, j = T.axis.remap("SR", [i1, i2]) + with T.init(): + o[e, i] = T.cast(T.float16(0), model_dtype) + o[e, i] += access_x(x, e, j) * y[i, j] + + return op.tensor_ir_op( + _func, + "moe_dequantize_gemv", + args=[x, w, scale, indptr], + out=Tensor.placeholder([experts_per_tok, out_features], model_dtype), + ) + + +def dequantize_float8_gemv( + x: Tensor, + w: Tensor, + scale: Optional[Tensor], + indptr: Tensor, + quantize_dtype: Literal["float8_e5m2", "float8_e4m3fn"], +) -> Tensor: + """GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized in + fp8 e5m2 or e4m3. It needs to be dequantized before the GEMV computation. + + Parameters + ---------- + x : Tensor + For project-in, the input tensor of shape (1, in_features); and for project-out, the input + shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated + experts per token. + + w : Tensor + The quantized weight tensor of shape (local_experts, out_features, in_features) + + scale : Optional[Tensor] + The optional scale tensor of shape (1,) + + indptr : Tensor + The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the + number of activated experts per token. + + quantize_dtype : Literal["float8_e5m2", "float8_e4m3fn"] + The quantize dtype of the weight tensor, which is either float8_e5m2 or float8_e4m3fn. + """ + (x_leading_dim, in_features), model_dtype = x.shape, x.dtype + (local_experts, out_features, _), storage_dtype = w.shape, w.dtype + _, experts_per_tok = indptr.shape + quantize_dtype_bits = DataType(quantize_dtype).bits + num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits + num_storage = tirx.ceildiv(in_features, num_elem_per_storage) + + def _dequantize(w, s, e, i, j): + if num_elem_per_storage == 1: + w = tirx.reinterpret(quantize_dtype, w[e, i, j]) + else: + assert DataType(storage_dtype).type_code == DataTypeCode.UINT + tir_bin_mask = tirx.const((2**quantize_dtype_bits) - 1, storage_dtype) + w = w[e, i, j // num_elem_per_storage] + shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype) + w = tirx.reinterpret( + quantize_dtype, + tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype("uint8"), + ) + w = w.astype(model_dtype) + if s is not None: + w = w * s[0] + return w + + def access_x(x, e, j): + return x[0, j] if x_leading_dim == 1 else x[e, j] + + @T.prim_func(private=True, s_tir=True) + def _func_with_scale( + x: T.Buffer((x_leading_dim, in_features), model_dtype), + w: T.Buffer((local_experts, out_features, num_storage), storage_dtype), + scale: T.Buffer((1,), "float32"), + indptr: T.Buffer((1, experts_per_tok), "int32"), + o: T.Buffer((experts_per_tok, out_features), model_dtype), + ): + T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable + for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"): + with T.sblock("gemv_o"): + e = T.axis.spatial(experts_per_tok, expert_id) + y = T.sblock_alloc_buffer((out_features, in_features), model_dtype) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("dequantize"): + i, j = T.axis.remap("SS", [i1, i2]) + y[i, j] = _dequantize(w, scale, indptr[0, e], i, j) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("gemv"): + i, j = T.axis.remap("SR", [i1, i2]) + with T.init(): + o[e, i] = T.cast(T.float16(0), model_dtype) + o[e, i] += access_x(x, e, j) * y[i, j] + + @T.prim_func(private=True, s_tir=True) + def _func_without_scale( + x: T.Buffer((x_leading_dim, in_features), model_dtype), + w: T.Buffer((local_experts, out_features, num_storage), storage_dtype), + indptr: T.Buffer((1, experts_per_tok), "int32"), + o: T.Buffer((experts_per_tok, out_features), model_dtype), + ): + T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable + for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"): + with T.sblock("gemv_o"): + e = T.axis.spatial(experts_per_tok, expert_id) + y = T.sblock_alloc_buffer((out_features, in_features), model_dtype) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("dequantize"): + i, j = T.axis.remap("SS", [i1, i2]) + y[i, j] = _dequantize(w, None, indptr[0, e], i, j) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("gemv"): + i, j = T.axis.remap("SR", [i1, i2]) + with T.init(): + o[e, i] = T.cast(T.float16(0), model_dtype) + o[e, i] += access_x(x, e, j) * y[i, j] + + if scale is not None: + return op.tensor_ir_op( + _func_with_scale, + "moe_dequantize_gemv", + args=[x, w, scale, indptr], + out=Tensor.placeholder([experts_per_tok, out_features], model_dtype), + ) + return op.tensor_ir_op( + _func_without_scale, + "moe_dequantize_gemv", + args=[x, w, indptr], + out=Tensor.placeholder([experts_per_tok, out_features], model_dtype), + ) + + +def dequantize_block_scale_float8_gemv( + x: Tensor, + w: Tensor, + w_scale: Tensor, + expert_indices: Tensor, + block_size: Tuple[int, int], # noqa: UP006 + out_dtype: str, +) -> Tensor: + """GEMV for project-in (e1-e3) or project-out (e2) in MLP but the weight is quantized in + fp8 e5m2 or e4m3. It needs to be dequantized before the GEMV computation. + + Parameters + ---------- + x : Tensor + For project-in, the input tensor of shape (1, in_features); and for project-out, the input + shape is (experts_per_tok, in_features), where `experts_per_tok` is the number of activated + experts per token. + + w : Tensor + The quantized weight tensor of shape (local_experts, out_features, in_features) + + w_scale : Tensor + The scale tensor of shape + (local_experts, out_features // block_size[0], in_features // block_size[1]) + + indptr : Tensor + The index pointer tensor of shape (1, experts_per_tok), where `experts_per_tok` is the + number of activated experts per token. + + block_size : Tuple[int, int] + The block size of the weight tensor. + + out_dtype : str + The output dtype of the GEMV computation. + """ + x_leading_dim, in_features = x.shape + local_experts, out_features, k = w.shape + _, experts_per_tok = expert_indices.shape + model_dtype = x.dtype + quantize_dtype = w.dtype + + assert out_features % block_size[0] == 0 + assert k % block_size[1] == 0 + + def _dequantize(w, s, e, i, j): + return w[e, i, j].astype(model_dtype) * s[e, i // block_size[0], j // block_size[1]].astype( + model_dtype + ) + + def load_x(x, e, j): + return x[0, j] if x_leading_dim == 1 else x[e, j] + + @T.prim_func(private=True, s_tir=True) + def _func( + x: T.Buffer((x_leading_dim, in_features), model_dtype), + w: T.Buffer((local_experts, out_features, k), quantize_dtype), + w_scale: T.Buffer( + (local_experts, out_features // block_size[0], k // block_size[1]), + "float32", + ), + expert_indices: T.Buffer((1, experts_per_tok), "int32"), + o: T.Buffer((experts_per_tok, out_features), out_dtype), + ): + T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable + for expert_id in T.thread_binding(experts_per_tok, thread="blockIdx.y"): + with T.sblock("gemv_o"): + e = T.axis.spatial(experts_per_tok, expert_id) + y = T.sblock_alloc_buffer((out_features, in_features), model_dtype) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("dequantize"): + i, j = T.axis.remap("SS", [i1, i2]) + y[i, j] = _dequantize(w, w_scale, expert_indices[0, e], i, j) + for i1, i2 in T.grid(out_features, in_features): + with T.sblock("gemv"): + i, j = T.axis.remap("SR", [i1, i2]) + with T.init(): + o[e, i] = T.cast(T.float16(0), out_dtype) + o[e, i] += (load_x(x, e, j) * y[i, j]).astype(out_dtype) + + return op.tensor_ir_op( + _func, + "moe_dequantize_gemv", + args=[x, w, w_scale, expert_indices], + out=Tensor.placeholder([experts_per_tok, out_features], out_dtype), + ) + + +def group_gemm(x: Tensor, w: Tensor, indptr: Tensor): + """Group GEMM in MoE models. + + Parameters + ---------- + x : Tensor + Input tensor of shape (batch_size, in_features), where `batch_size` could be dynamic shape. + + w : Tensor + Weight tensor of shape (num_local_experts, out_features, in_features). + `w[i, :, :]` is the weight matrix for the `i`-th local expert. + + indptr : Tensor + Index pointer tensor of shape (num_local_experts + 1, ). + `x[indptr[a] : indptr[a + 1]]` is the input for the `i`-th local expert. + + Returns + ------- + out : Tensor + Output tensor of shape (batch_size, out_features). + """ + # NOTE: Currently it assumes x.dtype == w.dtype, but the constraint can be relaxed easily. + (num_local_experts, out_features, in_features), dtype = w.shape, w.dtype + + assert x.shape[1:] == [in_features] and x.dtype == dtype + assert indptr.shape == [num_local_experts + 1] and indptr.dtype == "int32" + + Ne, N, K = num_local_experts, out_features, in_features + BLK_M, BLK_N, BLK_K = 8, 128, 32 + TX, TY, CTA_COUNT = 8, 32, 1024 + VEC_X, VEC_W, VEC_O, VEC_DOT = 1, 1, 1, 1 + UNROLL = 64 + STORAGE_ALIGN = False + assert BLK_K % 8 == 0 + tiles_per_row = (N + BLK_N - 1) // BLK_N + zero = tirx.const(0, dtype) + + @T.prim_func(private=True, s_tir=True) + def _func( + var_x: T.handle, + var_w: T.handle, + var_indptr: T.handle, + var_o: T.handle, + ): + T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True}) + B = T.int32() + X = T.match_buffer(var_x, (B, K), dtype) + W = T.match_buffer(var_w, (Ne, N, K), dtype) + indptr = T.match_buffer(var_indptr, (Ne + 1,), "int32") + out = T.match_buffer(var_o, (B, N), dtype) + + for _bx in T.thread_binding(CTA_COUNT, thread="blockIdx.x"): + with T.sblock("CTA"): + bx = T.axis.spatial(CTA_COUNT, _bx) + T.reads(indptr[:], X[:, :], W[:, :, :]) + T.writes(out[:, :]) + sum = T.sblock_alloc_buffer((2,), "int32", scope="local") + row = T.sblock_alloc_buffer((2,), "int32", scope="local") + cur_e = T.sblock_alloc_buffer((1,), "int32", scope="local") + tile_id = T.sblock_alloc_buffer((1,), "int32", scope="local") + sum[0] = 0 + sum[1] = T.ceildiv(indptr[1] - indptr[0], BLK_M) * tiles_per_row + row[0] = 0 + row[1] = indptr[1] - indptr[0] + cur_e[0] = 0 + tile_id[0] = bx + while T.tvm_thread_invariant(cur_e[0] < Ne): + # move to the current group + while sum[1] <= tile_id[0] and cur_e[0] < Ne: + cur_e[0] += 1 + if cur_e[0] < Ne: + e: T.int32 = cur_e[0] + delta: T.int32 = indptr[e + 1] - indptr[e] + sum[0] = sum[1] + sum[1] += T.ceildiv(delta, BLK_M) * tiles_per_row + row[0] = row[1] + row[1] += delta + # sync threads to make sure all threads have the same tile position + T.tvm_storage_sync("shared") + if T.tvm_thread_invariant(cur_e[0] < Ne): + # fetch current tile position + e: T.int32 = cur_e[0] + num_tiles: T.int32 = tile_id[0] - sum[0] + m_offset: T.int32 = BLK_M * T.floordiv(num_tiles, tiles_per_row) + row[0] + n_offset: T.int32 = BLK_N * T.floormod(num_tiles, tiles_per_row) + with T.sblock("gemm"): + T.reads( + row[1], + X[m_offset : m_offset + BLK_M, :], + W[e, n_offset : n_offset + BLK_N, :], + ) + T.writes( + out[ + m_offset : m_offset + BLK_M, + n_offset : n_offset + BLK_N, + ] + ) + X_tile = T.sblock_alloc_buffer((BLK_M, K), dtype, scope="shared") + W_tile = T.sblock_alloc_buffer((BLK_N, K), dtype, scope="shared") + O_tile = T.sblock_alloc_buffer((BLK_M, BLK_N), dtype, scope="local") + for a0, a1 in T.grid(BLK_M, K): + with T.sblock("X_shared"): + i, j = T.axis.remap("SS", [a0, a1]) + X_tile[i, j] = T.if_then_else( + m_offset + i < row[1], + X[m_offset + i, j], + zero, + ) + for a0, a1 in T.grid(BLK_N, K): + with T.sblock("W_shared"): + i, j = T.axis.remap("SS", [a0, a1]) + W_tile[i, j] = T.if_then_else( + n_offset + i < N, + W[e, n_offset + i, j], + zero, + ) + for a0, a1, a2 in T.grid(BLK_M, BLK_N, K): + with T.sblock("compute"): + i, j, k = T.axis.remap("SSR", [a0, a1, a2]) + with T.init(): + O_tile[i, j] = zero + O_tile[i, j] += X_tile[i, k] * W_tile[j, k] + for a0, a1 in T.grid(BLK_M, BLK_N): + with T.sblock("store"): + i, j = T.axis.remap("SS", [a0, a1]) + if m_offset + i < row[1] and n_offset + j < N: + out[m_offset + i, n_offset + j] = O_tile[i, j] + # move to next tile + tile_id[0] += CTA_COUNT + + def _schedule(): + sch = s_tir.Schedule(_func) + + def _cooperative_fetch(block, vec_len): + num_loops = len(sch.get_loops(block)) + sch.compute_at(block, ko, preserve_unit_loops=True) + loops = sch.get_loops(block)[-num_loops:] + ty, tx, _, vec = sch.split( + sch.fuse(*loops), + factors=[TY, TX, None, vec_len], + ) + sch.vectorize(vec) + sch.bind(ty, "threadIdx.y") + sch.bind(tx, "threadIdx.x") + if STORAGE_ALIGN: + sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len) + return block + + main_block = sch.get_sblock("compute") + x, y, k = sch.get_loops(main_block) + ty, yi = sch.split(y, [TY, None]) + tx, xi, vec_c = sch.split(x, [TX, None, VEC_DOT]) + ko, ki = sch.split(k, factors=[None, BLK_K]) + sch.reorder(ty, tx, ko, ki, yi, xi, vec_c) + sch.bind(ty, "threadIdx.y") + sch.bind(tx, "threadIdx.x") + sch.vectorize(vec_c) + if UNROLL > 0: + sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=UNROLL) + sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1) + l2g = sch.get_sblock("store") + sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True) + _, v = sch.split(sch.get_loops(l2g)[-1], [None, VEC_O]) + sch.vectorize(v) + _cooperative_fetch(sch.get_sblock("X_shared"), vec_len=VEC_X) + _cooperative_fetch(sch.get_sblock("W_shared"), vec_len=VEC_W) + sch.decompose_reduction(main_block, ko) + return sch.mod["main"] + + return op.tensor_ir_op( + _schedule(), + "group_gemm", + args=[x, w, indptr], + out=Tensor.placeholder([x.shape[0], out_features], dtype), + ) + + +def dequantize_group_gemm( + x: Tensor, + w: Tensor, + scale: Tensor, + indptr: Tensor, + quantize_dtype: str, + indptr_dtype: str, + group_size: int, +): + """Group GEMM in MoE models but the weight is quantized. + + Parameters + ---------- + x : Tensor + Input tensor of shape (batch_size, in_features), where `batch_size` could be dynamic shape. + + w : Tensor + Weight tensor of shape (num_local_experts, out_features, in_features // n), where n is the + number of elements per storage dtype, e.g. if the storage dtype is uint32, and the quantize + dtype is int4, then n is 8. + + scale : Tensor + The scale tensor of shape (num_local_experts, out_features, in_features // group_size). + + indptr : Tensor + Index pointer tensor of shape (num_local_experts + 1, ). `x[indptr[a] : indptr[a + 1]]` is + the input for the `i`-th local expert. + + group_size : int + The number of elements in each quantization group, e.g. 32 or 128. + + quantize_dtype : str + The quantize dtype of the weight tensor, which is usually int3, int4 or fp8, etc. + + indptr_dtype : str + The dtype of the index pointer tensor, which can be int32 or int64. + + Returns + ------- + out : Tensor + Output tensor of shape (batch_size, out_features). + """ + (_, in_features), model_dtype = x.shape, x.dtype + (num_local_experts, out_features, _), storage_dtype = w.shape, w.dtype + quantize_dtype_bits = DataType(quantize_dtype).bits + num_elem_per_storage = DataType(storage_dtype).bits // quantize_dtype_bits + num_group = (in_features + group_size - 1) // group_size + num_storage = group_size // num_elem_per_storage * num_group + + def _dequantize(w, s, e, i, j): + tir_bin_mask = tirx.const((1 << quantize_dtype_bits) - 1, storage_dtype) + tir_max_int = tirx.const((2 ** (quantize_dtype_bits - 1)) - 1, model_dtype) + w = w[e, i, j // num_elem_per_storage] + s = s[e, i, j // group_size] + shift = (j % num_elem_per_storage * quantize_dtype_bits).astype(storage_dtype) + w = tirx.bitwise_and(tirx.shift_right(w, shift), tir_bin_mask).astype(model_dtype) + return (w - tir_max_int) * s + + Ne, N, K = num_local_experts, out_features, in_features + BLK_M, BLK_N, BLK_K = 8, 128, 32 + TX, TY, CTA_COUNT = 8, 32, 1024 + VEC_X, VEC_W, VEC_O, VEC_DOT = 1, 1, 1, 1 + UNROLL = 64 + STORAGE_ALIGN = False + assert BLK_K % 8 == 0 + tiles_per_row = (N + BLK_N - 1) // BLK_N + zero = tirx.const(0, model_dtype) + if indptr_dtype == "int64": + indptr = op.pad(indptr, [1, 0], "constant", 0) + + @T.prim_func(private=True, s_tir=True) + def _func( + var_x: T.handle, + w: T.Buffer((Ne, N, num_storage), storage_dtype), + scale: T.Buffer((Ne, N, num_group), model_dtype), + indptr: T.Buffer((Ne + 1,), indptr_dtype), + var_o: T.handle, + ): + T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True}) + B = T.int32() + X = T.match_buffer(var_x, (B, K), model_dtype) + out = T.match_buffer(var_o, (B, N), model_dtype) + for _bx in T.thread_binding(CTA_COUNT, thread="blockIdx.x"): + with T.sblock("CTA"): + bx = T.axis.spatial(CTA_COUNT, _bx) + T.reads(X[:, :], w[:, :, :], scale[:, :, :], indptr[:]) + T.writes(out[:, :]) + sum = T.sblock_alloc_buffer((2,), indptr_dtype, scope="local") + row = T.sblock_alloc_buffer((2,), indptr_dtype, scope="local") + cur_e = T.sblock_alloc_buffer((1,), indptr_dtype, scope="local") + tile_id = T.sblock_alloc_buffer((1,), indptr_dtype, scope="local") + sum[0] = 0 + sum[1] = T.ceildiv(indptr[1] - indptr[0], BLK_M) * tiles_per_row + row[0] = 0 + row[1] = indptr[1] - indptr[0] + cur_e[0] = 0 + tile_id[0] = bx + while T.tvm_thread_invariant(cur_e[0] < Ne): + # move to the current group + while sum[1] <= tile_id[0] and cur_e[0] < Ne: + cur_e[0] += 1 + if cur_e[0] < Ne: + e = cur_e[0] + delta = indptr[e + 1] - indptr[e] + sum[0] = sum[1] + sum[1] += T.ceildiv(delta, BLK_M) * tiles_per_row + row[0] = row[1] + row[1] += delta + # sync threads to make sure all threads have the same tile position + T.tvm_storage_sync("shared") + if T.tvm_thread_invariant(cur_e[0] < Ne): + # fetch current tile position + e = cur_e[0] + num_tiles = tile_id[0] - sum[0] + m_offset = T.floordiv(num_tiles, tiles_per_row) * BLK_M + row[0] + n_offset = T.floormod(num_tiles, tiles_per_row) * BLK_N + with T.sblock("gemm"): + T.reads( + row[1], + X[m_offset : m_offset + BLK_M, :], + w[e, n_offset : n_offset + BLK_N, :], + scale[e, n_offset : n_offset + BLK_N, :], + ) + T.writes( + out[ + m_offset : m_offset + BLK_M, + n_offset : n_offset + BLK_N, + ] + ) + X_tile = T.sblock_alloc_buffer((BLK_M, K), model_dtype, scope="shared") + W_tile = T.sblock_alloc_buffer((BLK_N, K), model_dtype, scope="shared") + O_tile = T.sblock_alloc_buffer((BLK_M, BLK_N), "float32", scope="local") + for a0, a1 in T.grid(BLK_M, K): + with T.sblock("X_shared"): + i, j = T.axis.remap("SS", [a0, a1]) + X_tile[i, j] = T.if_then_else( + m_offset + i < row[1], + X[m_offset + i, j], + zero, + ) + for a0, a1 in T.grid(BLK_N, K): + with T.sblock("W_shared"): + i, j = T.axis.remap("SS", [a0, a1]) + W_tile[i, j] = T.if_then_else( + n_offset + i < N, + _dequantize(w, scale, e, n_offset + i, j), + zero, + ) + for a0, a1, a2 in T.grid(BLK_M, BLK_N, K): + with T.sblock("compute"): + i, j, k = T.axis.remap("SSR", [a0, a1, a2]) + with T.init(): + O_tile[i, j] = zero + O_tile[i, j] += X_tile[i, k] * W_tile[j, k] + for a0, a1 in T.grid(BLK_M, BLK_N): + with T.sblock("store"): + i, j = T.axis.remap("SS", [a0, a1]) + if m_offset + i < row[1] and n_offset + j < N: + out[m_offset + i, n_offset + j] = O_tile[i, j] + # move to next tile + tile_id[0] += CTA_COUNT + + def _schedule(): + sch = s_tir.Schedule(_func) + + def _cooperative_fetch(block, vec_len): + num_loops = len(sch.get_loops(block)) + sch.compute_at(block, ko, preserve_unit_loops=True) + loops = sch.get_loops(block)[-num_loops:] + ty, tx, _, vec = sch.split( + sch.fuse(*loops), + factors=[TY, TX, None, vec_len], + ) + sch.vectorize(vec) + sch.bind(ty, "threadIdx.y") + sch.bind(tx, "threadIdx.x") + if STORAGE_ALIGN: + sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len) + return block + + main_block = sch.get_sblock("compute") + x, y, k = sch.get_loops(main_block) + ty, yi = sch.split(y, [TY, None]) + tx, xi, vec_c = sch.split(x, [TX, None, VEC_DOT]) + ko, ki = sch.split(k, factors=[None, BLK_K]) + sch.reorder(ty, tx, ko, ki, yi, xi, vec_c) + sch.bind(ty, "threadIdx.y") + sch.bind(tx, "threadIdx.x") + sch.vectorize(vec_c) + if UNROLL > 0: + sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=UNROLL) + sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1) + l2g = sch.get_sblock("store") + sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True) + _, v = sch.split(sch.get_loops(l2g)[-1], [None, VEC_O]) + sch.vectorize(v) + _cooperative_fetch(sch.get_sblock("X_shared"), vec_len=VEC_X) + _cooperative_fetch(sch.get_sblock("W_shared"), vec_len=VEC_W) + sch.decompose_reduction(main_block, ko) + return sch.mod["main"] + + return op.tensor_ir_op( + _schedule(), + "dequantize_group_gemm", + args=[x, w, scale, indptr], + out=Tensor.placeholder([x.shape[0], out_features], model_dtype), + ) diff --git a/python/mlc_llm/op/moe_misc.py b/python/mlc_llm/op/moe_misc.py new file mode 100644 index 0000000..55e9872 --- /dev/null +++ b/python/mlc_llm/op/moe_misc.py @@ -0,0 +1,645 @@ +"""Mixture of Experts operators""" + +from functools import reduce +from typing import Literal, Optional, Tuple, Union # noqa: UP035 + +import numpy as np +from tvm import te, tirx +from tvm.relax.frontend.nn import IntExpr, Tensor, op +from tvm.script import tirx as T + +# mypy: disable-error-code="attr-defined,name-defined" + + +def moe_sum(x: Tensor, dim: int) -> Tensor: + """Compute the sum of the input tensor along the given axis. It is specialized for the MoE + case where `x.ndim == 3` and `x.shape[1] == num_experts_per_tok (which is 2)`. + """ + + if x.shape[1] == 1: + return x.reshape(x.shape[0], x.shape[2]) + + if x.ndim == 3 and x.shape[1] == 2: + return op.tensor_expr_op( + lambda x: te.compute( + (x.shape[0], x.shape[2]), + lambda i, j: x[i, 0, j] + x[i, 1, j], + name="sum_2", + ), + "sum", + args=[x], + ) + return op.sum(x, axis=dim) + + +def _gating_topk_init_local_top_k(k_val, dtype, local_top_k, local_top_k_index): + for t in range(k_val): + T.buffer_store(local_top_k, T.min_value(dtype), indices=[t]) + for t in range(k_val): + T.buffer_store(local_top_k_index, t, indices=[-1]) + + +def _gating_topk_process_value(k_val, x, local_top_k, local_top_k_index, vi, vk): + if_frames = [T.If(x[vi, vk] > local_top_k[i]) for i in range(k_val)] + then_frames = [T.Then() for _ in range(k_val)] + else_frames = [T.Else() for _ in range(k_val - 1)] + for i in range(k_val): + if_frames[i].__enter__() + with then_frames[i]: + for j in range(k_val - 1, i, -1): + T.buffer_store(local_top_k, local_top_k[j - 1], indices=[j]) + T.buffer_store(local_top_k_index, local_top_k_index[j - 1], indices=[j]) + T.buffer_store(local_top_k, x[vi, vk], indices=[i]) + T.buffer_store(local_top_k_index, vk, indices=[i]) + if i != k_val - 1: + else_frames[i].__enter__() + + for i in range(k_val - 1, -1, -1): + if i != k_val - 1: + else_frames[i].__exit__(None, None, None) + if_frames[i].__exit__(None, None, None) + + +def gating_topk(scores: Tensor, k: int) -> Tuple[Tensor, Tensor]: # noqa: UP006 + """Compute the top-k experts and their scores. + + Parameters + ---------- + scores : Tensor + The input tensor with shape [batch_size, num_local_experts]. + + k : int + The number of top elements to be selected, which is `num_experts_per_tok` in MoE. + + Returns + ------- + expert_weights: Tensor + The top-k expert scores with shape [batch_size, k]. + + expert_indices: Tensor + The top-k expert indices with shape [batch_size, k]. + """ + (batch_size, num_local_experts), dtype = scores.shape, scores.dtype + index_dtype = "int32" + + TX = 1024 + + def _get_topk_func(k_val: int): + @T.prim_func(private=True, s_tir=True) + def topk_func( + var_x: T.handle, + var_out: T.handle, + var_out_index: T.handle, + ) -> None: + T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": True}) + batch_size = T.int64() + x = T.match_buffer(var_x, (batch_size, num_local_experts), dtype) + out = T.match_buffer(var_out, (batch_size, k_val), dtype) + out_index = T.match_buffer(var_out_index, (batch_size, k_val), index_dtype) + local_top_k = T.sblock_alloc_buffer((k_val,), dtype=dtype, scope="local") + local_top_k_index = T.sblock_alloc_buffer((k_val,), dtype=index_dtype, scope="local") + for io in T.thread_binding(0, T.ceildiv(batch_size, TX), "blockIdx.x"): + for ii in T.thread_binding(0, TX, "threadIdx.x"): + with T.sblock("top_k"): + vi = T.axis.spatial(batch_size, io * TX + ii) + T.where(io * TX + ii < batch_size) + with T.sblock("init"): + _gating_topk_init_local_top_k( + k_val, dtype, local_top_k, local_top_k_index + ) + for k in range(num_local_experts): + with T.sblock("update"): + vk = T.axis.remap("S", [k]) + _gating_topk_process_value( + k_val, x, local_top_k, local_top_k_index, vi, vk + ) + for j in T.unroll(k_val): + with T.sblock("output"): + vj = T.axis.remap("S", [j]) + out[vi, vj] = local_top_k[vj] + out_index[vi, vj] = local_top_k_index[vj] + + return topk_func + + return op.tensor_ir_op( + _get_topk_func(k), + f"top{k}", + args=[scores], + out=( + Tensor.placeholder([batch_size, k], dtype), + Tensor.placeholder([batch_size, k], index_dtype), + ), + ) + + +def gating_softmax_topk(x: Tensor, k: int, norm_topk_prob=True) -> Tuple[Tensor, Tensor]: # noqa: UP006 + """Compute the softmax score, choose the top-k experts, and returns selected scores. + + Parameters + ---------- + x : Tensor + The input tensor with shape [batch_size, num_local_experts]. + + k : int + The number of top elements to be selected, which is `num_experts_per_tok` in MoE. + + norm_topk_prob : bool + Whether to normalize the top-k expert scores. + + Returns + ------- + expert_weights: Tensor + The top-k expert scores with shape [batch_size, k]. + + expert_indices: Tensor + The top-k expert indices with shape [batch_size, k]. + """ + (batch_size, num_local_experts), dtype = x.shape, x.dtype + index_dtype = "int32" + + TX = 1024 + + def _get_topk_softmax_norm_func(k_val: int): + def _nested_max(local_top_k_f32): + expr = local_top_k_f32[0] + for i in range(1, k_val): + expr = T.max(expr, local_top_k_f32[i]) + return expr + + def _nested_sum(local_top_k_f32, local_top_k_max): + expr = T.exp(local_top_k_f32[0] - local_top_k_max[0]) + for i in range(1, k_val): + expr = expr + T.exp(local_top_k_f32[i] - local_top_k_max[0]) + return expr + + @T.prim_func(private=True, s_tir=True) + def topk_softmax_norm_func( + var_x: T.handle, + var_out: T.handle, + var_out_index: T.handle, + ) -> None: + T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": True}) + batch_size = T.int64() + x = T.match_buffer(var_x, (batch_size, num_local_experts), dtype) + out = T.match_buffer(var_out, (batch_size, k_val), dtype) + out_index = T.match_buffer(var_out_index, (batch_size, k_val), index_dtype) + local_top_k = T.sblock_alloc_buffer((k_val,), dtype=dtype, scope="local") + local_top_k_index = T.sblock_alloc_buffer((k_val,), dtype=index_dtype, scope="local") + local_top_k_f32 = T.sblock_alloc_buffer((k_val,), dtype="float32", scope="local") + local_top_k_max = T.sblock_alloc_buffer((1,), dtype="float32", scope="local") + for io in T.thread_binding(0, T.ceildiv(batch_size, TX), "blockIdx.x"): + for ii in T.thread_binding(0, TX, "threadIdx.x"): + with T.sblock("top_k"): + vi = T.axis.spatial(batch_size, io * TX + ii) + T.where(io * TX + ii < batch_size) + with T.sblock("init"): + _gating_topk_init_local_top_k( + k_val, dtype, local_top_k, local_top_k_index + ) + for k in range(num_local_experts): + with T.sblock("update"): + vk = T.axis.remap("S", [k]) + _gating_topk_process_value( + k_val, x, local_top_k, local_top_k_index, vi, vk + ) + for j in T.unroll(k_val): + with T.sblock("cast"): + vj = T.axis.remap("S", [j]) + local_top_k_f32[vj] = T.cast(local_top_k[vj], "float32") + with T.sblock("max"): + local_top_k_max[0] = _nested_max(local_top_k_f32) + for j in T.unroll(k_val): + with T.sblock("output"): + vj = T.axis.remap("S", [j]) + out[vi, vj] = T.cast( + T.exp(local_top_k_f32[vj] - local_top_k_max[0]) + / _nested_sum(local_top_k_f32, local_top_k_max), + dtype, + ) + out_index[vi, vj] = local_top_k_index[vj] + + return topk_softmax_norm_func + + if norm_topk_prob: + return op.tensor_ir_op( + _get_topk_softmax_norm_func(k), + f"top{k}_softmax", + args=[x], + out=( + Tensor.placeholder([batch_size, k], dtype), + Tensor.placeholder([batch_size, k], index_dtype), + ), + ) + + expert_score = op.softmax(x.astype("float32"), axis=-1).astype(dtype) + return gating_topk(expert_score, k) + + +def group_limited_greedy_topk( + scores: Tensor, # (num_tokens, num_routed_experts) + top_k: int, + num_routed_experts: int, + n_group: int, + topk_group: int, + topk_method: Literal["group_limited_greedy", "noaux_tc"], + num_tokens: IntExpr, + e_score_correction_bias: Optional[Tensor], +) -> Tuple[Tensor, Tensor]: # noqa: UP006 + """Group-limited greedy top-k expert selection. + + Parameters + ---------- + scores : Tensor + The input tensor with shape [num_tokens, num_routed_experts]. + + top_k : int + The number of top elements to be selected, which is `num_experts_per_tok` in MoE. + + num_routed_experts : int + The number of routed experts. + + n_group : int + The number of groups. + + topk_group : int + The number of top-k groups to be selected. + + topk_method : Literal["group_limited_greedy", "noaux_tc"] + The method to select the top-k groups. + + num_tokens : IntExpr + The number of tokens. + + e_score_correction_bias : Optional[Tensor] + The bias of the expert scores. Only available for "noaux_tc". + + Returns + ------- + expert_weights : Tensor + The top-k expert scores with shape [num_tokens, top_k]. + + expert_indices : Tensor + The top-k expert indices with shape [num_tokens, top_k]. + """ + assert scores.dtype == "float32" + scores_for_choice = scores + if topk_method == "noaux_tc": + assert e_score_correction_bias is not None + assert e_score_correction_bias.dtype == "float32" + scores_for_choice = scores + e_score_correction_bias + group_size = num_routed_experts // n_group + if topk_method == "noaux_tc": + group_scores = op.sum( + gating_topk( + scores_for_choice.reshape(num_tokens * n_group, group_size), + 2, + )[0], + axis=-1, + ).reshape(num_tokens, n_group) + else: + group_scores = op.max( + scores_for_choice.reshape(num_tokens * n_group, group_size), axis=-1 + ).reshape(num_tokens, n_group) + group_idx = gating_topk(group_scores, topk_group)[1] # (num_tokens, top_k_group) + + @T.prim_func(private=True, s_tir=True) + def group_limited_mask_scores( + var_scores: T.handle, var_group_idx: T.handle, var_output: T.handle + ): + T.func_attr({"tirx.noalias": True}) + scores = T.match_buffer( + var_scores, (num_tokens, num_routed_experts), dtype=scores_for_choice.dtype + ) + group_idx_tir = T.match_buffer( + var_group_idx, (num_tokens, topk_group), dtype=group_idx.dtype + ) + output = T.match_buffer( + var_output, (num_tokens, num_routed_experts), dtype=scores_for_choice.dtype + ) + for i, j, k in T.grid(num_tokens, topk_group, group_size): + with T.sblock("mask_scores"): + vi, vj, vk = T.axis.remap("SSS", [i, j, k]) + output[vi, group_idx_tir[vi, vj] * group_size + vk] = scores[ + vi, group_idx_tir[vi, vj] * group_size + vk + ] + + tmp_scores = op.tensor_ir_inplace_op( + group_limited_mask_scores, + "group_limited_mask_scores", + args=[ + scores_for_choice, + group_idx, + op.full( + scores_for_choice.shape, + float(np.finfo("float32").min), + dtype=scores_for_choice.dtype, + ), + ], + inplace_indices=[2], + out=Tensor.placeholder(scores_for_choice.shape, scores_for_choice.dtype), + ) + + expert_weights, expert_indices = gating_topk(tmp_scores, top_k) + if topk_method == "noaux_tc": + + @T.prim_func(private=True, s_tir=True) + def gather_scores(var_scores: T.handle, var_expert_indices: T.handle, var_output: T.handle): + T.func_attr({"tirx.noalias": True}) + scores = T.match_buffer( + var_scores, + (num_tokens, num_routed_experts), + dtype=scores_for_choice.dtype, + ) + expert_indices_tir = T.match_buffer( + var_expert_indices, (num_tokens, top_k), dtype=expert_indices.dtype + ) + output = T.match_buffer(var_output, (num_tokens, top_k), dtype=scores_for_choice.dtype) + for i, j in T.grid(num_tokens, top_k): + with T.sblock("gather_scores"): + vi, vj = T.axis.remap("SS", [i, j]) + output[vi, vj] = scores[vi, expert_indices_tir[vi, vj]] + + expert_weights = op.tensor_ir_op( + gather_scores, + "gather_scores", + args=[scores, expert_indices], + out=Tensor.placeholder((num_tokens, top_k), scores_for_choice.dtype), + ) + return expert_weights, expert_indices + + +def moe_cumsum(expert_indices: Tensor, num_local_experts: int) -> Tensor: + """An operator that returns the cumsum array in MoE. + + The input `expert_indices` of shape [batch_size, experts_per_tok] indicates the indices of + the activated experts for each instance in a batch. This operator first converts it to + `expert_mask`, a boolean mask with shape [batch_size, num_local_experts], and then computes + cumsum over the transpose-then-flattened array of `expert_mask`. + + A position `(e, b)` in the result `cumsum`, where `e` is the expert id and `b` is the batch id, + indicates a shuffling plan that moves the `b`-th instance that ensures the inputs to the `e`-th + expert is contiguous. + + Parameters + ---------- + expert_indices : Tensor + The topk indices with shape [batch_size, experts_per_tok], int32, where + `experts_per_tok` is the number of activated experts. + + num_local_experts : int + The number of totally experts. + + Returns + ------- + cumsum: Tensor + The cumsum result with shape [num_local_experts * batch_size], int32. + + Example + ------- + Suppose `batch_size` is 4, `experts_per_tok` is 2, the total number of experts is 6, and + `expert_indices` is the 2D tensor below: + + [ + [0, 1], + [1, 2], + [3, 4], + [2, 5], + ] + + , then the `expert_mask` is a tensor of shape [batch_size, num_local_experts] below: + + [ + [1, 1, 0, 0, 0, 0], + [0, 1, 1, 0, 0, 0], + [0, 0, 0, 1, 1, 0], + [0, 0, 1, 0, 0, 1], + ] + + . The result cumsum of the transposed `expert_mask` is a flattened version of 2D tensor below: + + [ + [1, 1, 1, 1], + [2, 3, 3, 3], + [3, 4, 4, 5], + [5, 5, 6, 6], + [6, 6, 7, 7], + [7, 7, 7, 8], + ] + """ + batch_size, experts_per_tok = expert_indices.shape + expert_mask = ( + op.tensor_expr_op( + lambda expert_indices: te.compute( + (batch_size, num_local_experts), + lambda i, j: tirx.expr.Select( + reduce( + tirx.Or, + [expert_indices[i, k] == j for k in range(experts_per_tok)], + ), + true_value=tirx.const(1, "int32"), + false_value=tirx.const(0, "int32"), + ), + ), + "expert_mask", + args=[expert_indices], + ) + .permute_dims(1, 0) + .reshape(batch_size * num_local_experts) + ) + + return op.cumsum(expert_mask, axis=0, exclusive=False, dtype="int32") + + +def get_indices(cumsum: Tensor, expert_indices: Tensor) -> Tuple[Tensor, Tensor]: # noqa: UP006 + """Returns a 1D tensor of indices that represents the shuffling plan for each instance in a + batch, so that the inputs to each experts are contiguous and the indices for reverse permutation + (scatter) to the original order. + + If `reverse_indices[i] = (b, j)`, it means the `b`-th instance in the batch should be moved to the + `i`-th position in shuffling, and `j` doesn not matter only meaning `expert_indices[b, j]` + corresponds to the expert at position `i` in the shuffling plan. We also compute + `token_indices[i] = b` so that we can use `relax.op.take` for shuffling. + + Effectively it is equivalent to the following Python code: + + .. code-block:: python + + for b in range(batch_size): + for j in range(experts_per_tok): + e = expert_indices[b, j] + reverse_indices[cumsum[e * batch_size + b] - 1] = b * experts_per_tok + j + token_indices[cumsum[e * batch_size + b] - 1 + + Parameters + ---------- + cumsum : Tensor + A flattened 1D tensor whose original shape is [experts_per_tok, batch_size]. + + expert_indices : Tensor + The indices of the experts with shape [batch_size, experts_per_tok]. + + Returns + ------- + reverse_indices : Tensor + The indices for scattering with shape [batch_size * experts_per_tok]. + + token_indices : Tensor + The indices for shuffling with shape [batch_size * experts_per_tok]. + """ # noqa: E501 + TX = 1024 + batch_size, experts_per_tok = expert_indices.shape + + @T.prim_func(private=True, s_tir=True) + def _func( + var_cumsum: T.handle, + var_expert_indices: T.handle, + var_reverse_indices: T.handle, + var_token_indices: T.handle, + ): + T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True}) + batch_size = T.int32() + cumsum_len = T.int32() # [experts_per_tok * batch_size] + cumsum = T.match_buffer(var_cumsum, [cumsum_len], "int32") + expert_indices = T.match_buffer(var_expert_indices, [batch_size, experts_per_tok], "int32") + reverse_indices = T.match_buffer( + var_reverse_indices, [batch_size * experts_per_tok], "int32" + ) + token_indices = T.match_buffer(var_token_indices, [batch_size * experts_per_tok], "int32") + for bj_o in T.thread_binding(0, T.ceildiv(batch_size * experts_per_tok, TX), "blockIdx.x"): + for bj_i in T.thread_binding(0, TX, "threadIdx.x"): + with T.sblock("indices"): + T.reads(expert_indices[:, :], cumsum[:]) + T.writes(reverse_indices[:], token_indices[:]) + if bj_o * TX + bj_i < batch_size * experts_per_tok: + b: T.int32 = T.floordiv(bj_o * TX + bj_i, experts_per_tok) + j: T.int32 = T.floormod(bj_o * TX + bj_i, experts_per_tok) + e: T.int32 = expert_indices[b, j] + reverse_indices[cumsum[e * batch_size + b] - 1] = b * experts_per_tok + j + token_indices[cumsum[e * batch_size + b] - 1] = b + + return op.tensor_ir_op( + _func, + "get_indices", + args=[cumsum, expert_indices], + out=[Tensor.placeholder([batch_size * experts_per_tok], "int32") for _ in range(2)], + ) + + +def get_indptr( + cumsum: Tensor, + num_local_experts: int, + batch_size: Union[int, tirx.Var], + inclusive: bool, + out_dtype: str, +) -> Tensor: + """Extract the `indptr` array from MoE cumsum array. The MoE cumsum array is a flattened tensor + whose original shape is [num_local_experts, batch_size], and the `indptr` array is a 1D tensor + of length `num_local_experts + 1`. The range `[indptr[i], indptr[i + 1])` indicates instances in + the batch that corresponds to the `i`-th expert. + + Effectively, this operator is equivalent to the following numpy code: + + .. code-block:: python + + indptr = np.zeros(num_local_experts + 1, dtype=np.int32) + indptr[0] = 0 + for i in range(1, num_local_experts + 1): + indptr[i] = cumsum[i * batch_size - 1] + return indptr + + Parameters + ---------- + cumsum : Tensor + The prefix sum of the sparse array with shape [batch_size * num_local_experts], int32. + + num_local_experts : int + The number of experts. + + batch_size : int | tirx.Var + The batch size. Note that the batch size here refers to `batch_size * seq_len` in MoE, + and we name is `batch_size` for simplicity here only because the two dimensions are fused + in Mixtral. + + inclusive : bool + Whether to compute inclusive or exclusive prefix sum as the indptr. If `inclusive` is False, + the 0-th element of the `indptr` array, which always equals to 0, will be omitted. + + out_dtype : str + The output dtype. + + Returns + ------- + indptr : Tensor + The `indptr` array with shape [num_local_experts + 1] if `inclusive` is True, otherwise + [num_local_experts]. The `indptr` array is of type `out_dtype`. + """ + + out_shape = [num_local_experts if inclusive else num_local_experts + 1] + + @T.prim_func(private=True, s_tir=True) + def _func_exclusive(var_cumsum: T.handle, var_indptr: T.handle, batch_size: T.int64): + T.func_attr({"tirx.noalias": True}) + cumsum = T.match_buffer(var_cumsum, shape=[batch_size * num_local_experts], dtype="int32") + indptr = T.match_buffer(var_indptr, shape=out_shape, dtype=out_dtype) + for vi in T.serial(0, out_shape[0]): + with T.sblock("indptr"): + i = T.axis.spatial(out_shape[0], vi) + indptr[i] = T.Select(i > 0, cumsum[i * batch_size - 1], T.int32(0)) + + @T.prim_func(private=True, s_tir=True) + def _func_inclusive(var_cumsum: T.handle, var_indptr: T.handle, batch_size: T.int64): + T.func_attr({"tirx.noalias": True}) + cumsum = T.match_buffer(var_cumsum, shape=[batch_size * num_local_experts], dtype="int32") + indptr = T.match_buffer(var_indptr, shape=out_shape, dtype=out_dtype) + for vi in T.serial(0, out_shape[0]): + with T.sblock("indptr"): + i = T.axis.spatial(out_shape[0], vi) + indptr[i] = cumsum[(i + 1) * batch_size - 1] + + assert cumsum.ndim == 1 + return op.tensor_ir_op( + _func_inclusive if inclusive else _func_exclusive, + "get_expert_instance_indptr", + args=[cumsum, batch_size], + out=Tensor.placeholder(out_shape, out_dtype), + ) + + +def scatter_output(x: Tensor, indices: Tensor) -> Tensor: + """Scatter the output of MoE experts back to the original positions. + + Parameters + ---------- + x : Tensor + The output of MoE experts with shape [batch_size * num_experts_per_tok, hidden_size]. + + indices : Tensor + The indices of the experts with shape [batch_size * num_experts_per_tok]. + + Returns + ------- + out : Tensor + The output of MoE experts with shape [batch_size * num_experts_per_tok, hidden_size]. + """ + dtype = x.dtype + _, hidden_size = x.shape + + @T.prim_func(private=True, s_tir=True) + def _func(var_x: T.handle, var_indices: T.handle, var_out: T.handle): + T.func_attr({"tirx.noalias": True}) + indices_len = T.int64() + x = T.match_buffer(var_x, [indices_len, hidden_size], dtype) + indices = T.match_buffer(var_indices, [indices_len], "int32") + out = T.match_buffer(var_out, [indices_len, hidden_size], dtype) + for i in T.serial(0, indices_len): + for j in T.serial(0, hidden_size): + with T.sblock("scatter"): + vi, vj = T.axis.remap("SS", [i, j]) + out[indices[vi], vj] = x[vi, vj] + + return op.tensor_ir_op( + _func, + "scatter_output", + args=[x, indices], + out=Tensor.placeholder(x.shape, dtype), + ) diff --git a/python/mlc_llm/op/mrope.py b/python/mlc_llm/op/mrope.py new file mode 100644 index 0000000..784679e --- /dev/null +++ b/python/mlc_llm/op/mrope.py @@ -0,0 +1,442 @@ +"""Utilities for Multimodal Rotary Position Embeddings (MRoPE).""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import List, Optional, Tuple # noqa: UP035 + +import numpy as np +from tvm import te, tirx +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import Tensor, op + + +def _rotate_half(x: Tensor) -> Tensor: + """Rotate the last dimension of ``x`` by swapping pairs.""" + + x1, x2 = op.split(x, 2, axis=-1) + return op.concat([op.negative(x2), x1], dim=-1) + + +def _repeat_mrope_section(section: Sequence[int]) -> Tuple[int, ...]: # noqa: UP006 + if not section: + raise ValueError("mrope_section must not be empty.") + if any(s <= 0 for s in section): + raise ValueError(f"All mrope_section entries must be positive, got {section}.") + return tuple(section) * 2 + + +def _split_indices_from_sizes(sizes: Sequence[int]) -> List[int]: # noqa: UP006 + indices: List[int] = [] # noqa: UP006 + running = 0 + # Drop the final cumulative sum so split() keeps the last chunk. + for size in sizes[:-1]: + running += size + indices.append(running) + return indices + + +def _reorder_cos_sin( + tensor: Tensor, + split_sizes: Sequence[int], +) -> Tensor: + """Reorder cos/sin tensors so the head dimension follows T/H/W repeating sections.""" + + if not split_sizes: + raise ValueError("split_sizes must not be empty.") + split_points = _split_indices_from_sizes(split_sizes) + # relax.op.split returns a Python tuple, so we can iterate directly. + sections = op.split(tensor, indices_or_sections=split_points, axis=-1) + reordered = [] + for idx, chunk in enumerate(sections): + axis_selector = nn.Tensor.from_const(np.array([idx % 3], dtype="int32")) + axis_slice = op.take(chunk, axis_selector, axis=0) + reordered.append(nn.op.squeeze(axis_slice, 0)) + return op.concat(reordered, dim=-1) + + +class MultimodalRotaryEmbedding(nn.Module): + """Generate cosine/sine tables for multimodal rotary embeddings.""" + + def __init__( + self, + head_dim: int, + theta: float, + mrope_section: Sequence[int], + attention_scaling: float = 1.0, + ) -> None: + if head_dim % 2 != 0: + raise ValueError(f"head_dim must be even for RoPE, got {head_dim}.") + self.head_dim = head_dim + self.theta = theta + self.attention_scaling = attention_scaling + self.mrope_section = tuple(mrope_section) + self._inv_freq = 1.0 / ( + theta ** (np.arange(0, head_dim, 2, dtype="float32") / np.float32(head_dim)) + ) + + def forward(self, reference: Tensor, position_ids: Tensor) -> Tuple[Tensor, Tensor]: # noqa: UP006 + """Return ``(cos, sin)`` with shape ``(3, batch, seq, head_dim)``.""" + if len(position_ids.shape) != 3: + raise ValueError( + "position_ids must be rank-3 with either " + "(batch, seq, 3) or (3, batch, seq) layout, " + f"got shape {position_ids.shape}." + ) + if isinstance(position_ids.shape[0], int) and position_ids.shape[0] == 3: + batch_size, seq_len = position_ids.shape[1], position_ids.shape[2] + pos_tensor = op.reshape(position_ids, (3, batch_size, 1, seq_len)) + elif isinstance(position_ids.shape[-1], int) and position_ids.shape[-1] == 3: + batch_size, seq_len = position_ids.shape[0], position_ids.shape[1] + permuted_pos = op.permute_dims(position_ids, axes=[2, 0, 1]) + pos_tensor = op.reshape(permuted_pos, (3, batch_size, 1, seq_len)) + else: + raise ValueError( + "position_ids must have exactly one static dimension of size 3, " + f"got shape {position_ids.shape}." + ) + + dtype = reference.dtype + inv_freq_tensor = nn.Tensor.from_const(self._inv_freq.reshape(1, 1, -1, 1)) + inv_freq_tensor = op.broadcast_to(inv_freq_tensor, (3, batch_size, self._inv_freq.size, 1)) + + freqs = op.matmul(inv_freq_tensor.astype("float32"), pos_tensor.astype("float32")) + freqs = op.permute_dims(freqs, axes=[0, 1, 3, 2]) + emb = op.concat([freqs, freqs], dim=-1) + + def _apply_trig(func_name: str) -> Tensor: + def compute(x: te.Tensor): + return te.compute( + x.shape, + lambda *indices: getattr(tirx, func_name)(x[indices]), + name=f"mrope_{func_name}", + ) + + return op.tensor_expr_op(compute, f"mrope_{func_name}", [emb]) + + cos = _apply_trig("cos") * self.attention_scaling + sin = _apply_trig("sin") * self.attention_scaling + return cos.astype(dtype), sin.astype(dtype) + + +def apply_multimodal_rotary_pos_emb( + q: Tensor, + k: Tensor, + cos: Tensor, + sin: Tensor, + mrope_section: Sequence[int], + unsqueeze_dim: int = 2, +) -> Tuple[Tensor, Tensor]: # noqa: UP006 + """Apply multimodal rotary embedding to query and key tensors.""" + + split_sizes = _repeat_mrope_section(mrope_section) + reordered_cos = _reorder_cos_sin(cos, split_sizes) + reordered_sin = _reorder_cos_sin(sin, split_sizes) + cos_term = op.unsqueeze(reordered_cos, dim=unsqueeze_dim) + sin_term = op.unsqueeze(reordered_sin, dim=unsqueeze_dim) + cos_term = cos_term.astype(q.dtype) + sin_term = sin_term.astype(q.dtype) + q_embed = op.add(op.multiply(q, cos_term), op.multiply(_rotate_half(q), sin_term)) + k_embed = op.add(op.multiply(k, cos_term), op.multiply(_rotate_half(k), sin_term)) + return q_embed, k_embed + + +@dataclass +class VisionPositionMetadata: + """Metadata required to build multimodal position IDs.""" + + vision_start_token_id: int + image_token_id: int + video_token_id: int + spatial_merge_size: int + tokens_per_second: float + + def merged_hw(self, height: int, width: int) -> Tuple[int, int]: # noqa: UP006 + """Return merged height/width after applying ``spatial_merge_size``.""" + + if height % self.spatial_merge_size != 0 or width % self.spatial_merge_size != 0: + raise ValueError( + "Image or video grid is not divisible by spatial_merge_size " + f"(got h={height}, w={width}, merge={self.spatial_merge_size})." + ) + return height // self.spatial_merge_size, width // self.spatial_merge_size + + +def _text_chunk(length: int, offset: int) -> np.ndarray: + """Create a text-position chunk with a shared scalar offset for T/H/W axes.""" + + if length <= 0: + return np.zeros((3, 0), dtype=np.int64) + seq: np.ndarray = np.arange(length, dtype=np.int64) + chunk = np.broadcast_to(seq.reshape(1, -1), (3, length)) + return chunk + offset + + +def _grid_chunk( + grid_t: int, + grid_h: int, + grid_w: int, + offset: int, + tokens_per_second: float, + second_per_grid: float, +) -> np.ndarray: + if grid_t <= 0 or grid_h <= 0 or grid_w <= 0: + raise ValueError( + f"Invalid grid shape t={grid_t}, h={grid_h}, w={grid_w} for multimodal positions." + ) + time_axis = (np.arange(grid_t, dtype=np.float32) * second_per_grid * tokens_per_second).astype( + np.int64 + ) + t_index = np.repeat(time_axis, grid_h * grid_w) + h_index = np.tile(np.repeat(np.arange(grid_h, dtype=np.int64), grid_w), grid_t) + w_index = np.tile(np.tile(np.arange(grid_w, dtype=np.int64), grid_h), grid_t) + stacked = np.stack([t_index, h_index, w_index]) + return stacked + offset + + +def _find_token_index(tokens: Sequence[int], token_id: int, start: int) -> int: + for idx in range(start, len(tokens)): + if tokens[idx] == token_id: + return idx + return len(tokens) + + +def _next_chunk_offset(chunks: Sequence[np.ndarray]) -> int: + if not chunks: + return 0 + return int(chunks[-1].max()) + 1 + + +def _count_vision_items( + token_array: np.ndarray, + vision_start_token_id: int, + image_token_id: int, + video_token_id: int, +) -> Tuple[int, int]: # noqa: UP006 + vision_starts = np.where(token_array == vision_start_token_id)[0] + valid_starts = vision_starts[vision_starts + 1 < token_array.shape[0]] + following_tokens = token_array[valid_starts + 1] + image_count = int(np.sum(following_tokens == image_token_id)) + video_count = int(np.sum(following_tokens == video_token_id)) + return image_count, video_count + + +def _next_vision_block( + tokens: Sequence[int], + start: int, + meta: VisionPositionMetadata, + has_images: bool, + has_videos: bool, +) -> Tuple[str, int]: # noqa: UP006 + sentinel = len(tokens) + 1 + image_end = _find_token_index(tokens, meta.image_token_id, start) if has_images else sentinel + video_end = _find_token_index(tokens, meta.video_token_id, start) if has_videos else sentinel + if image_end < video_end: + return "image", image_end + return "video", video_end + + +def _load_grid_for_block( + block_kind: str, + image_grid_thw: Optional[np.ndarray], # noqa: UP045 + video_grid_thw: Optional[np.ndarray], # noqa: UP045 + second_per_grid_ts: Optional[np.ndarray], # noqa: UP045 + image_index: int, + video_index: int, +) -> Tuple[int, int, int, float, int, int]: # noqa: UP006 + if block_kind == "image": + if image_grid_thw is None: + raise ValueError("Image grids are required for sequences with image tokens.") + grid_t, grid_h, grid_w = image_grid_thw[image_index] + return int(grid_t), int(grid_h), int(grid_w), 0.0, image_index + 1, video_index + + if video_grid_thw is None: + raise ValueError("Video grids are required for sequences with video tokens.") + grid_t, grid_h, grid_w = video_grid_thw[video_index] + second_per_grid = ( + float(second_per_grid_ts[video_index]) if second_per_grid_ts is not None else 1.0 + ) + return int(grid_t), int(grid_h), int(grid_w), second_per_grid, image_index, video_index + 1 + + +def _build_sequence_position_ids( + input_tokens: Sequence[int], + meta: VisionPositionMetadata, + image_grid_thw: Optional[np.ndarray], # noqa: UP045 + video_grid_thw: Optional[np.ndarray], # noqa: UP045 + second_per_grid_ts: Optional[np.ndarray], # noqa: UP045 + image_index: int, + video_index: int, +) -> Tuple[np.ndarray, int, int, int]: # noqa: UP006 + token_array = np.asarray(input_tokens, dtype=np.int64) + image_count, video_count = _count_vision_items( + token_array, + vision_start_token_id=meta.vision_start_token_id, + image_token_id=meta.image_token_id, + video_token_id=meta.video_token_id, + ) + if image_count > 0 and image_grid_thw is None: + raise ValueError("Image grids are required for sequences with image tokens.") + if video_count > 0 and video_grid_thw is None: + raise ValueError("Video grids are required for sequences with video tokens.") + + llm_pos_ids_list: List[np.ndarray] = [] # noqa: UP006 + start = 0 + remain_images = image_count + remain_videos = video_count + for _ in range(image_count + video_count): + block_kind, block_end = _next_vision_block( + tokens=input_tokens, + start=start, + meta=meta, + has_images=remain_images > 0, + has_videos=remain_videos > 0, + ) + ( + grid_t, + grid_h, + grid_w, + second_per_grid, + image_index, + video_index, + ) = _load_grid_for_block( + block_kind=block_kind, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + image_index=image_index, + video_index=video_index, + ) + if block_kind == "image": + remain_images -= 1 + else: + remain_videos -= 1 + + llm_grid_h, llm_grid_w = meta.merged_hw(grid_h, grid_w) + text_len = block_end - start + text_offset = _next_chunk_offset(llm_pos_ids_list) + llm_pos_ids_list.append(_text_chunk(text_len, text_offset)) + grid_offset = text_offset + text_len + llm_pos_ids_list.append( + _grid_chunk( + grid_t=grid_t, + grid_h=llm_grid_h, + grid_w=llm_grid_w, + offset=grid_offset, + tokens_per_second=meta.tokens_per_second, + second_per_grid=second_per_grid, + ) + ) + start = block_end + grid_t * llm_grid_h * llm_grid_w + + if start < len(input_tokens): + tail_len = len(input_tokens) - start + tail_offset = _next_chunk_offset(llm_pos_ids_list) + llm_pos_ids_list.append(_text_chunk(tail_len, tail_offset)) + + if not llm_pos_ids_list: + empty_positions: np.ndarray = np.zeros((3, 0), dtype=np.int64) + return empty_positions, 0, image_index, video_index + llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) + delta = int(llm_positions.max()) + 1 - len(input_tokens) + return llm_positions, delta, image_index, video_index + + +def _text_only_position_ids( + input_ids: np.ndarray, + attention_mask: Optional[np.ndarray], # noqa: UP045 +) -> Tuple[np.ndarray, np.ndarray]: # noqa: UP006 + batch, seq_len = input_ids.shape + if attention_mask is None: + base: np.ndarray = np.arange(seq_len, dtype=np.int64).reshape(1, 1, -1) + tiled = np.broadcast_to(base, (3, batch, seq_len)) + return tiled, np.zeros((batch, 1), dtype=np.int64) + + position = attention_mask.cumsum(axis=-1) - 1 + position = np.where(attention_mask == 0, 1, position) + position = np.expand_dims(position, axis=0).repeat(3, axis=0) + max_pos = position.max(axis=0, keepdims=False).max(axis=-1, keepdims=True) + delta = (max_pos + 1 - seq_len).astype(np.int64) + return position.astype(np.int64), delta + + +def get_mrope_position_ids( + input_ids: np.ndarray, + meta: VisionPositionMetadata, + attention_mask: Optional[np.ndarray] = None, # noqa: UP045 + image_grid_thw: Optional[np.ndarray] = None, # noqa: UP045 + video_grid_thw: Optional[np.ndarray] = None, # noqa: UP045 + second_per_grid_ts: Optional[np.ndarray] = None, # noqa: UP045 +) -> Tuple[np.ndarray, np.ndarray]: # noqa: UP006 + """Generate 3D position IDs and deltas following Hugging Face Qwen2.5-VL.""" + + input_ids = np.asarray(input_ids, dtype=np.int64) + batch, seq_len = input_ids.shape + position_ids = np.ones((3, batch, seq_len), dtype=np.int64) + + attention = None + if attention_mask is not None: + attention_mask = np.asarray(attention_mask, dtype=np.int64) + if attention_mask.shape != input_ids.shape: + raise ValueError( + "attention_mask shape must match input_ids shape: " + f"{attention_mask.shape} vs {input_ids.shape}" + ) + attention = attention_mask.astype(bool) + + image_grid_thw = None if image_grid_thw is None else np.asarray(image_grid_thw, dtype=np.int64) + video_grid_thw = None if video_grid_thw is None else np.asarray(video_grid_thw, dtype=np.int64) + if second_per_grid_ts is not None: + second_per_grid_ts = np.asarray(second_per_grid_ts, dtype=np.float32) + + contains_image_tokens = bool(np.any(input_ids == meta.image_token_id)) + contains_video_tokens = bool(np.any(input_ids == meta.video_token_id)) + if contains_image_tokens and image_grid_thw is None: + raise ValueError("image_grid_thw must be provided when image tokens exist in input_ids.") + if contains_video_tokens and video_grid_thw is None: + raise ValueError("video_grid_thw must be provided when video tokens exist in input_ids.") + if ( + second_per_grid_ts is not None + and video_grid_thw is not None + and second_per_grid_ts.shape[0] != video_grid_thw.shape[0] + ): + raise ValueError( + "second_per_grid_ts length must match number of video grids " + f"({second_per_grid_ts.shape[0]} vs {video_grid_thw.shape[0]})." + ) + + if not (contains_image_tokens or contains_video_tokens): + return _text_only_position_ids(input_ids, attention_mask) + + image_index = 0 + video_index = 0 + deltas: List[int] = [] # noqa: UP006 + + for batch_idx in range(batch): + tokens = input_ids[batch_idx] + if attention is not None: + tokens = tokens[attention[batch_idx]] + token_values = np.asarray(tokens, dtype=np.int64).tolist() + input_tokens: List[int] = [int(token) for token in token_values] # noqa: UP006 + if not input_tokens: + deltas.append(0) + continue + + llm_positions, delta, image_index, video_index = _build_sequence_position_ids( + input_tokens=input_tokens, + meta=meta, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + image_index=image_index, + video_index=video_index, + ) + if attention is not None: + position_ids[:, batch_idx, attention[batch_idx]] = llm_positions + else: + position_ids[:, batch_idx, :] = llm_positions + deltas.append(delta) + + delta_array = np.asarray(deltas, dtype=np.int64).reshape(batch, 1) + return position_ids, delta_array diff --git a/python/mlc_llm/op/pipeline_parallel.py b/python/mlc_llm/op/pipeline_parallel.py new file mode 100644 index 0000000..dab0bc4 --- /dev/null +++ b/python/mlc_llm/op/pipeline_parallel.py @@ -0,0 +1,33 @@ +"""Operators for pipeline parallelism.""" + +from typing import List # noqa: UP035 + +from tvm import relax +from tvm.relax.frontend.nn import Tensor, op + + +def pipeline_stage_boundary(*tensors: Tensor) -> List[Tensor]: # noqa: UP006 + """Pipeline parallelism stage boundary mark operator in MLC. + + Parameters + ---------- + tensors : Tensor + The tensors to be passed to the next stage. + + Returns + ------- + tensors : List[Tensor] + The list of input tensors passed to the next stage. + """ + return op.wrap_nested( + relax.call_pure_packed( + "mlc.pipeline_parallel_stage_boundary", + *[tensor._expr for tensor in tensors], + ty_args=( + tensors[0]._expr.ty + if len(tensors) == 1 + else relax.TupleType([tensor._expr.ty for tensor in tensors]) + ), + ), + name="pipeline_stage_boundary", + ) diff --git a/python/mlc_llm/op/top_p_pivot.py b/python/mlc_llm/op/top_p_pivot.py new file mode 100644 index 0000000..2b28dea --- /dev/null +++ b/python/mlc_llm/op/top_p_pivot.py @@ -0,0 +1,335 @@ +"""Operators for choosing the pivot to cut-off top-p percentile""" + +import tvm +from tvm.script import tirx as T + +from mlc_llm.support.max_thread_check import get_max_num_threads_per_block + +# mypy: disable-error-code="attr-defined,valid-type,name-defined" + + +def top_p_pivot(pN, target: tvm.target.Target): + """Top-p pivot function. This function finds the pivot to cut-off top-p percentile. + + A valide pivot should satisfy the following conditions: + - lsum >= top_p + - top_p > lsum - cmin * lmin + where lsum is the sum of elements that are larger or equal to the pivot, + lmin is the minimum elements that is larger or equal to the pivot, + cmin is the count of elements that are equal to lmin, + + Parameters + ---------- + prob: + The probability vector + + top_p_arr: + The top-p threshold + + init_pivots: + The initial pivot candidates + + final_pivot: + The final pivot to cut-off top-p percentile + + final_lsum: + The final sum of the values after top-p filtering. + """ + TX = 1024 + K = 32 + eps_LR = 1e-7 + + max_num_threads_per_block = get_max_num_threads_per_block(target) + TX = min(TX, max_num_threads_per_block) + + def _var(dtype="int32"): + return T.sblock_alloc_buffer((1,), dtype, scope="local") + + def valid(lsum, lmin, cmin, top_p): + return tvm.tirx.all(lsum >= top_p, top_p > lsum - cmin * lmin) + + # fmt: off + @T.prim_func(private=True, s_tir=True) + def _func( + var_prob: T.handle, + var_top_p_arr: T.handle, + var_init_pivots: T.handle, + var_final_pivot: T.handle, + var_final_lsum: T.handle, + ): + T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True}) + B = T.int32() + N = T.int32() + prob = T.match_buffer(var_prob, (B, N,), "float32") + top_p_arr = T.match_buffer(var_top_p_arr, (B,), dtype="float32") + init_pivots = T.match_buffer(var_init_pivots, (B, pN), "float32") + final_pivot = T.match_buffer(var_final_pivot, (B,), "float32") + final_lsum = T.match_buffer(var_final_lsum, (B,), "float32") + + with T.sblock("kernel"): + pivot = T.sblock_alloc_buffer((pN,), "float32", scope="local") + top_p = _var("float32") + + L = T.sblock_alloc_buffer((1,), "float32", scope="shared") + R = T.sblock_alloc_buffer((1,), "float32", scope="shared") + L_local = _var("float32") + R_local = _var("float32") + + q = _var("float32") + lsum = T.sblock_alloc_buffer((pN,), "float32", scope="local") + lmin_broadcast = T.sblock_alloc_buffer((1), "float32", scope="shared") + lmin_broadcast_local = _var("float32") + lmin = T.sblock_alloc_buffer((pN,), "float32", scope="local") + cmin = T.sblock_alloc_buffer((pN,), "int32", scope="local") + total_sum = _var("float32") + + it = _var("int32") + es_local = _var("bool") + es = T.sblock_alloc_buffer((1,), "bool", scope="shared") + find_pivot_local = _var("bool") + find_pivot = T.sblock_alloc_buffer((1,), "bool", scope="shared") + + total_sum_reduce = _var("float32") + lsum_reduce = _var("float32") + lmin_reduce = _var("float32") + cmin_reduce = _var("int32") + + for _bx in T.thread_binding(0, B, thread="blockIdx.x"): + for _tx in T.thread_binding(0, TX, thread="threadIdx.x"): + with T.sblock("CTA"): + b, tx = T.axis.remap("SS", [_bx, _tx]) + + top_p[0] = top_p_arr[b] + + if tx == 0: + # leader thread initializes L, R + L[0] = 1.0 - top_p[0] + R[0] = eps_LR + find_pivot[0] = False + T.tvm_storage_sync("shared") + + L_local[0] = L[0] + R_local[0] = R[0] + for i in T.unroll(0, pN): + # pivots are in descending order + pivot[i] = init_pivots[b, i] + find_pivot_local[0] = False + if L_local[0] - R_local[0] <= eps_LR: + # When the initial value is too small, set the result directly. + if tx == 0: + final_lsum[b] = 1.0 + final_pivot[b] = 0.0 + find_pivot_local[0] = True + + while T.tvm_thread_invariant( + L_local[0] - R_local[0] > eps_LR + and T.Not(find_pivot_local[0]) + ): + # sync before each iteration + T.tvm_storage_sync("shared") + + ### get lsum, lmin, total_sum + for pidx in T.unroll(0, pN): + lsum[pidx] = 0.0 + lmin[pidx] = T.max_value("float32") + cmin[pidx] = 0 + total_sum[0] = 0.0 + it[0] = 0 + es_local[0] = False + while it[0] < T.ceildiv(N, TX) and T.Not(es_local[0]): + idx = T.meta_var(it[0] * TX + tx) + q[0] = T.if_then_else(idx < N, prob[b, idx], 0.0) + total_sum[0] += q[0] + for pidx in T.unroll(0, pN): + if q[0] >= pivot[pidx]: + lsum[pidx] += q[0] + if lmin[pidx] > q[0]: + lmin[pidx] = q[0] + cmin[pidx] = 1 + elif lmin[pidx] == q[0]: + cmin[pidx] += 1 + it[0] += 1 + + # early stop every K iterations + if it[0] % K == 0: + # reduce total_sum over tx + # T.tvm_storage_sync("shared") + with T.sblock("block_cross_thread"): + T.reads(total_sum[0]) + T.writes(total_sum_reduce[0]) + T.attr( + T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]), + "reduce_scope", + T.int32(0), + ) + T.tvm_thread_allreduce(T.uint32(1), total_sum[0], True, total_sum_reduce[0], tx, dtype="void") # noqa: E501 + # T.tvm_storage_sync("shared") + + if tx == 0: + # leader thread checks if we can stop early + es[0] = 1 - total_sum_reduce[0] < pivot[pN - 1] + T.tvm_storage_sync("shared") + es_local[0] = es[0] + + T.tvm_storage_sync("shared") + + # reduce lsum, lmin, cmin, over tx + for pidx in T.serial(0, pN): + # reduce lsum over tx for pivot[j] + with T.sblock("block_cross_thread"): + T.reads(lsum[pidx]) + T.writes(lsum_reduce[0]) + T.attr( + T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]), + "reduce_scope", + T.int32(0), + ) + T.tvm_thread_allreduce(T.uint32(1), lsum[pidx], True, lsum_reduce[0], tx, dtype="void") # noqa: E501 + + # reduce lmin over tx for pivot[j] + with T.sblock("block_cross_thread"): + T.reads(lmin[pidx]) + T.writes(lmin_reduce[0]) + T.attr( + T.comm_reducer(lambda x0, y0: T.min(x0, y0), [T.float32(0)]), # noqa: E501 + "reduce_scope", + T.int32(0), + ) + T.tvm_thread_allreduce(T.uint32(1), lmin[pidx], True, lmin_reduce[0], tx, dtype="void") # noqa: E501 + + if tx == 0: + # broadcast lmin to all threads + lmin_broadcast[0] = lmin_reduce[0] + T.tvm_storage_sync("shared") + lmin_broadcast_local[0] = lmin_broadcast[0] + if lmin[pidx] > lmin_broadcast_local[0]: + cmin[pidx] = 0 + if tx == 0: + # only the leader thread updates lsum, lmin + lsum[pidx] = lsum_reduce[0] + lmin[pidx] = lmin_reduce[0] + + # reduce cmin over tx for pivot[j] + with T.sblock("block_cross_thread"): + T.reads(cmin[pidx]) + T.writes(cmin_reduce[0]) + T.attr( + T.comm_reducer(lambda x0, y0: x0 + y0, [T.int32(0)]), + "reduce_scope", + T.int32(0), + ) + T.tvm_thread_allreduce(T.uint32(1), cmin[pidx], True, cmin_reduce[0], tx, dtype="void") # noqa: E501 + + if tx == 0: + # only the leader thread updates cmin + cmin[pidx] = cmin_reduce[0] + + T.tvm_storage_sync("shared") + + if tx == 0: + # leader thread checks if we have found the pivot, or updates L, R + it[0] = 0 + while it[0] < pN and T.Not(find_pivot_local[0]): + pidx = T.meta_var(it[0]) + if valid(lsum[pidx], lmin[pidx], cmin[pidx], top_p[0]): + find_pivot[0] = True + find_pivot_local[0] = True + # write back the pivot and lsum + final_pivot[b] = pivot[pidx] + final_lsum[b] = lsum[pidx] + elif lsum[pidx] - lmin[pidx] * cmin[pidx] >= top_p[0]: + R[0] = pivot[pidx] + final_lsum[b] = lsum[pidx] + elif lsum[pidx] < top_p[0]: + L[0] = pivot[pidx] + it[0] += 1 + + T.tvm_storage_sync("shared") + + L_local[0] = L[0] + R_local[0] = R[0] + find_pivot_local[0] = find_pivot[0] + # new pivots for next iteration + # uniform spacing between L and R + for pidx in T.unroll(0, pN): + pivot[pidx] = L[0] - (pidx + 1) * (L_local[0] - R_local[0]) / (pN + 1) # noqa: E501 + + if tx == 0: + # leader thread writes back the pivot + if T.Not(find_pivot_local[0]): + final_pivot[b] = R_local[0] + if R_local[0] == eps_LR: + final_lsum[b] = lsum[pN - 1] + # fmt: on + + return _func + + +def top_p_renorm(target: tvm.target.Target = None): + """Top-p renormalization function. This function renormalizes the probability vector. + + Given the pivot, the probability vector is renormalized as follows: + - if prob >= pivot, renorm_prob = prob / lsum + - otherwise, renorm_prob = 0 + + Parameters + ---------- + prob: + The probability vector + + final_pivot: + The final pivot to cut-off top-p percentile + + final_lsum: + The sum of elements that are larger or equal to the pivot + + renorm_prob: + The renormalized probability vector + """ + TX = 1024 + CTA_COUNT = 512 + + if target: + max_num_threads_per_block = get_max_num_threads_per_block(target) + TX = min(TX, max_num_threads_per_block) + + def _var(dtype="int32"): + return T.sblock_alloc_buffer((1,), dtype, scope="local") + + # fmt: off + @T.prim_func(private=True, s_tir=True) + def _func( + var_prob: T.handle, + var_final_pivot: T.handle, + var_final_lsum: T.handle, + var_renorm_prob: T.handle, + ): + T.func_attr({"tirx.is_scheduled": 1, "tirx.noalias": True}) + B = T.int32() + N = T.int32() + prob = T.match_buffer(var_prob, (B, N,), "float32") + final_pivot = T.match_buffer(var_final_pivot, (B,), "float32") + final_lsum = T.match_buffer(var_final_lsum, (B,), "float32") + renorm_prob = T.match_buffer(var_renorm_prob, (B, N,), "float32") + + with T.sblock("kernel"): + pivot = _var("float32") + lsum = _var("float32") + BX = T.meta_var(T.ceildiv(CTA_COUNT, B)) + + for _by in T.thread_binding(0, B, thread="blockIdx.y"): + for _bx in T.thread_binding(0, BX, thread="blockIdx.x"): + for _tx in T.thread_binding(0, TX, thread="threadIdx.x"): + with T.sblock("CTA"): + by, bx, tx = T.axis.remap("SSS", [_by, _bx, _tx]) + + pivot[0] = final_pivot[by] + lsum[0] = final_lsum[by] + + for i in T.serial(T.ceildiv(N, BX * TX)): + idx = T.meta_var(i * BX * TX + bx * TX + tx) + if idx < N: + renorm_prob[by, idx] = T.if_then_else(prob[by, idx] >= pivot[0], prob[by, idx] / lsum[0], 0.0) # noqa: E501 + # fmt: on + + return _func diff --git a/python/mlc_llm/op/triton.py b/python/mlc_llm/op/triton.py new file mode 100644 index 0000000..bc6cf0e --- /dev/null +++ b/python/mlc_llm/op/triton.py @@ -0,0 +1,762 @@ +"""Operators enabled by external modules.""" + +from typing import List, Literal, Tuple # noqa: UP035 + +import tvm +from tvm.relax.frontend import nn +from tvm.script import ir as I +from tvm.script import tirx as T + +try: + import triton + import triton.language as tl +except ImportError: + triton = None + tl = None + + +# We use a wrapper function to avoid type annotation issue of "tl.constexpr" when +# triton is not installed. +def _get_triton_w8a8_block_fp8_gemm(): + # Triton kernel adapted from SGLang project + # https://github.com/sgl-project/sglang/blob/v0.4.4/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py# noqa: E501 + def _triton_w8a8_block_fp8_gemm( + # Pointers to inputs and output + A, + B, + C, + As, + Bs, + # Shape for matmul + M, + N: tl.constexpr, + K: tl.constexpr, + # Stride for inputs and output + stride_am: tl.constexpr, + stride_ak: tl.constexpr, + stride_bk: tl.constexpr, + stride_bn: tl.constexpr, + stride_cm: tl.constexpr, + stride_cn: tl.constexpr, + stride_As_m: tl.constexpr, + stride_As_k: tl.constexpr, + stride_Bs_k: tl.constexpr, + stride_Bs_n: tl.constexpr, + # Block size for block-wise quantization + group_n: tl.constexpr, + group_k: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + ): + """Triton-accelerated function used to perform linear operations (dot + product) on input tensors `A` and `B` with block-wise quantization, + and store the result in output tensor `C`. + """ + + pid = tl.program_id(axis=0).to(tl.int64) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + As_ptrs = As + offs_am * stride_As_m + offs_bsn = offs_bn // group_n + Bs_ptrs = Bs + offs_bsn * stride_Bs_n + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_s = tl.load(As_ptrs + offs_ks * stride_As_k) + b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k) + + accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :] + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if C.dtype.element_ty == tl.bfloat16: + c = accumulator.to(tl.bfloat16) + elif C.dtype.element_ty == tl.float16: + c = accumulator.to(tl.float16) + else: + c = accumulator.to(tl.float32) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + return _triton_w8a8_block_fp8_gemm + + +# We use a wrapper function to avoid type annotation issue of "tl.constexpr" when +# triton is not installed. +def _get_triton_w8a8_block_fp8_group_gemm(): + # Triton kernel adapted from SGLang project + # https://github.com/sgl-project/sglang/blob/v0.4.4/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py# noqa: E501 + def _triton_w8a8_block_fp8_group_gemm( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + a_scale_ptr, + b_scale_ptr, + expert_ids_ptr, + indptr_ptr, + # Matrix dimensions + EM, + N: tl.constexpr, + K: tl.constexpr, + num_experts: tl.constexpr, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am: tl.constexpr, + stride_ak: tl.constexpr, + stride_be: tl.constexpr, + stride_bk: tl.constexpr, + stride_bn: tl.constexpr, + stride_cm: tl.constexpr, + stride_cn: tl.constexpr, + stride_asm: tl.constexpr, + stride_ask: tl.constexpr, + stride_bse: tl.constexpr, + stride_bsk: tl.constexpr, + stride_bsn: tl.constexpr, + # Block size for block-wise quantization + group_n: tl.constexpr, + group_k: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + even_Ks: tl.constexpr, + ): + """ + Implements the fused computation for a Mixture of Experts (MOE) using + token and expert matrices. + + Key Parameters: + - A: The input tensor representing tokens with shape (*, K), where '*' can + be any shape representing batches and K is the feature dimension of + each token. + - B: The stacked MOE weight tensor with shape (E, N, K), where E is + the number of experts, K is the input feature dimension, and N is + the output feature dimension. + - C: The output cache tensor with shape (*, N), where '*' means the + same shape as the input tensor A, and N is the output feature dimension. + - expert_ids: A tensor containing the indices of the expert for each + block. It determines which expert matrix from B should be used for + each block in A. + This kernel performs the multiplication of a token by its corresponding + expert matrix as determined by `expert_ids`. + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0).to(tl.int64) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_experts + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + expert_id = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if expert_id == -1: + return + + token_begin = tl.load(indptr_ptr + expert_id) + token_end = tl.load(indptr_ptr + expert_id + 1) + start_pid_m = tl.cdiv(token_begin, BLOCK_SIZE_M) + expert_id + offs_token_id = ( + token_begin + (pid_m - start_pid_m) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + ) + token_mask = offs_token_id < token_end + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + offs_token_id[:, None] * stride_am + offs_k[None, :] * stride_ak + + b_ptrs = ( + b_ptr + + expert_id * stride_be + + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + ) + + a_scale_ptrs = a_scale_ptr + offs_token_id * stride_asm + offs_bsn = offs_bn // group_n + b_scale_ptrs = b_scale_ptr + expert_id * stride_bse + offs_bsn * stride_bsn + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the + # K dimension. + if even_Ks: + a = tl.load( + a_ptrs, + mask=token_mask[:, None], + other=0.0, + ) + b = tl.load(b_ptrs) + else: + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + + # We accumulate along the K dimension. + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask, mask=token_mask, other=0.0) + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + + accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :] + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if c_ptr.dtype.element_ty == tl.bfloat16: + accumulator = accumulator.to(tl.bfloat16) + elif c_ptr.dtype.element_ty == tl.float16: + accumulator = accumulator.to(tl.float16) + else: + accumulator = accumulator.to(tl.float32) + + # ----------------------------------------------------------- + # Write back the block of the output + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token_id[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + return _triton_w8a8_block_fp8_group_gemm + + +def get_tir_w8a8_block_fp8_matmul( + N: int, + K: int, + block_n: int, + block_k: int, + in_dtype: Literal["float8_e4m3fn"], + out_dtype: Literal["float16", "bfloat16"], + BLOCK_SIZE_M: int, + BLOCK_SIZE_N: int, + BLOCK_SIZE_K: int, + GROUP_SIZE_M: int, + num_warps: int, + num_stages: int, + extern_mods: List[tvm.runtime.Module], # noqa: UP006 +): + """Get the TIR function for the w8a8_block_fp8_matmul kernel.""" + # NOTE: adding the type annotation of " -> Tuple[Optional[tvm.tirx.PrimFunc], str]" + # will cause the failure of the type resolution in mypy. + if triton is None: + raise RuntimeError("Triton is not installed. Please install it with `pip install triton`.") + + name_suffix = f"_N{N}_K{K}_block_n{block_n}_block_k{block_k}_in{in_dtype}_out{out_dtype}" + kernel_name = f"triton_w8a8_block_fp8_gemm{name_suffix}" + tir_name = f"tir_w8a8_block_fp8_matmul{name_suffix}" + for ext_mod in extern_mods: + if ext_mod.implements_function(kernel_name): + return [None, tir_name] + + triton_kernel = _get_triton_w8a8_block_fp8_gemm() + triton_kernel.__name__ = kernel_name + + @I.ir_module + class BlockFP8Matmul: + @T.prim_func(private=True, s_tir=True) + def tir_w8a8_block_fp8_matmul( + var_A: T.handle, + var_B: T.handle, + var_As: T.handle, + var_Bs: T.handle, + var_C: T.handle, + ): + T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1}) + M = T.int32() + A = T.match_buffer(var_A, (M, K), dtype=in_dtype) + B = T.match_buffer(var_B, (N, K), dtype=in_dtype) + As = T.match_buffer(var_As, (M, (K + block_k - 1) // block_k), "float32") + Bs = T.match_buffer( + var_Bs, + ((N + block_n - 1) // block_n, (K + block_k - 1) // block_k), + "float32", + ) + C = T.match_buffer(var_C, (M, N), dtype=out_dtype) + with T.sblock("root"): + T.reads( + A[0:M, 0:K], + B[0:N, 0:K], + As[0:M, 0 : (K + block_k - 1) // block_k], + Bs[ + 0 : (N + block_n - 1) // block_n, + 0 : (K + block_k - 1) // block_k, + ], + ) + T.writes(C[0:M, 0:N]) + T.call_kernel( + triton.jit(triton_kernel), + (T.ceildiv(M, BLOCK_SIZE_M) * T.ceildiv(N, BLOCK_SIZE_N),), + A.data, + B.data, + C.data, + As.data, + Bs.data, + M, + N, + K, + K, # stride_am + 1, # stride_ak + 1, # stride_bk + K, # stride_bn + N, # stride_cm + 1, # stride_cn + (K + block_k - 1) // block_k, # stride_As_m + 1, # stride_As_k + 1, # stride_Bs_k + (K + block_k - 1) // block_k, # stride_Bs_n + block_n, + block_k, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + num_warps=num_warps, + num_stages=num_stages, + ) + + new_ext_mods = BlockFP8Matmul.attrs["external_mods"] + assert len(new_ext_mods) == 1 + extern_mods.append(new_ext_mods[0]) + return BlockFP8Matmul["tir_w8a8_block_fp8_matmul"], tir_name + + +def get_tir_w8a8_block_fp8_group_matmul( + N: int, + K: int, + num_experts: int, + block_n: int, + block_k: int, + in_dtype: Literal["float8_e4m3fn"], + out_dtype: Literal["float16", "bfloat16"], + BLOCK_SIZE_M: int, + BLOCK_SIZE_N: int, + BLOCK_SIZE_K: int, + GROUP_SIZE_M: int, + num_warps: int, + num_stages: int, + extern_mods: List[tvm.runtime.Module], # noqa: UP006 +): + """Get the TIR function for the w8a8_block_fp8_group_gemm kernel.""" + if triton is None: + raise RuntimeError("Triton is not installed. Please install it with `pip install triton`.") + + name_suffix = ( + f"_N{N}_K{K}_num_experts{num_experts}_block_n{block_n}" + f"_block_k{block_k}_in{in_dtype}_out{out_dtype}" + ) + kernel_name = f"triton_w8a8_block_fp8_group_gemm{name_suffix}" + tir_name = f"tir_w8a8_block_fp8_group_gemm{name_suffix}" + for ext_mod in extern_mods: + if ext_mod.implements_function(kernel_name): + return [None, tir_name] + + triton_kernel = _get_triton_w8a8_block_fp8_group_gemm() + triton_kernel.__name__ = kernel_name + + @I.ir_module + class BlockFP8GroupMatmul: + @T.prim_func(private=True, s_tir=True) + def tir_w8a8_block_fp8_group_gemm( + var_A: T.handle, + var_B: T.handle, + var_As: T.handle, + var_Bs: T.handle, + var_expert_ids: T.handle, + var_indptr: T.handle, + var_C: T.handle, + ): + T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1}) + EM = T.int32() + A = T.match_buffer(var_A, (EM, K), dtype=in_dtype) + B = T.match_buffer(var_B, (num_experts, N, K), dtype=in_dtype) + As = T.match_buffer(var_As, (EM, (K + block_k - 1) // block_k), "float32") + Bs = T.match_buffer( + var_Bs, + ( + num_experts, + (N + block_n - 1) // block_n, + (K + block_k - 1) // block_k, + ), + "float32", + ) + expert_ids = T.match_buffer( + var_expert_ids, + ((EM + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,), + "int32", + ) + indptr = T.match_buffer(var_indptr, (num_experts + 1,), "int32") + C = T.match_buffer(var_C, (EM, N), dtype=out_dtype) + + with T.sblock("root"): + T.reads( + A[0:EM, 0:K], + B[0:num_experts, 0:N, 0:K], + As[0:EM, 0 : (K + block_k - 1) // block_k], + Bs[ + 0:num_experts, + 0 : (N + block_n - 1) // block_n, + 0 : (K + block_k - 1) // block_k, + ], + expert_ids[0 : (EM + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts], + indptr[0 : num_experts + 1], + ) + T.writes(C[0:EM, 0:N]) + T.call_kernel( + triton.jit(triton_kernel), + ((T.ceildiv(EM, BLOCK_SIZE_M) + num_experts) * T.ceildiv(N, BLOCK_SIZE_N),), + A.data, + B.data, + C.data, + As.data, + Bs.data, + expert_ids.data, + indptr.data, + EM, + N, + K, + num_experts, + K, # stride_am + 1, # stride_ak + N * K, # stride_be + 1, # stride_bk + K, # stride_bn + N, # stride_cm + 1, # stride_cn + (K + block_k - 1) // block_k, # stride_asm + 1, # stride_ask + ((N + block_n - 1) // block_n) * ((K + block_k - 1) // block_k), # stride_bse + 1, # stride_bsk + (K + block_k - 1) // block_k, # stride_Bs_n + block_n, + block_k, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + K % BLOCK_SIZE_K == 0, + num_warps=num_warps, + num_stages=num_stages, + ) + + new_ext_mods = BlockFP8GroupMatmul.attrs["external_mods"] + assert len(new_ext_mods) == 1 + extern_mods.append(new_ext_mods[0]) + return BlockFP8GroupMatmul["tir_w8a8_block_fp8_group_gemm"], tir_name + + +def _compute_expert_id_per_block( + indptr: nn.Tensor, + num_experts: int, + M: nn.IntExpr, + BLOCK_SIZE_M: int, +) -> nn.Tensor: + """Compute the expert id for each threadblock (CTA). + We assign an expert id to each threadblock, and the threadblock will + compute the gemm with regard to the specified expert. + + Parameters + ---------- + indptr : nn.Tensor + The indptr tensor of group gemm, with shape of [num_experts + 1,]. + + num_experts : int + The number of total experts. + + M : nn.IntExpr + The number of tokens. + + BLOCK_SIZE_M : int + The block size of the threadblock along the batch dimension. + + Returns + ------- + expert_ids : nn.Tensor + The expert id for each threadblock, with shape of + [(M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,]. + """ + + @T.prim_func(s_tir=True) + def tir_compute_expert_id_per_block( + var_indptr: T.handle, + var_expert_ids: T.handle, + M: T.int64, + ): + T.func_attr({"op_pattern": 8, "tirx.is_scheduled": 1}) + indptr = T.match_buffer(var_indptr, (num_experts + 1,), "int32") + expert_ids = T.match_buffer( + var_expert_ids, + ((M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,), + "int32", + ) + with T.sblock("root"): + for eid in T.thread_binding(0, num_experts, thread="threadIdx.x"): + start_block_id: T.int32 = (indptr[eid] + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + eid + num_blocks: T.int32 = ( + indptr[eid + 1] - indptr[eid] + BLOCK_SIZE_M - 1 + ) // BLOCK_SIZE_M + start_block_id_next_expert: T.int32 = ( + (indptr[eid + 1] + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + eid + 1 + ) + for block_id in T.serial(num_blocks): + expert_ids[start_block_id + block_id] = eid + for block_id in T.serial( + start_block_id_next_expert - (start_block_id + num_blocks) + ): + expert_ids[start_block_id + num_blocks + block_id] = -1 + + assert num_experts <= 1024 + return nn.tensor_ir_op( + tir_compute_expert_id_per_block, + "tir_compute_expert_id_per_block", + args=[indptr, M], + out=nn.Tensor.placeholder( + ((M + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + num_experts,), dtype="int32" + ), + ) + + +def fp8_groupwise_scaled_gemm( + x: nn.Tensor, + x_scale: nn.Tensor, + weight: nn.Tensor, + weight_scale: nn.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + out_dtype: str, +) -> nn.Tensor: + """Triton block-scale fp8 gemm operator. + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [m, k]. + + x_scale : nn.Tensor + The scale tensor, with shape of [m, k // block_size]. + + weight : nn.Tensor + The weight tensor, with shape of [n, k]. + + weight_scale : nn.Tensor + The scale tensor, with shape of [n // block_size, k // block_size]. + + block_size : Tuple[int, int] + The block size. + + out_dtype : str + The data type of the output tensor. + + Returns + ------- + out : nn.Tensor + The output tensor, with shape of [m, n] and dtype of `out_dtype`. + """ + assert x.ndim >= 2 + assert weight.ndim == 2 + assert x_scale.ndim == x.ndim + assert weight_scale.ndim == weight.ndim + assert x.shape[-1] == weight.shape[1] + assert x.shape[:-1] == x_scale.shape[:-1] + assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1] + assert (weight.shape[1] + block_size[1] - 1) // block_size[1] == weight_scale.shape[1] + assert (weight.shape[0] + block_size[0] - 1) // block_size[0] == weight_scale.shape[0] + + if x.dtype != "float8_e4m3fn" or weight.dtype != "float8_e4m3fn": + raise ValueError( + f"x and weight must be float8_e4m3fn, but got x={x.dtype}, weight={weight.dtype}" + ) + if x_scale.dtype != "float32" and weight_scale.dtype != "float32": + raise ValueError( + "x_scale and weight_scale must be float32, but got " + f"x_scale={x_scale.dtype}, weight_scale={weight_scale.dtype}" + ) + if out_dtype not in ["float16", "bfloat16"]: + raise ValueError(f"out_dtype must be float16 or bfloat16, but got {out_dtype}") + + M = x.shape[0] + for i in range(1, x.ndim - 1): + M *= x.shape[i] + N = weight.shape[0] + K = x.shape[-1] + + BLOCK_SIZE_M = 64 + BLOCK_SIZE_N = block_size[0] + BLOCK_SIZE_K = block_size[1] + GROUP_SIZE_M = 32 + num_warps = 4 + num_stages = 3 + + x_shape = x.shape + if x.ndim > 2: + x = x.reshape(M, K) + x_scale = x_scale.reshape(M, x_scale.shape[-1]) + + out = nn.extern( + "mlc.triton.w8a8_block_fp8_matmul", + args=[ + x, + weight, + x_scale, + weight_scale, + N, + K, + block_size[0], + block_size[1], + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + num_warps, + num_stages, + str(x.dtype), + str(out_dtype), + ], + out=nn.Tensor.placeholder((M, N), dtype=out_dtype), + ) + return out.reshape(*x_shape[:-1], N) if len(x_shape) > 2 else out + + +def fp8_groupwise_scaled_group_gemm( + x: nn.Tensor, + x_scale: nn.Tensor, + weight: nn.Tensor, + weight_scale: nn.Tensor, + indptr: nn.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + out_dtype: str, +): + """Triton block-scale fp8 group gemm operator. + + Parameters + ---------- + x : nn.Tensor + The input tensor, with shape of [m, k]. + + x_scale : nn.Tensor + The scale tensor, with shape of [m, k // block_size]. + + weight : nn.Tensor + The weight tensor, with shape of [num_experts, n, k]. + + weight_scale : nn.Tensor + The scale tensor, with shape of [num_experts, n // block_size, k // block_size]. + + indptr : nn.Tensor + The indptr tensor of group gemm, with shape of [num_experts + 1,]. + + block_size : Tuple[int, int] + The block size. + + out_dtype : str + The data type of the output tensor. + + Returns + ------- + out : nn.Tensor + The output tensor, with shape of [m, n] and dtype of `out_dtype`. + """ + assert x.ndim >= 2 + assert weight.ndim == 3 + assert x_scale.ndim == x.ndim + assert weight_scale.ndim == weight.ndim + assert x.shape[-1] == weight.shape[2] + assert (x.shape[-1] + block_size[1] - 1) // block_size[1] == x_scale.shape[-1] + assert (weight.shape[2] + block_size[1] - 1) // block_size[1] == weight_scale.shape[2] + assert (weight.shape[1] + block_size[0] - 1) // block_size[0] == weight_scale.shape[1] + + num_experts = weight.shape[0] + M = x.shape[0] + for i in range(1, x.ndim - 1): + M *= x.shape[i] + N = weight.shape[1] + K = x.shape[-1] + assert weight_scale.shape[0] == num_experts + assert indptr.ndim == 1 + assert indptr.shape[0] == num_experts + 1 + + BLOCK_SIZE_M = 64 + BLOCK_SIZE_N = block_size[0] + BLOCK_SIZE_K = block_size[1] + GROUP_SIZE_M = 32 + num_warps = 4 + num_stages = 3 + + x_shape = x.shape + if x.ndim > 2: + x = x.reshape(M, K) + x_scale = x_scale.reshape(M, x_scale.shape[-1]) + expert_ids = _compute_expert_id_per_block(indptr, num_experts, M, BLOCK_SIZE_M) + + out = nn.extern( + "mlc.triton.w8a8_block_fp8_group_matmul", + args=[ + x, + weight, + x_scale, + weight_scale, + expert_ids, + indptr, + N, + K, + num_experts, + block_size[0], + block_size[1], + BLOCK_SIZE_M, + BLOCK_SIZE_N, + BLOCK_SIZE_K, + GROUP_SIZE_M, + num_warps, + num_stages, + str(x.dtype), + str(out_dtype), + ], + out=nn.Tensor.placeholder((M, N), dtype=out_dtype), + ) + return out.reshape(*x_shape[:-1], N) if len(x_shape) > 2 else out diff --git a/python/mlc_llm/protocol/__init__.py b/python/mlc_llm/protocol/__init__.py new file mode 100644 index 0000000..b430746 --- /dev/null +++ b/python/mlc_llm/protocol/__init__.py @@ -0,0 +1,9 @@ +"""Definitions of pydantic models for API entry points and configurations + +Note +---- +We use the following convention + +- filename_protocol If the classes can appear in an API endpoint +- filename_config For other config classes +""" diff --git a/python/mlc_llm/protocol/conversation_protocol.py b/python/mlc_llm/protocol/conversation_protocol.py new file mode 100644 index 0000000..72afdd3 --- /dev/null +++ b/python/mlc_llm/protocol/conversation_protocol.py @@ -0,0 +1,278 @@ +"""The standard conversation protocol in MLC LLM""" + +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union # noqa: UP035 + +from pydantic import BaseModel, Field, field_validator + + +# The message placeholders in the message prompts according to roles. +class MessagePlaceholders(Enum): + """The message placeholders in the message prompts according to roles.""" + + SYSTEM = "{system_message}" + USER = "{user_message}" + ASSISTANT = "{assistant_message}" + TOOL = "{tool_message}" + FUNCTION = "{function_string}" + + +T = TypeVar("T", bound="BaseModel") + + +class Conversation(BaseModel): + """Class that specifies the convention template of conversation + and contains the conversation history. + + Given a conversation template, the corresponding prompt generated out + from it is usually in the following format: + + <><><><><> + <><><><> + ... + <><><><> + <><> + """ + + # Optional name of the template. + name: Optional[str] = None + # The system prompt template, it optionally contains the system + # message placeholder, and the placeholder will be replaced with + # the system message below. + system_template: str = MessagePlaceholders.SYSTEM.value + # The content of the system prompt (without the template format). + system_message: str = "" + # The system token ids to be prepended at the beginning of tokenized + # generated prompt. + system_prefix_token_ids: Optional[List[int]] = None # noqa: UP006 + # Whether or not to append user role and separator after the system message. + # This is mainly for [INST] [/INST] style prompt format + add_role_after_system_message: bool = True + + # The conversation roles + roles: Dict[str, str] # noqa: UP006 + + # The roles prompt template, it optionally contains the defaults + # message placeholders and will be replaced by actual content + role_templates: Dict[str, str] # noqa: UP006 + + # The conversation history messages. + # Each message is a pair of strings, denoting "(role, content)". + # The content can be None. + messages: List[Tuple[str, Optional[Union[str, List[Dict]]]]] = Field(default_factory=lambda: []) # noqa: UP006 + + # The separators between messages when concatenating into a single prompt. + # List size should be either 1 or 2. + # - When size is 1, the separator will be used between adjacent messages. + # - When size is 2, seps[0] is used after user message, and + # seps[1] is used after assistant message. + seps: List[str] # noqa: UP006 + + # The separator between the role and the content in a message. + role_content_sep: str = "" + # The separator between the role and empty contents. + role_empty_sep: str = "" + + # The stop criteria + stop_str: List[str] = Field(default_factory=lambda: []) # noqa: UP006 + stop_token_ids: List[int] = Field(default_factory=lambda: []) # noqa: UP006 + + # When True, strip `...` blocks (and any trailing whitespace) + # from historical assistant messages before rendering the prompt, mirroring + # Qwen3's official HF chat template. Only historical turns before the last + # user message are affected; reasoning on the most recent assistant turn is + # preserved for tool-call prefill scenarios. + strip_reasoning_in_history: bool = False + + # Function call fields + function_string: str = "" + # whether using function calling or not, helps check for output message format in API call + use_function_calling: bool = False + + def __init__(self, role_templates: Optional[Dict[str, str]] = None, **kwargs): # noqa: UP006 + # Defaults templates which would be overridden by model specific templates + _role_templates: Dict[str, str] = { # noqa: UP006 + "user": MessagePlaceholders.USER.value, + "assistant": MessagePlaceholders.ASSISTANT.value, + "tool": MessagePlaceholders.TOOL.value, + } + if role_templates is not None: + _role_templates.update(role_templates) + super().__init__(role_templates=_role_templates, **kwargs) + + @field_validator("seps") + @classmethod + def check_message_seps(cls, seps: List[str]) -> List[str]: # noqa: UP006 + """Check if the input message separators has size 1 or 2.""" + if len(seps) == 0 or len(seps) > 2: + raise ValueError("seps should have size 1 or 2.") + return seps + + def to_json_dict(self) -> Dict[str, Any]: # noqa: UP006 + """Convert to a json dictionary""" + return self.model_dump(by_alias=True, exclude_none=True) + + @classmethod + def from_json_dict(cls: Type[T], json_dict: Dict[str, Any]) -> T: # noqa: UP006 + """Convert from a json dictionary""" + return Conversation.model_validate(json_dict) + + def as_prompt(self, config=None) -> List[Any]: # noqa: UP006 + """Convert the conversation template and history messages to + a single prompt. + + Returns + ------- + prompts : List[Union[str, "mlc_llm.serve.data.Data"]] + The prompts converted from the conversation messages. + We use Any in the signature to avoid cyclic import. + """ + from ..serve import data + + # - Get the system message. + system_msg = self.system_template.replace( + MessagePlaceholders.SYSTEM.value, self.system_message + ) + + # - Get the message strings. + message_list: List[Union[str, data.Data]] = [] # noqa: UP006 + separators = list(self.seps) + if len(separators) == 1: + separators.append(separators[0]) + + if system_msg != "": + message_list.append(system_msg) + + messages = ( + _strip_reasoning_in_history(self.messages) + if self.strip_reasoning_in_history + else self.messages + ) + + for i, (role, content) in enumerate(messages): + if role not in self.roles.keys(): + raise ValueError(f'Role "{role}" is not a supported role in {self.roles.keys()}') + separator = separators[role == "assistant"] # check assistant role + + if content is None: + message_list.append(self.roles[role] + self.role_empty_sep) + continue + + role_prefix = ( + "" + # Do not append role prefix if this is the first message and there + # is already a system message + if (not self.add_role_after_system_message and system_msg != "" and i == 0) + else self.roles[role] + self.role_content_sep + ) + if isinstance(content, str): + message_list.append( + role_prefix + + self.role_templates[role].replace( + MessagePlaceholders[role.upper()].value, content + ) + + separator + ) + continue + + message_list.append(role_prefix) + + for item in content: + assert isinstance(item, dict), "Content should be a string or a list of dicts" + assert "type" in item, "Content item should have a type field" + if item["type"] == "text": + message = self.role_templates[role].replace( + MessagePlaceholders[role.upper()].value, item["text"] + ) + message_list.append(message) + elif item["type"] == "image_url": + assert config is not None, "Model config is required" + image_url = _get_url_from_item(item) + message_list.append(data.ImageData.from_url(image_url, config)) + message_list.append("\n") + else: + raise ValueError(f"Unsupported content type: {item['type']}") + + message_list.append(separator) + + prompt = _combine_consecutive_messages(message_list) + + if not any(isinstance(item, data.ImageData) for item in message_list): + # Replace the last function string placeholder with actual function string + prompt[0] = self.function_string.join( + prompt[0].rsplit(MessagePlaceholders.FUNCTION.value, 1) + ) + # Replace with remaining function string placeholders with empty string + prompt[0] = prompt[0].replace(MessagePlaceholders.FUNCTION.value, "") + + return prompt + + +def _get_url_from_item(item: Dict) -> str: # noqa: UP006 + image_url: str + assert "image_url" in item, "Content item should have an image_url field" + if isinstance(item["image_url"], str): + image_url = item["image_url"] + elif isinstance(item["image_url"], dict): + assert "url" in item["image_url"], ( + "Content image_url item should be a string or a dict with a url field" + ) + image_url = item["image_url"]["url"] + else: + raise ValueError( + "Content image_url item type not supported. " + "Should be a string or a dict with a url field." + ) + return image_url + + +def _strip_reasoning_in_history( + messages: List[Tuple[str, Optional[Union[str, List[Dict]]]]], # noqa: UP006 +) -> List[Tuple[str, Optional[Union[str, List[Dict]]]]]: # noqa: UP006 + """Strip `...` blocks from assistant messages that precede + the last user message, matching Qwen3's HF chat-template behavior. The last + assistant message (if any) is preserved so tool-call prefill continuations + keep their reasoning context. + """ + last_user_idx = -1 + for i, (role, _) in enumerate(messages): + if role == "user": + last_user_idx = i + + result: List[Tuple[str, Optional[Union[str, List[Dict]]]]] = [] # noqa: UP006 + for i, (role, content) in enumerate(messages): + if ( + role == "assistant" + and i < last_user_idx + and isinstance(content, str) + and "" in content + ): + content = content.split("")[-1].lstrip("\n") + result.append((role, content)) + return result + + +def _combine_consecutive_messages(messages: List[Any]) -> List[Any]: # noqa: UP006 + """Combining consecutive strings into one. + + Parameters + ---------- + messages : List[Union[str, "mlc_llm.serve.data.Data"]] + The input messages to be combined. + We use Any in the signature to avoid cyclic import. + + Returns + ------- + updated_messages : List[Union[str, "mlc_llm.serve.data.Data"]] + The combined messages + """ + if len(messages) == 0: + return [] + + combined_messages = [messages[0]] + for message in messages[1:]: + if isinstance(message, str) and isinstance(combined_messages[-1], str): + combined_messages[-1] += message + else: + combined_messages.append(message) + return combined_messages diff --git a/python/mlc_llm/protocol/debug_protocol.py b/python/mlc_llm/protocol/debug_protocol.py new file mode 100644 index 0000000..f233ce6 --- /dev/null +++ b/python/mlc_llm/protocol/debug_protocol.py @@ -0,0 +1,49 @@ +"""Debug protocols in MLC LLM""" + +from typing import Literal, Optional + +from pydantic import BaseModel + + +class DisaggConfig(BaseModel): + """The class of metadata used in microserving APIs.""" + + kind: Optional[Literal["prepare_receive", "remote_send", "start_generation"]] = None + # "kv_append_metadata" is base64-encoded and is thus a string. + kv_append_metadata: Optional[str] = None + # "kv_window_begin" and "kv_window_end" denote the KV interval of interests. + # "kv_window_end" supports Python style negative indexing. + # The concrete meaning varies for different special request kind: + # - For "prepare_receive", the begin is always 0, and "[0:end]" denotes + # the KV range to prefill on a prefill instance. + # - For "remote_send", "[begin:end]" means the KV range to compute prefill + # and send to the decode instance. + # - For "start_generation", the end is always None, and "[begin:]" denotes + # the KV range to prefill locally on the decode instance. + kv_window_begin: Optional[int] = None + kv_window_end: Optional[int] = None + # KV data destination group offset + dst_group_offset: Optional[int] = None + + +class DebugConfig(BaseModel): + """The class of debug options. + + These optionals are available to engine + but won't be available to serving endpoint + unless an explicit --enable-debug passed + """ + + ignore_eos: bool = False + pinned_system_prompt: bool = False + special_request: Optional[Literal["query_engine_metrics"]] = None + grammar_execution_mode: Literal["constraint", "jump_forward"] = "jump_forward" + disagg_config: Optional[DisaggConfig] = None + + """Special request indicators + + Special requests are handled by engine differently and do not go + through the normal engine step flow. + + The results to these requests are returned as field of "usage" + """ diff --git a/python/mlc_llm/protocol/error_protocol.py b/python/mlc_llm/protocol/error_protocol.py new file mode 100644 index 0000000..1dd1aaf --- /dev/null +++ b/python/mlc_llm/protocol/error_protocol.py @@ -0,0 +1,35 @@ +"""Error protocols in MLC LLM""" + +from http import HTTPStatus +from typing import Optional + +import fastapi +from pydantic import BaseModel + + +class BadRequestError(ValueError): + """The exception for bad requests in engines.""" + + def __init__(self, *args: object) -> None: + super().__init__(*args) + + +class ErrorResponse(BaseModel): + """The class of error response.""" + + object: str = "error" + message: str + code: Optional[int] = None + + +def create_error_response(status_code: HTTPStatus, message: str) -> fastapi.responses.JSONResponse: + """Create a JSON response that reports error with regarding the input message.""" + return fastapi.responses.JSONResponse( + ErrorResponse(message=message, code=status_code.value).model_dump_json(by_alias=True), + status_code=status_code.value, + ) + + +async def bad_request_error_handler(_request: fastapi.Request, e: BadRequestError): + """The handler of BadRequestError that converts an exception into error response.""" + return create_error_response(status_code=HTTPStatus.BAD_REQUEST, message=e.args[0]) diff --git a/python/mlc_llm/protocol/generation_config.py b/python/mlc_llm/protocol/generation_config.py new file mode 100644 index 0000000..ca828b6 --- /dev/null +++ b/python/mlc_llm/protocol/generation_config.py @@ -0,0 +1,32 @@ +"""Low-level generation config class""" + +from typing import Dict, List, Optional # noqa: UP035 + +from pydantic import BaseModel + +from .debug_protocol import DebugConfig +from .openai_api_protocol import RequestResponseFormat + + +class GenerationConfig(BaseModel): + """The generation configuration dataclass. + + This is a config class used by Engine internally. + """ + + n: int = 1 + temperature: Optional[float] = None + top_p: Optional[float] = None + frequency_penalty: Optional[float] = None + presence_penalty: Optional[float] = None + repetition_penalty: Optional[float] = None + logprobs: bool = False + top_logprobs: int = 0 + logit_bias: Optional[Dict[int, float]] = None # noqa: UP006 + # internally we use -1 to represent infinite + max_tokens: int = -1 + seed: Optional[int] = None + stop_strs: Optional[List[str]] = None # noqa: UP006 + stop_token_ids: Optional[List[int]] = None # noqa: UP006 + response_format: Optional[RequestResponseFormat] = None + debug_config: Optional[Optional[DebugConfig]] = None diff --git a/python/mlc_llm/protocol/microserving_protocol.py b/python/mlc_llm/protocol/microserving_protocol.py new file mode 100644 index 0000000..83bcc43 --- /dev/null +++ b/python/mlc_llm/protocol/microserving_protocol.py @@ -0,0 +1,72 @@ +"""Protocols in MLC LLM for MicroServing.""" + +from pydantic import BaseModel + +from mlc_llm.protocol.openai_api_protocol import CompletionRequest + + +class PrepRecvRequest(CompletionRequest): + """The extra request body for prep_recv request in MicroServing. + + Attributes + ---------- + kv_window_end : int + [0, kv_window_end] denotes the KV range of the prompt to prefill on + a prefill instance. + The entries of this KV range will be allocated on the decode instance. + """ + + end: int + + +class PrepRecvResponse(BaseModel): + """The response body for prep_recv request in MicroServing. + + Attributes + ---------- + prefix_matched_length : int + The matched common prefix length on the decode instance when + prefix cache is enabled, or 0 if there is no prefix cache. + + kv_append_metadata : str + The metadata of the KV range on the destination decode instance. + """ + + kv_append_metadata: str + prefix_matched_length: int + + +class RemoteSendRequest(CompletionRequest): + """The extra request body for remote_send request in MicroServing. + + Attributes + ---------- + kv_window_begin : int + Denote the start of the KV range to prefill. + + kv_window_end : int + Denote the end of the KV range to prefill. + + kv_append_metadata : str + The metadata of the KV range on the destination decode instance. + + dst_group_offset : int + The node group offset of the destination decode instance. + """ + + begin: int + end: int + kv_addr_info: str + recv_rank: int + + +class StartGenerateRequest(CompletionRequest): + """The extra request body for start_generate request in MicroServing. + + Attributes + ---------- + kv_window_begin : int + Denote the start of the KV range to prefill on the decode instance. + """ + + begin: int diff --git a/python/mlc_llm/protocol/mlc_chat_config.py b/python/mlc_llm/protocol/mlc_chat_config.py new file mode 100644 index 0000000..86df18c --- /dev/null +++ b/python/mlc_llm/protocol/mlc_chat_config.py @@ -0,0 +1,78 @@ +"""Schema for mlc-chat-config""" + +from typing import Any, Dict, List, Literal, Optional, Union # noqa: UP035 + +from pydantic import BaseModel, Field + +from mlc_llm.support.constants import MLC_CHAT_CONFIG_VERSION + +from .conversation_protocol import Conversation + +MLC_CHAT_SYSTEM_DEFAULT = { + "pad_token_id": 0, + "bos_token_id": 1, + "eos_token_id": 2, + "temperature": 1.0, + "presence_penalty": 0.0, + "frequency_penalty": 0.0, + "repetition_penalty": 1.0, + "top_p": 1.0, +} +"""system default values.""" + + +class MLCChatConfig(BaseModel): + """Fields in the dumped `mlc-chat-config.json` file.""" + + # Version control + version: str = MLC_CHAT_CONFIG_VERSION + + # use alias to avoid protected namespace conflict with pydantic + field_model_type: str = Field(alias="model_type") + quantization: str + # use alias to avoid protected namespace conflict with pydantic + field_model_config: Dict[str, Any] = Field(alias="model_config") # noqa: UP006 + vocab_size: int + context_window_size: int + sliding_window_size: int + prefill_chunk_size: int + attention_sink_size: int + tensor_parallel_shards: int + pipeline_parallel_stages: int = 1 + # Configuration of text generation + active_vocab_size: int = None + temperature: Optional[float] = None + presence_penalty: Optional[float] = None + frequency_penalty: Optional[float] = None + repetition_penalty: Optional[float] = None + top_p: Optional[float] = None + # Tokenizer configuration + tokenizer_files: List[str] = Field(default_factory=list) # noqa: UP006 + # The content of tokenizer.TokenizerInfo + tokenizer_info: Dict[str, Any] = Field(default_factory=dict) # noqa: UP006 + # conversation template + conv_template: Conversation + # extra fields from generation_config.json + # NOTE: they are not being used for now in MLCEngine + # but we keep them for book-keep purposes + pad_token_id: Optional[int] = None + bos_token_id: Optional[int] = None + eos_token_id: Optional[Union[int, List[int]]] = None # noqa: UP006 + + field_model_task: Literal["chat", "embedding"] = Field(default="chat", alias="model_task") + embedding_metadata: Optional[Dict[str, Any]] = None # noqa: UP006 + + def get_system_defaults_for_missing_fields(self) -> Dict[str, Any]: # noqa: UP006 + """Apply system default value for fields that are None + + Note + ---- + We implement default setting in this way so we can lazily create + MLCChatConfig, override its optional values then + apply_system_defaults in the end. + """ + res = {} + for key, value in MLC_CHAT_SYSTEM_DEFAULT.items(): + if getattr(self, key) is None: + res[key] = value + return res diff --git a/python/mlc_llm/protocol/openai_api_protocol.py b/python/mlc_llm/protocol/openai_api_protocol.py new file mode 100644 index 0000000..9471a7b --- /dev/null +++ b/python/mlc_llm/protocol/openai_api_protocol.py @@ -0,0 +1,462 @@ +"""Protocols in MLC LLM for OpenAI API. +Adapted from FastChat's OpenAI protocol: +https://github.com/lm-sys/FastChat/blob/main/fastchat/protocol/openai_api_protocol.py +""" + +import json +import time +from typing import Any, Dict, List, Literal, Optional, Tuple, Union # noqa: UP035 + +import shortuuid +from pydantic import BaseModel, Field, field_validator, model_validator + +from .conversation_protocol import Conversation +from .debug_protocol import DebugConfig +from .error_protocol import BadRequestError + +################ Commons ################ + + +# OPenAI API compatible limits +CHAT_COMPLETION_MAX_TOP_LOGPROBS = 20 +COMPLETION_MAX_TOP_LOGPROBS = 5 + + +class ListResponse(BaseModel): + object: str = "list" + data: List[Any] # noqa: UP006 + + +class TopLogProbs(BaseModel): + token: str + logprob: float + bytes: Optional[List[int]] # noqa: UP006 + + +class LogProbsContent(BaseModel): + token: str + logprob: float + bytes: Optional[List[int]] # noqa: UP006 + top_logprobs: List[TopLogProbs] = [] # noqa: UP006 + + +class LogProbs(BaseModel): + content: List[LogProbsContent] # noqa: UP006 + + +class CompletionLogProbs(BaseModel): + # The position of the token in the concatenated str: prompt + completion_text + # TODO(vvchernov): skip optional after support + text_offset: Optional[List[int]] # noqa: UP006 + token_logprobs: List[float] # noqa: UP006 + tokens: List[str] # noqa: UP006 + top_logprobs: List[Dict[str, float]] # noqa: UP006 + + +class CompletionUsage(BaseModel): + prompt_tokens: int + completion_tokens: int + total_tokens: int + extra: Optional[Dict[str, Any]] = None # noqa: UP006 + """Extra metrics and info that may be returned by debug_config + """ + + +class StreamOptions(BaseModel): + include_usage: Optional[bool] + + +################ v1/embeddings ################ + + +class EmbeddingRequest(BaseModel): + """OpenAI "v1/embeddings" request protocol. + API reference: https://platform.openai.com/docs/api-reference/embeddings/create + """ + + input: Union[str, List[str], List[int], List[List[int]]] # noqa: UP006 + model: Optional[str] = None + encoding_format: Literal["float", "base64"] = "float" + dimensions: Optional[int] = None + user: Optional[str] = None + + @field_validator("input") + @classmethod + def validate_input(cls, v): + """Check that the input is not an empty list. + + Note: empty strings are allowed — encoder models produce valid + embeddings from [CLS]+[SEP] tokens alone. + """ + if isinstance(v, list) and len(v) == 0: + raise ValueError("Input list must not be empty.") + return v + + +class EmbeddingObject(BaseModel): + object: str = "embedding" + embedding: Union[List[float], str] # noqa: UP006 + index: int + + +class EmbeddingUsage(BaseModel): + prompt_tokens: int + total_tokens: int + + +class EmbeddingResponse(BaseModel): + """OpenAI "v1/embeddings" response protocol. + API reference: https://platform.openai.com/docs/api-reference/embeddings/object + """ + + object: str = "list" + data: List[EmbeddingObject] # noqa: UP006 + model: Optional[str] = None + usage: EmbeddingUsage + + +################ v1/models ################ + + +class ModelResponse(BaseModel): + """OpenAI "v1/models" response protocol. + API reference: https://platform.openai.com/docs/api-reference/models/object + """ + + id: str + created: int = Field(default_factory=lambda: int(time.time())) + object: str = "model" + owned_by: str = "MLC-LLM" + + +################ v1/completions ################ + + +class RequestResponseFormat(BaseModel): + type: Literal["text", "json_object"] = "text" + json_schema: Optional[str] = Field(default=None, alias="schema") + """This field is named json_schema instead of schema because BaseModel defines a method called + schema. During construction of RequestResponseFormat, key "schema" still should be used: + `RequestResponseFormat(type="json_object", schema="{}")` + """ + + +class CompletionRequest(BaseModel): + """OpenAI completion request protocol. + API reference: https://platform.openai.com/docs/api-reference/completions/create + """ + + model: Optional[str] = None + prompt: Union[str, List[int]] # noqa: UP006 + best_of: int = 1 + echo: bool = False + frequency_penalty: Optional[float] = None + presence_penalty: Optional[float] = None + logprobs: Optional[int] = None + logit_bias: Optional[Dict[int, float]] = None # noqa: UP006 + max_tokens: Optional[int] = None + n: int = 1 + seed: Optional[int] = None + stop: Optional[Union[str, List[str]]] = None # noqa: UP006 + stream: bool = False + stream_options: Optional[StreamOptions] = None + suffix: Optional[str] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + user: Optional[str] = None + response_format: Optional[RequestResponseFormat] = None + debug_config: Optional[DebugConfig] = None + + @field_validator("frequency_penalty", "presence_penalty") + @classmethod + def check_penalty_range(cls, penalty_value: Optional[float]) -> Optional[float]: + """Check if the penalty value is in range [-2, 2].""" + if penalty_value and (penalty_value < -2 or penalty_value > 2): + raise ValueError("Penalty value should be in range [-2, 2].") + return penalty_value + + @field_validator("logit_bias") + @classmethod + def check_logit_bias( + cls, + logit_bias_value: Optional[Dict[int, float]], # noqa: UP006 + ) -> Optional[Dict[int, float]]: # noqa: UP006 + """Check if the logit bias key is given as an integer.""" + if logit_bias_value is None: + return None + for token_id, bias in logit_bias_value.items(): + if abs(bias) > 100: + raise ValueError( + "Logit bias value should be in range [-100, 100], while value " + f"{bias} is given for token id {token_id}" + ) + return logit_bias_value + + @model_validator(mode="after") + def check_logprobs(self) -> "CompletionRequest": + """Check if the logprobs requirements are valid.""" + if self.logprobs is not None and ( + self.logprobs < 0 or self.logprobs > COMPLETION_MAX_TOP_LOGPROBS + ): + raise ValueError(f'"logprobs" must be in range [0, {COMPLETION_MAX_TOP_LOGPROBS}]') + return self + + +class CompletionResponseChoice(BaseModel): + finish_reason: Optional[Literal["stop", "length", "preempt"]] = None + index: int = 0 + logprobs: Optional[CompletionLogProbs] = None + text: str + + +class CompletionResponse(BaseModel): + """OpenAI completion response protocol. + API reference: https://platform.openai.com/docs/api-reference/completions/object + """ + + id: str + choices: List[CompletionResponseChoice] # noqa: UP006 + created: int = Field(default_factory=lambda: int(time.time())) + model: Optional[str] = None + object: str = "text_completion" + usage: Optional[CompletionUsage] = None + + +################ v1/chat/completions ################ + + +class ChatFunction(BaseModel): + description: Optional[str] = None + name: str + parameters: Dict # noqa: UP006 + + +class ChatTool(BaseModel): + type: Literal["function"] + function: ChatFunction + + +class ChatFunctionCall(BaseModel): + name: str + arguments: Union[None, Dict[str, Any]] = None # noqa: UP006 + + +class ChatToolCall(BaseModel): + id: str = Field(default_factory=lambda: f"call_{shortuuid.random()}") + type: Literal["function"] + function: ChatFunctionCall + + +class ChatCompletionMessage(BaseModel): + content: Optional[Union[str, List[Dict]]] = None # noqa: UP006 + role: Literal["system", "user", "assistant", "tool"] + name: Optional[str] = None + tool_calls: Optional[List[ChatToolCall]] = None # noqa: UP006 + tool_call_id: Optional[str] = None + + +class ChatCompletionRequest(BaseModel): + """OpenAI chat completion request protocol. + API reference: https://platform.openai.com/docs/api-reference/chat/create + """ + + messages: List[ChatCompletionMessage] # noqa: UP006 + model: Optional[str] = None + frequency_penalty: Optional[float] = None + presence_penalty: Optional[float] = None + logprobs: bool = False + top_logprobs: int = 0 + logit_bias: Optional[Dict[int, float]] = None # noqa: UP006 + max_tokens: Optional[int] = None + n: int = 1 + seed: Optional[int] = None + stop: Optional[Union[str, List[str]]] = None # noqa: UP006 + stream: bool = False + stream_options: Optional[StreamOptions] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + tools: Optional[List[ChatTool]] = None # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None # noqa: UP006 + user: Optional[str] = None + response_format: Optional[RequestResponseFormat] = None + # NOTE: debug_config is not part of OpenAI protocol + # we add it to enable extra debug options + debug_config: Optional[DebugConfig] = None + + @field_validator("frequency_penalty", "presence_penalty") + @classmethod + def check_penalty_range(cls, penalty_value: Optional[float]) -> Optional[float]: + """Check if the penalty value is in range [-2, 2].""" + if penalty_value and (penalty_value < -2 or penalty_value > 2): + raise ValueError("Penalty value should be in range [-2, 2].") + return penalty_value + + @field_validator("logit_bias") + @classmethod + def check_logit_bias( + cls, + logit_bias_value: Optional[Dict[int, float]], # noqa: UP006 + ) -> Optional[Dict[int, float]]: # noqa: UP006 + """Check if the logit bias key is given as an integer.""" + if logit_bias_value is None: + return None + for token_id, bias in logit_bias_value.items(): + if abs(bias) > 100: + raise ValueError( + "Logit bias value should be in range [-100, 100], while value " + f"{bias} is given for token id {token_id}" + ) + return logit_bias_value + + @model_validator(mode="after") + def check_logprobs(self) -> "ChatCompletionRequest": + """Check if the logprobs requirements are valid.""" + if self.top_logprobs < 0 or self.top_logprobs > CHAT_COMPLETION_MAX_TOP_LOGPROBS: + raise ValueError( + f'"top_logprobs" must be in range [0, {CHAT_COMPLETION_MAX_TOP_LOGPROBS}]' + ) + if not self.logprobs and self.top_logprobs > 0: + raise ValueError('"logprobs" must be True to support "top_logprobs"') + return self + + @model_validator(mode="after") + def check_stream_options(self) -> "ChatCompletionRequest": + """Check stream options""" + if self.stream_options is None: + return self + if not self.stream: + raise ValueError("stream must be set to True when stream_options is present") + return self + + @model_validator(mode="after") + def check_debug_config(self) -> "ChatCompletionRequest": + """Check debug config""" + if self.debug_config is None: + return self + + if self.debug_config.special_request is None: + return self + + if not self.stream: + raise ValueError("DebugConfig.special_request requires stream=True") + + if self.stream_options is None or not self.stream_options.include_usage: + raise ValueError("DebugConfig.special_request requires include_usage in stream_options") + + return self + + def check_message_validity(self) -> None: + """Check if the given chat messages are valid. Return error message if invalid.""" + for i, message in enumerate(self.messages): + if message.role == "system" and i != 0: + raise BadRequestError( + f"System prompt at position {i} in the message list is invalid." + ) + if message.tool_call_id is not None: + if message.role != "tool": + raise BadRequestError("Non-tool message having `tool_call_id` is invalid.") + if isinstance(message.content, list): + if message.role != "user": + raise BadRequestError("Non-user message having a list of content is invalid.") + if message.tool_calls is not None: + if message.role != "assistant": + raise BadRequestError("Non-assistant message having `tool_calls` is invalid.") + raise BadRequestError("Assistant message having `tool_calls` is not supported yet.") + + def check_function_call_usage(self, conv_template: Conversation) -> None: + """Check if function calling is used and update the conversation template. + Return error message if invalid request format for function calling. + """ + + # return if no tools are provided or tool_choice is set to none + if self.tools is None or (isinstance(self.tool_choice, str) and self.tool_choice == "none"): + conv_template.use_function_calling = False + return + + # select the tool based on the tool_choice if specified + if isinstance(self.tool_choice, dict): + if self.tool_choice["type"] != "function": + raise BadRequestError("Only 'function' tool choice is supported") + + if len(self.tool_choice["function"]) > 1: + raise BadRequestError("Only one tool is supported when tool_choice is specified") + + for tool in self.tools: + if tool.function.name == self.tool_choice["function"]["name"]: + conv_template.use_function_calling = True + conv_template.function_string = tool.function.model_dump_json(by_alias=True) + return + + raise BadRequestError( + f"The tool_choice function {self.tool_choice['function']['name']}" + " is not found in the tools list" + ) + + if isinstance(self.tool_choice, str) and self.tool_choice != "auto": + raise BadRequestError(f"Invalid tool_choice value: {self.tool_choice}") + + function_list = [] + for tool in self.tools: + if tool.type != "function": + raise BadRequestError("Only 'function' tool type is supported") + function_list.append(tool.function.model_dump(by_alias=True)) + + conv_template.use_function_calling = True + conv_template.function_string = json.dumps(function_list) + + +class ChatCompletionResponseChoice(BaseModel): + finish_reason: Optional[Literal["stop", "length", "tool_calls", "error"]] = None + index: int = 0 + message: ChatCompletionMessage + logprobs: Optional[LogProbs] = None + + +class ChatCompletionStreamResponseChoice(BaseModel): + finish_reason: Optional[Literal["stop", "length", "tool_calls", "error"]] = None + index: int = 0 + delta: ChatCompletionMessage + logprobs: Optional[LogProbs] = None + + +class ChatCompletionResponse(BaseModel): + """OpenAI completion response protocol. + API reference: https://platform.openai.com/docs/api-reference/chat/object + """ + + id: str + choices: List[ChatCompletionResponseChoice] # noqa: UP006 + created: int = Field(default_factory=lambda: int(time.time())) + model: Optional[str] = None + system_fingerprint: str + object: Literal["chat.completion"] = "chat.completion" + usage: Optional[CompletionUsage] = None + + +class ChatCompletionStreamResponse(BaseModel): + """OpenAI completion stream response protocol. + API reference: https://platform.openai.com/docs/api-reference/chat/streaming + """ + + id: str + choices: List[ChatCompletionStreamResponseChoice] # noqa: UP006 + created: int = Field(default_factory=lambda: int(time.time())) + model: Optional[str] = None + system_fingerprint: str + object: Literal["chat.completion.chunk"] = "chat.completion.chunk" + usage: Optional[CompletionUsage] = None + + +def openai_api_get_unsupported_fields( + request: Union[CompletionRequest, ChatCompletionRequest], +) -> List[str]: # noqa: UP006 + """Get the unsupported fields in the request.""" + unsupported_field_default_values: List[Tuple[str, Any]] = [ # noqa: UP006 + ("best_of", 1), + ] + + unsupported_fields: List[str] = [] # noqa: UP006 + for field, value in unsupported_field_default_values: + if hasattr(request, field) and getattr(request, field) != value: + unsupported_fields.append(field) + return unsupported_fields diff --git a/python/mlc_llm/quantization/__init__.py b/python/mlc_llm/quantization/__init__.py new file mode 100644 index 0000000..a4bb1af --- /dev/null +++ b/python/mlc_llm/quantization/__init__.py @@ -0,0 +1,11 @@ +"""A subpackage for quantization and dequantization algorithms""" + +from .awq_quantization import AWQQuantize +from .block_scale_quantization import BlockScaleQuantize +from .fp8_quantization import FP8PerTensorQuantizeMixtralExperts +from .ft_quantization import FTQuantize +from .group_quantization import GroupQuantize +from .model_quantization import make_awq_quant, make_quantization_functions +from .no_quantization import NoQuantize +from .per_tensor_quantization import PerTensorQuantize +from .quantization import QUANTIZATION, Quantization diff --git a/python/mlc_llm/quantization/awq_quantization.py b/python/mlc_llm/quantization/awq_quantization.py new file mode 100644 index 0000000..7b70b7b --- /dev/null +++ b/python/mlc_llm/quantization/awq_quantization.py @@ -0,0 +1,282 @@ +"""AWQ Quantization""" + +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional # noqa: UP035 + +from tvm import DataType, DataTypeCode, te, tirx, topi +from tvm.relax.frontend import nn +from tvm.runtime import Tensor + +from mlc_llm.loader import QuantizeMapping + +from .utils import convert_uint_to_float, is_final_fc, is_moe_gate + + +def _make_divisible(c, divisor): + return (c + divisor - 1) // divisor + + +def _calculate_zeros_width(in_features, group_size=128, pack_num=8): + if group_size >= 128: + size_multiplier = 1 + elif group_size == 64: + size_multiplier = 2 + elif group_size == 32: + size_multiplier = 4 + else: + raise NotImplementedError + + base_width = _make_divisible(in_features // group_size, pack_num) + base_width = _make_divisible(base_width, size_multiplier) * size_multiplier + return base_width + + +@dataclass +class AWQQuantize: + """Configuration for AWQ quantization""" + + name: str + kind: str + group_size: int + quantize_dtype: str # "int3", "int4", "int8" + storage_dtype: str # "uint32" + model_dtype: str # "float16", "float32" + + num_elem_per_storage: int = 0 + num_storage_per_group: int = 0 + max_int_value: int = 0 + + prebuilt_quantize_func: Dict[str, Callable[[Tensor], Tensor]] = field( # noqa: UP006 + default_factory=lambda: {} + ) + + def __post_init__(self): + assert self.kind == "awq" + quantize_dtype = DataType(self.quantize_dtype) + storage_dtype = DataType(self.storage_dtype) + model_dtype = DataType(self.model_dtype) + assert quantize_dtype.type_code == DataTypeCode.INT + assert storage_dtype.type_code == DataTypeCode.UINT + assert model_dtype.type_code == DataTypeCode.FLOAT + if storage_dtype.bits < quantize_dtype.bits: + raise ValueError("Storage unit should be greater or equal to quantized element") + + self.num_elem_per_storage = storage_dtype.bits // quantize_dtype.bits + if self.group_size % self.num_elem_per_storage != 0: + raise ValueError("Group size should be divisible by numbers of elements per storage") + self.num_storage_per_group = self.group_size // self.num_elem_per_storage + self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1 + + def quantize_model( + self, + model: nn.Module, + quant_map: QuantizeMapping, + name_prefix: str, + ) -> nn.Module: + """ + Quantize model with awq quantization. + + Parameters + ---------- + model : nn.Module + The non-quantized nn.Module. + + quant_map : QuantizeMapping + The quantize mapping with name mapping and func mapping. + + name_prefix : str + The name prefix for visited weight. + + Returns + ------- + ret : nn.Module + The quantized nn.Module. + """ + + class _Mutator(nn.Mutator): + def __init__(self, config: AWQQuantize, quant_map: QuantizeMapping) -> None: + super().__init__() + self.config = config + self.quant_map = quant_map + + def visit_module(self, name: str, node: nn.Module) -> Any: + """ + The visiting method for awq quantization of nn.Module nodes. + + Parameters + ---------- + name : str + The name of the current node + + node : nn.Module + The current node of nn.Module to mutate. + + Returns + ------- + ret_node : Any + The new node to replace current node. + """ + + if ( + isinstance(node, nn.Linear) + and not is_final_fc(name) + and not is_moe_gate(name, node) + ): + return AWQQuantizeLinear.from_linear(node, self.config) + return self.visit(name, node) + + model.to(dtype=self.model_dtype) + mutator = _Mutator(self, quant_map) + model = mutator.visit(name_prefix, model) + return model + + def _dequantize( + self, + weight: te.Tensor, + zeros: te.Tensor, + scale: te.Tensor, + out_shape: Optional[List[tirx.Expr]] = None, # noqa: UP006 + ): + float_weight = convert_uint_to_float( + weight, + DataType(self.quantize_dtype).bits, + self.num_elem_per_storage, + self.storage_dtype, + self.model_dtype, + out_shape=[weight.shape[0], weight.shape[1] * self.num_elem_per_storage], + ft_reorder=True, + ) + float_zeros = convert_uint_to_float( + zeros, + DataType(self.quantize_dtype).bits, + self.num_elem_per_storage, + self.storage_dtype, + self.model_dtype, + out_shape=[zeros.shape[0], zeros.shape[1] * self.num_elem_per_storage], + ft_reorder=True, + ) + float_weight = topi.transpose(float_weight) + float_zeros = topi.transpose(float_zeros) + scale = topi.transpose(scale) + return te.compute( + shape=( + [weight.shape[0], weight.shape[1] * self.num_elem_per_storage] + if out_shape is None + else out_shape + ), + fcompute=lambda i, j: tirx.Mul( + tirx.Sub(float_weight[i, j], float_zeros[i, j // self.group_size]), + scale[i, j // self.group_size], + ), + name="dequantize", + ) + + +class AWQQuantizeLinear(nn.Module): + """An nn.Linear module with AWQ quantization""" + + def __init__( + self, + in_features: int, + out_features: int, + config: AWQQuantize, + bias: bool = True, + out_dtype: Optional[str] = None, + ) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.out_dtype = out_dtype + self.config = config + self.qweight = nn.Parameter( + (in_features, out_features // config.num_elem_per_storage), + config.storage_dtype, + ) + self.qzeros = nn.Parameter( + ( + in_features // config.group_size, + out_features // config.num_elem_per_storage, + ), + config.storage_dtype, + ) + self.scales = nn.Parameter( + (in_features // config.group_size, out_features), config.model_dtype + ) + if bias: + self.bias = nn.Parameter( + (out_features,), config.model_dtype if out_dtype is None else out_dtype + ) + else: + self.bias = None + + @staticmethod + def from_linear(linear: nn.Linear, config: AWQQuantize) -> "AWQQuantizeLinear": + """ + Converts a non-quantized nn.Linear to a group quantized AWQQuantizeLinear + + Parameters + ---------- + linear : nn.Linear + The non-quantized nn.Linear. + + config : AWQQuantize + The awq quantization config. + + Returns + ------- + ret : GroupQuantizeLinear + The awq quantized AWQQuantizeLinear layer. + """ + return AWQQuantizeLinear( + in_features=linear.in_features, + out_features=linear.out_features, + config=config, + bias=getattr(linear, "bias", None) is not None, + out_dtype=linear.out_dtype, + ) + + def forward(self, x: nn.Tensor) -> nn.Tensor: + """ + Forward method for awq quantized linear layer + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the group quantized linear layer. + """ + w = nn.op.tensor_expr_op( + lambda weight, zeros, scale: self.config._dequantize( + weight, + zeros, + scale, + [ + tirx.IntImm("int64", self.out_features), + tirx.IntImm("int64", self.in_features), + ], + ), + name_hint="dequantize", + args=[self.qweight, self.qzeros, self.scales], + ) + w = nn.op.permute_dims(w) + x = nn.op.matmul(x, w, out_dtype=self.out_dtype) + if self.bias is not None: + x = x + self.bias + return x + + def to(self, dtype: Optional[str] = None) -> None: + """ + Override to() such that we do not convert bias if there is an out_dtype. + Otherwise, we might run into dtype mismatch when computing x + self.bias. + """ + self.qweight.to(dtype=dtype) + self.qzeros.to(dtype=dtype) + self.scales.to(dtype=dtype) + if self.bias is not None and self.out_dtype is None: + self.bias.to(dtype=dtype) + if dtype is not None and isinstance(getattr(self, "dtype", None), str): + self.dtype = dtype diff --git a/python/mlc_llm/quantization/block_scale_quantization.py b/python/mlc_llm/quantization/block_scale_quantization.py new file mode 100644 index 0000000..c81eb84 --- /dev/null +++ b/python/mlc_llm/quantization/block_scale_quantization.py @@ -0,0 +1,828 @@ +"""The block-scale quantization config""" + +from dataclasses import dataclass +from typing import Any, Literal, Optional, Tuple # noqa: UP035 + +import tvm +from tvm import DataType, DataTypeCode, te, tirx +from tvm.relax.frontend import nn +from tvm.script import tirx as T + +from mlc_llm.loader import QuantizeMapping +from mlc_llm.nn import MixtralExperts +from mlc_llm.op import cutlass, extern, moe_matmul, triton +from mlc_llm.support import logging +from mlc_llm.support import tensor_parallel as tp + +from .utils import apply_sharding, is_final_fc, is_moe_gate + +logger = logging.getLogger(__name__) + + +@dataclass +class BlockScaleQuantize: + """Configuration for block-scale quantization""" + + name: str + kind: str = "block-scale" + weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"] = "float8_e4m3fn" + model_dtype: Literal["float16", "bfloat16"] = "bfloat16" + quantize_linear: bool = True + weight_block_size: Optional[Tuple[int, int]] = None # noqa: UP006 + use_activation_scale: bool = False + + def __post_init__(self): + assert self.kind == "block-scale-quant" + weight_dtype = DataType(self.weight_dtype) + model_dtype = DataType(self.model_dtype) + assert weight_dtype.type_code in [ + DataTypeCode.Float8E4M3FN, + DataTypeCode.Float8E5M2, + ] + assert model_dtype.type_code in [ + DataTypeCode.FLOAT, + DataTypeCode.BFLOAT, + ] + + def quantize_model( + self, + model: nn.Module, + quant_map: QuantizeMapping, + name_prefix: str, + ) -> nn.Module: + """Quantize model with block-scale quantization + + Parameters + ---------- + model : nn.Module + The non-quantized nn.Module. + + quant_map : QuantizeMapping + The quantize mapping with name mapping and func mapping. + + name_prefix : str + The name prefix for visited weight. + + Returns + ------- + ret : nn.Module + The quantized nn.Module. + """ + + weight_block_size = model.weight_block_size + + class _Mutator(nn.Mutator): + def __init__(self, config: BlockScaleQuantize, quant_map: QuantizeMapping) -> None: + super().__init__() + self.config = config + self.quant_map = quant_map + + def visit_module(self, name: str, node: nn.Module) -> Any: + """The visiting method for block-scale quantization of nn.Module nodes. + + Parameters + ---------- + name : str + The name of the current node. + + node : nn.Module + The current node of nn.Module to mutate. + + Returns + ------ + ret : Any + """ + if getattr(node, "no_quantization", False): + return node + + if hasattr(node, "w_uk"): + assert hasattr(node, "w_uv") + assert node.block_size == weight_block_size + if ( + node.qk_nope_head_dim % node.block_size[0] != 0 + or node.v_head_dim % node.block_size[1] != 0 + ): + raise ValueError( + "Invalid DeepSeek model config: " + "qk_nope_head_dim must be multiple of weight_block_size[0], and " + "v_head_dim must be multiple of weight_block_size[1]. " + f"However, qk_nope_head_dim is {node.qk_nope_head_dim}, " + f"v_head_dim is {node.v_head_dim}, " + f"weight_block_size is {node.block_size}." + ) + w_uk_shard_strategy = node.w_uk.attrs.get("shard_strategy", None) + w_uv_shard_strategy = node.w_uv.attrs.get("shard_strategy", None) + node.w_uk = nn.Parameter( + (node.num_heads, node.kv_lora_rank, node.qk_nope_head_dim), + self.config.weight_dtype, + ) + node.w_uv = nn.Parameter( + (node.num_heads, node.v_head_dim, node.kv_lora_rank), + self.config.weight_dtype, + ) + node.w_uk_scale_inv = nn.Parameter( + ( + node.num_heads, + node.kv_lora_rank // node.block_size[1], + node.qk_nope_head_dim // node.block_size[0], + ), + "float32", + ) + node.w_uv_scale_inv = nn.Parameter( + ( + node.num_heads, + node.v_head_dim // node.block_size[0], + node.kv_lora_rank // node.block_size[1], + ), + "float32", + ) + if w_uk_shard_strategy is not None: + assert w_uk_shard_strategy.segs is None + apply_sharding(w_uk_shard_strategy, w_uk_shard_strategy.name, node.w_uk) + apply_sharding( + w_uk_shard_strategy, + f"{w_uk_shard_strategy.name}_scale_inv", + node.w_uk_scale_inv, + ) + if w_uv_shard_strategy is not None: + assert w_uv_shard_strategy.segs is None + apply_sharding(w_uv_shard_strategy, w_uv_shard_strategy.name, node.w_uv) + apply_sharding( + w_uv_shard_strategy, + f"{w_uv_shard_strategy.name}_scale_inv", + node.w_uv_scale_inv, + ) + + if ( + isinstance(node, nn.Linear) + and not is_final_fc(name) + and not is_moe_gate(name, node) + ): + if self.config.use_activation_scale: + return BlockScaleQuantizeLinearStaticActivation.from_linear( + node, self.config, weight_block_size + ) + return BlockScaleQuantizeLinear.from_linear( + node, self.config, weight_block_size + ) + if isinstance(node, MixtralExperts): + return BlockScaleQuantizeMixtralExperts.from_mixtral_experts( + node, self.config, weight_block_size + ) + return self.visit(name, node) + + model.to(dtype=self.model_dtype) + mutator = _Mutator(self, quant_map) + model = mutator.visit(name_prefix, model) + self.weight_block_size = weight_block_size + return model + + +class BlockScaleQuantizeLinear(nn.Module): + """Block-scale quantization for Linear""" + + def __init__( + self, + in_features: int, + out_features: int, + weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"], + block_size: Tuple[int, int], # noqa: UP006 + bias: bool = True, + dtype: Optional[str] = None, + out_dtype: Optional[str] = None, + ) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.out_dtype = out_dtype + self.weight = nn.Parameter((out_features, in_features), weight_dtype) + self.weight_scale_inv = nn.Parameter( + ( + (out_features + block_size[0] - 1) // block_size[0], + (in_features + block_size[1] - 1) // block_size[1], + ), + "float32", + ) + self.weight_dtype = weight_dtype + self.block_size = block_size + if bias: + self.bias = nn.Parameter((out_features,), dtype if out_dtype is None else out_dtype) + else: + self.bias = None + + @staticmethod + def from_linear( + src: nn.Linear, + config: BlockScaleQuantize, + weight_block_size: Optional[Tuple[int, int]], # noqa: UP006 + ) -> "BlockScaleQuantizeLinear": + """ + Converts a non-quantized nn.Linear to a block-scale quantized BlockScaleQuantizeLinear + + Parameters + ---------- + src : nn.Linear + The non-quantized nn.Linear. + + config : BlockScaleQuantize + The block-scale quantization config. + + weight_block_size : Optional[Tuple[int, int]] + The weight block size. + + Returns + ------- + ret : BlockScaleQuantizeLinear + The block-scale quantized BlockScaleQuantizeLinear. + """ + assert weight_block_size is not None + out_features, in_features = src.weight.shape + quantized_linear = BlockScaleQuantizeLinear( + in_features=in_features, + out_features=out_features, + weight_dtype=config.weight_dtype, + block_size=weight_block_size, + bias=getattr(src, "bias", None) is not None, + dtype=config.model_dtype, + out_dtype=src.out_dtype, + ) + if quantized_linear.bias is not None: + quantized_linear.bias.attrs = src.bias.attrs + if "shard_strategy" in src.weight.attrs: + shard = src.weight.attrs["shard_strategy"] + apply_sharding(shard, shard.name, quantized_linear.weight) + if isinstance(shard, tp.ShardSingleDim) and shard.segs is not None: + shard.segs = [x // weight_block_size[shard.dim] for x in shard.segs] + apply_sharding(shard, f"{shard.name}_scale_inv", quantized_linear.weight_scale_inv) + return quantized_linear + + def forward(self, x: nn.Tensor) -> nn.Tensor: + """Forward pass of the block-scale quantized linear layer. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor. + """ + m = 1 + for i in range(x.ndim - 1): + m *= x.shape[i] + if m == 1: + x_shape = x.shape + return dequantize_float8_groupwise_scaled_gemv( + x.reshape(1, x.shape[-1]), + self.weight, + self.weight_scale_inv, + self.block_size, + self.out_dtype if self.out_dtype is not None else x.dtype, + ).reshape(*x_shape[:-1], -1) + + shape_supported_by_cutlass = ( # noqa: F841 + self.weight.shape[0] % 128 == 0 and self.weight.shape[1] % 128 == 0 + ) + # Todo: check "shape supported by cutlass" for Hopper + if ( + extern.get_store().cutlass_gemm + and tvm.get_global_func( + "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn", allow_missing=True + ) + is not None + ): + x_fp8, x_scale = rowwise_group_quant_fp8( + x, self.block_size[1], self.weight_dtype, transpose_scale=True + ) + x = cutlass.fp8_groupwise_scaled_gemm( + x_fp8, + x_scale, + self.weight, + self.weight_scale_inv, + self.block_size, + self.out_dtype if self.out_dtype is not None else x.dtype, + ) + else: + x_fp8, x_scale = rowwise_group_quant_fp8( + x, self.block_size[1], self.weight_dtype, transpose_scale=False + ) + x = triton.fp8_groupwise_scaled_gemm( + x_fp8, + x_scale, + self.weight, + self.weight_scale_inv, + self.block_size, + self.out_dtype if self.out_dtype is not None else x.dtype, + ) + if self.bias is not None: + x = x + self.bias + return x + + def to(self, dtype: Optional[str] = None) -> None: + """ + Override to() such that we do not convert bias if there is an out_dtype. + Otherwise, we might run into dtype mismatch when computing x + self.bias. + """ + if self.bias is not None and self.out_dtype is None: + self.bias.to(dtype=dtype) + if dtype is not None and isinstance(getattr(self, "dtype", None), str): + self.dtype = dtype + + +class BlockScaleQuantizeLinearStaticActivation(BlockScaleQuantizeLinear): + """Block-scale quantization for static activation FP8.""" + + def __init__( + self, + in_features: int, + out_features: int, + weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"], + block_size: Tuple[int, int], # noqa: UP006 + bias: bool = True, + dtype: Optional[str] = None, + out_dtype: Optional[str] = None, + ) -> None: + super().__init__( + in_features=in_features, + out_features=out_features, + weight_dtype=weight_dtype, + block_size=block_size, + bias=bias, + dtype=dtype, + out_dtype=out_dtype, + ) + num_in_groups = (in_features + block_size[1] - 1) // block_size[1] + self.activation_scale = nn.Parameter((num_in_groups,), "float32") + + @staticmethod + def from_linear( + src: nn.Linear, + config: BlockScaleQuantize, + weight_block_size: Optional[Tuple[int, int]], # noqa: UP006 + ) -> "BlockScaleQuantizeLinearStaticActivation": + """ + Convert a non-quantized nn.Linear to a block-scale quantized BlockScaleQuantizeLinearStaticActivation. + + Parameters + ---------- + src : nn.Linear + The non-quantized nn.Linear. + + config : BlockScaleQuantize + The block-scale quantization config. + + weight_block_size : Optional[Tuple[int, int]] + The weight block size. + + Returns + ------- + ret : BlockScaleQuantizeLinearStaticActivation + The block-scale quantized BlockScaleQuantizeLinearStaticActivation + """ # noqa: E501 + assert weight_block_size is not None + out_features, in_features = src.weight.shape + quantized_linear = BlockScaleQuantizeLinearStaticActivation( + in_features=in_features, + out_features=out_features, + weight_dtype=config.weight_dtype, + block_size=weight_block_size, + bias=getattr(src, "bias", None) is not None, + dtype=config.model_dtype, + out_dtype=src.out_dtype, + ) + if quantized_linear.bias is not None: + quantized_linear.bias.attrs = src.bias.attrs + if "shard_strategy" in src.weight.attrs: + shard = src.weight.attrs["shard_strategy"] + apply_sharding(shard, shard.name, quantized_linear.weight) + if isinstance(shard, tp.ShardSingleDim) and shard.segs is not None: + shard.segs = [x // weight_block_size[shard.dim] for x in shard.segs] + apply_sharding(shard, f"{shard.name}_scale_inv", quantized_linear.weight_scale_inv) + apply_sharding( + shard, + f"{shard.name}_activation_scale", + quantized_linear.activation_scale, + ) + return quantized_linear + + def forward(self, x: nn.Tensor) -> nn.Tensor: + x_fp8 = static_activation_group_quant_fp8( + x, + self.activation_scale, + self.block_size[1], + self.weight_dtype, + ) + shape_supported_by_cutlass = ( + self.weight.shape[0] % 128 == 0 and self.weight.shape[1] % 128 == 0 + ) + if ( + extern.get_store().cutlass_gemm + and shape_supported_by_cutlass + and tvm.get_global_func( + "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn", allow_missing=True + ) + is not None + ): + x_scale = broadcast_activation_scale( + x, + self.activation_scale, + transpose=True, + ) + out = cutlass.fp8_groupwise_scaled_gemm( + x_fp8, + x_scale, + self.weight, + self.weight_scale_inv, + self.block_size, + self.out_dtype if self.out_dtype is not None else x.dtype, + ) + else: + x_scale_triton = broadcast_activation_scale( + x, + self.activation_scale, + transpose=False, + ) + out = triton.fp8_groupwise_scaled_gemm( + x_fp8, + x_scale_triton, + self.weight, + self.weight_scale_inv, + self.block_size, + self.out_dtype if self.out_dtype is not None else x.dtype, + ) + if self.bias is not None: + out = out + self.bias + return out + + +class BlockScaleQuantizeMixtralExperts(nn.Module): + """Block-scale quantization for MoE experts""" + + def __init__( + self, + num_local_experts: int, + in_features: int, + out_features: int, + weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"], + block_size: Tuple[int, int], # noqa: UP006 + ) -> None: + super().__init__() + self.num_local_experts = num_local_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter((num_local_experts, out_features, in_features), weight_dtype) + self.weight_scale_inv = nn.Parameter( + ( + num_local_experts, + (out_features + block_size[0] - 1) // block_size[0], + (in_features + block_size[1] - 1) // block_size[1], + ), + "float32", + ) + self.weight_dtype = weight_dtype + self.block_size = block_size + + @staticmethod + def from_mixtral_experts( + src: "MixtralExperts", + config: BlockScaleQuantize, + weight_block_size: Optional[Tuple[int, int]], # noqa: UP006 + ) -> "BlockScaleQuantizeMixtralExperts": + """ + Converts a non-quantized MixtralExperts to a block-scale + quantized BlockScaleQuantizeMixtralExperts + + Parameters + ---------- + src : MixtralExperts + The non-quantized MixtralExperts + + config : BlockScaleQuantize + The block-scale quantization config. + + weight_block_size : Optional[Tuple[int, int]] + The weight block size. + + Returns + ------- + ret : BlockScaleQuantizeMixtralExperts + The block-scale quantized BlockScaleQuantizeMixtralExperts layer. + """ + assert weight_block_size is not None + quantized_mistral_experts = BlockScaleQuantizeMixtralExperts( + num_local_experts=src.num_local_experts, + in_features=src.in_features, + out_features=src.out_features, + weight_dtype=config.weight_dtype, + block_size=weight_block_size, + ) + if "shard_strategy" in src.weight.attrs: + shard = src.weight.attrs["shard_strategy"] + apply_sharding(shard, shard.name, quantized_mistral_experts.weight) + if isinstance(shard, tp.ShardSingleDim) and shard.segs is not None: + shard.segs = [x // weight_block_size[shard.dim - 1] for x in shard.segs] + apply_sharding( + shard, + f"{shard.name}_scale_inv", + quantized_mistral_experts.weight_scale_inv, + ) + return quantized_mistral_experts + + def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor: + """Forward pass of the block-scale quantized MixtralExperts. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + indptr : nn.Tensor + The indptr tensor of group gemm, with shape of [num_experts + 1,]. + + Returns + ------- + ret : nn.Tensor + The output tensor. + """ + if indptr.ndim == 2: + # The input is for single token, which does not need group gemm + # and can be specialized. + expert_indices = indptr + assert expert_indices.shape[0] == 1 + return moe_matmul.dequantize_block_scale_float8_gemv( + x, + self.weight, + self.weight_scale_inv, + expert_indices, + self.block_size, + x.dtype, + ) + + x_fp8, x_scale = rowwise_group_quant_fp8( + x, self.block_size[1], self.weight_dtype, transpose_scale=False + ) + if ( + extern.get_store().cutlass_gemm + and tvm.get_global_func( + "cutlass.groupwise_scaled_group_gemm_e4m3fn_e4m3fn", allow_missing=True + ) + is not None + ): + x = cutlass.fp8_groupwise_scaled_group_gemm( + x_fp8, + x_scale, + self.weight, + self.weight_scale_inv, + indptr, + self.block_size, + x.dtype, + ) + else: + x = triton.fp8_groupwise_scaled_group_gemm( + x_fp8, + x_scale, + self.weight, + self.weight_scale_inv, + indptr, + self.block_size, + x.dtype, + ) + return x + + def to(self, dtype: Optional[str] = None) -> None: + """ + Override to() such that we do not convert bias if there is an out_dtype. + Otherwise, we might run into dtype mismatch when computing x + self.bias. + """ + if dtype is not None and isinstance(getattr(self, "dtype", None), str): + self.dtype = dtype + + +def rowwise_group_quant_fp8( + x: nn.Tensor, + group_size: int, + dtype: Literal["float8_e4m3fn", "float8_e5m2"], + transpose_scale: bool, + eps: float = 1e-10, + keep_first_batch_dim: bool = False, +) -> Tuple[nn.Tensor, nn.Tensor]: # noqa: UP006 + """Rowwise group quantization of fp8 tensor. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + group_size : int + The group size per row for quantization. + + transpose_scale : bool + Whether return the transposed scales or not. + + Returns + ------- + x_fp8 : nn.Tensor + The quantized tensor. + + x_scale : nn.Tensor + The scales of the quantized tensor. + If transpose_scale is True, the shape is + (*x.shape[:-2], ceildiv(x.shape[-1], group_size), x.shape[-2]). + Otherwise, the shape is (*x.shape[:-1], ceildiv(x.shape[-1], group_size)). + """ + assert x.ndim >= 2 + assert group_size > 0 + + def quantize(x: te.Tensor): + num_group = tirx.ceildiv(x.shape[-1], group_size) + max_abs_shape = (*x.shape[:-1], num_group) + max_abs_reduce_axis = te.reduce_axis((0, group_size), name="r") + scale_dtype = "float32" + max_abs = te.compute( + shape=max_abs_shape, + fcompute=lambda *idx: te.max( + tirx.if_then_else( + idx[-1] * group_size + max_abs_reduce_axis < x.shape[-1], + tirx.Max( + te.abs( + x(*idx[:-1], idx[-1] * group_size + max_abs_reduce_axis).astype( + scale_dtype + ) + ), + eps, + ), + tirx.min_value(scale_dtype), + ), + axis=max_abs_reduce_axis, + ), + name="max_abs", + ) + assert dtype in ["float8_e4m3fn", "float8_e5m2"] + fp8_max = 448.0 if dtype == "float8_e4m3fn" else 57344.0 + fp8_min = -fp8_max + scale = te.compute( + shape=max_abs_shape, + fcompute=lambda *idx: max_abs(*idx) / tirx.const(fp8_max, scale_dtype), + name="scale", + ) + x_quantized = te.compute( + shape=x.shape, + fcompute=lambda *idx: tirx.max( + tirx.min( + x(*idx).astype(scale_dtype) / scale(*idx[:-1], idx[-1] // group_size), + fp8_max, + ), + fp8_min, + ).astype(dtype), + name="x_quantized", + ) + if transpose_scale: + if not keep_first_batch_dim: + scale = te.compute( + shape=(num_group, *x.shape[:-1]), + fcompute=lambda *idx: scale(*idx[1:], idx[0]), + name="scale", + ) + else: + assert len(x.shape) > 2 + scale = te.compute( + shape=(x.shape[0], num_group, *x.shape[1:-1]), + fcompute=lambda *idx: scale(idx[0], *idx[2:], idx[1]), + name="scale", + ) + return x_quantized, scale + + x_quantized, scale = nn.tensor_expr_op(quantize, name_hint="rowwise_group_quant_fp8", args=[x]) + return x_quantized, scale + + +def static_activation_group_quant_fp8( + x: nn.Tensor, + activation_scale: nn.Tensor, + group_size: int, + dtype: Literal["float8_e4m3fn", "float8_e5m2"], +) -> nn.Tensor: + """Quantize activations with a pre-computed scale.""" + + assert activation_scale.ndim == 1 + + def quantize(x: te.Tensor, scale: te.Tensor): + fp8_max = 448.0 if dtype == "float8_e4m3fn" else 57344.0 + fp8_min = -fp8_max + + def fcompute(*idx): + group_idx = tirx.indexdiv(idx[-1], group_size) + return tirx.max( + tirx.min( + x(*idx).astype("float32") / scale(group_idx), + fp8_max, + ), + fp8_min, + ).astype(dtype) + + return te.compute(shape=x.shape, fcompute=fcompute, name="static_activation_group_fp8") + + quantized = nn.tensor_expr_op( + quantize, + name_hint="static_activation_group_fp8", + args=[x, activation_scale], + ) + return quantized + + +def broadcast_activation_scale( + x: nn.Tensor, + activation_scale: nn.Tensor, + transpose: bool, +) -> nn.Tensor: + """Broadcast stored activation scales.""" + + reshape_shape = (1,) * (x.ndim - 1) + (activation_scale.shape[0],) + scale = nn.op.reshape(activation_scale, reshape_shape) + scale = nn.op.broadcast_to(scale, (*x.shape[:-1], activation_scale.shape[0])) + if transpose: + axes = list(range(scale.ndim)) + axes[-1], axes[-2] = axes[-2], axes[-1] + scale = nn.op.permute_dims(scale, axes=axes) + return scale + + +def dequantize_float8_groupwise_scaled_gemv( + x: nn.Tensor, + w: nn.Tensor, + w_scale: nn.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + out_dtype: str, +) -> nn.Tensor: + """GEMV for FP8 groupwise scaled quantization. + + Parameters + ---------- + x : Tensor + The input tensor of shape (k,) + + w : Tensor + The quantized weight tensor of shape (n, k) + + w_scale : Tensor + The scale tensor of shape + (n // block_size[0], k // block_size[1]) + + block_size : Tuple[int, int] + The block size of the weight tensor. + + out_dtype : str + The output dtype of the GEMV computation. + """ + assert x.ndim == 2 + assert w.ndim == 2 + assert w_scale.ndim == 2 + assert x.shape[0] == 1 + assert x.shape[1] == w.shape[1] + _, k = x.shape + n, _ = w.shape + model_dtype = x.dtype + quantize_dtype = w.dtype + + assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[0] + assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[1] + + def _dequantize(w, s, i, j): + return w[i, j].astype(model_dtype) * s[i // block_size[0], j // block_size[1]].astype( + model_dtype + ) + + @T.prim_func(private=True, s_tir=True) + def _func( + x: T.Buffer((1, k), model_dtype), + w: T.Buffer((n, k), quantize_dtype), + w_scale: T.Buffer( + ( + (n + block_size[0] - 1) // block_size[0], + (k + block_size[1] - 1) // block_size[1], + ), + "float32", + ), + o: T.Buffer((n,), out_dtype), + ): + T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # kOutEWiseFusable + y = T.sblock_alloc_buffer((n, k), model_dtype) + for i1, i2 in T.grid(n, k): + with T.sblock("dequantize"): + i, j = T.axis.remap("SS", [i1, i2]) + y[i, j] = _dequantize(w, w_scale, i, j) + for i1, i2 in T.grid(n, k): + with T.sblock("gemv"): + i, j = T.axis.remap("SR", [i1, i2]) + with T.init(): + o[i] = T.cast(T.float16(0), out_dtype) + o[i] += (x[0, j] * y[i, j]).astype(out_dtype) + + return nn.op.tensor_ir_op( + _func, + "moe_dequantize_gemv", + args=[x, w, w_scale], + out=nn.Tensor.placeholder([n], out_dtype), + ) diff --git a/python/mlc_llm/quantization/fp8_quantization.py b/python/mlc_llm/quantization/fp8_quantization.py new file mode 100644 index 0000000..c6c0f03 --- /dev/null +++ b/python/mlc_llm/quantization/fp8_quantization.py @@ -0,0 +1,122 @@ +"""Quantization techniques for FP8""" + +import numpy as np +from tvm import relax, runtime +from tvm.relax.frontend import nn + +from mlc_llm.nn import MixtralExperts + +from ..op import cutlass, extern, moe_matmul +from . import per_tensor_quantization as ptq +from .utils import apply_sharding + + +class FP8PerTensorQuantizeMixtralExperts(ptq.PerTensorQuantizeMixtralExperts): + """MixtralExperts with per-tensor quantization in FP8.""" + + def __init__( + self, + num_local_experts, + in_features, + out_features, + config: ptq.PerTensorQuantize, + name: str, + tensor_parallel_shards=1, + ): + super().__init__(num_local_experts, in_features, out_features, config, name) + self.tensor_parallel_shards = tensor_parallel_shards + + @staticmethod + def from_mixtral_experts( + src: "MixtralExperts", + config: ptq.PerTensorQuantize, + name: str, + ) -> "FP8PerTensorQuantizeMixtralExperts": + """ + Converts a non-quantized MixtralExperts to a per-tensor quantized MixtralExperts. + + Parameters + ---------- + src : MixtralExperts + The non-quantized MixtralExperts + + config : PerTensorQuantize + The FP8 quantization weight_config. + + name : str + The name of the layer. + + Returns + ------- + ret : MixtralExpertsFP8 + The per-tensor quantized MixtralExperts. + """ + quantized_mistral_experts = FP8PerTensorQuantizeMixtralExperts( + num_local_experts=src.num_local_experts, + in_features=src.in_features, + out_features=src.out_features, + config=config, + name=name, + tensor_parallel_shards=src.tensor_parallel_shards, + ) + + if "shard_strategy" in src.weight.attrs: + shard = src.weight.attrs["shard_strategy"] + apply_sharding(shard, f"{shard.name}_q_weight", quantized_mistral_experts.q_weight) + # scale doesn't need to be sharded since it's the same for all shards + + return quantized_mistral_experts + + def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor: + w = self.q_weight + + if self.config.calibration_mode == "max": + _, x_scale = self.config.quantize_float8( + x, + quantize_dtype=self.config.activation_dtype, + storage_dtype=self.config.activation_dtype, + ) + if self.config.tensor_parallel_shards > 1: + x_scale = nn.ccl_allreduce(x_scale, "max") + x_scale = nn.extern( + "mlc_llm.calibration_observer", + [f"{self.name}.q_calibration_scale", "max", x_scale], + out=nn.Tensor.placeholder(x_scale.shape, x_scale.dtype), + ) + x_q = (x / x_scale.astype(x.dtype)).astype(self.config.activation_dtype) + x = x_q.astype(self.config.model_dtype) * x_scale.astype(self.config.model_dtype) + + if indptr.ndim == 2: + assert indptr.shape[0] == 1 + return moe_matmul.dequantize_float8_gemv( + x, w, self.q_scale, indptr, self.config.weight_dtype + ) + + if extern.get_store().cutlass_group_gemm: + if self.config.calibration_mode == "inference": + if self.q_calibration_scale is not None: + x /= self.q_calibration_scale.astype(x.dtype) + x_q = nn.op.astype(x, dtype=self.config.activation_dtype) + x_scale = self.q_calibration_scale + + scale = ( + x_scale * self.q_scale + if self.q_scale is not None + else nn.wrap_nested( + relax.Constant(runtime.tensor(np.array([1.0]).astype("float32"))), + "scale", + ) + ) + return cutlass.group_gemm( + x_q, w, indptr, scale, self.config.weight_dtype, self.config.model_dtype + ) + # Note: convert_weight is target agnostic, so a fallback must be provided + w = nn.tensor_expr_op( + self.config.dequantize_float8, + "dequantize", + args=[w, self.q_scale, self.config.weight_dtype], + ) + return moe_matmul.group_gemm(x, w, indptr) + + +ptq.PerTensorQuantizeMixtralExperts._IMPL["fp8"] = FP8PerTensorQuantizeMixtralExperts diff --git a/python/mlc_llm/quantization/ft_quantization.py b/python/mlc_llm/quantization/ft_quantization.py new file mode 100644 index 0000000..d8e919a --- /dev/null +++ b/python/mlc_llm/quantization/ft_quantization.py @@ -0,0 +1,404 @@ +"""The FasterTransformer quantization config""" + +from dataclasses import dataclass +from typing import Any, Callable, List, Literal, Optional, Tuple # noqa: UP035 + +import tvm +from tvm import DataType, DataTypeCode, IRModule, relax, te, tirx +from tvm.relax.frontend import nn +from tvm.runtime import Tensor +from tvm.s_tir import dlight as dl +from tvm.target import Target + +from ..loader import QuantizeMapping +from ..op import faster_transformer_dequantize_gemm +from ..support import logging +from ..support.auto_target import detect_cuda_arch_list +from ..support.style import bold +from .group_quantization import ( + GroupQuantize, + GroupQuantizeEmbedding, + GroupQuantizeLinear, +) +from .utils import is_final_fc, is_moe_gate + +logger = logging.getLogger(__name__) + + +@dataclass +class FTQuantize: + """Configuration for FasterTransformer quantization""" + + name: str + kind: str + quantize_dtype: Literal["int4", "int8"] + storage_dtype: Literal["int8"] + model_dtype: Literal["float16"] + group_size: Optional[int] = None + + num_elem_per_storage: int = 0 + max_int_value: int = 0 + + def fallback_group_quantize(self) -> GroupQuantize: + """ + The fallback group quantization config for other parameters. + + Returns + ------ + quantize: GroupQuantize + The group quantization config to fallback. + """ + return GroupQuantize( + name=self.name, + kind="group-quant", + group_size=32, # hardcoded to 32 as only supporting int4 quantization + quantize_dtype=self.quantize_dtype, + storage_dtype="uint32", + model_dtype=self.model_dtype, + linear_weight_layout="NK", + ) + + def __post_init__(self): + assert self.kind == "ft-quant" + quantize_dtype = DataType(self.quantize_dtype) + storage_dtype = DataType(self.storage_dtype) + assert self.quantize_dtype in ["int4", "int8"] + assert storage_dtype.type_code == DataTypeCode.INT + assert self.model_dtype == "float16" + assert self.group_size in [None, 64, 128] + if storage_dtype.bits < quantize_dtype.bits: + raise ValueError("Storage unit should be greater or equal to quantized element") + + self.num_elem_per_storage = storage_dtype.bits // quantize_dtype.bits + self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1 + self._quantize_func_cache = {} + + def quantize_model( + self, + model: nn.Module, + quant_map: QuantizeMapping, + name_prefix: str, + ) -> nn.Module: + """ + Quantize model with FasterTransformer quantization + + Parameters + ---------- + model : nn.Module + The non-quantized nn.Module. + + quant_map : QuantizeMapping + The quantize mapping with name mapping and func mapping. + + name_prefix : str + The name prefix for visited weight. + + Returns + ------- + ret : nn.Module + The quantized nn.Module. + """ + + class _Mutator(nn.Mutator): + def __init__(self, config: FTQuantize, quant_map: QuantizeMapping) -> None: + super().__init__() + self.config = config + self.quant_map = quant_map + + def visit_module(self, name: str, node: nn.Module) -> Any: + """ + The visiting method for FasterTransformer quantization of nn.Module nodes. + + Parameters + ---------- + name : str + The name of the current node. + + node : nn.Module + The current node of nn.Module to mutate. + + Returns + ------ + ret_node: Any + The new node to replace current node. + """ + if isinstance(node, nn.Linear): + weight_name = f"{name}.weight" + self.quant_map.param_map[weight_name] = [ + f"{name}.q_weight", + f"{name}.q_scale", + ] + if ( + is_final_fc(name) + or node.out_dtype == "float32" + or (self.config.quantize_dtype == "int4" and node.out_features % 8 != 0) + or (self.config.quantize_dtype == "int8" and node.out_features % 4 != 0) + ): + # Under any of the conditions we fall back to GroupQuantize + # For `is_final_fc()` see https://github.com/mlc-ai/mlc-llm/issues/1723 + # If simply skipping lm_head quantization degrades performance + # Other requirements are from CUTLASS + logger.info( + 'Fallback to GroupQuantize for nn.Linear: "%s", ' + + "weight.shape: %s, out_dtype: %s", + bold(name), + node.weight.shape, + node.out_dtype, + ) + group_quantize = self.config.fallback_group_quantize() + self.quant_map.map_func[weight_name] = group_quantize.quantize_weight + return GroupQuantizeLinear.from_linear(node, group_quantize) + if not is_moe_gate(name, node): + self.quant_map.map_func[weight_name] = self.config.quantize_weight + return FTQuantizeLinear.from_linear(node, self.config) + if isinstance(node, nn.Embedding): + weight_name = f"{name}.weight" + self.quant_map.param_map[weight_name] = [ + f"{name}.q_weight", + f"{name}.q_scale", + ] + group_quantize = self.config.fallback_group_quantize() + self.quant_map.map_func[weight_name] = group_quantize.quantize_weight + return GroupQuantizeEmbedding.from_embedding(node, group_quantize) + return self.visit(name, node) + + model.to(dtype=self.model_dtype) + mutator = _Mutator(self, quant_map) + model = mutator.visit(name_prefix, model) + return model + + def quantize_weight(self, weight: Tensor) -> List[Tensor]: # noqa: UP006 + """ + Quantize weight with FasterTransformer quantization + + Parameters + ---------- + weight : Tensor + The original weight. + + Returns + ------ + ret: List[Tensor] + The list of FasterTransformer quantized weights. + """ + assert tvm.get_global_func("relax.ext.cutlass", True), ( + "Cutlass should be enabled in TVM runtime to quantize weight, " + "but not enabled in current TVM runtime environment. " + "To enable Cutlass in TVM runtime, please `set(USE_CUTLASS ON)` " + "in config.cmake when compiling TVM from source" + ) + assert len(weight.shape) == 2 + device = weight.device + device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()] + if device_type == "cuda": + target = Target.current() + if target is None: + target = Target.from_device(device) + with target: + + def _create_quantize_func() -> IRModule: + bb = relax.BlockBuilder() + weight_var = relax.Var("weight", relax.TensorType(weight.shape, weight.dtype)) + with bb.function(name="main", params=[weight_var]): + with bb.dataflow(): + lv0 = bb.emit_te(self._quantize, weight_var) + lv1 = bb.normalize(lv0[0]) + lv2 = bb.emit( + relax.call_pure_packed( + "cutlass.ft_preprocess_weight", + lv1, + detect_cuda_arch_list(target=target)[0], + DataType(self.quantize_dtype).bits == 4, + ty_args=lv1.ty, + ) + ) + gv = bb.emit_output(relax.Tuple([lv2, lv0[1]])) + bb.emit_func_output(gv) + return bb.finalize() + + def _compile_quantize_func(mod: IRModule) -> Callable: + mod = dl.ApplyDefaultSchedule( + dl.gpu.Reduction(), + dl.gpu.GeneralReduction(), + dl.gpu.Fallback(), + )(mod) + ex = relax.build(mod, target=target) + vm = relax.VirtualMachine(ex, device) + return vm["main"] + + key = str( + ( + int(weight.shape[0]), + int(weight.shape[1]), + weight.dtype, + device_type, + ) + ) + quantize_func = self._quantize_func_cache.get(key, None) + if quantize_func is None: + logger.info("Compiling quantize function for key: %s", key) + quantize_func = _compile_quantize_func(_create_quantize_func()) + self._quantize_func_cache[key] = quantize_func + data = quantize_func(weight) + return data + else: + raise NotImplementedError(f"Device type {device_type} is not supported") + + def _quantize( + self, + weight: te.Tensor, + ) -> Tuple[te.Tensor, te.Tensor]: # noqa: UP006 + """FasterTransformer quantization for weight tensor, defined in tensor expression.""" + assert len(weight.shape) == 2 + n, k = weight.shape + + cur_group_size = k if not self.group_size else self.group_size + scale_shape = (tirx.ceildiv(k, cur_group_size), n) + r = te.reduce_axis((0, cur_group_size), name="r") + + max_abs = te.compute( + shape=scale_shape, + fcompute=lambda j, i: te.max( + tirx.if_then_else( + j * cur_group_size + r < k, + te.abs(weight[i, j * cur_group_size + r]), + te.min_value(self.model_dtype), + ), + axis=r, + ), + name="max_abs_value", + ) + max_int = tirx.const(self.max_int_value, self.model_dtype) + scale = te.compute( + scale_shape, + lambda i, j: max_abs[i, j].astype(self.model_dtype) / max_int, + name="scale", + ) + # compute scaled weight + quantize_dtype = DataType(self.quantize_dtype) + bin_mask = tirx.const((1 << quantize_dtype.bits) - 1, self.storage_dtype) + scaled_weight = te.compute( + shape=weight.shape, + fcompute=lambda i, j: ( + tirx.min( + tirx.max( + tirx.round(weight[i, j] / scale[j // cur_group_size, i]), + -max_int - 1, + ), + max_int, + ).astype(self.storage_dtype) + & bin_mask + ), + ) + + quantized_weight_shape = (k, tirx.ceildiv(n, self.num_elem_per_storage)) + r = te.reduce_axis((0, self.num_elem_per_storage), name="r") + quantized_weight = te.compute( + shape=quantized_weight_shape, + fcompute=lambda j, i: tirx.sum( + tirx.if_then_else( + i * self.num_elem_per_storage + r < n, + scaled_weight[i * self.num_elem_per_storage + r, j] + << ( + r.astype(self.storage_dtype) + * tirx.const(quantize_dtype.bits, self.storage_dtype) + ), + tirx.const(0, self.storage_dtype), + ), + axis=r, + ), + name="weight", + ) + + return quantized_weight, scale + + +class FTQuantizeLinear(nn.Module): + """An nn.Linear module with FasterTransformer quantization""" + + def __init__( + self, + in_features: int, + out_features: int, + config: FTQuantize, + bias: bool = True, + out_dtype: Optional[str] = None, + ) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.out_dtype = out_dtype + self.config = config + cur_group_size = in_features if not config.group_size else config.group_size + self.q_weight = nn.Parameter( + (in_features, tirx.ceildiv(out_features, config.num_elem_per_storage)), + config.storage_dtype, + ) + self.q_scale = nn.Parameter( + (tirx.ceildiv(in_features, cur_group_size), out_features), config.model_dtype + ) + if bias: + self.bias = nn.Parameter( + (out_features,), config.model_dtype if out_dtype is None else out_dtype + ) + else: + self.bias = None + + @staticmethod + def from_linear(src: nn.Linear, config: FTQuantize) -> "FTQuantizeLinear": + """ + Converts a non-quantized nn.Linear to a FasterTransformer quantized FTQuantizeLinear + + Parameters + ---------- + src : nn.Linear + The non-quantized nn.Linear. + + config : FTQuantize + The FasterTransformer quantization config. + + Returns + ------- + ret : FTQuantizeLinear + The FasterTransformer quantized FTQuantizeLinear layer. + """ + quantized_linear = FTQuantizeLinear( + in_features=src.in_features, + out_features=src.out_features, + config=config, + bias=getattr(src, "bias", None) is not None, + out_dtype=src.out_dtype, + ) + if quantized_linear.bias is not None: + quantized_linear.bias.attrs = src.bias.attrs + return quantized_linear + + def forward(self, x: nn.Tensor) -> nn.Tensor: + """ + Forward method for FasterTransformer quantized linear layer. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the FasterTransformer quantized linear layer. + """ + return faster_transformer_dequantize_gemm( + x, self.q_weight, self.q_scale, self.bias, group_size=self.config.group_size + ) + + def to(self, dtype: Optional[str] = None) -> None: + """ + Override to() such that we do not convert bias if there is an out_dtype. + Otherwise, we might run into dtype mismatch when computing x + self.bias. + """ + self.q_weight.to(dtype=dtype) + self.q_scale.to(dtype=dtype) + if self.bias is not None and self.out_dtype is None: + self.bias.to(dtype=dtype) + if dtype is not None and isinstance(getattr(self, "dtype", None), str): + self.dtype = dtype diff --git a/python/mlc_llm/quantization/group_quantization.py b/python/mlc_llm/quantization/group_quantization.py new file mode 100644 index 0000000..bf7e63a --- /dev/null +++ b/python/mlc_llm/quantization/group_quantization.py @@ -0,0 +1,660 @@ +"""The group quantization config""" + +from dataclasses import dataclass +from functools import partial +from typing import Any, List, Literal, Optional, Tuple, Union # noqa: UP035 + +from tvm import DataType, DataTypeCode, IRModule, relax, te, tirx, topi +from tvm.relax.frontend import nn +from tvm.runtime import Tensor + +from mlc_llm.loader import QuantizeMapping +from mlc_llm.nn import MixtralExperts +from mlc_llm.support import logging + +from .utils import ( + apply_sharding, + compile_quantize_func, + convert_uint_to_float, + is_final_fc, + is_moe_gate, + pack_weight, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class GroupQuantize: + """Configuration for group quantization""" + + name: str + kind: str + group_size: int + quantize_dtype: Literal["int3", "int4", "int8"] + storage_dtype: Literal["uint32"] + model_dtype: Literal["float16", "float32", "bfloat16"] + linear_weight_layout: Literal["KN", "NK"] + quantize_embedding: bool = True + quantize_final_fc: bool = True + + num_elem_per_storage: int = 0 + num_storage_per_group: int = 0 + max_int_value: int = 0 + tensor_parallel_shards: int = 0 + + def __post_init__(self): + assert self.kind == "group-quant" + quantize_dtype = DataType(self.quantize_dtype) + storage_dtype = DataType(self.storage_dtype) + model_dtype = DataType(self.model_dtype) + assert quantize_dtype.type_code == DataTypeCode.INT + assert storage_dtype.type_code == DataTypeCode.UINT + assert model_dtype.type_code in (DataTypeCode.FLOAT, DataTypeCode.BFLOAT) + if storage_dtype.bits < quantize_dtype.bits: + raise ValueError("Storage unit should be greater or equal to quantized element") + + self.num_elem_per_storage = storage_dtype.bits // quantize_dtype.bits + if self.group_size % self.num_elem_per_storage != 0: + raise ValueError("Group size should be divisible by numbers of elements per storage") + self.num_storage_per_group = self.group_size // self.num_elem_per_storage + self.max_int_value = (2 ** (quantize_dtype.bits - 1)) - 1 + self.linear_quant_axis = 0 if self.linear_weight_layout == "KN" else 1 + self._quantize_func_cache = {} + + def quantize_model( + self, + model: nn.Module, + quant_map: QuantizeMapping, + name_prefix: str, + ) -> nn.Module: + """ + Quantize model with group quantization + + Parameters + ---------- + model : nn.Module + The non-quantized nn.Module. + + quant_map : QuantizeMapping + The quantize mapping with name mapping and func mapping. + + name_prefix : str + The name prefix for visited weight. + + Returns + ------- + ret : nn.Module + The quantized nn.Module. + """ + + class _Mutator(nn.Mutator): + def __init__(self, config: GroupQuantize, quant_map: QuantizeMapping) -> None: + super().__init__() + self.config = config + self.quant_map = quant_map + + def visit_module(self, name: str, node: nn.Module) -> Any: + """ + The visiting method for group quantization of nn.Module nodes. + + Parameters + ---------- + name : str + The name of the current node. + + node : nn.Module + The current node of nn.Module to mutate. + + Returns + ------ + ret_node: Any + The new node to replace current node. + """ + if getattr(node, "no_quantization", False): + return node + + if ( + isinstance(node, nn.Linear) + and (not is_final_fc(name) or self.config.quantize_final_fc) + and not is_moe_gate(name, node) + ): + weight_name = f"{name}.weight" + self.quant_map.param_map[weight_name] = [ + f"{name}.q_weight", + f"{name}.q_scale", + ] + self.quant_map.map_func[weight_name] = partial( + self.config.quantize_weight, + output_transpose=self.config.linear_weight_layout == "KN", + ) + return GroupQuantizeLinear.from_linear(node, self.config) + if isinstance(node, nn.Embedding) and self.config.quantize_embedding: + weight_name = f"{name}.weight" + self.quant_map.param_map[weight_name] = [ + f"{name}.q_weight", + f"{name}.q_scale", + ] + self.quant_map.map_func[weight_name] = self.config.quantize_weight + return GroupQuantizeEmbedding.from_embedding(node, self.config) + if isinstance(node, MixtralExperts): + weight_name = f"{name}.weight" + self.quant_map.param_map[weight_name] = [ + f"{name}.q_weight", + f"{name}.q_scale", + ] + self.quant_map.map_func[weight_name] = self.config.quantize_weight + return GroupQuantizeMixtralExperts.from_mixtral_experts(node, self.config) + return self.visit(name, node) + + model.to(dtype=self.model_dtype) + mutator = _Mutator(self, quant_map) + model = mutator.visit(name_prefix, model) + return model + + def _dequantize( + self, + weight: te.Tensor, + scale: te.Tensor, + axis: int, + out_shape: Optional[List[tirx.Expr]] = None, # noqa: UP006 + ): + tir_max_int = tirx.const(self.max_int_value, self.model_dtype) + float_weight = convert_uint_to_float( + weight, + DataType(self.quantize_dtype).bits, + self.num_elem_per_storage, + self.storage_dtype, + self.model_dtype, + axis=axis, + out_shape=out_shape, + ) + if out_shape is None: + out_shape = weight.shape + out_shape[axis] *= self.num_elem_per_storage + axis = axis if axis >= 0 else len(out_shape) + axis + return te.compute( + shape=out_shape, + fcompute=lambda *idx: tirx.Mul( + tirx.Sub( + float_weight(*idx), + tir_max_int, + ), + scale(*idx[:axis], idx[axis] // self.group_size, *idx[axis + 1 :]), + ), + name="dequantize", + ) + + def quantize_weight( + self, weight: Tensor, axis: int = -1, output_transpose: bool = False + ) -> List[Tensor]: # noqa: UP006 + """ + Quantize weight with group quantization + + Parameters + ---------- + weight : Tensor + The original weight. + + axis : int + The group axis. + + output_transpose : bool + Whether to transpose the output quantized weight. Only 2D weight is supported. + + Returns + ------ + ret: List[Tensor] + The list of group quantized weights. + """ + device = weight.device + device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()] + axis = axis if axis >= 0 else len(weight.shape) + axis + + def _create_quantize_func() -> IRModule: + bb = relax.BlockBuilder() + weight_var = relax.Var("weight", relax.TensorType(weight.shape, weight.dtype)) + with bb.function(name="main", params=[weight_var]): + with bb.dataflow(): + lv = bb.emit_te(self._quantize, weight_var, axis, output_transpose) + gv = bb.emit_output(lv) + bb.emit_func_output(gv) + return bb.finalize() + + key = ( + f"({weight.shape}, {weight.dtype}, {device_type}, " + f"axis={axis}, output_transpose={output_transpose})" + ) + quantize_func = self._quantize_func_cache.get(key, None) + if quantize_func is None: + logger.info("Compiling quantize function for key: %s", key) + quantize_func = compile_quantize_func(_create_quantize_func(), device=device) + self._quantize_func_cache[key] = quantize_func + return quantize_func(weight) + + def _quantize( + self, + weight: te.Tensor, + axis: int = -1, + output_transpose: bool = False, + ) -> Tuple[te.Tensor, te.Tensor]: # noqa: UP006 + """Group quantization for weight tensor, defined in tensor expression.""" + max_int = tirx.const(self.max_int_value, self.model_dtype) + shape = weight.shape + axis = axis if axis >= 0 else len(shape) + axis + k = shape[axis] + # compute scale per group + r = te.reduce_axis((0, self.group_size), name="r") + num_group = tirx.ceildiv(k, self.group_size) + scale_shape = (*shape[:axis], num_group, *shape[axis + 1 :]) + max_abs = te.compute( + shape=scale_shape, + fcompute=lambda *idx: te.max( + tirx.if_then_else( + idx[axis] * self.group_size + r < k, + te.abs( + weight( + *idx[:axis], + idx[axis] * self.group_size + r, + *idx[axis + 1 :], + ) + ), + te.min_value(self.model_dtype), + ), + axis=r, + ), + name="max_abs_value", + ) + scale = te.compute( + scale_shape, + lambda *idx: max_abs(*idx).astype(self.model_dtype) / max_int, + name="scale", + ) + # compute scaled weight + scaled_weight = te.compute( + shape=weight.shape, + fcompute=lambda *idx: tirx.min( + tirx.max( + tirx.round( + weight(*idx) + / scale(*idx[:axis], idx[axis] // self.group_size, *idx[axis + 1 :]) + + max_int + ), + tirx.const(0, self.model_dtype), + ), + max_int * 2, + ).astype(self.storage_dtype), + ) + # compute quantized weight per storage + num_storage = self.num_storage_per_group * num_group + quantized_weight_shape = (*shape[:axis], num_storage, *shape[axis + 1 :]) + quantized_weight = pack_weight( + scaled_weight, + axis=axis, + num_elem_per_storage=self.num_elem_per_storage, + weight_dtype=self.quantize_dtype, + storage_dtype=self.storage_dtype, + out_shape=quantized_weight_shape, + ) + if output_transpose: + if len(quantized_weight.shape) != 2 or len(scale.shape) != 2: + raise ValueError( + "Does not support transpose output quantized weight with ndim != 2" + ) + quantized_weight = topi.transpose(quantized_weight) + scale = topi.transpose(scale) + return quantized_weight, scale + + +class GroupQuantizeLinear(nn.Module): + """An nn.Linear module with group quantization""" + + def __init__( + self, + in_features: int, + out_features: Union[int, tirx.Var], + config: GroupQuantize, + bias: bool = True, + out_dtype: Optional[str] = None, + ) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.out_dtype = out_dtype + self.config = config + num_group = tirx.ceildiv(in_features, config.group_size) + num_shards = config.tensor_parallel_shards + if num_shards > 1 and (in_features * num_shards // config.group_size) % num_shards != 0: + raise ValueError( + f"The linear dimension {in_features * num_shards} has " + f"{in_features * num_shards // config.group_size} groups under group size " + f"{config.group_size}. The groups cannot be evenly distributed on " + f"{num_shards} GPUs.\n" + "Possible solutions: reduce number of GPUs, or use quantization with smaller " + "group size." + ) + if config.linear_weight_layout == "KN": + self.q_weight = nn.Parameter( + (config.num_storage_per_group * num_group, out_features), + config.storage_dtype, + ) + self.q_scale = nn.Parameter((num_group, out_features), config.model_dtype) + else: + self.q_weight = nn.Parameter( + (out_features, config.num_storage_per_group * num_group), + config.storage_dtype, + ) + self.q_scale = nn.Parameter((out_features, num_group), config.model_dtype) + if bias: + self.bias = nn.Parameter( + (out_features,), config.model_dtype if out_dtype is None else out_dtype + ) + else: + self.bias = None + + @staticmethod + def from_linear(src: nn.Linear, config: GroupQuantize) -> "GroupQuantizeLinear": + """ + Converts a non-quantized nn.Linear to a group quantized GroupQuantizeLinear + + Parameters + ---------- + src : nn.Linear + The non-quantized nn.Linear. + + config : GroupQuantize + The group quantization config. + + Returns + ------- + ret : GroupQuantizeLinear + The group quantized GroupQuantizeLinear layer. + """ + # For dynamic shape, src.out_features is `"name"`; src.weight.shape[0] is `tirx.Var("name")` + out_features, in_features = src.weight.shape + quantized_linear = GroupQuantizeLinear( + in_features=in_features, + out_features=out_features, + config=config, + bias=getattr(src, "bias", None) is not None, + out_dtype=src.out_dtype, + ) + if quantized_linear.bias is not None: + quantized_linear.bias.attrs = src.bias.attrs + if "shard_strategy" in src.weight.attrs: + shard = src.weight.attrs["shard_strategy"] + apply_sharding(shard, f"{shard.name}_q_weight", quantized_linear.q_weight) + apply_sharding(shard, f"{shard.name}_q_scale", quantized_linear.q_scale) + return quantized_linear + + def forward(self, x: nn.Tensor) -> nn.Tensor: + """ + Forward method for group quantized linear layer. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the group quantized linear layer. + """ + w = nn.op.tensor_expr_op( + lambda weight, scale: self.config._dequantize( + weight, + scale, + axis=self.config.linear_quant_axis, + out_shape=( + [ + ( + tirx.IntImm("int64", self.out_features) + if isinstance(self.out_features, int) + else weight.shape[0] + ), # Reuse same tirx.Var for symbolic shape (after Exporter) + tirx.IntImm("int64", self.in_features), + ] + if self.config.linear_weight_layout == "NK" + else [ + tirx.IntImm("int64", self.in_features), + ( + tirx.IntImm("int64", self.out_features) + if isinstance(self.out_features, int) + else weight.shape[1] + ), # Reuse same tirx.Var for symbolic shape (after Exporter) + ] + ), + ), + name_hint="dequantize", + args=[self.q_weight, self.q_scale], + ) + if self.config.linear_weight_layout == "NK": + w = nn.op.permute_dims(w) + x = nn.op.matmul(x, w, out_dtype=self.out_dtype) + if self.bias is not None: + x = x + self.bias + return x + + def to(self, dtype: Optional[str] = None) -> None: + """ + Override to() such that we do not convert bias if there is an out_dtype. + Otherwise, we might run into dtype mismatch when computing x + self.bias. + """ + self.q_weight.to(dtype=dtype) + self.q_scale.to(dtype=dtype) + if self.bias is not None and self.out_dtype is None: + self.bias.to(dtype=dtype) + if dtype is not None and isinstance(getattr(self, "dtype", None), str): + self.dtype = dtype + + +class GroupQuantizeEmbedding(nn.Module): + """An nn.Embedding module with group quantization""" + + def __init__(self, num: Union[int, tirx.Var], dim: int, config: GroupQuantize): + self.num = num + self.dim = dim + self.config = config + num_group = tirx.ceildiv(dim, config.group_size) + self.q_weight = nn.Parameter( + (num, config.num_storage_per_group * num_group), config.storage_dtype + ) + self.q_scale = nn.Parameter((num, num_group), config.model_dtype) + + @staticmethod + def from_embedding(embedding: nn.Embedding, config: GroupQuantize) -> "GroupQuantizeEmbedding": + """ + Converts a non-quantized nn.Embedding to a group quantized GroupQuantizeEmbedding + + Parameters + ---------- + linear : nn.Embedding + The non-quantized nn.Embedding. + + config : GroupQuantize + The group quantization config. + + Returns + ------- + ret : GroupQuantizeEmbedding + The group quantized GroupQuantizeEmbedding layer. + """ + num, dim = embedding.weight.shape + return GroupQuantizeEmbedding(num, dim, config) + + def forward(self, x: nn.Tensor): + """ + Forward method for group quantized embedding layer. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the embedding layer. + """ + w = nn.op.tensor_expr_op( + lambda weight, scale: self.config._dequantize( + weight, + scale, + axis=-1, + out_shape=[ + ( + tirx.IntImm("int64", self.num) + if isinstance(self.num, int) + else weight.shape[0] + ), # Reuse same tirx.Var for symbolic shape (after Exporter) + tirx.IntImm("int64", self.dim), + ], + ), + name_hint="dequantize", + args=[self.q_weight, self.q_scale], + ) + if x.ndim == 1: + return nn.op.take(w, x, axis=0) + return nn.op.reshape( + nn.op.take(w, nn.op.reshape(x, shape=[-1]), axis=0), + shape=[*x.shape, self.dim], + ) + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which dequantizes the weight + and multiplies it with the input tensor. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the lm_head layer. + """ + w = nn.op.tensor_expr_op( + lambda weight, scale: self.config._dequantize( + weight, + scale, + axis=-1, + out_shape=[ + ( + tirx.IntImm("int64", self.num) + if isinstance(self.num, int) + else weight.shape[0] + ), + tirx.IntImm("int64", self.dim), + ], + ), + name_hint="dequantize", + args=[self.q_weight, self.q_scale], + ) + w = nn.op.permute_dims(w) + return nn.op.matmul(x, w, out_dtype="float32") + + +class GroupQuantizeMixtralExperts(nn.Module): + """An MixtralExperts module with group quantization""" + + def __init__( + self, + num_local_experts, + in_features, + out_features, + config: GroupQuantize, + ): + self.num_local_experts = num_local_experts + self.in_features = in_features + self.out_features = out_features + self.config = config + num_group = tirx.ceildiv(in_features, config.group_size) + self.q_weight = nn.Parameter( + (num_local_experts, out_features, config.num_storage_per_group * num_group), + config.storage_dtype, + ) + self.q_scale = nn.Parameter( + (num_local_experts, out_features, num_group), config.model_dtype + ) + self.quantize_dtype = config.quantize_dtype + self.group_size = config.group_size + self.dtype = config.model_dtype + if config.linear_weight_layout == "KN": + raise NotImplementedError("GroupQuantizeMixtralExperts does not support KN layout now.") + + @staticmethod + def from_mixtral_experts( + src: "MixtralExperts", config: GroupQuantize + ) -> "GroupQuantizeMixtralExperts": + """ + Converts a non-quantized MixtralExperts to a group quantized GroupQuantizeMixtralExperts + + Parameters + ---------- + src : MixtralExperts + The non-quantized MixtralExperts + + config : GroupQuantize + The group quantization config. + + Returns + ------- + ret : GroupQuantizeMixtralExperts + The group quantized GroupQuantizeMixtralExperts layer. + """ + quantized_mistral_experts = GroupQuantizeMixtralExperts( + num_local_experts=src.num_local_experts, + in_features=src.in_features, + out_features=src.out_features, + config=config, + ) + if "shard_strategy" in src.weight.attrs: + shard = src.weight.attrs["shard_strategy"] + apply_sharding(shard, f"{shard.name}_q_weight", quantized_mistral_experts.q_weight) + apply_sharding(shard, f"{shard.name}_q_scale", quantized_mistral_experts.q_scale) + return quantized_mistral_experts + + def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor: + """Forward method for group quantized mistral experts. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + indptr: nn.Tensor + The indptr tensor + + single_batch_decode: bool + Whether to use single-batch decode + + Returns + ------- + ret : nn.Tensor + The output tensor for the group quantized mistral experts layer. + """ + from mlc_llm.op import moe_matmul + + assert x.ndim == 2 + if indptr.ndim == 2: # single-batch + assert indptr.shape[0] == 1 + return moe_matmul.dequantize_gemv( + x, + self.q_weight, + self.q_scale, + indptr, + quantize_dtype=self.quantize_dtype, + group_size=self.group_size, + ) + assert indptr.ndim == 1 + return moe_matmul.dequantize_group_gemm( + x, + self.q_weight, + self.q_scale, + indptr, + quantize_dtype=self.quantize_dtype, + indptr_dtype=indptr.dtype, + group_size=self.group_size, + ) diff --git a/python/mlc_llm/quantization/model_quantization.py b/python/mlc_llm/quantization/model_quantization.py new file mode 100644 index 0000000..d314aa8 --- /dev/null +++ b/python/mlc_llm/quantization/model_quantization.py @@ -0,0 +1,157 @@ +"""Quantization factory utilities for model quantization.""" + +from typing import Any, Callable, Dict, Optional, Tuple, Type # noqa: UP035 + +from tvm.relax.frontend import nn + +from mlc_llm.loader import QuantizeMapping + +from .awq_quantization import AWQQuantize +from .block_scale_quantization import BlockScaleQuantize +from .ft_quantization import FTQuantize +from .group_quantization import GroupQuantize +from .no_quantization import NoQuantize +from .per_tensor_quantization import PerTensorQuantize +from .quantization import Quantization + +FuncQuantization = Callable[[Any, Quantization], Tuple[nn.Module, QuantizeMapping]] # noqa: UP006 + + +def make_quantization_functions( + model_cls: Type[nn.Module], # noqa: UP006 + *, + model_ctor: Optional[Callable[[Any], nn.Module]] = None, + supports_group_quant: bool = True, + supports_ft_quant: bool = True, + supports_awq: bool = False, + awq_unsupported_message: Optional[str] = None, + supports_per_tensor: bool = False, + supports_block_scale: bool = False, + set_tensor_parallel_shards: bool = True, + per_tensor_use_shards: bool = True, +) -> Dict[str, FuncQuantization]: # noqa: UP006 + """Create standard quantization function implementations for a model class.""" + + def _create_model(model_config: Any) -> nn.Module: + if model_ctor is not None: + return model_ctor(model_config) + return model_cls(model_config) + + def _no_quant(model_config: Any, quantization: NoQuantize) -> Tuple[nn.Module, QuantizeMapping]: # noqa: UP006 + model = _create_model(model_config) + model.to(quantization.model_dtype) + return model, QuantizeMapping({}, {}) + + def _group_quant( + model_config: Any, + quantization: GroupQuantize, + ) -> Tuple[nn.Module, QuantizeMapping]: # noqa: UP006 + model = _create_model(model_config) + model.to(quantization.model_dtype) + quant_map = QuantizeMapping({}, {}) + if set_tensor_parallel_shards: + if not hasattr(model_config, "tensor_parallel_shards"): + raise AttributeError( + "model_config is missing required " + "attribute 'tensor_parallel_shards' for group quantization" + ) + quantization.tensor_parallel_shards = getattr(model_config, "tensor_parallel_shards") + model = quantization.quantize_model( + model, + quant_map, + "", + ) + return model, quant_map + + def _ft_quant(model_config: Any, quantization: FTQuantize) -> Tuple[nn.Module, QuantizeMapping]: # noqa: UP006 + model = _create_model(model_config) + model.to(quantization.model_dtype) + quant_map = QuantizeMapping({}, {}) + model = quantization.quantize_model( + model, + quant_map, + "", + ) + return model, quant_map + + def _awq_quant( + model_config: Any, quantization: AWQQuantize + ) -> Tuple[nn.Module, QuantizeMapping]: # noqa: UP006 + if awq_unsupported_message is not None: + raise NotImplementedError(awq_unsupported_message) + model = _create_model(model_config) + model.to(quantization.model_dtype) + quant_map = QuantizeMapping({}, {}) + model = quantization.quantize_model( + model, + quant_map, + "", + ) + return model, quant_map + + def _per_tensor_quant( + model_config: Any, + quantization: PerTensorQuantize, + ) -> Tuple[nn.Module, QuantizeMapping]: # noqa: UP006 + model = _create_model(model_config) + model.to(quantization.model_dtype) + quant_map = QuantizeMapping({}, {}) + kwargs = {} + if per_tensor_use_shards: + if not hasattr(model_config, "tensor_parallel_shards"): + raise AttributeError( + "model_config is missing required attribute " + "'tensor_parallel_shards' for per-tensor quantization" + ) + kwargs["tensor_parallel_shards"] = getattr(model_config, "tensor_parallel_shards") + model = quantization.quantize_model( + model, + quant_map, + "", + **kwargs, + ) + return model, quant_map + + def _block_scale_quant( + model_config: Any, + quantization: BlockScaleQuantize, + ) -> Tuple[nn.Module, QuantizeMapping]: # noqa: UP006 + model = _create_model(model_config) + model.to(quantization.model_dtype) + quant_map = QuantizeMapping({}, {}) + model = quantization.quantize_model(model, quant_map, "") + return model, quant_map + + quantize_fns: Dict[str, FuncQuantization] = {"no-quant": _no_quant} # noqa: UP006 + if supports_group_quant: + quantize_fns["group-quant"] = _group_quant + if supports_ft_quant: + quantize_fns["ft-quant"] = _ft_quant + if supports_awq: + quantize_fns["awq"] = _awq_quant + if supports_per_tensor: + quantize_fns["per-tensor-quant"] = _per_tensor_quant + if supports_block_scale: + quantize_fns["block-scale-quant"] = _block_scale_quant + return quantize_fns + + +def make_awq_quant( + model_cls: Type[nn.Module], # noqa: UP006 +) -> Callable[[Any, AWQQuantize], Tuple[nn.Module, QuantizeMapping]]: # noqa: UP006 + """Create a standard AWQ quantization function for loaders.""" + + def awq_quant( + model_config: Any, quantization: AWQQuantize + ) -> Tuple[nn.Module, QuantizeMapping]: # noqa: UP006 + model = model_cls(model_config) + model.to(quantization.model_dtype) + quant_map = QuantizeMapping({}, {}) + model = quantization.quantize_model( + model, + quant_map, + "", + ) + return model, quant_map + + return awq_quant diff --git a/python/mlc_llm/quantization/no_quantization.py b/python/mlc_llm/quantization/no_quantization.py new file mode 100644 index 0000000..d443e94 --- /dev/null +++ b/python/mlc_llm/quantization/no_quantization.py @@ -0,0 +1,15 @@ +"""The no quantization config""" + +from dataclasses import dataclass + + +@dataclass +class NoQuantize: + """Configuration for no quantization""" + + name: str + kind: str + model_dtype: str # "float16", "float32" + + def __post_init__(self): + assert self.kind == "no-quant" diff --git a/python/mlc_llm/quantization/per_tensor_quantization.py b/python/mlc_llm/quantization/per_tensor_quantization.py new file mode 100644 index 0000000..aafccd4 --- /dev/null +++ b/python/mlc_llm/quantization/per_tensor_quantization.py @@ -0,0 +1,700 @@ +"""The per-tensor quantization config""" + +import functools +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Type, Union # noqa: UP035 + +import numpy as np +from tvm import DataType, DataTypeCode, IRModule, relax, runtime, te, tirx, topi +from tvm.relax.frontend import nn +from tvm.runtime import Tensor + +from mlc_llm.loader import QuantizeMapping +from mlc_llm.nn import MixtralExperts +from mlc_llm.op import cutlass, extern +from mlc_llm.support import logging + +from .utils import ( + apply_sharding, + compile_quantize_func, + convert_uint_packed_fp8_to_float, + is_final_fc, + is_moe_gate, + pack_weight, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class PerTensorQuantize: + """Configuration for per-tensor quantization""" + + name: str + kind: str + activation_dtype: Literal["float8_e4m3fn", "float8_e5m2"] + weight_dtype: Literal["float8_e4m3fn", "float8_e5m2"] + storage_dtype: Literal["uint32", "float8_e4m3fn", "float8_e5m2"] + model_dtype: Literal["float16"] + quantize_embedding: bool = True + quantize_final_fc: bool = True + quantize_linear: bool = True + + num_elem_per_storage: int = 0 + max_int_value: int = 0 + use_scale: bool = True + # The calibration mode for quantization. If set to "inference", the model is built for + # inference. This should be used after calibration is done. + # If set to "max", the model is built for calibration that computes the scale using max value of + # the activations. + calibration_mode: Literal["inference", "max"] = "inference" + tensor_parallel_shards: int = 1 + + def __post_init__(self): + assert self.kind == "per-tensor-quant" + self.num_elem_per_storage = ( + DataType(self.storage_dtype).bits // DataType(self.weight_dtype).bits + ) + self.max_int_value = int(tirx.max_value(self.weight_dtype).value) + self._quantize_func_cache = {} + + def quantize_model( + self, + model: nn.Module, + quant_map: QuantizeMapping, + name_prefix: str, + tensor_parallel_shards: int, + ) -> nn.Module: + """ + Quantize model with per-tensor quantization + + Parameters + ---------- + model : nn.Module + The non-quantized nn.Module. + + quant_map : QuantizeMapping + The quantize mapping with name mapping and func mapping. + + name_prefix : str + The name prefix for visited weight. + + tensor_parallel_shards : int + The number of tensor parallel shards. + + Returns + ------- + ret : nn.Module + The quantized nn.Module. + """ + + self.tensor_parallel_shards = tensor_parallel_shards + + class _Mutator(nn.Mutator): + def __init__(self, config: PerTensorQuantize, quant_map: QuantizeMapping) -> None: + super().__init__() + self.config = config + self.quant_map = quant_map + + def visit_module(self, name: str, node: nn.Module) -> Any: + """ + The visiting method for per-tensor quantization of nn.Module nodes. + + Parameters + ---------- + name : str + The name of the current node. + + node : nn.Module + The current node of nn.Module to mutate. + + Returns + ------ + ret_node: Any + The new node to replace current node. + """ + weight_name = f"{name}.weight" + param_names = ( + [f"{name}.q_weight", f"{name}.q_scale"] + if self.config.use_scale + else [ + f"{name}.q_weight", + ] + ) + if ( + isinstance(node, nn.Linear) + and self.config.quantize_linear + and (not is_final_fc(name) or self.config.quantize_final_fc) + and not is_moe_gate(name, node) + ): + self.quant_map.param_map[weight_name] = param_names + self.quant_map.map_func[weight_name] = self.config.quantize_weight + op = PerTensorQuantizeLinear.from_linear(node, self.config, name) + elif isinstance(node, nn.Embedding) and self.config.quantize_embedding: + self.quant_map.param_map[weight_name] = param_names + self.quant_map.map_func[weight_name] = self.config.quantize_weight + op = PerTensorQuantizeEmbedding.from_embedding(node, self.config) + elif isinstance(node, MixtralExperts): + self.quant_map.param_map[weight_name] = param_names + self.quant_map.map_func[weight_name] = self.config.quantize_weight + op = PerTensorQuantizeMixtralExperts.from_mixtral_experts( + node, self.config, name + ) + else: + return self.visit(name, node) + + if hasattr(op, "q_calibration_scale") and op.q_calibration_scale: + # update quant_map for calibration scale + param_name = f"{name}.q_calibration_scale" + old_map_func = self.quant_map.map_func[weight_name] + + def map_func(*args, **kwargs): + # placeholder for calibration scale, the actual value will be set after + # calibration. + scale = runtime.empty( + shape=op.q_calibration_scale.shape, + dtype=op.q_calibration_scale.dtype, + ) + return [*old_map_func(*args, **kwargs), scale] + + self.quant_map.param_map[weight_name].append(param_name) + self.quant_map.map_func[weight_name] = map_func + return op + + model.to(dtype=self.model_dtype) + mutator = _Mutator(self, quant_map) + model = mutator.visit(name_prefix, model) + return model + + def quantize_weight(self, weight) -> List[Tensor]: # noqa: UP006 + """ + Quantize weight with per-tensor quantization. + + Parameters + ---------- + weight : Tensor + The weight to quantize. + + Returns + ------- + ret : List[Tensor] + The quantized weight and the scale if use_scale is True. + """ + device = weight.device + device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()] + + def _create_quantize_func() -> IRModule: + if DataType(self.weight_dtype).type_code in [ + DataTypeCode.Float8E4M3FN, + DataTypeCode.Float8E5M2, + ]: + quantize_func = functools.partial( + self.quantize_float8, + quantize_dtype=self.weight_dtype, + storage_dtype=self.storage_dtype, + ) + else: + assert NotImplementedError() + + class Quantizer(nn.Module): + """Quantizer module for per-tensor quantization.""" + + def main(self, weight: nn.Tensor): + return quantize_func(weight) + + mod = Quantizer() + mod, _ = mod.export_tvm( + spec={"main": {"weight": nn.spec.Tensor(weight.shape, weight.dtype)}} + ) + return mod + + key = f"({weight.shape}, {weight.dtype}, {device_type}" + quantize_func = self._quantize_func_cache.get(key, None) + if quantize_func is None: + logger.info("Compiling quantize function for key: %s", key) + quantize_func = compile_quantize_func(_create_quantize_func(), device) + self._quantize_func_cache[key] = quantize_func + return quantize_func(weight) + + def quantize_float8( + self, + tensor: nn.Tensor, + quantize_dtype: str, + storage_dtype: str, + ) -> Union[Tuple[nn.Tensor], Tuple[nn.Tensor, nn.Tensor]]: # noqa: UP006 + """Per-tensor quantization for weight tensor, defined in tensor expression.""" + + if self.use_scale: + # min_scaling_factor taken from TRT-LLM + def _compute_scale(x: te.Tensor) -> te.Tensor: + max_abs = topi.max(topi.abs(x)) + min_scaling_factor = tirx.const( + 1.0 / (self.max_int_value * 512.0), self.model_dtype + ) + scale = topi.maximum( + max_abs.astype(self.model_dtype) / self.max_int_value, + min_scaling_factor, + ).astype("float32") + scale = topi.expand_dims(scale, axis=0) + return scale + + scale = nn.tensor_expr_op(_compute_scale, "compute_scale", args=[tensor]) + else: + scale = None + + def _compute_quantized_tensor(weight: te.Tensor, scale: Optional[te.Tensor]) -> te.Tensor: + elem_storage_dtype = ( + f"uint{DataType(quantize_dtype).bits}" + if DataType(storage_dtype).type_code == DataTypeCode.UINT + else quantize_dtype + ) + scaled_tensor = te.compute( + shape=weight.shape, + fcompute=lambda *idx: tirx.Cast( + self.storage_dtype, + tirx.reinterpret( + elem_storage_dtype, + tirx.Cast( + quantize_dtype, + weight(*idx) / scale(0) if scale is not None else weight(*idx), + ), + ), + ), + ) + + if quantize_dtype == self.storage_dtype: + return scaled_tensor + + packed_weight = pack_weight( + scaled_tensor, + axis=-1, + num_elem_per_storage=self.num_elem_per_storage, + weight_dtype=self.weight_dtype, + storage_dtype=self.storage_dtype, + ) + + return packed_weight + + quantized_tensor = nn.tensor_expr_op( + _compute_quantized_tensor, "compute_quantized_tensor", args=[tensor, scale] + ) + + if self.use_scale: + return quantized_tensor, scale + return (quantized_tensor,) + + def _dequantize( + self, + q_weight: te.Tensor, + scale: Optional[te.Tensor] = None, + out_shape: Optional[Sequence[tirx.Expr]] = None, + ) -> te.Tensor: + if self.use_scale: + assert scale is not None + if DataType(self.weight_dtype).type_code in [ + DataTypeCode.Float8E4M3FN, + DataTypeCode.Float8E5M2, + ]: + return self.dequantize_float8(q_weight, scale, self.weight_dtype, out_shape) + raise NotImplementedError() + + def dequantize_float8( + self, + q_tensor: te.Tensor, + scale: Optional[te.Tensor], + quantize_dtype: str, + out_shape: Optional[Sequence[tirx.Expr]] = None, + ) -> te.Tensor: + """Dequantize a fp8 tensor (input or weight) to higher-precision float.""" + if quantize_dtype != self.storage_dtype: + dequantized_tensor = convert_uint_packed_fp8_to_float( + q_tensor, + self.num_elem_per_storage, + self.storage_dtype, + self.model_dtype, + quantize_dtype, + axis=-1, + out_shape=out_shape, + ) + else: + dequantized_tensor = q_tensor.astype(self.model_dtype) + if scale is not None: + dequantized_tensor = dequantized_tensor * scale.astype(dequantized_tensor.dtype) + return dequantized_tensor + + +class PerTensorQuantizeLinear(nn.Module): + """An nn.Linear module with per-tensor quantization.""" + + def __init__( + self, + in_features: int, + out_features: Union[int, tirx.Var], + config: PerTensorQuantize, + name: str, + bias: bool = True, + out_dtype: Optional[str] = None, + ) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.out_dtype = out_dtype or config.model_dtype + self.config = config + self.name = name + self.q_weight = nn.Parameter( + (out_features, tirx.ceildiv(in_features, config.num_elem_per_storage)), + config.storage_dtype, + ) + self.q_calibration_scale = None + if config.use_scale: + self.q_scale = nn.Parameter((1,), "float32") + if config.calibration_mode == "inference": + self.q_calibration_scale = nn.Parameter((1,), "float32") + else: + self.q_scale = None + if bias: + self.bias = nn.Parameter( + (out_features,), config.model_dtype if out_dtype is None else out_dtype + ) + else: + self.bias = None + + @classmethod + def from_linear( + cls, src: nn.Linear, config: PerTensorQuantize, name: str + ) -> "PerTensorQuantizeLinear": + """ + Converts a non-quantized nn.Linear to a per-tensor quantized PerTensorQuantizeLinear + + Parameters + ---------- + src : nn.Linear + The non-quantized nn.Linear. + + config : PerTensorQuantize + The per-tensor quantization config. + + name: str + The name of the layer. + + Returns + ------- + ret : PerTensorQuantizeLinear + The per-tensor quantized PerTensorQuantizeLinear layer. + """ + out_features, in_features = src.weight.shape + quantized_linear = cls( + in_features=in_features, + out_features=out_features, + config=config, + name=name, + bias=getattr(src, "bias", None) is not None, + out_dtype=src.out_dtype, + ) + if quantized_linear.bias is not None: + quantized_linear.bias.attrs = src.bias.attrs + if "shard_strategy" in src.weight.attrs: + shard = src.weight.attrs["shard_strategy"] + apply_sharding(shard, f"{shard.name}_q_weight", quantized_linear.q_weight) + # scale doesn't need to be sharded since it's the same for all shards + return quantized_linear + + def forward(self, x: nn.Tensor) -> nn.Tensor: + """ + Forward method for per-tensor quantized linear layer. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the per-tensor quantized linear layer. + """ + # Note: Use calibration scale when calibration is enabled + if self.config.calibration_mode == "inference": + if self.q_calibration_scale: + x /= self.q_calibration_scale.astype(x.dtype) + x_q = x.astype(self.config.activation_dtype) + x_scale = self.q_calibration_scale + elif self.config.calibration_mode == "max": + _, x_scale = self.config.quantize_float8( + x, + quantize_dtype=self.config.activation_dtype, + storage_dtype=self.config.storage_dtype, + ) + if self.config.tensor_parallel_shards > 1: + x_scale = nn.ccl_allreduce(x_scale, "max") + x_scale = nn.extern( + "mlc_llm.calibration_observer", + [f"{self.name}.q_calibration_scale", "max", x_scale], + out=nn.Tensor.placeholder(x_scale.shape, x_scale.dtype), + ) + x_q = (x / x_scale.astype(x.dtype)).astype(self.config.activation_dtype) + x = x_q.astype(self.config.model_dtype) * x_scale.astype(self.config.model_dtype) + else: + raise ValueError(f"Unknown calibration mode: {self.config.calibration_mode}") + + if ( + self.config.weight_dtype == self.config.storage_dtype + and self.config.calibration_mode == "inference" + ): + if ( + extern.get_store().cutlass_gemm + and functools.reduce(lambda x, y: x * y, x_q.shape[:-1]) != 1 + ): + # Dispatch to cutlass kernel for gemm when cutlass is available. + scale = ( + x_scale * self.q_scale + if self.config.use_scale + else nn.wrap_nested( + relax.Constant(runtime.tensor(np.array([1.0]).astype("float32"))), + "scale", + ) + ) + return cutlass.fp8_gemm( + x_q, + self.q_weight, + scale, + self.config.weight_dtype, + self.config.model_dtype, + ) + x = nn.op.matmul(x_q, nn.permute_dims(self.q_weight), out_dtype="float32") + if self.config.use_scale: + scale = x_scale * self.q_scale + x = x * scale + x = x.astype(self.out_dtype) + else: + w = nn.op.tensor_expr_op( + lambda weight, scale: self.config._dequantize( + weight, + scale, + out_shape=[ + ( + tirx.IntImm("int64", self.out_features) + if isinstance(self.out_features, int) + else weight.shape[0] + ), + tirx.IntImm("int64", self.in_features), + ], + ), + "dequantize", + args=[self.q_weight, self.q_scale], + ) + x = nn.op.matmul(x, nn.permute_dims(w), out_dtype=self.out_dtype) + if self.bias is not None: + x = x + self.bias + return x + + def to(self, dtype: Optional[str] = None) -> None: + """ + Override to() such that we do not convert bias if there is an out_dtype. + Otherwise, we might run into dtype mismatch when computing x + self.bias. + """ + self.q_weight.to(dtype=dtype) + if self.q_scale: + self.q_scale.to(dtype=dtype) + if self.bias is not None and self.out_dtype is None: + self.bias.to(dtype=dtype) + if dtype is not None and isinstance(getattr(self, "dtype", None), str): + self.dtype = dtype + + +class PerTensorQuantizeEmbedding(nn.Module): + """An nn.Embedding module with group quantization""" + + def __init__(self, num: Union[int, tirx.Var], dim: int, config: PerTensorQuantize): + self.num = num + self.dim = dim + self.config = config + self.q_weight = nn.Parameter( + (num, tirx.ceildiv(dim, config.num_elem_per_storage)), config.storage_dtype + ) + if self.config.use_scale: + self.q_scale = nn.Parameter((1,), "float32") + else: + self.q_scale = None + + @staticmethod + def from_embedding( + embedding: nn.Embedding, config: PerTensorQuantize + ) -> "PerTensorQuantizeEmbedding": + """ + Converts a non-quantized nn.Embedding to a per-tensor quantized PerTensorQuantizeEmbedding + + Parameters + ---------- + linear : nn.Embedding + The non-quantized nn.Embedding. + + config : PerTensorQuantize + The per-tensor quantization config. + + Returns + ------- + ret : PerTensorQuantizeEmbedding + The per-tensor quantized embedding layer. + """ + num, dim = embedding.weight.shape + return PerTensorQuantizeEmbedding(num, dim, config) + + def forward(self, x: nn.Tensor): + """ + Forward method for per-tensor quantized embedding layer. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the embedding layer. + """ + w = nn.op.tensor_expr_op( + lambda weight, scale: self.config._dequantize( + weight, + scale, + out_shape=[ + ( + tirx.IntImm("int64", self.num) + if isinstance(self.num, int) + else weight.shape[0] + ), + tirx.IntImm("int64", self.dim), + ], + ), + "dequantize", + args=[self.q_weight, self.q_scale], + ) + if x.ndim == 1: + return nn.op.take(w, x, axis=0) + return nn.op.reshape( + nn.op.take(w, nn.op.reshape(x, shape=[-1]), axis=0), + shape=[*x.shape, self.dim], + ) + + def lm_head_forward(self, x: nn.Tensor): + """The lm_head forwarding, which dequantizes the weight + and multiplies it with the input tensor. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + Returns + ------- + ret : nn.Tensor + The output tensor for the lm_head layer. + """ + w = nn.op.tensor_expr_op( + lambda weight, scale: self.config._dequantize( + weight, + scale, + out_shape=[ + ( + tirx.IntImm("int64", self.num) + if isinstance(self.num, int) + else weight.shape[0] + ), + tirx.IntImm("int64", self.dim), + ], + ), + "dequantize", + args=[self.q_weight, self.q_scale], + ) + w = nn.op.permute_dims(w) + return nn.op.matmul(x, w, out_dtype="float32") + + +class PerTensorQuantizeMixtralExperts(nn.Module): + """An MixtralExperts module with group quantization""" + + _IMPL: ClassVar[Dict[str, Type["PerTensorQuantizeMixtralExperts"]]] = {} # noqa: UP006 + + def __init__( + self, + num_local_experts, + in_features, + out_features, + config: PerTensorQuantize, + name: str, + ): + self.num_local_experts = num_local_experts + self.in_features = in_features + self.out_features = out_features + self.config = config + self.name = name + self.q_weight = nn.Parameter( + ( + num_local_experts, + out_features, + tirx.ceildiv(in_features, config.num_elem_per_storage), + ), + config.storage_dtype, + ) + self.q_calibration_scale = None + if config.use_scale: + self.q_scale = nn.Parameter((1,), "float32") + if config.calibration_mode == "inference": + self.q_calibration_scale = nn.Parameter((1,), "float32") + else: + self.q_scale = None + + @staticmethod + def from_mixtral_experts( + src: "MixtralExperts", + config: PerTensorQuantize, + name: str, + ) -> "PerTensorQuantizeMixtralExperts": + """ + Converts a non-quantized MixtralExperts to a per-tensor quantized + PerTensorQuantizeMixtralExperts + + Parameters + ---------- + src : MixtralExperts + The non-quantized MixtralExperts + + config : PerTensorQuantize + The per-tensor quantization config + + name: str + The name of the layer. + + Returns + ------- + ret : PerTensorQuantizeMixtralExperts + The per-tensor quantized MixtralExperts layer + """ + if DataType(config.weight_dtype).type_code in [ + DataTypeCode.Float8E4M3FN, + DataTypeCode.Float8E5M2, + ]: + return PerTensorQuantizeMixtralExperts._IMPL["fp8"].from_mixtral_experts( + src, config, name + ) + raise NotImplementedError() + + def forward(self, x: nn.Tensor, indptr: nn.Tensor) -> nn.Tensor: + """Forward method for per-tensor quantized mistral experts. + + Parameters + ---------- + x : nn.Tensor + The input tensor. + + indptr: nn.Tensor + The indptr tensor + + Returns + ------- + ret : nn.Tensor + The output tensor for the per-tensor quantized mistral experts layer. + """ + raise NotImplementedError() diff --git a/python/mlc_llm/quantization/quantization.py b/python/mlc_llm/quantization/quantization.py new file mode 100644 index 0000000..fb6d9f7 --- /dev/null +++ b/python/mlc_llm/quantization/quantization.py @@ -0,0 +1,201 @@ +"""A centralized registry of all existing quantization methods and their configurations.""" + +from typing import Any, Dict # noqa: UP035 + +from .awq_quantization import AWQQuantize +from .block_scale_quantization import BlockScaleQuantize +from .ft_quantization import FTQuantize +from .group_quantization import GroupQuantize +from .no_quantization import NoQuantize +from .per_tensor_quantization import PerTensorQuantize + +Quantization = Any +"""Quantization is an object that represents an quantization algorithm. It is required to +have the following fields: + + name : str + The name of the quantization algorithm, for example, "q4f16_1". + + kind : str + The kind of quantization algorithm, for example, "group-quant", "faster-transformer". + +It is also required to have the following method: + + def quantize_model(self, module: nn.Module) -> nn.Module: + ... + + def quantize_weight(self, weight: tvm.runtime.Tensor) -> List[tvm.runtime.Tensor]: + ... +""" + +QUANTIZATION: Dict[str, Quantization] = { # noqa: UP006 + "q0f16": NoQuantize( + name="q0f16", + kind="no-quant", + model_dtype="float16", + ), + "q0bf16": NoQuantize( + name="q0bf16", + kind="no-quant", + model_dtype="bfloat16", + ), + "q0f32": NoQuantize( + name="q0f32", + kind="no-quant", + model_dtype="float32", + ), + "q3f16_0": GroupQuantize( + name="q3f16_0", + kind="group-quant", + group_size=40, + quantize_dtype="int3", + storage_dtype="uint32", + model_dtype="float16", + linear_weight_layout="KN", + quantize_embedding=True, + quantize_final_fc=True, + ), + "q3f16_1": GroupQuantize( + name="q3f16_1", + kind="group-quant", + group_size=40, + quantize_dtype="int3", + storage_dtype="uint32", + model_dtype="float16", + linear_weight_layout="NK", + quantize_embedding=True, + quantize_final_fc=True, + ), + "q4f16_0": GroupQuantize( + name="q4f16_0", + kind="group-quant", + group_size=32, + quantize_dtype="int4", + storage_dtype="uint32", + model_dtype="float16", + linear_weight_layout="KN", + quantize_embedding=True, + quantize_final_fc=True, + ), + "q4f16_1": GroupQuantize( + name="q4f16_1", + kind="group-quant", + group_size=32, + quantize_dtype="int4", + storage_dtype="uint32", + model_dtype="float16", + linear_weight_layout="NK", + quantize_embedding=True, + quantize_final_fc=True, + ), + "q4bf16_0": GroupQuantize( + name="q4bf16_0", + kind="group-quant", + group_size=32, + quantize_dtype="int4", + storage_dtype="uint32", + model_dtype="bfloat16", + linear_weight_layout="KN", + quantize_embedding=True, + quantize_final_fc=True, + ), + "q4bf16_1": GroupQuantize( + name="q4bf16_1", + kind="group-quant", + group_size=32, + quantize_dtype="int4", + storage_dtype="uint32", + model_dtype="bfloat16", + linear_weight_layout="NK", + quantize_embedding=True, + quantize_final_fc=True, + ), + "q4f32_1": GroupQuantize( + name="q4f32_1", + kind="group-quant", + group_size=32, + quantize_dtype="int4", + storage_dtype="uint32", + model_dtype="float32", + linear_weight_layout="NK", + quantize_embedding=True, + quantize_final_fc=True, + ), + "q4f16_2": GroupQuantize( + name="q4f16_2", + kind="group-quant", + group_size=32, + quantize_dtype="int4", + storage_dtype="uint32", + model_dtype="float16", + linear_weight_layout="NK", + quantize_embedding=False, + quantize_final_fc=False, + ), + "q4f16_autoawq": AWQQuantize( + name="q4f16_autoawq", + kind="awq", + group_size=128, + quantize_dtype="int4", + storage_dtype="uint32", + model_dtype="float16", + ), + "q4f16_ft": FTQuantize( + name="q4f16_ft", + kind="ft-quant", + quantize_dtype="int4", + storage_dtype="int8", + model_dtype="float16", + ), + "e5m2_e5m2_f16": PerTensorQuantize( + name="e5m2_e5m2_f16", + kind="per-tensor-quant", + activation_dtype="float8_e5m2", + weight_dtype="float8_e5m2", + storage_dtype="float8_e5m2", + model_dtype="float16", + quantize_final_fc=False, + quantize_embedding=False, + quantize_linear=True, + use_scale=False, + ), + "e4m3_e4m3_f16": PerTensorQuantize( + name="e4m3_e4m3_f16", + kind="per-tensor-quant", + activation_dtype="float8_e4m3fn", + weight_dtype="float8_e4m3fn", + storage_dtype="float8_e4m3fn", + model_dtype="float16", + quantize_final_fc=False, + quantize_embedding=False, + quantize_linear=True, + use_scale=True, + calibration_mode="inference", + ), + "e4m3_e4m3_f16_max_calibrate": PerTensorQuantize( + name="e4m3_e4m3_f16_max_calibrate", + kind="per-tensor-quant", + activation_dtype="float8_e4m3fn", + weight_dtype="float8_e4m3fn", + storage_dtype="float8_e4m3fn", + model_dtype="float16", + quantize_final_fc=False, + quantize_embedding=False, + quantize_linear=True, + use_scale=True, + calibration_mode="max", + ), + "fp8_e4m3fn_bf16_block_scale": BlockScaleQuantize( + name="fp8_e4m3fn_bf16_block_scale", + kind="block-scale-quant", + weight_dtype="float8_e4m3fn", + model_dtype="bfloat16", + ), + "fp8_e4m3fn_bf16_block_scale_static_activation": BlockScaleQuantize( + name="fp8_e4m3fn_bf16_block_scale_static_activation", + kind="block-scale-quant", + weight_dtype="float8_e4m3fn", + model_dtype="bfloat16", + use_activation_scale=True, + ), +} diff --git a/python/mlc_llm/quantization/utils.py b/python/mlc_llm/quantization/utils.py new file mode 100644 index 0000000..429d593 --- /dev/null +++ b/python/mlc_llm/quantization/utils.py @@ -0,0 +1,188 @@ +"""Common utilities for quantization""" + +from collections.abc import Sequence +from typing import Callable, List, Optional # noqa: UP035 + +from tvm import IRModule, relax, te, tirx +from tvm.relax.frontend import nn +from tvm.runtime import DataType, DataTypeCode +from tvm.s_tir import dlight as dl +from tvm.target import Target + +from mlc_llm.support import tensor_parallel as tp + + +def convert_uint_to_float( + weight: te.Tensor, + bits: int, + num_elem_per_storage: int, + storage_dtype: str, + model_dtype: str, + axis: int = -1, + out_shape: Optional[List[tirx.Expr]] = None, # noqa: UP006 + ft_reorder: Optional[bool] = False, +) -> te.Tensor: + """Convert a quantized uint weight to an unquantized float weight.""" + tir_bin_mask = tirx.const((1 << bits) - 1, storage_dtype) + if out_shape is None: + out_shape = weight.shape + out_shape[axis] *= num_elem_per_storage + axis = axis if axis >= 0 else len(out_shape) + axis + return te.compute( + shape=out_shape, + fcompute=lambda *idx: tirx.bitwise_and( + tirx.shift_right( + weight(*idx[:axis], idx[axis] // num_elem_per_storage, *idx[axis + 1 :]), + ( + ( + (idx[axis] % num_elem_per_storage) % 2 * 4 + + (idx[axis] % num_elem_per_storage) // 2 + ) + * bits + if ft_reorder + else (idx[axis] % num_elem_per_storage) * bits + ).astype(storage_dtype), + ), + tir_bin_mask, + ).astype(model_dtype), + ) + + +def is_final_fc(name: str) -> bool: + """Determines whether the parameter is the last layer based on its name.""" + # TODO: use more specious condition to determine final fc + return name in ["head", "lm_head", "lm_head.linear", "embed_out"] + + +def is_moe_gate(name: str, node: nn.Linear) -> bool: + """Check whether the parameter is the MoE gate layer.""" + return name.endswith("gate") and isinstance(node.out_features, int) and node.out_features <= 256 + + +def compile_quantize_func(mod: IRModule, device) -> Callable: + """Compile a quantization function for a given device.""" + device_type = device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()] + if device_type in ["cuda", "rocm", "metal", "vulkan", "opencl"]: + target = Target.current() + if target is None: + target = Target.from_device(device) + with target: + mod = dl.ApplyDefaultSchedule( + dl.gpu.Reduction(), + dl.gpu.GeneralReduction(), + dl.gpu.Fallback(), + )(mod) + elif device_type == "cpu": + target = "llvm" + mod = relax.transform.LegalizeOps()(mod) + else: + raise NotImplementedError(f"Device type {device_type} is not supported") + ex = relax.build(mod, target=target) + vm = relax.VirtualMachine(ex, device) + return vm["main"] + + +def apply_sharding(shard_strategy, name: str, weight: nn.Parameter): + """Apply sharding strategy to a weight.""" + if isinstance(shard_strategy, tp.ShardSingleDim): + weight.attrs["shard_strategy"] = tp.ShardSingleDim( + name=name, + dim=shard_strategy.dim, + segs=shard_strategy.segs, + ) + else: + raise NotImplementedError(f"Unknowing sharding strategy: {shard_strategy}") + + +def convert_uint_packed_fp8_to_float( + weight: te.Tensor, + num_elem_per_storage: int, + storage_dtype: str, + model_dtype: str, + quant_dtype: str, + axis: int = -1, + out_shape: Optional[Sequence[tirx.Expr]] = None, +) -> te.Tensor: + """Unpack a fp8 value from the storage dtype and convert to float.""" + assert quant_dtype in ["float8_e4m3fn", "float8_e5m2"] + assert DataType(storage_dtype).type_code == DataTypeCode.UINT + bits = DataType(quant_dtype).bits + elem_storage_dtype = DataType(f"uint{bits}") + tir_bin_mask = tirx.const((1 << bits) - 1, "uint8") + if axis < 0: + axis += len(weight.shape) + if out_shape is None: + out_shape = ( + *weight.shape[:axis], + weight.shape[axis] * num_elem_per_storage, + *weight.shape[axis + 1 :], + ) + axis = axis if axis >= 0 else len(out_shape) + axis + return te.compute( + shape=out_shape, + fcompute=lambda *idx: tirx.reinterpret( + quant_dtype, + tirx.bitwise_and( + tirx.shift_right( + weight(*idx[:axis], idx[axis] // num_elem_per_storage, *idx[axis + 1 :]), + ((idx[axis] % num_elem_per_storage) * bits).astype(storage_dtype), + ).astype(elem_storage_dtype), + tir_bin_mask, + ), + ).astype(model_dtype), + ) + + +def pack_weight( + weight: te.Tensor, + axis: int, + num_elem_per_storage: int, + weight_dtype: str, + storage_dtype: str, + out_shape: Optional[Sequence[tirx.Expr]] = None, +): + """Convert a tensor to a packed format by packing consecutive bits. + This can be useful for sub-byte quantization. + + Parameters + ---------- + weight : te.Tensor + The weight + axis : int + The axis to pack. + num_elem_per_storage : int + The number of elements per storage. + weight_dtype : str + The dtype of the input tensor. + storage_dtype : str + The dtype of the packed tensor. + out_shape : Optional[Sequence[tirx.Expr]] + The output shape of the packed tensor. Zero-padding is added if needed. + """ + assert weight.dtype == storage_dtype + shape = weight.shape + if axis < 0: + axis += len(shape) + k = shape[axis] + axis = axis if axis >= 0 else len(shape) + axis + if out_shape is None: + out_shape = ( + *shape[:axis], + tirx.ceildiv(k, num_elem_per_storage), + *shape[axis + 1 :], + ) + r = te.reduce_axis((0, num_elem_per_storage), name="r") + packed_weight = te.compute( + shape=out_shape, + fcompute=lambda *idx: tirx.sum( + tirx.if_then_else( + idx[axis] * num_elem_per_storage + r < k, + weight(*idx[:axis], idx[axis] * num_elem_per_storage + r, *idx[axis + 1 :]) + << (r * DataType(weight_dtype).bits), + tirx.const(0, storage_dtype), + ), + axis=r, + ), + name="packed_weight", + ).astype(storage_dtype) + return packed_weight diff --git a/python/mlc_llm/router/__init__.py b/python/mlc_llm/router/__init__.py new file mode 100644 index 0000000..921aa96 --- /dev/null +++ b/python/mlc_llm/router/__init__.py @@ -0,0 +1,4 @@ +"""Subdirectory of router, which routes to multiple engine endpoints.""" + +from .. import base +from .router import Router diff --git a/python/mlc_llm/router/router.py b/python/mlc_llm/router/router.py new file mode 100644 index 0000000..115b446 --- /dev/null +++ b/python/mlc_llm/router/router.py @@ -0,0 +1,390 @@ +"""Programmable router for dispatching OpenAI API to Microserving API""" + +import json +import math +import threading +from collections.abc import AsyncGenerator, Iterable +from typing import Any, List, Literal, Optional, Tuple # noqa: UP035 + +import aiohttp +import tvm + +from mlc_llm.protocol import openai_api_protocol +from mlc_llm.serve import EngineConfig, PopenServer +from mlc_llm.serve.entrypoints import microserving_entrypoints +from mlc_llm.tokenizers import Tokenizer + + +class Router: + """Programmable Router Implementation""" + + def __init__( + self, + model: str, + model_lib: Optional[str] = None, + hosts: Optional[List[str]] = None, # noqa: UP006 + ports: Optional[List[int]] = None, # noqa: UP006 + num_gpus: Optional[List[int]] = None, # noqa: UP006 + enable_prefix_cache: bool = False, + router_mode: Literal["disagg", "round-robin"] = "disagg", + pd_balance_factor: float = 0.0, + ): + """ + Spawn len(host_list) server endpoints with Popen. + """ + if hosts is None: + hosts = ["127.0.0.1"] + if ports is None: + ports = [8080] + if num_gpus is None: + num_gpus = [1] + + self.router_mode = router_mode + self.pd_balance_factor = pd_balance_factor + # Get endpoint urls + self.num_servers = len(hosts) + assert self.num_servers == len(ports) == len(num_gpus) + self.hosts = hosts + self.ports = ports + self.server_urls = [] + for i in range(self.num_servers): + self.server_urls.append(f"http://{hosts[i]}:{ports[i]}") + + # Misc + self.headers = {"Content-Type": "application/json"} + self.num_running_requests = [0] * self.num_servers + + # Call nvshmem_init here to get uid, then pass to env variables to server.start() below + f_init_nvshmem_uid = tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid") + uid = list(f_init_nvshmem_uid()) + + # Start underlying servers concurrently. Otherwise 1 server cannot start on its own + # since initializing nvhsmem world requires all GPUs. + self.servers: List[PopenServer] = [] # noqa: UP006 + + self.device_id_starts = [0] + for num_gpus_val in num_gpus: + self.device_id_starts.append(self.device_id_starts[-1] + num_gpus_val) + # device_id_starts[-1] is the total number of GPUs. + + def start_server(i: int): + nvshmem_config = { + "uid": uid, + "npes": self.device_id_starts[-1], # total number of workers in the nvshmem world + "pe_start": self.device_id_starts[i], # start of PE for this endpoint's workers + } + + server = PopenServer( + model=model, + model_lib=model_lib, + host=hosts[i], + port=ports[i], + enable_debug=True, + device=f"cuda:{self.device_id_starts[i]}", + mode="server", + engine_config=EngineConfig( + prefix_cache_mode="radix" if enable_prefix_cache else "disable", + gpu_memory_utilization=0.8, + ), + ) + self.servers.append(server) + server.start(extra_env={"MLC_NVSHMEM_INIT_CONFIG_JSON_STR": json.dumps(nvshmem_config)}) + + threads = [] + num_used_gpus = 0 + for i in range(self.num_servers): + thread = threading.Thread( + target=start_server, + args=[i], + ) + num_used_gpus += num_gpus[i] + thread.start() + threads.append(thread) + for thread in threads: + thread.join() + self.tokenizer = Tokenizer(model) + + def terminate(self): + """Terminate the underlying servers""" + for server in self.servers: + server.terminate() + + async def handle_completion( + self, + request: openai_api_protocol.CompletionRequest, + request_id: str, + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + """ + Handle a completion request from API with a schedule. + """ + if isinstance(request.prompt, str): + request.prompt = self.tokenizer.encode(request.prompt) + # Add a debugConfig if not present + if request.debug_config is None: + request.debug_config = openai_api_protocol.DebugConfig() + completed = False + while not completed: + completed = True + async for response in self.translate_request(request, request_id): + if response is None: + completed = False + break + yield response + + async def translate_request( + self, request: openai_api_protocol.CompletionRequest, request_id: str + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + """ + Translate OpenAI API request to microserving API calls. + """ + if self.router_mode == "disagg": + async for response in self._handle_completion_disagg( + request, request_id, pd_balance_factor=self.pd_balance_factor + ): + yield response + elif self.router_mode == "round-robin": + async for response in self._handle_completion_round_robin(request): + yield response + else: + raise ValueError("Cannot reach here") + + def _pick_endpoint(self, endpoint_ids: Iterable[int]) -> int: + # Pick the least congested endpoint. + endpoint_id = -1 + min_running_req = int(1e9) + for candidate_id in endpoint_ids: + if self.num_running_requests[candidate_id] < min_running_req: + min_running_req = self.num_running_requests[candidate_id] + endpoint_id = candidate_id + assert endpoint_id != -1 + return endpoint_id + + async def _handle_completion_round_robin( + self, + request: openai_api_protocol.CompletionRequest, + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + """ + Handle a completion request from API. Given a streaming request, yields multiple response + chunks. Given a non-streaming request, yield a single response. Dispatch request to + endpoints with round-robin scheduling at a request level. + """ + # Round robin + cur_endpoint = self._pick_endpoint(range(self.num_servers)) + self.num_running_requests[cur_endpoint] += 1 + payload = request.model_dump() + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=3 * 3600), trust_env=True + ) as session: + # todo: replace this with start_generate + async with session.post( + self.server_urls[cur_endpoint] + "/v1/completions", + json=payload, + headers=self.headers, + ) as response: + assert response.status == 200, await response.text() + if payload["stream"]: + async for chunk in response.content: + # Convert raw bytes to CompletionResponse + chunk = chunk.strip() + if not chunk or chunk == b"\n": + continue + # Get rid of the prefix "data: " and suffix "\n" + raw_data = chunk[6:].strip() + if raw_data == b"[DONE]": + continue + data = json.loads(raw_data) + # Commented because we still want usage chunk to be passed back + # if not data["choices"]: + # continue + response = openai_api_protocol.CompletionResponse.model_validate(data) + if response.choices: + reason = response.choices[0].finish_reason + if reason == "preempt": + yield None + yield response + else: + data = await response.json() + response = openai_api_protocol.CompletionResponse.model_validate(data) + if response.choices: + reason = response.choices[0].finish_reason + if reason == "preempt": + yield None + yield response + self.num_running_requests[cur_endpoint] -= 1 + + # Below methods are for disaggregated serving + # Note that only _handle_completion_disagg() has scheduling logics. The other three + # helper methods only reflect our flow. + async def _handle_completion_disagg( + self, + original_request: openai_api_protocol.CompletionRequest, + request_id: str, + pd_balance_factor=0, + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + """ + Handle a completion request from API with disaggregated scheduling. Given two servers + P (prefill) and D (decode), the router does the following: + 1. Ask D to prepare metadata, receive D's metadata + (prefix cache, KV append positions, etc.) + 2. Send P the prefill request and D's metadata, receive ack + 3. Ask D to start decoding, receive response as a normal streaming + """ + original_request.user = request_id + # Arbitrarily determine server 0 is P, other servers are D + prefill_server_id = 0 + decode_server_id = self._pick_endpoint(range(1, self.num_servers)) + + # Tell D to prepare metadata for prompt[0:kv_window_end]. + # P does not need to sample. Ask D to treat the last + # token like the first sampled token. + kv_window_end = ( + -1 + if math.fabs(pd_balance_factor) < 1e-5 + else int((1 - pd_balance_factor) * len(original_request.prompt)) + ) + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=3 * 3600), trust_env=True + ) as session: + self.num_running_requests[decode_server_id] += 1 + try: + # 1. Ask D to prepare metadata + prep_recv_request = microserving_entrypoints.PrepRecvRequest( + **original_request.model_dump(), end=kv_window_end + ) + ( + kv_append_metadata_base64, + prefix_matched_length, + ) = await self.send_prepare_receive( + session=session, + request=prep_recv_request, + server_url=self.server_urls[decode_server_id], + ) + + kv_window_end = ( + len(original_request.prompt) + kv_window_end + if kv_window_end < 0 + else kv_window_end + ) + assert prefix_matched_length <= kv_window_end + + # 2. Send P the prefill request and D's metadata. When it returns, it means that + # KV transfer has finished prefilling and transferring the KV of + # prompt[prefix_matched_length:kv_window_end]. So D is ready to decode. + if prefix_matched_length < kv_window_end: + remote_send_request = microserving_entrypoints.RemoteSendRequest( + **original_request.model_dump(), + begin=prefix_matched_length, + end=kv_window_end, + kv_addr_info=kv_append_metadata_base64, + recv_rank=self.device_id_starts[decode_server_id], + ) + await self.send_remote_send( + session=session, + request=remote_send_request, + server_url=self.server_urls[prefill_server_id], + ) + + # 3. Start decoding, receive and yield back response as a normal request + # The kv window passed through denotes the range to prefill on the + # decode server, which should be [-1:] here. + start_generate_request = microserving_entrypoints.StartGenerateRequest( + **original_request.model_dump(), + begin=kv_window_end, + ) + async for response in self.send_start_generate( + session=session, + request=start_generate_request, + server_url=self.server_urls[decode_server_id], + ): + if len(response.choices) > 0: + finish_reason = response.choices[0].finish_reason + if finish_reason == "preempt": + yield None + yield response + except Exception as e: + self.num_running_requests[decode_server_id] -= 1 + raise e + self.num_running_requests[decode_server_id] -= 1 + + async def send_prepare_receive( + self, + session: aiohttp.ClientSession, + request: openai_api_protocol.CompletionRequest, + server_url: str, + ) -> Tuple[str, int]: # noqa: UP006 + """ + Performs step 1 of disaggregated serving: ask D to prepare metadata. + Returns: + The metadata received from D, which is a tuple of 2 elements: + - kv_append_metadata_base64: str, info about KV append encoded in base64 string + - prefix_matched_length: int, length of the matched prefix. + i.e. prompt[0:prefix_matched_length] is the matched prefix + """ + # Send request to the decode server for receive preparation. + # Get the prompt length, matched prefix length and the KV metadata. + async with session.post( + server_url + "/microserving/prep_recv", + json=request.model_dump(), + headers=self.headers, + ) as response: + assert response.status == 200, await response.text() + data = await response.json() + + return ( + data["kv_append_metadata"], + data["prefix_matched_length"], + ) + + async def send_remote_send( + self, + session: aiohttp.ClientSession, + request: openai_api_protocol.CompletionRequest, + server_url: str, + ) -> None: + """ + Performs step 2 of disaggregated serving: ask P to prefill and transfer KV to D. + P returns an empty chunk to acknowledge completion. + """ + # Send request to P and get ack + async with session.post( + server_url + "/microserving/remote_send", + json=request.model_dump(), + headers=self.headers, + ) as response: + assert response.status == 200, await response.text() + await response.json() + + async def send_start_generate( + self, + session: aiohttp.ClientSession, + request: openai_api_protocol.CompletionRequest, + server_url: str, + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + """ + Performs step 3 of disaggregated serving: ask D to decode and return normal response. + """ + # Todo: return string directly to reduce str->json->str roundtrip overhead + async with session.post( + server_url + "/microserving/start_generate", + json=request.model_dump(), + headers=self.headers, + ) as response: + assert response.status == 200, await response.text() + if request.stream: + async for chunk in response.content: + # Convert raw bytes to CompletionResponse + chunk = chunk.strip() + if not chunk or chunk == b"\n": + continue + # Get rid of the prefix "data: " and suffix "\n" + raw_data = chunk[6:].strip() + if raw_data == b"[DONE]": + continue + data = json.loads(raw_data) + # Commented because we still want usage chunk to be passed back + # if not data["choices"]: + # continue + yield openai_api_protocol.CompletionResponse.model_validate(data) + else: + data = await response.json() + yield openai_api_protocol.CompletionResponse.model_validate(data) diff --git a/python/mlc_llm/serve/__init__.py b/python/mlc_llm/serve/__init__.py new file mode 100644 index 0000000..fc1e81f --- /dev/null +++ b/python/mlc_llm/serve/__init__.py @@ -0,0 +1,11 @@ +"""Subdirectory of serving.""" + +# Load MLC LLM library by importing base +from .. import base +from .config import EngineConfig +from .data import Data, ImageData, RequestStreamOutput, TextData, TokenData +from .embedding_engine import AsyncEmbeddingEngine +from .engine import AsyncMLCEngine, MLCEngine +from .radix_tree import PagedRadixTree +from .request import Request +from .server import PopenServer diff --git a/python/mlc_llm/serve/_ffi_api.py b/python/mlc_llm/serve/_ffi_api.py new file mode 100644 index 0000000..feab878 --- /dev/null +++ b/python/mlc_llm/serve/_ffi_api.py @@ -0,0 +1,7 @@ +"""FFI APIs for mlc_llm.serve""" + +import tvm_ffi + +# Exports functions registered via TVM_FFI_REGISTER_GLOBAL with the "mlc.serve" prefix. +# e.g. TVM_FFI_REGISTER_GLOBAL("mlc.serve.TextData") +tvm_ffi.init_ffi_api("mlc.serve", __name__) diff --git a/python/mlc_llm/serve/config.py b/python/mlc_llm/serve/config.py new file mode 100644 index 0000000..b13ca1c --- /dev/null +++ b/python/mlc_llm/serve/config.py @@ -0,0 +1,169 @@ +"""Configuration dataclasses used in MLC LLM serving""" + +import json +from dataclasses import asdict, dataclass, field +from typing import List, Literal, Optional, Tuple, Union # noqa: UP035 + + +@dataclass +class EngineConfig: + """The class of MLCEngine execution configuration. + + Parameters + ---------- + model : str + The path to the model directory. + + model_lib : str + The path to the model library. + + additional_models : List[Union[str, Tuple[str, str]]] + The paths to the additional models' directories (and model libraries). + Each element is a single string (denoting the model directory) + or a tuple of two strings (denoting the model directory and model lib path). + + mode : Literal["local", "interactive", "server"] + The engine mode in MLC LLM. + We provide three preset modes: "local", "interactive" and "server". + The default mode is "local". + The choice of mode decides the values of "max_num_sequence", "max_total_sequence_length" + and "prefill_chunk_size" when they are not explicitly specified. + 1. Mode "local" refers to the local server deployment which has low + request concurrency. So the max batch size will be set to 4, and max + total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + 2. Mode "interactive" refers to the interactive use of server, which + has at most 1 concurrent request. So the max batch size will be set to 1, + and max total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + 3. Mode "server" refers to the large server use case which may handle + many concurrent request and want to use GPU memory as much as possible. + In this mode, we will automatically infer the largest possible max batch + size and max total sequence length. + + You can manually specify arguments "max_num_sequence", "max_total_sequence_length" and + "prefill_chunk_size" to override the automatic inferred values. + + tensor_parallel_shards : Optional[int] + Number of shards to split the model into in tensor parallelism multi-gpu inference. + When "model_lib" is given, this field will be ignored, and the tensor_parallel_shards + in the model_lib metadata will be used. + + pipeline_parallel_stages : Optional[int] + Number of pipeline stages to split the model layers for pipeline parallelism. + When "model_lib" is given, this field will be ignored, and the pipeline_parallel_stages + in the model_lib metadata will be used. + + opt : Optional[str] + The optimization flags for JIT compilation. + When "model_lib" is given, this field will be ignored. + MLC LLM maintains a predefined set of optimization flags, + denoted as O0, O1, O2, O3, where O0 means no optimization, O2 means majority of them, + and O3 represents extreme optimization that could potentially break the system. + Meanwhile, optimization flags could be explicitly specified via details knobs, e.g. + "cublas_gemm=1;cudagraph=0". + + gpu_memory_utilization : Optional[float] + A number in (0, 1) denoting the fraction of GPU memory used by the server in total. + It is used to infer to maximum possible KV cache capacity. + When it is unspecified, it defaults to 0.85. + Under mode "local" or "interactive", the actual memory usage may be + significantly smaller than this number. Under mode "server", the actual + memory usage may be slightly larger than this number. + + kv_cache_page_size : int + The number of consecutive tokens handled in each page in paged KV cache. + + max_num_sequence : Optional[int] + The maximum number of sequences that are allowed to be + processed by the KV cache at any time. + + max_total_sequence_length : Optional[int] + The maximum total number of tokens whose KV data are allowed + to exist in the KV cache at any time. + + max_single_sequence_length : Optional[int] + The maximum length allowed for a single sequence in the engine. + + prefill_chunk_size : Optional[int] + The maximum total sequence length in a prefill. + + sliding_window_size : Optional[int] + The sliding window size in sliding window attention (SWA). + + attention_sink_size : Optional[int] + The number of attention sinks when sliding window is enabled.. + + max_history_size: Optional[int] + The maximum history size for RNN state to roll back. + + kv_state_kind: Optional[Literal["kv_cache", "rnn_state"]] + The kind of cache. + + speculative_mode : Literal["disable", "small_draft", "eagle", "medusa"] + The speculative mode. + "disable" means speculative decoding is disabled. + "small_draft" means the normal speculative decoding (small draft) mode. + "eagle" means the eagle-style speculative decoding. + "medusa" means the medusa-style speculative decoding. + + spec_draft_length : int + The number of tokens to generate in speculative proposal (draft). + Being 0 means to enable adaptive speculative mode, where the draft length + will be automatically adjusted based on engine state. + + spec_tree_width : int + The width of the speculative decoding tree. + + prefix_cache_mode : Literal["disable", "radix"] + The prefix cache mode. + "disable" means no prefix cache is disabled. + "radix" means the paged radix tree based prefix cache mode. + + prefix_cache_max_num_recycling_seqs: Optional[int] + The maximum number of recycling sequences in prefix cache, default as max_num_sequence. + And set 0 to disable prefix cache, set -1 to have infinite capacity prefix cache. + + prefill_mode : Literal["chunked", "hybrid"] + The prefill mode. + "chunked" means the basic prefill with chunked input enabled. + "hybrid" means the hybrid prefill or split-fuse, + so that decode step will be converted into prefill. + + verbose : bool + A boolean indicating whether to print logging info in engine. + """ + + model: Optional[str] = None + model_lib: Optional[str] = None + additional_models: List[Union[str, Tuple[str, str]]] = field(default_factory=list) # noqa: UP006 + mode: Optional[Literal["local", "interactive", "server"]] = None + tensor_parallel_shards: Optional[int] = None + pipeline_parallel_stages: Optional[int] = None + opt: Optional[str] = None + gpu_memory_utilization: Optional[float] = None + kv_cache_page_size: int = 16 + max_num_sequence: Optional[int] = None + max_total_sequence_length: Optional[int] = None + max_single_sequence_length: Optional[int] = None + prefill_chunk_size: Optional[int] = None + sliding_window_size: Optional[int] = None + attention_sink_size: Optional[int] = None + max_history_size: Optional[int] = None + kv_state_kind: Optional[Literal["kv_cache", "rnn_state"]] = None + speculative_mode: Literal["disable", "small_draft", "eagle", "medusa"] = "disable" + spec_draft_length: int = 0 + spec_tree_width: int = 1 + prefix_cache_mode: Literal["disable", "radix"] = "radix" + prefix_cache_max_num_recycling_seqs: Optional[int] = None + prefill_mode: Literal["chunked", "hybrid"] = "hybrid" + verbose: bool = True + + def asjson(self) -> str: + """Return the config in string of JSON format.""" + return json.dumps(asdict(self)) + + @staticmethod + def from_json(json_str: str) -> "EngineConfig": + """Construct a config from JSON string.""" + return EngineConfig(**json.loads(json_str)) diff --git a/python/mlc_llm/serve/data.py b/python/mlc_llm/serve/data.py new file mode 100644 index 0000000..e1bc4a7 --- /dev/null +++ b/python/mlc_llm/serve/data.py @@ -0,0 +1,214 @@ +"""Classes denoting multi-modality data used in MLC LLM serving""" + +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple # noqa: UP035 + +import tvm +import tvm_ffi +from tvm.runtime import Object, Tensor + +from . import _ffi_api + + +@tvm_ffi.register_object("mlc.serve.Data") +class Data(Object): + """The base class of multi-modality data (text, tokens, embedding, etc).""" + + def __init__(self): + pass + + +@tvm_ffi.register_object("mlc.serve.TextData") +class TextData(Data): + """The class of text data, containing a text string. + + Parameters + ---------- + text : str + The text string. + """ + + def __init__(self, text: str): + self.__init_handle_by_constructor__(_ffi_api.TextData, text) + + @property + def text(self) -> str: + """The text data in `str`.""" + return str(_ffi_api.TextDataGetTextString(self)) + + def __str__(self) -> str: + return self.text + + +@tvm_ffi.register_object("mlc.serve.TokenData") +class TokenData(Data): + """The class of token data, containing a list of token ids. + + Parameters + ---------- + token_ids : List[int] + The list of token ids. + """ + + def __init__(self, token_ids: List[int]): # noqa: UP006 + self.__init_handle_by_constructor__(_ffi_api.TokenData, *token_ids) + + @property + def token_ids(self) -> List[int]: # noqa: UP006 + """Return the token ids of the TokenData.""" + return list(_ffi_api.TokenDataGetTokenIds(self)) + + +# mypy: disable-error-code="attr-defined" +@tvm_ffi.register_object("mlc.serve.ImageData") +class ImageData(Data): + """The class of image data, containing the image as Tensor. + + Parameters + ---------- + image : tvm.runtime.Tensor + The image data. + """ + + def __init__(self, image: Tensor, embed_size: int): + self.embed_size = embed_size + self.__init_handle_by_constructor__(_ffi_api.ImageData, image, embed_size) + + @property + def image(self) -> Tensor: + """Return the image data.""" + return _ffi_api.ImageDataGetImage(self) + + def __len__(self): + return self.embed_size + + @staticmethod + def from_url(url: str, config: Dict) -> "ImageData": # noqa: UP006 + """Get the image from the given URL, process and return the image tensor as TVM Tensor.""" + + import base64 + from io import BytesIO + + import numpy as np + import requests + from PIL import Image + + if url.startswith("data:image"): + # The image is encoded in base64 format + base64_image = url.split(",")[1] + image_data = base64.b64decode(base64_image) + image_tensor = Image.open(BytesIO(image_data)).convert("RGB") + elif url.startswith("http"): + response = requests.get(url, timeout=5) + image_tensor = Image.open(BytesIO(response.content)).convert("RGB") + else: + raise ValueError(f"Unsupported image URL format: {url}") + + # image_embed_size = ImageData.get_embed_size(config) + # TODO: fix these hard-coded values for phi3.5-vision and llava + image_embed_size = 576 + if config["model_type"] == "phi3_v": + image_embed_size = 1921 + image_tensor = np.expand_dims(image_tensor, axis=0) # HWC -> NHWC + image_features = tvm.runtime.tensor(image_tensor) + image_data = ImageData(image_features, image_embed_size) + return image_data + + @staticmethod + def get_embed_size(config: Dict) -> int: # noqa: UP006 + """Get the image embedding size from the model config file.""" + image_size = config["model_config"]["vision_config"]["image_size"] + patch_size = config["model_config"]["vision_config"]["patch_size"] + embed_size = (image_size // patch_size) ** 2 + return embed_size + + @staticmethod + def get_input_size(config: Dict) -> int: # noqa: UP006 + """Get the image input size from the model config file.""" + image_size = config["model_config"]["vision_config"]["image_size"] + return image_size + + +@dataclass +class SingleRequestStreamOutput: + """The request stream output of a single request. + + Attributes + ---------- + delta_token_ids : List[int] + The new generated tokens since the last callback invocation + for the input request. + + delta_logprob_json_strs : Optional[List[str]] + The logprobs JSON strings of the new generated tokens + since last invocation. + + finish_reason : Optional[str] + The finish reason of the request when it is finished, + of None if the request has not finished yet. + """ + + delta_token_ids: List[int] # noqa: UP006 + delta_logprob_json_strs: Optional[List[str]] # noqa: UP006 + finish_reason: Optional[str] + request_final_usage_json_str: Optional[str] + extra_prefix_string: str + + +@tvm_ffi.register_object("mlc.serve.RequestStreamOutput") +class RequestStreamOutput(Object): + """The generated delta request output that is streamed back + through callback stream function. + It contains four fields (in order): + + request_id : str + The id of the request that the function is invoked for. + + stream_outputs : List[SingleRequestStreamOutput] + The output instances, one for a request. + + Note + ---- + We do not provide constructor, since in practice only C++ side + instantiates this class. + """ + + def unpack(self) -> Tuple[str, List[SingleRequestStreamOutput]]: # noqa: UP006 + """Return the fields of the delta output in a tuple. + + Returns + ------- + request_id : str + The id of the request that the function is invoked for. + + stream_outputs : List[SingleRequestStreamOutput] + The output instances, one for a request. + """ + fields = _ffi_api.RequestStreamOutputUnpack(self) + request_final_usage_json_str = fields[4] + request_id = str(fields[0]) + if request_final_usage_json_str is not None: + return ( + request_id, + [SingleRequestStreamOutput([], None, None, request_final_usage_json_str, "")], + ) + + stream_outputs = [] + for i, (delta_token_ids, finish_reason, extra_prefix_string) in enumerate( + zip(fields[1], fields[3], fields[5]) + ): + delta_logprob_json_strs = ( + [str(logprob_json_str) for logprob_json_str in fields[2][i]] + if fields[2] is not None + else None + ) + stream_outputs.append( + SingleRequestStreamOutput( + delta_token_ids=list(delta_token_ids), + delta_logprob_json_strs=delta_logprob_json_strs, + finish_reason=str(finish_reason) if finish_reason is not None else None, + request_final_usage_json_str=None, + extra_prefix_string=str(extra_prefix_string), + ) + ) + return request_id, stream_outputs diff --git a/python/mlc_llm/serve/embedding_engine.py b/python/mlc_llm/serve/embedding_engine.py new file mode 100644 index 0000000..782c1eb --- /dev/null +++ b/python/mlc_llm/serve/embedding_engine.py @@ -0,0 +1,490 @@ +"""Asynchronous embedding inference engine for encoder and decoder models.""" + +import asyncio +import concurrent.futures +import json +import os +from typing import List, Literal, Optional, Tuple, Union # noqa: UP035 + +import numpy as np +import tvm +from tvm import relax +from tvm.runtime import Device +from tvm_ffi import Shape + +from mlc_llm.serve import engine_utils +from mlc_llm.support.auto_device import detect_device +from mlc_llm.tokenizers import Tokenizer + + +class AsyncEmbeddingEngine: + """Asynchronous embedding inference engine. + + Supports both encoder models (BERT-style) and decoder-only embedding models + (e.g. Qwen3-Embeddings). Uses a ThreadPoolExecutor for background inference + so that the asyncio event loop is not blocked. + + Parameters + ---------- + model : str + Path to the model weight directory. + + model_lib : str + Path to the compiled model library (.so/.dylib file). + + device : Union[str, Device] + Device string, e.g. "auto", "cuda:0", "metal". + + pooling_strategy : Optional[str] + Pooling strategy: "cls" (first token), "mean" (masked average), + or "last" (last token). If None, auto-detected based on model type: + encoder -> "cls", decoder -> "last". + """ + + def __init__( + self, + model: str, + model_lib: str, + device: Union[str, Device] = "auto", + *, + pooling_strategy: Optional[str] = None, + ) -> None: + # Reuse existing utility: device detection + self.device = detect_device(device) if isinstance(device, str) else device + # Reuse existing utility: tokenizer + self.tokenizer = Tokenizer(model) + + # Load TVM module, metadata, and params via engine_utils helpers + ex = tvm.runtime.load_module(model_lib) + vm = relax.VirtualMachine(ex, device=self.device) + self._mod = vm.module + self._metadata = json.loads(self._mod["_metadata"]()) + self._params = engine_utils.load_embedding_params(model, self.device, self._metadata) + + # Detect model type and set pooling strategy + self.embedding_metadata = engine_utils.get_embedding_metadata(self._metadata) + if self.embedding_metadata: + self.model_type = self.embedding_metadata["model_type"] + self.pooling_strategy = self.embedding_metadata["pooling_strategy"] + self.normalize = self.embedding_metadata["normalize"] + else: + self.model_type = engine_utils.detect_embedding_model_type(self._mod) + self.pooling_strategy = "cls" if self.model_type == "encoder" else "last" + self.normalize = True + # Allow caller to override pooling strategy + if pooling_strategy: + self.pooling_strategy = pooling_strategy + + # Initialize model-type-specific functions + if self.model_type == "encoder": + self._init_encoder(model) + else: + self._init_decoder(model) + + # Background thread pool (1 worker = serialized GPU inference) + self._executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="embedding" + ) + self._terminated = False + + def _init_encoder(self, model: str) -> None: + """Initialize encoder (BERT-style) model functions and special tokens.""" + self._prefill_func = self._mod["prefill"] + self._cls_token_id: Optional[int] = None + self._sep_token_id: Optional[int] = None + tok_config_path = os.path.join(model, "tokenizer_config.json") + if os.path.exists(tok_config_path): + with open(tok_config_path, encoding="utf-8") as f: + tok_config = json.load(f) + # Try added_tokens_decoder first (newer HF format) + added = tok_config.get("added_tokens_decoder", {}) + for tid, info in added.items(): + if info.get("content") == tok_config.get("cls_token"): + self._cls_token_id = int(tid) + if info.get("content") == tok_config.get("sep_token"): + self._sep_token_id = int(tid) + # Fallback: encode the special token strings via tokenizer + if self._cls_token_id is None and tok_config.get("cls_token"): + ids = list(self.tokenizer.encode(tok_config["cls_token"])) + if len(ids) == 1: + self._cls_token_id = ids[0] + if self._sep_token_id is None and tok_config.get("sep_token"): + ids = list(self.tokenizer.encode(tok_config["sep_token"])) + if len(ids) == 1: + self._sep_token_id = ids[0] + + def _init_decoder(self, model: str) -> None: + """Initialize decoder (Qwen3-Embeddings style) model functions.""" + # Prefer tokenizer post-processing (HF-style) for terminal/pooling token handling. + # Only fall back to manual EOS append when tokenizer does not define a post-processor + # that actually appends a token at the end of the sequence. + self._decoder_tokenizer_appends_eos = False + tokenizer_json_path = os.path.join(model, "tokenizer.json") + if os.path.exists(tokenizer_json_path): + with open(tokenizer_json_path, encoding="utf-8") as f: + tokenizer_json = json.load(f) + post_proc = tokenizer_json.get("post_processor") + if post_proc is not None: + # Check if the post-processor actually appends a special token at the end + # (e.g. TemplateProcessing with "$A <|endoftext|>"). We verify by encoding + # a test string and checking if the last token is a known special token. + test_tokens = list(self.tokenizer.encode("test")) + if len(test_tokens) > 0: + vocab = tokenizer_json.get("added_tokens", []) + special_ids = {t["id"] for t in vocab if t.get("special", False)} + if test_tokens[-1] in special_ids: + self._decoder_tokenizer_appends_eos = True + + # Read EOS token from config — fallback only when tokenizer does not auto-append. + self._decoder_eos_token_id: Optional[int] = None + config_path = os.path.join(model, "mlc-chat-config.json") + if os.path.exists(config_path): + with open(config_path, encoding="utf-8") as f: + chat_config = json.load(f) + eos = chat_config.get("eos_token_id") + if isinstance(eos, list): + self._decoder_eos_token_id = eos[0] + elif isinstance(eos, int): + self._decoder_eos_token_id = eos + + self._embed_func = self._mod["embed"] + self._prefill_to_hidden_func = self._mod["prefill_to_last_hidden_states"] + self._batch_prefill_to_hidden_func = self._mod["batch_prefill_to_last_hidden_states"] + if self._mod.implements_function("create_tir_paged_kv_cache"): + self._create_kv_cache_func = self._mod["create_tir_paged_kv_cache"] + elif self._mod.implements_function("create_flashinfer_paged_kv_cache"): + self._create_kv_cache_func = self._mod["create_flashinfer_paged_kv_cache"] + else: + raise RuntimeError("Cannot find KV cache creation function in model library.") + self._kv_state_add_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence") + self._kv_state_remove_sequence = tvm.get_global_func("vm.builtin.kv_state_remove_sequence") + self._kv_state_begin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward") + self._kv_state_end_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward") + self._nd_reshape = tvm.get_global_func("vm.builtin.reshape") + + def embed(self, inputs: List[str]) -> Tuple[List[List[float]], int]: # noqa: UP006 + """Compute embeddings for a list of input strings (synchronous). + + Parameters + ---------- + inputs : List[str] + The input strings to embed. + + Returns + ------- + embeddings : List[List[float]] + The L2-normalized embedding vectors. + total_tokens : int + Total number of tokens processed. + """ + if self.model_type == "encoder": + return self._embed_encoder(inputs) + return self._embed_decoder(inputs) + + async def async_embed(self, inputs: List[str]) -> Tuple[List[List[float]], int]: # noqa: UP006 + """Compute embeddings asynchronously in a background thread. + + This method does not block the asyncio event loop. + + Parameters + ---------- + inputs : List[str] + The input strings to embed. + + Returns + ------- + embeddings : List[List[float]] + The L2-normalized embedding vectors. + total_tokens : int + Total number of tokens processed. + """ + loop = asyncio.get_running_loop() + return await loop.run_in_executor(self._executor, self.embed, inputs) + + def _embed_encoder( + self, + inputs: List[str], # noqa: UP006 + ) -> Tuple[List[List[float]], int]: # noqa: UP006 + """Encoder model embedding (BERT-style). + + Processes each input individually to avoid batch padding artifacts. + + Encoder uses bidirectional attention, so chunked prefill is NOT possible + (each token must attend to all other tokens in the full sequence). + Inputs exceeding prefill_chunk_size are truncated. + + (Additional Strategy) + TODO: For better long-text support, implement sliding window + mean pooling: + 1. Split text into overlapping windows of prefill_chunk_size (stride=chunk/2) + 2. Encode each window independently + 3. Mean-pool all window embeddings → final embedding → L2 normalize + This preserves information from the full text at the cost of N× compute. + """ # noqa: RUF002 + embeddings: List[List[float]] = [] # noqa: UP006 + total_tokens = 0 + prefill_chunk = self._metadata.get("prefill_chunk_size", 512) + + for text in inputs: + tokens = list(self.tokenizer.encode(text)) + # Add [CLS] and [SEP] if needed + if self._cls_token_id is not None and ( + len(tokens) == 0 or tokens[0] != self._cls_token_id + ): + tokens = [self._cls_token_id, *tokens] + if self._sep_token_id is not None and ( + len(tokens) == 0 or tokens[-1] != self._sep_token_id + ): + tokens = [*tokens, self._sep_token_id] + + # Truncate to compiled buffer limit (keep [CLS] at start, [SEP] at end) + if len(tokens) > prefill_chunk: + tokens = tokens[:prefill_chunk] + if self._sep_token_id is not None: + tokens[-1] = self._sep_token_id + + seq_len = len(tokens) + total_tokens += seq_len + + token_ids = np.array([tokens], dtype=np.int32) # [1, seq_len] + attention_mask: np.ndarray = np.ones((1, seq_len), dtype=np.int32) # [1, seq_len] + + tokens_tvm = tvm.runtime.tensor(token_ids, device=self.device) + mask_tvm = tvm.runtime.tensor(attention_mask, device=self.device) + + output = self._prefill_func(tokens_tvm, mask_tvm, self._params) + # .numpy() copies to CPU, escaping TVM workspace buffer reuse across calls. + output_np = output.numpy() # [1, seq_len, hidden_size] + + # Pooling + if self.pooling_strategy == "cls": + pooled = output_np[0, 0, :] + elif self.pooling_strategy == "mean": + pooled = output_np[0].mean(axis=0) + else: # "last" + pooled = output_np[0, -1, :] + + # L2 normalize + pooled = pooled.astype(np.float32) + if self.normalize: + norm = np.linalg.norm(pooled) + if norm > 1e-12: + pooled = pooled / norm + + embeddings.append(pooled.tolist()) + + return embeddings, total_tokens + + def _embed_decoder(self, inputs: List[str]) -> Tuple[List[List[float]], int]: # noqa: UP006 + """Decoder model embedding with batch prefill optimization. + + When total tokens fit within prefill_chunk_size, all inputs are processed + in a single batch forward pass using shared KV cache. Otherwise, falls back + to sequential chunked prefill per input. + """ + # Read KV cache config from metadata + prefill_chunk = self._metadata.get("prefill_chunk_size", 2048) + max_seq_len = self._metadata.get("context_window_size", 32768) + if max_seq_len == -1: + max_seq_len = self._metadata.get("sliding_window_size", -1) + assert max_seq_len > 0, f"max_seq_len must be positive, got {max_seq_len}" + support_sliding = int(self._metadata.get("sliding_window_size", -1) != -1) + + # Tokenize all inputs. Prefer tokenizer post-processor output. If absent (older models), + # fall back to appending eos_token_id when missing. + token_lists: List[List[int]] = [] # noqa: UP006 + for text in inputs: + tokens = list(self.tokenizer.encode(text)) + if ( + not self._decoder_tokenizer_appends_eos + and self._decoder_eos_token_id is not None + and (len(tokens) == 0 or tokens[-1] != self._decoder_eos_token_id) + ): + tokens.append(self._decoder_eos_token_id) + if len(tokens) > max_seq_len: + tokens = tokens[:max_seq_len] + token_lists.append(tokens) + + total_tokens = sum(len(t) for t in token_lists) + + # Fast path: all tokens fit in one prefill chunk → batch forward + if total_tokens <= prefill_chunk and all(len(t) > 0 for t in token_lists): + return self._batch_embed_decoder( + token_lists, total_tokens, max_seq_len, prefill_chunk, support_sliding + ) + + # Greedy sub-batching: pack texts into sub-batches that fit within + # prefill_chunk, preserving input order. Oversize texts (single text + # exceeding prefill_chunk) fall back to sequential chunked prefill. + sub_batches = self._build_sub_batches(token_lists, prefill_chunk) + all_embeddings: List[List[float]] = [] # noqa: UP006 + for batch_type, batch, batch_total in sub_batches: + if batch_type == "batch": + embs, _ = self._batch_embed_decoder( + batch, batch_total, max_seq_len, prefill_chunk, support_sliding + ) + else: + embs, _ = self._sequential_embed_decoder( + batch, batch_total, max_seq_len, prefill_chunk, support_sliding + ) + all_embeddings.extend(embs) + + return all_embeddings, total_tokens + + @staticmethod + def _build_sub_batches( + token_lists: List[List[int]], # noqa: UP006 + prefill_chunk: int, + ) -> List[Tuple[Literal["batch", "sequential"], List[List[int]], int]]: # noqa: UP006 + """Partition token lists into sub-batches that fit within prefill_chunk. + + Each sub-batch is a tuple of (mode, token_lists, total_token_count). + Empty token lists are skipped to avoid invalid batch processing. + """ + sub_batches: List[Tuple[Literal["batch", "sequential"], List[List[int]], int]] = [] # noqa: UP006 + current_batch: List[List[int]] = [] # noqa: UP006 + current_tokens = 0 + + for tokens in token_lists: + if not tokens: + continue + token_len = len(tokens) + is_oversized = token_len > prefill_chunk + if current_batch and (is_oversized or current_tokens + token_len > prefill_chunk): + sub_batches.append(("batch", current_batch, current_tokens)) + current_batch, current_tokens = [], 0 + if is_oversized: + sub_batches.append(("sequential", [tokens], token_len)) + else: + current_batch.append(tokens) + current_tokens += token_len + if current_batch: + sub_batches.append(("batch", current_batch, current_tokens)) + + return sub_batches + + def _batch_embed_decoder( + self, + token_lists: List[List[int]], # noqa: UP006 + total_tokens: int, + max_seq_len: int, + prefill_chunk: int, + support_sliding: int, + ) -> Tuple[List[List[float]], int]: # noqa: UP006 + """Batch prefill: process all inputs in a single forward pass.""" + batch_size = len(token_lists) + + # Create KV cache for the entire batch + kv_cache = self._create_kv_cache_func( + Shape([batch_size]), + Shape([max_seq_len]), + Shape([prefill_chunk]), + Shape([16]), + Shape([support_sliding]), + ) + + # Register all sequences + seq_ids = list(range(batch_size)) + seq_lens = [len(t) for t in token_lists] + for sid in seq_ids: + self._kv_state_add_sequence(kv_cache, sid) + + # Begin forward with all sequences at once + self._kv_state_begin_forward(kv_cache, Shape(seq_ids), Shape(seq_lens)) + + # Concatenate all tokens → embed → batch prefill + all_tokens = [] + for tokens in token_lists: + all_tokens.extend(tokens) + token_ids = tvm.runtime.tensor(np.array(all_tokens, dtype=np.int32), device=self.device) + all_embed = self._embed_func(token_ids, self._params) + all_embed = self._nd_reshape(all_embed, Shape([1, total_tokens, all_embed.shape[-1]])) + + hidden_states, _ = self._batch_prefill_to_hidden_func(all_embed, kv_cache, self._params) + # .numpy() copies to CPU, escaping TVM workspace buffer reuse across calls. + # (torch.from_dlpack is zero-copy and hits aliasing bugs on 2nd+ invocation.) + hidden_np = hidden_states.numpy() + self._kv_state_end_forward(kv_cache) + for sid in seq_ids: + self._kv_state_remove_sequence(kv_cache, sid) + + # Extract last token hidden state per sequence + embeddings: List[List[float]] = [] # noqa: UP006 + offset = 0 + for tokens in token_lists: + last_pos = offset + len(tokens) - 1 + pooled = hidden_np[0, last_pos, :].astype(np.float32) + if self.normalize: + norm = np.linalg.norm(pooled) + if norm > 1e-12: + pooled = pooled / norm + embeddings.append(pooled.tolist()) + offset += len(tokens) + + return embeddings, total_tokens + + def _sequential_embed_decoder( + self, + token_lists: List[List[int]], # noqa: UP006 + total_tokens: int, + max_seq_len: int, + prefill_chunk: int, + support_sliding: int, + ) -> Tuple[List[List[float]], int]: # noqa: UP006 + """Sequential chunked prefill: process each input independently.""" + embeddings: List[List[float]] = [] # noqa: UP006 + + for tokens in token_lists: + if len(tokens) == 0: + continue + + # Create KV cache for this single sequence + kv_cache = self._create_kv_cache_func( + Shape([1]), + Shape([max_seq_len]), + Shape([prefill_chunk]), + Shape([16]), + Shape([support_sliding]), + ) + self._kv_state_add_sequence(kv_cache, 0) + + # Process tokens in chunks + hidden = None + for chunk_start in range(0, len(tokens), prefill_chunk): + chunk_end = min(chunk_start + prefill_chunk, len(tokens)) + chunk_tokens = tokens[chunk_start:chunk_end] + chunk_len = len(chunk_tokens) + + token_ids = tvm.runtime.tensor( + np.array(chunk_tokens, dtype=np.int32), device=self.device + ) + chunk_embed = self._embed_func(token_ids, self._params) + chunk_embed = self._nd_reshape( + chunk_embed, Shape([1, chunk_len, chunk_embed.shape[-1]]) + ) + self._kv_state_begin_forward(kv_cache, Shape([0]), Shape([chunk_len])) + hidden, kv_cache = self._prefill_to_hidden_func(chunk_embed, kv_cache, self._params) + # .numpy() copies to CPU, escaping TVM buffer aliasing. + hidden_np = hidden.numpy() + self._kv_state_end_forward(kv_cache) + + self._kv_state_remove_sequence(kv_cache, 0) + + pooled = hidden_np[0, -1, :] if hidden_np.ndim == 3 else hidden_np[-1, :] + pooled = pooled.astype(np.float32) + if self.normalize: + norm = np.linalg.norm(pooled) + if norm > 1e-12: + pooled = pooled / norm + embeddings.append(pooled.tolist()) + + return embeddings, total_tokens + + def terminate(self) -> None: + """Terminate the engine and clean up the thread pool.""" + if getattr(self, "_terminated", True): + return + self._terminated = True + self._executor.shutdown(wait=False) + + def __del__(self): + self.terminate() diff --git a/python/mlc_llm/serve/engine.py b/python/mlc_llm/serve/engine.py new file mode 100644 index 0000000..5e17925 --- /dev/null +++ b/python/mlc_llm/serve/engine.py @@ -0,0 +1,1946 @@ +"""The MLC LLM Serving Engine.""" + +import asyncio +import queue +import weakref +from collections.abc import AsyncGenerator, Iterator +from typing import ( # noqa: UP035 + Any, + Dict, + List, + Literal, + Optional, + Tuple, + Union, + overload, +) + +from tvm.runtime import Device + +from mlc_llm.protocol import debug_protocol, openai_api_protocol +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import data, engine_utils +from mlc_llm.serve.config import EngineConfig +from mlc_llm.support import logging +from mlc_llm.tokenizers import TextStreamer + +from . import engine_base + +logger = logging.getLogger(__name__) + + +# Note: we define both AsyncChat and Chat for Python type analysis. +class AsyncChat: + """The proxy class to direct to async chat completions.""" + + def __init__(self, engine: weakref.ReferenceType) -> None: + assert isinstance(engine(), AsyncMLCEngine) + self.completions = AsyncChatCompletion(engine) + + +class Chat: + """The proxy class to direct to chat completions.""" + + def __init__(self, engine: weakref.ReferenceType) -> None: + assert isinstance(engine(), MLCEngine) + self.completions = ChatCompletion(engine) + + +class AsyncChatCompletion: + """The proxy class to direct to async chat completions.""" + + engine: weakref.ReferenceType["AsyncMLCEngine"] + + def __init__(self, engine: weakref.ReferenceType) -> None: + self.engine = engine + + @overload + async def create( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + stream: Literal[True], + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any]: + """Asynchronous streaming chat completion interface with OpenAI API compatibility. + The method is a coroutine that streams ChatCompletionStreamResponse + that conforms to OpenAI API one at a time via yield. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Yields + ------ + stream_response : ChatCompletionStreamResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/chat/streaming for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + @overload + async def create( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: Literal[False] = False, + stream_options: Literal[None] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> openai_api_protocol.ChatCompletionResponse: + """Asynchronous non-streaming chat completion interface with OpenAI API compatibility. + The method is a coroutine that streams ChatCompletionStreamResponse + that conforms to OpenAI API one at a time via yield. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Returns + ------- + response : ChatCompletionResponse + The chat completion response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/chat/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + async def create( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any], + openai_api_protocol.ChatCompletionResponse, + ]: + """Asynchronous chat completion interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + return await self.engine()._chat_completion( + messages=messages, + model=model, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + top_logprobs=top_logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + temperature=temperature, + top_p=top_p, + tools=tools, + tool_choice=tool_choice, + user=user, + response_format=response_format, + request_id=request_id, + debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), + ) + + +class ChatCompletion: + """The proxy class to direct to chat completions.""" + + engine: weakref.ReferenceType["MLCEngine"] + + def __init__(self, engine: weakref.ReferenceType) -> None: + self.engine = engine + + @overload + def create( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + stream: Literal[True], + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]: + """Synchronous streaming chat completion interface with OpenAI API compatibility. + The method streams back ChatCompletionStreamResponse that conforms to + OpenAI API one at a time via yield. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Yields + ------ + stream_response : ChatCompletionStreamResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/chat/streaming for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + @overload + def create( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: Literal[False] = False, + stream_options: Literal[None] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> openai_api_protocol.ChatCompletionResponse: + """Synchronous non-streaming chat completion interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Returns + ------ + response : ChatCompletionResponse + The chat completion response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/chat/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + def create( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + Iterator[openai_api_protocol.ChatCompletionStreamResponse], + openai_api_protocol.ChatCompletionResponse, + ]: + """Synchronous chat completion interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + return self.engine()._chat_completion( + messages=messages, + model=model, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + top_logprobs=top_logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + temperature=temperature, + top_p=top_p, + tools=tools, + tool_choice=tool_choice, + user=user, + response_format=response_format, + request_id=request_id, + debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), + ) + + +class AsyncCompletion: + """The proxy class to direct to async completions.""" + + engine: weakref.ReferenceType["AsyncMLCEngine"] + + def __init__(self, engine: weakref.ReferenceType) -> None: + self.engine = engine + + @overload + async def create( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + stream: Literal[True], + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + """Asynchronous streaming completion interface with OpenAI API compatibility. + The method is a coroutine that streams CompletionResponse + that conforms to OpenAI API one at a time via yield. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Yields + ------ + stream_response : CompletionResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/completions/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + @overload + async def create( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: Literal[False] = False, + stream_options: Literal[None] = None, + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> openai_api_protocol.CompletionResponse: + """Asynchronous non-streaming completion interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Returns + ------ + response : CompletionResponse + The completion response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/completions/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + async def create( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + AsyncGenerator[openai_api_protocol.CompletionResponse, Any], + openai_api_protocol.CompletionResponse, + ]: + """Asynchronous completion interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + return await self.engine()._completion( + model=model, + prompt=prompt, + best_of=best_of, + echo=echo, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + suffix=suffix, + temperature=temperature, + top_p=top_p, + user=user, + response_format=response_format, + request_id=request_id, + debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), + ) + + +class Completion: + """The proxy class to direct to completions.""" + + engine: weakref.ReferenceType["MLCEngine"] + + def __init__(self, engine: weakref.ReferenceType) -> None: + self.engine = engine + + @overload + def create( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + stream: Literal[True], + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Iterator[openai_api_protocol.CompletionResponse]: + """Synchronous streaming completion interface with OpenAI API compatibility. + The method streams back CompletionResponse that conforms to + OpenAI API one at a time via yield. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Yields + ------ + stream_response : CompletionResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/completions/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + @overload + def create( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: Literal[False] = False, + stream_options: Literal[None] = None, + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> openai_api_protocol.CompletionResponse: + """Synchronous non-streaming completion interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Returns + ------- + response : CompletionResponse + The completion response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/completions/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + + def create( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + extra_body: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + Iterator[openai_api_protocol.CompletionResponse], + openai_api_protocol.CompletionResponse, + ]: + """Synchronous completion interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + extra_body: Optional[Dict[str, Any]] = None, + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + return self.engine()._completion( + model=model, + prompt=prompt, + best_of=best_of, + echo=echo, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + suffix=suffix, + temperature=temperature, + top_p=top_p, + user=user, + response_format=response_format, + request_id=request_id, + debug_config=(extra_body.get("debug_config", None) if extra_body is not None else None), + ) + + +class AsyncMLCEngine(engine_base.MLCEngineBase): + """The AsyncMLCEngine in MLC LLM that provides the asynchronous + interfaces with regard to OpenAI API. + + Parameters + ---------- + model : str + A path to ``mlc-chat-config.json``, or an MLC model directory that contains + `mlc-chat-config.json`. + It can also be a link to a HF repository pointing to an MLC compiled model. + + device: Union[str, Device] + The device used to deploy the model such as "cuda" or "cuda:0". + Will default to "auto" and detect from local available GPUs if not specified. + + model_lib : Optional[str] + The full path to the model library file to use (e.g. a ``.so`` file). + If unspecified, we will use the provided ``model`` to search over possible paths. + It the model lib is not found, it will be compiled in a JIT manner. + + mode : Literal["local", "interactive", "server"] + The engine mode in MLC LLM. + We provide three preset modes: "local", "interactive" and "server". + The default mode is "local". + The choice of mode decides the values of "max_num_sequence", "max_total_sequence_length" + and "prefill_chunk_size" when they are not explicitly specified. + 1. Mode "local" refers to the local server deployment which has low + request concurrency. So the max batch size will be set to 4, and max + total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + 2. Mode "interactive" refers to the interactive use of server, which + has at most 1 concurrent request. So the max batch size will be set to 1, + and max total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + 3. Mode "server" refers to the large server use case which may handle + many concurrent request and want to use GPU memory as much as possible. + In this mode, we will automatically infer the largest possible max batch + size and max total sequence length. + + You can manually specify arguments "max_num_sequence", "max_total_sequence_length" and + "prefill_chunk_size" to override the automatic inferred values. + + engine_config : Optional[EngineConfig] + Additional configurable arguments of MLC engine. + See class "EngineConfig" for more detail. + + enable_tracing : bool + A boolean indicating if to enable event logging for requests. + """ + + def __init__( + self, + model: str, + device: Union[str, Device] = "auto", + *, + model_lib: Optional[str] = None, + mode: Literal["local", "interactive", "server"] = "local", + engine_config: Optional[EngineConfig] = None, + enable_tracing: bool = False, + ) -> None: + super().__init__( + "async", + model=model, + device=device, + model_lib=model_lib, + mode=mode, + engine_config=engine_config, + enable_tracing=enable_tracing, + ) + self.chat = AsyncChat(weakref.ref(self)) + self.completions = AsyncCompletion(weakref.ref(self)) + + async def abort(self, request_id: str) -> None: + """Generation abortion interface. + + Parameters + --------- + request_id : str + The id of the request to abort. + """ + self._abort(request_id) + + async def metrics(self) -> engine_base.EngineMetrics: + """Get engine metrics + + Returns + ------- + metrics: EngineMetrics + The engine metrics + """ + return await engine_base._async_query_engine_metrics(self) + + async def _chat_completion( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any], + openai_api_protocol.ChatCompletionResponse, + ]: + """Asynchronous chat completion internal interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + Extra body options to pass to the request. + Can be used to pass debug config as extra_body["debug_config"] + + debug_config: Optional[Dict[str, Any]] = None, + Debug config body options to pass to the request. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + if request_id is None: + request_id = f"chatcmpl-{engine_utils.random_uuid()}" + + chatcmpl_generator = self._handle_chat_completion( + openai_api_protocol.ChatCompletionRequest( + messages=[ + openai_api_protocol.ChatCompletionMessage.model_validate(message) + for message in messages + ], + model=model, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + top_logprobs=top_logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + temperature=temperature, + top_p=top_p, + tools=( + [openai_api_protocol.ChatTool.model_validate(tool) for tool in tools] + if tools is not None + else None + ), + tool_choice=tool_choice, + user=user, + response_format=( + openai_api_protocol.RequestResponseFormat.model_validate(response_format) + if response_format is not None + else None + ), + debug_config=( + debug_protocol.DebugConfig.model_validate(debug_config) + if debug_config is not None + else None + ), + ), + request_id=request_id, + request_final_usage_include_extra=True, + ) + if stream: + # Stream response. + return chatcmpl_generator + # Normal response. + output_texts = ["" for _ in range(n)] + finish_reasons: List[Optional[str]] = [None for _ in range(n)] # noqa: UP006 + logprob_results: Optional[List[List[openai_api_protocol.LogProbsContent]]] = ( # noqa: UP006 + [[] for _ in range(n)] if logprobs else None + ) + request_final_usage = None + try: + async for response in chatcmpl_generator: + # when usage is not None this is the last chunk + if response.usage is not None: + request_final_usage = response.usage + continue + for choice in response.choices: + assert isinstance(choice.delta.content, str) + output_texts[choice.index] += choice.delta.content + if choice.finish_reason is not None and finish_reasons[choice.index] is None: + finish_reasons[choice.index] = choice.finish_reason + if choice.logprobs is not None: + assert logprob_results is not None + logprob_results[choice.index] += choice.logprobs.content + except asyncio.CancelledError: + # for cancelled error, we can simply pass it through + raise + except Exception as err: + logger.error("Error in chat completion with request ID %s: %s", request_id, err) + raise + + assert all(finish_reason is not None for finish_reason in finish_reasons) + use_function_calling, tool_calls_list = engine_base.process_function_call_output( + output_texts, finish_reasons + ) + return engine_base.wrap_chat_completion_response( + request_id=request_id, + model=model, + output_texts=output_texts, + finish_reasons=finish_reasons, + tool_calls_list=tool_calls_list, + logprob_results=logprob_results, + use_function_calling=use_function_calling, + usage=request_final_usage, + ) + + async def _completion( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + AsyncGenerator[openai_api_protocol.CompletionResponse, Any], + openai_api_protocol.CompletionResponse, + ]: + """Asynchronous completion internal interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + debug_config: Optional[Dict[str, Any]] = None, + Extra debug options to pass to the request. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + if request_id is None: + request_id = f"cmpl-{engine_utils.random_uuid()}" + cmpl_generator = self._handle_completion( + openai_api_protocol.CompletionRequest( + model=model, + prompt=prompt, + best_of=best_of, + echo=echo, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + suffix=suffix, + temperature=temperature, + top_p=top_p, + user=user, + response_format=( + openai_api_protocol.RequestResponseFormat.model_validate(response_format) + if response_format is not None + else None + ), + debug_config=( + debug_protocol.DebugConfig.model_validate(debug_config) + if debug_config is not None + else None + ), + ), + request_id=request_id, + request_final_usage_include_extra=True, + ) + if stream: + # Stream response. + return cmpl_generator + # Normal response. + request_final_usage = None + output_texts = [""] * n + finish_reasons: List[Optional[str]] = [None] * n # noqa: UP006 + logprob_results: List[Optional[openai_api_protocol.CompletionLogProbs]] = [None] * n # noqa: UP006 + + async for response in cmpl_generator: + # this is the final chunk + if response.usage is not None: + request_final_usage = response.usage + continue + for choice in response.choices: + output_texts[choice.index] += choice.text + if choice.finish_reason is not None and finish_reasons[choice.index] is None: + finish_reasons[choice.index] = choice.finish_reason + if choice.logprobs is not None: + logprob_results[choice.index] = choice.logprobs + + assert all(finish_reason is not None for finish_reason in finish_reasons) + + return engine_base.wrap_completion_response( + request_id=request_id, + model=model, + output_texts=output_texts, + finish_reasons=finish_reasons, + logprob_results=logprob_results, + usage=request_final_usage, + ) + + async def _handle_chat_completion( + self, + request: openai_api_protocol.ChatCompletionRequest, + request_id: str, + request_final_usage_include_extra: bool, + ) -> AsyncGenerator[openai_api_protocol.ChatCompletionStreamResponse, Any]: + """The implementation fo asynchronous ChatCompletionRequest handling. + + Yields + ------ + stream_response : ChatCompletionStreamResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/chat/streaming for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + ( + prompts, + generation_cfg, + use_function_calling, + prompt_length, + ) = engine_base.process_chat_completion_request( + request, + request_id, + self.state, + self.model_config_dicts[0], + self.tokenizer.encode, + self.max_input_sequence_length, + self.conv_template.model_copy(deep=True), + ) + # prompt length is not used + _ = prompt_length + finish_reasons: List[Optional[str]] = [None for _ in range(generation_cfg.n)] # noqa: UP006 + self.state.record_event(request_id, event="invoke generate") + try: + async for delta_outputs in self._generate( + prompts, + generation_cfg, + request_id, + ): + response = engine_base.process_chat_completion_stream_output( + delta_outputs, + request, + request_id, + self.state, + use_function_calling, + finish_reasons, + ) + + if response is not None: + if response.usage is not None: + if not request_final_usage_include_extra: + response.usage.extra = None + yield response + self.state.record_event(request_id, event="finish") + except asyncio.CancelledError: + # for cancelled error, we can simply pass it through + raise + except Exception as err: + logger.error("Error in _handle_chat_completion for request %s: %s", request_id, err) + raise + + async def _handle_completion( + self, + request: openai_api_protocol.CompletionRequest, + request_id: str, + request_final_usage_include_extra: bool, + ) -> AsyncGenerator[openai_api_protocol.CompletionResponse, Any]: + """The implementation fo asynchronous CompletionRequest handling. + + Yields + ------ + stream_response : CompletionResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/completions/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + ( + prompt, + generation_cfg, + prompt_length, + echo_response, + ) = engine_base.process_completion_request( + request, + request_id, + self.state, + self.tokenizer, + self.max_input_sequence_length, + self.conv_template.model_copy(deep=True), + ) + _ = prompt_length + if echo_response is not None: + yield echo_response + + finish_reasons: List[Optional[str]] = [None] * generation_cfg.n # noqa: UP006 + self.state.record_event(request_id, event="invoke generate") + try: + async for delta_outputs in self._generate( + prompt, + generation_cfg, + request_id, + ): + response = engine_base.process_completion_stream_output( + delta_outputs, + request, + request_id, + self.state, + finish_reasons, + ) + + if response is not None: + if response.usage is not None: + if not request_final_usage_include_extra: + response.usage.extra = None + yield response + + suffix_response = engine_base.create_completion_suffix_response( + request, request_id, finish_reasons + ) + if suffix_response is not None: + yield suffix_response + self.state.record_event(request_id, event="finish") + except asyncio.CancelledError: + # for cancelled error, we can simply pass it through + raise + except Exception as err: + logger.error("Error in _handle_completion for request %s: %s", request_id, err) + raise + + async def _generate( + self, + prompt: Union[str, List[int], List[Union[str, List[int], data.Data]]], # noqa: UP006 + generation_config: GenerationConfig, + request_id: str, + ) -> AsyncGenerator[List[engine_base.CallbackStreamOutput], Any]: # noqa: UP006 + """Internal asynchronous text generation interface of AsyncMLCEngine. + The method is a coroutine that streams a list of CallbackStreamOutput + at a time via yield. The returned list length is the number of + parallel generations specified by `generation_config.n`. + + Parameters + ---------- + prompt : Union[str, List[int], List[Union[str, List[int], data.Data]]] + The input prompt in forms of text strings, lists of token ids or data. + + generation_config : GenerationConfig + The generation config of the request. + + request_id : str + The unique identifier (in string) or this generation request. + + Yields + ------ + request_output : List[engine_base.CallbackStreamOutput] + The delta generated outputs in a list. + The number of list elements equals to `generation_config.n`, + and each element corresponds to the delta output of a parallel + generation. + """ + if self._terminated: + raise ValueError("The AsyncThreadedEngine has terminated.") + self.state.async_lazy_init_event_loop() + + # Create the request with the given id, input data, generation + # config and the created callback. + input_data = engine_utils.convert_prompts_to_data(prompt) + request = self._ffi["create_request"]( + request_id, input_data, generation_config.model_dump_json(by_alias=True) + ) + + # Create the unique async request stream of the request. + stream = engine_base.AsyncRequestStream() + if request_id in self.state.async_streamers: + # Report error in the stream if the request id already exists. + stream.push( + RuntimeError( + f'The request id "{request_id} already exists. ' + 'Please make sure the request id is unique."' + ) + ) + else: + # Record the stream in the tracker + self.state.async_streamers[request_id] = ( + stream, + [TextStreamer(self.tokenizer) for _ in range(generation_config.n)], + ) + self._ffi["add_request"](request) + + def abort_request(): + """clean up""" + self._abort(request_id) + logger.info("request %s cancelled", request_id) + + with engine_utils.ErrorCleanupScope(abort_request): + # Iterate the stream asynchronously and yield the output. + try: + async for request_output in stream: + yield request_output + except asyncio.CancelledError: + # for cancelled error, we can simply pass it through + raise + except Exception as exception: + logger.error("Exception in _generate for request %s: %s", request_id, exception) + raise + + def _abort(self, request_id: str): + """Internal implementation of request abortion.""" + self.state.async_streamers.pop(request_id, None) + self._ffi["abort_request"](request_id) + + +class MLCEngine(engine_base.MLCEngineBase): + """The MLCEngine in MLC LLM that provides the synchronous + interfaces with regard to OpenAI API. + + Parameters + ---------- + model : str + A path to ``mlc-chat-config.json``, or an MLC model directory that contains + `mlc-chat-config.json`. + It can also be a link to a HF repository pointing to an MLC compiled model. + + device: Union[str, Device] + The device used to deploy the model such as "cuda" or "cuda:0". + Will default to "auto" and detect from local available GPUs if not specified. + + model_lib : Optional[str] + The full path to the model library file to use (e.g. a ``.so`` file). + If unspecified, we will use the provided ``model`` to search over possible paths. + It the model lib is not found, it will be compiled in a JIT manner. + + mode : Literal["local", "interactive", "server"] + The engine mode in MLC LLM. + We provide three preset modes: "local", "interactive" and "server". + The default mode is "local". + The choice of mode decides the values of "max_num_sequence", "max_total_sequence_length" + and "prefill_chunk_size" when they are not explicitly specified. + 1. Mode "local" refers to the local server deployment which has low + request concurrency. So the max batch size will be set to 4, and max + total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + 2. Mode "interactive" refers to the interactive use of server, which + has at most 1 concurrent request. So the max batch size will be set to 1, + and max total sequence length and prefill chunk size are set to the context + window size (or sliding window size) of the model. + 3. Mode "server" refers to the large server use case which may handle + many concurrent request and want to use GPU memory as much as possible. + In this mode, we will automatically infer the largest possible max batch + size and max total sequence length. + + You can manually specify arguments "max_num_sequence", "max_total_sequence_length" and + "prefill_chunk_size" to override the automatic inferred values. + + engine_config : Optional[EngineConfig] + Additional configurable arguments of MLC engine. + See class "EngineConfig" for more detail. + + enable_tracing : bool + A boolean indicating if to enable event logging for requests. + """ + + def __init__( + self, + model: str, + device: Union[str, Device] = "auto", + *, + model_lib: Optional[str] = None, + mode: Literal["local", "interactive", "server"] = "local", + engine_config: Optional[EngineConfig] = None, + enable_tracing: bool = False, + ) -> None: + super().__init__( + "sync", + model=model, + device=device, + model_lib=model_lib, + mode=mode, + engine_config=engine_config, + enable_tracing=enable_tracing, + ) + self.chat = Chat(weakref.ref(self)) + self.completions = Completion(weakref.ref(self)) + + def abort(self, request_id: str) -> None: + """Generation abortion interface. + + Parameters + --------- + request_id : str + The id of the request to abort. + """ + self._ffi["abort_request"](request_id) + + def metrics(self) -> engine_base.EngineMetrics: + """Get engine metrics + + Returns + ------- + metrics: EngineMetrics + The engine metrics + """ + return engine_base._query_engine_metrics(self) + + def _chat_completion( + self, + *, + messages: List[Dict[str, Any]], # noqa: UP006 + model: Optional[str] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: bool = False, + top_logprobs: int = 0, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + temperature: Optional[float] = None, + top_p: Optional[float] = None, + tools: Optional[List[Dict[str, Any]]] = None, # noqa: UP006 + tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None, # noqa: UP006 + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + Iterator[openai_api_protocol.ChatCompletionStreamResponse], + openai_api_protocol.ChatCompletionResponse, + ]: + """Synchronous chat completion internal interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/chat/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + debug_config: Optional[Dict[str, Any]] = None, + Extra debug options to pass to the request. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + if request_id is None: + request_id = f"chatcmpl-{engine_utils.random_uuid()}" + + chatcmpl_generator = self._handle_chat_completion( + openai_api_protocol.ChatCompletionRequest( + messages=[ + openai_api_protocol.ChatCompletionMessage.model_validate(message) + for message in messages + ], + model=model, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + top_logprobs=top_logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + temperature=temperature, + top_p=top_p, + tools=( + [openai_api_protocol.ChatTool.model_validate(tool) for tool in tools] + if tools is not None + else None + ), + tool_choice=tool_choice, + user=user, + response_format=( + openai_api_protocol.RequestResponseFormat.model_validate(response_format) + if response_format is not None + else None + ), + debug_config=( + debug_protocol.DebugConfig.model_validate(debug_config) + if debug_config is not None + else None + ), + ), + request_id=request_id, + ) + if stream: + # Stream response. + return chatcmpl_generator + # Normal response. + request_final_usage = None + output_texts = ["" for _ in range(n)] + finish_reasons: List[Optional[str]] = [None for _ in range(n)] # noqa: UP006 + logprob_results: Optional[List[List[openai_api_protocol.LogProbsContent]]] = ( # noqa: UP006 + [[] for _ in range(n)] if logprobs else None + ) + for response in chatcmpl_generator: + # if usage is not None, this is the last chunk + if response.usage is not None: + request_final_usage = response.usage + continue + for choice in response.choices: + assert isinstance(choice.delta.content, str) + output_texts[choice.index] += choice.delta.content + if choice.finish_reason is not None and finish_reasons[choice.index] is None: + finish_reasons[choice.index] = choice.finish_reason + if choice.logprobs is not None: + assert logprob_results is not None + logprob_results[choice.index] += choice.logprobs.content + + assert all(finish_reason is not None for finish_reason in finish_reasons) + use_function_calling, tool_calls_list = engine_base.process_function_call_output( + output_texts, finish_reasons + ) + return engine_base.wrap_chat_completion_response( + request_id=request_id, + model=model, + output_texts=output_texts, + finish_reasons=finish_reasons, + tool_calls_list=tool_calls_list, + logprob_results=logprob_results, + use_function_calling=use_function_calling, + usage=request_final_usage, + ) + + def _completion( + self, + *, + prompt: Union[str, List[int]], # noqa: UP006 + model: Optional[str] = None, + best_of: int = 1, + echo: bool = False, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + logprobs: Optional[int] = None, + logit_bias: Optional[Dict[int, float]] = None, # noqa: UP006 + max_tokens: Optional[int] = None, + n: int = 1, + seed: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, # noqa: UP006 + stream: bool = False, + stream_options: Optional[Dict[str, Any]] = None, # noqa: UP006 + suffix: Optional[str] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + user: Optional[str] = None, + response_format: Optional[Dict[str, Any]] = None, # noqa: UP006 + request_id: Optional[str] = None, + debug_config: Optional[Dict[str, Any]] = None, # noqa: UP006 + ) -> Union[ + Iterator[openai_api_protocol.CompletionResponse], + openai_api_protocol.CompletionResponse, + ]: + """Synchronous completion internal interface with OpenAI API compatibility. + + See https://platform.openai.com/docs/api-reference/completions/create for specification. + + Parameters + ---------- + request_id : Optional[str] + The optional request id. + A random one will be generated if it is not given. + + debug_config: Optional[Dict[str, Any]] = None, + Extra debug options to pass to the request. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + if request_id is None: + request_id = f"cmpl-{engine_utils.random_uuid()}" + + cmpl_generator = self._handle_completion( + openai_api_protocol.CompletionRequest( + model=model, + prompt=prompt, + best_of=best_of, + echo=echo, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + logprobs=logprobs, + logit_bias=logit_bias, + max_tokens=max_tokens, + n=n, + seed=seed, + stop=stop, + stream=stream, + stream_options=( + openai_api_protocol.StreamOptions.model_validate(stream_options) + if stream_options is not None + else None + ), + suffix=suffix, + temperature=temperature, + top_p=top_p, + user=user, + response_format=( + openai_api_protocol.RequestResponseFormat.model_validate(response_format) + if response_format is not None + else None + ), + debug_config=( + debug_protocol.DebugConfig.model_validate(debug_config) + if debug_config is not None + else None + ), + ), + request_id=request_id, + ) + if stream: + # Stream response. + return cmpl_generator + # Normal response. + request_final_usage = None + output_texts = [""] * n + finish_reasons: List[Optional[str]] = [None] * n # noqa: UP006 + logprob_results: List[Optional[openai_api_protocol.CompletionLogProbs]] = [None] * n # noqa: UP006 + + for response in cmpl_generator: + # this is the final chunk + if response.usage is not None: + request_final_usage = response.usage + continue + for choice in response.choices: + output_texts[choice.index] += choice.text + if choice.finish_reason is not None and finish_reasons[choice.index] is None: + finish_reasons[choice.index] = choice.finish_reason + if choice.logprobs is not None: + logprob_results[choice.index] = choice.logprobs + + assert all(finish_reason is not None for finish_reason in finish_reasons) + return engine_base.wrap_completion_response( + request_id=request_id, + model=model, + output_texts=output_texts, + finish_reasons=finish_reasons, + logprob_results=logprob_results, + usage=request_final_usage, + ) + + def _handle_chat_completion( + self, request: openai_api_protocol.ChatCompletionRequest, request_id: str + ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]: + """The implementation fo synchronous ChatCompletionRequest handling. + + Yields + ------ + stream_response : CompletionResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/chat/streaming for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + ( + prompts, + generation_cfg, + use_function_calling, + prompt_length, + ) = engine_base.process_chat_completion_request( + request, + request_id, + self.state, + self.model_config_dicts[0], + self.tokenizer.encode, + self.max_input_sequence_length, + self.conv_template.model_copy(deep=True), + ) + _ = prompt_length + + finish_reasons: List[Optional[str]] = [None for _ in range(generation_cfg.n)] # noqa: UP006 + self.state.record_event(request_id, event="invoke generate") + for delta_outputs in self._generate(prompts, generation_cfg, request_id): + response = engine_base.process_chat_completion_stream_output( + delta_outputs, + request, + request_id, + self.state, + use_function_calling, + finish_reasons, + ) + if response is not None: + yield response + self.state.record_event(request_id, event="finish") + + def _handle_completion( + self, request: openai_api_protocol.CompletionRequest, request_id: str + ) -> Iterator[openai_api_protocol.CompletionResponse]: + """The implementation for synchronous CompletionRequest handling. + + Yields + ------ + stream_response : CompletionResponse + The stream response conforming to OpenAI API. + See mlc_llm/protocol/openai_api_protocol.py or + https://platform.openai.com/docs/api-reference/completions/object for specification. + + Raises + ------ + e : BadRequestError + BadRequestError is raised when the request is invalid. + """ + ( + prompt, + generation_cfg, + prompt_length, + echo_response, + ) = engine_base.process_completion_request( + request, + request_id, + self.state, + self.tokenizer, + self.max_input_sequence_length, + self.conv_template.model_copy(deep=True), + ) + _ = prompt_length + if echo_response is not None: + yield echo_response + + finish_reasons: List[Optional[str]] = [None for _ in range(generation_cfg.n)] # noqa: UP006 + self.state.record_event(request_id, event="invoke generate") + for delta_outputs in self._generate(prompt, generation_cfg, request_id): + response = engine_base.process_completion_stream_output( + delta_outputs, + request, + request_id, + self.state, + finish_reasons, + ) + if response is not None: + yield response + + suffix_response = engine_base.create_completion_suffix_response( + request, request_id, finish_reasons + ) + if suffix_response is not None: + yield suffix_response + self.state.record_event(request_id, event="finish") + + def _generate( + self, + prompt: Union[str, List[int], List[Union[str, List[int], data.Data]]], # noqa: UP006 + generation_config: GenerationConfig, + request_id: str, + ) -> Iterator[List[engine_base.CallbackStreamOutput]]: # noqa: UP006 + """Internal synchronous text generation interface of MLCEngine. + The method is a coroutine that streams a list of CallbackStreamOutput + at a time via yield. The returned list length is the number of + parallel generations specified by `generation_config.n` + except for the final chunk(which is always an List of size 1 and comes with usage) + + Parameters + ---------- + prompt : Union[str, List[int], List[Union[str, List[int], data.Data]]] + The input prompt in forms of text strings, lists of token ids or data. + + generation_config : GenerationConfig + The generation config of the request. + + request_id : str + The unique identifier (in string) or this generation request. + + Yields + ------ + request_output : List[engine_base.CallbackStreamOutput] + The delta generated outputs in a list. + Except for the final chunk, the number of list elements equals to `generation_config.n`, + and each element corresponds to the delta output of a parallel generation. + """ + if self._terminated: + raise ValueError("The engine has terminated.") + + # Create the request with the given id, input data, generation + # config and the created callback. + input_data = engine_utils.convert_prompts_to_data(prompt) + request = self._ffi["create_request"]( + request_id, input_data, generation_config.model_dump_json(by_alias=True) + ) + + # Record the stream in the tracker + self.state.sync_output_queue = queue.Queue() + self.state.sync_text_streamers = [ + TextStreamer(self.tokenizer) for _ in range(generation_config.n) + ] + self._ffi["add_request"](request) + + def abort_request(): + """clean up request if exception happens""" + self.abort(request_id) + + # Iterate the stream asynchronously and yield the token. + with engine_utils.ErrorCleanupScope(abort_request): + while True: + delta_outputs = self.state.sync_output_queue.get() + request_outputs, request_final_usage_json_str = self._request_stream_callback_impl( + delta_outputs + ) + for request_output in request_outputs: + yield request_output + + if request_final_usage_json_str is not None: + # final chunk, we can break + output = engine_base.CallbackStreamOutput( + delta_text="", + delta_logprob_json_strs=None, + finish_reason=None, + request_final_usage_json_str=request_final_usage_json_str, + ) + yield [output] + break + + def _request_stream_callback_impl( + self, + delta_outputs: List[data.RequestStreamOutput], # noqa: UP006 + ) -> Tuple[List[List[engine_base.CallbackStreamOutput]], Optional[str]]: # noqa: UP006 + """The underlying implementation of request stream callback of MLCEngine.""" + batch_outputs: List[List[engine_base.CallbackStreamOutput]] = [] # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + self.state.record_event(request_id, event="start callback") + + # final chunk is now always indicated by a chunk + # where usage json is present + # the backend engine always streams back this chunk + # regardless of include_usage option + is_final_chunk = stream_outputs[0].request_final_usage_json_str is not None + if is_final_chunk: + return (batch_outputs, stream_outputs[0].request_final_usage_json_str) + + outputs: List[engine_base.CallbackStreamOutput] = [] # noqa: UP006 + for stream_output, text_streamer in zip(stream_outputs, self.state.sync_text_streamers): + self.state.record_event(request_id, event="start detokenization") + delta_text = stream_output.extra_prefix_string + ( + text_streamer.put(stream_output.delta_token_ids) + if len(stream_output.delta_token_ids) > 0 + else "" + ) + if stream_output.finish_reason is not None: + delta_text += text_streamer.finish() + self.state.record_event(request_id, event="finish detokenization") + + outputs.append( + engine_base.CallbackStreamOutput( + delta_text=delta_text, + delta_logprob_json_strs=stream_output.delta_logprob_json_strs, + finish_reason=stream_output.finish_reason, + request_final_usage_json_str=None, + ) + ) + batch_outputs.append(outputs) + self.state.record_event(request_id, event="finish callback") + return (batch_outputs, None) diff --git a/python/mlc_llm/serve/engine_base.py b/python/mlc_llm/serve/engine_base.py new file mode 100644 index 0000000..a8198e4 --- /dev/null +++ b/python/mlc_llm/serve/engine_base.py @@ -0,0 +1,1278 @@ +"""The MLC LLM Serving engine base class.""" + +import ast +import asyncio +import json +import numbers +import queue +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import ( # noqa: UP035 + Any, + Callable, + ClassVar, + Dict, + List, + Literal, + Optional, + Tuple, + Union, +) + +import tvm +from tvm.runtime import Device + +from mlc_llm.protocol import openai_api_protocol +from mlc_llm.protocol.conversation_protocol import Conversation +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.protocol.mlc_chat_config import MLCChatConfig +from mlc_llm.serve import data, engine_utils +from mlc_llm.serve.config import EngineConfig +from mlc_llm.serve.event_trace_recorder import EventTraceRecorder +from mlc_llm.support import download_cache, logging +from mlc_llm.support.auto_device import detect_device +from mlc_llm.support.style import green +from mlc_llm.tokenizers import TextStreamer, Tokenizer + +logger = logging.getLogger(__name__) + + +@dataclass +class ModelInfo: + """The model info dataclass. + + Parameters + ---------- + model : str + The identifier of the input model. + It may be a compiled model's id (e.g., "Llama-2-7b-chat-hf-q4f16_1"), + or a full path to a model directory + (e.g., "dist/prebuilt/mlc-chat-Llama-2-7b-chat-hf-q4f16_1") + + model_lib : Optional[str] + The path to the compiled library of the model. + E.g., "dist/prebuilt/lib/Llama-2-7b-chat-hf-q4f16_1-cuda.so" + """ + + model: str + model_lib: Optional[str] = None + + +def _check_engine_config( + model: str, + model_lib: Optional[str], + mode: Literal["local", "interactive", "server"], + engine_config: EngineConfig, +) -> None: + """Check if the given engine config is valid.""" + if engine_config.model is not None and engine_config.model != model: + raise ValueError( + f'The argument "model" of engine constructor is "{model}", while the "model" ' + f'field in argument "engine_config" is "{engine_config.model}". ' + 'Please set the "engine_config.model" to None or set it to the same as the ' + 'argument "model".' + ) + if ( + engine_config.model_lib is not None + and model_lib is not None + and engine_config.model_lib != model_lib + ): + raise ValueError( + f'The argument "model_lib" of engine constructor is "{model_lib}", while the ' + f'"model_lib" field in argument "engine_config" is "{engine_config.model_lib}". ' + 'Please set the "engine_config.model_lib" to None or set it to the same as the ' + 'argument "model_lib".' + ) + if engine_config.mode is not None and engine_config.mode != mode: + raise ValueError( + f'The argument "mode" of engine constructor is "{mode}", while the ' + f'"mode" field in argument "engine_config" is "{engine_config.mode}". ' + 'Please set the "engine_config.mode" to None or set it to the same as the ' + 'argument "mode".' + ) + if engine_config.kv_cache_page_size != 16: + raise ValueError( + 'KV cache only supports page size 16, while the "kv_cache_page_size" field in ' + f'argument "engine_config" is "{engine_config.kv_cache_page_size}". ' + 'Please set "engine_config.kv_cache_page_size" to 16.' + ) + + +def _parse_models( + model: str, + model_lib: Optional[str], + additional_models: List[Union[str, Tuple[str, str]]], # noqa: UP006 +) -> List[ModelInfo]: # noqa: UP006 + """Parse the specified model paths and model libs. + Return a list of ModelInfo, which is a wrapper class of the model path + lib path. + """ + models = [ModelInfo(model, model_lib)] + for additional_model in additional_models: + if isinstance(additional_model, str): + models.append(ModelInfo(additional_model)) + else: + models.append(ModelInfo(additional_model[0], additional_model[1])) + return models + + +def _process_model_args( + models: List[ModelInfo], # noqa: UP006 + device: tvm.runtime.Device, + engine_config: EngineConfig, +) -> Tuple[List[Tuple[str, str]], List[str], Conversation]: # noqa: UP006 + """Process the input ModelInfo to get the engine initialization arguments.""" + conversation: Optional[Conversation] = None + config_file_paths: List[str] = [] # noqa: UP006 + + def _convert_model_info(model: ModelInfo) -> Tuple[str, str]: # noqa: UP006 + nonlocal conversation + + model_path = download_cache.get_or_download_model(model.model) + mlc_config_path = model_path / "mlc-chat-config.json" + config_file_paths.append(str(mlc_config_path)) + + with open(mlc_config_path, encoding="utf-8") as file: + mlc_chat_config = MLCChatConfig.model_validate_json(file.read()) + + if conversation is None: + conversation = mlc_chat_config.conv_template + + if model.model_lib is not None: + # do model lib search if the model lib is provided + # error out if file not found + if model.model_lib.startswith("mock://"): + model_lib = model.model_lib + logger.info("[DEBUG] mock test: %s", model_lib) + elif Path(model.model_lib).is_file(): + model_lib = model.model_lib + logger.info("Using library model: %s", model_lib) + else: + raise FileNotFoundError( + f"The `model_lib` you passed in is not a file: {model.model_lib}.\n" + ) + else: + # Run jit if model_lib is not provided + # NOTE: we only import jit when necessary + # so the engine do not have to depend on compilation + from mlc_llm.interface import jit + + model_compile_overrides = { + "context_window_size": engine_config.max_single_sequence_length, + "prefill_chunk_size": engine_config.prefill_chunk_size, + "sliding_window_size": engine_config.sliding_window_size, + "attention_sink_size": engine_config.attention_sink_size, + "tensor_parallel_shards": engine_config.tensor_parallel_shards, + "pipeline_parallel_stages": engine_config.pipeline_parallel_stages, + "max_batch_size": engine_config.max_num_sequence, + "opt": engine_config.opt, + } + + model_lib = jit.jit( + model_path=model_path, + overrides=model_compile_overrides, + device=device, + ).model_lib_path + return str(model_path), model_lib + + model_args: List[Tuple[str, str]] = [_convert_model_info(model) for model in models] # noqa: UP006 + + assert conversation is not None + return model_args, config_file_paths, conversation + + +def _print_engine_mode_logging_msg( + mode: Literal["local", "interactive", "server"], +) -> None: + """Print the logging info for engine mode selection.""" + if mode == "local": + logger.info( + "The selected engine mode is %s. " + "We choose small max batch size and KV cache capacity to use less GPU memory.", + green(mode), + ) + elif mode == "interactive": + logger.info( + "The selected engine mode is %s. " + "We fix max batch size to 1 for interactive single sequence use.", + green(mode), + ) + else: + logger.info( + "The selected engine mode is %s. " + "We use as much GPU memory as possible (within the limit " + "of gpu_memory_utilization).", + green(mode), + ) + + if mode != "local": + logger.info( + "If you have low concurrent requests and want to use less GPU memory, " + 'please select mode "local".' + ) + if mode != "interactive": + logger.info( + "If you don't have concurrent requests and only use the engine interactively, " + 'please select mode "interactive".' + ) + if mode != "server": + logger.info( + "If you have high concurrent requests and want to maximize the GPU memory utilization, " + 'please select mode "server".' + ) + + +class EngineMetrics: + """Class to store the result returned by engine metrics""" + + metrics: dict + + def __init__(self, metrics): + self.metrics = metrics + + def __str__(self): + return self.metrics.__str__() + + def __repr__(self): + return self.metrics.__repr__() + + def __getitem__(self, key): + return self.metrics[key] + + def prometheus_text(self) -> str: + """Convert engine metrics into prometheus text format + + Returns + ------- + text: str + The metrics in prometheus text format + """ + output_lines = [ + "# NOTE: these metrics count token in the unit of serving model's tokenization", + "# be careful when comparing them to client-side metrics that may use", + "# different tokenization to standardize across models.\n", + ] + + def traverse(comment_scope, key_prefix, curr_value): + if isinstance(curr_value, dict): + if comment_scope: + output_lines.append(f"\n# {comment_scope}") + # first prioritize metrics in current scope + for key, value in curr_value.items(): + if isinstance(value, numbers.Number): + output_lines.append(f"{key_prefix}{key}\t{value}") + # then look into nested scopes if any + for key, value in curr_value.items(): + if isinstance(value, dict) and len(value) != 0: + traverse(f"{comment_scope}/{key}", f"{key_prefix}{key}_", value) + + traverse("", "", self.metrics) + return "\n".join(output_lines) + + +def _query_engine_metrics(engine): + """Query engine metrics via debug options""" + dummy_message = {"role": "user", "context": ""} + for response in engine.chat.completions.create( + messages=[dummy_message], + model="model", + stream=True, + stream_options={"include_usage": True}, + extra_body={"debug_config": {"special_request": "query_engine_metrics"}}, + ): + if response.usage is not None: + return EngineMetrics(response.usage.extra) + raise RuntimeError("query_engine metrics did not get metrics back") + + +async def _async_query_engine_metrics(engine): + """Query engine metrics via debug options""" + dummy_message = {"role": "user", "context": ""} + result = None + async for response in await engine.chat.completions.create( + messages=[dummy_message], + model="model", + stream=True, + stream_options={"include_usage": True}, + extra_body={"debug_config": {"special_request": "query_engine_metrics"}}, + ): + if response.usage is not None: + assert result is None + result = EngineMetrics(response.usage.extra) + + if result is not None: + return result + raise RuntimeError("query_engine metrics did not get metrics back") + + +@dataclass +class CallbackStreamOutput: + """The output of MLCEngine._generate and AsyncMLCEngine._generate + + Attributes + ---------- + delta_text : str + The delta text generated since the last output. + + delta_logprob_json_strs : Optional[List[str]] + The list of logprob JSON strings since the last output, + or None if the request does not require logprobs. + + finish_reason : Optional[str] + The finish reason of the request, or None if unfinished. + + request_final_usage_json_str: Optional[str] + The usage json which appears in last chunk, + when it appears all other fields will be empty + """ + + delta_text: str + delta_logprob_json_strs: Optional[List[str]] # noqa: UP006 + finish_reason: Optional[str] + request_final_usage_json_str: Optional[str] + + +class AsyncRequestStream: + """The asynchronous stream for requests in AsyncMLCEngine. + + Each request has its own unique stream. + The stream exposes the method `push` for engine to push new generated + delta text to the stream, and the method `finish` for engine to mark + the finish of generation. + + The stream implements `__aiter__` and `__anext__`, which the engine + can use to iterates all the generated tokens in order asynchronously. + """ + + # The asynchronous queue to hold elements of either a list of + # CallbackStreamOutput or an exception. + _queue: asyncio.Queue[ + Union[List[CallbackStreamOutput], Exception] # noqa: UP006 + ] + # The finish flag. + _finished: bool + + def __init__(self) -> None: + self._queue = asyncio.Queue() + self._finished = False + + def push(self, item_or_exception: Union[List[CallbackStreamOutput], Exception]) -> None: # noqa: UP006 + """Push a new token to the stream.""" + if self._finished: + # No new item is expected after finish. + self._queue.put_nowait( + RuntimeError( + "The request has already finished. " + "The stream is not supposed to accept new items." + ) + ) + return + self._queue.put_nowait(item_or_exception) + + def finish(self) -> None: + """Mark the finish of the generation in the stream.""" + self._queue.put_nowait(StopIteration()) + self._finished = True + + def __aiter__(self): + return self + + async def __anext__(self) -> List[CallbackStreamOutput]: # noqa: UP006 + result = await self._queue.get() + if isinstance(result, StopIteration): + raise StopAsyncIteration + if isinstance(result, Exception): + raise result + return result + + +class EngineState: + """The engine states that the request stream callback function may use. + + This class is used for both AsyncMLCEngine and MLCEngine. + AsyncMLCEngine uses the fields and methods starting with "async", + and MLCEngine uses the ones starting with "sync". + + - For AsyncMLCEngine, the state contains an asynchronous event loop, + the streamers and the number of unfinished generations for each request + being processed. + - For MLCEngine, the state contains a callback output blocking queue, + the text streamers and the number of unfinished requests. + + We use this state class to avoid the callback function from capturing + the AsyncMLCEngine. + + The state also optionally maintains an event trace recorder, which can + provide Chrome tracing when enabled. + """ + + trace_recorder = None + # States used for AsyncMLCEngine + async_event_loop: Optional[asyncio.AbstractEventLoop] = None + async_streamers: ClassVar[Dict[str, Tuple[AsyncRequestStream, List[TextStreamer]]]] = {} # noqa: UP006 + # States used for MLCEngine + sync_output_queue: queue.Queue = queue.Queue() + sync_text_streamers: ClassVar[List[TextStreamer]] = [] # noqa: UP006 + + def __init__(self, enable_tracing: bool) -> None: + """Constructor.""" + if enable_tracing: + self.trace_recorder = EventTraceRecorder() + + def record_event(self, request_id: str, event: str) -> None: + """Record a event for the input request in the trace + recorder when the recorder exists. + + Parameters + ---------- + request_id : str + The subject request of the event. + + event : str + The event in a string name. + It can have one of the following patterns: + - "start xxx", which marks the start of event "xxx", + - "finish xxx", which marks the finish of event "xxx", + - "yyy", which marks the instant event "yyy". + The "starts" and "finishes" will be automatically paired in the trace recorder. + """ + if self.trace_recorder is None: + return + self.trace_recorder.add_event(request_id, event) + + def get_request_stream_callback( + self, kind: Literal["async", "sync"] + ) -> Callable[[List[data.RequestStreamOutput]], None]: # noqa: UP006 + """Construct a callback function and return. + + The callback function has signature + "Callable[[List[data.RequestStreamOutput]], None]", + whose input is a list of "data.RequestStreamOutput". + Each "data.RequestStreamOutput" is the delta output of a request, + generated from the engine. + """ + + f_callback = ( + self._async_request_stream_callback + if kind == "async" + else self._sync_request_stream_callback + ) + + def _callback(delta_outputs: List[data.RequestStreamOutput]) -> None: # noqa: UP006 + f_callback(delta_outputs) + + return _callback + + def async_lazy_init_event_loop(self) -> None: + """Lazily set the asyncio event loop so that the event + loop is the main driving event loop of the process. + """ + if self.async_event_loop is None: + self.async_event_loop = asyncio.get_event_loop() + + def _async_request_stream_callback(self, delta_outputs: List[data.RequestStreamOutput]) -> None: # noqa: UP006 + """The request stream callback function for AsyncMLCEngine to stream back + the request generation results. + + Note + ---- + This callback function uses `call_soon_threadsafe` in asyncio to + schedule the invocation in the event loop, so that the underlying + callback logic will be executed asynchronously in the future rather + than right now. + """ + + # Schedule a callback run in the event loop without executing right now. + # NOTE: This function causes GIL during execution. + self.async_event_loop.call_soon_threadsafe( + self._async_request_stream_callback_impl, delta_outputs + ) + + def _async_request_stream_callback_impl( + self, + delta_outputs: List[data.RequestStreamOutput], # noqa: UP006 + ) -> None: + """The underlying implementation of request stream callback for AsyncMLCEngine.""" + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + streamers = self.async_streamers.get(request_id, None) + if streamers is None: + continue + + self.record_event(request_id, event="start callback") + stream, text_streamers = streamers + + # final chunk is now always indicated by a chunk + # where usage json is present + # the backend engine always streams back this chunk + # regardless of include_usage option + is_final_chunk = stream_outputs[0].request_final_usage_json_str is not None + if is_final_chunk: + # stream back this final usage chunk + output = CallbackStreamOutput( + delta_text="", + delta_logprob_json_strs=None, + finish_reason=None, + request_final_usage_json_str=stream_outputs[0].request_final_usage_json_str, + ) + stream.push([output]) + stream.finish() + self.async_streamers.pop(request_id, None) + continue + + outputs = [] + for stream_output, text_streamer in zip(stream_outputs, text_streamers): + self.record_event(request_id, event="start detokenization") + delta_text = stream_output.extra_prefix_string + ( + text_streamer.put(stream_output.delta_token_ids) + if len(stream_output.delta_token_ids) > 0 + else "" + ) + if stream_output.finish_reason is not None: + delta_text += text_streamer.finish() + self.record_event(request_id, event="finish detokenization") + + outputs.append( + CallbackStreamOutput( + delta_text=delta_text, + delta_logprob_json_strs=stream_output.delta_logprob_json_strs, + finish_reason=stream_output.finish_reason, + request_final_usage_json_str=None, + ) + ) + + # Push new delta text to the stream. + stream.push(outputs) + self.record_event(request_id, event="finish callback") + + def _sync_request_stream_callback(self, delta_outputs: List[data.RequestStreamOutput]) -> None: # noqa: UP006 + """The request stream callback function for MLCEngine to stream back + the request generation results. + """ + # Put the delta outputs to the queue in the unblocking way. + self.sync_output_queue.put_nowait(delta_outputs) + + +class MLCEngineBase: + """The base engine class, which implements common functions that + are shared by MLCEngine and AsyncMLCEngine. + + This class wraps a threaded engine that runs on a standalone + thread inside and streams back the delta generated results via + callback functions. The internal threaded engine keeps running an + loop that drives the engine. + + MLCEngine and AsyncMLCEngine inherits this MLCEngineBase class, and implements + their own methods to process the delta generated results received + from callback functions and yield the processed delta results in + the forms of standard API protocols. + + Checkout subclasses AsyncMLCEngine/MLCEngine for the docstring of constructor parameters. + """ + + def __init__( + self, + kind: Literal["async", "sync"], + model: str, + device: Union[str, tvm.runtime.Device], + model_lib: Optional[str], + mode: Literal["local", "interactive", "server"], + engine_config: Optional[EngineConfig], + enable_tracing: bool, + ) -> None: + # - Check the fields fields of `engine_config`. + if engine_config is None: + engine_config = EngineConfig() + _check_engine_config(model, model_lib, mode, engine_config) + + # - Initialize model loading info. + models = _parse_models(model, model_lib, engine_config.additional_models) + if isinstance(device, str): + device = detect_device(device) + assert isinstance(device, Device) + ( + model_args, + model_config_paths, + self.conv_template, + ) = _process_model_args(models, device, engine_config) + + # - Load the raw model config into dict + self.model_config_dicts = [] + for i, model_info in enumerate(models): + model_info.model_lib = model_args[i][1] + with open(model_config_paths[i], encoding="utf-8") as file: + self.model_config_dicts.append(json.load(file)) + + # - Print logging info for regarding the mode selection. + if engine_config.verbose: + _print_engine_mode_logging_msg(mode) + + # - Initialize engine state and engine. + self.state = EngineState(enable_tracing) + module = tvm.get_global_func("mlc.serve.create_threaded_engine", allow_missing=False)() + self._ffi = { + key: module[key] + for key in [ + "add_request", + "abort_request", + "run_background_loop", + "run_background_stream_back_loop", + "reload", + "init_threaded_engine", + "exit_background_loop", + "create_request", + "get_complete_engine_config", + "reset", + "debug_call_func_on_all_worker", + ] + } + self.tokenizer = Tokenizer(model_args[0][0]) + self._ffi["init_threaded_engine"]( + device, + self.state.get_request_stream_callback(kind), + self.state.trace_recorder, + ) + + background_loop = self._ffi["run_background_loop"] + background_stream_back_loop = self._ffi["run_background_stream_back_loop"] + + # - Create the background engine-driving thread and start the loop. + self._background_loop_thread: threading.Thread = threading.Thread(target=background_loop) + self._background_stream_back_loop_thread: threading.Thread = threading.Thread( + target=background_stream_back_loop + ) + self._background_loop_thread.start() + self._background_stream_back_loop_thread.start() + self._terminated = False + + engine_config.model = model_args[0][0] + engine_config.model_lib = model_args[0][1] + engine_config.additional_models = model_args[1:] + engine_config.mode = mode + self._ffi["reload"](engine_config.asjson()) + self.engine_config = EngineConfig.from_json(self._ffi["get_complete_engine_config"]()) + self.max_input_sequence_length = min( + self.engine_config.max_single_sequence_length, + self.engine_config.max_total_sequence_length, + ) + + def __del__(self): + """deleter, auto terminate""" + self.terminate() + + def terminate(self): + """Terminate the engine.""" + if hasattr(self, "_terminated") and self._terminated: + return + self._terminated = True + if not hasattr(self, "_ffi"): + return + self._ffi["exit_background_loop"]() + if hasattr(self, "_background_loop_thread"): + self._background_loop_thread.join() + if hasattr(self, "_background_stream_back_loop_thread"): + self._background_stream_back_loop_thread.join() + + def _debug_call_func_on_all_worker( + self, func_name: str, func_args: Optional[str] = None + ) -> None: + """Call the given global function on all workers. Only for debug purpose.""" + self._ffi["debug_call_func_on_all_worker"](func_name, func_args) + + def reset(self): + """Reset the engine, clear the running data and metrics.""" + return self._ffi["reset"]() + + +def process_chat_completion_request( + request: openai_api_protocol.ChatCompletionRequest, + request_id: str, + engine_state: EngineState, + model_config: Dict[str, Any], # noqa: UP006 + f_tokenize: Callable[[str], List[int]], # noqa: UP006 + max_input_sequence_length: int, + conv_template: Conversation, +) -> Tuple[List[Union[List[int], data.Data]], GenerationConfig, bool, int]: # noqa: UP006 + """Process the given ChatCompletionRequest, apply request validity + checks, and return the processed prompts, and other info. + + Parameters + ---------- + request : openai_api_protocol.ChatCompletionRequest + The request to be processed and checked. + + request_id : str + The id of the request. + + engine_state : EngineState + The state of the engine. + + model_config : Dict[str, Any] + The model configuration dictionary. + + f_tokenize : Callable[[str], List[int]] + The tokenizer encode function. + + max_input_sequence_length : int + The maximum allowed total prompt length. + + conv_template : Conversation + The conversation template of the model. + + Returns + ------- + prompts : List[Union[List[int], data.Data]] + The prompts, in a list. + Each element is a list of token ids or a "data.Data" instance. + + generation_cfg : GenerationConfig + The generation config of the request got from the input request. + + use_function_calling : bool + A boolean flag indicating if the request uses function call. + + prompt_length : int + The total prompt length. + """ + engine_state.record_event(request_id, event="receive request") + # - Check if unsupported arguments are specified. + engine_utils.check_unsupported_fields(request) + + # - Process messages and update the conversation template in three steps: + # i. Check the message validity. + # ii. Add the input messages to the conversation template. + # iii. Add the additional message for the assistant. + request.check_message_validity() + # - Check for function calling usage and update the conversation template + request.check_function_call_usage(conv_template) + + for message in request.messages: + role = message.role + content = message.content + if role == "system": + assert isinstance(content, str) + conv_template.system_message = content if content is not None else "" + continue + conv_template.messages.append((role, content)) + conv_template.messages.append(("assistant", None)) + + # - Get the prompt from template, and encode to token ids. + # - Check prompt length + engine_state.record_event(request_id, event="start tokenization") + prompts = engine_utils.process_prompts(conv_template.as_prompt(model_config), f_tokenize) + engine_state.record_event(request_id, event="finish tokenization") + + if conv_template.system_prefix_token_ids is not None: + if isinstance(prompts[0], list): + prompts[0] = conv_template.system_prefix_token_ids + prompts[0] + else: + prompts.insert(0, conv_template.system_prefix_token_ids) + prompt_length = engine_utils.check_and_get_prompts_length(prompts, max_input_sequence_length) + + # Process generation config. Create request id. + generation_cfg = engine_utils.get_generation_config( + request, + extra_stop_token_ids=conv_template.stop_token_ids, + extra_stop_str=conv_template.stop_str, + ) + return prompts, generation_cfg, conv_template.use_function_calling, prompt_length + + +def process_chat_completion_stream_output( + delta_outputs: List[CallbackStreamOutput], # noqa: UP006 + request: openai_api_protocol.ChatCompletionRequest, + request_id: str, + engine_state: EngineState, + use_function_calling: bool, + finish_reasons: List[Optional[str]], # noqa: UP006 +) -> Optional[openai_api_protocol.ChatCompletionStreamResponse]: + """Process the delta outputs of a single request of ChatCompletion, + convert the delta output to ChatCompletionStreamResponse and return. + + Parameters + ---------- + delta_outputs : List[CallbackStreamOutput] + The delta outputs of a request. + The list length is the number of parallel generation specified by "n". + Each element corresponds to a generation. + + request_id : str + The id of the request. + + engine_state : EngineState + The state of the engine. + + use_function_calling : bool + A boolean flag indicating if the request uses function call. + + finish_reasons : List[Optional[str]] + The list of finish reasons of each generation. + The list length is the number of parallel generation specified by "n". + This list is updated in place. + + Returns + ------- + response : Optional[openai_api_protocol.ChatCompletionStreamResponse] + The converted OpenAI API ChatCompletionStreamResponse instance. + It can be none when there is no content. + """ + # we always stream back the final chunk with usage + is_final_chunk = delta_outputs[0].request_final_usage_json_str is not None + if is_final_chunk: + assert len(delta_outputs) == 1 + engine_state.record_event(request_id, event="yield final usage") + response = openai_api_protocol.ChatCompletionStreamResponse( + id=request_id, + choices=[], + model=request.model, + system_fingerprint="", + usage=openai_api_protocol.CompletionUsage.model_validate_json( + delta_outputs[0].request_final_usage_json_str + ), + ) + # non streaming mode always comes with usage + if not request.stream: + return response + # skip usage if stream option does not indicate include usage + if request.stream_options is None: + return None + if not request.stream_options.include_usage: + return None + return response + + # normal chunk + assert len(delta_outputs) == request.n + choices = [] + for i, delta_output in enumerate(delta_outputs): + finish_reason_updated = False + if delta_output.finish_reason is not None and finish_reasons[i] is None: + finish_reasons[i] = ( + delta_output.finish_reason if not use_function_calling else "tool_calls" + ) + finish_reason_updated = True + if not finish_reason_updated and delta_output.delta_text == "": + # Ignore empty delta text when finish reason is not updated. + engine_state.record_event(request_id, event="skip empty delta text") + continue + + choices.append( + openai_api_protocol.ChatCompletionStreamResponseChoice( + index=i, + finish_reason=finish_reasons[i], + delta=openai_api_protocol.ChatCompletionMessage( + content=delta_output.delta_text, role="assistant" + ), + logprobs=( + openai_api_protocol.LogProbs( + content=[ + openai_api_protocol.LogProbsContent.model_validate_json( + logprob_json_str + ) + for logprob_json_str in delta_output.delta_logprob_json_strs + ] + ) + if delta_output.delta_logprob_json_strs is not None + else None + ), + ) + ) + + if len(choices) == 0: + # Skip return when there is no delta output and no number of completion tokens. + return None + response = openai_api_protocol.ChatCompletionStreamResponse( + id=request_id, choices=choices, model=request.model, system_fingerprint="" + ) + engine_state.record_event(request_id, event="yield delta output") + return response + + +def process_completion_request( + request: openai_api_protocol.CompletionRequest, + request_id: str, + engine_state: EngineState, + tokenizer: Tokenizer, + max_input_sequence_length: int, + conv_template: Conversation, +) -> Tuple[List[int], GenerationConfig, int, Optional[openai_api_protocol.CompletionResponse]]: # noqa: UP006 + """Process the given CompletionRequest, apply request validity + checks, and return the processed prompts, and other info. + + Parameters + ---------- + request : openai_api_protocol.CompletionRequest + The request to be processed and checked. + + request_id : str + The id of the request. + + engine_state : EngineState + The state of the engine. + + tokenizer : Tokenizer + The tokenizer instance of the model. + + max_input_sequence_length : int + The maximum allowed total prompt length. + + conv_template : Conversation + The conversation template of the model. + + Returns + ------- + prompt : List[int] + The prompt in a list of token ids. + + generation_cfg : GenerationConfig + The generation config of the request got from the input request. + + prompt_length : int + The total prompt length. + + echo_response : Optional[openai_api_protocol.CompletionResponse] + The CompletionResponse of the echoing part, when argument "echo" + of the input request is specified. + """ + engine_state.record_event(request_id, event="receive request") + # - Check if unsupported arguments are specified. + engine_utils.check_unsupported_fields(request) + + # - Process prompt and check validity. + engine_state.record_event(request_id, event="start tokenization") + prompts = engine_utils.process_prompts(request.prompt, tokenizer.encode) + engine_state.record_event(request_id, event="finish tokenization") + prompt_length = engine_utils.check_and_get_prompts_length(prompts, max_input_sequence_length) + prompt = prompts[0] + assert isinstance(prompt, list) + + # Process generation config. Create request id. + generation_cfg = engine_utils.get_generation_config( + request, + extra_stop_token_ids=conv_template.stop_token_ids, + extra_stop_str=conv_template.stop_str, + ) + + # - Echo back the prompt. + echo_response = None + if request.echo: + text = tokenizer.decode(prompt) + response = openai_api_protocol.CompletionResponse( + id=request_id, + choices=[ + openai_api_protocol.CompletionResponseChoice(index=i, text=text) + for i in range(generation_cfg.n) + ], + model=request.model, + usage=None, + ) + echo_response = response + return prompt, generation_cfg, prompt_length, echo_response + + +def get_logprobs_from_delta( + delta_logprob_json_strs: List[str], # noqa: UP006 +) -> openai_api_protocol.CompletionLogProbs: + """Convert json strings containing logprobs information to + completion response format (OpenAI API compatible) + + Parameters + ---------- + delta_logprob_json_strs : List[str] + Logprobs information packed in json strings and + kept in the delta outputs of a request. + + Returns + ------- + logprobs : openai_api_protocol.CompletionLogProbs + Logprobs information extracted from json string and converted to completion response format + """ + token_logprobs = [] + tokens = [] + top_logprobs = [] + for logprob_json_str in delta_logprob_json_strs: + content = openai_api_protocol.LogProbsContent.model_validate_json(logprob_json_str) + tokens.append(content.token) + token_logprobs.append(content.logprob) + top_logprob_dict = {} + for top_logprob in content.top_logprobs: + top_logprob_dict[top_logprob.token] = top_logprob.logprob + top_logprobs.append(top_logprob_dict) + return openai_api_protocol.CompletionLogProbs( + # TODO(vvchernov): support text_offset + text_offset=None, + token_logprobs=token_logprobs, + tokens=tokens, + top_logprobs=top_logprobs, + ) + + +def process_completion_stream_output( + delta_outputs: List[CallbackStreamOutput], # noqa: UP006 + request: openai_api_protocol.CompletionRequest, + request_id: str, + engine_state: EngineState, + finish_reasons: List[Optional[str]], # noqa: UP006 +) -> Optional[openai_api_protocol.CompletionResponse]: + """Process the delta outputs of a single request of Completion, + convert the delta output to CompletionResponse and return. + + Parameters + ---------- + delta_outputs : List[CallbackStreamOutput] + The delta outputs of a request. + The list length is the number of parallel generation specified by "n". + Each element corresponds to a generation. + + request: openai_api_protocol.CompletionRequest + Information about the request + + request_id : str + The id of the request. + + engine_state : EngineState + The state of the engine. + + finish_reasons : List[Optional[str]] + The list of finish reasons of each generation. + The list length is the number of parallel generation specified by "n". + This list is updated in place. + + Returns + ------- + response : Optional[openai_api_protocol.CompletionResponse] + The converted OpenAI API CompletionResponse instance. + It can be none when there is no content. + """ + # we always stream back the final chunk with usage + is_final_chunk = delta_outputs[0].request_final_usage_json_str is not None + if is_final_chunk: + assert len(delta_outputs) == 1 + engine_state.record_event(request_id, event="yield final usage") + response = openai_api_protocol.CompletionResponse( + id=request_id, + choices=[], + model=request.model, + system_fingerprint="", + usage=openai_api_protocol.CompletionUsage.model_validate_json( + delta_outputs[0].request_final_usage_json_str + ), + ) + # non streaming mode always comes with usage + if not request.stream: + return response + if request.stream_options is None: + return None + if not request.stream_options.include_usage: + return None + return response + + # normal chunk + assert len(delta_outputs) == request.n + choices = [] + for i, delta_output in enumerate(delta_outputs): + finish_reason_updated = False + if delta_output.finish_reason is not None and finish_reasons[i] is None: + finish_reasons[i] = delta_output.finish_reason + finish_reason_updated = True + if not finish_reason_updated and delta_output.delta_text == "": + # Ignore empty delta text when finish reason is not updated. + continue + + if delta_output.delta_logprob_json_strs is not None: + logprobs = get_logprobs_from_delta(delta_output.delta_logprob_json_strs) + else: + logprobs = None + choices.append( + openai_api_protocol.CompletionResponseChoice( + index=i, + finish_reason=finish_reasons[i], + text=delta_output.delta_text, + logprobs=logprobs, + ) + ) + + if len(choices) == 0: + # Skip return when there is no delta output and no number of completion tokens. + return None + response = openai_api_protocol.CompletionResponse( + id=request_id, + choices=choices, + model=request.model, + usage=None, + ) + engine_state.record_event(request_id, event="yield delta output") + return response + + +def create_completion_suffix_response( + request: openai_api_protocol.CompletionRequest, + request_id: str, + finish_reasons: List[Optional[str]], # noqa: UP006 +) -> Optional[openai_api_protocol.CompletionResponse]: + """Create the suffix response of Completion request + when the request requires suffix. + + Parameters + ---------- + request : openai_api_protocol.CompletionRequest + The request whose suffix response if to be created. + + request_id : str + The id of the request. + + finish_reasons : List[Optional[str]] + The list of finish reasons of each generation. + The list length is the number of parallel generation specified by "n". + This list is updated in place. + + Returns + ------- + suffix_response : Optional[openai_api_protocol.CompletionResponse] + The created OpenAI API CompletionResponse instance for the suffix. + Or None if the request does not require suffix. + """ + # - Echo the suffix. + if request.suffix is None: + return None + assert all(finish_reason is not None for finish_reason in finish_reasons) + response = openai_api_protocol.CompletionResponse( + id=request_id, + choices=[ + openai_api_protocol.CompletionResponseChoice( + index=i, + finish_reason=finish_reason, + text=request.suffix, + ) + for i, finish_reason in enumerate(finish_reasons) + ], + model=request.model, + usage=None, + ) + return response + + +def convert_function_str_to_json(stringified_calls: str) -> List[Union[Dict, None]]: # noqa: UP006 + """Convert a (possibly list) of function call string to a list of json objects. + Return None for invalid function call string.""" + + def parse_function_call(call_str: str): + node = ast.parse(call_str, mode="eval") + call_node = node.body + if isinstance(call_node, ast.Call) and isinstance(call_node.func, ast.Name): + name = call_node.func.id + arguments = {} + for keyword in call_node.keywords: + arguments[keyword.arg] = ast.literal_eval(keyword.value) + return {"name": name, "arguments": arguments} + return None + + if ( + stringified_calls[0] == "[" and stringified_calls[-1] == "]" + ): # hacky way to check if string list + calls = ast.literal_eval(stringified_calls) + else: + calls = [stringified_calls] + function_calls_json = [parse_function_call(call_str) for call_str in calls] + return function_calls_json + + +def process_function_call_output( + output_texts: List[str], # noqa: UP006 + finish_reasons: List[str], # noqa: UP006 +) -> Tuple[bool, List[List[openai_api_protocol.ChatToolCall]]]: # noqa: UP006 + """Process the potential function call results outputted by model, + according to the finish reasons. + Return whether the output has function call, and the list of tool calls. + """ + n = len(output_texts) + tool_calls_list: List[List[openai_api_protocol.ChatToolCall]] = [[] for _ in range(n)] # noqa: UP006 + use_function_calling = any(finish_reason == "tool_calls" for finish_reason in finish_reasons) + if use_function_calling: + for i, output_text in enumerate(output_texts): + try: + fn_json_list = convert_function_str_to_json(output_text) + except (SyntaxError, ValueError): + output_text = "Got an invalid function call output from model" + finish_reasons[i] = "error" + else: + tool_calls_list[i] = [ + openai_api_protocol.ChatToolCall( + type="function", + function=openai_api_protocol.ChatFunctionCall( + name=fn_json_obj["name"], arguments=fn_json_obj["arguments"] + ), + ) + for fn_json_obj in fn_json_list + if fn_json_obj is not None + ] + if len(tool_calls_list[i]) == 0: + output_texts[i] = "Got an invalid function call output from model" + finish_reasons[i] = "error" + else: + finish_reasons[i] = "tool_calls" + return use_function_calling, tool_calls_list + + +def wrap_chat_completion_response( + request_id: str, + model: str, + output_texts: List[str], # noqa: UP006 + finish_reasons: List[str], # noqa: UP006 + tool_calls_list: List[List[openai_api_protocol.ChatToolCall]], # noqa: UP006 + logprob_results: Optional[List[List[openai_api_protocol.LogProbsContent]]], # noqa: UP006 + use_function_calling: bool, + usage: Optional[Dict[str, Any]], # noqa: UP006 +) -> openai_api_protocol.ChatCompletionResponse: + """Wrap the non-streaming chat completion results to ChatCompletionResponse instance.""" + return openai_api_protocol.ChatCompletionResponse( + id=request_id, + choices=[ + openai_api_protocol.ChatCompletionResponseChoice( + index=i, + finish_reason=finish_reasons[i], + message=( + openai_api_protocol.ChatCompletionMessage(role="assistant", content=output_text) + if not use_function_calling or finish_reason == "error" + else openai_api_protocol.ChatCompletionMessage( + role="assistant", tool_calls=tool_calls + ) + ), + logprobs=( + openai_api_protocol.LogProbs(content=logprob_results[i]) + if logprob_results is not None + else None + ), + ) + for i, (output_text, finish_reason, tool_calls) in enumerate( + zip(output_texts, finish_reasons, tool_calls_list) + ) + ], + model=model, + system_fingerprint="", + usage=usage, + ) + + +def wrap_completion_response( + request_id: str, + model: str, + output_texts: List[str], # noqa: UP006 + finish_reasons: List[str], # noqa: UP006 + logprob_results: List[Optional[openai_api_protocol.CompletionLogProbs]], # noqa: UP006 + usage: openai_api_protocol.CompletionUsage, +) -> openai_api_protocol.CompletionResponse: + """Wrap the non-streaming completion results to CompletionResponse instance.""" + return openai_api_protocol.CompletionResponse( + id=request_id, + choices=[ + openai_api_protocol.CompletionResponseChoice( + index=i, + finish_reason=finish_reason, + text=output_text, + logprobs=logprob_results[i], + ) + for i, (output_text, finish_reason) in enumerate(zip(output_texts, finish_reasons)) + ], + model=model, + usage=usage, + ) diff --git a/python/mlc_llm/serve/engine_utils.py b/python/mlc_llm/serve/engine_utils.py new file mode 100644 index 0000000..7228865 --- /dev/null +++ b/python/mlc_llm/serve/engine_utils.py @@ -0,0 +1,337 @@ +"""Utility functions for MLC Serve engine""" + +import uuid +from typing import Any, Callable, Dict, List, Literal, Optional, Union # noqa: UP035 + +from mlc_llm.protocol import error_protocol, openai_api_protocol +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import data + +RequestProtocol = Union[ + openai_api_protocol.CompletionRequest, openai_api_protocol.ChatCompletionRequest +] + + +def get_unsupported_fields(request: RequestProtocol) -> List[str]: # noqa: UP006 + """Get the unsupported fields of the request. + Return the list of unsupported field names. + """ + if isinstance( + request, + ( + openai_api_protocol.CompletionRequest, + openai_api_protocol.ChatCompletionRequest, + ), + ): + return openai_api_protocol.openai_api_get_unsupported_fields(request) + raise RuntimeError("Cannot reach here") + + +def openai_api_get_generation_config(request: RequestProtocol) -> Dict[str, Any]: # noqa: UP006 + """Create the generation config from the given request.""" + kwargs: Dict[str, Any] = {} # noqa: UP006 + arg_names = [ + "n", + "temperature", + "top_p", + "max_tokens", + "frequency_penalty", + "presence_penalty", + "logit_bias", + "seed", + "response_format", + "debug_config", + ] + for arg_name in arg_names: + kwargs[arg_name] = getattr(request, arg_name) + if kwargs["max_tokens"] is None: + # Setting to -1 means the generation will not stop until + # exceeding model capability or hit any stop criteria. + kwargs["max_tokens"] = -1 + if request.stop is not None: + kwargs["stop_strs"] = [request.stop] if isinstance(request.stop, str) else request.stop + if isinstance(request, openai_api_protocol.ChatCompletionRequest): + kwargs["logprobs"] = request.logprobs + kwargs["top_logprobs"] = request.top_logprobs + else: + logprobs = request.logprobs is not None + kwargs["logprobs"] = logprobs + kwargs["top_logprobs"] = request.logprobs if logprobs else 0 + return kwargs + + +def get_generation_config( + request: RequestProtocol, + extra_stop_token_ids: Optional[List[int]] = None, # noqa: UP006 + extra_stop_str: Optional[List[str]] = None, # noqa: UP006 +) -> GenerationConfig: + """Create the generation config in MLC LLM out from the input request protocol.""" + kwargs: Dict[str, Any] # noqa: UP006 + if isinstance( + request, + ( + openai_api_protocol.CompletionRequest, + openai_api_protocol.ChatCompletionRequest, + ), + ): + kwargs = openai_api_get_generation_config(request) + else: + raise RuntimeError("Cannot reach here") + + if extra_stop_token_ids is not None: + stop_token_ids = kwargs.get("stop_token_ids", []) + assert isinstance(stop_token_ids, list) + stop_token_ids += extra_stop_token_ids + kwargs["stop_token_ids"] = stop_token_ids + + if extra_stop_str is not None: + stop_strs = kwargs.get("stop_strs", []) + assert isinstance(stop_strs, list) + stop_strs += extra_stop_str + kwargs["stop_strs"] = stop_strs + + return GenerationConfig(**kwargs) + + +def random_uuid() -> str: + """Generate a random id in hexadecimal string.""" + return uuid.uuid4().hex + + +def check_unsupported_fields(request: RequestProtocol) -> None: + """Check if the request has unsupported fields. Raise BadRequestError if so.""" + unsupported_fields = get_unsupported_fields(request) + if len(unsupported_fields) != 0: + unsupported_fields = [f'"{field}"' for field in unsupported_fields] + raise error_protocol.BadRequestError( + f"Request fields {', '.join(unsupported_fields)} are not supported right now.", + ) + + +def check_and_get_prompts_length( + prompts: List[Union[List[int], data.ImageData]], # noqa: UP006 + max_input_sequence_length: int, +) -> int: + """Check if the total prompt length exceeds the max single sequence + sequence length allowed by the served model. Raise BadRequestError if so. + Return the total prompt length. + """ + total_length: int = 0 + for prompt in prompts: + total_length += len(prompt) + if total_length > max_input_sequence_length: + raise error_protocol.BadRequestError( + f"Request prompt has {total_length} tokens in total," + f" larger than the model input length limit {max_input_sequence_length}.", + ) + return total_length + + +def process_prompts( + input_prompts: Union[str, List[int], List[Union[str, List[int], data.ImageData]]], # noqa: UP006 + ftokenize: Callable[[str], List[int]], # noqa: UP006 +) -> List[Union[List[int], data.ImageData]]: # noqa: UP006 + """Convert all input tokens to list of token ids with regard to the + given tokenization function. + For each input prompt, return the list of token ids after tokenization. + """ + error_msg = f"Invalid request prompt {input_prompts}" + + # Case 1. The prompt is a single string. + if isinstance(input_prompts, str): + return [ftokenize(input_prompts)] + + assert isinstance(input_prompts, list) + if len(input_prompts) == 0: + raise error_protocol.BadRequestError(error_msg) + + # Case 2. The prompt is a list of token ids. + if isinstance(input_prompts[0], int): + assert isinstance(input_prompts, list) + if not all(isinstance(token_id, int) for token_id in input_prompts): + raise error_protocol.BadRequestError(error_msg) + return [input_prompts] + + # Case 3. A list of prompts. + output_prompts: List[Union[List[int], data.ImageData]] = [] # noqa: UP006 + for input_prompt in input_prompts: + if isinstance(input_prompt, str): + output_prompts.append(ftokenize(input_prompt)) + elif isinstance(input_prompt, list) and all( + isinstance(token_id, int) for token_id in input_prompt + ): + output_prompts.append(input_prompt) + elif isinstance(input_prompt, data.ImageData): + output_prompts.append(input_prompt) + else: + raise error_protocol.BadRequestError(error_msg) + return output_prompts + + +def convert_prompts_to_data( + prompts: Union[str, List[int], List[Union[str, List[int], data.Data]]], # noqa: UP006 +) -> List[data.Data]: # noqa: UP006 + """Convert the given prompts in the combination of token id lists + and/or data to all data.""" + if isinstance(prompts, data.Data): + return [prompts] + if isinstance(prompts, str): + return [data.TextData(prompts)] + if isinstance(prompts[0], int): + assert isinstance(prompts, list) and all(isinstance(token_id, int) for token_id in prompts) + return [data.TokenData(prompts)] + return [convert_prompts_to_data(x)[0] for x in prompts] + + +class ErrorCleanupScope: + """Scope to call cleanup when an error is thrown. + + This class provides an important pattern properly cleanup + when async scope CancelledError or other exception happens. + + Parameters + ---------- + cleanup : Callable + A callable function to trigger at scope exit during an exception. + + Note + ---- + This helper is motivated by the need to properly + abort an async generator and trigger corresponding + cleanup functions. Naively use the try except + pattern will results in bug when we chain up + async generators. + + .. code:: python + + class EngineNotSafe: + async def _inner_gen(self, request): + request_id = self.get_request_id() + self.add_request(request) + try: + async for res in await producer_stream: + yield res + except asyncio.CancelledError: + self.abort(request_id) + + async def generate(self, request): + async for res in await self._inner_gen(request): + # async error can he raised in here + # this will cause + res = await process(res) + yield res + + The above except pattern is not safe. + This is because CancelledError may also be raised + outside _inner_gen during the process of generate + function in between iterations. + + Instead, we use ErrorCleanupScope to safeguard the + generation process. The scope will always properly + cleanup in exit function when the exception is raised + + .. code:: python + + class EngineSafe: + async def _inner_gen(self, request): + request_id = self.get_request_id() + self.add_request(request) + with ErrorCleanupScope(lambda: self.abort(request_id)) + async for res in await producer_stream: + yield res + + async def generate(self, request): + async for res in await self._inner_gen(request): + # even if async error is raised here + # it will cleanup the ErrorCleanupScope + # properly during function exit + res = await process(res) + yield res + """ + + cleanup: Callable + + def __init__(self, cleanup: Callable): + self.cleanup = cleanup + + def __enter__(self): + pass + + def __exit__(self, exc_type, exc_value, traceback) -> None: + # only cleanup when exc type is not none + if exc_type is not None: + self.cleanup() + + +# ====== Embedding Engine Utilities ====== + + +def load_embedding_params(model_weight_path, device, model_metadata) -> list: + """Load embedding model parameters from weight directory. + + Parameters + ---------- + model_weight_path : str + Path to the model weight directory. + device : tvm.runtime.Device + The target device. + model_metadata : dict + The model metadata dictionary containing param info. + + Returns + ------- + params : list + List of tvm.runtime.Tensor parameters in metadata order. + """ + from tvm.contrib import tvmjs + + params, meta = tvmjs.load_tensor_cache(model_weight_path, device) + param_names = [param["name"] for param in model_metadata["params"]] + assert len(param_names) == meta["ParamSize"] + return [params[name] for name in param_names] + + +def get_embedding_metadata(config: Dict[str, Any]) -> Optional[Dict[str, Any]]: # noqa: UP006 + """Read emedding metadata from mlc-chat-config or model lib metadata. + + Parameters + ---------- + config : Dict[str, Any] + The configuration dictionary containing model metadata. + + Returns + ------- + embedding_metadata : Optional[Dict[str, Any]] = None if it's not an embedding model. + The embedding metadata dictionary. + """ + if config.get("model_task") == "embedding": + return config.get("embedding_metadata") + return None + + +def detect_embedding_model_type(mod) -> Literal["encoder", "decoder"]: + """Detect embedding model type from compiled TVM module functions. + + Parameters + ---------- + mod : tvm.runtime.Module + The VM module with model functions. + + Returns + ------- + model_type : str + "encoder" for BERT-style models, "decoder" for Qwen3-Embeddings style. + """ + has_embed = mod.implements_function("embed") + has_prefill_to_hidden = mod.implements_function("prefill_to_last_hidden_states") + has_prefill = mod.implements_function("prefill") + + if has_embed and has_prefill_to_hidden: + return "decoder" + if has_prefill: + return "encoder" + raise ValueError( + "Model does not support embedding inference. " + "Expected 'embed' + 'prefill_to_last_hidden_states' (decoder) " + "or 'prefill' (encoder)." + ) diff --git a/python/mlc_llm/serve/entrypoints/__init__.py b/python/mlc_llm/serve/entrypoints/__init__.py new file mode 100644 index 0000000..3af7104 --- /dev/null +++ b/python/mlc_llm/serve/entrypoints/__init__.py @@ -0,0 +1,8 @@ +"""The entrypoints for MLC LLM server.""" + +from . import ( + debug_entrypoints, + metrics_entrypoints, + microserving_entrypoints, + openai_entrypoints, +) diff --git a/python/mlc_llm/serve/entrypoints/debug_entrypoints.py b/python/mlc_llm/serve/entrypoints/debug_entrypoints.py new file mode 100644 index 0000000..210c8eb --- /dev/null +++ b/python/mlc_llm/serve/entrypoints/debug_entrypoints.py @@ -0,0 +1,128 @@ +"""MLC LLM server debug entrypoints""" + +import json +from http import HTTPStatus + +import fastapi + +from mlc_llm.protocol import error_protocol +from mlc_llm.serve.server import ServerContext + +app = fastapi.APIRouter() + +################ /debug/dump_event_trace ################ + + +@app.post("/debug/dump_event_trace") +async def debug_dump_event_trace(request: fastapi.Request): + """Return the recorded events in Chrome Trace Event Format in JSON string. + The input request payload should have only one field, specifying the + model to query. For example: `{"model": "Llama-2-7b-chat-hf-q0f16"}`. + """ + # Get the raw request body as bytes + request_raw_data = await request.body() + request_json_str = request_raw_data.decode("utf-8") + try: + # Parse the JSON string + request_dict = json.loads(request_json_str) + except json.JSONDecodeError: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}" + ) + if "model" not in request_dict: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}" + ) + + # Check the requested model. + model = request_dict["model"] + + server_context: ServerContext = ServerContext.current() + async_engine = server_context.get_engine(model) + + if async_engine is None: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, + message=f'The requested model "{model}" is not served.', + ) + if async_engine.state.trace_recorder is None: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, + message=f'The requested model "{model}" does not enable tracing', + ) + + return json.loads(async_engine.state.trace_recorder.dump_json()) + + +################ /debug/cuda_profiler_start/end ################ + + +@app.post("/debug/cuda_profiler_start") +async def debug_cuda_profiler_start(_request: fastapi.Request): + """Start the cuda profiler for the engine. Only for debug purpose.""" + server_context: ServerContext = ServerContext.current() + # Since the CUDA profiler is process-wise, call the function for one model is sufficient. + for model in server_context.get_model_list(): + async_engine = server_context.get_engine(model) + async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_start") + break + + +@app.post("/debug/cuda_profiler_stop") +async def debug_cuda_profiler_stop(_request: fastapi.Request): + """Stop the cuda profiler for the engine. Only for debug purpose.""" + server_context: ServerContext = ServerContext.current() + # Since the CUDA profiler is process-wise, call the function for one model is sufficient. + for model in server_context.get_model_list(): + async_engine = server_context.get_engine(model) + async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_stop") + break + + +@app.post("/debug/dump_engine_metrics") +async def debug_dump_engine_metrics(request: fastapi.Request): + """Dump the engine metrics for the engine. Only for debug purpose.""" + # Get the raw request body as bytes + request_raw_data = await request.body() + request_json_str = request_raw_data.decode("utf-8") + try: + # Parse the JSON string + request_dict = json.loads(request_json_str) + except json.JSONDecodeError: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}" + ) + + # Check the requested model. + model = request_dict.get("model", None) + + server_context: ServerContext = ServerContext.current() + async_engine = server_context.get_engine(model) + res = await async_engine.metrics() + return res + + +@app.post("/debug/reset_engine") +async def debug_reset_engine_stats(request: fastapi.Request): + """Reset the engine, clean up all running data and metrics.""" + # Get the raw request body as bytes + request_raw_data = await request.body() + request_json_str = request_raw_data.decode("utf-8") + try: + # Parse the JSON string + request_dict = json.loads(request_json_str) + except json.JSONDecodeError: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}" + ) + if "model" not in request_dict: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}" + ) + + # Check the requested model. + model = request_dict["model"] + + server_context: ServerContext = ServerContext.current() + async_engine = server_context.get_engine(model) + async_engine.reset() diff --git a/python/mlc_llm/serve/entrypoints/metrics_entrypoints.py b/python/mlc_llm/serve/entrypoints/metrics_entrypoints.py new file mode 100644 index 0000000..71ee65d --- /dev/null +++ b/python/mlc_llm/serve/entrypoints/metrics_entrypoints.py @@ -0,0 +1,23 @@ +"""MLC LLM server metrics entrypoints""" + +import fastapi +from fastapi.responses import PlainTextResponse + +from mlc_llm.serve.server import ServerContext + +app = fastapi.APIRouter() + +################ /metrics ################ + + +@app.get("/metrics", response_class=PlainTextResponse) +async def metrics(_request: fastapi.Request): + """Start the cuda profiler for the engine. Only for debug purpose.""" + server_context: ServerContext = ServerContext.current() + # Use the metrics from first engine for now + # TODO(mlc-team): consider refactor server context to + # single engine since multiple AsyncMLCEngine do not work well with each other + # We need to work within the internal engine instead. + for model in server_context.get_model_list(): + async_engine = server_context.get_engine(model) + return (await async_engine.metrics()).prometheus_text() diff --git a/python/mlc_llm/serve/entrypoints/microserving_entrypoints.py b/python/mlc_llm/serve/entrypoints/microserving_entrypoints.py new file mode 100644 index 0000000..1371fdf --- /dev/null +++ b/python/mlc_llm/serve/entrypoints/microserving_entrypoints.py @@ -0,0 +1,73 @@ +"""MicroServing server entrypoints in MLC LLM""" + +import fastapi + +from mlc_llm.protocol.debug_protocol import DisaggConfig +from mlc_llm.protocol.microserving_protocol import ( + PrepRecvRequest, + PrepRecvResponse, + RemoteSendRequest, + StartGenerateRequest, +) +from mlc_llm.protocol.openai_api_protocol import StreamOptions + +from .openai_entrypoints import request_completion + +app = fastapi.APIRouter() + + +################ MicroServing Endpoints ################ + + +@app.post("/microserving/prep_recv") +async def prep_recv(request: PrepRecvRequest, raw_request: fastapi.Request) -> PrepRecvResponse: + """Handle the microserving request for receive preparation. + Match the prompt in the prefix cache (when enabled), + allocate entries in the KV cache to prepare receiving the KV data of the prompt. + Return the matched prefix length and the allocated KV entry metadata. + """ + request.debug_config.disagg_config = DisaggConfig( + kind="prepare_receive", + kv_window_begin=0, # always zero for prepare_receive + kv_window_end=request.end, + ) + request.stream_options = StreamOptions(include_usage=True) + request.stream = False + + response = await request_completion(request=request, raw_request=raw_request) + assert response.usage is not None + assert response.usage.extra is not None + assert "prefix_matched_length" in response.usage.extra + assert "kv_append_metadata" in response.usage.extra + return PrepRecvResponse( + prefix_matched_length=response.usage.extra["prefix_matched_length"], + kv_append_metadata=response.usage.extra["kv_append_metadata"], + ) + + +@app.post("/microserving/remote_send") +async def remote_send(request: RemoteSendRequest, raw_request: fastapi.Request): + """Compute and generate the KV data of the prompt in the specified KV window. + Send the KV data to the destination server.""" + request.debug_config.disagg_config = DisaggConfig( + kind="remote_send", + kv_window_begin=request.begin, + kv_window_end=request.end, + kv_append_metadata=request.kv_addr_info, + dst_group_offset=request.recv_rank, + ) + request.stream_options = StreamOptions(include_usage=True) + request.stream = False + + await request_completion(request=request, raw_request=raw_request) + return {} + + +@app.post("/microserving/start_generate") +async def start_generate(request: StartGenerateRequest, raw_request: fastapi.Request): + """Prefill the prompt in the specified KV window, and start decode.""" + request.debug_config.disagg_config = DisaggConfig( + kind="start_generation", + kv_window_begin=request.begin, + ) + return await request_completion(request=request, raw_request=raw_request) diff --git a/python/mlc_llm/serve/entrypoints/openai_entrypoints.py b/python/mlc_llm/serve/entrypoints/openai_entrypoints.py new file mode 100644 index 0000000..639e8a6 --- /dev/null +++ b/python/mlc_llm/serve/entrypoints/openai_entrypoints.py @@ -0,0 +1,340 @@ +"""OpenAI API-compatible server entrypoints in MLC LLM""" + +import base64 +import struct +from collections.abc import AsyncGenerator +from http import HTTPStatus +from typing import List, Optional # noqa: UP035 + +import fastapi +import numpy as np + +from mlc_llm.protocol import error_protocol +from mlc_llm.protocol.openai_api_protocol import ( + ChatCompletionRequest, + CompletionLogProbs, + CompletionRequest, + EmbeddingObject, + EmbeddingRequest, + EmbeddingResponse, + EmbeddingUsage, + ListResponse, + LogProbsContent, + ModelResponse, +) +from mlc_llm.serve import engine_base, engine_utils +from mlc_llm.serve.server import ServerContext + + +def verify_api_key(request: fastapi.Request): + """Function to verify API key""" + server_context = ServerContext.current() + # Only perform verification when API key is configured + if server_context is not None and server_context.api_key is not None: + provided_key = request.headers.get("Authorization", "").replace("Bearer ", "") + if provided_key != server_context.api_key: + raise fastapi.HTTPException(status_code=401, detail="Invalid API Key") + + +app = fastapi.APIRouter(dependencies=[fastapi.Depends(verify_api_key)]) + + +################ v1/embeddings ################ + + +@app.post("/v1/embeddings") +async def request_embedding(request: EmbeddingRequest): + """OpenAI-compatible embedding API. + API reference: https://platform.openai.com/docs/api-reference/embeddings/create + """ + server_context: ServerContext = ServerContext.current() + embedding_engine = server_context.get_embedding_engine(request.model) + if embedding_engine is None: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, + message=f'The requested model "{request.model}" is not served as an embedding model.', + ) + + # Normalize input to List[str] + inputs: List[str] # noqa: UP006 + if isinstance(request.input, str): + inputs = [request.input] + elif ( + isinstance(request.input, list) + and len(request.input) > 0 + and isinstance(request.input[0], str) + ): + inputs = list(request.input) + else: + # Token ID inputs (List[int] or List[List[int]]) — decode back to strings + if isinstance(request.input[0], int): + inputs = [embedding_engine.tokenizer.decode(request.input)] + else: + inputs = [embedding_engine.tokenizer.decode(ids) for ids in request.input] + + # Run embedding inference (async — does not block the event loop) + try: + embeddings, total_tokens = await embedding_engine.async_embed(inputs) + except Exception as exc: + return error_protocol.create_error_response( + HTTPStatus.INTERNAL_SERVER_ERROR, + message=f"Embedding inference failed: {exc}", + ) + + # Optional: truncate dimensions (Matryoshka-style). + # This is API-level renormalization after dimension truncation, + # independent of model metadata normalize. Always renormalize + # truncated vectors to maintain unit length per OpenAI API contract. + if request.dimensions is not None: + for i, emb in enumerate(embeddings): + vec = np.array(emb[: request.dimensions], dtype=np.float32) + norm = np.linalg.norm(vec) + if norm > 1e-12: + vec = vec / norm + embeddings[i] = vec.tolist() + + # Build response data + resp_data = [] + for i, emb in enumerate(embeddings): + if request.encoding_format == "base64": + binary = struct.pack(f"<{len(emb)}f", *emb) + resp_data.append( + EmbeddingObject( + embedding=base64.b64encode(binary).decode("utf-8"), + index=i, + ) + ) + else: + resp_data.append(EmbeddingObject(embedding=emb, index=i)) + + return EmbeddingResponse( + data=resp_data, + model=request.model, + usage=EmbeddingUsage(prompt_tokens=total_tokens, total_tokens=total_tokens), + ) + + +################ v1/models ################ + + +@app.get("/v1/models") +async def request_models() -> ListResponse: + """OpenAI-compatible served model query API. + API reference: https://platform.openai.com/docs/api-reference/models + """ + server_context: ServerContext = ServerContext.current() + return ListResponse(data=[ModelResponse(id=model) for model in server_context.get_model_list()]) + + +################ v1/completions ################ + + +@app.post("/v1/completions") +async def request_completion(request: CompletionRequest, raw_request: fastapi.Request): + """OpenAI-compatible completion API. + API reference: https://platform.openai.com/docs/api-reference/completions/create + """ + # - Check the requested model. + server_context: ServerContext = ServerContext.current() + request_final_usage_include_extra = server_context.enable_debug + request_include_debug_config = server_context.enable_debug + + if not request_include_debug_config: + request.debug_config = None + + async_engine = server_context.get_engine(request.model) + if async_engine is None: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, + message=f'The requested model "{request.model}" is not served.', + ) + # FIXME: This is a temporary solution to make sure + # prep_recv, remote_send and start_generation process the same request + request_id = request.user if request.user is not None else f"cmpl-{engine_utils.random_uuid()}" + + # Streaming response. + if request.stream: + # We manually get the first response from generator to + # capture potential exceptions in this scope, rather then + # the StreamingResponse scope. + stream_generator = async_engine._handle_completion( + request, + request_id, + request_final_usage_include_extra=request_final_usage_include_extra, + ) + first_response = await anext( # noqa: F821 + stream_generator + ) + + async def completion_stream_generator() -> AsyncGenerator[str, None]: + if isinstance(first_response, StopAsyncIteration): + yield "data: [DONE]\n\n" + return + yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n" + async for response in stream_generator: + yield f"data: {response.model_dump_json(by_alias=True)}\n\n" + yield "data: [DONE]\n\n" + + return fastapi.responses.StreamingResponse( + completion_stream_generator(), media_type="text/event-stream" + ) + + # Normal response. + request_final_usage = None + output_texts = [""] * request.n + finish_reasons: List[Optional[str]] = [None] * request.n # noqa: UP006 + logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n # noqa: UP006 + + async for response in async_engine._handle_completion( + request, + request_id, + request_final_usage_include_extra=request_final_usage_include_extra, + ): + if await raw_request.is_disconnected(): + # In non-streaming cases, the engine will not be notified + # when the request is disconnected. + # Therefore, we check if it is disconnected each time, + # and explicitly return. + # Note that requesta abort is triggered when the async for and funciton scope ends. + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message="The request has disconnected" + ) + # this is the final chunk + if response.usage is not None: + request_final_usage = response.usage + # remove extra information if debug is not enabled + if not server_context.enable_debug: + request_final_usage.extra = None + continue + for choice in response.choices: + output_texts[choice.index] += choice.text + if choice.finish_reason is not None and finish_reasons[choice.index] is None: + finish_reasons[choice.index] = choice.finish_reason + if choice.logprobs is not None: + if logprob_results[choice.index] is None: + logprob_results[choice.index] = choice.logprobs + else: + logprob_results[choice.index].token_logprobs.extend( + choice.logprobs.token_logprobs + ) + logprob_results[choice.index].tokens.extend(choice.logprobs.tokens) + logprob_results[choice.index].top_logprobs.extend(choice.logprobs.top_logprobs) + + return engine_base.wrap_completion_response( + request_id=request_id, + model=request.model, + output_texts=output_texts, + finish_reasons=finish_reasons, + logprob_results=logprob_results, + usage=request_final_usage, + ) + + +################ v1/chat/completions ################ + + +@app.post("/v1/chat/completions") +async def request_chat_completion(request: ChatCompletionRequest, raw_request: fastapi.Request): + """OpenAI-compatible chat completion API. + API reference: https://platform.openai.com/docs/api-reference/chat + """ + # - Check the requested model. + server_context: ServerContext = ServerContext.current() + request_final_usage_include_extra = server_context.enable_debug + request_include_debug_config = server_context.enable_debug + + if not request_include_debug_config: + request.debug_config = None + + async_engine = server_context.get_engine(request.model) + if async_engine is None: + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, + message=f'The requested model "{request.model}" is not served.', + ) + # FIXME: This is a temporary solution to make sure + # prep_recv, remote_send and start_generation process the same request + request_id = ( + request.user if request.user is not None else f"chatcmpl-{engine_utils.random_uuid()}" + ) + + # Streaming response. + if request.stream: + # We manually get the first response from generator to + # capture potential exceptions in this scope, rather then + # the StreamingResponse scope. + stream_generator = async_engine._handle_chat_completion( + request, + request_id, + request_final_usage_include_extra=request_final_usage_include_extra, + ) + first_response = await anext( # noqa: F821 + stream_generator + ) + + async def completion_stream_generator() -> AsyncGenerator[str, None]: + if isinstance(first_response, StopAsyncIteration): + yield "data: [DONE]\n\n" + return + yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n" + async for response in stream_generator: + yield f"data: {response.model_dump_json(by_alias=True)}\n\n" + yield "data: [DONE]\n\n" + + return fastapi.responses.StreamingResponse( + completion_stream_generator(), media_type="text/event-stream" + ) + + # Normal response. + request_final_usage = None + output_texts = ["" for _ in range(request.n)] + finish_reasons: List[Optional[str]] = [None for _ in range(request.n)] # noqa: UP006 + logprob_results: Optional[List[List[LogProbsContent]]] = ( # noqa: UP006 + [[] for _ in range(request.n)] if request.logprobs else None + ) + + async for response in async_engine._handle_chat_completion( + request, + request_id, + request_final_usage_include_extra=request_final_usage_include_extra, + ): + if await raw_request.is_disconnected(): + # In non-streaming cases, the engine will not be notified + # when the request is disconnected. + # Therefore, we check if it is disconnected each time, + # no need to explicitly abort, as the chat completion + # return will trigger abort call + return error_protocol.create_error_response( + HTTPStatus.BAD_REQUEST, message="The request has disconnected" + ) + # usage is always the last chunk + if response.usage is not None: + request_final_usage = response.usage + # remove extra information if debug is not enabled + if not server_context.enable_debug: + request_final_usage.extra = None + + for choice in response.choices: + assert isinstance(choice.delta.content, str) + output_texts[choice.index] += choice.delta.content + if choice.finish_reason is not None and finish_reasons[choice.index] is None: + finish_reasons[choice.index] = choice.finish_reason + if choice.logprobs is not None: + assert logprob_results is not None + logprob_results[choice.index] += choice.logprobs.content + + assert all(finish_reason is not None for finish_reason in finish_reasons) + use_function_calling, tool_calls_list = engine_base.process_function_call_output( + output_texts, finish_reasons + ) + + return engine_base.wrap_chat_completion_response( + request_id=request_id, + model=request.model, + output_texts=output_texts, + finish_reasons=finish_reasons, + tool_calls_list=tool_calls_list, + logprob_results=logprob_results, + use_function_calling=use_function_calling, + usage=request_final_usage, + ) diff --git a/python/mlc_llm/serve/event_trace_recorder.py b/python/mlc_llm/serve/event_trace_recorder.py new file mode 100644 index 0000000..416e905 --- /dev/null +++ b/python/mlc_llm/serve/event_trace_recorder.py @@ -0,0 +1,37 @@ +"""The event trace recorder in MLC LLM serving""" + +import tvm_ffi +from tvm.runtime import Object + +from . import _ffi_api + + +@tvm_ffi.register_object("mlc.serve.EventTraceRecorder") +class EventTraceRecorder(Object): + """The event trace recorder for requests.""" + + def __init__(self) -> None: + """Initialize a trace recorder.""" + self.__init_handle_by_constructor__(_ffi_api.EventTraceRecorder) + + def add_event(self, request_id: str, event: str) -> None: + """Record a event for the input request in the trace recorder. + + Parameters + ---------- + request_id : str + The subject request of the event. + + event : str + The event in a string name. + It can have one of the following patterns: + - "start xxx", which marks the start of event "xxx", + - "finish xxx", which marks the finish of event "xxx", + - "yyy", which marks the instant event "yyy". + The "starts" and "finishes" will be automatically paired in the trace recorder. + """ + return _ffi_api.EventTraceRecorderAddEvent(self, request_id, event) + + def dump_json(self) -> str: + """Dump the logged events in Chrome Trace Event Format in JSON string.""" + return _ffi_api.EventTraceRecorderDumpJSON(self) diff --git a/python/mlc_llm/serve/radix_tree.py b/python/mlc_llm/serve/radix_tree.py new file mode 100644 index 0000000..cccd3fe --- /dev/null +++ b/python/mlc_llm/serve/radix_tree.py @@ -0,0 +1,154 @@ +"""The Paged Radix Tree class.""" + +from typing import List, Tuple, Union # noqa: UP035 + +import tvm_ffi +from tvm.runtime import Object +from tvm_ffi import Shape + +from . import _ffi_api + + +@tvm_ffi.register_object("mlc.serve.PagedRadixTree") +class PagedRadixTree(Object): + """The paged radix tree to manage prefix and sequence.""" + + def __init__(self): + """ + Constructor of paged radix tree. + """ + self.__init_handle_by_constructor__(_ffi_api.PagedRadixTree) + + def match(self, tokens: Union[Shape, List, Tuple]) -> Tuple[int, Shape]: # noqa: UP006 + """ + Get all sequences with longest common prefix with given prefix tokens. + + Parameters + ---------- + tokens : Union[Shape, List, Tuple] + The prefix tokens for reference. + + Returns + ------ + matched_offset : int + The matched prefix length. + seq_ids : Shape + The array of matched sequence indice. + """ + if isinstance(tokens, (list, tuple)): + tokens = Shape(tokens) + output = _ffi_api.PagedRadixTreeMatchPrefix(self, tokens) + if len(output) == 1: + return output[0], [] + return output[0], output[1:] + + def add(self, seq_id: int) -> None: + """ + Add an empty sequence. + + Parameters + ---------- + seq_id : int + The sequence ID for index. + """ + _ffi_api.PagedRadixTreeAddSequence(self, seq_id) + + def remove(self, seq_id: int) -> None: + """ + Remove a sequence. + + Parameters + ---------- + seq_id : int + The sequence ID to remove. + """ + _ffi_api.PagedRadixTreeRemoveSequence(self, seq_id) + + def extend(self, seq_id: int, tokens: Union[Shape, List, Tuple]) -> None: # noqa: UP006 + """ + Extend a sequence with given tokens. + + Parameters + ---------- + seq_id : int + The sequence ID for index. + tokens : Union[Shape, List, Tuple] + The given tokens to extend. + """ + if isinstance(tokens, (list, tuple)): + tokens = Shape(tokens) + _ffi_api.PagedRadixTreeExtendSequence(self, seq_id, tokens) + + def rollback(self, seq_id: int, num_tokens: int) -> None: + """ + Roll back a sequence by number of tokens. + + Parameters + ---------- + seq_id : int + The sequence ID for index. + num_tokens : int + The number of tokens to be rolled back. + """ + _ffi_api.PagedRadixTreeRollBackSequence(self, seq_id, num_tokens) + + def fork(self, seq_id: int, parent_seq_id: int, forked_offset: int) -> None: + """ + Fork a sequence from parent sequence at given position. + + Parameters + ---------- + seq_id : int + The new sequence ID. + parent_seq_id : int + The parent sequence ID to fork from. + forked_offset : int + The position of parent sequence to fork at. + The valid value is [1, length of forked sequence]. + If the position equals the length of forked sequence, + the new sequence will copy the entire forked sequence. + """ + _ffi_api.PagedRadixTreeForkSequence(self, seq_id, parent_seq_id, forked_offset) + + def get(self, seq_id: int) -> Shape: + """ + Get a sequence's all tokens. + + Parameters + ---------- + seq_id : int + The sequence ID for index. + + Returns + ------ + tokens : Shape + The sequence tokens. + """ + return _ffi_api.PagedRadixTreeGetSequence(self, seq_id) + + def get_length(self, seq_id: int) -> int: + """ + Get a sequence's length. + + Parameters + ---------- + seq_id : int + The sequence ID for index. + + Returns + ------ + length : int + The sequence length. + """ + return _ffi_api.PagedRadixTreeGetSequenceLength(self, seq_id) + + def free_capacity(self) -> int: + """ + Get the remaining token capacity of the paged radix tree. + + Returns + ------ + capacity : int + The remaining token capacity of the paged radix tree. + """ + return _ffi_api.PagedRadixTreeFreeCapacity(self) diff --git a/python/mlc_llm/serve/request.py b/python/mlc_llm/serve/request.py new file mode 100644 index 0000000..ceae4ee --- /dev/null +++ b/python/mlc_llm/serve/request.py @@ -0,0 +1,34 @@ +"""The request class in MLC LLM serving""" + +from typing import List # noqa: UP035 + +import tvm_ffi +from tvm.runtime import Object + +from mlc_llm.protocol.generation_config import GenerationConfig + +from . import _ffi_api +from .data import Data + + +@tvm_ffi.register_object("mlc.serve.Request") +class Request(Object): + """The user submitted text-generation request, which contains + a unique request id, a list of multi-modal inputs, a set of + generation configuration parameters. + + Note + ---- + Do not explicitly construct this class. + Construct this object via engine.create_request functions. + """ + + @property + def inputs(self) -> List[Data]: # noqa: UP006 + """The inputs of the request.""" + return _ffi_api.RequestGetInputs(self) + + @property + def generation_config(self) -> GenerationConfig: + """The generation config of the request.""" + return GenerationConfig.model_validate_json(_ffi_api.RequestGetGenerationConfigJSON(self)) diff --git a/python/mlc_llm/serve/server/__init__.py b/python/mlc_llm/serve/server/__init__.py new file mode 100644 index 0000000..3f12704 --- /dev/null +++ b/python/mlc_llm/serve/server/__init__.py @@ -0,0 +1,4 @@ +"""The server related data structure and tools in MLC LLM serve.""" + +from .popen_server import PopenServer +from .server_context import ServerContext diff --git a/python/mlc_llm/serve/server/popen_server.py b/python/mlc_llm/serve/server/popen_server.py new file mode 100644 index 0000000..7d72022 --- /dev/null +++ b/python/mlc_llm/serve/server/popen_server.py @@ -0,0 +1,200 @@ +"""The MLC LLM server launched in a subprocess.""" + +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Literal, Optional, Union + +import psutil +import requests +from tvm.runtime import Device + +from mlc_llm.serve.config import EngineConfig +from mlc_llm.serve.engine_base import _check_engine_config + + +class PopenServer: + """The wrapper of MLC LLM server, which runs the server in + a background subprocess. + + This server can be used for debugging purposes. + """ + + def __init__( + self, + model: str, + device: Union[str, Device] = "auto", + *, + model_lib: Optional[str] = None, + mode: Literal["local", "interactive", "server"] = "local", + engine_config: Optional[EngineConfig] = None, + enable_debug: bool = True, + enable_tracing: bool = False, + host: str = "127.0.0.1", + port: int = 8082, + ) -> None: + """Please check out `python/mlc_llm/cli/serve.py` for the server arguments.""" + # - Check the fields fields of `engine_config`. + if engine_config is None: + engine_config = EngineConfig() + _check_engine_config(model, model_lib, mode, engine_config) + + self.model = model + self.model_lib = model_lib + self.device = device + self.mode = mode + self.enable_debug = enable_debug + self.engine_config = engine_config + self.enable_tracing = enable_tracing + self.enable_debug = enable_debug + self.host = host + self.port = port + self._proc: Optional[subprocess.Popen] = None + + self.base_url = "" + self.openai_v1_base_url = "" + + def start(self, extra_env=None) -> None: + """Launch the server in a popen subprocess. + Wait until the server becomes ready before return. + """ + extra_env = extra_env or {} + cmd = [sys.executable] + cmd += ["-m", "mlc_llm", "serve", self.model] + if self.model_lib is not None: + cmd += ["--model-lib", self.model_lib] + cmd += ["--device", self.device] + + if self.enable_debug: + cmd += ["--enable-debug"] + + if self.mode is not None: + cmd += ["--mode", self.mode] + + if len(self.engine_config.additional_models) > 0: + args_additional_model = [] + for additional_model in self.engine_config.additional_models: + if isinstance(additional_model, str): + args_additional_model.append(additional_model) + else: + args_additional_model.append(additional_model[0] + "," + additional_model[1]) + cmd += ["--additional-models", *args_additional_model] + cmd += ["--speculative-mode", self.engine_config.speculative_mode] + cmd += ["--prefix-cache-mode", self.engine_config.prefix_cache_mode] + + args_overrides = [] + if self.engine_config.max_num_sequence is not None: + args_overrides.append(f"max_num_sequence={self.engine_config.max_num_sequence}") + if self.engine_config.max_total_sequence_length is not None: + args_overrides.append( + f"max_total_seq_length={self.engine_config.max_total_sequence_length}" + ) + if self.engine_config.prefill_chunk_size is not None: + args_overrides.append(f"prefill_chunk_size={self.engine_config.prefill_chunk_size}") + if self.engine_config.max_history_size is not None: + args_overrides.append(f"max_history_size={self.engine_config.max_history_size}") + if self.engine_config.gpu_memory_utilization is not None: + args_overrides.append( + f"gpu_memory_utilization={self.engine_config.gpu_memory_utilization}" + ) + if self.engine_config.spec_draft_length is not None: + args_overrides.append(f"spec_draft_length={self.engine_config.spec_draft_length}") + if self.engine_config.prefix_cache_max_num_recycling_seqs is not None: + args_overrides.append( + "prefix_cache_max_num_recycling_seqs=" + + str(self.engine_config.prefix_cache_max_num_recycling_seqs) + ) + if len(args_overrides) > 0: + cmd += ["--overrides", ";".join(args_overrides)] + + if self.enable_tracing: + cmd += ["--enable-tracing"] + if self.enable_debug: + cmd += ["--enable-debug"] + + cmd += ["--host", self.host] + cmd += ["--port", str(self.port)] + process_path = str(Path(__file__).resolve().parents[4]) + final_env = os.environ.copy() + for key, value in extra_env.items(): + final_env[key] = value + self._proc = subprocess.Popen(cmd, cwd=process_path, env=final_env) + # NOTE: DO NOT USE `stdout=subprocess.PIPE, stderr=subprocess.PIPE` + # in subprocess.Popen here. PIPE has a fixed-size buffer with may block + # and hang forever. + + # Try to query the server until it is ready. + self.base_url = f"http://{self.host}:{str(self.port)}" + self.openai_v1_base_url = f"http://{self.host}:{str(self.port)}/v1" + openai_v1_models_url = f"{self.base_url}/v1/models" + + query_result = None + timeout = 120 + attempts = 0.0 + while query_result is None and attempts < timeout: + try: + query_result = requests.get(openai_v1_models_url, timeout=60) + if query_result.status_code != 200: + query_result = None + attempts += 0.1 + time.sleep(0.1) + except Exception: + attempts += 0.1 + time.sleep(0.1) + + # Check if the subprocess terminates unexpectedly or + # the queries reach the timeout. + process_return_code = self._proc.poll() + if process_return_code is not None: + raise RuntimeError( + "The server fails to launch. " + f'Please check if "{self.model}" is a valid model compiled by MLC LLM.' + ) + if attempts == timeout: + self.terminate() + raise RuntimeError(f"The server fails to launch in {timeout} seconds.") + + def terminate(self) -> None: + """Terminate the server subprocess.""" + if self._proc is None: + return + + # Kill all the child processes. + def kill_child_processes(): + try: + parent = psutil.Process(self._proc.pid) + children = parent.children(recursive=True) + except psutil.NoSuchProcess: + return + + for process in children: + try: + process.kill() + except psutil.NoSuchProcess: + pass + + kill_child_processes() + + # Kill the process. + try: + self._proc.kill() + except OSError: + pass + + # Join the process to avoid zombies. + try: + self._proc.wait(timeout=10.0) + except subprocess.TimeoutExpired: + pass + self._proc = None + + def __enter__(self): + """Start the server.""" + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Terminate the server.""" + self.terminate() diff --git a/python/mlc_llm/serve/server/server_context.py b/python/mlc_llm/serve/server/server_context.py new file mode 100644 index 0000000..725527b --- /dev/null +++ b/python/mlc_llm/serve/server/server_context.py @@ -0,0 +1,72 @@ +"""Server context that shared by multiple entrypoint files.""" + +from typing import TYPE_CHECKING, Dict, List, Optional # noqa: UP035 + +from ..engine import AsyncMLCEngine + +if TYPE_CHECKING: + from ..embedding_engine import AsyncEmbeddingEngine + + +class ServerContext: + """The global server context, including the running models + and corresponding async engines. + """ + + server_context: Optional["ServerContext"] = None + enable_debug: bool = False + + def __init__(self) -> None: + self._models: Dict[str, AsyncMLCEngine] = {} # noqa: UP006 + self._embedding_engines: Dict[str, AsyncEmbeddingEngine] = {} # noqa: UP006 + self.api_key: Optional[str] = None + + def __enter__(self): + if ServerContext.server_context is not None: + raise RuntimeError("Server context already exists.") + ServerContext.server_context = self + return self + + def __exit__(self, exc_type, exc_value, traceback): + for model_engine in self._models.values(): + model_engine.terminate() + for emb_engine in self._embedding_engines.values(): + emb_engine.terminate() + self._models.clear() + self._embedding_engines.clear() + ServerContext.server_context = None + + @staticmethod + def current(): + """Returns the current ServerContext.""" + return ServerContext.server_context + + def add_model(self, hosted_model: str, engine: AsyncMLCEngine) -> None: + """Add a new model to the server context together with the engine.""" + if hosted_model in self._models: + raise RuntimeError(f"Model {hosted_model} already running.") + self._models[hosted_model] = engine + + def get_engine(self, model: Optional[str]) -> Optional[AsyncMLCEngine]: + """Get the async engine of the requested model, or the unique async engine + if only one engine is served.""" + if len(self._models) == 1: + return next(iter(self._models.values())) + return self._models.get(model, None) + + def get_model_list(self) -> List[str]: # noqa: UP006 + """Get the list of all models on serve, including embedding models.""" + return list(self._models.keys()) + list(self._embedding_engines.keys()) + + def add_embedding_engine(self, hosted_model: str, engine: "AsyncEmbeddingEngine") -> None: + """Add a new embedding model to the server context.""" + if hosted_model in self._embedding_engines: + raise RuntimeError(f"Embedding model {hosted_model} already running.") + self._embedding_engines[hosted_model] = engine + + def get_embedding_engine(self, model: Optional[str]) -> Optional["AsyncEmbeddingEngine"]: + """Get the embedding engine of the requested model, or the unique + embedding engine if only one is served.""" + if len(self._embedding_engines) == 1: + return next(iter(self._embedding_engines.values())) + return self._embedding_engines.get(model, None) diff --git a/python/mlc_llm/serve/sync_engine.py b/python/mlc_llm/serve/sync_engine.py new file mode 100644 index 0000000..e3f7fb9 --- /dev/null +++ b/python/mlc_llm/serve/sync_engine.py @@ -0,0 +1,363 @@ +"""The MLC LLM synchronized engine. + +NOTE: This engine defined in this file directly wraps the underlying +Engine implementation in C++, is not optimized by multi-threading and +does not offer standard OpenAI API interface. + +We do not expose it and use it by default. As of now it mainly serves +the test and debug purpose because of its simplicity. +""" + +import json +from collections.abc import Sequence +from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union # noqa: UP035 + +import tvm + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import data +from mlc_llm.serve.config import EngineConfig +from mlc_llm.serve.engine_base import ( + EngineMetrics, + _check_engine_config, + _parse_models, + _print_engine_mode_logging_msg, + _process_model_args, + detect_device, +) +from mlc_llm.serve.event_trace_recorder import EventTraceRecorder +from mlc_llm.serve.request import Request +from mlc_llm.support import logging +from mlc_llm.tokenizers import TextStreamer, Tokenizer + +logger = logging.getLogger(__name__) + + +def _create_tvm_module( + creator: str, + ffi_funcs: Sequence[str], + creator_args: Optional[List[Any]] = None, # noqa: UP006 +) -> Dict[str, Callable]: # noqa: UP006 + """Internal method to create a module.""" + if creator_args is None: + creator_args = [] + module = tvm.get_global_func(creator, allow_missing=False)(*creator_args) + return {key: module[key] for key in ffi_funcs} + + +class SyncMLCEngine: + """The Python interface of synchronize request serving engine for MLC LLM. + + The engine receives requests from the "add_request" method. For + an given request, the engine will keep generating new tokens for + the request until finish (under certain criterion). After finish, + the engine will return the generation result through the callback + function provided by the request. + + NOTE: This engine directly wraps the underlying Engine implementation + in C++, is not optimized by multi-threading and does not offer standard + OpenAI API interface. We do not expose it and use it by default. + As of now it mainly serves the test and debug purpose because of its + simplicity. + + Parameters + ---------- + engine_config : Optional[EngineConfig] + Additional configurable arguments of MLC engine. + See class "EngineConfig" for more detail. + + enable_tracing : bool + A boolean indicating if to enable event logging for requests. + + request_stream_callback : Optional[Callable[[str, data.TokenData, Optional[str]], None]] + The provided callback function to handle the generation + output. It has the signature of `(str, data.TokenData, bool) -> None`, + where + - the first string is the request id, + - the TokenData contains the generated **delta** token ids since + the last invocation of the callback on the specific request, + - the optional string value denotes the finish reason if the + generation of the request is finished, or None if it has not finished. + + The callback function is optional at construction, but it needs to + be set before the engine executing requests. This can be done via + the `set_request_stream_callback` method. Otherwise, the engine will raise + exception. + """ + + def __init__( + self, + model: str, + device: Union[str, tvm.runtime.Device] = "auto", + *, + model_lib: Optional[str] = None, + mode: Literal["local", "interactive", "server"] = "local", + engine_config: Optional[EngineConfig] = None, + enable_tracing: bool = False, + request_stream_callback: Optional[Callable[[List[data.RequestStreamOutput]], None]] = None, # noqa: UP006 + ): + # - Check the fields fields of `engine_config`. + if engine_config is None: + engine_config = EngineConfig() + _check_engine_config( + model, + model_lib, + mode, + engine_config, + ) + + # - Initialize model loading info. + models = _parse_models(model, model_lib, engine_config.additional_models) + if isinstance(device, str): + device = detect_device(device) + assert isinstance(device, tvm.runtime.Device) + ( + model_args, + model_config_paths, + self.conv_template, + ) = _process_model_args(models, device, engine_config) + + # - Load the raw model config into dict + self.model_config_dicts = [] + for i, model_info in enumerate(models): + model_info.model_lib = model_args[i][1] + with open(model_config_paths[i], encoding="utf-8") as file: + self.model_config_dicts.append(json.load(file)) + + # - Print logging info for regarding the mode selection. + if engine_config.verbose: + _print_engine_mode_logging_msg(mode) + + self._ffi = _create_tvm_module( + "mlc.serve.create_engine", + ffi_funcs=[ + "init", + "add_request", + "abort_request", + "step", + "reset", + "json_metrics", + "get_request_stream_callback", + "set_request_stream_callback", + "create_request", + ], + ) + self.trace_recorder = EventTraceRecorder() if enable_tracing else None + + engine_config.model = model_args[0][0] + engine_config.model_lib = model_args[0][1] + engine_config.additional_models = model_args[1:] + engine_config.mode = mode + self._ffi["init"]( + engine_config.asjson(), + device, + request_stream_callback, + self.trace_recorder, + ) + self.tokenizer = Tokenizer(model_args[0][0]) + + def generate( + self, + prompts: Union[str, List[str], List[int], List[List[int]], List[List[data.Data]]], # noqa: UP006 + generation_config: Union[GenerationConfig, List[GenerationConfig]], # noqa: UP006 + ) -> Tuple[List[List[str]], List[Optional[List[List[str]]]]]: # noqa: UP006 + """Generate texts for a list of input prompts. + Each prompt can be a string or a list of token ids. + The generation for each prompt is independent. + Return the generation results, one for each prompt. + + Parameters + ---------- + prompts : Union[str, List[str], List[int], List[List[int]]] + One or a list of input prompts for text generation. + Each prompt can be a string or a list of token ids. + + generation_config : Union[GenerationConfig, List[GenerationConfig]] + The generation config for each requests. + If the it is a single GenerationConfig instance, + this config will be shared by all the prompts. + Otherwise, one generation config is required for every + prompt. + + Returns + ------- + output_text : List[List[str]] + The text generation results, one list of strings for each input prompt. + The length of each list is the parallel generation `n` in + generation config. + + output_logprobs_str : List[Optional[List[List[str]]]] + The logprob strings of each token for each input prompt, or None + if an input prompt does not require logprobs. + """ + if isinstance(prompts, str): + # `prompts` is a single string. + prompts = [prompts] + else: + assert isinstance(prompts, list), ( + "Input `prompts` is expected to be a string, a list of " + "str, a list of token ids or multiple lists of token ids. " + ) + if len(prompts) == 0: + return [], [] + if isinstance(prompts[0], int): + # `prompts` is a list of token ids + prompts = [prompts] + + num_requests = len(prompts) + if not isinstance(generation_config, list): + generation_config = [generation_config] * num_requests + + assert len(generation_config) == num_requests, ( + "Number of generation config and number of prompts mismatch" + ) + + num_finished_generations = 0 + output_texts: List[List[str]] = [] # noqa: UP006 + output_logprobs_str: List[Optional[List[List[str]]]] = [] # noqa: UP006 + text_streamers: List[List[TextStreamer]] = [] # noqa: UP006 + for i in range(num_requests): + output_texts.append([]) + output_logprobs_str.append([] if generation_config[i].logprobs else None) + text_streamers.append([]) + for _ in range(generation_config[i].n): + output_texts[i].append("") + text_streamers[i].append(TextStreamer(self.tokenizer)) + if output_logprobs_str[i] is not None: + output_logprobs_str[i].append([]) + + num_total_generations = sum(cfg.n for cfg in generation_config) + + # Save a copy of the original function callback since `generate` + # overrides the callback function. + # The original callback will be set back later on. + original_callback = self._ffi["get_request_stream_callback"]() + + # Define the callback function for request generation results + def request_stream_callback(delta_outputs: List[data.RequestStreamOutput]): # noqa: UP006 + nonlocal num_finished_generations + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + rid = int(request_id) + + assert len(stream_outputs) == generation_config[rid].n + for i, (stream_output, text_streamer) in enumerate( + zip(stream_outputs, text_streamers[rid]) + ): + if output_logprobs_str[rid] is not None: + assert stream_output.delta_logprob_json_strs is not None + output_logprobs_str[rid][i] += stream_output.delta_logprob_json_strs + + delta_text = stream_output.extra_prefix_string + ( + text_streamer.put(stream_output.delta_token_ids) + if len(stream_output.delta_token_ids) > 0 + else "" + ) + if stream_output.finish_reason is not None: + delta_text += text_streamer.finish() + + output_texts[rid][i] += delta_text + if stream_output.finish_reason is not None: + num_finished_generations += 1 + + # Override the callback function in engine. + self._ffi["set_request_stream_callback"](request_stream_callback) + + def convert_to_data( + prompt: Union[str, List[int], List[data.Data]], # noqa: UP006 + ) -> List[data.Data]: # noqa: UP006 + if isinstance(prompt, str): + return [data.TextData(prompt)] + if isinstance(prompt[0], int): + return [data.TokenData(prompt)] + return prompt + + # Add requests to engine. + for req_id, (prompt, generation_cfg) in enumerate(zip(prompts, generation_config)): + input_data = convert_to_data(prompt) + self.add_request( + self.create_request( + request_id=str(req_id), + inputs=input_data, + generation_config=generation_cfg, + ) + ) + + while num_finished_generations != num_total_generations: + self.step() + + # Restore the callback function in engine. + self._ffi["set_request_stream_callback"](original_callback) + return output_texts, output_logprobs_str + + def create_request( + self, + request_id: str, + inputs: Union[data.Data, List[data.Data]], # noqa: UP006 + generation_config: GenerationConfig, + ): + """Create a new request that can be added to engine. + + Parameters + ---------- + request_id : str + The unique identifier of the request. + Different requests should have different ids. + + inputs : List[Data] + The user inputs of a request. Input may have multi-modality. + + generation_config : GenerationConfig + The generation configuration of the request. + + Note + ---- + engine may fill in default generation config of the model. + """ + if not isinstance(inputs, list): + inputs = [inputs] + return self._ffi["create_request"]( + request_id, inputs, generation_config.model_dump_json(by_alias=True) + ) + + def add_request(self, request: Request) -> None: + """Add a new request to the engine. + + Parameters + ---------- + request : Request + The request to add. + """ + self._ffi["add_request"](request) + + def abort_request(self, request_id: str) -> None: + """Abort the generation of the request corresponding to the input request id. + + Parameters + ---------- + request_id : str + The unique id of the request to abort. + """ + self._ffi["abort_request"](request_id) + + def step(self) -> None: + """The main function that the engine takes a step of action. + + At each step, the engine may decide to + - run prefill for one (or more) requests, + - run one-step decode for the all existing requests + ... + + In the end of certain actions (e.g., decode), the engine will + check if any request has finished, and will return the + generation results for those finished requests. + """ + self._ffi["step"]() + + def reset(self) -> None: + """Reset the engine, clean up all running data and metrics.""" + self._ffi["reset"]() + + def metrics(self) -> EngineMetrics: + """Reset the engine, clean up all running data and metrics.""" + return EngineMetrics(json.loads(self._ffi["json_metrics"]())) diff --git a/python/mlc_llm/support/__init__.py b/python/mlc_llm/support/__init__.py new file mode 100644 index 0000000..ca5d7a6 --- /dev/null +++ b/python/mlc_llm/support/__init__.py @@ -0,0 +1,4 @@ +""" +Common utilities used in the Python package. Do not import anything by default, +as they may introduce unnecessary dependencies. +""" diff --git a/python/mlc_llm/support/argparse.py b/python/mlc_llm/support/argparse.py new file mode 100644 index 0000000..6d36f43 --- /dev/null +++ b/python/mlc_llm/support/argparse.py @@ -0,0 +1,16 @@ +"""An enhanced argument parser for mlc-chat.""" + +import argparse +import sys + + +class ArgumentParser(argparse.ArgumentParser): + """An enhanced argument parser for mlc-chat.""" + + def error(self, message): + """Overrides the behavior when erroring out""" + print("-" * 25 + " Usage " + "-" * 25) + self.print_help() + print("-" * 25 + " Error " + "-" * 25) + print(message, file=sys.stderr) + sys.exit(2) diff --git a/python/mlc_llm/support/auto_config.py b/python/mlc_llm/support/auto_config.py new file mode 100644 index 0000000..1470b9d --- /dev/null +++ b/python/mlc_llm/support/auto_config.py @@ -0,0 +1,190 @@ +"""Help function for detecting the model configuration file `config.json`""" + +import json +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING + +from . import logging +from .style import bold, green + +if TYPE_CHECKING: + from mlc_llm.model import Model + from mlc_llm.quantization import Quantization + + +logger = logging.getLogger(__name__) + +FOUND = green("Found") + + +def detect_mlc_chat_config(mlc_chat_config: str) -> Path: + """Detect and return the path that points to mlc-chat-config.json. + If `mlc_chat_config` is a directory, it looks for mlc-chat-config.json below it. + + Parameters + --------- + mlc_chat_config : str + The path to `mlc-chat-config.json`, or the directory containing + `mlc-chat-config.json`. + + Returns + ------- + mlc_chat_config_json_path : pathlib.Path + The path points to mlc_chat_config.json. + """ + from mlc_llm.model import MODEL_PRESETS + + from .download_cache import download_and_cache_mlc_weights + + if mlc_chat_config.startswith("HF://") or mlc_chat_config.startswith("http"): + mlc_chat_config_path = Path(download_and_cache_mlc_weights(model_url=mlc_chat_config)) + elif isinstance(mlc_chat_config, str) and mlc_chat_config in MODEL_PRESETS: + logger.info("%s mlc preset model: %s", FOUND, mlc_chat_config) + content = MODEL_PRESETS[mlc_chat_config].copy() + content["model_preset_tag"] = mlc_chat_config + temp_file = tempfile.NamedTemporaryFile( + suffix=".json", + delete=False, + ) + logger.info("Dumping config to: %s", temp_file.name) + mlc_chat_config_path = Path(temp_file.name) + with mlc_chat_config_path.open("w", encoding="utf-8") as mlc_chat_config_file: + json.dump(content, mlc_chat_config_file, indent=2) + else: + mlc_chat_config_path = Path(mlc_chat_config) + if not mlc_chat_config_path.exists(): + raise ValueError(f"{mlc_chat_config_path} does not exist.") + + if mlc_chat_config_path.is_dir(): + # search mlc-chat-config.json under path + mlc_chat_config_json_path = mlc_chat_config_path / "mlc-chat-config.json" + if not mlc_chat_config_json_path.exists(): + raise ValueError(f"Fail to find mlc-chat-config.json under {mlc_chat_config_path}.") + else: + mlc_chat_config_json_path = mlc_chat_config_path + + logger.info("%s model configuration: %s", FOUND, mlc_chat_config_json_path) + return mlc_chat_config_json_path + + +def detect_config(config: str) -> Path: + """Detect and return the path that points to config.json. If `config` is a directory, + it looks for config.json below it. + + Parameters + --------- + config : str + The preset name of the model, or the path to `config.json`, or the directory containing + `config.json`. + + Returns + ------- + config_json_path : pathlib.Path + The path points to config.json. + """ + from mlc_llm.model import MODEL_PRESETS + + if isinstance(config, str) and config in MODEL_PRESETS: + logger.info("%s preset model: %s", FOUND, config) + content = MODEL_PRESETS[config].copy() + content["model_preset_tag"] = config + temp_file = tempfile.NamedTemporaryFile( + suffix=".json", + delete=False, + ) + logger.info("Dumping config to: %s", temp_file.name) + config_path = Path(temp_file.name) + with config_path.open("w", encoding="utf-8") as config_file: + json.dump(content, config_file, indent=2) + else: + config_path = Path(config) + if not config_path.exists(): + raise ValueError(f"{config_path} does not exist.") + + if config_path.is_dir(): + # search config.json under config path + config_json_path = config_path / "config.json" + if not config_json_path.exists(): + raise ValueError(f"Fail to find config.json under {config_path}.") + else: + config_json_path = config_path + + logger.info("%s model configuration: %s", FOUND, config_json_path) + return config_json_path + + +def detect_model_type(model_type: str, config: Path) -> "Model": + """Detect the model type from the configuration file. If `model_type` is "auto", it will be + inferred from the configuration file. Otherwise, it will be used as the model type, and sanity + check will be performed. + + Parameters + ---------- + model_type : str + The model type, for example, "llama". + + config : pathlib.Path + The path to config.json. + + Returns + ------- + model : mlc_llm.compiler.Model + The model type. + """ + + from mlc_llm.model import MODELS + + if model_type == "auto": + with open(config, encoding="utf-8") as config_file: + cfg = json.load(config_file) + if "model_type" not in cfg and ( + "model_config" not in cfg or "model_type" not in cfg["model_config"] + ): + raise ValueError( + f"'model_type' not found in: {config}. " + f"Please explicitly specify `--model-type` instead." + ) + model_type = cfg["model_type"] if "model_type" in cfg else cfg["model_config"]["model_type"] + if model_type in ["mixformer-sequential"]: + model_type = "phi-msft" + logger.info("%s model type: %s. Use `--model-type` to override.", FOUND, bold(model_type)) + if model_type not in MODELS: + raise ValueError(f"Unknown model type: {model_type}. Available ones: {list(MODELS.keys())}") + return MODELS[model_type] + + +def detect_quantization(quantization_arg: str, config: Path) -> "Quantization": + """Detect the model quantization scheme from the configuration file or `--quantization` + argument. If `--quantization` is provided, it will override the value on the configuration + file. + + Parameters + ---------- + quantization_arg : str + The quantization scheme, for example, "q4f16_1". + + config : pathlib.Path + The path to mlc-chat-config.json. + + Returns + ------- + quantization : mlc_llm.quantization.Quantization + The model quantization scheme. + """ + from mlc_llm.quantization import ( + QUANTIZATION, + ) + + with open(config, encoding="utf-8") as config_file: + cfg = json.load(config_file) + if quantization_arg is not None: + quantization = QUANTIZATION[quantization_arg] + elif "quantization" in cfg: + quantization = QUANTIZATION[cfg["quantization"]] + else: + raise ValueError( + f"'quantization' not found in: {config}. " + f"Please explicitly specify `--quantization` instead." + ) + return quantization diff --git a/python/mlc_llm/support/auto_device.py b/python/mlc_llm/support/auto_device.py new file mode 100644 index 0000000..fc811e8 --- /dev/null +++ b/python/mlc_llm/support/auto_device.py @@ -0,0 +1,93 @@ +"""Automatic detection of the device available on the local machine.""" + +import os +import subprocess +import sys +from typing import Dict, Optional # noqa: UP035 + +import tvm +from tvm.runtime import Device +from tvm_ffi import DLDeviceType + +from . import logging +from .style import bold, green, red + +FOUND = green("Found") +NOT_FOUND = red("Not found") +AUTO_DETECT_DEVICES = ["cuda", "rocm", "metal", "vulkan", "opencl", "cpu"] +_RESULT_CACHE: Dict[str, bool] = {} # noqa: UP006 + + +logger = logging.getLogger(__name__) + + +def detect_device(device_hint: str) -> Optional[Device]: + """Detect locally available device from string hint.""" + if device_hint == "auto": + device = None + for device_type in AUTO_DETECT_DEVICES: + cur_device = tvm.device(device_type=device_type, index=0) + if _device_exists(cur_device): + if device is None: + device = cur_device + if device is None: + logger.info("%s: No available device detected", NOT_FOUND) + return None + logger.info("Using device: %s", bold(device2str(device))) + return device + try: + device = tvm.device(device_hint) + except Exception as err: + raise ValueError(f"Invalid device name: {device_hint}") from err + if not _device_exists(device): + raise ValueError(f"Device is not found on your local environment: {device_hint}") + return device + + +def device2str(device: Device) -> str: + """Convert a TVM device object to string.""" + return f"{tvm.runtime.Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]}:{device.index}" + + +def _device_exists(device: Device) -> bool: + device_type = tvm.runtime.Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()] + device_str = device2str(device) + if device_str in _RESULT_CACHE: + return _RESULT_CACHE[device_str] + cmd = [ + sys.executable, + "-m", + "mlc_llm.cli.check_device", + device_type, + ] + prefix = "check_device:" + subproc_outputs = [ + line[len(prefix) :].strip() + for line in subprocess.run( + cmd, + capture_output=True, + text=True, + check=False, + env=os.environ, + ) + .stdout.strip() + .splitlines() + if line.startswith(prefix) + ] + if subproc_outputs: + if subproc_outputs[0]: + for i in subproc_outputs[0].split(","): + logger.info("%s device: %s:%s", FOUND, device_type, i) + _RESULT_CACHE[f"{device_type}:{i}"] = True + if device.dlpack_device_type() == DLDeviceType.kDLCPU: + break + else: + logger.error( + "GPU device detection failed. Please report this issue with the output of command: %s", + " ".join(cmd), + ) + if device_str in _RESULT_CACHE: + return _RESULT_CACHE[device_str] + logger.info("%s device: %s", NOT_FOUND, device_str) + _RESULT_CACHE[device_str] = False + return False diff --git a/python/mlc_llm/support/auto_target.py b/python/mlc_llm/support/auto_target.py new file mode 100644 index 0000000..55aca59 --- /dev/null +++ b/python/mlc_llm/support/auto_target.py @@ -0,0 +1,554 @@ +"""Helper functions for target auto-detection.""" + +import os +from pathlib import Path +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple # noqa: UP035 + +from tvm import IRModule, relax +from tvm.ir.transform import Pass +from tvm.support import ndk, tar, xcode +from tvm.target import Target +from tvm_ffi import get_global_func, register_global_func + +from . import logging +from .auto_device import AUTO_DETECT_DEVICES, detect_device, device2str +from .constants import MLC_MULTI_ARCH +from .style import bold, green, red + +if TYPE_CHECKING: + from mlc_llm.compiler.compile import CompileArgs + + +logger = logging.getLogger(__name__) + +# TODO: add help message on how to specify the target manually +HELP_MSG = """TBD""" +FOUND = green("Found") +NOT_FOUND = red("Not found") +BuildFunc = Callable[[IRModule, "CompileArgs", Pass], None] + + +def detect_target_and_host( + target_hint: str, + host_hint: str = "auto", + enable_subgroups: Optional[bool] = None, +) -> Tuple[Target, BuildFunc]: # noqa: UP006 + """Detect the configuration for the target device and its host, for example, target GPU and + the host CPU. + + Parameters + ---------- + target_hint : str + The hint for the target device. + + host_hint : str + The hint for the host CPU, default is "auto". + """ + target, build_func = _detect_target_gpu(target_hint) + if target.host is None: + target = Target(target, host=_detect_target_host(host_hint)) + target = _apply_webgpu_subgroups(target, enable_subgroups) + if target.kind.name == "cuda": + # Enable thrust for CUDA + target_dict = dict(target.export()) + target_dict["libs"] = ( + (target_dict["libs"] + ["thrust"]) if "libs" in target_dict else ["thrust"] + ) + target = Target(target_dict) + _register_cuda_hook(target) + elif target.kind.name == "rocm": + target_dict = dict(target.export()) + extra_libs = ["thrust", "rocblas", "miopen", "hipblas"] + target_dict["libs"] = ( + (target_dict["libs"] + extra_libs) if "libs" in target_dict else extra_libs + ) + target = Target(target_dict) + return target, build_func + + +def _apply_webgpu_subgroups(target: Target, enable_subgroups: Optional[bool]) -> Target: + if not enable_subgroups: + return target + if target.kind.name != "webgpu": + logger.warning( + "--enable-subgroups is only supported for WebGPU targets; ignoring for %s", + target.kind.name, + ) + return target + target_dict = dict(target.export()) + target_dict["supports_subgroups"] = True + return Target(target_dict) + + +def _detect_target_gpu(hint: str) -> Tuple[Target, BuildFunc]: # noqa: UP006 + if hint in ["iphone", "macabi", "android", "webgpu", "mali", "opencl"]: + hint += ":generic" + if hint == "auto" or hint in AUTO_DETECT_DEVICES: + target: Optional[Target] = None + device = detect_device(hint) + if device is not None: + device_str = device2str(device) + try: + target = Target.from_device(device) + except ValueError: + logger.info("%s: Cannot detect target from device: %s", NOT_FOUND, device_str) + if target is None: + raise ValueError(f"No target detected from device: {hint}. Please specify explicitly") + logger.info( + '%s configuration of target device "%s": %s', + FOUND, + bold(device_str), + target.export(), + ) + return target, _build_default() + if hint in PRESET: + preset = PRESET[hint] + target = Target(preset["target"]) + build = preset.get("build", _build_default) + return target, build() + if _is_device(hint): + logger.info("Detecting target device: %s", hint) + target = Target.from_device(hint) + logger.info("%s target: %s", FOUND, target.export()) + return target, _build_default() + try: + logger.info("Try creating device target from string: %s", hint) + target = Target(hint) + logger.info("%s target: %s", FOUND, target.export()) + return target, _build_default() + except Exception as err: + logger.info("%s: Failed to create target", NOT_FOUND) + raise ValueError(f"Invalid target: {hint}") from err + + +def _detect_target_host(hint: str) -> Target: + """Detect the host CPU architecture.""" + if hint == "auto": + target_triple = get_global_func("tvm.codegen.llvm.GetDefaultTargetTriple")() + target = Target.from_device("cpu") + logger.info("%s host LLVM triple: %s", FOUND, bold(target.attrs["mtriple"])) + logger.info("%s host LLVM CPU: %s", FOUND, bold(target.attrs["mcpu"])) + return target + target_triple = hint + logger.info("Using LLVM triple specified by --host: %s", bold(target_triple)) + return Target({"kind": "llvm", "mtriple": target_triple}) + + +def _is_device(device: str): + if " " in device: + return False + if device.count(":") != 1: + return False + return True + + +def _add_system_lib_prefix(mod: IRModule, prefix: str, is_system_lib: bool) -> IRModule: + if is_system_lib and prefix: + mod = mod.with_attrs({"system_lib_prefix": prefix}) + elif is_system_lib: + logger.warning( + "%s is not specified when building a static library", + bold("--system-lib-prefix"), + ) + elif prefix: + logger.warning( + "--system-lib-prefix is specified, but it will not take any effect " + "when building the shared library" + ) + return mod + + +def _build_metal_x86_64(): + def build(mod: IRModule, args: "CompileArgs", pipeline=None): + output = args.output + mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=False) + assert output.suffix == ".dylib" + relax.build( + mod, + target=args.target, + relax_pipeline=pipeline, + ).export_library( + str(output), + fcompile=xcode.create_dylib, + sdk="macosx", + arch="x86_64", + ) + + return build + + +def _build_iphone(): + @register_global_func("tvm_callback_metal_compile", override=True) + def compile_metal(src, target): + libs = target.attrs.get("libs", None) + if libs: + return xcode.compile_metal(src, sdk=libs[0]) + return xcode.compile_metal(src) + + def build(mod: IRModule, args: "CompileArgs", pipeline=None): + output = args.output + mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True) + assert output.suffix == ".tar" + relax.build( + mod, + target=args.target, + relax_pipeline=pipeline, + system_lib=True, + ).export_library( + str(output), + fcompile=tar.tar, + ) + + return build + + +def _build_android(): + def build(mod: IRModule, args: "CompileArgs", pipeline=None): + output = args.output + mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True) + assert output.suffix == ".tar" + ex = relax.build( + mod, + target=args.target, + relax_pipeline=pipeline, + system_lib=True, + ) + ex.export_library( + str(output), + fcompile=tar.tar, + ) + if args.debug_dump is not None: + source = ex.mod.imports[0].imports[0].inspect_source() + with open(args.debug_dump / "kernel.cl", "w", encoding="utf-8") as f: + f.write(source) + + return build + + +def _build_android_so(): + def build(mod: IRModule, args: "CompileArgs", pipeline=None): + output = args.output + mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=False) + assert output.suffix == ".so" + ex = relax.build( + mod, + target=args.target, + relax_pipeline=pipeline, + system_lib=False, + ) + ex.export_library( + str(output), + fcompile=ndk.create_shared, + ) + if args.debug_dump is not None: + source = ex.mod.imports[0].imports[0].inspect_source() + with open(args.debug_dump / "kernel.cl", "w", encoding="utf-8") as f: + f.write(source) + + return build + + +def _build_webgpu(): + def build(mod: IRModule, args: "CompileArgs", pipeline=None): + output = args.output + mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True) + assert output.suffix == ".wasm" + + # Try to locate `mlc_wasm_runtime.bc` + bc_path = None + bc_candidates = ["web/dist/wasm/mlc_wasm_runtime.bc"] + if os.environ.get("MLC_LLM_SOURCE_DIR", None): + mlc_source_home_dir = os.environ["MLC_LLM_SOURCE_DIR"] + bc_candidates.append( + os.path.join(mlc_source_home_dir, "web", "dist", "wasm", "mlc_wasm_runtime.bc") + ) + error_info = ( + "Cannot find library: mlc_wasm_runtime.bc\n" + + "Make sure you have run `./web/prep_emcc_deps.sh` and " + + "`export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm` so that we can locate the file. " + + "We tried to look at candidate paths:\n" + ) + for candidate in bc_candidates: + error_info += candidate + "\n" + if Path(candidate).exists(): + bc_path = candidate + if not bc_path: + raise RuntimeError(error_info) + + relax.build( + mod, + target=args.target, + relax_pipeline=pipeline, + system_lib=True, + ).export_library( + str(output), + libs=[bc_path], + ) + + return build + + +def _build_mali(): + def build(mod: IRModule, args: "CompileArgs", pipeline=None): + output = args.output + mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True) + assert output.suffix == ".so" + mod = relax.build( + mod, + target=args.target, + relax_pipeline=pipeline, + system_lib=True, + ) + if "TVM_NDK_CC" in os.environ: + mod.export_library(str(output), fcompile=ndk.create_shared) + else: + mod.export_library(str(output)) + + return build + + +def _build_default(): + def build(mod: IRModule, args: "CompileArgs", pipeline=None): + output = args.output + if output.suffix in [".tar", ".lib"]: + system_lib = True + elif output.suffix in [".so", ".dylib", ".dll"]: + system_lib = False + else: + logger.warning("Unknown output suffix: %s. Assuming shared library.", output.suffix) + system_lib = False + mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=system_lib) + relax.build( + mod, + target=args.target, + relax_pipeline=pipeline, + system_lib=system_lib, + ).export_library( + str(output), + ) + + return build + + +def detect_cuda_arch_list(target: Target) -> List[int]: # noqa: UP006 + """Detect the CUDA architecture list from the target.""" + + def convert_to_num(arch_str): + arch_num_str = "".join(filter(str.isdigit, arch_str)) + assert arch_num_str, f"'{arch_str}' does not contain any digits" + return int(arch_num_str) + + assert target.kind.name == "cuda", f"Expect target to be CUDA, but got {target}" + if MLC_MULTI_ARCH is not None: + multi_arch = [convert_to_num(x) for x in MLC_MULTI_ARCH.split(",")] + else: + assert target.attrs.get("arch", "").startswith("sm_") + multi_arch = [convert_to_num(target.attrs.get("arch")[3:])] + multi_arch = list(set(multi_arch)) + return multi_arch + + +def _register_cuda_hook(target: Target): + if MLC_MULTI_ARCH is None: + default_arch = target.attrs.get("arch", None) + logger.info("Generating code for CUDA architecture: %s", bold(default_arch)) + logger.info( + "To produce multi-arch fatbin, set environment variable %s. " + "Example: MLC_MULTI_ARCH=70,72,75,80,86,87,89,90a", + bold("MLC_MULTI_ARCH"), + ) + multi_arch = None + else: + logger.info("%s %s: %s", FOUND, bold("MLC_MULTI_ARCH"), MLC_MULTI_ARCH) + multi_arch = [x.strip() for x in MLC_MULTI_ARCH.split(",")] + logger.info("Generating code for CUDA architecture: %s", multi_arch) + + @register_global_func("tvm_callback_cuda_compile", override=True) + def tvm_callback_cuda_compile(code): + """use nvcc to generate fatbin code for better optimization""" + from tvm.support import nvcc + + if multi_arch is None: + ptx = nvcc.compile_cuda(code, target_format="fatbin", compiler="nvcc") + else: + arch = [] + for compute_version in multi_arch: + arch += [ + "-gencode", + f"arch=compute_{compute_version},code=sm_{compute_version}", + ] + ptx = nvcc.compile_cuda(code, target_format="fatbin", arch=arch, compiler="nvcc") + return ptx + + +def detect_system_lib_prefix( + target_hint: str, prefix_hint: str, model_name: str, quantization: str +) -> str: + """Detect the iOS / Android system lib prefix to identify the library needed to load the app. + + Parameters + ---------- + target_hint : str + The hint for the target device. + + prefix_hint : str + The hint for the system lib prefix. + """ + if prefix_hint == "auto" and ( + target_hint.startswith("iphone") + or target_hint.startswith("macabi") + or target_hint.startswith("android") + ): + prefix = f"{model_name}_{quantization}_".replace("-", "_") + logger.warning( + "%s is automatically picked from the filename, %s, this allows us to use the filename " + "as the model_lib in android/iOS builds. Please avoid renaming the .tar file when " + "uploading the prebuilt.", + bold("--system-lib-prefix"), + bold(prefix), + ) + return prefix + if target_hint not in ["iphone", "macabi", "android"]: + return "" + return prefix_hint + + +_MACABI_ARCH = os.environ.get("MLC_MACABI_ARCH", "").strip() or "arm64" +if _MACABI_ARCH not in ["arm64", "x86_64"]: + _MACABI_ARCH = "arm64" +_MACABI_MTRIPLE = ( + "x86_64-apple-ios18.0-macabi" if _MACABI_ARCH == "x86_64" else "arm64-apple-ios18.0-macabi" +) + +PRESET = { + "iphone:generic": { + "target": { + "kind": "metal", + "max_threads_per_block": 256, + "max_shared_memory_per_block": 32768, + "thread_warp_size": 1, + "libs": ["iphoneos"], + "host": { + "kind": "llvm", + "mtriple": "arm64-apple-darwin", + }, + }, + "build": _build_iphone, + }, + "macabi:generic": { + "target": { + "kind": "metal", + "max_threads_per_block": 256, + "max_shared_memory_per_block": 32768, + "thread_warp_size": 1, + "libs": ["macosx"], + "host": { + "kind": "llvm", + "mtriple": _MACABI_MTRIPLE, + }, + }, + "build": _build_iphone, + }, + "android:generic": { + "target": { + "kind": "opencl", + "host": { + "kind": "llvm", + "mtriple": "aarch64-linux-android", + }, + }, + "build": _build_android, + }, + "android:adreno": { + "target": { + "kind": "opencl", + "device": "adreno", + "max_threads_per_block": 512, + "host": { + "kind": "llvm", + "mtriple": "aarch64-linux-android", + }, + }, + "build": _build_android, + }, + "android:adreno-so": { + "target": { + "kind": "opencl", + "device": "adreno", + "max_threads_per_block": 512, + "host": { + "kind": "llvm", + "mtriple": "aarch64-linux-android", + }, + }, + "build": _build_android_so, + }, + "windows:adreno_x86": { + "target": { + "kind": "opencl", + "device": "adreno", + "max_threads_per_block": 512, + "host": { + "kind": "llvm", + "mtriple": "x86_64-pc-windows-msvc", + }, + }, + }, + "metal:x86-64": { + "target": { + "kind": "metal", + "max_threads_per_block": 256, + "max_shared_memory_per_block": 32768, + "thread_warp_size": 1, + }, + "build": _build_metal_x86_64, + }, + "webgpu:generic": { + "target": { + "kind": "webgpu", + "host": { + "kind": "llvm", + "mtriple": "wasm32-unknown-unknown-wasm", + }, + }, + "build": _build_webgpu, + }, + "opencl:generic": { + "target": { + "kind": "opencl", + }, + }, + "mali:generic": { + "target": { + "kind": "opencl", + "host": { + "kind": "llvm", + "mtriple": "aarch64-linux-gnu", + }, + }, + "build": _build_mali, + }, + "metal:generic": { + "target": { + "kind": "metal", + "max_threads_per_block": 256, + "max_shared_memory_per_block": 32768, + "thread_warp_size": 1, + }, + }, + "vulkan:generic": { + "target": { + "kind": "vulkan", + "max_threads_per_block": 256, + "max_shared_memory_per_block": 32768, + "thread_warp_size": 1, + "supports_float16": 1, + "supports_int64": 1, + "supports_int16": 1, + "supports_int8": 1, + "supports_8bit_buffer": 1, + "supports_16bit_buffer": 1, + "supports_storage_buffer_storage_class": 1, + }, + }, +} diff --git a/python/mlc_llm/support/auto_weight.py b/python/mlc_llm/support/auto_weight.py new file mode 100644 index 0000000..90e6ad6 --- /dev/null +++ b/python/mlc_llm/support/auto_weight.py @@ -0,0 +1,178 @@ +"""Help functions for detecting weight paths and weight formats.""" + +import json +from pathlib import Path +from typing import List, Optional, Tuple # noqa: UP035 + +from . import logging +from .style import bold, green, red + +logger = logging.getLogger(__name__) + +FOUND = green("Found") +NOT_FOUND = red("Not found") + + +def detect_weight( + weight_path: Path, + config_json_path: Path, + weight_format: str = "auto", +) -> Tuple[Path, str]: # noqa: UP006 + """Detect the weight directory, and detect the weight format. + + Parameters + --------- + weight_path : pathlib.Path + The path to weight files. If `weight_path` is not None, check if it exists. Otherwise, find + `weight_path` in `config.json` or use the same directory as `config.json`. + + config_json_path: pathlib.Path + The path to `config.json`. + + weight_format : str + The hint for the weight format. If it is "auto", guess the weight format. + Otherwise, check the weights are in that format. + Available weight formats: + - auto (guess the weight format) + - huggingface-torch (validate via checking pytorch_model.bin.index.json) + - huggingface-safetensor (validate via checking model.safetensors.index.json) + - awq + - ggml + - gguf + + Returns + ------- + weight_config_path : pathlib.Path + The path that points to the weights config file or the weights directory. + + weight_format : str + The valid weight format. + """ + if weight_path is None: + assert config_json_path is not None and config_json_path.exists(), ( + "Please provide config.json path." + ) + + # 1. Find the weight_path in config.json + with open(config_json_path, encoding="utf-8") as i_f: + config = json.load(i_f) + if "weight_path" in config: + weight_path = Path(config["weight_path"]) + logger.info('Found "weight_path" in config.json: %s', weight_path) + if not weight_path.exists(): + raise ValueError(f"weight_path doesn't exist: {weight_path}") + else: + # 2. Find the weights file in the same directory as config.json + weight_path = config_json_path.parent + else: + if not weight_path.exists(): + raise ValueError(f"weight_path doesn't exist: {weight_path}") + + logger.info("Finding weights in: %s", weight_path) + + # check weight format + # weight_format = "auto", guess the weight format. + # otherwise, check the weight format is valid. + if weight_format == "auto": + return _guess_weight_format(weight_path) + + if weight_format not in AVAILABLE_WEIGHT_FORMAT: + raise ValueError( + f"Available weight format list: {AVAILABLE_WEIGHT_FORMAT}, but got {weight_format}" + ) + if weight_format in CHECK_FORMAT_METHODS: + check_func = CHECK_FORMAT_METHODS[weight_format] + weight_config_path = check_func(weight_path) + if not weight_config_path: + raise ValueError(f"The weight is not in {weight_format} format.") + else: + weight_config_path = weight_path + return weight_config_path, weight_format + + +def _guess_weight_format(weight_path: Path) -> Tuple[Path, str]: # noqa: UP006 + possible_formats: List[Tuple[Path, str]] = [] # noqa: UP006 + for weight_format, check_func in CHECK_FORMAT_METHODS.items(): + weight_config_path = check_func(weight_path) + if weight_config_path: + possible_formats.append((weight_config_path, weight_format)) + + if len(possible_formats) == 0: + raise ValueError( + "Fail to detect source weight format. " + "Use `--source-format` to explicitly specify the format." + ) + + weight_config_path, selected_format = possible_formats[0] + logger.info( + "Using source weight configuration: %s. Use `--source` to override.", + bold(str(weight_config_path)), + ) + logger.info( + "Using source weight format: %s. Use `--source-format` to override.", + bold(selected_format), + ) + return weight_config_path, selected_format + + +def _check_pytorch(weight_path: Path) -> Optional[Path]: + pytorch_json_path = weight_path / "pytorch_model.bin.index.json" + if pytorch_json_path.exists(): + logger.info( + "%s source weight format: huggingface-torch. Source configuration: %s", + FOUND, + pytorch_json_path, + ) + return pytorch_json_path + + pytorch_file_path = weight_path / "pytorch_model.bin" + if pytorch_file_path.exists(): + logger.info( + "%s source weight format: huggingface-torch. Source configuration: %s", + FOUND, + pytorch_file_path, + ) + return pytorch_file_path + + logger.info("%s Huggingface PyTorch", NOT_FOUND) + return None + + +def _check_safetensor(weight_path: Path) -> Optional[Path]: + safetensor_json_path = weight_path / "model.safetensors.index.json" + if safetensor_json_path.exists(): + logger.info( + "%s source weight format: huggingface-safetensor. Source configuration: %s", + FOUND, + safetensor_json_path, + ) + return safetensor_json_path + + safetensor_file_path = weight_path / "model.safetensors" + if safetensor_file_path.exists(): + from safetensors.torch import ( + load_file, + ) + + weights = load_file(safetensor_file_path, device="cpu") + weight_map = {key: "model.safetensors" for key in weights} + with open(safetensor_json_path, "w", encoding="utf-8") as file: + json.dump({"weight_map": weight_map}, file, indent=2) + logger.info( + "%s source weight format: huggingface-safetensor. Source configuration: %s", + FOUND, + safetensor_json_path, + ) + return safetensor_json_path + + logger.info("%s Huggingface Safetensor", NOT_FOUND) + return None + + +CHECK_FORMAT_METHODS = { + "huggingface-torch": _check_pytorch, + "huggingface-safetensor": _check_safetensor, +} + +# "ggml", "gguf" are not supported yet. +AVAILABLE_WEIGHT_FORMAT = ["huggingface-torch", "huggingface-safetensor", "awq"] diff --git a/python/mlc_llm/support/config.py b/python/mlc_llm/support/config.py new file mode 100644 index 0000000..cc7a81a --- /dev/null +++ b/python/mlc_llm/support/config.py @@ -0,0 +1,111 @@ +""" +A common base class for configuration. A configuration could be initialized from its constructor, +a JSON string or a JSON file, and irrelevant fields during initialization are automatically moved +to the `kwargs` field. + +Take model configuration as an example: it is usually a JSON file in HuggingFace that contains +the model's hyperparameters. For instance, Vicuna-13b-v1.5-16k contains the following +[JSON file](https://huggingface.co/lmsys/vicuna-13b-v1.5-16k/blob/main/config.json). +The base class allows us to load the configuration from this JSON file, moving irrelevant fields +into `kwargs`, such as `transformers_version` and `use_cache`. +""" + +import dataclasses +import json +from pathlib import Path +from typing import Any, Dict, Type, TypeVar # noqa: UP035 + +from . import logging +from .style import bold, red + +logger = logging.getLogger(__name__) + +ConfigClass = TypeVar("ConfigClass", bound="ConfigBase") + + +@dataclasses.dataclass +class ConfigBase: + """Base class for configurations, providing a common interface for loading configs from a + JSON file or a dict. It requires the subclasses to be dataclasses, and has an `kwargs` field + that stores the extra fields that are not defined in the dataclass. + """ + + @classmethod + def from_dict(cls: Type[ConfigClass], source: Dict[str, Any]) -> ConfigClass: # noqa: UP006 + """Create a config object from a dictionary. + + Parameters + ---------- + source : Dict[str, Any] + Source to create config from, usually loaded from `config.json` in HuggingFace style. + + Returns + ------- + cfg : ConfigClass + An instance of the config object. + """ + field_names = [field.name for field in dataclasses.fields(cls)] + fields = {k: v for k, v in source.items() if k in field_names} + kwargs = {k: v for k, v in source.items() if k not in field_names} + return cls(**fields, kwargs=kwargs) + + @classmethod + def from_file(cls: Type[ConfigClass], source: Path) -> ConfigClass: # noqa: UP006 + """Create a config object from a file. + + Parameters + ---------- + cfg_cls : Type[ConfigClass] + The config class to create, for example, LlamaConfig. + + source : pathlib.Path + Path to the source file, usually `config.json` in HuggingFace repo. + + Returns + ------- + cfg : ConfigClass + An instance of the config object. + """ + with source.open("r", encoding="utf-8") as in_file: + return cls.from_dict(json.load(in_file)) + + def asdict(self): + """Convert the config object to a dictionary. + + Returns + ------- + Dict[str, Any] + A dictionary representation of the config object. + """ + result = dataclasses.asdict(self) + result.pop("kwargs") + return result + + +class ConfigOverrideBase: + """Base class for ConfigOverride, providing a common interface for overriding configs. + It requires the subclasses to be dataclasses. + """ + + def apply(self, config): + """Apply the overrides to the given config.""" + updated = config.asdict() + for field in dataclasses.fields(self): + key = field.name + value = getattr(self, key) + if value is None: + continue + if key not in updated: + logger.warning( + "%s: Cannot override %s, because %s does not have this field", + red("Warning"), + bold(key), + bold(type(config).__name__), + ) + else: + logger.info(f"Overriding {bold(key)} from {updated[key]} to {value}") + updated[key] = value + return type(config).from_dict(updated) + + +__all__ = ["ConfigBase", "ConfigOverrideBase"] diff --git a/python/mlc_llm/support/constants.py b/python/mlc_llm/support/constants.py new file mode 100644 index 0000000..bd1d2da --- /dev/null +++ b/python/mlc_llm/support/constants.py @@ -0,0 +1,90 @@ +"""Environment variables used by the MLC LLM.""" + +import os +import sys +from pathlib import Path +from typing import List # noqa: UP035 + +MLC_CHAT_CONFIG_VERSION = "0.1.0" + + +def _check(): + if MLC_JIT_POLICY not in ["ON", "OFF", "REDO", "READONLY"]: + raise ValueError( + 'Invalid MLC_JIT_POLICY. It has to be one of "ON", "OFF", "REDO", "READONLY"' + f"but got {MLC_JIT_POLICY}." + ) + + if MLC_DOWNLOAD_CACHE_POLICY not in ["ON", "OFF", "REDO", "READONLY"]: + raise ValueError( + "Invalid MLC_AUTO_DOWNLOAD_POLICY. " + 'It has to be one of "ON", "OFF", "REDO", "READONLY"' + f"but got {MLC_DOWNLOAD_CACHE_POLICY}." + ) + + +def _get_cache_dir() -> Path: + if "MLC_LLM_HOME" in os.environ: + result = Path(os.environ["MLC_LLM_HOME"]) + elif sys.platform == "win32": + result = Path(os.environ["LOCALAPPDATA"]) + result = result / "mlc_llm" + elif os.getenv("XDG_CACHE_HOME", None) is not None: + result = Path(os.getenv("XDG_CACHE_HOME")) + result = result / "mlc_llm" + else: + result = Path(os.path.expanduser("~/.cache")) + result = result / "mlc_llm" + result.mkdir(parents=True, exist_ok=True) + if not result.is_dir(): + raise ValueError( + f"The default cache directory is not a directory: {result}. " + "Use environment variable MLC_LLM_HOME to specify a valid cache directory." + ) + (result / "model_weights").mkdir(parents=True, exist_ok=True) + (result / "model_lib").mkdir(parents=True, exist_ok=True) + return result + + +def _get_dso_suffix() -> str: + if "MLC_DSO_SUFFIX" in os.environ: + return os.environ["MLC_DSO_SUFFIX"] + if sys.platform == "win32": + return "dll" + if sys.platform == "darwin": + return "dylib" + return "so" + + +def _get_test_model_path() -> List[Path]: # noqa: UP006 + paths = [] + if "MLC_LLM_TEST_MODEL_PATH" in os.environ: + paths += [Path(p) for p in os.environ["MLC_LLM_TEST_MODEL_PATH"].split(os.pathsep)] + # by default, we reuse the cache dir via mlc_llm chat + # note that we do not auto download for testcase + # to avoid networking dependencies + base_list = ["hf"] + paths += [_get_cache_dir() / "model_weights" / base / "mlc-ai" for base in base_list] + [ + Path(os.path.abspath(os.path.curdir)), + Path(os.path.abspath(os.path.curdir)) / "dist", + ] + return paths + + +def _get_read_only_weight_caches() -> List[Path]: # noqa: UP006 + if "MLC_LLM_READONLY_WEIGHT_CACHE" in os.environ: + return [Path(p) for p in os.environ["MLC_LLM_READONLY_WEIGHT_CACHE"].split(os.pathsep)] + return [] + + +MLC_TEMP_DIR = os.getenv("MLC_TEMP_DIR", None) +MLC_MULTI_ARCH = os.environ.get("MLC_MULTI_ARCH", None) +MLC_JIT_POLICY = os.environ.get("MLC_JIT_POLICY", "ON") +MLC_DSO_SUFFIX = _get_dso_suffix() +MLC_TEST_MODEL_PATH: List[Path] = _get_test_model_path() # noqa: UP006 + +MLC_DOWNLOAD_CACHE_POLICY = os.environ.get("MLC_DOWNLOAD_CACHE_POLICY", "ON") +MLC_LLM_HOME: Path = _get_cache_dir() +MLC_LLM_READONLY_WEIGHT_CACHE = _get_read_only_weight_caches() + +_check() diff --git a/python/mlc_llm/support/convert_tiktoken.py b/python/mlc_llm/support/convert_tiktoken.py new file mode 100644 index 0000000..2432e78 --- /dev/null +++ b/python/mlc_llm/support/convert_tiktoken.py @@ -0,0 +1,171 @@ +""" +Adapted from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee +Generator of mlc-chat-config.json and tokenizer configuration. +""" + +# isort: off +import json +import os +from typing import Dict, List, Optional # noqa: UP035 + + +def bpe( + mergeable_ranks: Dict[bytes, int], # noqa: UP006 + token: bytes, + max_rank: Optional[int] = None, +) -> List[bytes]: # noqa: UP006 + """Adapted from https://github.com/openai/tiktoken/issues/60#issuecomment-1499977960""" + parts = [bytes([b]) for b in token] + while True: + min_idx = None + min_rank = None + for i, pair in enumerate(zip(parts[:-1], parts[1:])): + rank = mergeable_ranks.get(pair[0] + pair[1]) + if rank is not None and (min_rank is None or rank < min_rank): + min_idx = i + min_rank = rank + if min_rank is None or (max_rank is not None and min_rank >= max_rank): + break + assert min_idx is not None + parts = [*parts[:min_idx], parts[min_idx] + parts[min_idx + 1], *parts[min_idx + 2 :]] + return parts + + +def generate_vocab_and_merges(encoder, mergeable_ranks): + """Generate vocab and merges in huggingface tokenizers format""" + + from transformers.models.gpt2.tokenization_gpt2 import ( + bytes_to_unicode, + ) + + byte_encoder = bytes_to_unicode() + + def token_bytes_to_string(b): + """Convert a token from bytes to a string""" + return "".join([byte_encoder[ord(char)] for char in b.decode("latin-1")]) + + merges = [] + vocab = {} + for token, rank in mergeable_ranks.items(): + vocab[token_bytes_to_string(token)] = rank + + if len(token) == 1: + continue + merged = tuple(bpe(mergeable_ranks, token, max_rank=rank)) + assert len(merged) == 2 + + merges.append(" ".join(map(token_bytes_to_string, merged))) + + # Also add special tokens + vocab.update(encoder._special_tokens) + + return vocab, merges + + +def convert_tiktoken(model_path, output_dir, context_window_size=None): + """Convert tiktoken tokenizers to huggingface tokenizers style""" + try: + from transformers import AutoTokenizer + except ImportError: + raise ImportError( + 'Converting tiktoken tokenizer requires the "transformers" package.' + 'Please install the "transformers" package to convert toktoken tokenizer' + ) + + tiktoken_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + encoder = tiktoken_tokenizer.tokenizer + + vocab, merges = generate_vocab_and_merges(encoder, tiktoken_tokenizer.get_vocab()) + + added_tokens = [ + { + "id": id, + "content": content, + "single_word": False, + "lstrip": False, + "rstrip": False, + "normalized": False, + "special": True, + } + for content, id in encoder._special_tokens.items() + ] + + tokenizer_template = { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": added_tokens, + "normalizer": None, + "pre_tokenizer": { + "type": "ByteLevel", + "add_prefix_space": False, + "trim_offsets": True, + "use_regex": True, + }, + "post_processor": { + "type": "ByteLevel", + "add_prefix_space": True, + "trim_offsets": False, + "use_regex": True, + }, + "decoder": { + "type": "ByteLevel", + "add_prefix_space": True, + "trim_offsets": True, + "use_regex": True, + }, + "model": { + "type": "BPE", + "dropout": None, + "unk_token": None, + "continuing_subword_prefix": "", + "end_of_word_suffix": "", + "fuse_unk": False, + "byte_fallback": False, + "vocab": vocab, + "merges": merges, + }, + } + + tokenizer_config_template = { + "add_prefix_space": False, + "bos_token": "<|endoftext|>", + "clean_up_tokenization_spaces": True, + "eos_token": "<|endoftext|>", + "unk_token": "<|endoftext|>", + } + + tokenizer_name = type(tiktoken_tokenizer).__name__ + + tokenizer_config_template["tokenizer_class"] = tokenizer_name + if context_window_size: + tokenizer_config_template["model_max_length"] = context_window_size + tokenizer_config_template = dict(sorted(tokenizer_config_template.items(), key=lambda x: x[0])) + + os.makedirs(output_dir, exist_ok=True) + + # Save to files + with open(os.path.join(output_dir, "vocab.json"), "w", encoding="utf-8") as fp: + json.dump(vocab, fp, indent=2, ensure_ascii=False) + + with open(os.path.join(output_dir, "tokenizer.json"), "w", encoding="utf-8") as fp: + json.dump(tokenizer_template, fp, indent=2, ensure_ascii=False) + + with open(os.path.join(output_dir, "tokenizer_config.json"), "w", encoding="utf-8") as fp: + json.dump(tokenizer_config_template, fp, indent=2, ensure_ascii=False) + + with open(os.path.join(output_dir, "special_tokens_map.json"), "w", encoding="utf-8") as fp: + json.dump( + { + "bos_token": "<|endoftext|>", + "eos_token": "<|endoftext|>", + "unk_token": "<|endoftext|>", + }, + fp, + indent=2, + ensure_ascii=False, + ) + + with open(os.path.join(output_dir, "merges.txt"), "w", encoding="utf-8") as fp: + fp.write("#version: 0.2\n") + fp.write("\n".join(merges)) diff --git a/python/mlc_llm/support/download_cache.py b/python/mlc_llm/support/download_cache.py new file mode 100644 index 0000000..fd262cb --- /dev/null +++ b/python/mlc_llm/support/download_cache.py @@ -0,0 +1,237 @@ +"""Common utilities for downloading files from HuggingFace or other URLs online.""" + +import concurrent.futures as cf +import hashlib +import json +import os +import shutil +import subprocess +import tempfile +from pathlib import Path +from typing import List, Optional, Tuple # noqa: UP035 + +import requests + +from . import logging, tqdm +from .constants import ( + MLC_DOWNLOAD_CACHE_POLICY, + MLC_LLM_HOME, + MLC_LLM_READONLY_WEIGHT_CACHE, + MLC_TEMP_DIR, +) +from .style import bold + +logger = logging.getLogger(__name__) + + +def log_download_cache_policy(): + """log current download policy""" + logger.info( + "%s = %s. Can be one of: ON, OFF, REDO, READONLY", + bold("MLC_DOWNLOAD_CACHE_POLICY"), + MLC_DOWNLOAD_CACHE_POLICY, + ) + + +def _ensure_directory_not_exist(path: Path, force_redo: bool) -> None: + if path.exists(): + if force_redo: + logger.info("Deleting existing directory: %s", path) + shutil.rmtree(path) + else: + raise ValueError(f"Directory already exists: {path}") + else: + path.parent.mkdir(parents=True, exist_ok=True) + + +def git_clone(url: str, destination: Path, ignore_lfs: bool) -> None: + """Clone a git repository into a directory.""" + repo_name = ".tmp" + command = ["git", "clone", url, repo_name] + _ensure_directory_not_exist(destination, force_redo=False) + try: + env = os.environ.copy() + env["GIT_LFS_SKIP_SMUDGE"] = "1" + with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir: + logger.info("[Git] Cloning %s to %s", bold(url), destination) + subprocess.run( + command, + env=env, + cwd=tmp_dir, + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + git_dir = os.path.join(tmp_dir, repo_name) + if not ignore_lfs: + git_lfs_pull(Path(git_dir)) + shutil.move(git_dir, str(destination)) + except subprocess.CalledProcessError as error: + raise ValueError( + f"Git clone failed with return code {error.returncode}: {error.stderr}. " + f"The command was: {command}" + ) from error + + +def git_lfs_pull(repo_dir: Path, ignore_extensions: Optional[List[str]] = None) -> None: # noqa: UP006 + """Pull files with Git LFS.""" + filenames = ( + subprocess.check_output( + ["git", "-C", str(repo_dir), "lfs", "ls-files", "-n"], + stderr=subprocess.STDOUT, + ) + .decode("utf-8") + .splitlines() + ) + if ignore_extensions is not None: + filenames = [ + filename + for filename in filenames + if not any(filename.endswith(extension) for extension in ignore_extensions) + ] + logger.info("[Git LFS] Downloading %d files with Git LFS: %s", len(filenames), filenames) + with tqdm.redirect(): + for file in tqdm.tqdm(filenames): + logger.info("[Git LFS] Downloading %s", file) + subprocess.check_output( + ["git", "-C", str(repo_dir), "lfs", "pull", "--include", file], + stderr=subprocess.STDOUT, + ) + + +def download_file( + url: str, + destination: Path, + md5sum: Optional[str], +) -> Tuple[str, Path]: # noqa: UP006 + """Download a file from a URL to a destination file.""" + with requests.get(url, stream=True, timeout=30) as response: + response.raise_for_status() + with destination.open("wb") as file: + for chunk in response.iter_content(chunk_size=8192): + file.write(chunk) + if md5sum is not None: + hash_md5 = hashlib.md5() + with destination.open("rb") as file: + for chunk in iter(lambda: file.read(8192), b""): + hash_md5.update(chunk) + file_md5 = hash_md5.hexdigest() + if file_md5 != md5sum: + raise ValueError( + f"MD5 checksum mismatch for downloaded file: {destination}. " + f"Expected {md5sum}, got {file_md5}" + ) + return url, destination + + +def download_and_cache_mlc_weights( + model_url: str, + num_processes: int = 4, + force_redo: Optional[bool] = None, +) -> Path: + """Download weights for a model from the HuggingFace Git LFS repo.""" + log_download_cache_policy() + if MLC_DOWNLOAD_CACHE_POLICY == "OFF": + raise RuntimeError(f"Cannot download {model_url} as MLC_DOWNLOAD_CACHE_POLICY=OFF") + + prefixes, mlc_prefix = ["HF://", "https://huggingface.co/"], "" + mlc_prefix = next(p for p in prefixes if model_url.startswith(p)) + assert mlc_prefix + + git_url_template = "https://huggingface.co/{user}/{repo}" + bin_url_template = "https://huggingface.co/{user}/{repo}/resolve/main/{record_name}" + + if model_url.count("/") != 1 + mlc_prefix.count("/") or not model_url.startswith(mlc_prefix): + raise ValueError(f"Invalid model URL: {model_url}") + user, repo = model_url[len(mlc_prefix) :].split("/") + domain = "hf" + + readonly_cache_dirs = [] + for base in MLC_LLM_READONLY_WEIGHT_CACHE: + cache_dir = base / domain / user / repo + readonly_cache_dirs.append(str(cache_dir)) + if (cache_dir / "mlc-chat-config.json").is_file(): + logger.info("Use cached weight: %s", bold(str(cache_dir))) + return cache_dir + + if force_redo is None: + force_redo = MLC_DOWNLOAD_CACHE_POLICY == "REDO" + + git_dir = MLC_LLM_HOME / "model_weights" / domain / user / repo + readonly_cache_dirs.append(str(git_dir)) + + try: + _ensure_directory_not_exist(git_dir, force_redo=force_redo) + except ValueError: + logger.info("Weights already downloaded: %s", bold(str(git_dir))) + return git_dir + + if MLC_DOWNLOAD_CACHE_POLICY == "READONLY": + raise RuntimeError( + f"Cannot find cache for {model_url}, " + "cannot proceed to download as MLC_DOWNLOAD_CACHE_POLICY=READONLY, " + "please check settings MLC_LLM_READONLY_WEIGHT_CACHE, " + f"local path candidates: {readonly_cache_dirs}" + ) + + with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir_prefix: + tmp_dir = Path(tmp_dir_prefix) / "tmp" + git_url = git_url_template.format(user=user, repo=repo) + git_clone(git_url, tmp_dir, ignore_lfs=True) + git_lfs_pull(tmp_dir, ignore_extensions=[".bin"]) + shutil.rmtree(tmp_dir / ".git", ignore_errors=True) + with (tmp_dir / "tensor-cache.json").open(encoding="utf-8") as in_file: + param_metadata = json.load(in_file)["records"] + with cf.ProcessPoolExecutor(max_workers=num_processes) as executor: + futures = [] + for record in param_metadata: + record_name = record["dataPath"] + file_url = bin_url_template.format(user=user, repo=repo, record_name=record_name) + file_dest = tmp_dir / record_name + file_md5 = record.get("md5sum", None) + futures.append(executor.submit(download_file, file_url, file_dest, file_md5)) + with tqdm.redirect(): + for future in tqdm.tqdm(cf.as_completed(futures), total=len(futures)): + file_url, file_dest = future.result() + logger.info("Downloaded %s to %s", file_url, file_dest) + logger.info("Moving %s to %s", tmp_dir, bold(str(git_dir))) + shutil.move(str(tmp_dir), str(git_dir)) + return git_dir + + +def get_or_download_model(model: str) -> Path: + """Use user-provided argument ``model`` to get model_path + + We define "valid" as having an ``mlc-chat-config.json`` right under the folder. + + Parameters + ---------- + model : str + User's input; may a path or url + + Returns + ------ + model_path : Path + A "valid" path to model folder, with + ``(model_path / "mlc-chat-config.json").is_file`` being True + + Note + ---- + This function may perform additional download and caching + + Raises + ------ + FileNotFoundError: if we cannot find a valid `model_path`. + """ + if model.startswith("HF://"): + logger.info("Downloading model from HuggingFace: %s", model) + model_path = download_and_cache_mlc_weights(model) + else: + model_path = Path(model) + + if not model_path.is_dir(): + raise FileNotFoundError(f"Cannot find model {model}, directory does not exist") + mlc_config_path = model_path / "mlc-chat-config.json" + if mlc_config_path.is_file(): + return model_path + raise FileNotFoundError(f"Cannot find {str(mlc_config_path)} in the model directory provided") diff --git a/python/mlc_llm/support/logging.py b/python/mlc_llm/support/logging.py new file mode 100644 index 0000000..e8240f5 --- /dev/null +++ b/python/mlc_llm/support/logging.py @@ -0,0 +1,24 @@ +""" +Logging support for MLC. It derives from Python's logging module, and in the future, +it can be easily replaced by other logging modules such as structlog. +""" + +import logging +import os + + +def enable_logging(): + """Enable MLC's default logging format""" + if os.getenv("MLC_UNSET_LOGGING"): + return + logging.basicConfig( + level=logging.INFO, + style="{", + datefmt="%Y-%m-%d %H:%M:%S", + format="[{asctime}] {levelname} {filename}:{lineno}: {message}", + ) + + +def getLogger(name: str): + """Get a logger according to the given name""" + return logging.getLogger(name) diff --git a/python/mlc_llm/support/max_thread_check.py b/python/mlc_llm/support/max_thread_check.py new file mode 100644 index 0000000..2d570c8 --- /dev/null +++ b/python/mlc_llm/support/max_thread_check.py @@ -0,0 +1,38 @@ +"""Helper functions for checking max num thread.""" + +from tvm.target import Target + + +def get_max_num_threads_per_block(target: Target) -> int: + """ + max(max_num_threads, max_threads_per_block); if latter does not exist, return max_num_threads. + We add this method since some targets have both fields and `max_threads_per_block` is larger. + """ + max_num_threads = target.attrs.get("max_num_threads") + max_threads_per_block = target.attrs.get("max_threads_per_block", None) + if max_threads_per_block is None: + return max_num_threads + return max(max_num_threads, max_threads_per_block) + + +def check_thread_limits(target: Target, bdx: int, bdy: int, bdz: int, gdz: int): + """ + Check whether max num threads exceeded given a target. + + Parameters + ---------- + bdx: threadIdx.x + bdy: threadIdx.y + bdz: threadIdx.z + gdz: blockIdx.z + """ + max_num_threads_per_block = get_max_num_threads_per_block(target) + + assert bdx * bdy * bdz <= max_num_threads_per_block, ( + f"{target.kind} max num threads exceeded: {bdx}*{bdy}*{bdz}>{max_num_threads_per_block}" + ) + + if target.kind.name == "webgpu": + # https://gpuweb.github.io/gpuweb/#dom-supported-limits-maxcomputeworkgroupsizez + assert bdz <= 64, f"webgpu's threadIdx.z cannot exceed 64, but got bdz={bdz}" + assert gdz == 1, f"webgpu's blockIdx.z should be 1, but got gdz={gdz}" diff --git a/python/mlc_llm/support/preshard.py b/python/mlc_llm/support/preshard.py new file mode 100644 index 0000000..b0b5b19 --- /dev/null +++ b/python/mlc_llm/support/preshard.py @@ -0,0 +1,123 @@ +"""Functions for pre-sharding weights""" + +import logging +from collections.abc import Sequence +from typing import Any, Callable, Dict, Tuple # noqa: UP035 + +from tvm import IRModule, relax +from tvm.relax.frontend import nn +from tvm.runtime import Device, Tensor +from tvm.s_tir import dlight as dl +from tvm.target import Target + +logger = logging.getLogger("preshard") + + +def _sharded_param_name(param_name, worker_id): + return f"{param_name}_shard-{worker_id}" + + +def _create_shard_func(bb: relax.BlockBuilder, param: nn.Parameter, tensor_parallel_shards: int): + shard_strategy = param.attrs.get("shard_strategy", None) + # generate tirx shard function + tir_func = shard_strategy.gen_tir(shards=tensor_parallel_shards, weight=param) + tir_func = tir_func.with_attr("global_symbol", f"{shard_strategy.name}_tir") + # add tirx shard function to the IRModule + tir_gvar = bb.add_func(tir_func, func_name=f"{shard_strategy.name}_tir") + # create relax function that + # 1. shard weight with tirx shard function, result: [num_shards, *sharded_weight_shape] + # 2. split the sharded weight along dim 0, result: num_shards * [1, *sharded_weight_shape] + # 3. squeeze the 0th-dim of all shards, result: num_shards * [*sharded_weight_shape] + weight_shape = param.shape + weight_shape[shard_strategy.dim] = weight_shape[shard_strategy.dim] * tensor_parallel_shards + sharded_weight_shape = [tensor_parallel_shards, *param.shape] + weight_var = relax.Var("weight", relax.TensorType(weight_shape, param.dtype)) + with bb.function(name=shard_strategy.name, params=[weight_var]): + with bb.dataflow(): + lv0 = bb.emit( + relax.call_tir( + tir_gvar, + weight_var, + out_ty=relax.TensorType(sharded_weight_shape, param.dtype), + ) + ) + lv1 = bb.emit(relax.op.split(lv0, indices_or_sections=tensor_parallel_shards, axis=0)) + output_vars = [] + for i in range(tensor_parallel_shards): + lvi = bb.emit(relax.TupleGetItem(lv1, i)) + squeezed_lvi = bb.emit(relax.op.squeeze(lvi, 0)) + output_vars.append(squeezed_lvi) + gv = bb.emit_output(output_vars) + bb.emit_func_output(gv) + + +def _compile_shard_funcs(mod: IRModule, device: Device): + target = Target.from_device(device) + with target: + mod = relax.transform.LegalizeOps()(mod) + mod = dl.ApplyDefaultSchedule( + dl.gpu.Matmul(), + dl.gpu.GEMV(), + dl.gpu.Reduction(), + dl.gpu.GeneralReduction(), + dl.gpu.Fallback(), + )(mod) + ex = relax.build(mod, target=target) + vm = relax.VirtualMachine(ex, device) + return vm + + +def apply_preshard( + named_params: Dict[str, nn.Parameter], # noqa: UP006 + tensor_parallel_shards: int, + args: Any, +) -> Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]]: # noqa: UP006 + """Apply pre-sharding to the named parameters. + + Parameters + ---------- + named_params : Dict[str, nn.Parameter] + The named parameters of the model. If the model is quantized, the named parameters should + the state dictionary of the quantized model. + tensor_parallel_shards : int + The number of tensor parallel shards. + args : Any + The parsed arguments of weight conversion. + + Returns + ------- + Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]] + The updated named parameters and the mapping from parameter name to the shard function. + """ + bb = relax.BlockBuilder() + param_to_shard_func = {} + shard_func_names = set() + new_named_params: Dict[str, nn.Parameter] = {} # noqa: UP006 + has_shard_strategy = False + for name, param in named_params.items(): + shard_strategy = param.attrs.get("shard_strategy", None) + if shard_strategy is not None: + has_shard_strategy = True + for i in range(tensor_parallel_shards): + new_named_params[_sharded_param_name(name, i)] = param + # create shard functions + param_to_shard_func[name] = shard_strategy.name + if shard_strategy.name not in shard_func_names: + _create_shard_func(bb, param, tensor_parallel_shards) + shard_func_names.add(shard_strategy.name) + else: + new_named_params[name] = param + + if not has_shard_strategy: + logger.warning( + "No parameters with 'shard_strategy' found." + "At least one parameter must have a 'shard_strategy' for presharding. " + "The model will continue to convert weights in a non-presharded manner." + ) + + mod = bb.finalize() + vm = _compile_shard_funcs(mod, args.device) + + for name in param_to_shard_func: + param_to_shard_func[name] = vm[param_to_shard_func[name]] + return new_named_params, param_to_shard_func diff --git a/python/mlc_llm/support/random.py b/python/mlc_llm/support/random.py new file mode 100644 index 0000000..9c142ed --- /dev/null +++ b/python/mlc_llm/support/random.py @@ -0,0 +1,17 @@ +"""Utility functions for random number generation.""" + +import sys + + +def set_global_random_seed(seed): + """Set global random seed for python, numpy, torch and tvm.""" + if "numpy" in sys.modules: + sys.modules["numpy"].random.seed(seed) + if "torch" in sys.modules: + sys.modules["torch"].manual_seed(seed) + if "random" in sys.modules: + sys.modules["random"].seed(seed) + if "tvm" in sys.modules: + set_seed = sys.modules["tvm"].get_global_func("mlc.random.set_seed") + if set_seed: + set_seed(seed) diff --git a/python/mlc_llm/support/style.py b/python/mlc_llm/support/style.py new file mode 100644 index 0000000..5b2272e --- /dev/null +++ b/python/mlc_llm/support/style.py @@ -0,0 +1,62 @@ +"""Printing styles.""" + +from enum import Enum + + +class Styles(Enum): + """Predefined set of styles to be used. + + Reference: + - https://en.wikipedia.org/wiki/ANSI_escape_code#3-bit_and_4-bit + - https://stackoverflow.com/a/17303428 + """ + + RED = "\033[91m" + GREEN = "\033[92m" + YELLOW = "\033[93m" + BLUE = "\033[94m" + PURPLE = "\033[95m" + CYAN = "\033[96m" + BOLD = "\033[1m" + UNDERLINE = "\033[4m" + END = "\033[0m" + + +def red(text: str) -> str: + """Return red text.""" + return f"{Styles.RED.value}{text}{Styles.END.value}" + + +def green(text: str) -> str: + """Return green text.""" + return f"{Styles.GREEN.value}{text}{Styles.END.value}" + + +def yellow(text: str) -> str: + """Return yellow text.""" + return f"{Styles.YELLOW.value}{text}{Styles.END.value}" + + +def blue(text: str) -> str: + """Return blue text.""" + return f"{Styles.BLUE.value}{text}{Styles.END.value}" + + +def purple(text: str) -> str: + """Return purple text.""" + return f"{Styles.PURPLE.value}{text}{Styles.END.value}" + + +def cyan(text: str) -> str: + """Return cyan text.""" + return f"{Styles.CYAN.value}{text}{Styles.END.value}" + + +def bold(text: str) -> str: + """Return bold text.""" + return f"{Styles.BOLD.value}{text}{Styles.END.value}" + + +def underline(text: str) -> str: + """Return underlined text.""" + return f"{Styles.UNDERLINE.value}{text}{Styles.END.value}" diff --git a/python/mlc_llm/support/tensor_parallel.py b/python/mlc_llm/support/tensor_parallel.py new file mode 100644 index 0000000..44d66fe --- /dev/null +++ b/python/mlc_llm/support/tensor_parallel.py @@ -0,0 +1,118 @@ +"""Sharding operators for tensor parallelism.""" + +import dataclasses +from contextlib import contextmanager +from typing import Any, Dict, List, Optional # noqa: UP035 + +from tvm import te, tirx, topi +from tvm.relax.frontend import nn + + +@dataclasses.dataclass +class ShardSingleDim: + """ + Shard a tensor by a single dimension. + + + Parameters + ---------- + name : str + The name of the shard func + + dim : int + The dimension to shard + + segs : Optional[List[int]] + The length of segments along `dim`. Default to None. If specified, + shard a tensor by its "segmented" dimension, where each segment has a different length + and sharded evenly on each worker. + + """ + + name: str + dim: int + segs: Optional[List[int]] = None # noqa: UP006 + + def gen_tir(self, shards: int, weight: nn.Tensor) -> tirx.PrimFunc: + """Generate a TIR function that shards the weight tensor by its rows.""" + shape = weight.shape + segs = self.segs or [shape[self.dim]] + assert sum(segs) == shape[self.dim] + # NOTE: we use int64 to prevent int32 overflow + shape = [tirx.IntImm("int64", v) for v in shape] + segs = [tirx.IntImm("int64", v) for v in segs] + w = te.placeholder( + [tirx.IntImm("int64", v) for v in self._compute_in_shape(shards, weight)], + weight.dtype, + name="w", + ) + ws: List[te.Tensor] = [] # noqa: UP006 + offset = 0 + for idx, sub_seg in enumerate(segs): + ws.append( + topi.transpose( + topi.reshape( + te.compute( + ( + *shape[: self.dim], + sub_seg * shards, + *shape[self.dim + 1 :], + ), + lambda *idx: w[ + ( + *idx[: self.dim], + idx[self.dim] + offset, + *idx[self.dim + 1 :], + ) + ], + name=f"w_{idx}", + ), + ( + *shape[: self.dim], + tirx.IntImm("int64", shards), + sub_seg, + *shape[self.dim + 1 :], + ), + ), + [self.dim, *range(self.dim), *range(self.dim + 1, len(shape) + 1)], + ) + ) + offset += sub_seg * shards + o = topi.concatenate(ws, axis=1 + self.dim) + func = te.create_prim_func([w, o]) + return func + + def gen_shard_info(self, shards: int, weight: nn.Tensor) -> Dict[str, Any]: # noqa: UP006 + """Generate shard info for this sharding strategy.""" + return { + "func_name": self.name, + "in_shape": self._compute_in_shape(shards, weight), + "out_shape": (shards, *weight.shape), + "out_dtype": str(weight.dtype), + } + + def _compute_in_shape(self, shards: int, weight: nn.Tensor) -> List[int]: # noqa: UP006 + """Compute the weight shape before sharding.""" + shape = weight.shape + return [*shape[: self.dim], shape[self.dim] * shards, *shape[self.dim + 1 :]] + + +@contextmanager +def shard_bias(linear: nn.Linear, tensor_parallel_shards: int): + """ + A context manager to shard the bias of a linear into `tensor_parallel_shards` shards. + + + Parameters + ---------- + linear : nn.Linear + The linear layer whose bias would be sharded. + + tensor_parallel_shards : int + The number of shards. + """ + original_bias = linear.bias + if tensor_parallel_shards > 1: + linear.bias = linear.bias / tensor_parallel_shards + yield + linear.bias = original_bias diff --git a/python/mlc_llm/support/tqdm.py b/python/mlc_llm/support/tqdm.py new file mode 100644 index 0000000..8a4d09c --- /dev/null +++ b/python/mlc_llm/support/tqdm.py @@ -0,0 +1,39 @@ +"""Utils to better use tqdm""" + +import contextlib +import inspect +import io + +from tqdm import tqdm +from tqdm.contrib.logging import logging_redirect_tqdm as _redirect_logging + + +@contextlib.contextmanager +def _redirect_print(): + old_print = print + + def new_print(*args, **kwargs): + with io.StringIO() as output: + kwargs["file"] = output + kwargs["end"] = "" + old_print(*args, **kwargs) + content = output.getvalue() + tqdm.write(content) + + try: + inspect.builtins.print = new_print + yield + finally: + inspect.builtins.print = old_print + + +@contextlib.contextmanager +def redirect(): + """Redirect tqdm output to logging and print.""" + + with _redirect_logging(): + with _redirect_print(): + yield + + +__all__ = ["redirect", "tqdm"] diff --git a/python/mlc_llm/testing/__init__.py b/python/mlc_llm/testing/__init__.py new file mode 100644 index 0000000..ef1c388 --- /dev/null +++ b/python/mlc_llm/testing/__init__.py @@ -0,0 +1,5 @@ +""" +Test and debug tools for MLC LLM +""" + +from .pytest_utils import require_test_model, require_test_tokenizers diff --git a/python/mlc_llm/testing/debug_chat.py b/python/mlc_llm/testing/debug_chat.py new file mode 100644 index 0000000..9bee3de --- /dev/null +++ b/python/mlc_llm/testing/debug_chat.py @@ -0,0 +1,560 @@ +"""Debug compiled models with TVM instrument""" + +import json +import random +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union # noqa: UP035 + +import numpy as np +import tvm +import tvm_ffi +from tvm import DataType, relax +from tvm.contrib import tvmjs +from tvm.runtime import Device, Module, Object +from tvm.runtime.vm import VirtualMachine +from tvm_ffi import Shape + +from mlc_llm.conversation_template import ConvTemplateRegistry +from mlc_llm.interface.help import HELP +from mlc_llm.protocol.mlc_chat_config import MLCChatConfig +from mlc_llm.serve import data, engine_utils +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.support.auto_device import detect_device +from mlc_llm.support.style import green, red +from mlc_llm.tokenizers import Tokenizer + + +def _extract_metadata(mod: Module): + return json.loads(VirtualMachine(mod, tvm.runtime.device("cpu"))["_metadata"]()) + + +def _load_params( + model_weight_path: str, + device: Device, + model_metadata: Dict[str, Any], # noqa: UP006 +) -> List[tvm.runtime.Tensor]: # noqa: UP006 + params, meta = tvmjs.load_tensor_cache(model_weight_path, device) + param_names = [param["name"] for param in model_metadata["params"]] + assert len(param_names) == meta["ParamSize"] + + plist = [] + for param_name in param_names: + plist.append(params[param_name]) + return plist + + +def _get_tvm_module( + model_weight_path: str, + lib_path: str, + device: Device, + instrument: Union[tvm_ffi.Function, None], +): + ex = tvm.runtime.load_module(lib_path) + vm = relax.VirtualMachine(ex, device) + if instrument is not None: + vm.set_instrument(instrument) + metadata = _extract_metadata(ex) + params = _load_params(model_weight_path, device, metadata) + return vm.module, params, metadata + + +class DefaultDebugInstrument: + """The default debug instrument to use if users don't specify + a customized one. + + This debug instrument will dump the arguments and output of each + VM Call instruction into a .npz file. It will also alert the user + if any function outputs are NaN or INF. + """ + + def __init__(self, debug_out: Path): + """Constructor + + Parameters + ---------- + debug_out : Path + the directory to dump the .npz files + """ + self.counter = 0 + self.first_nan_occurred = False + self.first_inf_occurred = False + self.debug_out = debug_out + debug_out.mkdir(exist_ok=True, parents=True) + + def reset(self, debug_out: Path): + """Reset the state of the Instrument class + + Parameters + ---------- + debug_out : Path + the directory to dump the .npz files + """ + self.counter = 0 + self.first_nan_occurred = False + self.first_inf_occurred = False + self.debug_out = debug_out + debug_out.mkdir(exist_ok=True, parents=True) + + def __call__(self, func, name, before_run, ret_val, *args): + # Determine what functions to look at + if before_run: # Whether before the function is called or after + return + if self.first_nan_occurred: + return + if self.first_inf_occurred: + return + if ( + name.startswith("vm.builtin.") + and "call_tir_dyn" not in name + and "attention_with_fused_qkv" not in name + and "self_attention" not in name + and "cross_attention" not in name + ): + return + + # Decide what to print or save about the function's arguments (where args[-1] is the + # buffer we write the result to) + func_name = f"f{self.counter}_{name}" + + # Write your own behavior below. For example, we can count the number of INF/NaN in args[-1] + def _check_nan_inf(npy): + num_nans = np.sum(np.isnan(npy)) + num_infs = np.sum(np.isinf(npy)) + if num_nans > 0: + print(f"{red(f'{func_name} has NaN')}: {num_nans}") + self.first_nan_occurred = True + if num_infs > 0: + print(f"{red(f'{func_name} has INF')}: {num_infs}") + self.first_inf_occurred = True + + # Save the arguments to npz + arg_dict = {} + for i, arg in enumerate(args): + if isinstance(arg, tvm.runtime.Tensor): + if np.prod(arg.shape) * (DataType(arg.dtype).bits // 8) > 2147483648: + # We skip dump large tensors + arg_dict[f"arg_{i}"] = np.zeros(()) + elif arg.dtype in ["bfloat16", "float8_e4m3fn"]: + arg_dict[f"arg_{i}"] = arg.numpy().astype(np.float32) + else: + arg_dict[f"arg_{i}"] = arg.numpy() + _check_nan_inf(arg.numpy()) + np.savez(self.debug_out / f"{func_name}.npz", **arg_dict) + + self.counter += 1 + + +class DebugChat: + """A chat interface used only for debugging purpose. + + It debugs auto-regressive decoding fully in Python via the prefill and + decode interface. It supports debugging instrument (either default or + customized) to dump intermediate values for each VM function call. + + Given a prompt, it also prints out the parsed prompt, input tokens, output + tokens and output text. + + Sample usage: + + dc = DebugChat( + model="./dist/Llama-2-7b-chat-hf-q4f16_1-MLC", + debug_dir=Path("./debug-llama-2"), + model_lib="./dist/llama-2-7b-chat-q4f16_1-metal.so", + ) + dc.generate("hello world", 3) + """ + + def __init__( + self, + model: str, + model_lib: str, + debug_dir: Path, + device: Optional[str] = "auto", + debug_instrument: Optional[Any] = None, + is_image_model: Optional[bool] = False, + disable_instrument: Optional[bool] = False, + ): + """_summary_ + + Parameters + ---------- + model: str + The model folder after compiling with MLC-LLM build process. The parameter + can either be the model name with its quantization scheme + (e.g. ``Llama-2-7b-chat-hf-q4f16_1``), or a full path to the model + folder. In the former case, we will use the provided name to search + for the model folder over possible paths. + + model_lib : str + The full path to the model library file to use (e.g. a ``.so`` file). + + debug_dir: Path + The output folder to store the dumped debug files. + + device : Optional[str] + The description of the device to run on. User should provide a string in the + form of 'device_name:device_id' or 'device_name', where 'device_name' is one of + 'cuda', 'metal', 'vulkan', 'rocm', 'opencl', 'auto' (automatically detect the + local device), and 'device_id' is the device id to run on. If no 'device_id' + is provided, it will be set to 0 by default. + + chat_config : Optional[ChatConfig] + A ``ChatConfig`` instance partially filled. Will be used to override the + ``mlc-chat-config.json``. + + debug_instrument : Optional[Any] + An instrument function that will be called before/after each Call instruction. + The function have the following signature: + + .. code:: python + + def instrument( + func: Union[VMClosure, Function], + func_symbol: str, + before_run: bool, + ret_value: any, + *args) -> bool: + pass + + The instrument takes the following parameters: + - func: function object to be called. + - func_symbol: the symbol name of the function. + - before_run: whether it is before or after call. + - ret_value: the return value of the call, only valid after run. + - args: the arguments being passed to call. + + is_image_model: Optional[bool] + Whether the model support image input. If so, will look for image embedding method. + Default to False. + + disable_instrument: Optional[bool] + If true, will not use debug instrument for faster generation. Default to False. + """ + self.debug_dir = debug_dir + self.device = detect_device(device) + if disable_instrument: + self.instrument = None + else: + self.instrument = ( + debug_instrument + if debug_instrument + else DefaultDebugInstrument(debug_dir / "prefill") + ) + self.mod, self.params, self.metadata = _get_tvm_module( + model, model_lib, self.device, self.instrument + ) + self.model_path = Path(model) + self.config_file_path = self.model_path / "mlc-chat-config.json" + with open(self.config_file_path, encoding="utf-8") as file: + self.chat_config = MLCChatConfig.model_validate_json(file.read()) + + conv_template = self.chat_config.conv_template + + self.conversation = ( + ConvTemplateRegistry.get_conv_template(conv_template) + if isinstance(conv_template, str) + else conv_template + ) + self.tokenizer = Tokenizer(str(self.model_path)) + + self.add_sequence_func = tvm.get_global_func("vm.builtin.kv_state_add_sequence") + self.begin_forward_func = tvm.get_global_func("vm.builtin.kv_state_begin_forward") + self.end_forward_func = tvm.get_global_func("vm.builtin.kv_state_end_forward") + self.nd_view_func = tvm.get_global_func("vm.builtin.reshape") + self.sample_topp_from_prob_func = tvm.get_global_func("vm.builtin.sample_top_p_from_prob") + + try: + self.embed_func = self.mod["embed"] + except AttributeError as exc: + raise RuntimeError("DebugChat only supports separate embedding layer") from exc + + if is_image_model: + try: + self.embed_image_func = self.mod["image_embed"] + except AttributeError as exc: + raise RuntimeError( + "Expect the model to be an image model, but cannot find `image_embed`." + ) from exc + + self.prefill_func = self.mod["prefill"] + self.decode_func = self.mod["decode"] + self.create_kv_cache_func = None + if self.mod.implements_function("create_flashinfer_paged_kv_cache"): + self.create_kv_cache_func = self.mod["create_flashinfer_paged_kv_cache"] + elif self.mod.implements_function("create_tir_paged_kv_cache"): + self.create_kv_cache_func = self.mod["create_tir_paged_kv_cache"] + else: + # TODO: Support RNN KVState + raise RuntimeError("DebugChat cannot find create KV cache function") + + self.appeared_token_freq: Dict[int, int] = {} # noqa: UP006 + + def _preprocess_prompts( + self, prompt: str, image_url: Optional[str] = None + ) -> List[Union[List[int], data.ImageData]]: # noqa: UP006 + print("======================= Starts Tokenization & Embedding =======================") + # Step 0. Generate prompt string using conversation template + if image_url is None: + self.conversation.messages.append(("user", prompt)) + else: + self.conversation.messages.append( + ( + "user", + [ + {"type": "image_url", "image_url": image_url}, + {"type": "text", "text": prompt}, + ], + ) + ) + self.conversation.messages.append(("assistant", None)) + + with open(self.config_file_path, encoding="utf-8") as file: + config = json.load(file) + parsed_prompt = self.conversation.as_prompt(config) + print( + "Parsed prompt using conversation template " + f"{green(self.conversation.name)}: {parsed_prompt}" + ) + tokens = engine_utils.process_prompts(parsed_prompt, self.tokenizer.encode) + + if self.conversation.system_prefix_token_ids is not None: + tokens[0] = self.conversation.system_prefix_token_ids + tokens[0] + + return tokens + + def _embed( + self, + data_inputs: List[Union[List[int], data.ImageData]], # noqa: UP006 + ) -> Tuple[tvm.runtime.Tensor, int]: # noqa: UP006 + # We currently convert to numpy after embedded, concat in numpy, then convert back to + # tvm tensor; could be more optimized; but may suffice for debug purposes. + embeddings = [] + for data_input in data_inputs: + if isinstance(data_input, data.ImageData): + # Process image data + # print(f"data_input.get_embed_size(): {data_input.embed_size}") + image_input = data_input.image + if data_input.image.device != self.device: + image_input = data_input.image.copyto(self.device) + embeddings.append(self.embed_image_func(image_input, self.params).asnumpy()) + else: + # Process token data + data_input = tvm.runtime.tensor( + np.array(data_input).astype("int32"), device=self.device + ) + embeddings.append(self.embed_func(data_input, self.params).asnumpy()) + # for embedding in embeddings: + # print(f"embedding.shape: {embedding.shape}") + + # Concatenate + concat_embeddings = tvm.runtime.tensor( + np.concatenate(embeddings, axis=0), device=self.device + ) + concat_embeddings = self.nd_view_func( + concat_embeddings, + Shape([1, concat_embeddings.shape[0], concat_embeddings.shape[1]]), + ) + input_len = concat_embeddings.shape[1] + + return concat_embeddings, input_len + + def _prefill(self, embedding: tvm.runtime.Tensor, input_len: int): + print("======================= Starts Prefill =======================") + seq_len_shape = Shape([input_len]) + max_num_sequence = 1 + page_size = 16 + sliding_window_size = ( + self.chat_config.sliding_window_size + if self.chat_config.sliding_window_size + else self.metadata["sliding_window_size"] + ) + context_window_size = ( + self.chat_config.context_window_size + if self.chat_config.context_window_size + else self.metadata["context_window_size"] + ) + prefill_chunk_size = ( + self.chat_config.prefill_chunk_size + if self.chat_config.prefill_chunk_size + else self.metadata["prefill_chunk_size"] + ) + max_total_sequence_length = ( + sliding_window_size if context_window_size == -1 else context_window_size + ) + support_sliding_window = int(sliding_window_size != -1) + + kv_caches = self.create_kv_cache_func( + Shape([max_num_sequence]), + Shape([max_total_sequence_length]), + Shape([prefill_chunk_size]), + Shape([page_size]), + Shape([support_sliding_window]), + ) + self.add_sequence_func(kv_caches, 0) + self.begin_forward_func(kv_caches, Shape([0]), seq_len_shape) + logits, kv_caches = self.prefill_func(embedding, kv_caches, self.params) + self.end_forward_func(kv_caches) + return logits, kv_caches + + def _decode(self, token: int, kv_caches: Object): + embedding, _ = self._embed([[token]]) + self.begin_forward_func(kv_caches, Shape([0]), Shape([1])) + logits, kv_caches = self.decode_func(embedding, kv_caches, self.params) + self.end_forward_func(kv_caches) + return logits + + def _softmax_with_temperature(self, logits: np.ndarray, temperature: float): + # Adjust logits based on the temperature + logits = np.array(logits) / temperature + logits -= np.max(logits, axis=-1, keepdims=True) + + exp_logits = np.exp(logits, logits) + exp_logits /= np.sum(exp_logits, axis=-1, keepdims=True) + return exp_logits + + def _apply_presence_and_freq_penalty( + self, logits: np.ndarray, presence_penalty: float, freq_penalty: float + ): + for token_id, freq in self.appeared_token_freq.items(): + logits[:, :, token_id] -= freq * freq_penalty + presence_penalty + + def _sample_token_from_logits( + self, + logits: tvm.runtime.Tensor, + *, + temperature=1.0, + top_p=1.0, + presence_penalty=0.0, + frequency_penalty=0.0, + ): + logits_np = logits.numpy() + + if presence_penalty != 0.0 or frequency_penalty != 0.0: + self._apply_presence_and_freq_penalty(logits_np, presence_penalty, frequency_penalty) + + logits_np = self._softmax_with_temperature(logits_np, temperature) + if self.instrument is not None: + np.savez(self.instrument.debug_out / "logits.npz", logits_np) + + logits = logits.copyfrom(logits_np) + next_token = self.sample_topp_from_prob_func(logits, top_p, random.random()) + return next_token + + def generate( + self, + prompt: str, + generate_length: int, + image_url: Optional[str] = None, + ): + """Generates the response from the model given a user prompt. User will need to + specify the generation length for debugging purpose. For example, a generation + length of 3 will include 1 prefill step and 2 decode steps. + + Parameters + ---------- + prompt : str + The user input prompt. + + generate_length : int + How many tokens to generate. + """ + out_tokens = [] + + data_inputs = self._preprocess_prompts(prompt, image_url) + print(f"{green('Data inputs: ')}: {data_inputs}") + embedding, input_len = self._embed(data_inputs) + logits, kv_caches = self._prefill(embedding, input_len) + next_token = self._sample_token_from_logits(logits) + out_tokens.append(next_token) + if self.instrument is not None: + path_str = (self.debug_dir / "prefill").as_posix() + print(f"Debug instrument output dumped to {green(path_str)}") + + print("======================= Starts Decode =======================") + for i in range(generate_length - 1): + if self.instrument is not None: + self.instrument.reset(self.debug_dir / f"decode_{i}") + logits = self._decode(next_token, kv_caches) + next_token = self._sample_token_from_logits(logits) + out_tokens.append(next_token) + if self.instrument is not None: + path_str = (self.debug_dir / f"decode_{i}").as_posix() + print(f"Debug instrument output dumped to {green(path_str)}") + + if next_token in self.conversation.stop_token_ids: + break + + print(f"{green('Generated output tokens')}: {np.array(out_tokens)}") + + out_text = self.tokenizer.decode(out_tokens) + print(f"{green('Generated output text')}: {out_text}") + + +def main(): + """The main function to start a DebugChat CLI""" + + parser = ArgumentParser("MLC LLM Chat Debug Tool") + parser.add_argument( + "prompt", + type=str, + help="The user input prompt.", + ) + parser.add_argument( + "--generate-len", + type=int, + help="Number of output tokens to generate.", + required=True, + ) + parser.add_argument( + "--model", + type=str, + help="An MLC model directory that contains `mlc-chat-config.json`", + required=True, + ) + parser.add_argument( + "--model-lib", + type=str, + help="The full path to the model library file to use (e.g. a ``.so`` file).", + required=True, + ) + parser.add_argument( + "--debug-dir", + type=str, + help="The output folder to store the dumped debug files.", + required=True, + ) + parser.add_argument( + "--device", + type=str, + default="auto", + help=HELP["device_compile"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--image-url", + type=str, + required=False, + help="Image to prefill into the model, can only be set for image models", + ) + parser.add_argument( + "--disable-instrument", + action="store_true", + help=( + "Disable dumping customizable detailed information of kernel input " + + "and output, hence making generation faster." + ), + ) + parsed = parser.parse_args() + dc = DebugChat( + model=parsed.model, + model_lib=parsed.model_lib, + debug_dir=Path(parsed.debug_dir), + device=parsed.device, + is_image_model=parsed.image_url is not None, + disable_instrument=parsed.disable_instrument, + ) + + dc.generate(parsed.prompt, parsed.generate_len, parsed.image_url) + + +if __name__ == "__main__": + main() diff --git a/python/mlc_llm/testing/debug_compare.py b/python/mlc_llm/testing/debug_compare.py new file mode 100644 index 0000000..ab83731 --- /dev/null +++ b/python/mlc_llm/testing/debug_compare.py @@ -0,0 +1,256 @@ +"""Debug compiled models with TVM instrument""" + +import os +from pathlib import Path +from typing import Dict, List, Set, Tuple # noqa: UP035 + +import tvm +from tvm import rpc, runtime +from tvm.relax.testing.lib_comparator import LibCompareVMInstrument + +from mlc_llm.interface.help import HELP +from mlc_llm.support.argparse import ArgumentParser +from mlc_llm.testing.debug_chat import DebugChat + + +def _print_as_table(sorted_list): + print("=" * 100) + print( + "Name".ljust(50) + + "Time (ms)".ljust(12) + + "Count".ljust(8) + + "Total time (ms)".ljust(18) + + "Percentage (%)" + ) + total_time = sum(record[1][0] * record[1][1] for record in sorted_list) * 1000 + for record in sorted_list: + time = record[1][0] * 1000 + weighted_time = time * record[1][1] + percentage = weighted_time / total_time * 100 + print( + record[0].ljust(50) + + f"{time:.4f}".ljust(12) + + str(record[1][1]).ljust(8) + + f"{weighted_time:.4f}".ljust(18) + + f"{percentage:.2f}" + ) + print(f"Total time: {total_time:.4f} ms") + + +class LibCompare(LibCompareVMInstrument): + """The default debug instrument to use if users don't specify + a customized one. + + This debug instrument will dump the arguments and output of each + VM Call instruction into a .npz file. It will also alert the user + if any function outputs are NaN or INF. + + Parameters + ---------- + mod: runtime.Module + The module of interest to be validated. + + device: runtime.Device + The device to run the target module on. + + time_eval: bool + Whether to time evaluate the functions. + + rtol: float + rtol used in validation + + atol: float + atol used in validation + """ + + def __init__( + self, + mod: runtime.Module, + device: runtime.Device, + debug_out: Path, + time_eval: bool = True, + rtol: float = 1e-2, + atol: float = 1, + skip_rounds: int = 0, + ): + super().__init__(mod, device, True, rtol, atol) + self.debug_out = debug_out + self.time_eval = time_eval + self.time_eval_results: Dict[str, Tuple[float, int]] = {} # noqa: UP006 + self.visited: Set[str] = set([]) # noqa: UP006 + self.skip_rounds = skip_rounds + self.counter = 0 + debug_out.mkdir(exist_ok=True, parents=True) + + def reset(self, debug_out: Path): + """Reset the state of the Instrument class + + Note + ---- + `debug_out` is not used in this class. + + Parameters + ---------- + debug_out : Path + the directory to dump the .npz files + """ + self.debug_out = debug_out + _print_as_table( + sorted( + self.time_eval_results.items(), + key=lambda x: -(x[1][0] * x[1][1]), + ) + ) + self.time_eval_results = {} + self.visited = set([]) + self.counter = 0 + debug_out.mkdir(exist_ok=True, parents=True) + + def skip_instrument(self, func, name, before_run, ret_val, *args): + if name.startswith("shape_func"): + return True + if self.counter < self.skip_rounds: + self.counter += 1 + print(f"[{self.counter}] Skip validating {name}..") + return True + if name in self.visited: + if self.time_eval and name in self.time_eval_results: + record = self.time_eval_results[name] + self.time_eval_results[name] = (record[0], record[1] + 1) + return True + self.visited.add(name) + return False + + def compare( + self, + name: str, + ref_args: List[tvm.runtime.Tensor], # noqa: UP006 + new_args: List[tvm.runtime.Tensor], # noqa: UP006 + ret_indices: List[int], # noqa: UP006 + ): + super().compare(name, ref_args, new_args, ret_indices) + + if self.time_eval and name not in self.time_eval_results: + res = self.mod.time_evaluator( + name, + self.device, + number=20, + repeat=3, + min_repeat_ms=100, + # cache_flush_bytes=256 * 10**6 + )(*new_args) + self.time_eval_results[name] = (res.mean, 1) + print(f"Time-eval result {name} on {self.device}:\n {res}") + + +def get_instrument(args): + """Get the debug instrument from the CLI arguments""" + if args.cmp_device is None: + assert args.cmp_lib_path is None, "cmp_lib_path must be None if cmp_device is None" + args.cmp_device = args.device + args.cmp_lib_path = args.model_lib + + if args.cmp_device == "iphone": + assert args.cmp_lib_path.endswith(".dylib"), "Require a dylib file for iPhone" + proxy_host = os.environ.get("TVM_RPC_PROXY_HOST", "127.0.0.1") + proxy_port = int(os.environ.get("TVM_RPC_PROXY_PORT", "9090")) + sess = rpc.connect(proxy_host, proxy_port, "iphone") + sess.upload(args.cmp_lib_path) + lib = sess.load_module(os.path.basename(args.cmp_lib_path)) + cmp_device = sess.metal() + elif args.cmp_device == "android": + assert args.cmp_lib_path.endswith(".so"), "Require a so file for Android" + tracker_host = os.environ.get("TVM_TRACKER_HOST", "0.0.0.0") + tracker_port = int(os.environ.get("TVM_TRACKER_PORT", "9190")) + tracker = rpc.connect_tracker(tracker_host, tracker_port) + sess = tracker.request("android") + sess.upload(args.cmp_lib_path) + lib = sess.load_module(os.path.basename(args.cmp_lib_path)) + cmp_device = sess.cl(0) + else: + lib = tvm.runtime.load_module(args.cmp_lib_path) + cmp_device = tvm.device(args.cmp_device) + + return LibCompare( + lib, + cmp_device, + time_eval=args.time_eval, + debug_out=Path(args.debug_dir), + ) + + +def main(): + """The main function to start a DebugChat CLI""" + + parser = ArgumentParser("MLC LLM Chat Debug Tool") + parser.add_argument( + "prompt", + type=str, + help="The user input prompt.", + ) + parser.add_argument( + "--generate-len", + type=int, + help="Number of output tokens to generate.", + required=True, + ) + parser.add_argument( + "--model", + type=str, + help="An MLC model directory that contains `mlc-chat-config.json`", + required=True, + ) + parser.add_argument( + "--model-lib", + type=str, + help="The full path to the model library file to use (e.g. a ``.so`` file).", + required=True, + ) + parser.add_argument( + "--debug-dir", + type=str, + help="The output folder to store the dumped debug files.", + required=True, + ) + parser.add_argument( + "--device", + type=str, + default="auto", + help=HELP["device_compile"] + ' (default: "%(default)s")', + ) + parser.add_argument( + "--cmp-device", + type=str, + default="none", + ) + parser.add_argument( + "--cmp-lib-path", + type=str, + default="none", + ) + parser.add_argument( + "--time-eval", + action="store_true", + help="Whether to time evaluate the functions.", + ) + parsed = parser.parse_args() + instrument = get_instrument(parsed) + debug_chat = DebugChat( + model=parsed.model, + model_lib=parsed.model_lib, + debug_dir=Path(parsed.debug_dir), + device=parsed.device, + debug_instrument=instrument, + ) + debug_chat.generate(parsed.prompt, parsed.generate_len) + # Only print decode for now + _print_as_table( + sorted( + instrument.time_eval_results.items(), + key=lambda x: -(x[1][0] * x[1][1]), + ) + ) + + +if __name__ == "__main__": + main() diff --git a/python/mlc_llm/testing/pytest_utils.py b/python/mlc_llm/testing/pytest_utils.py new file mode 100644 index 0000000..b55a280 --- /dev/null +++ b/python/mlc_llm/testing/pytest_utils.py @@ -0,0 +1,86 @@ +"""Extra utilities to mark tests""" + +import functools +import inspect +from pathlib import Path +from typing import Callable + +import pytest + +from mlc_llm.support.constants import MLC_TEST_MODEL_PATH + + +def require_test_model(*models: str): + """Testcase decorator to require a model + + Examples + -------- + .. code:: + + @require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") + def test_reload_reset_unload(model): + # model now points to the right path + # specified by MLC_TEST_MODEL_PATH + engine = mlc_llm.MLCEngine(model) + # test code follows + + Parameters + ---------- + models : List[str] + The model directories or URLs. + """ + model_paths = [] + missing_models = [] + + for model in models: + model_path = None + for base_path in MLC_TEST_MODEL_PATH: + if (base_path / model / "mlc-chat-config.json").is_file(): + model_path = base_path / model + break + if model_path is None and (Path(model) / "mlc-chat-config.json").is_file(): + model_path = Path(model) + + if model_path is None: + missing_models.append(model) + else: + model_paths.append(str(model_path)) + + message = ( + f"Model {', '.join(missing_models)} not found in candidate paths " + f"{[str(p) for p in MLC_TEST_MODEL_PATH]}," + " if you set MLC_TEST_MODEL_PATH, please ensure model paths are in the right location," + " by default we reuse cache, try to run mlc_llm chat to download right set of models." + ) + + def _decorator(func: Callable[..., None]): + wrapped = functools.partial(func, *model_paths) + wrapped.__name__ = func.__name__ + + if inspect.iscoroutinefunction(wrapped): + # The function is a coroutine function ("async def func(...)") + @functools.wraps(wrapped) + async def wrapper(*args, **kwargs): + if len(missing_models) > 0: + print(f"{message} skipping...") + return + await wrapped(*args, **kwargs) + + else: + # The function is a normal function ("def func(...)") + @functools.wraps(wrapped) + def wrapper(*args, **kwargs): + if len(missing_models) > 0: + print(f"{message} skipping...") + return + wrapped(*args, **kwargs) + + return pytest.mark.skipif(len(missing_models) > 0, reason=message)(wrapper) + + return _decorator + + +def require_test_tokenizers(*models: str): + """Testcase decorator to require a path to tokenizers""" + # redirect to require models for now + return require_test_model(*models) diff --git a/python/mlc_llm/tokenizers/__init__.py b/python/mlc_llm/tokenizers/__init__.py new file mode 100644 index 0000000..88704b4 --- /dev/null +++ b/python/mlc_llm/tokenizers/__init__.py @@ -0,0 +1,4 @@ +"""Namespace for tokenizer rleated utilities""" + +from .streamer import StopStrHandler, TextStreamer +from .tokenizers import Tokenizer diff --git a/python/mlc_llm/tokenizers/_ffi_api.py b/python/mlc_llm/tokenizers/_ffi_api.py new file mode 100644 index 0000000..c2c29ca --- /dev/null +++ b/python/mlc_llm/tokenizers/_ffi_api.py @@ -0,0 +1,7 @@ +"""FFI APIs for mlc_llm""" + +import tvm_ffi + +# Exports functions registered via TVM_FFI_REGISTER_GLOBAL with the "mlc" prefix. +# e.g. TVM_FFI_REGISTER_GLOBAL("mlc.Tokenizer") +tvm_ffi.init_ffi_api("mlc.tokenizers", __name__) diff --git a/python/mlc_llm/tokenizers/streamer.py b/python/mlc_llm/tokenizers/streamer.py new file mode 100644 index 0000000..7554397 --- /dev/null +++ b/python/mlc_llm/tokenizers/streamer.py @@ -0,0 +1,83 @@ +"""Streamers in MLC LLM.""" + +from typing import List, Union # noqa: UP035 + +import tvm_ffi +from tvm.runtime import Object +from tvm_ffi import Shape + +from . import _ffi_api +from .tokenizers import Tokenizer + + +@tvm_ffi.register_object("mlc.TextStreamer") +class TextStreamer(Object): + """The class that streams back validated utf-8 text strings + that generated by tokenizer. + """ + + def __init__(self, tokenizer: Tokenizer) -> None: + """Create the text streamer from tokenizer""" + self.__init_handle_by_constructor__( + _ffi_api.TextStreamer, + tokenizer, + ) + + def put(self, delta_tokens: Union[List[int], Shape]) -> str: # noqa: UP006 + """Put new delta tokens into the streamer, and get the UTF-8-valid + delta string. The text streamer may hold some of the input delta tokens + which cannot decode into valid UTF-8 strings. The returned string + is always guaranteed to be UTF-8 valid. + + Parameters + ---------- + delta_tokens : Union[List[int], Shape] + The new tokens to put into the streamer. + + Returns + ------- + delta_text : str + The decoded delta string after putting the input new tokens. + """ + if isinstance(delta_tokens, list): + delta_tokens = Shape(delta_tokens) + return _ffi_api.TextStreamerPut(self, delta_tokens) + + def finish(self) -> str: + """Return the string decoded by remaining tokens.""" + return _ffi_api.TextStreamerFinish(self) + + +@tvm_ffi.register_object("mlc.StopStrHandler") +class StopStrHandler(Object): + """The stop string handler in MLC LLM, which takes input delta tokens + one at a time, and return the output delta token before stopping due to + stop strings.""" + + def __init__( + self, + stop_strs: List[str], # noqa: UP006 + tokenizer: Tokenizer, + ) -> None: + self.__init_handle_by_constructor__( + _ffi_api.StopStrHandler, + stop_strs, + tokenizer, + ) + + def put(self, token_id: int) -> List[int]: # noqa: UP006 + """Add new input delta token to the handler, return output + delta tokens before stopping. The stop string handler may hold + some of the input delta token which may be part of a stop string. + The returned tokens are always guaranteed not to be part of stop string. + """ + return list(_ffi_api.StopStrHandlerPut(self, token_id)) + + def finish(self) -> List[int]: # noqa: UP006 + """Stop string handling has finished, return remaining cached token ids.""" + return list(_ffi_api.StopStringHandlerFinish(self)) + + @property + def stop_triggered(self) -> bool: + """Check if the generation has stopped due to stop string.""" + return _ffi_api.StopStrHandlerStopTriggered(self) diff --git a/python/mlc_llm/tokenizers/tokenizers.py b/python/mlc_llm/tokenizers/tokenizers.py new file mode 100644 index 0000000..3bbd23f --- /dev/null +++ b/python/mlc_llm/tokenizers/tokenizers.py @@ -0,0 +1,127 @@ +"""The tokenizer and related tools in MLC LLM. +This tokenizer essentially wraps and binds the HuggingFace tokenizer +library and sentencepiece. +Reference: https://github.com/mlc-ai/tokenizers-cpp +""" + +import json +from dataclasses import asdict, dataclass +from typing import List, Literal # noqa: UP035 + +import tvm_ffi +from tvm.runtime import Object + +from . import _ffi_api + + +@dataclass +class TokenizerInfo: + """Useful information of the tokenizer during generation. + + Attributes + ---------- + token_postproc_method : Literal["byte_fallback", "byte_level"] + The method to post-process the tokens to their original strings. + Possible values (each refers to a kind of tokenizer): + - "byte_fallback": The same as the byte-fallback BPE tokenizer, including LLaMA-2, + Mixtral-7b, etc. E.g. "▁of" -> " of", "<0x1B>" -> "\x1b". + This method: + 1) Transform tokens like <0x1B> to hex char byte 1B. (so-called byte-fallback) + 2) Replace \\u2581 "▁" with space. + - "byte_level": The same as the byte-level BPE tokenizer, including LLaMA-3, GPT-2, + Phi-2, etc. E.g. "Ġin" -> " in", "ě" -> "\x1b" + This method inverses the bytes-to-unicode transformation in the encoding process in + https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59 + + prepend_space_in_encode : bool + Whether to prepend a space during encoding. + + strip_space_in_decode : bool + Whether to strip the first space during decoding. + """ + + token_postproc_method: Literal["byte_fallback", "byte_level"] = "byte_fallback" + prepend_space_in_encode: bool = False + strip_space_in_decode: bool = False + + def asjson(self) -> str: + """Return the config in string of JSON format.""" + return json.dumps(asdict(self)) + + @staticmethod + def from_json(json_str: str) -> "TokenizerInfo": + """Construct a config from JSON string.""" + return TokenizerInfo(**json.loads(json_str)) + + +@tvm_ffi.register_object("mlc.Tokenizer") +class Tokenizer(Object): + """The tokenizer class in MLC LLM.""" + + def __init__(self, tokenizer_path: str) -> None: + """Create the tokenizer from tokenizer directory path.""" + self.__init_handle_by_constructor__( + _ffi_api.Tokenizer, + tokenizer_path, + ) + + def encode(self, text: str) -> List[int]: # noqa: UP006 + """Encode text into ids. + + Parameters + ---------- + text : str + The text string to encode. + + Returns + ------- + token_ids : List[int] + The list of encoded token ids. + """ + return list(_ffi_api.TokenizerEncode(self, text)) + + def encode_batch(self, texts: List[str]) -> List[List[int]]: # noqa: UP006 + """Encode a batch of texts into ids. + + Parameters + ---------- + texts : List[str] + The list of text strings to encode. + + Returns + ------- + token_ids : List[List[int]] + The list of list of encoded token ids. + """ + return list(_ffi_api.TokenizerEncodeBatch(self, texts)) + + def decode(self, token_ids: List[int]) -> str: # noqa: UP006 + """Decode token ids into text. + + Parameters + ---------- + token_ids : List[int] + The token ids to decode to string. + + Returns + ------- + text : str + The decoded text string. + """ + return _ffi_api.TokenizerDecode(self, tvm_ffi.Shape(token_ids)) + + @staticmethod + def detect_tokenizer_info(tokenizer_path: str) -> TokenizerInfo: + """Detect the tokenizer info from the given path of the tokenizer. + + Parameters + ---------- + tokenizer_path : str + The tokenizer directory path. + + Returns + ------- + tokenizer_info : str + The detected tokenizer info in JSON string. + """ + return TokenizerInfo.from_json(_ffi_api.DetectTokenizerInfo(tokenizer_path)) diff --git a/python/requirements.txt b/python/requirements.txt new file mode 100644 index 0000000..a2a0973 --- /dev/null +++ b/python/requirements.txt @@ -0,0 +1,17 @@ +apache-tvm-ffi +datasets +fastapi +flashinfer-python; sys_platform == "linux" +ml_dtypes>=0.5.1 +openai +pandas +prompt_toolkit +requests +safetensors +sentencepiece +shortuuid +tiktoken +torch +tqdm +transformers +uvicorn diff --git a/python/setup.py b/python/setup.py new file mode 100644 index 0000000..2b0ee7d --- /dev/null +++ b/python/setup.py @@ -0,0 +1,141 @@ +"""Setup MLC LLM package.""" + +import os +import shutil + +from setuptools import find_packages, setup +from setuptools.dist import Distribution + +CURRENT_DIR = os.path.dirname(__file__) +CONDA_BUILD = os.getenv("CONDA_BUILD") is not None + + +def get_lib_path(): + """Get library path, name and version""" + # Directly exec libinfo to get the right setup + libinfo_py = os.path.join(CURRENT_DIR, "./mlc_llm/libinfo.py") + libinfo = {"__file__": libinfo_py} + with open(libinfo_py, "rb") as f: + exec(compile(f.read(), libinfo_py, "exec"), libinfo, libinfo) + version = libinfo["__version__"] + + # conda installs libraries into env instead of packaging with pip + if not CONDA_BUILD: + libs = [ + libinfo["find_lib_path"]("mlc_llm")[0], + libinfo["find_lib_path"]("mlc_llm_module")[0], + ] + else: + libs = None + + return libs, version + + +def git_describe_version(original_version): + """Get git describe version.""" + ver_py = os.path.join(CURRENT_DIR, "..", "version.py") + libver = {"__file__": ver_py} + with open(ver_py, "rb") as f: + exec(compile(f.read(), ver_py, "exec"), libver, libver) + _, gd_version = libver["git_describe_version"]() + if gd_version is not None and gd_version != original_version: + print(f"Use git describe based version {gd_version}") + if gd_version is None: + print(f"Use original version {original_version}") + return original_version + return gd_version + + +def parse_requirements(filename: os.PathLike): + """Parse requirements.txt.""" + with open(filename, encoding="utf-8") as f: + requirements = f.read().splitlines() + + def extract_url(line): + return next(filter(lambda x: x[0] != "-", line.split())) + + extra_URLs = [] + deps = [] + for line in requirements: + if line.startswith(("#", "-r")): + continue + + # handle -i and --extra-index-url options + if "-i " in line or "--extra-index-url" in line: + extra_URLs.append(extract_url(line)) + else: + deps.append(line) + return deps, extra_URLs + + +LIB_LIST, __version__ = get_lib_path() +__version__ = git_describe_version(__version__) + + +class BinaryDistribution(Distribution): + """This class is needed in order to create OS specific wheels.""" + + def has_ext_modules(self): + """Return True for binary distribution.""" + return True + + def is_pure(self): + """Return False for binary distribution.""" + return False + + +def main(): + """The main entrypoint.""" + setup_kwargs = {} + if not CONDA_BUILD: + with open("MANIFEST.in", "w", encoding="utf-8") as fo: + for path in LIB_LIST: + if os.path.isfile(path): + shutil.copy(path, os.path.join(CURRENT_DIR, "mlc_llm")) + _, libname = os.path.split(path) + fo.write(f"include mlc_llm/{libname}\n") + setup_kwargs = {"include_package_data": True} + + setup( + name="mlc_llm", + version=__version__, + description="MLC LLM: an universal LLM deployment engine via ML compilation.", + url="https://llm.mlc.ai/", + author="MLC LLM Contributors", + license="Apache 2.0", + # See https://pypi.org/classifiers/ + classifiers=[ + "License :: OSI Approved :: Apache Software License", + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + ], + keywords="machine learning", + zip_safe=False, + packages=find_packages(), + entry_points={ + "console_scripts": ["mlc_llm = mlc_llm.__main__:main"], + }, + package_dir={"mlc_llm": "mlc_llm"}, + install_requires=parse_requirements("requirements.txt")[0], + distclass=BinaryDistribution, + **setup_kwargs, + ) + + def _remove_path(path): + if os.path.exists(path): + if os.path.isfile(path): + os.remove(path) + elif os.path.isdir(path): + shutil.rmtree(path) + + if not CONDA_BUILD: + # Wheel cleanup + os.remove("MANIFEST.in") + for path in LIB_LIST: + _, libname = os.path.split(path) + _remove_path(f"mlc_llm/{libname}") + + +main() diff --git a/scripts/build_mlc_for_docs.sh b/scripts/build_mlc_for_docs.sh new file mode 100755 index 0000000..70f098a --- /dev/null +++ b/scripts/build_mlc_for_docs.sh @@ -0,0 +1,8 @@ +#!/bin/bash +set -euxo pipefail + +mkdir -p build +cd build +cmake .. -DCMAKE_POLICY_VERSION_MINIMUM=3.5 +make -j$(nproc) +cd - diff --git a/scripts/build_site.sh b/scripts/build_site.sh new file mode 100755 index 0000000..062f809 --- /dev/null +++ b/scripts/build_site.sh @@ -0,0 +1,10 @@ +#!/bin/bash +set -euxo pipefail + +export PYTHONPATH=$PWD/python +cd docs && make html && cd .. + +cd site && jekyll b && cd .. + +rm -rf site/_site/docs +cp -r docs/_build/html site/_site/docs diff --git a/scripts/check_url_validity.py b/scripts/check_url_validity.py new file mode 100644 index 0000000..d1192d6 --- /dev/null +++ b/scripts/check_url_validity.py @@ -0,0 +1,41 @@ +import argparse +import re +from pathlib import Path + +import requests + + +def find_urls_in_file(file_path): + with open(file_path) as file: + content = file.read() + + # Regular expression pattern to match URLs + url_pattern = re.compile( + r"http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+" + ) + + # Find all matches of URLs in the content + urls = re.findall(url_pattern, content) + return [url.strip(">") for url in urls] + + +def main(): + parser = argparse.ArgumentParser(description="Check validity of links in documentation") + parser.add_argument("--directory", type=str, default="docs", help="Directory of documentation.") + args = parser.parse_args() + + # traversal the directory and find all rst files + doc_directory = Path(args.directory) + for file_path in doc_directory.glob("**/*.rst"): + print(f"Checking {file_path}...") + for url in find_urls_in_file(file_path): + try: + r = requests.get(url) + if r.status_code == 404: + print(f"404 not found: {url}") + except Exception as e: + print(f"Error connecting {url}, error: {e}") + + +if __name__ == "__main__": + main() diff --git a/scripts/gh_deploy_site.sh b/scripts/gh_deploy_site.sh new file mode 100755 index 0000000..326c280 --- /dev/null +++ b/scripts/gh_deploy_site.sh @@ -0,0 +1,21 @@ +#!/bin/bash +# NOTE: this script is triggered by github action automatically +# when megred into main + +set -euxo pipefail + +scripts/build_mlc_for_docs.sh +scripts/build_site.sh + +git fetch +git checkout -B gh-pages origin/gh-pages +rm -rf docs .gitignore +mkdir -p docs +cp -rf site/_site/* docs +touch docs/.nojekyll + +DATE=`date` +git add docs && git commit -am "Build at ${DATE}" +git push origin gh-pages +git checkout main && git submodule update +echo "Finish deployment at ${DATE}" diff --git a/scripts/local_deploy_site.sh b/scripts/local_deploy_site.sh new file mode 100755 index 0000000..52ba40b --- /dev/null +++ b/scripts/local_deploy_site.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# NOTE: use this script to check local site + +set -euxo pipefail + +scripts/build_site.sh + +cd site && jekyll serve --skip-initial-build --host localhost --baseurl / --port 8888 diff --git a/site/.gitignore b/site/.gitignore new file mode 100644 index 0000000..6f86b47 --- /dev/null +++ b/site/.gitignore @@ -0,0 +1,5 @@ +dist +llm-chat-config.json +_includes/stable_diffusion.html +_site +.jekyll-cache diff --git a/site/CNAME b/site/CNAME new file mode 100644 index 0000000..25dbc6d --- /dev/null +++ b/site/CNAME @@ -0,0 +1 @@ +llm.mlc.ai diff --git a/site/Gemfile b/site/Gemfile new file mode 100644 index 0000000..d8b7dbb --- /dev/null +++ b/site/Gemfile @@ -0,0 +1,7 @@ +# frozen_string_literal: true + +source "https://rubygems.org" + +# gem "rails" +gem "jekyll-remote-theme" +gem "jekyll-sass-converter" diff --git a/site/_config.yml b/site/_config.yml new file mode 100644 index 0000000..9806232 --- /dev/null +++ b/site/_config.yml @@ -0,0 +1,42 @@ +name: "MLC LLM" +short_name: "MLC LLM" + +url: https://llm.mlc.ai/ + +exclude: [README.md, serve_local.sh] + +plugins: + - jekyll-remote-theme + +remote_theme: mlc-ai/jekyll-theme-mlc + + +# Colorize code snippets with the rogue module if we want to deploy on GH. +highlighter: rouge + +markdown: kramdown + +# The path structure for blog posts. +permalink: /blog/:year/:month/:day/:title.html + +# Number of news stories on the front page. +front_page_news: 8 + +# Base pathname for links. +base: '' + +# make pages for the _projects folder +collections: + projects: + output: true + +course_title: + +# Navigation bar links. +navigation: + - title: Home + link: / + - title: Docs + link: /docs + - title: Github + link: https://github.com/mlc-ai/mlc-llm diff --git a/site/_includes/arrow.svg b/site/_includes/arrow.svg new file mode 100644 index 0000000..1883ca7 --- /dev/null +++ b/site/_includes/arrow.svg @@ -0,0 +1,21 @@ + diff --git a/site/_includes/github.svg b/site/_includes/github.svg new file mode 100644 index 0000000..09a5c74 --- /dev/null +++ b/site/_includes/github.svg @@ -0,0 +1,8 @@ + + + diff --git a/site/_includes/head.html b/site/_includes/head.html new file mode 100644 index 0000000..28c070c --- /dev/null +++ b/site/_includes/head.html @@ -0,0 +1,11 @@ + + + + + diff --git a/site/_includes/hero.html b/site/_includes/hero.html new file mode 100644 index 0000000..90a27bb --- /dev/null +++ b/site/_includes/hero.html @@ -0,0 +1,43 @@ +
+
+

MLC LLM: Universal LLM Deployment Engine With ML Compilation

+
+ + + {% include github.svg %} + GitHub + {% include arrow.svg %} + + + + + + Get Started + {% include arrow.svg %} + + +
+
+
+ + + {% include project-workflow.svg %} +
+
+ + diff --git a/site/_includes/project-workflow.svg b/site/_includes/project-workflow.svg new file mode 100644 index 0000000..7541c41 --- /dev/null +++ b/site/_includes/project-workflow.svg @@ -0,0 +1,1176 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/site/assets/css/hero.scss b/site/assets/css/hero.scss new file mode 100644 index 0000000..5179a0b --- /dev/null +++ b/site/assets/css/hero.scss @@ -0,0 +1,238 @@ +--- +--- + +#hero { + background: radial-gradient(100% 50rem at center 50rem, #3351cb50, #ffffff); + padding: 3rem; + width: 100vw; + margin-left: calc(50% - 50vw); + margin-top: -20px; + display: flex; + flex-direction: column; + align-items: center; + + a { + color: black; + } + + .heading-container { + display: flex; + flex-direction: column; + align-items: center; + font-family: "Mona Sans", "MonaSansFallback", -apple-system, BlinkMacSystemFont, "Segoe UI", Helvetica, Arial, sans-serif, "Apple Color Emoji", "Segoe UI Emoji"; + margin: auto; + + a { + min-width: fit-content; + max-width: 16rem; + flex-grow: 1; + } + + h1 { + text-align: center; + font-size: 2rem; + font-weight: 700; + } + + .link-container { + display: flex; + margin-top: 2rem; + align-items: center; + flex-wrap: wrap; + font-size: 1rem; + word-break: keep-all; + font-weight: 600; + gap: 1rem; + justify-content: center; + + .github-link { + display: inline-flex; + gap: 1rem; + border-radius: 9999px; + vertical-align: middle; + align-items: center; + justify-content: center; + text-decoration: none; + cursor: pointer; + height: fit-content; + // padding: .25rem; + + .github-link-content { + width: 100%; + height: 100%; + z-index: 1; + border-radius: 9999px; + padding: 1rem 1.75rem; + background-color: #000000; + display: inline-flex; + gap: .5rem; + display: inline-flex; + justify-content: center; + color: rgb(229 229 229); + + .icon { + display: inline-flex; + align-items: center; + margin-right: .5rem; + + svg { + height: 1.5rem; + } + } + } + } + + .get-start-link { + display: inline-flex; + gap: 1rem; + background-color: white; + border-radius: 9999px; + vertical-align: middle; + align-items: center; + justify-content: center; + text-decoration: none; + cursor: pointer; + height: fit-content; + padding: .25rem; + + .get-start-link-content { + width: 100%; + height: 100%; + z-index: 1; + border-radius: 9999px; + padding: 1rem 1.75rem; + background-color: white; + display: inline-flex; + justify-content: center; + } + } + + .arrow-container { + margin-left: .25rem; + display: inline-flex; + align-items: center; + } + } + } + + .arrow-expandable { + stroke-dasharray: 10; + stroke-dashoffset: 10; + transition: stroke-dashoffset 200ms; + } + + .expanded { + .arrow-expandable { + stroke-dashoffset: 20; + } + } + + .demo-container { + position: relative; + margin-top: 96px; + width: calc(100% + 4rem); + max-width: 1024px; + flex-shrink: 0; + padding: 2rem; + + svg { + height: auto; + width: 100%; + border-radius: inherit; + } + } +} + +.moving-border { + overflow: hidden; + position: relative; + + .border { + position: absolute; + inset: -1000%; + animation: spin 3s linear infinite; + border-radius: 1rem; + background-image: conic-gradient(from 90deg at 50% 50%, #e2cbff 0, #393bb2 50%, #e2cbff 100%); + } +} + +@media screen and (min-width:640px) { + #hero { + padding: 6rem; + + .heading-container { + max-width: 40rem; + + h1 { + font-size: 3rem; + } + } + + .demo-container { + width: calc(100% + 10rem); + } + } +} + + +@media screen and (min-width:768px) { + #hero { + .heading-container { + max-width: 45rem; + + h1 { + font-size: 3.2rem; + } + + .link-container { + font-size: 1.2rem; + } + } + } +} + +@media screen and (min-width:1024px) { + #hero { + padding: 8rem; + + .heading-container { + max-width: 50rem; + + h1 { + font-size: 3.5rem; + } + } + + .demo-container { + width: 100%; + } + } + +} + +@media screen and (min-width:1280px) { + #hero { + .heading-container { + max-width: 60rem; + + h1 { + font-size: 4rem; + } + } + } +} + +@media screen and (min-width:1760px) { + #hero { + background: radial-gradient(100% 50rem at center 50rem, #3351cb50, #ffffff); + + gap: 4rem; + padding-bottom: 12rem; + } +} + +@keyframes spin { + 100% { + transform: rotate(1turn); + } +} diff --git a/site/assets/gif/android-demo.gif b/site/assets/gif/android-demo.gif new file mode 100644 index 0000000..aec883f Binary files /dev/null and b/site/assets/gif/android-demo.gif differ diff --git a/site/assets/gif/ios-demo.gif b/site/assets/gif/ios-demo.gif new file mode 100644 index 0000000..7256afe Binary files /dev/null and b/site/assets/gif/ios-demo.gif differ diff --git a/site/assets/gif/linux-demo.gif b/site/assets/gif/linux-demo.gif new file mode 100644 index 0000000..15cfc9d Binary files /dev/null and b/site/assets/gif/linux-demo.gif differ diff --git a/site/assets/img/android/android-diagram.png b/site/assets/img/android/android-diagram.png new file mode 100644 index 0000000..5f49f7c Binary files /dev/null and b/site/assets/img/android/android-diagram.png differ diff --git a/site/assets/img/android/android-studio.png b/site/assets/img/android/android-studio.png new file mode 100644 index 0000000..7c40215 Binary files /dev/null and b/site/assets/img/android/android-studio.png differ diff --git a/site/assets/img/android/android-vs-ios.png b/site/assets/img/android/android-vs-ios.png new file mode 100644 index 0000000..2436797 Binary files /dev/null and b/site/assets/img/android/android-vs-ios.png differ diff --git a/site/assets/img/android/local-advantage.png b/site/assets/img/android/local-advantage.png new file mode 100644 index 0000000..854864f Binary files /dev/null and b/site/assets/img/android/local-advantage.png differ diff --git a/site/assets/img/diag.svg b/site/assets/img/diag.svg new file mode 100644 index 0000000..c04ff4c --- /dev/null +++ b/site/assets/img/diag.svg @@ -0,0 +1 @@ + diff --git a/site/assets/img/multi-gpu/figure-1.svg b/site/assets/img/multi-gpu/figure-1.svg new file mode 100644 index 0000000..d3083cf --- /dev/null +++ b/site/assets/img/multi-gpu/figure-1.svg @@ -0,0 +1,247 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/site/assets/img/multi-gpu/figure-2.svg b/site/assets/img/multi-gpu/figure-2.svg new file mode 100644 index 0000000..70d35f5 --- /dev/null +++ b/site/assets/img/multi-gpu/figure-2.svg @@ -0,0 +1,418 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/site/assets/img/multi-gpu/figure-3.svg b/site/assets/img/multi-gpu/figure-3.svg new file mode 100644 index 0000000..078231f --- /dev/null +++ b/site/assets/img/multi-gpu/figure-3.svg @@ -0,0 +1,167 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/site/index.md b/site/index.md new file mode 100644 index 0000000..50c5eea --- /dev/null +++ b/site/index.md @@ -0,0 +1,24 @@ +--- +layout: default +title: Home +notitle: true +--- + +{% include hero.html %} + +## Overview + +MLC LLM is a machine learning compiler and high-performance deployment engine for large language models. The mission of this project is to enable everyone to develop, optimize, and deploy AI models natively on everyone's platforms.  + +MLC LLM compiles and runs code on MLCEngine -- a unified high-performance LLM inference engine across the above platforms. MLCEngine provides OpenAI-compatible API available through REST server, python, javascript, iOS, Android, all backed by the same engine and compiler that we keep improving with the community. + +## Get Started + +Please visit our [documentation](https://llm.mlc.ai/docs/) to get started with MLC LLM. +- [Installation](https://llm.mlc.ai/docs/install/mlc_llm) +- [Quick start](https://llm.mlc.ai/docs/get_started/quick_start) +- [Introduction](https://llm.mlc.ai/docs/get_started/introduction) + +## Links +- [MLC LLM Github](https://github.com/mlc-ai/mlc-llm) +- [WebLLM Project](https://webllm.mlc.ai) diff --git a/site/privacy.md b/site/privacy.md new file mode 100644 index 0000000..f7f2d29 --- /dev/null +++ b/site/privacy.md @@ -0,0 +1,10 @@ +--- +layout: default +title: Home +notitle: true +--- + +# MLC Chat App Privacy + +MLC Chat run all generation locally. +All data stays in users' device and is not collected by the app. diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..a16e4ca --- /dev/null +++ b/tests/README.md @@ -0,0 +1,8 @@ +# MLC LLM Tests + +We primarily relies on pytest to test our engine. +Most of the unit functionalities in C++ can be exposed via TVM FFI, +and tested through python environment. + +We categorize the test cases by adding `pytestmark = [pytest.mark.category_name]`. +Checkout [python/conftest.py](python/conftest.py) for categories. diff --git a/tests/cpp/conv_template_unittest.cc b/tests/cpp/conv_template_unittest.cc new file mode 100644 index 0000000..7c29c5b --- /dev/null +++ b/tests/cpp/conv_template_unittest.cc @@ -0,0 +1,129 @@ +#include "json_ffi/conv_template.h" + +#include + +namespace mlc { +namespace llm { +namespace json_ffi { + +void _TestConvTemplateLoadJSONTextContent() { + std::string conv_template = + "{\n" + " \"name\": \"test\",\n" + " \"system_template\": \"abc{system_message}\",\n" + " \"system_message\": \"de\",\n" + " \"roles\": {\n" + " \"user\": \"Instruct\",\n" + " \"assistant\": \"Output\",\n" + " \"tool\": \"Instruct\"\n" + " },\n" + " \"role_templates\": {\n" + " \"user\": \"{user_message}\",\n" + " \"assistant\": \"{assistant_message}\",\n" + " \"tool\": \"{tool_message}\"\n" + " },\n" + " \"messages\": [[\"Instruct\", \"Hello\"], [\"Output\", \"Hey\"]],\n" + " \"seps\": [\n" + " \"\\n\"\n" + " ],\n" + " \"role_content_sep\": \": \",\n" + " \"role_empty_sep\": \":\",\n" + " \"stop_str\": [\n" + " \"<|endoftext|>\"\n" + " ],\n" + " \"add_role_after_system_message\": false,\n" + " \"stop_token_ids\": [\n" + " 50256\n" + " ]" + "}"; + + auto res = Conversation::FromJSON(conv_template).IsOk(); + ASSERT_TRUE(res); + const Conversation& conv = Conversation::FromJSON(conv_template).Unwrap(); + ASSERT_EQ(conv.name, "test"); + ASSERT_EQ(conv.system_template, "abc{system_message}"); + ASSERT_EQ(conv.system_message, "de"); + ASSERT_EQ(conv.roles.at("user"), "Instruct"); + ASSERT_EQ(conv.roles.at("assistant"), "Output"); + ASSERT_EQ(conv.roles.at("tool"), "Instruct"); + ASSERT_EQ(conv.role_templates.at("user"), "{user_message}"); + ASSERT_EQ(conv.role_templates.at("assistant"), "{assistant_message}"); + ASSERT_EQ(conv.role_templates.at("tool"), "{tool_message}"); + ASSERT_EQ(conv.messages.at(0).role, "Instruct"); + ASSERT_EQ(conv.messages.at(0).content.Text(), "Hello"); + ASSERT_EQ(conv.messages.at(1).role, "Output"); + ASSERT_EQ(conv.messages.at(1).content.Text(), "Hey"); + ASSERT_EQ(conv.seps.at(0), "\n"); + ASSERT_EQ(conv.role_content_sep, ": "); + ASSERT_EQ(conv.role_empty_sep, ":"); + ASSERT_EQ(conv.stop_str.at(0), "<|endoftext|>"); + ASSERT_EQ(conv.add_role_after_system_message, false); + ASSERT_EQ(conv.stop_token_ids.at(0), 50256); +} + +void _TestConvTemplateLoadJSONPartsContent() { + std::string conv_template = + "{\n" + " \"name\": \"test\",\n" + " \"system_template\": \"abc{system_message}\",\n" + " \"system_message\": \"de\",\n" + " \"roles\": {\n" + " \"user\": \"Instruct\",\n" + " \"assistant\": \"Output\",\n" + " \"tool\": \"Instruct\"\n" + " },\n" + " \"role_templates\": {\n" + " \"user\": \"{user_message}\",\n" + " \"assistant\": \"{assistant_message}\",\n" + " \"tool\": \"{tool_message}\"\n" + " },\n" + " \"messages\": [[\"Instruct\", " + " [{\"type\": \"text\", \"text\": \"What's in the image?\"},\n" + " {\"type\": \"image_url\", \"image_url\": \"https://example.com/image.jpg\"}]\n" + " ]],\n" + " \"seps\": [\n" + " \"\\n\"\n" + " ],\n" + " \"role_content_sep\": \": \",\n" + " \"role_empty_sep\": \":\",\n" + " \"stop_str\": [\n" + " \"<|endoftext|>\"\n" + " ],\n" + " \"add_role_after_system_message\": false,\n" + " \"stop_token_ids\": [\n" + " 50256\n" + " ]" + "}"; + + auto res = Conversation::FromJSON(conv_template).IsOk(); + ASSERT_TRUE(res); + const Conversation& conv = Conversation::FromJSON(conv_template).Unwrap(); + ASSERT_EQ(conv.name, "test"); + ASSERT_EQ(conv.system_template, "abc{system_message}"); + ASSERT_EQ(conv.system_message, "de"); + ASSERT_EQ(conv.roles.at("user"), "Instruct"); + ASSERT_EQ(conv.roles.at("assistant"), "Output"); + ASSERT_EQ(conv.roles.at("tool"), "Instruct"); + ASSERT_EQ(conv.role_templates.at("user"), "{user_message}"); + ASSERT_EQ(conv.role_templates.at("assistant"), "{assistant_message}"); + ASSERT_EQ(conv.role_templates.at("tool"), "{tool_message}"); + ASSERT_EQ(conv.messages.at(0).role, "Instruct"); + ASSERT_EQ(conv.messages.at(0).content.Parts().at(0).at("type"), "text"); + ASSERT_EQ(conv.messages.at(0).content.Parts().at(0).at("text"), "What's in the image?"); + ASSERT_EQ(conv.messages.at(0).content.Parts().at(1).at("type"), "image_url"); + ASSERT_EQ(conv.messages.at(0).content.Parts().at(1).at("image_url"), + "https://example.com/image.jpg"); + ASSERT_EQ(conv.seps.at(0), "\n"); + ASSERT_EQ(conv.role_content_sep, ": "); + ASSERT_EQ(conv.role_empty_sep, ":"); + ASSERT_EQ(conv.stop_str.at(0), "<|endoftext|>"); + ASSERT_EQ(conv.add_role_after_system_message, false); + ASSERT_EQ(conv.stop_token_ids.at(0), 50256); +} + +TEST(JsonFFIConvTest, LoadJSONTextContentTest) { _TestConvTemplateLoadJSONTextContent(); } +TEST(JsonFFIConvTest, LoadJSONPartsContentTest) { _TestConvTemplateLoadJSONPartsContent(); } + +} // namespace json_ffi +} // namespace llm +} // namespace mlc diff --git a/tests/python/__init__.py b/tests/python/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/python/compiler_pass/test_fuse_ft_dequantize_matmul_epilogue.py b/tests/python/compiler_pass/test_fuse_ft_dequantize_matmul_epilogue.py new file mode 100644 index 0000000..3dbf0a5 --- /dev/null +++ b/tests/python/compiler_pass/test_fuse_ft_dequantize_matmul_epilogue.py @@ -0,0 +1,341 @@ +import tvm +from tvm.ir import assert_structural_equal +from tvm.script import ir as I +from tvm.script import relax as R + +from mlc_llm.compiler_pass.fuse_ft_dequantize_matmul_epilogue import ( + FuseFTDequantizeEpilogue, +) + + +def test_fuse_bias(): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + ): + with R.dataflow(): + lv1 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int", + ( + x, + weight, + scale, + "identity", + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + lv2 = R.add(lv1, bias) + R.output(lv2) + return lv2 + + @I.ir_module + class After: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + ) -> R.Tensor((1, 1, 1024), "float16"): + with R.dataflow(): + lv2 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias", + ( + x, + weight, + scale, + bias, + R.str("identity"), + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + R.prim_value(0), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + R.output(lv2) + return lv2 + + seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()]) + mod = seq(Before) + assert_structural_equal(mod, After) + + +def test_fuse_activation(): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + ): + with R.dataflow(): + lv1 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int", + ( + x, + weight, + scale, + "identity", + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + lv2 = R.nn.silu(lv1) + R.output(lv2) + return lv2 + + @I.ir_module + class After: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + ) -> R.Tensor((1, 1, 1024), "float16"): + with R.dataflow(): + lv2 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int", + ( + x, + weight, + scale, + R.str("silu"), + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + R.output(lv2) + return lv2 + + seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()]) + mod = seq(Before) + assert_structural_equal(mod, After) + + +def test_fuse_bias_activation(): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + ): + with R.dataflow(): + lv1 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int", + ( + x, + weight, + scale, + "identity", + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + lv2 = R.add(lv1, bias) + lv3 = R.nn.relu(lv2) + R.output(lv3) + return lv3 + + @I.ir_module + class After: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + ) -> R.Tensor((1, 1, 1024), "float16"): + with R.dataflow(): + lv2 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias", + ( + x, + weight, + scale, + bias, + R.str("relu"), + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + R.prim_value(0), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + R.output(lv2) + return lv2 + + seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()]) + mod = seq(Before) + assert_structural_equal(mod, After) + + +def test_fuse_residual_binary(): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + residual: R.Tensor((1, 1, 1024), "float16"), + ): + with R.dataflow(): + lv1 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int", + ( + x, + weight, + scale, + "identity", + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + lv2 = R.add(lv1, bias) + lv3 = R.nn.relu(lv2) + lv4 = R.multiply(lv3, residual) + R.output(lv4) + return lv4 + + @I.ir_module + class After: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + residual: R.Tensor((1, 1, 1024), "float16"), + ) -> R.Tensor((1, 1, 1024), "float16"): + with R.dataflow(): + lv2 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias_residual", + ( + x, + weight, + scale, + bias, + residual, + R.str("relu"), + R.str("multiply"), + R.str("identity"), + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + R.output(lv2) + return lv2 + + seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()]) + mod = seq(Before) + assert_structural_equal(mod, After) + + +def test_fuse_residual_unary(): + @I.ir_module + class Before: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + residual: R.Tensor((1, 1, 1024), "float16"), + ): + with R.dataflow(): + lv1 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int", + ( + x, + weight, + scale, + "identity", + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + lv2 = R.add(lv1, bias) + lv3 = R.nn.relu(lv2) + lv4 = R.add(lv3, residual) + lv5 = R.nn.gelu(lv4) + R.output(lv5) + return lv5 + + @I.ir_module + class After: + @R.function + def main( + x: R.Tensor((1, 1, 4096), "float16"), + weight: R.Tensor((4096, 512), "int8"), + scale: R.Tensor((1, 1024), "float16"), + bias: R.Tensor((1, 1, 1024), "float16"), + residual: R.Tensor((1, 1, 1024), "float16"), + ) -> R.Tensor((1, 1, 1024), "float16"): + with R.dataflow(): + lv2 = R.call_dps_packed( + "fastertransformer.gemm_fp16_int_bias_residual", + ( + x, + weight, + scale, + bias, + residual, + R.str("relu"), + R.str("plus"), + R.str("gelu"), + R.prim_value(1), + R.prim_value(1024), + R.prim_value(4096), + R.prim_value(4096), + ), + out_sinfo=R.Tensor((1, 1, 1024), "float16"), + ) + R.output(lv2) + return lv2 + + seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()]) + mod = seq(Before) + assert_structural_equal(mod, After) + + +if __name__ == "__main__": + test_fuse_bias() + test_fuse_activation() + test_fuse_bias_activation() + test_fuse_residual_binary() + test_fuse_residual_unary() diff --git a/tests/python/conftest.py b/tests/python/conftest.py new file mode 100644 index 0000000..93f0021 --- /dev/null +++ b/tests/python/conftest.py @@ -0,0 +1,41 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +def pytest_configure(config): + """Register markers""" + config.addinivalue_line( + "markers", + "unittest: unittests for modules, do not require GPU, usually run fast", + ) + config.addinivalue_line("markers", "op_correctness: unittest for op corectness, requires GPU") + config.addinivalue_line( + "markers", + ( + "engine: testing engine feature functionalities, requires model and GPU, " + "note: for most request related tests, use endpoint test instead." + ), + ) + config.addinivalue_line( + "markers", + ( + "endpoint: sending requests to a global endpoint fixture(can be an rest or API), " + "tests compatibilities of API behaviors" + ), + ) + config.addinivalue_line( + "markers", + "uncategorized: this test is not yet categorized, team should work to categorize it", + ) diff --git a/tests/python/conversation_template/test_conversation_protocol.py b/tests/python/conversation_template/test_conversation_protocol.py new file mode 100644 index 0000000..33506e2 --- /dev/null +++ b/tests/python/conversation_template/test_conversation_protocol.py @@ -0,0 +1,83 @@ +import pytest + +from mlc_llm.conversation_template import ConvTemplateRegistry +from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders + + +def get_conv_templates(): + return [ + "llama-3", + "llama-2", + "mistral_default", + "gorilla", + "gorilla-openfunctions-v2", + "chatml", + "phi-2", + "codellama_completion", + "codellama_instruct", + "rwkv_world", + ] + + +@pytest.mark.parametrize("conv_template_name", get_conv_templates()) +def test_json(conv_template_name): + template = ConvTemplateRegistry.get_conv_template(conv_template_name) + j = template.to_json_dict() + template_parsed = Conversation.from_json_dict(j) + assert template == template_parsed + + +@pytest.mark.parametrize("conv_template_name", get_conv_templates()) +def test_prompt(conv_template_name): + conversation = ConvTemplateRegistry.get_conv_template(conv_template_name) + user_msg = "test1" + assistant_msg = "test2" + prompt = "test3" + + expected_user_msg = ( + conversation.role_templates["user"] + .replace(MessagePlaceholders.USER.value, user_msg) + .replace(MessagePlaceholders.FUNCTION.value, "") + ) + + expected_prompt = ( + conversation.role_templates["user"] + .replace(MessagePlaceholders.USER.value, prompt) + .replace(MessagePlaceholders.FUNCTION.value, "") + ) + + conversation.messages.append(("user", user_msg)) + conversation.messages.append(("assistant", assistant_msg)) + conversation.messages.append(("user", prompt)) + conversation.messages.append(("assistant", None)) + res = conversation.as_prompt() + + system_msg = conversation.system_template.replace( + MessagePlaceholders.SYSTEM.value, conversation.system_message + ) + expected_final_prompt = ( + system_msg + + (conversation.seps[0] if system_msg != "" else "") + + ( + conversation.roles["user"] + conversation.role_content_sep + if conversation.add_role_after_system_message + else "" + ) + + expected_user_msg + + conversation.seps[0 % len(conversation.seps)] + + conversation.roles["assistant"] + + conversation.role_content_sep + + assistant_msg + + conversation.seps[1 % len(conversation.seps)] + + conversation.roles["user"] + + conversation.role_content_sep + + expected_prompt + + conversation.seps[0 % len(conversation.seps)] + + conversation.roles["assistant"] + + conversation.role_empty_sep + ) + assert res == expected_final_prompt + + +if __name__ == "__main__": + test_json("llama-3") diff --git a/tests/python/conversation_template/test_llama_template.py b/tests/python/conversation_template/test_llama_template.py new file mode 100644 index 0000000..cfa17fe --- /dev/null +++ b/tests/python/conversation_template/test_llama_template.py @@ -0,0 +1,43 @@ +import pytest + +from mlc_llm.conversation_template import ConvTemplateRegistry + +pytestmark = [pytest.mark.runtime_unittest] + + +# From the official Llama-3 example: +# https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-3/ +def test_llama3_prompt(): + conversation = ConvTemplateRegistry.get_conv_template("llama-3") + system_msg = "You are a helpful AI assistant for travel tips and recommendations" + user_msg1 = "What is France's capital?" + assistant_msg1 = "Bonjour! The capital of France is Paris!" + user_msg2 = "What can I do there?" + assistant_msg2 = "Paris, the City of Light, offers a romantic getaway with must-see attractions like the Eiffel Tower and Louvre Museum, romantic experiences like river cruises and charming neighborhoods, and delicious food and drink options, with helpful tips for making the most of your trip." # noqa: E501 + prompt = "Give me a detailed list of the attractions I should visit, and time it takes in each one, to plan my trip accordingly." # noqa: E501 + + conversation.system_message = system_msg + conversation.messages.append(("user", user_msg1)) + conversation.messages.append(("assistant", assistant_msg1)) + conversation.messages.append(("user", user_msg2)) + conversation.messages.append(("assistant", assistant_msg2)) + conversation.messages.append(("user", prompt)) + conversation.messages.append(("assistant", None)) + res = conversation.as_prompt() + + expected = ( + "<|start_header_id|>system<|end_header_id|>\n\n" + "You are a helpful AI assistant for travel tips and recommendations<|eot_id|>\n" + "<|start_header_id|>user<|end_header_id|>\n\n" + "What is France's capital?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + "Bonjour! The capital of France is Paris!<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n" # noqa: E501 + "What can I do there?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + "Paris, the City of Light, offers a romantic getaway with must-see attractions like the Eiffel Tower and Louvre Museum, romantic experiences like river cruises and charming neighborhoods, and delicious food and drink options, with helpful tips for making the most of your trip.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n" # noqa: E501 + "Give me a detailed list of the attractions I should visit, and time it takes in each one, to plan my trip accordingly.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" # noqa: E501 + ) + + assert res[0] == expected + + +if __name__ == "__main__": + test_llama3_prompt() diff --git a/tests/python/integration/test_model_compile.py b/tests/python/integration/test_model_compile.py new file mode 100644 index 0000000..476ea39 --- /dev/null +++ b/tests/python/integration/test_model_compile.py @@ -0,0 +1,169 @@ +import concurrent.futures as cf +import os +import shlex +import subprocess +import sys +import tempfile +from itertools import product + +import tvm + +from mlc_llm.model import MODEL_PRESETS +from mlc_llm.model import MODELS as SUPPORTED_MODELS +from mlc_llm.quantization import QUANTIZATION as SUPPORTED_QUANTS +from mlc_llm.support.constants import MLC_TEMP_DIR + +OPT_LEVEL = "O2" +DEVICE2TARGET = { + "cuda": { + "kind": "cuda", + "arch": "sm_86", + "max_threads_per_block": 1024, + "max_num_threads": 1024, + "max_shared_memory_per_block": 49152, + "thread_warp_size": 32, + }, + "rocm": { + "kind": "rocm", + "mtriple": "amdgcn-amd-amdhsa-hcc", + "mcpu": "gfx1100", + "thread_warp_size": 32, + "max_threads_per_block": 1024, + "max_num_threads": 256, + "max_shared_memory_per_block": 65536, + }, + "vulkan": { + "kind": "vulkan", + "max_threads_per_block": 1024, + "max_num_threads": 256, + "max_shared_memory_per_block": 32768, + "thread_warp_size": 1, + "supports_float32": 1, + "supports_float16": 1, + "supports_int64": 1, + "supports_int32": 1, + "supports_int16": 1, + "supports_int8": 1, + "supports_16bit_buffer": 1, + }, + "metal": "metal", + "wasm": "webgpu", + "android": "android", + "ios": "iphone", +} +DEVICE2SUFFIX = { + "cuda": "so", + "rocm": "so", + "vulkan": "so", + "metal": "dylib", + "wasm": "wasm", + "android": "tar", + "ios": "tar", +} +MODELS = list(MODEL_PRESETS.keys()) +QUANTS = [ # TODO(@junrushao): use `list(mlc_llm.quantization.QUANTIZATION.keys())` + "q0f16", + "q0f32", + "q3f16_1", + "q4f16_1", + "q4f32_1", + "q4f16_ft", +] +TENSOR_PARALLEL_SHARDS = [ + 1, +] + + +def run_command(log_file, cmd): + with open(log_file, "w", encoding="utf-8") as file: + subprocess.check_call( + cmd, + stdout=file, + stderr=subprocess.STDOUT, + ) + + +def test_model_compile(): + device = sys.argv[1] + num_workers = int(sys.argv[2]) + target = DEVICE2TARGET[device] + if not isinstance(target, str): + target = str(tvm.target.Target(target)) + suffix = DEVICE2SUFFIX[device] + + passed_cmds = [] + failed_cmds = [] + with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir: + with cf.ProcessPoolExecutor(max_workers=num_workers) as executor: + log_files = [] + cmds = [] + futures = [] + for idx, (model, quant, tp_shard) in enumerate( + product( + MODELS, + QUANTS, + TENSOR_PARALLEL_SHARDS, + ) + ): + if ( + SUPPORTED_QUANTS[quant].kind + not in SUPPORTED_MODELS[MODEL_PRESETS[model]["model_type"]].quantize + ): + continue + if not target.startswith("cuda") and quant == "q4f16_ft": + # FasterTransformer only works with cuda + continue + if "deepseek_v2" in model and "32" in quant: + # Skip f32 for deepseek v2 model for now. + continue + log_file = os.path.join(tmp_dir, f"lib{idx}.log") + cmd = [ + sys.executable, + "-m", + "mlc_llm", + "compile", + model, + "--quantization", + quant, + "--overrides", + f"tensor_parallel_shards={tp_shard}", + "--device", + target, + "--opt", + OPT_LEVEL, + "-o", + os.path.join(tmp_dir, f"lib{idx}.{suffix}"), + ] + future = executor.submit(run_command, log_file, cmd) + log_files.append(log_file) + cmds.append(cmd) + futures.append(future) + for log_file, cmd, future in zip(log_files, cmds, futures): + cmd = shlex.join(cmd) + try: + future.result() + passed_cmds.append(cmd) + print(f"[PASS] {cmd}") + except Exception: + failed_cmds.append(cmd) + print("-------------------------------") + print(f"[FAIL] {cmd}") + with open(log_file, encoding="utf-8") as file: + print(file.read()) + print("-------------------------------") + print("-------------------------------") + print(f"Total {len(passed_cmds)} passed, {len(failed_cmds)} failed.") + print("-------------------------------") + print("Passed commands:") + for cmd in passed_cmds: + print(cmd) + if failed_cmds: + print("-------------------------------") + print("Failed commands:") + for cmd in failed_cmds: + print(cmd) + sys.exit(1) + + +if __name__ == "__main__": + test_model_compile() diff --git a/tests/python/json_ffi/test_json_ffi_engine.py b/tests/python/json_ffi/test_json_ffi_engine.py new file mode 100644 index 0000000..c3c43a3 --- /dev/null +++ b/tests/python/json_ffi/test_json_ffi_engine.py @@ -0,0 +1,197 @@ +import json +from typing import Dict, List, Optional # noqa: UP035 + +import pytest +from pydantic import BaseModel + +from mlc_llm.json_ffi import JSONFFIEngine +from mlc_llm.testing import require_test_model + +# test category "engine_feature" +pytestmark = [pytest.mark.engine_feature] + + +chat_completion_prompts = [ + "What is the meaning of life?", + "Introduce the history of Pittsburgh to me. Please elaborate in detail.", + "Write a three-day Seattle travel plan. Please elaborate in detail.", + "What is Alaska famous of? Please elaborate in detail.", + "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501 + "What are the necessary components to assemble a desktop computer? Please elaborate in detail.", + "Why is Vitamin D important to human beings? Please elaborate in detail.", + "Where is milk tea originated from? Please elaborate in detail.", + "Where is the southernmost place in United States? Please elaborate in detail.", + "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501 +] + +function_calling_prompts = [ + "What is the temperature in Pittsburgh, PA?", + "What is the temperature in Tokyo, JP?", + "What is the temperature in Pittsburgh, PA and Tokyo, JP?", +] + +tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + }, + } +] + + +def run_chat_completion( + engine: JSONFFIEngine, + model: str, + prompts: List[str] = chat_completion_prompts, # noqa: UP006 + tools: Optional[List[Dict]] = None, # noqa: UP006 +): + num_requests = 2 + max_tokens = 64 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + for rid in range(num_requests): + print(f"chat completion for request {rid}") + for response in engine.chat.completions.create( + messages=[{"role": "user", "content": [{"type": "text", "text": prompts[rid]}]}], + model=model, + max_tokens=max_tokens, + n=n, + request_id=str(rid), + tools=tools, + ): + for choice in response.choices: + assert choice.delta.role == "assistant" + assert isinstance(choice.delta.content, str) + output_texts[rid][choice.index] += choice.delta.content + + # Print output. + print("Chat completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + +def run_json_schema_function_calling( + engine: JSONFFIEngine, + model: str, + prompts: List[str] = function_calling_prompts, # noqa: UP006 + tools: Optional[List[Dict]] = None, # noqa: UP006 +): + num_requests = 2 + max_tokens = 64 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + class ToolCall(BaseModel): + name: str + arguments: Dict[str, str] # noqa: UP006 + + class Schema(BaseModel): + tool_calls: List[ToolCall] # noqa: UP006 + + schema_str = json.dumps(Schema.model_json_schema()) + print("Schema str", schema_str) + + for rid in range(num_requests): + print(f"chat completion for request {rid}") + for response in engine.chat.completions.create( + messages=[ + { + "role": "system", + "content": "You are a function calling AI model. You are provided with function signatures within " # noqa: E501 + " XML tags. You may call one or more functions to assist with the user query. Don't make " # noqa: E501 + f"assumptions about what values to plug into functions. Here are the available tools: {json.dumps(tools)} " # noqa: E501 + "Do not stop calling functions until the task has been accomplished or you've reached max iteration of 10. " # noqa: E501 + "Calling multiple functions at once can overload the system and increase cost so call one function at a time please. " # noqa: E501 + "If you plan to continue with analysis, always call another function. Return a valid json object (using double " # noqa: E501 + f"quotes) in the following schema: {schema_str}", + }, + {"role": "user", "content": [{"type": "text", "text": prompts[rid]}]}, + ], + model=model, + max_tokens=max_tokens, + n=n, + request_id=str(rid), + response_format={"type": "json_object", "schema": schema_str}, + ): + for choice in response.choices: + assert choice.delta.role == "assistant" + assert isinstance(choice.delta.content, str) + output_texts[rid][choice.index] += choice.delta.content + + # Print output. + print("Chat completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + +@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_chat_completion(model): + # Create engine. + engine = JSONFFIEngine(model) + + run_chat_completion(engine, model) + + # Test malformed requests. + for response in engine._raw_chat_completion( + "malformed_string", include_usage=False, request_id="123" + ): + assert len(response.choices) == 1 + assert response.choices[0].finish_reason == "error" + + engine.terminate() + + +@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_reload_reset_unload(model): + # Create engine. + engine = JSONFFIEngine(model) + + # Run chat completion before and after reload/reset. + run_chat_completion(engine, model) + engine._test_reload() + run_chat_completion(engine, model) + engine._test_reset() + run_chat_completion(engine, model) + engine._test_unload() + + engine.terminate() + + +@require_test_model("Hermes-2-Pro-Mistral-7B-q4f16_1-MLC") +def test_json_schema_with_system_prompt(model): + engine = JSONFFIEngine(model) + + # run function calling + run_json_schema_function_calling(engine, model, function_calling_prompts, tools) + + engine.terminate() + + +if __name__ == "__main__": + test_chat_completion() + test_reload_reset_unload() + test_json_schema_with_system_prompt() diff --git a/tests/python/json_ffi/test_json_ffi_engine_image.py b/tests/python/json_ffi/test_json_ffi_engine_image.py new file mode 100644 index 0000000..9ff009c --- /dev/null +++ b/tests/python/json_ffi/test_json_ffi_engine_image.py @@ -0,0 +1,92 @@ +import base64 +from typing import Dict, List, Optional # noqa: UP035 + +import requests + +from mlc_llm.json_ffi import JSONFFIEngine +from mlc_llm.testing import require_test_model + + +def base64_encode_image(url: str) -> str: + response = requests.get(url) + response.raise_for_status() # Ensure we got a successful response + image_data = base64.b64encode(response.content) + image_data_str = image_data.decode("utf-8") + data_url = f"data:image/jpeg;base64,{image_data_str}" + return data_url + + +image_prompts = [ + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": f"{base64_encode_image('https://llava-vl.github.io/static/images/view.jpg')}", + }, + {"type": "text", "text": "What does the image represent?"}, + ], + } + ] +] + + +def run_chat_completion( + engine: JSONFFIEngine, + model: str, + prompts: List[List[Dict]] = image_prompts, # noqa: UP006 + tools: Optional[List[Dict]] = None, # noqa: UP006 +): + num_requests = 1 + max_tokens = 64 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + for rid in range(num_requests): + print(f"chat completion for request {rid}") + for response in engine.chat.completions.create( + messages=prompts[rid], + model=model, + max_tokens=max_tokens, + n=n, + request_id=str(rid), + tools=tools, + ): + for choice in response.choices: + assert choice.delta.role == "assistant" + assert isinstance(choice.delta.content[0], Dict) # noqa: UP006 + assert choice.delta.content[0]["type"] == "text" + output_texts[rid][choice.index] += choice.delta.content[0]["text"] + + # Print output. + print("Chat completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + +@require_test_model("llava-1.5-7b-hf-q4f16_1-MLC") +def test_chat_completion(): + # Create engine. + engine = JSONFFIEngine( + model, # noqa: F821 + max_total_sequence_length=1024, + ) + + run_chat_completion(engine, model) # noqa: F821 + + # Test malformed requests. + for response in engine._raw_chat_completion("malformed_string", n=1, request_id="123"): + assert len(response.choices) == 1 + assert response.choices[0].finish_reason == "error" + + engine.terminate() + + +if __name__ == "__main__": + test_chat_completion() diff --git a/tests/python/json_ffi/test_json_ffi_engine_mock.py b/tests/python/json_ffi/test_json_ffi_engine_mock.py new file mode 100644 index 0000000..78ff5ad --- /dev/null +++ b/tests/python/json_ffi/test_json_ffi_engine_mock.py @@ -0,0 +1,105 @@ +import json + +import pytest +import tvm + +from mlc_llm.json_ffi import JSONFFIEngine +from mlc_llm.testing import require_test_model + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +def check_error_handling(engine, expect_str, **params): + """Check error handling in raw completion API""" + body = { + "messages": [{"role": "user", "content": "hello"}], + "stream_options": {"include_usage": True}, + } + body.update(params) + + for response in engine._raw_chat_completion( + json.dumps(body), include_usage=False, request_id="123" + ): + if response.choices[0].finish_reason is not None: + break + if response.choices[0].finish_reason != "error": + raise RuntimeError(f"expect the request {params} to hit an error") + + if expect_str not in response.choices[0].delta.content: + raise RuntimeError( + f"expect '{expect_str}' in error msg, but get '{response.choices[0].delta.content}'" + ) + + +# NOTE: we only need tokenizers in folder +# launch time of mock test is fast so we can put it in unittest +@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC") +def test_chat_completion_misuse(model: str): + engine = JSONFFIEngine(model, tvm.cpu(), model_lib="mock://echo") + # Test malformed requests. + for response in engine._raw_chat_completion( + "malformed_string", include_usage=False, request_id="123" + ): + assert len(response.choices) == 1 + assert response.choices[0].finish_reason == "error" + # check parameters + check_error_handling(engine, "should be non-negative", temperature=-1) + check_error_handling(engine, "in range [0, 1]", top_p=100) + check_error_handling(engine, "frequency_penalty", frequency_penalty=100) + + +def check_normal_param_passing(engine): + json_schema = """ + {"properties": {"result": {"items": {"type": "Integer"}, "title": "Result", "type": "array"}}, + "required": ["result"], "title": "Output", "type": "object"} + """ + param_dict = { + "top_p": 0.6, + "temperature": 0.8, + "frequency_penalty": 0.1, + "presence_penalty": 0.1, + } + usage = None + for response in engine.chat.completions.create( + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, + response_format={"type": "json_object", "schema": json_schema}, + **param_dict, + ): + if response.usage is not None: + usage = response.usage + + # echo mock will echo back the generation config + for k, v in param_dict.items(): + assert usage.extra[k] == v, f"{k} mismatch" + assert "response_format" in usage.extra + assert usage.extra["response_format"]["type"] == "json_object" + assert "schema" in usage.extra["response_format"] + + +def check_n_generation(engine): + hit_set = set() + for response in engine.chat.completions.create( + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, + n=3, + ): + for choice in response.choices: + hit_set.add(choice.index) + for i in range(3): + assert i in hit_set, f"{i} not in n generation" + + +@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC") +def test_chat_completion_api(model: str): + engine = JSONFFIEngine(model, tvm.cpu(), model_lib="mock://echo") + check_normal_param_passing(engine) + check_n_generation(engine) + + +if __name__ == "__main__": + test_chat_completion_api() + test_chat_completion_misuse() diff --git a/tests/python/loader/test_awq.py b/tests/python/loader/test_awq.py new file mode 100644 index 0000000..4717b66 --- /dev/null +++ b/tests/python/loader/test_awq.py @@ -0,0 +1,39 @@ +from pathlib import Path +from typing import Union + +import pytest +import tvm + +from mlc_llm.loader import HuggingFaceLoader +from mlc_llm.model import MODEL_PRESETS, MODELS +from mlc_llm.quantization import QUANTIZATION +from mlc_llm.support import logging, tqdm + +logging.enable_logging() + + +@pytest.mark.parametrize( + "param_path", + [ + "./dist/models/llama-2-7b-w4-g128-awq.pt", + "./dist/models/Llama-2-7B-AWQ/model.safetensors", + ], +) +def test_load_llama(param_path: Union[str, Path]): + path_params = Path(param_path) + + model = MODELS["llama"] + quantization = QUANTIZATION["q4f16_awq"] + config = model.config.from_dict(MODEL_PRESETS["llama2_7b"]) + loader = HuggingFaceLoader( + path=path_params, + extern_param_map=model.source["awq"](config, quantization), + ) + with tqdm.redirect(): + for _name, _param in loader.load(tvm.device("cpu")): + ... + + +if __name__ == "__main__": + test_load_llama(param_path="./dist/models/llama-2-7b-w4-g128-awq.pt") + test_load_llama(param_path="./dist/models/Llama-2-7B-AWQ/model.safetensors") diff --git a/tests/python/loader/test_huggingface.py b/tests/python/loader/test_huggingface.py new file mode 100644 index 0000000..9ec812b --- /dev/null +++ b/tests/python/loader/test_huggingface.py @@ -0,0 +1,68 @@ +from pathlib import Path +from typing import Union + +import pytest +import tvm + +from mlc_llm.loader import HuggingFaceLoader +from mlc_llm.model import MODELS +from mlc_llm.support import logging, tqdm + +logging.enable_logging() + + +@pytest.mark.parametrize( + "base_path", + [ + "./dist/models/Llama-2-7b-hf", + "./dist/models/Llama-2-13b-hf", + "./dist/models/Llama-2-70b-hf", + ], +) +def test_load_torch_llama(base_path: Union[str, Path]): + base_path = Path(base_path) + path_config = base_path / "config.json" + path_params = base_path / "pytorch_model.bin.index.json" + + model = MODELS["llama"] + config = model.config.from_file(path_config) + loader = HuggingFaceLoader( + path=path_params, + extern_param_map=model.source["huggingface-torch"](config, None), + ) + with tqdm.redirect(): + for _name, _param in loader.load(device=tvm.device("cpu")): + return # To reduce the time of the test + + +@pytest.mark.parametrize( + "base_path", + [ + "./dist/models/Llama-2-7b-hf", + "./dist/models/Llama-2-13b-hf", + "./dist/models/Llama-2-70b-hf", + ], +) +def test_load_safetensor_llama(base_path: Union[str, Path]): + base_path = Path(base_path) + path_config = base_path / "config.json" + path_params = base_path / "model.safetensors.index.json" + + model = MODELS["llama"] + config = model.config.from_file(path_config) + loader = HuggingFaceLoader( + path=path_params, + extern_param_map=model.source["huggingface-safetensor"](config, None), + ) + with tqdm.redirect(): + for _name, _param in loader.load(device=tvm.device("cpu")): + return # To reduce the time of the test + + +if __name__ == "__main__": + test_load_torch_llama(base_path="./dist/models/Llama-2-7b-hf") + test_load_torch_llama(base_path="./dist/models/Llama-2-13b-hf") + test_load_torch_llama(base_path="./dist/models/Llama-2-70b-hf") + test_load_safetensor_llama(base_path="./dist/models/Llama-2-7b-hf") + test_load_safetensor_llama(base_path="./dist/models/Llama-2-13b-hf") + test_load_safetensor_llama(base_path="./dist/models/Llama-2-70b-hf") diff --git a/tests/python/model/test_gemma3.py b/tests/python/model/test_gemma3.py new file mode 100644 index 0000000..cec70f0 --- /dev/null +++ b/tests/python/model/test_gemma3.py @@ -0,0 +1,70 @@ +"""Unit tests for Gemma3 model architecture.""" + +import pytest + +from mlc_llm.model import MODEL_PRESETS, MODELS + + +def test_gemma3_model_registered(): + """Verify Gemma3 model is in the registry.""" + assert "gemma3" in MODELS, "gemma3 should be registered in MODELS" + + +@pytest.mark.parametrize( + "model_name", + [ + "gemma3_2b", + "gemma3_9b", + ], +) +def test_gemma3_creation(model_name: str): + """Test Gemma3 model creation and export to TVM IR. + + Verifies: + - Config can be loaded from preset + - Model instance can be created + - Model exports to TVM IR successfully + - Named parameters are extracted + """ + model_info = MODELS["gemma3"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + model = model_info.model(config) + mod, named_params = model.export_tvm( + spec=model.get_default_spec(), + ) + + # Verify export succeeded + assert mod is not None + assert len(named_params) > 0 + + # Optional: show module structure + mod.show(black_format=False) + + # Print parameters for debugging + for name, param in named_params: + print(name, param.shape, param.dtype) + + +def test_gemma3_config_validation(): + """Test Gemma3 configuration has required fields.""" + model_info = MODELS["gemma3"] + config = model_info.config.from_dict(MODEL_PRESETS["gemma3_2b"]) + + # Check required config parameters + assert hasattr(config, "hidden_size") and config.hidden_size > 0 + assert hasattr(config, "num_hidden_layers") and config.num_hidden_layers > 0 + assert hasattr(config, "num_attention_heads") and config.num_attention_heads > 0 + assert hasattr(config, "vocab_size") and config.vocab_size > 0 + + print( + f"Gemma3 Config: hidden_size={config.hidden_size}, " + f"layers={config.num_hidden_layers}, " + f"heads={config.num_attention_heads}, " + f"vocab={config.vocab_size}" + ) + + +if __name__ == "__main__": + # Allow running tests directly + test_gemma3_creation("gemma3_2b") + test_gemma3_creation("gemma3_9b") diff --git a/tests/python/model/test_gpt2.py b/tests/python/model/test_gpt2.py new file mode 100644 index 0000000..5423b48 --- /dev/null +++ b/tests/python/model/test_gpt2.py @@ -0,0 +1,20 @@ +import pytest + +from mlc_llm.model import MODEL_PRESETS, MODELS + + +@pytest.mark.parametrize("model_name", ["gpt2"]) +def test_gpt2_creation(model_name: str): + model_info = MODELS["gpt2"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + model = model_info.model(config) + mod, named_params = model.export_tvm( + spec=model.get_default_spec(), + ) + mod.show(black_format=False) + for name, param in named_params: + print(name, param.shape, param.dtype) + + +if __name__ == "__main__": + test_gpt2_creation("gpt2") diff --git a/tests/python/model/test_gptNeox.py b/tests/python/model/test_gptNeox.py new file mode 100644 index 0000000..0f54da3 --- /dev/null +++ b/tests/python/model/test_gptNeox.py @@ -0,0 +1,20 @@ +import pytest + +from mlc_llm.model import MODEL_PRESETS, MODELS + + +@pytest.mark.parametrize("model_name", ["redpajama_3b_v1"]) +def test_mistral_creation(model_name: str): + model_info = MODELS["gpt_neox"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + model = model_info.model(config) + mod, named_params = model.export_tvm( + spec=model.get_default_spec(), + ) + mod.show(black_format=False) + for name, param in named_params: + print(name, param.shape, param.dtype) + + +if __name__ == "__main__": + test_mistral_creation("redpajama_3b_v1") diff --git a/tests/python/model/test_kv_cache.py b/tests/python/model/test_kv_cache.py new file mode 100644 index 0000000..1f98548 --- /dev/null +++ b/tests/python/model/test_kv_cache.py @@ -0,0 +1,114 @@ +import tvm +from tvm import tirx +from tvm.relax.frontend.nn import core, modules, spec +from tvm.script import ir as I +from tvm.script import relax as R +from tvm.script import tirx as T + +from mlc_llm.nn.kv_cache import PagedKVCache, RopeMode + +# mypy: disable-error-code="attr-defined" + + +def test_nn_module_paged_kv_cache(): + # fmt: off + @I.ir_module + class Module: + @R.function + def create_paged_kv_cache( + max_batch_size: R.Shape(["max_batch_size_1"]), + max_total_seq_len: R.Shape(["max_total_seq_len_1"]), + prefill_chunk_size: R.Shape(["prefill_chunk_size_1"]), + page_size: R.Shape(["page_size_1"]), + support_sliding_window: R.Shape(["support_sliding_window_1"]), + ) -> R.Object: + max_batch_size_1 = T.int64() + max_total_seq_len_1 = T.int64() + prefill_chunk_size_1 = T.int64() + page_size_1 = T.int64() + support_sliding_window_1 = T.int64() + R.func_attr({"num_input": 5}) + with R.dataflow(): + paged_kv_cache: R.Object = R.call_pure_packed("mlc.create_paged_kv_cache_generic", R.shape([max_batch_size_1, max_total_seq_len_1, prefill_chunk_size_1, page_size_1, support_sliding_window_1]), R.prim_value(32), R.prim_value(32), R.prim_value(32), R.prim_value(128), R.prim_value(1), R.prim_value(1), R.prim_value(10000), R.prim_value(128), R.dtype("float16"), sinfo_args=(R.Object,)) # noqa: E501 + gv1: R.Object = paged_kv_cache + R.output(gv1) + return gv1 + + @R.function + def forward( + cache: R.Object, qkv: R.Tensor((1, 100, 96, 128), dtype="float16") + ) -> R.Tensor((1, 100, 32, 128), dtype="float16"): + R.func_attr({"num_input": 2}) + with R.dataflow(): + reshape: R.Tensor((100, 96, 128), dtype="float16") = R.reshape( + qkv, R.shape([100, 96, 128]) + ) + lv = R.call_dps_packed( + "vm.builtin.attention_kv_cache_attention_with_fused_qkv", + (cache, R.prim_value(0), R.prim_value(T.float32(1)), reshape), + out_sinfo=R.Tensor((100, 32, 128), dtype="float16"), + ) + reshape1: R.Tensor((1, 100, 32, 128), dtype="float16") = R.reshape( + lv, R.shape([1, 100, 32, 128]) + ) + gv: R.Tensor((1, 100, 32, 128), dtype="float16") = reshape1 + R.output(gv) + return gv + # fmt: on + + class PagedKVCacheTest(modules.Module): + def forward( + self, + cache: PagedKVCache, + qkv: core.Tensor, + ) -> core.Tensor: + return cache.attention_with_fused_qkv(0, qkv, num_qo_heads=32, sm_scale=128**-0.5) + + def create_paged_kv_cache( + self, + max_batch_size: tirx.Var, + max_total_seq_len: tirx.Var, + prefill_chunk_size: tirx.Var, + page_size: tirx.Var, + support_sliding_window: tirx.Var, + ) -> PagedKVCache: + return PagedKVCache.create_generic( + attn_kind="mha", + max_batch_size=max_batch_size, + max_total_seq_len=max_total_seq_len, + prefill_chunk_size=prefill_chunk_size, + page_size=page_size, + support_sliding_window=support_sliding_window, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + qk_head_dim=128, + v_head_dim=128, + rope_mode=RopeMode.NORMAL, + rope_scale=1, + rope_theta=10000, + rotary_dim=128, + dtype="float16", + ) + + export_results = PagedKVCacheTest().export_tvm( + spec={ + "forward": { + "cache": spec.Object(object_type=PagedKVCache), + "qkv": spec.Tensor((1, 100, 96, 128), "float16"), + }, + "create_paged_kv_cache": { + "max_batch_size": int, + "max_total_seq_len": int, + "prefill_chunk_size": int, + "page_size": int, + "support_sliding_window": int, + }, + }, + ) + tvm_mod = export_results[0] + tvm.ir.assert_structural_equal(tvm_mod, Module, True) + + +if __name__ == "__main__": + test_nn_module_paged_kv_cache() diff --git a/tests/python/model/test_llama.py b/tests/python/model/test_llama.py new file mode 100644 index 0000000..bdcc156 --- /dev/null +++ b/tests/python/model/test_llama.py @@ -0,0 +1,25 @@ +import pytest + +from mlc_llm.model import MODEL_PRESETS, MODELS + + +@pytest.mark.parametrize( + "model_name", ["llama2_7b", "llama2_13b", "llama2_70b", "tinyllama_1b_chat_v1.0"] +) +def test_llama2_creation(model_name: str): + model_info = MODELS["llama"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + model = model_info.model(config) + mod, named_params = model.export_tvm( + spec=model.get_default_spec(), + ) + mod.show(black_format=False) + for name, param in named_params: + print(name, param.shape, param.dtype) + + +if __name__ == "__main__": + test_llama2_creation("llama2_7b") + test_llama2_creation("llama2_13b") + test_llama2_creation("llama2_70b") + test_llama2_creation("tinyllama_1b_chat_v1") diff --git a/tests/python/model/test_llama_quantization.py b/tests/python/model/test_llama_quantization.py new file mode 100644 index 0000000..5bae2cd --- /dev/null +++ b/tests/python/model/test_llama_quantization.py @@ -0,0 +1,72 @@ +import pytest + +from mlc_llm.model import MODEL_PRESETS, MODELS +from mlc_llm.quantization import QUANTIZATION +from mlc_llm.quantization.group_quantization import ( + GroupQuantizeEmbedding, + GroupQuantizeLinear, +) + + +@pytest.mark.parametrize( + "model_name", + ["llama2_7b", "llama2_13b", "llama2_70b"], +) +@pytest.mark.parametrize( + "quant_name", + ["q3f16_1", "q4f16_1", "q4f32_1"], +) +def test_llama2_group_quantization(model_name: str, quant_name: str): + model_info = MODELS["llama"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + model, quant_map = model_info.quantize["group-quant"](config, QUANTIZATION[quant_name]) + assert "model.embed_tokens.weight" in quant_map.param_map + assert isinstance( + model.model.embed_tokens, + GroupQuantizeEmbedding, + ) + assert "lm_head.weight" in quant_map.param_map + assert isinstance(model.lm_head, GroupQuantizeLinear) + for i in range(config.num_hidden_layers): + assert f"model.layers.{i}.self_attn.qkv_proj.weight" in quant_map.param_map + assert isinstance( + model.model.layers[i].self_attn.qkv_proj, + GroupQuantizeLinear, + ) + assert f"model.layers.{i}.self_attn.o_proj.weight" in quant_map.param_map + assert isinstance( + model.model.layers[i].self_attn.o_proj, + GroupQuantizeLinear, + ) + assert f"model.layers.{i}.mlp.gate_up_proj.weight" in quant_map.param_map + assert isinstance( + model.model.layers[i].mlp.gate_up_proj, + GroupQuantizeLinear, + ) + assert f"model.layers.{i}.mlp.down_proj.weight" in quant_map.param_map + assert isinstance( + model.model.layers[i].mlp.down_proj, + GroupQuantizeLinear, + ) + + +@pytest.mark.parametrize( + "model_name", + ["llama2_7b", "llama2_13b", "llama2_70b"], +) +@pytest.mark.parametrize( + "quant_name", + ["q0f16", "q0f32"], +) +def test_llama2_no_quantization(model_name: str, quant_name: str): + model_info = MODELS["llama"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + _, quant_map = model_info.quantize["no-quant"](config, QUANTIZATION[quant_name]) + assert len(quant_map.param_map) == 0 + assert len(quant_map.map_func) == 0 + + +if __name__ == "__main__": + test_llama2_group_quantization("llama2_7b", "q4f16_1") + test_llama2_group_quantization("llama2_13b", "q4f16_1") + test_llama2_group_quantization("llama2_70b", "q4f16_1") diff --git a/tests/python/model/test_mistral.py b/tests/python/model/test_mistral.py new file mode 100644 index 0000000..20fa019 --- /dev/null +++ b/tests/python/model/test_mistral.py @@ -0,0 +1,20 @@ +import pytest + +from mlc_llm.model import MODEL_PRESETS, MODELS + + +@pytest.mark.parametrize("model_name", ["mistral_7b"]) +def test_mistral_creation(model_name: str): + model_info = MODELS["mistral"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + model = model_info.model(config) + mod, named_params = model.export_tvm( + spec=model.get_default_spec(), + ) + mod.show(black_format=False) + for name, param in named_params: + print(name, param.shape, param.dtype) + + +if __name__ == "__main__": + test_mistral_creation("mistral_7b") diff --git a/tests/python/model/test_phi.py b/tests/python/model/test_phi.py new file mode 100644 index 0000000..0f6b506 --- /dev/null +++ b/tests/python/model/test_phi.py @@ -0,0 +1,21 @@ +import pytest + +from mlc_llm.model import MODEL_PRESETS, MODELS + + +@pytest.mark.parametrize("model_name", ["phi-1_5", "phi-2"]) +def test_phi_creation(model_name: str): + model_info = MODELS["phi-msft"] + config = model_info.config.from_dict(MODEL_PRESETS[model_name]) + model = model_info.model(config) + mod, named_params = model.export_tvm( + spec=model.get_default_spec(), + ) + mod.show(black_format=False) + for name, param in named_params: + print(name, param.shape, param.dtype) + + +if __name__ == "__main__": + test_phi_creation("phi-1_5") + test_phi_creation("phi-2") diff --git a/tests/python/model/test_qwen3_embedding.py b/tests/python/model/test_qwen3_embedding.py new file mode 100644 index 0000000..57e9a71 --- /dev/null +++ b/tests/python/model/test_qwen3_embedding.py @@ -0,0 +1,147 @@ +import json +import os + +import numpy as np +import pytest +import torch +import tvm +from safetensors import safe_open +from transformers import AutoModel, AutoTokenizer +from tvm import relax +from tvm.contrib import tvmjs +from tvm.runtime import ShapeTuple +from tvm.runtime.vm import VirtualMachine + +MLC_QWEN3_EMB_HF_DIR = os.environ.get("MLC_QWEN3_EMB_HF_DIR") +MLC_QWEN3_EMB_MODEL_DIR = os.environ.get("MLC_QWEN3_EMB_MODEL_DIR") +MLC_QWEN3_EMB_MODEL_LIB = os.environ.get("MLC_QWEN3_EMB_MODEL_LIB") +MLC_QWEN3_EMB_DEVICE = os.environ.get("MLC_QWEN3_EMB_DEVICE", "cuda") + +_skip = not all([MLC_QWEN3_EMB_HF_DIR, MLC_QWEN3_EMB_MODEL_DIR, MLC_QWEN3_EMB_MODEL_LIB]) +_skip_reason = ( + "Set MLC_QWEN3_EMB_HF_DIR, MLC_QWEN3_EMB_MODEL_DIR, MLC_QWEN3_EMB_MODEL_LIB to run this test" +) + +TEST_TEXTS = [ + "What is machine learning?", + "CMU is Carnegie Mellon University", + "机器学习是人工智能的一个分支", + "量子コンピュータの基本原理を説明してください", + "머신러닝은 인공지능의 한 분야입니다.", + ( + "Instruct: Given a web search query, retrieve relevant passages " + "that answer the query\nQuery: What is the capital of China?" + ), + ( + "The Transformer architecture, introduced in the paper Attention Is All You Need, " + "revolutionized natural language processing by replacing recurrent layers with " + "self-attention mechanisms. This allows the model to process all positions in a " + "sequence simultaneously rather than sequentially, leading to significant improvements " + "in both training efficiency and the ability to capture long-range dependencies. " + "The key innovation is the multi-head attention mechanism, which allows the model " + "to jointly attend to information from different representation subspaces at " + "different positions." + ), + "Hello", + "def fibonacci(n): return n if n <= 1 else fibonacci(n-1) + fibonacci(n-2)", +] + + +def _load_embed_weight(hf_dir): + safetensor_files = [f for f in os.listdir(hf_dir) if f.endswith(".safetensors")] + for sf in safetensor_files: + with safe_open(os.path.join(hf_dir, sf), framework="pt", device="cpu") as f: + if "embed_tokens.weight" in f.keys(): + return f.get_tensor("embed_tokens.weight") + raise FileNotFoundError(f"embed_tokens.weight not found in {hf_dir}") + + +def _hf_logits(text, tokenizer, hf_model, embed_weight): + inputs = tokenizer(text, return_tensors="pt") + with torch.no_grad(): + hidden = hf_model(**inputs).last_hidden_state.float() + logits = hidden @ embed_weight.float().T + return logits[0, -1, :].numpy() + + +def _mlc_logits(text, tokenizer, mlc_module, params, metadata, dev, embed_weight): + input_ids = tokenizer(text, return_tensors="pt")["input_ids"][0].numpy().astype(np.int32) + seq_len = len(input_ids) + + embed_func = mlc_module["embed"] + prefill_func = mlc_module["prefill_to_last_hidden_states"] + + if mlc_module.implements_function("create_flashinfer_paged_kv_cache"): + create_kv = mlc_module["create_flashinfer_paged_kv_cache"] + elif mlc_module.implements_function("create_tir_paged_kv_cache"): + create_kv = mlc_module["create_tir_paged_kv_cache"] + else: + raise RuntimeError("Cannot find KV cache creation function") + + sliding_window = metadata.get("sliding_window_size", -1) + context_window = metadata.get("context_window_size", 32768) + prefill_chunk = metadata.get("prefill_chunk_size", 2048) + max_seq_len = sliding_window if context_window == -1 else context_window + + kv_cache = create_kv( + ShapeTuple([1]), + ShapeTuple([max_seq_len]), + ShapeTuple([prefill_chunk]), + ShapeTuple([16]), + ShapeTuple([int(sliding_window != -1)]), + ) + + nd_view = tvm.get_global_func("vm.builtin.reshape") + add_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence") + begin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward") + end_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward") + + tokens_tvm = tvm.runtime.tensor(input_ids, device=dev) + embedding = embed_func(tokens_tvm, params) + embedding = nd_view(embedding, ShapeTuple([1, seq_len, embedding.shape[-1]])) + + add_sequence(kv_cache, 0) + begin_forward(kv_cache, ShapeTuple([0]), ShapeTuple([seq_len])) + hidden_states, _ = prefill_func(embedding, kv_cache, params) + end_forward(kv_cache) + + # Compute logits from hidden states using embed_tokens weight (tie_word_embeddings) + hidden = hidden_states.numpy().astype(np.float32) + logits = hidden @ embed_weight.float().numpy().T + return logits[0, -1, :] + + +@pytest.mark.skipif(_skip, reason=_skip_reason) +def test_mlc_hf_logit_match(): + tokenizer = AutoTokenizer.from_pretrained(MLC_QWEN3_EMB_HF_DIR, padding_side="left") + hf_model = AutoModel.from_pretrained(MLC_QWEN3_EMB_HF_DIR) + embed_weight = _load_embed_weight(MLC_QWEN3_EMB_HF_DIR) + + dev = tvm.runtime.device(MLC_QWEN3_EMB_DEVICE, 0) + ex = tvm.runtime.load_module(MLC_QWEN3_EMB_MODEL_LIB) + vm = relax.VirtualMachine(ex, dev) + mlc_module = vm.module + + metadata = json.loads(VirtualMachine(ex, tvm.runtime.device("cpu"))["_metadata"]()) + params_dict, _ = tvmjs.load_tensor_cache(MLC_QWEN3_EMB_MODEL_DIR, dev) + param_names = [p["name"] for p in metadata["params"]] + params = [params_dict[name] for name in param_names] + + for text in TEST_TEXTS: + hf = _hf_logits(text, tokenizer, hf_model, embed_weight) + mlc = _mlc_logits(text, tokenizer, mlc_module, params, metadata, dev, embed_weight) + + cos_sim = np.dot(hf, mlc) / (np.linalg.norm(hf) * np.linalg.norm(mlc)) + assert cos_sim > 0.99, f"[{text[:30]}] Cosine similarity {cos_sim:.6f} below 0.99" + + max_diff = np.max(np.abs(hf - mlc)) + assert max_diff < 1.0, f"[{text[:30]}] Max absolute diff {max_diff:.6e} exceeds 1.0" + + hf_top10 = set(np.argsort(hf)[-10:]) + mlc_top10 = set(np.argsort(mlc)[-10:]) + overlap = len(hf_top10 & mlc_top10) + assert overlap >= 7, f"[{text[:30]}] Top-10 overlap {overlap}/10 below 7" + + +if __name__ == "__main__": + test_mlc_hf_logit_match() diff --git a/tests/python/op/test_batch_spec_verify.py b/tests/python/op/test_batch_spec_verify.py new file mode 100644 index 0000000..40bd292 --- /dev/null +++ b/tests/python/op/test_batch_spec_verify.py @@ -0,0 +1,163 @@ +import numpy as np +import pytest +import tvm +import tvm.testing + +from mlc_llm.op.batch_spec_verify import batch_spec_verify + +# test category "op_correctness" +pytestmark = [pytest.mark.op_correctness] + + +@pytest.mark.parametrize("nbatch", [32, 64]) +@pytest.mark.parametrize("vocab", [3, 32, 64, 32000, 33, 65, 32001, 128000]) +@pytest.mark.parametrize("plist", [[0.5, 0.5], [1, 0], [0, 1]]) +def test_batch_spec_verify(nbatch, vocab, plist): + def numpy_reference( + draft_probs, + draft_tokens, + model_probs, + token_tree_first_child, + token_tree_next_sibling, + uniform_samples, + token_tree_parent_ptr, + ): + nbatch = token_tree_parent_ptr.shape[0] + for b in range(nbatch): + parent_ptr = token_tree_parent_ptr[b] + child_ptr = token_tree_first_child[parent_ptr] + while child_ptr != -1: + child_token = draft_tokens[child_ptr] + p_child = model_probs[parent_ptr, child_token] + q_child = draft_probs[child_ptr, child_token] + uniform_sample = uniform_samples[child_ptr] + if p_child / q_child >= uniform_sample: + parent_ptr = child_ptr + child_ptr = token_tree_first_child[child_ptr] + else: + model_probs[parent_ptr, :] = np.maximum( + model_probs[parent_ptr, :] - draft_probs[child_ptr, :], 0.0 + ) + psum = np.sum(model_probs[parent_ptr, :]) + model_probs[parent_ptr, :] /= psum + child_ptr = token_tree_next_sibling[child_ptr] + token_tree_parent_ptr[b] = parent_ptr + + np.random.seed(0) + + def gen_chain(num_nodes, base): + token_tree_first_child = list() + token_tree_next_sibling = list() + for i in range(num_nodes): + token_tree_first_child.append(base + i + 1 if i + 1 < num_nodes else -1) + token_tree_next_sibling.append(-1) + return token_tree_first_child, token_tree_next_sibling, base, base + 1 + + def gen_full_binary_tree(height, base): + token_tree_first_child = list() + token_tree_next_sibling = list() + num_nodes = 2**height - 1 + for i in range(num_nodes): + token_tree_first_child.append(base + i * 2 + 1 if i * 2 + 1 < num_nodes else -1) + token_tree_next_sibling.append(base + i * 2 + 2 if i * 2 + 2 < num_nodes else -1) + return token_tree_first_child, token_tree_next_sibling, base, base + 1 + + ### Inputs + num_nodes = 0 + token_tree_first_child = list() + token_tree_next_sibling = list() + token_tree_parent_ptr = list() + + for _ in range(nbatch): + choice = np.random.choice(2, 1, p=plist) + if choice == 0: + nodes_batch = np.random.randint(3, 32) + res = gen_chain(nodes_batch, num_nodes) + num_nodes += nodes_batch + else: + height = np.random.randint(3, 5) + res = gen_full_binary_tree(height, num_nodes) + num_nodes += 2**height - 1 + token_tree_first_child.extend(res[0]) + token_tree_next_sibling.extend(res[1]) + token_tree_parent_ptr.append(res[2]) + + token_tree_first_child = np.array(token_tree_first_child).astype("int32") + token_tree_next_sibling = np.array(token_tree_next_sibling).astype("int32") + token_tree_parent_ptr = np.array(token_tree_parent_ptr).astype("int32") + + draft_probs = np.random.rand(num_nodes, vocab).astype("float32") + draft_probs /= np.sum(draft_probs, axis=1, keepdims=True) + draft_tokens = np.random.randint(0, vocab, num_nodes).astype("int32") + model_probs = np.random.rand(num_nodes, vocab).astype("float32") + model_probs /= np.sum(model_probs, axis=1, keepdims=True) + uniform_samples = np.random.rand(num_nodes).astype("float32") + + ### TVM Inputs + dev = tvm.cuda(0) + draft_probs_tvm = tvm.runtime.tensor(draft_probs, dev) + draft_tokens_tvm = tvm.runtime.tensor(draft_tokens, dev) + model_probs_tvm = tvm.runtime.tensor(model_probs, dev) + token_tree_first_child_tvm = tvm.runtime.tensor(token_tree_first_child, dev) + token_tree_next_sibling_tvm = tvm.runtime.tensor(token_tree_next_sibling, dev) + uniform_samples_tvm = tvm.runtime.tensor(uniform_samples, dev) + token_tree_parent_ptr_tvm = tvm.runtime.tensor(token_tree_parent_ptr, dev) + + # print("draft_probs", draft_probs) + # print("draft_tokens", draft_tokens) + # print("model_probs", model_probs) + # print("token_tree_first_child", token_tree_first_child) + # print("token_tree_next_sibling", token_tree_next_sibling) + # print("uniform_samples", uniform_samples) + # print("token_tree_parent_ptr", token_tree_parent_ptr) + + ### Numpy reference + numpy_reference( + draft_probs, + draft_tokens, + model_probs, + token_tree_first_child, + token_tree_next_sibling, + uniform_samples, + token_tree_parent_ptr, + ) + # print("model_probs", model_probs) + # print("token_tree_parent_ptr", token_tree_parent_ptr) + + ### TVM + kernel = batch_spec_verify(vocab) + mod = tvm.build(kernel, target="cuda") + mod( + draft_probs_tvm, + draft_tokens_tvm, + model_probs_tvm, + token_tree_first_child_tvm, + token_tree_next_sibling_tvm, + uniform_samples_tvm, + token_tree_parent_ptr_tvm, + ) + # print("model_probs", model_probs_tvm.asnumpy()) + # print("token_tree_parent_ptr", token_tree_parent_ptr_tvm.asnumpy()) + + tvm.testing.assert_allclose(model_probs, model_probs_tvm.asnumpy()) + tvm.testing.assert_allclose( + token_tree_parent_ptr, token_tree_parent_ptr_tvm.asnumpy(), rtol=0, atol=0 + ) + + time_evaluator = mod.time_evaluator(mod.entry_name, dev, number=10, repeat=3) + print(f"batch_size: {nbatch}, vocab_size: {vocab}, tree_structure: {plist}") + print( + time_evaluator( + draft_probs_tvm, + draft_tokens_tvm, + model_probs_tvm, + token_tree_first_child_tvm, + token_tree_next_sibling_tvm, + uniform_samples_tvm, + token_tree_parent_ptr_tvm, + ) + ) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/op/test_fp8_block_matmul.py b/tests/python/op/test_fp8_block_matmul.py new file mode 100644 index 0000000..35f4f42 --- /dev/null +++ b/tests/python/op/test_fp8_block_matmul.py @@ -0,0 +1,1037 @@ +from itertools import product +from typing import Tuple # noqa: UP035 + +import ml_dtypes +import numpy as np +import pytest +import torch +import tvm +from tvm import relax +from tvm.relax.frontend import nn +from tvm.relax.frontend.nn import spec +from tvm.s_tir import dlight as dl + +from mlc_llm.compiler_pass.dispatch_triton_kernel import DispatchTritonKernel +from mlc_llm.op import batch_matmul, cutlass, moe_matmul, triton +from mlc_llm.quantization.block_scale_quantization import rowwise_group_quant_fp8 + +# test category "op_correctness" +pytestmark = [pytest.mark.op_correctness] + +block_size = (128, 128) +fp8_dtype = "float8_e4m3fn" + +torch_fp8_dtype = torch.float8_e4m3fn +torch_device = torch.device("cuda") + +torch.set_grad_enabled(False) + + +def test_fp8_block_matmul_cutlass(M: int, N: int, K: int, dtype: str): + class TestModule(nn.Module): + def __init__(self): + pass + + def cutlass_gemm(self, x: nn.Tensor, w: nn.Tensor, w_scale: nn.Tensor): + n, k = w.shape + # assert n % block_size[0] == 0 + assert k % block_size[1] == 0 + assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[0] + assert k // block_size[1] == w_scale.shape[1] + assert x.shape[1] == k + x_fp8, x_scale = rowwise_group_quant_fp8( + x, block_size[1], w.dtype, transpose_scale=True + ) + assert x_fp8.dtype == w.dtype + assert x_scale.dtype == "float32" + o = cutlass.fp8_groupwise_scaled_gemm(x_fp8, x_scale, w, w_scale, block_size, x.dtype) + return x_fp8, x_scale, o + + mod, _, ext_mods = TestModule().export_tvm( + spec={ + "cutlass_gemm": { + "x": spec.Tensor(("m", K), dtype), + "w": spec.Tensor((N, K), fp8_dtype), + "w_scale": spec.Tensor( + ( + (N + block_size[0] - 1) // block_size[0], + (K + block_size[1] - 1) // block_size[1], + ), + "float32", + ), + }, + }, + allow_extern=True, + ) + device = tvm.cuda() + target = tvm.target.Target.from_device(device) + exec = relax.build( + mod, + target=target, + relax_pipeline=relax.backend.cuda.get_default_pipeline(target), + ) + vm = relax.VirtualMachine(exec, device) + + x_torch = torch.rand(M, K, dtype=getattr(torch, dtype), device=torch_device) * 2 - 1 + w_full_torch = torch.rand(N, K, dtype=getattr(torch, dtype), device=torch_device) * 2 - 1 + w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype) + x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype) + o_torch = blockwise_matmul(x_fp8_torch, x_scale_torch, w_torch, w_scale_torch, x_torch.dtype) + x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device) + w_tvm = tvm.runtime.tensor( + w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device + ) + w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device) + x_fp8_tvm, x_scale_tvm, o_tvm = vm["cutlass_gemm"](x_tvm, w_tvm, w_scale_tvm) + x_fp8_tvm = x_fp8_tvm.numpy() + x_scale_tvm = x_scale_tvm.numpy() + o_tvm = o_tvm.numpy() + + np.testing.assert_allclose( + x_fp8_tvm, + x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), + atol=1e-1, + rtol=1e-1, + ) + np.testing.assert_allclose(x_scale_tvm.T, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5) + atol = 0.5 + rtol = 1e-4 + o_tvm_flat = o_tvm.flatten() + o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten() + failed_indices = np.where( + np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat)) + )[0] + if len(failed_indices) > 0: + print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}") + print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}") + print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}") + np.testing.assert_allclose( + o_tvm, + o_torch.view(torch.float16).cpu().numpy().view(dtype), + atol=atol, + rtol=rtol, + ) + + +def test_fp8_block_matmul_triton(M: int, N: int, K: int, dtype: str): + device = tvm.cuda() + target = tvm.target.Target.from_device(device) + + class TestModule(nn.Module): + def __init__(self): + pass + + def triton_gemm(self, x: nn.Tensor, w: nn.Tensor, w_scale: nn.Tensor): + n, k = w.shape + assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[0] + assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[1] + assert x.shape[1] == k + x_fp8, x_scale = rowwise_group_quant_fp8( + x, block_size[1], w.dtype, transpose_scale=False + ) + assert x_fp8.dtype == w.dtype + assert x_scale.dtype == "float32" + o = triton.fp8_groupwise_scaled_gemm( + x_fp8, + x_scale, + w, + w_scale, + block_size, + x.dtype, + ) + return x_fp8, x_scale, o + + mod, _, ext_mods = TestModule().export_tvm( + spec={ + "triton_gemm": { + "x": spec.Tensor(("m", K), dtype), + "w": spec.Tensor((N, K), fp8_dtype), + "w_scale": spec.Tensor( + ( + (N + block_size[0] - 1) // block_size[0], + (K + block_size[1] - 1) // block_size[1], + ), + "float32", + ), + }, + }, + allow_extern=True, + ) + mod = DispatchTritonKernel(target)(mod) + exec = relax.build( + mod, + target=target, + relax_pipeline=relax.backend.cuda.get_default_pipeline(target), + ) + vm = relax.VirtualMachine(exec, device) + + x_torch = torch.randn(M, K, dtype=getattr(torch, dtype), device=torch_device) + w_full_torch = torch.randn(N, K, dtype=getattr(torch, dtype), device=torch_device) + w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype) + x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype) + o_torch = blockwise_matmul(x_fp8_torch, x_scale_torch, w_torch, w_scale_torch, x_torch.dtype) + x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device) + w_tvm = tvm.runtime.tensor( + w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device + ) + w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device) + x_fp8_tvm, x_scale_tvm, o_tvm = vm["triton_gemm"](x_tvm, w_tvm, w_scale_tvm) + x_fp8_tvm = x_fp8_tvm.numpy() + x_scale_tvm = x_scale_tvm.numpy() + o_tvm = o_tvm.numpy() + np.testing.assert_allclose( + x_fp8_tvm, + x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), + atol=1e-1, + rtol=1e-1, + ) + np.testing.assert_allclose(x_scale_tvm, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5) + atol = 0.5 + rtol = 1e-4 + o_tvm_flat = o_tvm.flatten() + o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten() + failed_indices = np.where( + np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat)) + )[0] + if len(failed_indices) > 0: + print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}") + print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}") + print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}") + np.testing.assert_allclose( + o_tvm, + o_torch.view(torch.float16).cpu().numpy().view(dtype), + atol=atol, + rtol=rtol, + ) + + +def test_fp8_block_group_matmul_cutlass(M: int, N: int, K: int, dtype: str): + num_experts = 256 + top_k = 8 + + device = tvm.cuda() + target = tvm.target.Target.from_device(device) + + class TestModule(nn.Module): + def __init__(self): + pass + + def cutlass_group_gemm( + self, + x: nn.Tensor, + w: nn.Tensor, + w_scale: nn.Tensor, + indptr: nn.Tensor, + ): + e, n, k = w.shape + assert e == num_experts + assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[1] + assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[2] + assert x.shape[1] == k + x_fp8, x_scale = rowwise_group_quant_fp8( + x, block_size[1], w.dtype, transpose_scale=False + ) + assert x_fp8.dtype == w.dtype + assert x_scale.dtype == "float32" + o = cutlass.fp8_groupwise_scaled_group_gemm( + x_fp8, + x_scale, + w, + w_scale, + indptr, + block_size, + x.dtype, + ) + return x_fp8, x_scale, o + + mod, _, ext_mods = TestModule().export_tvm( + spec={ + "cutlass_group_gemm": { + "x": spec.Tensor(("m", K), dtype), + "w": spec.Tensor((num_experts, N, K), fp8_dtype), + "w_scale": spec.Tensor( + ( + num_experts, + (N + block_size[0] - 1) // block_size[0], + (K + block_size[1] - 1) // block_size[1], + ), + "float32", + ), + "indptr": spec.Tensor((num_experts,), "int64"), + }, + }, + allow_extern=True, + ) + exec = relax.build( + mod, + target=target, + relax_pipeline=relax.backend.cuda.get_default_pipeline(target), + ) + vm = relax.VirtualMachine(exec, device) + + # Randomly sample `top_k` experts for each token with pytorch + expert_choices = torch.randint( + 0, num_experts, (M * top_k,), device=torch_device, dtype=torch.int32 + ) + + factor = 1 + # Balance so that the number of tokens for each expert is a multiple of `factor` + token_balance = 0 + num_tokens_list = [int((expert_choices == i).sum().to("cpu")) for i in range(num_experts)] + for i in range(num_experts): + if token_balance > 0: + diff = min(token_balance, num_tokens_list[i]) + num_tokens_list[i] -= diff + token_balance -= diff + if num_tokens_list[i] % factor != 0: + token_balance += factor - num_tokens_list[i] % factor + num_tokens_list[i] += factor - num_tokens_list[i] % factor + assert sum(num_tokens_list) == M * top_k + + indptr = torch.zeros(num_experts + 1, device=torch_device, dtype=torch.int64) + for i in range(num_experts): + indptr[i + 1] = indptr[i] + (expert_choices == i).sum() + token_ids_list = [] + for i in range(num_experts): + # Get the indices of the tokens that belong to the i-th expert + token_ids = torch.where(expert_choices == i)[0] + token_ids_list.append(token_ids) + + x_torch = torch.randn(M * top_k, K, dtype=getattr(torch, dtype), device=torch_device) + w_full_torch = torch.randn(num_experts, N, K, dtype=getattr(torch, dtype), device=torch_device) + w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype) + x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype) + o_torch = blockwise_group_matmul( + x_fp8_torch, + x_scale_torch, + w_torch, + w_scale_torch, + indptr, + x_torch.dtype, + ) + x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device) + w_tvm = tvm.runtime.tensor( + w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device + ) + w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device) + indptr_tvm = tvm.runtime.tensor(indptr[1:].cpu().numpy(), device=device) + x_fp8_tvm, x_scale_tvm, o_tvm = vm["cutlass_group_gemm"]( + x_tvm, + w_tvm, + w_scale_tvm, + indptr_tvm, + ) + x_fp8_tvm = x_fp8_tvm.numpy() + x_scale_tvm = x_scale_tvm.numpy() + o_tvm = o_tvm.numpy() + np.testing.assert_allclose( + x_fp8_tvm, + x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), + atol=1e-1, + rtol=1e-1, + ) + np.testing.assert_allclose(x_scale_tvm, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5) + atol = 0.5 + rtol = 1e-4 + o_tvm_flat = o_tvm.flatten() + o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten() + failed_indices = np.where( + np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat)) + )[0] + if len(failed_indices) > 0: + print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}") + print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}") + print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}") + np.testing.assert_allclose( + o_tvm, + o_torch.view(torch.float16).cpu().numpy().view(dtype), + atol=atol, + rtol=rtol, + ) + + +def test_fp8_block_group_matmul_triton(M: int, N: int, K: int, dtype: str): + num_experts = 256 + top_k = 8 + + device = tvm.cuda() + target = tvm.target.Target.from_device(device) + + class TestModule(nn.Module): + def __init__(self): + pass + + def triton_group_gemm( + self, + x: nn.Tensor, + w: nn.Tensor, + w_scale: nn.Tensor, + indptr: nn.Tensor, + ): + e, n, k = w.shape + assert e == num_experts + assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[1] + assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[2] + assert x.shape[1] == k + x_fp8, x_scale = rowwise_group_quant_fp8( + x, block_size[1], w.dtype, transpose_scale=False + ) + assert x_fp8.dtype == w.dtype + assert x_scale.dtype == "float32" + o = triton.fp8_groupwise_scaled_group_gemm( + x_fp8, + x_scale, + w, + w_scale, + indptr, + block_size, + x.dtype, + ) + return x_fp8, x_scale, o + + mod, _, ext_mods = TestModule().export_tvm( + spec={ + "triton_group_gemm": { + "x": spec.Tensor(("m", K), dtype), + "w": spec.Tensor((num_experts, N, K), fp8_dtype), + "w_scale": spec.Tensor( + ( + num_experts, + (N + block_size[0] - 1) // block_size[0], + (K + block_size[1] - 1) // block_size[1], + ), + "float32", + ), + "indptr": spec.Tensor((num_experts + 1,), "int32"), + }, + }, + allow_extern=True, + ) + mod = DispatchTritonKernel(target)(mod) + exec = relax.build( + mod, + target=target, + relax_pipeline=relax.backend.cuda.get_default_pipeline(target), + ) + vm = relax.VirtualMachine(exec, device) + + # Randomly sample `top_k` experts for each token with pytorch + expert_choices = torch.randint( + 0, num_experts, (M * top_k,), device=torch_device, dtype=torch.int32 + ) + + indptr = torch.zeros(num_experts + 1, device=torch_device, dtype=torch.int32) + for i in range(num_experts): + indptr[i + 1] = indptr[i] + (expert_choices == i).sum() + token_ids_list = [] + for i in range(num_experts): + # Get the indices of the tokens that belong to the i-th expert + token_ids = torch.where(expert_choices == i)[0] + token_ids_list.append(token_ids) + + x_torch = torch.randn(M * top_k, K, dtype=getattr(torch, dtype), device=torch_device) + w_full_torch = torch.randn(num_experts, N, K, dtype=getattr(torch, dtype), device=torch_device) + w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype) + x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype) + o_torch = blockwise_group_matmul( + x_fp8_torch, + x_scale_torch, + w_torch, + w_scale_torch, + indptr, + x_torch.dtype, + ) + x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device) + w_tvm = tvm.runtime.tensor( + w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device + ) + w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device) + indptr_tvm = tvm.runtime.tensor(indptr.cpu().numpy(), device=device) + x_fp8_tvm, x_scale_tvm, o_tvm = vm["triton_group_gemm"]( + x_tvm, + w_tvm, + w_scale_tvm, + indptr_tvm, + ) + x_fp8_tvm = x_fp8_tvm.numpy() + x_scale_tvm = x_scale_tvm.numpy() + o_tvm = o_tvm.numpy() + np.testing.assert_allclose( + x_fp8_tvm, + x_fp8_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), + atol=1e-1, + rtol=1e-1, + ) + np.testing.assert_allclose(x_scale_tvm, x_scale_torch.cpu().numpy(), atol=1e-5, rtol=1e-5) + atol = 0.5 + rtol = 1e-4 + o_tvm_flat = o_tvm.flatten() + o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten() + failed_indices = np.where( + np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat)) + )[0] + if len(failed_indices) > 0: + print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}") + print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}") + print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}") + np.testing.assert_allclose( + o_tvm, + o_torch.view(torch.float16).cpu().numpy().view(dtype), + atol=atol, + rtol=rtol, + ) + + +def test_fp8_block_bmm_cutlass(M: int, N: int, K: int, H: int, dtype: str): + class TestModule(nn.Module): + def __init__(self): + pass + + def cutlass_bmm(self, x: nn.Tensor, w: nn.Tensor, w_scale: nn.Tensor): + _, n, k = w.shape + assert w.shape[0] == x.shape[0] == H + assert n % block_size[0] == 0 + assert k % block_size[1] == 0 + assert n // block_size[0] == w_scale.shape[1] + assert k // block_size[1] == w_scale.shape[2] + assert x.shape[2] == k + o = batch_matmul.quantized_bmm(x, w, w_scale, block_size) + return o + + mod, _, ext_mods = TestModule().export_tvm( + spec={ + "cutlass_bmm": { + "x": spec.Tensor((H, "m", K), dtype), + "w": spec.Tensor((H, N, K), fp8_dtype), + "w_scale": spec.Tensor( + ( + H, + (N + block_size[0] - 1) // block_size[0], + (K + block_size[1] - 1) // block_size[1], + ), + "float32", + ), + }, + }, + allow_extern=True, + ) + device = tvm.cuda() + target = tvm.target.Target.from_device(device) + exec = relax.build( + mod, + target=target, + relax_pipeline=relax.backend.cuda.get_default_pipeline(target), + ) + vm = relax.VirtualMachine(exec, device) + + x_torch = torch.randn(H, M, K, dtype=getattr(torch, dtype), device=torch_device) + w_full_torch = torch.randn(H, N, K, dtype=getattr(torch, dtype), device=torch_device) + w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype) + x_torch, x_fp8_torch, x_scale_torch = rowwise_quant_fp8(x_torch, block_size, torch_fp8_dtype) + o_torch = blockwise_bmm(x_fp8_torch, x_scale_torch, w_torch, w_scale_torch, x_torch.dtype) + x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device) + w_tvm = tvm.runtime.tensor( + w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device + ) + w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device) + o_tvm = vm["cutlass_bmm"](x_tvm, w_tvm, w_scale_tvm) + o_tvm = o_tvm.numpy() + atol = 0.5 + rtol = 1e-4 + o_tvm_flat = o_tvm.flatten() + o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten() + failed_indices = np.where( + np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat)) + )[0] + if len(failed_indices) > 0: + print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}") + print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}") + print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}") + np.testing.assert_allclose( + o_tvm, + o_torch.view(torch.float16).cpu().numpy().view(dtype), + atol=atol, + rtol=rtol, + ) + + +def test_fp8_block_gemv_tir(N: int, K: int, up: bool, dtype: str): + num_experts = 256 + top_k = 8 + M = 1 if up else top_k + + device = tvm.cuda() + target = tvm.target.Target.from_device(device) + + class TestModule(nn.Module): + def __init__(self): + pass + + def tir_moe_gemv( + self, + x: nn.Tensor, + w: nn.Tensor, + w_scale: nn.Tensor, + expert_indices: nn.Tensor, + ): + e, n, k = w.shape + assert e == num_experts + assert (n + block_size[0] - 1) // block_size[0] == w_scale.shape[1] + assert (k + block_size[1] - 1) // block_size[1] == w_scale.shape[2] + assert x.shape[1] == k + o = moe_matmul.dequantize_block_scale_float8_gemv( + x, w, w_scale, expert_indices, block_size, x.dtype + ) + return o + + mod, _, ext_mods = TestModule().export_tvm( + spec={ + "tir_moe_gemv": { + "x": spec.Tensor((M, K), dtype), + "w": spec.Tensor((num_experts, N, K), fp8_dtype), + "w_scale": spec.Tensor( + ( + num_experts, + (N + block_size[0] - 1) // block_size[0], + (K + block_size[1] - 1) // block_size[1], + ), + "float32", + ), + "expert_indices": spec.Tensor((1, top_k), "int32"), + }, + }, + allow_extern=True, + ) + with target: + mod = dl.ApplyDefaultSchedule( + dl.gpu.Matmul(), + dl.gpu.GEMV(), + dl.gpu.Reduction(), + dl.gpu.GeneralReduction(), + dl.gpu.Fallback(), + )(mod) + exec = relax.build( + mod, + target=target, + relax_pipeline=relax.backend.cuda.get_default_pipeline(target), + ) + vm = relax.VirtualMachine(exec, device) + + # Randomly sample `top_k` experts for each token with pytorch + expert_choices = torch.randint(0, num_experts, (top_k,), device=torch_device, dtype=torch.int32) + indptr = torch.zeros(num_experts + 1, device=torch_device, dtype=torch.int32) + for i in range(num_experts): + indptr[i + 1] = indptr[i] + (expert_choices == i).sum() + token_ids_list = [] + for i in range(num_experts): + # Get the indices of the tokens that belong to the i-th expert + token_ids = torch.where(expert_choices == i)[0] + token_ids_list.append(token_ids) + + x_torch = torch.randn(M, K, dtype=getattr(torch, dtype), device=torch_device) + w_full_torch = torch.randn(num_experts, N, K, dtype=getattr(torch, dtype), device=torch_device) + w_torch, w_scale_torch = blockwise_quant_fp8(w_full_torch, block_size, torch_fp8_dtype) + x_input_torch = torch.repeat_interleave(x_torch, top_k, dim=0) if up else x_torch + o_torch = blockwise_group_matmul_unquantized( + x_input_torch, w_torch, w_scale_torch, expert_choices + ) + x_tvm = tvm.runtime.tensor(x_torch.view(torch.float16).cpu().numpy().view(dtype), device=device) + w_tvm = tvm.runtime.tensor( + w_torch.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device + ) + w_scale_tvm = tvm.runtime.tensor(w_scale_torch.cpu().numpy(), device=device) + expert_choices = tvm.runtime.tensor( + expert_choices.reshape(1, top_k).cpu().numpy(), device=device + ) + o_tvm = vm["tir_moe_gemv"](x_tvm, w_tvm, w_scale_tvm, expert_choices) + o_tvm = o_tvm.numpy() + atol = 0.5 + rtol = 1e-4 + o_tvm_flat = o_tvm.flatten() + o_torch_flat = o_torch.view(torch.float16).cpu().numpy().view(dtype).flatten() + failed_indices = np.where( + np.abs(o_tvm_flat - o_torch_flat) > (atol + rtol * np.abs(o_torch_flat)) + )[0] + if len(failed_indices) > 0: + print(f"failed_indices: {failed_indices}, size: {len(failed_indices)}") + print(f"o_tvm_flat[failed_indices]: {o_tvm_flat[failed_indices]}") + print(f"o_torch_flat[failed_indices]: {o_torch_flat[failed_indices]}") + np.testing.assert_allclose( + o_tvm, + o_torch.view(torch.float16).cpu().numpy().view(dtype), + atol=atol, + rtol=rtol, + ) + + +def blockwise_matmul( + x_fp8_torch: torch.Tensor, + x_scale_torch: torch.Tensor, + w_torch: torch.Tensor, + w_scale_torch: torch.Tensor, + dtype, +): + o_torch = torch.zeros( + (x_fp8_torch.shape[0], w_torch.shape[0]), dtype=dtype, device=torch_device + ) + for j in range(w_scale_torch.shape[0]): + for k in range(w_scale_torch.shape[1]): + o_torch[ + :, + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[0]), + ] += ( + torch.matmul( + x_fp8_torch[ + :, + k * block_size[1] : min((k + 1) * block_size[1], x_fp8_torch.shape[1]), + ].to(dtype), + w_torch[ + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[0]), + k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[1]), + ].T.to(dtype), + ) + * x_scale_torch[:, k : k + 1] + * w_scale_torch[j, k] + ) + return o_torch + + +def blockwise_group_matmul( + x_fp8_torch: torch.Tensor, + x_scale_torch: torch.Tensor, + w_torch: torch.Tensor, + w_scale_torch: torch.Tensor, + indptr: torch.Tensor, + dtype, +): + o_torch = torch.zeros( + (x_fp8_torch.shape[0], w_torch.shape[1]), dtype=dtype, device=torch_device + ) + for e in range(w_scale_torch.shape[0]): + if indptr[e + 1] - indptr[e] == 0: + continue + indices = slice(indptr[e], indptr[e + 1]) + for j in range(w_scale_torch.shape[1]): + for k in range(w_scale_torch.shape[2]): + o_torch[ + indices, + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]), + ] += ( + torch.matmul( + x_fp8_torch.to(dtype)[ + indices, + k * block_size[1] : min((k + 1) * block_size[1], x_fp8_torch.shape[1]), + ], + w_torch[ + e, + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]), + k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[2]), + ].T.to(dtype), + ) + * x_scale_torch[indices, k : k + 1] + * w_scale_torch[e, j, k] + ) + return o_torch + + +def blockwise_group_matmul_unquantized( + x_torch: torch.Tensor, + w_torch: torch.Tensor, + w_scale_torch: torch.Tensor, + expert_choices: torch.Tensor, +): + o_torch = torch.zeros( + (x_torch.shape[0], w_torch.shape[1]), dtype=x_torch.dtype, device=torch_device + ) + for i, e in enumerate(expert_choices): + for j in range(w_scale_torch.shape[1]): + for k in range(w_scale_torch.shape[2]): + o_torch[ + i, + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]), + ] += torch.matmul( + x_torch[ + i, + k * block_size[1] : min((k + 1) * block_size[1], x_torch.shape[1]), + ], + w_torch[ + e, + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]), + k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[2]), + ].T.to(x_torch.dtype) + * w_scale_torch[e, j, k].to(x_torch.dtype), + ) + return o_torch + + +def blockwise_bmm( + x_fp8_torch: torch.Tensor, + x_scale_torch: torch.Tensor, + w_torch: torch.Tensor, + w_scale_torch: torch.Tensor, + dtype, +): + o_torch = torch.zeros( + (x_fp8_torch.shape[0], x_fp8_torch.shape[1], w_torch.shape[1]), + dtype=dtype, + device=torch_device, + ) + for j in range(w_scale_torch.shape[1]): + for k in range(w_scale_torch.shape[2]): + o_torch[ + ..., + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]), + ] += ( + torch.bmm( + x_fp8_torch[ + ..., + k * block_size[1] : min((k + 1) * block_size[1], x_fp8_torch.shape[2]), + ].to(dtype), + w_torch[ + ..., + j * block_size[0] : min((j + 1) * block_size[0], w_torch.shape[1]), + k * block_size[1] : min((k + 1) * block_size[1], w_torch.shape[2]), + ] + .transpose(1, 2) + .to(dtype), + ) + * x_scale_torch[..., k : k + 1] + * w_scale_torch[..., j : j + 1, k : k + 1] + ) + return o_torch + + +def blockwise_quant_fp8( + w_full_torch: torch.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + quant_dtype: torch.dtype, +): + w_scale_shape = ( + *w_full_torch.shape[:-2], + (w_full_torch.shape[-2] + block_size[0] - 1) // block_size[0], + (w_full_torch.shape[-1] + block_size[1] - 1) // block_size[1], + ) + # For each (block_size[0], block_size[1]) block, compute the max abs value of `w_full_torch` + w_max_abs_torch = torch.zeros(w_scale_shape, dtype=torch.float32, device=torch_device) + for i in range(w_scale_shape[-2]): + for j in range(w_scale_shape[-1]): + w_max_abs_torch[..., i, j] = torch.max( + torch.abs( + w_full_torch[ + ..., + i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]), + j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]), + ] + ).flatten(-2, -1), + dim=-1, + )[0] + # Scale is the `w_max_abs_torch` divided by the max value of quant_dtype in ml_dtypes + fp8_max = float(ml_dtypes.finfo(fp8_dtype).max) + w_scale_torch = w_max_abs_torch / fp8_max + # `w_torch` is the `w_full_torch` divided by the `w_scale_torch` (with block awareness), + # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype` + w_torch = torch.zeros_like(w_full_torch, dtype=quant_dtype, device=torch_device) + if len(w_scale_shape) == 2: + for i in range(w_scale_shape[-2]): + for j in range(w_scale_shape[-1]): + w_torch[ + i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]), + j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]), + ] = torch.clamp( + w_full_torch[ + i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]), + j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]), + ] + / w_scale_torch[..., i, j], + -fp8_max, + fp8_max, + ) + else: + for e in range(w_scale_shape[0]): + for i in range(w_scale_shape[-2]): + for j in range(w_scale_shape[-1]): + w_torch[ + e, + i * block_size[0] : min((i + 1) * block_size[0], w_full_torch.shape[-2]), + j * block_size[1] : min((j + 1) * block_size[1], w_full_torch.shape[-1]), + ] = torch.clamp( + w_full_torch[ + e, + i * block_size[0] : min( + (i + 1) * block_size[0], w_full_torch.shape[-2] + ), + j * block_size[1] : min( + (j + 1) * block_size[1], w_full_torch.shape[-1] + ), + ] + / w_scale_torch[e, i, j], + -fp8_max, + fp8_max, + ) + + w_scale_torch = ( + torch.rand(w_scale_torch.shape, dtype=torch.float32, device=torch_device) / fp8_max + ) + return w_torch, w_scale_torch + + +def rowwise_quant_fp8( + x_full_torch: torch.Tensor, + block_size: Tuple[int, int], # noqa: UP006 + quant_dtype: torch.dtype, +): + x_scale_shape = ( + *x_full_torch.shape[:-1], + (x_full_torch.shape[-1] + block_size[1] - 1) // block_size[1], + ) + # For each (block_size[1]) block, compute the max abs value of `w_full_torch` + x_max_abs_torch = torch.zeros(x_scale_shape, dtype=torch.float32, device=torch_device) + for i in range(x_scale_shape[-1]): + x_max_abs_torch[..., i] = torch.max( + torch.abs( + x_full_torch[ + ..., + i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]), + ] + ), + dim=-1, + )[0] + # Scale is the `x_max_abs_torch` divided by the max value of quant_dtype in ml_dtypes + fp8_max = float(ml_dtypes.finfo(fp8_dtype).max) + x_scale_torch = x_max_abs_torch / fp8_max + # `x_torch` is the `x_full_torch` divided by the `x_scale_torch` (with block awareness), + # clamped to (-fp8_max, fp8_max), and cast to `quant_dtype` + x_torch = torch.zeros_like(x_full_torch, dtype=quant_dtype, device=torch_device) + for i in range(x_scale_shape[-1]): + x_torch[ + ..., + i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]), + ] = torch.clamp( + x_full_torch[ + ..., + i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]), + ] + / x_scale_torch[..., i : i + 1], + -fp8_max, + fp8_max, + ) + + x_scale_torch = ( + torch.rand(x_scale_torch.shape, dtype=torch.float32, device=torch_device) / fp8_max + ) + for i in range(x_scale_shape[-1]): + x_full_torch[ + ..., + i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]), + ] = ( + x_torch[ + ..., + i * block_size[1] : min((i + 1) * block_size[1], x_full_torch.shape[-1]), + ].to(x_scale_torch.dtype) + * x_scale_torch[..., i : i + 1] + ) + return x_full_torch, x_torch, x_scale_torch + + +@pytest.mark.skip(reason="Test requiring SM90a") +def test_cutlass_gemm(): + # Cutlass GEMM + for M, (N, K), dtype in product( + [4, 128, 256, 1024, 2112], + [ + (4608, 896), + (896, 2304), + (3072, 896), + (512, 896), + (3072, 512), + (4096, 512), + (896, 2048), + (129280, 896), + ], + ["bfloat16"], + ): + print(f"Cutlass, M: {M}, N: {N}, K: {K}, dtype: {dtype}") + test_fp8_block_matmul_cutlass(M, N, K, dtype) + + +@pytest.mark.skip(reason="Test requiring SM90a") +def test_triton_gemm(): + # Triton GEMM + for M, (N, K), dtype in product( + [1, 128, 256, 1024, 2111], + [ + (4608, 896), + (896, 576), + (896, 2304), + ], + ["bfloat16"], + ): + print(f"Triton, M: {M}, N: {N}, K: {K}, dtype: {dtype}") + test_fp8_block_matmul_triton(M, N, K, dtype) + + +@pytest.mark.skip(reason="Test requiring SM90a") +def test_cutlass_group_gemm(): + # Cutlass group GEMM + for M, (N, K), dtype in product( + [1, 128, 256, 1024, 2111], + [ + (512, 896), + (896, 256), + ], + ["bfloat16"], + ): + print(f"Cutlass group gemm, M: {M}, N: {N}, K: {K}, dtype: {dtype}") + test_fp8_block_group_matmul_cutlass(M, N, K, dtype) + + +@pytest.mark.skip(reason="Test requiring SM90a") +def test_triton_group_gemm(): + # Triton group GEMM + for M, (N, K), dtype in product( + [1, 128, 256, 1024, 2111], + [ + (512, 896), + (896, 256), + ], + ["bfloat16"], + ): + print(f"Triton group gemm, M: {M}, N: {N}, K: {K}, dtype: {dtype}") + test_fp8_block_group_matmul_triton(M, N, K, dtype) + + +@pytest.mark.skip(reason="Test requiring SM90a") +def test_cutlass_bmm(): + # Cutlass BMM + for M, H, (N, K), dtype in product( + [4, 128, 256, 1024, 2112], + [16, 64, 128], + [ + (512, 128), + (128, 512), + ], + ["bfloat16"], + ): + print(f"Cutlass BMM, M: {M}, N: {N}, K: {K}, H: {H}, dtype: {dtype}") + test_fp8_block_bmm_cutlass(M, N, K, H, dtype) + + +@pytest.mark.skip(reason="Test requiring SM90a") +def test_tir_moe_gemv(): + # TIR MoE GEMV + for (N, K), up, dtype in product( + [(512, 896), (896, 256)], + [True, False], + ["bfloat16"], + ): + print(f"TIR MoE GEMV, N: {N}, K: {K}, up: {up}, dtype: {dtype}") + test_fp8_block_gemv_tir(N, K, up, dtype) + + +if __name__ == "__main__": + test_cutlass_gemm() + test_triton_gemm() + test_cutlass_group_gemm() + test_triton_group_gemm() + test_cutlass_bmm() + test_tir_moe_gemv() diff --git a/tests/python/op/test_mrope.py b/tests/python/op/test_mrope.py new file mode 100644 index 0000000..ce062d4 --- /dev/null +++ b/tests/python/op/test_mrope.py @@ -0,0 +1,251 @@ +import numpy as np +import pytest + +tvm = pytest.importorskip("tvm") +from tvm import relax # noqa: E402 +from tvm.relax.frontend import nn # noqa: E402 +from tvm.relax.frontend.nn import spec # noqa: E402 +from tvm.runtime import tensor as tvm_tensor # noqa: E402 + +from mlc_llm.op import ( # noqa: E402 + MultimodalRotaryEmbedding, + VisionPositionMetadata, + apply_multimodal_rotary_pos_emb, + get_mrope_position_ids, +) + + +def _numpy_rotate_half(x: np.ndarray) -> np.ndarray: + x1, x2 = np.split(x, 2, axis=-1) + return np.concatenate([-x2, x1], axis=-1) + + +def _numpy_apply_mrope( + q: np.ndarray, + k: np.ndarray, + position_ids: np.ndarray, + theta: float, + mrope_section: tuple[int, ...], +) -> tuple[np.ndarray, np.ndarray]: + if position_ids.ndim != 3: + raise ValueError(f"position_ids must be rank-3, got shape {position_ids.shape}") + if position_ids.shape[0] == 3: + position_ids = np.transpose(position_ids, (1, 2, 0)) + elif position_ids.shape[-1] != 3: + raise ValueError( + "position_ids must have shape (batch, seq, 3) or (3, batch, seq), " + f"got {position_ids.shape}" + ) + + head_dim = q.shape[-1] + inv_freq = 1.0 / (theta ** (np.arange(0, head_dim, 2, dtype=np.float32) / float(head_dim))) + pos = np.transpose(position_ids, (2, 0, 1)) + inv = inv_freq.reshape(1, 1, -1, 1).astype(np.float32) + inv = np.broadcast_to(inv, (3, pos.shape[1], inv_freq.size, 1)) + pos = pos.reshape(3, pos.shape[1], 1, pos.shape[2]).astype(np.float32) + freqs = np.matmul(inv, pos) + freqs = np.transpose(freqs, (0, 1, 3, 2)) + emb = np.concatenate([freqs, freqs], axis=-1) + cos = np.cos(emb) + sin = np.sin(emb) + split_sizes = list(mrope_section) * 2 + split_points = np.cumsum(split_sizes)[:-1] + cos_chunks = np.split(cos, split_points, axis=-1) + sin_chunks = np.split(sin, split_points, axis=-1) + cos = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(cos_chunks)], axis=-1) + sin = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(sin_chunks)], axis=-1) + cos = np.expand_dims(cos, axis=2) + sin = np.expand_dims(sin, axis=2) + q_out = q * cos + _numpy_rotate_half(q) * sin + k_out = k * cos + _numpy_rotate_half(k) * sin + return q_out, k_out + + +def _evaluate_tensor(expr): + mod = tvm.IRModule.from_expr(expr) + target = tvm.target.Target("llvm") + ex = tvm.relax.build(mod, target) + vm = tvm.relax.VirtualMachine(ex, tvm.cpu()) + return vm["main"]().numpy() + + +def _run_mlc_mrope( + q_np: np.ndarray, + k_np: np.ndarray, + position_ids_np: np.ndarray, + theta: float, + mrope_section: tuple[int, ...], +) -> tuple[np.ndarray, np.ndarray]: + class RopeModule(nn.Module): + def __init__(self): + super().__init__() + self.rotary = MultimodalRotaryEmbedding(q_np.shape[-1], theta, mrope_section) + + def forward( + self, + q: nn.Tensor, + k: nn.Tensor, + pos: nn.Tensor, + ): + """Run MRoPE on test tensors and return rotated query/key outputs.""" + cos, sin = self.rotary(q, pos) + return apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section) + + module = RopeModule() + mod, _, _ = module.export_tvm( + spec={ + "forward": { + "q": spec.Tensor(q_np.shape, "float32"), + "k": spec.Tensor(k_np.shape, "float32"), + "pos": spec.Tensor(position_ids_np.shape, "int64"), + } + }, + allow_extern=True, + ) + target = tvm.target.Target("llvm") + exec_mod = relax.build(mod, target=target) + vm = relax.VirtualMachine(exec_mod, tvm.cpu()) + device = tvm.cpu() + q_nd = tvm_tensor(q_np.astype("float32"), device=device) + k_nd = tvm_tensor(k_np.astype("float32"), device=device) + pos_nd = tvm_tensor(position_ids_np.astype("int64"), device=device) + out_q, out_k = vm["forward"](q_nd, k_nd, pos_nd) + return out_q.numpy(), out_k.numpy() + + +def test_apply_mrope_matches_numpy_reference(): + theta = 10000.0 + mrope_section = (2, 2, 2) + batch, seq_len, heads, head_dim = 1, 4, 2, 12 + rng = np.random.default_rng(0) + q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32) + k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32) + position_ids = np.zeros((batch, seq_len, 3), dtype=np.int64) + position_ids[0, :, 0] = np.arange(seq_len) + position_ids[0, :, 1] = np.arange(seq_len) * 2 + position_ids[0, :, 2] = np.arange(seq_len) * 3 + + mlc_q, mlc_k = _run_mlc_mrope(q_np, k_np, position_ids, theta, mrope_section) + ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids, theta, mrope_section) + + np.testing.assert_allclose(mlc_q, ref_q, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(mlc_k, ref_k, rtol=1e-5, atol=1e-5) + + +def test_get_mrope_position_ids_text_only(): + input_ids = np.array([[1, 2, 3, 0, 0]], dtype=np.int64) + attention_mask = np.array([[1, 1, 1, 0, 0]], dtype=np.int64) + meta = VisionPositionMetadata( + vision_start_token_id=1000, + image_token_id=1001, + video_token_id=1002, + spatial_merge_size=2, + tokens_per_second=4.0, + ) + position_ids, deltas = get_mrope_position_ids( + input_ids, + meta, + attention_mask=attention_mask, + image_grid_thw=None, + video_grid_thw=None, + second_per_grid_ts=None, + ) + expected = attention_mask.cumsum(axis=-1) - 1 + expected = np.where(attention_mask == 0, 1, expected) + expected = np.expand_dims(expected, axis=0).repeat(3, axis=0) + np.testing.assert_array_equal(position_ids, expected) + np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64)) + + +def test_get_mrope_position_ids_single_image_block(): + meta = VisionPositionMetadata( + vision_start_token_id=5000, + image_token_id=5001, + video_token_id=6000, + spatial_merge_size=2, + tokens_per_second=4.0, + ) + input_ids = np.array( + [[11, 12, 5000, 5001, 21, 22, 23, 24, 31, 32]], + dtype=np.int64, + ) + attention_mask = np.ones_like(input_ids, dtype=np.int64) + image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64) + position_ids, deltas = get_mrope_position_ids( + input_ids, + meta, + attention_mask=attention_mask, + image_grid_thw=image_grid_thw, + video_grid_thw=None, + second_per_grid_ts=None, + ) + expected = np.array( + [ + [0, 1, 2, 3, 3, 3, 3, 5, 6, 7], + [0, 1, 2, 3, 3, 4, 4, 5, 6, 7], + [0, 1, 2, 3, 4, 3, 4, 5, 6, 7], + ], + dtype=np.int64, + ).reshape(3, 1, -1) + np.testing.assert_array_equal(position_ids, expected) + np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64)) + + +def test_apply_mrope_accepts_3_batch_seq_layout(): + theta = 10000.0 + mrope_section = (2, 2, 2) + batch, seq_len, heads, head_dim = 1, 4, 2, 12 + rng = np.random.default_rng(1) + q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32) + k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32) + + position_ids_bsc = np.zeros((batch, seq_len, 3), dtype=np.int64) + position_ids_bsc[0, :, 0] = np.arange(seq_len) + position_ids_bsc[0, :, 1] = np.arange(seq_len) * 2 + position_ids_bsc[0, :, 2] = np.arange(seq_len) * 3 + position_ids_3bs = np.transpose(position_ids_bsc, (2, 0, 1)) + + mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section) + mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section) + ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section) + + np.testing.assert_allclose(mlc_q_bsc, ref_q, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(mlc_k_bsc, ref_k, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(mlc_q_3bs, ref_q, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(mlc_k_3bs, ref_k, rtol=1e-5, atol=1e-5) + + +def test_get_mrope_position_ids_output_is_directly_usable(): + theta = 10000.0 + mrope_section = (2, 2, 2) + meta = VisionPositionMetadata( + vision_start_token_id=7000, + image_token_id=7001, + video_token_id=7002, + spatial_merge_size=2, + tokens_per_second=4.0, + ) + input_ids = np.array([[11, 12, 7000, 7001, 21, 22, 23, 24, 31, 32]], dtype=np.int64) + attention_mask = np.ones_like(input_ids, dtype=np.int64) + image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64) + position_ids_3bs, _ = get_mrope_position_ids( + input_ids, + meta, + attention_mask=attention_mask, + image_grid_thw=image_grid_thw, + video_grid_thw=None, + second_per_grid_ts=None, + ) + position_ids_bsc = np.transpose(position_ids_3bs, (1, 2, 0)) + + batch, seq_len = input_ids.shape + heads, head_dim = 2, 12 + rng = np.random.default_rng(2) + q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32) + k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32) + + mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section) + mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section) + + np.testing.assert_allclose(mlc_q_3bs, mlc_q_bsc, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(mlc_k_3bs, mlc_k_bsc, rtol=1e-5, atol=1e-5) diff --git a/tests/python/op/test_top_p_pivot.py b/tests/python/op/test_top_p_pivot.py new file mode 100644 index 0000000..33daeef --- /dev/null +++ b/tests/python/op/test_top_p_pivot.py @@ -0,0 +1,86 @@ +import numpy as np +import pytest +import tvm +import tvm.testing + +from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm + +# mypy: disable-error-code="var-annotated" + +# test category "op_correctness" +pytestmark = [pytest.mark.op_correctness] + + +@pytest.mark.parametrize("batch_size", [32, 64]) +@pytest.mark.parametrize("vocab", [3, 32, 64, 128]) +def test_top_p_renorm(batch_size, vocab): + top_p = 0.95 + init_pivots_np = np.array([1 - top_p, 0.02, 0.01]).astype(np.float32) + top_p_np = np.array([top_p]).astype(np.float32) + + p_np = np.random.exponential(3, size=(batch_size, vocab)).astype(np.float32) + p_np /= np.sum(p_np, axis=-1, keepdims=True) + final_pivot_np = np.zeros(batch_size).astype(np.float32) + final_lsum_np = np.zeros(batch_size).astype(np.float32) + + dev = tvm.cuda(0) + var_prob = tvm.runtime.tensor(p_np, dev) + var_init_pivots = tvm.runtime.tensor(init_pivots_np, dev) + top_p_global = tvm.runtime.tensor(top_p_np, dev) + var_final_pivot = tvm.runtime.tensor(final_pivot_np, dev) + var_final_lsum = tvm.runtime.tensor(final_lsum_np, dev) + + kernel = top_p_pivot(init_pivots_np.shape[0]) + mod = tvm.build(kernel, target="cuda") + mod(var_prob, top_p_global, var_init_pivots, var_final_pivot, var_final_lsum) + + final_pivot = var_final_pivot.asnumpy() + final_lsum = var_final_lsum.asnumpy() + + renorm_np = p_np.copy() + var_renorm = tvm.runtime.tensor(renorm_np, dev) + + kernel_renorm = top_p_renorm() + mod_renorm = tvm.build(kernel_renorm, target="cuda") + mod_renorm(var_prob, var_final_pivot, var_final_lsum, var_renorm) + + renorm = var_renorm.asnumpy() + + def verify_pivot(probs: np.ndarray, pivot: float, lsum: float, renorm: np.ndarray): + sorted_probs = np.sort(probs, axis=-1)[::-1] + num_larger_than_pivot = np.sum(sorted_probs >= pivot) + filtered_sorted_probs = sorted_probs[:num_larger_than_pivot] + min_larger_than_pivot = min(filtered_sorted_probs) + + sum_larger_than_pivot = np.sum(np.where(sorted_probs >= pivot, sorted_probs, 0)) + sum_larger_than_pivot_exclude_min = np.sum( + np.where(filtered_sorted_probs != min_larger_than_pivot, filtered_sorted_probs, 0) + ) + + probs[probs < pivot] = 0 + renorm_prob = probs / np.sum(probs, axis=-1, keepdims=True) + try: + assert sum_larger_than_pivot >= top_p + assert sum_larger_than_pivot_exclude_min < top_p + assert abs(lsum - sum_larger_than_pivot) < 1e-6 + assert np.allclose(renorm, renorm_prob, atol=1e-6, rtol=1e-6) + except AssertionError: + print("Failed") + print("probs:", repr(probs)) + print("pivot:", pivot) + print("sorted_probs:", sorted_probs) + print("num_larger_than_pivot:", num_larger_than_pivot) + print("filtered_sorted_probs:", filtered_sorted_probs) + print("min_larger_than_pivot:", min_larger_than_pivot) + print("sum_larger_than_pivot:", sum_larger_than_pivot) + print("sum_larger_than_pivot_exclude_min:", sum_larger_than_pivot_exclude_min) + print("renom_prob:", renorm_prob) + print("renorm:", renorm) + raise + + for i in range(batch_size): + verify_pivot(p_np[i], final_pivot[i], final_lsum[i], renorm[i]) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/op/test_tree_attn.py b/tests/python/op/test_tree_attn.py new file mode 100644 index 0000000..3fc8593 --- /dev/null +++ b/tests/python/op/test_tree_attn.py @@ -0,0 +1,243 @@ +import math + +import numpy as np +import pytest +import tvm +import tvm.testing +from tvm.relax.frontend.nn.llm import tree_attn + +# test category "op_correctness" +pytestmark = [pytest.mark.op_correctness] + + +@pytest.mark.parametrize("nbatch", [1, 4, 32]) +@pytest.mark.parametrize("h_q", [8, 16]) +@pytest.mark.parametrize("h_kv", [4, 8]) +@pytest.mark.parametrize("d", [128]) +@pytest.mark.parametrize("rotary_mode", [0, 1]) +def test_tree_attn(nbatch, h_q, h_kv, d, rotary_mode): + np.random.seed(0) + np.set_printoptions(linewidth=10000) + + def gen_chain(num_nodes): + mask = np.tril(np.ones((num_nodes, num_nodes))) + return num_nodes, list(mask.flatten()), np.arange(num_nodes) + + def gen_full_binary_tree(height): + mask = list() + pos = list() + num_nodes = 2**height - 1 + for i in range(num_nodes): + if i == 0: + mask_0 = [0] * num_nodes + mask_0[0] = 1 + mask.append(mask_0) + pos.append(0) + else: + mask_i = mask[(i + 1) // 2 - 1].copy() + mask_i[i] = 1 + mask.append(mask_i) + pos.append(pos[(i + 1) // 2 - 1] + 1) + return num_nodes, list(np.array(mask).flatten()), pos + + ### Inputs + num_nodes = 0 + m_list = list() + mn_list = list() + mask_list = list() + q_pos_list = list() + + mn_list.append(0) + + for _ in range(nbatch): + choice = np.random.choice(2, 1, p=[1, 0]) + if choice == 0: + nodes_batch = np.random.randint(3, 32) + res = gen_chain(nodes_batch) + num_nodes += nodes_batch + else: + height = np.random.randint(2, 6) + res = gen_full_binary_tree(height) + num_nodes += 2**height - 1 + m_list.append(res[0]) + mn_list.append(res[0] ** 2) + mask_list.extend(res[1]) + q_pos_list.extend(res[2]) + + qkv_indptr = np.array(np.cumsum([0, *m_list])).astype(np.int32) + m_list = np.array(m_list).astype(np.int32) + mn_list = np.array(mn_list).astype(np.int32) + mn_list = np.cumsum(mn_list).astype(np.int32) + mask_list = np.array(mask_list).astype(np.int32) + q_pos_list = np.array(q_pos_list).astype(np.int32) + + # print("qkv_indptr:", qkv_indptr) + # print("m_list:", m_list) + # print("mn_list:", mn_list) + # for num_nodes, base in zip(m_list, mn_list): + # print("num_nodes:", num_nodes) + # print("indptr:", base) + # print( + # "mask:", + # mask_list[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes), + # ) + # print("q_pos:", q_pos_list[base : base + num_nodes]) + + q = np.random.rand(num_nodes, h_q, d).astype(np.float16) + q_indptr = qkv_indptr + k = np.random.rand(num_nodes, h_kv, d).astype(np.float16) + v = np.random.rand(num_nodes, h_kv, d).astype(np.float16) + kv_indptr = qkv_indptr + q_rope_position = q_pos_list + m_arr = m_list + mn_indptr = mn_list + mask = mask_list + output = np.zeros((num_nodes, h_q, d), dtype=np.float16) + lse = np.zeros((num_nodes, h_q), dtype=np.float32) + rotary_scale = 1.0 + rotary_theta = 10000.0 + attn_score_scaling_factor = 1.0 + + ### TVM Inputs + dev = tvm.cuda(0) + q_tvm = tvm.runtime.tensor(q, dev) + q_indptr_tvm = tvm.runtime.tensor(q_indptr, dev) + k_tvm = tvm.runtime.tensor(k, dev) + v_tvm = tvm.runtime.tensor(v, dev) + kv_indptr_tvm = tvm.runtime.tensor(kv_indptr, dev) + q_rope_position_tvm = tvm.runtime.tensor(q_rope_position, dev) + # m_arr_tvm = tvm.runtime.tensor(m_arr, dev) + mn_indptr_tvm = tvm.runtime.tensor(mn_indptr, dev) + mask_tvm = tvm.runtime.tensor(mask, dev) + output_tvm = tvm.runtime.tensor(output, dev) + lse_tvm = tvm.runtime.tensor(lse, dev) + + target = tvm.target.Target("cuda") + kernel = tree_attn(h_kv=h_kv, h_q=h_q, d=d, dtype="float16", rope_scaling={}, target=target) + mod = tvm.build(kernel, target=target) + mod( + q_tvm, + q_indptr_tvm, + k_tvm, + v_tvm, + kv_indptr_tvm, + q_rope_position_tvm, + # m_arr_tvm, + mn_indptr_tvm, + mask_tvm, + output_tvm, + lse_tvm, + rotary_mode, + rotary_scale, + rotary_theta, + attn_score_scaling_factor, + nbatch, + ) + + ### Numpy reference + def numpy_reference( + q, + q_indptr, + k, + v, + kv_indptr, + q_rope_position, + m_arr, + mn_indptr, + mask, + rotary_mode, + rotary_scale, + rotary_theta, + attn_score_scaling_factor, + output_tvm, + ): + def rope_freq(s, d, d_range, theta, dtype): + freq = s / math.pow(theta, (d * 2 % d_range) / float(d_range)) + cos_freq = np.cos(freq).astype(dtype) + sin_freq = np.sin(freq).astype(dtype) + return cos_freq, sin_freq + + def rope(buffer, offset, rotary_dim, theta, scale, dtype): + result = buffer.copy() + for pos, h, d in np.ndindex(buffer.shape): + cos_freq, sin_freq = rope_freq(offset[pos] * scale, d, rotary_dim, theta, dtype) + cos = cos_freq * buffer[pos, h, d] + sin = sin_freq * ( + -buffer[pos, h, d + rotary_dim // 2] + if d < rotary_dim // 2 + else buffer[pos, h, d - rotary_dim // 2] + ) + result[pos, h, d] = cos + sin + return result + + for i in range(len(m_arr)): + num_nodes = m_arr[i] + base = mn_indptr[i] + q_base = q_indptr[i] + kv_base = kv_indptr[i] + q_pos = q_rope_position[q_base : q_base + num_nodes] # (num_nodes,) + q_i = q[q_base : q_base + num_nodes] # (num_nodes, h_q, d) + k_i = k[kv_base : kv_base + num_nodes] # (num_nodes, h_kv, d) + v_i = v[kv_base : kv_base + num_nodes] # (num_nodes, h_kv, d) + mask_i = mask[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes) + + if rotary_mode == 1: + q_i = rope(q_i, q_pos, d, rotary_theta, rotary_scale, q_i.dtype) + k_i = rope(k_i, q_pos, d, rotary_theta, rotary_scale, k_i.dtype) + + # group attention + # q: (num_nodes, h_q, d) + # k: (num_nodes, h_kv, d) + # v: (num_nodes, h_kv, d) + group_size = h_q // h_kv + q_reshape = q_i.transpose(1, 0, 2) # (h_q, num_nodes, d) + k_reshape = k_i.transpose(1, 2, 0) # (h_kv, d, num_nodes) + v_reshape = v_i.transpose(1, 0, 2) # (h_kv, num_nodes, d) + # expand k_reshape + k_reshape = k_reshape.reshape(h_kv, 1, d, num_nodes) + k_reshape = np.repeat(k_reshape, group_size, axis=1) + k_reshape = k_reshape.reshape(h_q, d, num_nodes) + # expand v_reshape + v_reshape = v_reshape.reshape(h_kv, 1, num_nodes, d) + v_reshape = np.repeat(v_reshape, group_size, axis=1) + v_reshape = v_reshape.reshape(h_q, num_nodes, d) + # print("q_reshape:", q_reshape.shape) + # print("k_reshape:", k_reshape.shape) + # print("v_reshape:", v_reshape.shape) + + # qk: (h_q, num_nodes, num_nodes) + qk = np.matmul(q_reshape, k_reshape) * attn_score_scaling_factor / math.sqrt(float(d)) + # softmax(qk, axis=-1), numerical stability + qk[:, mask_i == 0] = -np.inf + qk_max = np.max(qk, axis=-1, keepdims=True) + qk = np.exp(qk - qk_max) + qk = qk / np.sum(qk, axis=-1, keepdims=True) + + # attention + output_i = np.matmul(qk, v_reshape).transpose(1, 0, 2) # (num_nodes, h_q, d) + # print(output_i) + + tvm.testing.assert_allclose( + output_i, output_tvm[q_base : q_base + num_nodes], rtol=1e-3, atol=1e-3 + ) + + numpy_reference( + q, + q_indptr, + k, + v, + kv_indptr, + q_rope_position, + m_arr, + mn_indptr, + mask, + rotary_mode, + rotary_scale, + rotary_theta, + attn_score_scaling_factor, + output_tvm.numpy(), + ) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/op/test_two_stage_softmax.py b/tests/python/op/test_two_stage_softmax.py new file mode 100644 index 0000000..5816a33 --- /dev/null +++ b/tests/python/op/test_two_stage_softmax.py @@ -0,0 +1,51 @@ +import numpy as np +import pytest +import scipy.special +import tvm +from tvm.s_tir import dlight + +# test category "op_correctness" +pytestmark = [pytest.mark.op_correctness] + + +def test_two_stage_softmax(): + from mlc_llm.compiler_pass.rewrite_softmax import _get_lse_and_softmax_func + + chunk_size = 4096 + target = tvm.target.Target("cuda") + f_chunk_lse, f_softmax_with_lse = _get_lse_and_softmax_func(target, chunk_size) + mod = tvm.IRModule({"chunk_lse": f_chunk_lse, "softmax_with_chunked_lse": f_softmax_with_lse}) + with target: + mod = dlight.ApplyDefaultSchedule(dlight.gpu.GeneralReduction())(mod) + + runtime_mod = tvm.build(mod, target=target) + device = tvm.cuda() + + num_runs = 5 + vocab_size = 128256 + for batch_size in [1, 2, 4, 8, 16, 32, 64, 128]: + for _ in range(num_runs): + x_np = np.random.uniform(low=-10, high=10, size=(batch_size, vocab_size)).astype( + "float32" + ) + y_np = scipy.special.softmax(x_np, axis=-1) + + x_nd = tvm.runtime.tensor(x_np, device=device) + r_nd = tvm.runtime.empty( + (batch_size, (vocab_size + chunk_size - 1) // chunk_size), + x_np.dtype, + device=device, + ) + y_nd = tvm.runtime.empty(x_np.shape, x_np.dtype, device=device) + + runtime_mod["chunk_lse"](x_nd, r_nd) + runtime_mod["softmax_with_chunked_lse"](x_nd, r_nd, y_nd) + + y_nd_arr = y_nd.numpy() + np.testing.assert_allclose(y_nd_arr, y_np, atol=1e-6, rtol=1e-6) + + print(f"pass batch size {batch_size}") + + +if __name__ == "__main__": + test_two_stage_softmax() diff --git a/tests/python/quantization/test_awq_quantization.py b/tests/python/quantization/test_awq_quantization.py new file mode 100644 index 0000000..8d65bdd --- /dev/null +++ b/tests/python/quantization/test_awq_quantization.py @@ -0,0 +1,86 @@ +from typing import List # noqa: UP035 + +import numpy as np +import pytest +import torch +import tvm +import tvm.testing +from tvm import DataType +from tvm.relax.frontend import nn + +from mlc_llm.loader import QuantizeMapping +from mlc_llm.quantization import QUANTIZATION, AWQQuantize + + +def dequantize_np( + config: AWQQuantize, + weight: np.ndarray, + zeros: np.ndarray, + scale: np.ndarray, +) -> np.ndarray: + def decode_int_arr(int_arr: np.ndarray, num_elem_per_storage: int, bits: int): + bin_mask = (1 << bits) - 1 + int_arr_repeated = np.repeat(int_arr, num_elem_per_storage, axis=-1) + indice_j = np.indices(int_arr_repeated.shape)[1] + arr_bin = np.bitwise_and( + np.right_shift( + int_arr_repeated, + (indice_j % num_elem_per_storage) * bits, + ), + bin_mask, + ) + return arr_bin + + weight_bin = decode_int_arr( + weight, config.num_elem_per_storage, DataType(config.quantize_dtype).bits + ) + zero_bin = decode_int_arr( + zeros, config.num_elem_per_storage, DataType(config.quantize_dtype).bits + ) + scale_repeated = np.repeat(scale, config.group_size, axis=-1) + zero_bin_repeated = np.repeat(zero_bin, config.group_size, axis=-1) + return (weight_bin - zero_bin_repeated) * scale_repeated + + +@pytest.mark.parametrize( + "quant_name, shape, dtype", + [ + ("q4f16_awq", [2, 4096], "float16"), + ], +) +def test_dequantize_weight(quant_name: str, shape: List[int], dtype: str): # noqa: UP006 + class Test(nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(shape[1], shape[0], bias=False, dtype=dtype) + + def forward(self, x: nn.Tensor): + return self.linear(x) + + config = QUANTIZATION[quant_name] + assert isinstance(config, AWQQuantize) + weight_np = np.random.randint( + np.iinfo(config.storage_dtype).min, + np.iinfo(config.storage_dtype).max, + (shape[0], shape[1] // config.num_elem_per_storage), + ).astype(config.storage_dtype) + zeros_np = np.random.randint( + np.iinfo(config.storage_dtype).min, + np.iinfo(config.storage_dtype).max, + (shape[0], shape[1] // config.num_elem_per_storage // config.group_size), + ).astype(config.storage_dtype) + scale_np = np.random.random((shape[0], shape[1] // config.group_size)).astype( + config.model_dtype + ) + mod = config.quantize_model(Test(), QuantizeMapping({}, {}), "") + mod.linear.qweight.data = weight_np + mod.linear.qzeros.data = zeros_np + mod.linear.scales.data = scale_np + model = mod.jit(spec={"forward": {"x": nn.spec.Tensor((shape[1], shape[1]), dtype)}}) + out = model["forward"](torch.from_numpy(np.diag(np.ones(shape[1]).astype(dtype)))) + ref = dequantize_np(config, weight_np, zeros_np, scale_np).T + tvm.testing.assert_allclose(out, ref, rtol=1e-3, atol=1e-3) + + +if __name__ == "__main__": + test_dequantize_weight("q4f16_awq", [2, 4096], "float16") diff --git a/tests/python/quantization/test_group_quantization.py b/tests/python/quantization/test_group_quantization.py new file mode 100644 index 0000000..24f5412 --- /dev/null +++ b/tests/python/quantization/test_group_quantization.py @@ -0,0 +1,188 @@ +from typing import List, Optional # noqa: UP035 + +import numpy as np +import pytest +import torch +import tvm +import tvm.testing +from tvm import DataType +from tvm.relax.frontend import nn + +from mlc_llm.loader import QuantizeMapping +from mlc_llm.quantization import QUANTIZATION +from mlc_llm.quantization.group_quantization import ( + GroupQuantize, + GroupQuantizeEmbedding, + GroupQuantizeLinear, +) + + +def quantize_np(config: GroupQuantize, weight: np.ndarray): + n, k = weight.shape + weight_padded = np.pad( + weight, + ((0, 0), (0, (config.group_size - k % config.group_size) % config.group_size)), + ) + n, k = weight_padded.shape + weight_reshaped = np.reshape(weight_padded, (n, k // config.group_size, config.group_size)) + max_abs = np.maximum(np.max(np.abs(weight_reshaped), axis=-1), 1e-4) + scale = np.divide(max_abs, config.max_int_value) + scale_reshaped = np.reshape(scale, (*scale.shape, 1)) + weight_scaled_reshaped = np.clip( + np.add( + np.round(np.divide(weight_reshaped, scale_reshaped)), + config.max_int_value, + ), + 0, + config.max_int_value * 2, + ).astype(config.storage_dtype) + weight_filtered = np.reshape(weight_scaled_reshaped, (n, k)) + weight_filtered[..., weight.shape[1] :] = 0 + weight_scaled = np.reshape( + weight_filtered, + (n, k // config.num_elem_per_storage, config.num_elem_per_storage), + ) + indice_k = np.indices(weight_scaled.shape, dtype=config.storage_dtype)[-1] + quantized_weight = np.sum( + np.left_shift(weight_scaled, indice_k * DataType(config.quantize_dtype).bits), + axis=-1, + dtype=config.storage_dtype, + ) + return quantized_weight, scale + + +def dequantize_np( + config: GroupQuantize, + weight: np.ndarray, + scale: np.ndarray, + out_shape: Optional[List[int]] = None, # noqa: UP006 +): + assert weight.shape[0] == scale.shape[0] + bin_mask = (1 << DataType(config.quantize_dtype).bits) - 1 + max_int = config.max_int_value + out_shape = ( + [weight.shape[0], weight.shape[1] * config.num_elem_per_storage] + if out_shape is None + else out_shape + ) + weight_repeated = np.repeat(weight, config.num_elem_per_storage, axis=-1) + scale_repeated = np.repeat(scale, config.group_size, axis=-1) + indice_j = np.indices(weight_repeated.shape)[1] + weight_bin = np.bitwise_and( + np.right_shift( + weight_repeated, + (indice_j % config.num_elem_per_storage) * DataType(config.quantize_dtype).bits, + ), + bin_mask, + ) + assert weight_bin.shape[1] <= scale_repeated.shape[1] + return ((weight_bin - max_int) * scale_repeated[..., : weight_bin.shape[1]])[ + : out_shape[0], : out_shape[1] + ] + + +@pytest.mark.parametrize( + "quant_name, shape, dtype, device", + [ + ("q3f16_1", [2, 13], "float16", "cpu"), + ("q3f16_1", [16, 120], "float16", "cpu"), + ("q4f16_1", [2, 13], "float16", "cpu"), + ("q4f16_1", [16, 128], "float16", "cpu"), + ("q4f32_1", [2, 13], "float32", "cpu"), + ("q4f32_1", [16, 128], "float32", "cpu"), + ], +) +def test_quantize_weight(quant_name: str, shape: List[int], dtype: str, device: str): # noqa: UP006 + config = QUANTIZATION[quant_name] + assert isinstance(config, GroupQuantize) + weight_np = np.random.random(shape).astype(dtype) + output = config.quantize_weight(tvm.runtime.tensor(weight_np, device=tvm.device(device))) + quantized_weight, scale = output[0].numpy(), output[1].numpy() + quantized_weight_ref, scale_ref = quantize_np(config, weight_np) + tvm.testing.assert_allclose(scale, scale_ref, rtol=1e-3, atol=1e-3) + tvm.testing.assert_allclose( + dequantize_np(config, quantized_weight, scale, shape), + dequantize_np(config, quantized_weight_ref, scale_ref, shape), + rtol=1e-2 if quant_name.startswith("q3") else 1e-3, + atol=0.4 if quant_name.startswith("q3") else 0.2, + ) + + +@pytest.mark.parametrize( + "quant_name, shape, dtype", + [ + ("q3f16_1", [2, 13], "float16"), + ("q3f16_1", [16, 120], "float16"), + ("q4f16_1", [2, 13], "float16"), + ("q4f16_1", [16, 128], "float16"), + ("q4f32_1", [2, 13], "float32"), + ("q4f32_1", [16, 128], "float32"), + ], +) +def test_dequantize_weight(quant_name: str, shape: List[int], dtype: str): # noqa: UP006 + class Test(nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(shape[1], shape[0], bias=False, dtype=dtype) + + def forward(self, x: nn.Tensor): + return self.linear(x) + + config = QUANTIZATION[quant_name] + assert isinstance(config, GroupQuantize) + num_group = -(shape[1] // -config.group_size) + weight_np = np.random.randint( + np.iinfo(config.storage_dtype).min, + np.iinfo(config.storage_dtype).max, + (shape[0], config.num_storage_per_group * num_group), + ).astype(config.storage_dtype) + scale_np = np.random.random((shape[0], num_group)).astype(config.model_dtype) + mod = config.quantize_model(Test(), QuantizeMapping({}, {}), "") + mod.linear.q_weight.data = weight_np + mod.linear.q_scale.data = scale_np + model = mod.jit(spec={"forward": {"x": nn.spec.Tensor((shape[1], shape[1]), dtype)}}) + out = model["forward"](torch.from_numpy(np.diag(np.ones(shape[1]).astype(dtype)))) + ref = dequantize_np(config, weight_np, scale_np, shape).T + tvm.testing.assert_allclose(out, ref, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize( + "quant_name, shape, dtype", + [ + ("q3f16_1", [16, 128], "float16"), + ("q4f16_1", [16, 128], "float16"), + ("q4f32_1", [16, 128], "float32"), + ], +) +def test_quantize_model(quant_name: str, shape: List[int], dtype: str): # noqa: UP006 + class Test(nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(shape[0], shape[1], dtype=dtype) + self.embedding = nn.Embedding(shape[0], shape[1], dtype=dtype) + + def forward(self, x: nn.Tensor): + return self.linear(x) + + config = QUANTIZATION[quant_name] + assert isinstance(config, GroupQuantize) + quant_map = QuantizeMapping({}, {}) + mod = config.quantize_model(Test(), quant_map, "model") + assert quant_map.param_map["model.linear.weight"] == [ + "model.linear.q_weight", + "model.linear.q_scale", + ] + assert quant_map.map_func["model.linear.weight"] == config.quantize_weight + assert isinstance(mod.linear, GroupQuantizeLinear) + assert quant_map.param_map["model.embedding.weight"] == [ + "model.embedding.q_weight", + "model.embedding.q_scale", + ] + assert quant_map.map_func["model.embedding.weight"] == config.quantize_weight + assert isinstance(mod.embedding, GroupQuantizeEmbedding) + + +if __name__ == "__main__": + test_quantize_weight("q4f16_1", [16, 128], "float16", "llvm") + test_quantize_model("q4f16_1", [16, 128], "float16") + test_dequantize_weight("q4f16_1", [16, 128], "float16") diff --git a/tests/python/router/test_router.py b/tests/python/router/test_router.py new file mode 100644 index 0000000..067ed3c --- /dev/null +++ b/tests/python/router/test_router.py @@ -0,0 +1,101 @@ +import asyncio + +from mlc_llm.protocol import openai_api_protocol +from mlc_llm.router import Router + +model_tp1 = "./dist/Llama-3.2-1B-Instruct-q0f16-MLC/" +model_lib_tp1 = "./dist/lib/Llama-3.2-1B-q0f16-cuda.so" +# model_lib_tp1 = None + +model_tp2 = "./dist/Llama-3.2-1B-Instruct-q0f16-MLC-tp2/" +model_lib_tp2 = "./dist/lib/Llama-3.2-1B-q0f16-cuda-tp2.so" +# model_lib_tp2 = None + + +def get_router_1tp1(): + return ( + Router( + model_tp1, + model_lib=model_lib_tp1, + hosts=["127.0.0.1"], + ports=[8080], + ), + model_tp1, + ) + + +def get_router_2tp1(): + return ( + Router( + model_tp1, + model_lib=model_lib_tp1, + hosts=["127.0.0.1", "127.0.0.1"], + ports=[8080, 8081], + device_id_starts=[0, 1], + npes=2, + ), + model_tp1, + ) + + +def get_router_1tp2(): + return ( + Router( + model_tp2, + model_lib=model_lib_tp2, + hosts=["127.0.0.1"], + ports=[8080], + npes=2, + ), + model_tp2, + ) + + +def get_router_2tp2(): + return ( + Router( + model_tp2, + model_lib=model_lib_tp2, + hosts=["127.0.0.1", "127.0.0.1"], + ports=[8080, 8081], + device_id_starts=[0, 2], + npes=4, + ), + model_tp2, + ) + + +CONFIG_TO_ROUTER = { + "1tp1": get_router_1tp1, + "2tp1": get_router_2tp1, + "1tp2": get_router_1tp2, + "2tp2": get_router_2tp2, +} + + +async def test_router(schedule: str = "round_robin", endpoints_config: str = "1tp1"): + router, model_id = CONFIG_TO_ROUTER[endpoints_config]() + + request = openai_api_protocol.CompletionRequest( + prompt="The meaning of life ", + model=model_id, + stream=True, + max_tokens=64, + stream_options=openai_api_protocol.StreamOptions(include_usage=True), + ) + if schedule == "round_robin": + async for chunk in router._handle_completion_round_robin(request, "1"): + print(chunk) + elif schedule == "disagg": + async for chunk in router._handle_completion_disagg(request, "1"): + print(chunk) + else: + raise ValueError(f"Unknown scheduling method: {schedule}") + router.terminate() + + +if __name__ == "__main__": + # asyncio.run(test_router("round_robin", endpoints_config="1tp1")) + # asyncio.run(test_router("round_robin", endpoints_config="1tp2")) + # asyncio.run(test_router("round_robin", endpoints_config="2tp1")) + asyncio.run(test_router("round_robin", endpoints_config="2tp2")) diff --git a/tests/python/serve/evaluate_engine.py b/tests/python/serve/evaluate_engine.py new file mode 100644 index 0000000..5af73e8 --- /dev/null +++ b/tests/python/serve/evaluate_engine.py @@ -0,0 +1,73 @@ +import argparse +import os +import random +from typing import List, Tuple # noqa: UP035 + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine + + +def _parse_args(): + args = argparse.ArgumentParser() + args.add_argument("--model-lib", type=str) + args.add_argument("--device", type=str, default="auto") + args.add_argument("--batch-size", type=int, default=80) + args.add_argument("--max-total-seq-length", type=int) + args.add_argument("--seed", type=int, default=0) + + parsed = args.parse_args() + parsed.model = os.path.dirname(parsed.model_lib) + assert parsed.batch_size % 16 == 0 + return parsed + + +def generate_requests( + num_requests: int, input_length: int, output_length: int +) -> Tuple[List[List[int]], List[GenerationConfig]]: # noqa: UP006 + prompt_ids = [] + for _ in range(num_requests): + token_ids = [] + for _ in range(input_length): + token_ids.append(random.randint(0, 30000)) + prompt_ids.append(token_ids) + generation_config_list = [ + GenerationConfig(temperature=1.0, top_p=1.0, max_tokens=output_length) + ] * num_requests + return prompt_ids, generation_config_list + + +def benchmark(args: argparse.Namespace): + random.seed(args.seed) + + # Create engine + engine = SyncMLCEngine( + model=args.model, + device=args.device, + model_lib=args.model_lib, + mode="server", + engine_config=EngineConfig( + max_num_sequence=args.batch_size, + max_total_sequence_length=args.max_total_seq_length, + ), + ) + + print(args) + for num_requests in [1, 2, 4, 8, 16, 32, 64]: + if num_requests > args.batch_size: + continue + for input_length in [64, 128, 256, 512, 1024]: + if num_requests * input_length >= 16384: + continue + for output_length in [4]: + print(f"nreq={num_requests}\tin={input_length}\tout={output_length}") + prompt_ids, generation_config = generate_requests( + num_requests, input_length, output_length + ) + engine.reset() + engine.generate(prompt_ids, generation_config) + print() + + +if __name__ == "__main__": + ARGS = _parse_args() + benchmark(ARGS) diff --git a/tests/python/serve/server/conftest.py b/tests/python/serve/server/conftest.py new file mode 100644 index 0000000..a9c052c --- /dev/null +++ b/tests/python/serve/server/conftest.py @@ -0,0 +1,34 @@ +import os +from typing import Tuple # noqa: UP035 + +import pytest + +from mlc_llm.serve import PopenServer + + +@pytest.fixture(scope="session") +def served_model() -> Tuple[str, str]: # noqa: UP006 + model_lib = os.environ.get("MLC_SERVE_MODEL_LIB") + if model_lib is None: + raise ValueError( + 'Environment variable "MLC_SERVE_MODEL_LIB" not found. ' + "Please set it to model lib compiled by MLC LLM " + "(e.g., `dist/Llama-2-7b-chat-hf-q0f16-MLC/Llama-2-7b-chat-hf-q0f16-MLC-cuda.so`)." + ) + model = os.path.dirname(model_lib) + return model, model_lib + + +@pytest.fixture(scope="session") +def launch_server(served_model): + """A pytest session-level fixture which launches the server in a subprocess.""" + server = PopenServer( + model=served_model[0], + model_lib=served_model[1], + enable_tracing=True, + enable_debug=True, + port=8000, + ) + + with server: + yield diff --git a/tests/python/serve/server/test_embedding_server.py b/tests/python/serve/server/test_embedding_server.py new file mode 100644 index 0000000..423c758 --- /dev/null +++ b/tests/python/serve/server/test_embedding_server.py @@ -0,0 +1,369 @@ +"""Embedding server endpoint tests in MLC LLM. + +Tests the /v1/embeddings endpoint via HTTP using the OpenAI client, +following the same patterns as test_server.py. + +Reuses MLC LLM test infrastructure: + - Pytest markers (endpoint) + - expect_error() response validation pattern from test_server.py + - OpenAI client usage pattern from test_server.py + - Session-scoped server fixture pattern from conftest.py + +Run (launches its own embedding-only server): + MLC_SERVE_EMBEDDING_MODEL_LIB="path/to/model.dylib" \ + pytest -m endpoint tests/python/serve/server/test_embedding_server.py -v + +Environment variables: + MLC_SERVE_EMBEDDING_MODEL_LIB Path to compiled embedding model library (required) + MLC_SERVE_EMBEDDING_MODEL Path to embedding model weight directory + (optional, defaults to dirname of model lib) +""" + +import json +import os +import signal +import subprocess +import sys +import time +from pathlib import Path +from typing import Dict, Optional # noqa: UP035 + +import numpy as np +import pytest +import requests +from openai import OpenAI + +# Reuse MLC LLM marker system +pytestmark = [pytest.mark.endpoint] + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + +EMBEDDING_MODEL_LIB = os.environ.get("MLC_SERVE_EMBEDDING_MODEL_LIB") +EMBEDDING_MODEL_DIR = os.environ.get( + "MLC_SERVE_EMBEDDING_MODEL", + os.path.dirname(EMBEDDING_MODEL_LIB) if EMBEDDING_MODEL_LIB else None, +) +EMBEDDING_SERVER_HOST = "127.0.0.1" +EMBEDDING_SERVER_PORT = 8321 +EMBEDDING_BASE_URL = f"http://{EMBEDDING_SERVER_HOST}:{EMBEDDING_SERVER_PORT}/v1" +EMBEDDING_MODEL_NAME = "embedding" + + +def _skip_if_no_model(): + if EMBEDDING_MODEL_LIB is None: + pytest.skip( + 'Environment variable "MLC_SERVE_EMBEDDING_MODEL_LIB" not found. ' + "Set it to a compiled embedding model library." + ) + if not os.path.isfile(EMBEDDING_MODEL_LIB): + pytest.skip(f"Embedding model library not found at: {EMBEDDING_MODEL_LIB}") + if EMBEDDING_MODEL_DIR is None or not os.path.isdir(EMBEDDING_MODEL_DIR): + pytest.skip(f"Embedding model directory not found at: {EMBEDDING_MODEL_DIR}") + + +# --------------------------------------------------------------------------- +# Response validation helpers — adapted from test_server.py patterns +# --------------------------------------------------------------------------- + + +def check_embedding_response( + response: Dict, # noqa: UP006 + *, + model: str, + num_embeddings: int, + expected_dim: Optional[int] = None, + check_unit_norm: bool = True, +): + """Validate an OpenAI-compatible embedding response. + + Adapted from check_openai_nonstream_response() in test_server.py, + specialized for embedding responses. + """ + assert response["object"] == "list" + assert response["model"] == model + + data = response["data"] + assert isinstance(data, list) + assert len(data) == num_embeddings + + for item in data: + assert item["object"] == "embedding" + assert isinstance(item["index"], int) + emb = item["embedding"] + assert isinstance(emb, list) + assert len(emb) > 0 + + if expected_dim is not None: + assert len(emb) == expected_dim, f"Expected dim={expected_dim}, got {len(emb)}" + + if check_unit_norm: + norm = float(np.linalg.norm(emb)) + assert abs(norm - 1.0) < 1e-3, f"Expected unit norm, got {norm}" + + # Usage validation — same pattern as test_server.py + usage = response["usage"] + assert isinstance(usage, dict) + assert usage["prompt_tokens"] > 0 + assert usage["total_tokens"] == usage["prompt_tokens"] + + +def expect_error(response_str: str, msg_prefix: Optional[str] = None): + """Validate error response — reused directly from test_server.py.""" + response = json.loads(response_str) + assert response["object"] == "error" + assert isinstance(response["message"], str) + if msg_prefix is not None: + assert response["message"].startswith(msg_prefix) + + +# --------------------------------------------------------------------------- +# Server fixture — follows PopenServer/launch_server pattern from conftest.py +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def launch_embedding_server(): + """Launch an embedding-only server as a subprocess. + + Follows the same lifecycle pattern as the launch_server fixture + in serve/server/conftest.py, but uses a lightweight embedding-only + server since PopenServer doesn't support --embedding-model yet. + """ + _skip_if_no_model() + + mlc_llm_path = str(Path(__file__).resolve().parents[4] / "python") + server_code = f""" +import sys +sys.path.insert(0, "{mlc_llm_path}") + +import fastapi +import uvicorn +from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine +from mlc_llm.serve.server import ServerContext +from mlc_llm.serve.entrypoints import openai_entrypoints + +app = fastapi.FastAPI() +app.include_router(openai_entrypoints.app) + +engine = AsyncEmbeddingEngine( + model="{EMBEDDING_MODEL_DIR}", + model_lib="{EMBEDDING_MODEL_LIB}", + device="auto", +) +ctx = ServerContext() +ServerContext.server_context = ctx +ctx.add_embedding_engine("{EMBEDDING_MODEL_NAME}", engine) + +uvicorn.run(app, host="{EMBEDDING_SERVER_HOST}", port={EMBEDDING_SERVER_PORT}, log_level="info") +""" + with subprocess.Popen( + [sys.executable, "-c", server_code], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) as proc: + # Wait for server readiness — same polling pattern as PopenServer.start() + timeout = 120 + attempts = 0.0 + ready = False + while attempts < timeout: + try: + response = requests.get(f"{EMBEDDING_BASE_URL}/models", timeout=2) + if response.status_code == 200: + ready = True + break + except requests.RequestException: + pass + attempts += 0.5 + time.sleep(0.5) + + if not ready: + stderr = proc.stderr.read().decode() if proc.stderr else "" + proc.kill() + raise RuntimeError(f"Embedding server failed to start in {timeout}s.\nStderr: {stderr}") + + yield proc + + # Cleanup — same pattern as PopenServer.terminate() + proc.send_signal(signal.SIGINT) + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + + +@pytest.fixture(scope="module") +def client(launch_embedding_server): + """OpenAI client connected to the embedding server.""" + assert launch_embedding_server is not None + return OpenAI(base_url=EMBEDDING_BASE_URL, api_key="none") + + +# =================================================================== +# /v1/models +# =================================================================== + + +@pytest.mark.usefixtures("client") +def test_models_endpoint(): + """The /v1/models endpoint lists the embedding model.""" + resp = requests.get(f"{EMBEDDING_BASE_URL}/models", timeout=5) + assert resp.status_code == 200 + data = resp.json() + assert isinstance(data["data"], list) + + +# =================================================================== +# Single input +# =================================================================== + + +def test_single_string_input(client): + """Single string input returns one embedding.""" + resp = client.embeddings.create(input="What is machine learning?", model=EMBEDDING_MODEL_NAME) + raw = resp.model_dump() + check_embedding_response(raw, model=EMBEDDING_MODEL_NAME, num_embeddings=1) + + +# =================================================================== +# Batch input +# =================================================================== + +BATCH_INPUTS = [ + "What is machine learning?", + "How to brew coffee?", + "ML is a subset of AI.", +] + + +def test_batch_string_input(client): + """List of strings returns one embedding per input.""" + resp = client.embeddings.create(input=BATCH_INPUTS, model=EMBEDDING_MODEL_NAME) + raw = resp.model_dump() + check_embedding_response(raw, model=EMBEDDING_MODEL_NAME, num_embeddings=len(BATCH_INPUTS)) + + +def test_batch_index_ordering(client): + """Embedding indices are sequential.""" + resp = client.embeddings.create(input=BATCH_INPUTS, model=EMBEDDING_MODEL_NAME) + indices = [d.index for d in resp.data] + assert indices == list(range(len(BATCH_INPUTS))) + + +# =================================================================== +# Cosine similarity — semantic quality via endpoint +# =================================================================== + + +def test_cosine_similarity_via_endpoint(client): + """Related texts have higher similarity than unrelated (end-to-end).""" + resp = client.embeddings.create( + input=[ + "What is machine learning?", + "Explain deep learning", + "Order a pizza", + ], + model=EMBEDDING_MODEL_NAME, + ) + e0, e1, e2 = [np.array(d.embedding) for d in resp.data] + sim_related = float(np.dot(e0, e1)) + sim_unrelated = float(np.dot(e0, e2)) + assert sim_related > sim_unrelated, ( + f"Related ({sim_related:.4f}) should > unrelated ({sim_unrelated:.4f})" + ) + + +# =================================================================== +# Dimension truncation (Matryoshka) +# =================================================================== + + +def test_dimension_truncation(client): + """dimensions parameter truncates and re-normalizes output.""" + target_dim = 256 + resp = client.embeddings.create( + input="Hello world", model=EMBEDDING_MODEL_NAME, dimensions=target_dim + ) + raw = resp.model_dump() + check_embedding_response( + raw, + model=EMBEDDING_MODEL_NAME, + num_embeddings=1, + expected_dim=target_dim, + ) + + +# =================================================================== +# Encoding format +# =================================================================== + + +@pytest.mark.usefixtures("launch_embedding_server") +def test_base64_encoding(): + """base64 encoding format returns base64-encoded embeddings.""" + resp = requests.post( + f"{EMBEDDING_BASE_URL}/embeddings", + json={ + "input": "Hello world", + "model": EMBEDDING_MODEL_NAME, + "encoding_format": "base64", + }, + timeout=5, + ) + assert resp.status_code == 200 + data = resp.json() + assert data["data"][0]["object"] == "embedding" + # base64 string should be a non-empty string (not a list) + emb = data["data"][0]["embedding"] + assert isinstance(emb, str) and len(emb) > 0 + + +# =================================================================== +# Error handling — reuses expect_error() pattern from test_server.py +# =================================================================== + + +@pytest.mark.usefixtures("launch_embedding_server") +def test_any_model_name_works_with_single_engine(): + """When only one embedding engine is served, any model name works. + + This mirrors ServerContext.get_engine() behavior: a single served + model is returned regardless of the requested model name. + """ + resp = requests.post( + f"{EMBEDDING_BASE_URL}/embeddings", + json={"input": "test", "model": "any-name-works"}, + timeout=5, + ) + assert resp.status_code == 200 + data = resp.json() + assert len(data["data"]) == 1 + + +# =================================================================== +# Standalone runner (same pattern as test_server.py __main__) +# =================================================================== + +if __name__ == "__main__": + _skip_if_no_model() + + print(f"Using model: {EMBEDDING_MODEL_DIR}") + print(f"Using model lib: {EMBEDDING_MODEL_LIB}") + print(f"Server URL: {EMBEDDING_BASE_URL}") + print( + "\nMake sure the embedding server is running, or set env vars " + "and use pytest to auto-launch." + ) + + # Allow running against an already-running server + c = OpenAI(base_url=EMBEDDING_BASE_URL, api_key="none") + test_models_endpoint() + test_single_string_input(c) + test_batch_string_input(c) + test_batch_index_ordering(c) + test_cosine_similarity_via_endpoint(c) + test_dimension_truncation(c) + test_base64_encoding() + test_any_model_name_works_with_single_engine() + print("\nAll embedding server tests passed!") diff --git a/tests/python/serve/server/test_server.py b/tests/python/serve/server/test_server.py new file mode 100644 index 0000000..ed0642d --- /dev/null +++ b/tests/python/serve/server/test_server.py @@ -0,0 +1,1386 @@ +"""Server tests in MLC LLM. +Before running any test, we use pytest fixtures to launch a +test-session-wide server in a subprocess, and then execute the tests. + +The recommended way to run the tests is to use the following command: + MLC_SERVE_MODEL_LIB="YOUR_MODEL_LIB" pytest -vv tests/python/serve/server/test_server.py + +Here "YOUR_MODEL_LIB" is a compiled model library like +`dist/Llama-2-7b-chat-hf-q4f16_1/Llama-2-7b-chat-hf-q4f16_1-cuda.so`, +as long as the model is built with batching and embedding separation enabled. + +To directly run the Python file (a.k.a., not using pytest), you need to +launch the server in ahead before running this file. This can be done in +two steps: +- start a new shell session, run + python -m mlc_llm.serve.server --model "YOUR_MODEL_LIB" +- start another shell session, run this file + MLC_SERVE_MODEL_LIB="YOUR_MODEL_LIB" python tests/python/serve/server/test_server.py +""" + +import json +import os +from http import HTTPStatus +from typing import Dict, List, Optional, Tuple # noqa: UP035 + +import pytest +import regex +import requests +from openai import OpenAI +from pydantic import BaseModel + +from mlc_llm.protocol.openai_api_protocol import ( + CHAT_COMPLETION_MAX_TOP_LOGPROBS, + COMPLETION_MAX_TOP_LOGPROBS, +) + +OPENAI_BASE_URL = "http://127.0.0.1:8000/v1" +OPENAI_V1_MODELS_URL = "http://127.0.0.1:8000/v1/models" +OPENAI_V1_COMPLETION_URL = "http://127.0.0.1:8000/v1/completions" +OPENAI_V1_CHAT_COMPLETION_URL = "http://127.0.0.1:8000/v1/chat/completions" +DEBUG_DUMP_EVENT_TRACE_URL = "http://127.0.0.1:8000/debug/dump_event_trace" +METRICS_URL = "http://127.0.0.1:8000/metrics" + + +JSON_TOKEN_PATTERN = ( + r"((-?(?:0|[1-9]\d*))(\.\d+)?([eE][-+]?\d+)?)|null|true|false|" + r'("((\\["\\\/bfnrt])|(\\u[0-9a-fA-F]{4})|[^"\\\x00-\x1f])*")' +) +JSON_TOKEN_RE = regex.compile(JSON_TOKEN_PATTERN) + + +def is_json(s: str) -> bool: + try: + json.loads(s) + return True + except json.JSONDecodeError: + return False + + +def is_json_prefix(s: str) -> bool: + try: + json.loads(s) + return True + except json.JSONDecodeError as e: + # If the JSON decoder reaches the end of s, it is a prefix of a JSON string. + if e.pos == len(s): + return True + # Since json.loads is token-based instead of char-based, there may remain half a token after + # the matching position. + # If the left part is a prefix of a valid JSON token, the output is also valid + regex_match = JSON_TOKEN_RE.fullmatch(s[e.pos :], partial=True) + return regex_match is not None + + +def check_openai_nonstream_response( + response: Dict, # noqa: UP006 + *, + is_chat_completion: bool, + model: str, + object_str: str, + num_choices: int, + finish_reasons: List[str], # noqa: UP006 + completion_tokens: Optional[int] = None, + echo_prompt: Optional[str] = None, + suffix: Optional[str] = None, + stop: Optional[List[str]] = None, # noqa: UP006 + require_substr: Optional[List[str]] = None, # noqa: UP006 + check_json_output: bool = False, +): + assert response["model"] == model + assert response["object"] == object_str + + choices = response["choices"] + assert isinstance(choices, list) + assert len(choices) <= num_choices + texts: List[str] = ["" for _ in range(num_choices)] # noqa: UP006 + for choice in choices: + idx = choice["index"] + assert choice["finish_reason"] in finish_reasons + + if not is_chat_completion: + assert isinstance(choice["text"], str) + texts[idx] = choice["text"] + if echo_prompt is not None: + assert texts[idx] + if suffix is not None: + assert texts[idx] + else: + message = choice["message"] + assert message["role"] == "assistant" + assert isinstance(message["content"], str) + texts[idx] = message["content"] + + if stop is not None: + for stop_str in stop: + assert stop_str not in texts[idx] + if require_substr is not None: + for substr in require_substr: + assert substr in texts[idx] + if check_json_output: + # the output should be json or a prefix of a json string + # if the output is a prefix of a json string, the output must exceed the max output + # length + output_is_json = is_json(texts[idx]) + output_is_json_prefix = is_json_prefix(texts[idx]) + assert output_is_json or output_is_json_prefix + if not output_is_json and output_is_json_prefix: + assert choice["finish_reason"] == "length" + + usage = response["usage"] + if usage is not None: + assert isinstance(usage, dict) + assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"] + assert usage["prompt_tokens"] > 0 + if completion_tokens is not None: + assert usage["completion_tokens"] == completion_tokens + + +def check_openai_stream_response( + responses: List[Dict], # noqa: UP006 + *, + is_chat_completion: bool, + model: str, + object_str: str, + num_choices: int, + finish_reasons: List[str], # noqa: UP006 + completion_tokens: Optional[int] = None, + echo_prompt: Optional[str] = None, + suffix: Optional[str] = None, + stop: Optional[List[str]] = None, # noqa: UP006 + require_substr: Optional[List[str]] = None, # noqa: UP006 + check_json_output: bool = False, +): + assert len(responses) > 0 + + finished = [False for _ in range(num_choices)] + outputs = ["" for _ in range(num_choices)] + finish_reason_list = ["" for _ in range(num_choices)] + for response in responses: + assert response["model"] == model + assert response["object"] == object_str + + choices = response["choices"] + assert isinstance(choices, list) + assert len(choices) <= num_choices + for choice in choices: + idx = choice["index"] + + if not is_chat_completion: + assert isinstance(choice["text"], str) + outputs[idx] += choice["text"] + else: + delta = choice["delta"] + assert delta["role"] == "assistant" + assert isinstance(delta["content"], str) + outputs[idx] += delta["content"] + + if finished[idx]: + assert choice["finish_reason"] in finish_reasons + finish_reason_list[idx] = choice["finish_reason"] + elif choice["finish_reason"] is not None: + assert choice["finish_reason"] in finish_reasons + finish_reason_list[idx] = choice["finish_reason"] + finished[idx] = True + + if not is_chat_completion: + usage = response["usage"] + if usage is not None: + assert isinstance(usage, dict) + assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"] + assert usage["prompt_tokens"] >= 0 + if completion_tokens is not None: + assert usage["completion_tokens"] <= completion_tokens + + if not is_chat_completion: + if completion_tokens is not None and responses[-1]["usage"] is not None: + assert responses[-1]["usage"]["completion_tokens"] == completion_tokens + + for i, (output, finish_reason) in enumerate(zip(outputs, finish_reason_list)): + if echo_prompt is not None: + assert output.startswith(echo_prompt) + if suffix is not None: + assert output.endswith(suffix) + if stop is not None: + for stop_str in stop: + assert stop_str not in output + if require_substr is not None: + for substr in require_substr: + assert substr in output + if check_json_output: + # the output should be json or a prefix of a json string + # if the output is a prefix of a json string, the output must exceed the max output + # length + output_is_json = is_json(output) + output_is_json_prefix = is_json_prefix(output) + assert output_is_json or output_is_json_prefix + if not output_is_json and output_is_json_prefix: + assert finish_reason == "length" + + +def expect_error(response_str: str, msg_prefix: Optional[str] = None): + response = json.loads(response_str) + assert response["object"] == "error" + assert isinstance(response["message"], str) + if msg_prefix is not None: + assert response["message"].startswith(msg_prefix) + + +def test_openai_v1_models( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + response = requests.get(OPENAI_V1_MODELS_URL, timeout=180).json() + assert response["object"] == "list" + models = response["data"] + assert isinstance(models, list) + assert len(models) == 1 + + model_card = models[0] + assert isinstance(model_card, dict) + assert model_card["id"] == served_model[0], f"{model_card['id']} {served_model[0]}" + assert model_card["object"] == "model" + assert model_card["owned_by"] == "MLC-LLM" + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = "What is the meaning of life?" + max_tokens = 256 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stream": stream, + "debug_config": {"ignore_eos": True}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_openai_package( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + client = OpenAI(base_url=OPENAI_BASE_URL, api_key="None") + prompt = "What is the meaning of life?" + max_tokens = 256 + response = client.completions.create( + model=served_model[0], + prompt=prompt, + max_tokens=max_tokens, + stream=stream, + ) + if not stream: + check_openai_nonstream_response( + response.model_dump(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length", "stop"], + completion_tokens=max_tokens, + ) + else: + responses = [] + for chunk in response: + responses.append(chunk.model_dump()) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length", "stop"], + completion_tokens=max_tokens, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_echo( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = "What is the meaning of life?" + max_tokens = 256 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "echo": True, + "stream": stream, + "debug_config": {"ignore_eos": True}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + echo_prompt=prompt, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + echo_prompt=prompt, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_suffix( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = "What is the meaning of life?" + suffix = "Hello, world!" + max_tokens = 256 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "suffix": suffix, + "stream": stream, + "debug_config": {"ignore_eos": True}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + suffix=suffix, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + suffix=suffix, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_stop_str( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + # Choose "in" as the stop string since it is very unlikely that + # "in" does not appear in the generated output. + prompt = "What is the meaning of life?" + stop = ["in"] + max_tokens = 256 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stop": stop, + "stream": stream, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["stop", "length"], + stop=stop, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["stop", "length"], + stop=stop, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_temperature( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = "What's the meaning of life?" + max_tokens = 128 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stream": stream, + "temperature": 0.0, + "debug_config": {"ignore_eos": True}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_json( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = "Response with a json object:" + max_tokens = 128 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stream": stream, + "response_format": {"type": "json_object"}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=60) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_json_schema( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = ( + "Generate a json containing three fields: an integer field named size, a " + "boolean field named is_accepted, and a float field named num:" + ) + max_tokens = 128 + + class Schema(BaseModel): + size: int + is_accepted: bool + num: float + + schema_str = json.dumps(Schema.model_json_schema()) + + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stream": stream, + "response_format": {"type": "json_object", "schema": schema_str}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=60) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_logit_bias( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + # NOTE: This test only tests that the system does not break on logit bias. + # The test does not promise the correctness of logit bias handling. + + prompt = "What's the meaning of life?" + max_tokens = 128 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stream": stream, + "logit_bias": {338: -100}, # 338 is " is" in Llama tokenizer. + "debug_config": {"ignore_eos": True}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_presence_frequency_penalty( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = "What's the meaning of life?" + max_tokens = 128 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stream": stream, + "frequency_penalty": 2.0, + "presence_penalty": 2.0, + "debug_config": {"ignore_eos": True}, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + ) + + +def test_openai_v1_completions_seed( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = "What's the meaning of life?" + max_tokens = 128 + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": max_tokens, + "stream": False, + "seed": 233, + "debug_config": {"ignore_eos": True}, + } + + response1 = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + response2 = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + for response in [response1, response2]: + check_openai_nonstream_response( + response.json(), + is_chat_completion=False, + model=served_model[0], + object_str="text_completion", + num_choices=1, + finish_reasons=["length"], + ) + + text1 = response1.json()["choices"][0]["text"] + text2 = response2.json()["choices"][0]["text"] + assert text1 == text2 + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_prompt_overlong( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + num_tokens = 1000000 + prompt = [128] * num_tokens + payload = { + "model": served_model[0], + "prompt": prompt, + "max_tokens": 256, + "stream": stream, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + error_msg_prefix = ( + f"Request prompt has {num_tokens} tokens in total, larger than the model input length limit" + ) + if not stream: + expect_error(response.json(), msg_prefix=error_msg_prefix) + else: + num_chunks = 0 + for chunk in response.iter_lines(chunk_size=512): + if not chunk: + continue + num_chunks += 1 + expect_error(json.loads(chunk.decode("utf-8")), msg_prefix=error_msg_prefix) + assert num_chunks == 1 + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_completions_invalid_logprobs( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + payload = { + "model": served_model[0], + "prompt": "What is the meaning of life?", + "max_tokens": 256, + "stream": stream, + "logprobs": COMPLETION_MAX_TOP_LOGPROBS + 1, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + assert response.json()["detail"][0]["msg"].endswith( + f'"top_logprobs" must be in range [0, {COMPLETION_MAX_TOP_LOGPROBS}]' + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_chat_completions_invalid_logprobs( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + payload = { + "model": served_model[0], + "messages": [{"role": "user", "content": "Hello! Our project is MLC LLM."}], + "max_tokens": 256, + "stream": stream, + "logprobs": False, + "top_logprobs": CHAT_COMPLETION_MAX_TOP_LOGPROBS - 1, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + assert response.json()["detail"][0]["msg"].endswith( + '"logprobs" must be True to support "top_logprobs"' + ) + + payload["logprobs"] = True + payload["top_logprobs"] = CHAT_COMPLETION_MAX_TOP_LOGPROBS + 1 + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + assert response.status_code == HTTPStatus.UNPROCESSABLE_ENTITY + assert response.json()["detail"][0]["msg"].endswith( + f'"top_logprobs" must be in range [0, {CHAT_COMPLETION_MAX_TOP_LOGPROBS}]' + ) + + +def test_openai_v1_completions_unsupported_args( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + # Right now "best_of" is unsupported. + best_of = 2 + payload = { + "model": served_model[0], + "prompt": "What is the meaning of life?", + "max_tokens": 256, + "best_of": best_of, + } + + response = requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=180) + error_msg_prefix = 'Request fields "best_of" are not supported right now.' + expect_error(response.json(), msg_prefix=error_msg_prefix) + + +def test_openai_v1_completions_request_cancellation( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + # Use a large max_tokens and small timeout to force timeouts. + payload = { + "model": served_model[0], + "prompt": "What is the meaning of life?", + "max_tokens": 2048, + "stream": False, + } + with pytest.raises(requests.exceptions.Timeout): + requests.post(OPENAI_V1_COMPLETION_URL, json=payload, timeout=1) + + # The server should still be alive after a request cancelled. + # We query `v1/models` to validate the server liveness. + response = requests.get(OPENAI_V1_MODELS_URL, timeout=180).json() + + assert response["object"] == "list" + models = response["data"] + assert isinstance(models, list) + assert len(models) == 1 + + model_card = models[0] + assert isinstance(model_card, dict) + assert model_card["id"] == served_model[0] + assert model_card["object"] == "model" + assert model_card["owned_by"] == "MLC-LLM" + + +CHAT_COMPLETION_MESSAGES = [ + # messages #0 + [{"role": "user", "content": "Hello! Our project is MLC LLM."}], + # messages #1 + [ + {"role": "user", "content": "Hello! Our project is MLC LLM."}, + { + "role": "assistant", + "content": "Hello! It's great to hear about your project, MLC LLM.", + }, + {"role": "user", "content": "What is the name of our project?"}, + ], + # messages #2 + [ + { + "role": "system", + "content": "You are a helpful, respectful and honest assistant. " + "You always ends your response with an emoji.", + }, + {"role": "user", "content": "Hello! Our project is MLC LLM."}, + ], +] + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES) +def test_openai_v1_chat_completions( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, + messages: List[Dict[str, str]], # noqa: UP006 +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reasons=["stop"], + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reasons=["stop"], + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES) +def test_openai_v1_chat_completions_n( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, + messages: List[Dict[str, str]], # noqa: UP006 +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + n = 3 + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + "n": n, + "max_tokens": 300, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=n, + finish_reasons=["stop", "length"], + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=n, + finish_reasons=["stop", "length"], + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES) +def test_openai_v1_chat_completions_openai_package( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, + messages: List[Dict[str, str]], # noqa: UP006 +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + client = OpenAI(base_url=OPENAI_BASE_URL, api_key="None") + response = client.chat.completions.create( + model=served_model[0], + messages=messages, + stream=stream, + logprobs=True, + top_logprobs=2, + ) + if not stream: + check_openai_nonstream_response( + response.model_dump(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reasons=["stop"], + ) + else: + responses = [] + for chunk in response: + responses.append(chunk.model_dump()) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reasons=["stop"], + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_chat_completions_max_tokens( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + messages = [{"role": "user", "content": "Write a novel with at least 500 words."}] + max_tokens = 16 + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + "max_tokens": max_tokens, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_chat_completions_json( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + messages = [{"role": "user", "content": "Response with a json object:"}] + max_tokens = 128 + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + "max_tokens": max_tokens, + "response_format": {"type": "json_object"}, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=60) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_chat_completions_json_schema( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + prompt = ( + "Generate a json containing three fields: an integer field named size, a " + "boolean field named is_accepted, and a float field named num:" + ) + messages = [{"role": "user", "content": prompt}] + max_tokens = 128 + + class Schema(BaseModel): + size: int + is_accepted: bool + num: float + + schema_str = json.dumps(Schema.model_json_schema()) + + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + "max_tokens": max_tokens, + "response_format": {"type": "json_object", "schema": schema_str}, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=60) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reasons=["length", "stop"], + check_json_output=True, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_chat_completions_ignore_eos( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + messages = [{"role": "user", "content": "Write a sentence with less than 20 words."}] + max_tokens = 128 + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + "max_tokens": max_tokens, + "debug_config": {"ignore_eos": True}, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reasons=["length"], + completion_tokens=max_tokens, + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_openai_v1_chat_completions_system_prompt_wrong_pos( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + messages = [ + {"role": "user", "content": "Hello! Our project is MLC LLM."}, + { + "role": "system", + "content": "You are a helpful, respectful and honest assistant. " + "You always ends your response with an emoji.", + }, + ] + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=180) + error_msg = "System prompt at position 1 in the message list is invalid." + if not stream: + expect_error(response.json(), msg_prefix=error_msg) + else: + num_chunks = 0 + for chunk in response.iter_lines(chunk_size=512): + if not chunk: + continue + num_chunks += 1 + expect_error(json.loads(chunk.decode("utf-8")), msg_prefix=error_msg) + assert num_chunks == 1 + + +def test_debug_dump_event_trace( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + # We only check that the request does not fail. + payload = {"model": served_model[0]} + response = requests.post(DEBUG_DUMP_EVENT_TRACE_URL, json=payload, timeout=180) + assert response.status_code == HTTPStatus.OK + + +def test_metrics( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + # We only check that the request does not fail. + metrics_text = requests.get(METRICS_URL, timeout=180).text + assert "engine_prefill_time_sum" in metrics_text + + +if __name__ == "__main__": + model_lib = os.environ.get("MLC_SERVE_MODEL_LIB") + if model_lib is None: + raise ValueError( + 'Environment variable "MLC_SERVE_MODEL_LIB" not found. ' + "Please set it to model lib compiled by MLC LLM " + "(e.g., `dist/Llama-2-7b-chat-hf-q0f16-MLC/Llama-2-7b-chat-hf-q0f16-MLC-cuda.so`)." + ) + MODEL = (os.path.dirname(model_lib), model_lib) + + test_openai_v1_models(MODEL, None) + + test_openai_v1_completions(MODEL, None, stream=False) + test_openai_v1_completions(MODEL, None, stream=True) + test_openai_v1_completions_openai_package(MODEL, None, stream=False) + test_openai_v1_completions_openai_package(MODEL, None, stream=True) + test_openai_v1_completions_echo(MODEL, None, stream=False) + test_openai_v1_completions_echo(MODEL, None, stream=True) + test_openai_v1_completions_suffix(MODEL, None, stream=False) + test_openai_v1_completions_suffix(MODEL, None, stream=True) + test_openai_v1_completions_stop_str(MODEL, None, stream=False) + test_openai_v1_completions_stop_str(MODEL, None, stream=True) + test_openai_v1_completions_temperature(MODEL, None, stream=False) + test_openai_v1_completions_temperature(MODEL, None, stream=True) + test_openai_v1_completions_logit_bias(MODEL, None, stream=False) + test_openai_v1_completions_logit_bias(MODEL, None, stream=True) + test_openai_v1_completions_presence_frequency_penalty(MODEL, None, stream=False) + test_openai_v1_completions_presence_frequency_penalty(MODEL, None, stream=True) + test_openai_v1_completions_seed(MODEL, None) + test_openai_v1_completions_prompt_overlong(MODEL, None, stream=False) + test_openai_v1_completions_prompt_overlong(MODEL, None, stream=True) + test_openai_v1_completions_invalid_logprobs(MODEL, None, stream=False) + test_openai_v1_completions_invalid_logprobs(MODEL, None, stream=True) + test_openai_v1_completions_unsupported_args(MODEL, None) + test_openai_v1_completions_request_cancellation(MODEL, None) + + for msg in CHAT_COMPLETION_MESSAGES: + test_openai_v1_chat_completions(MODEL, None, stream=False, messages=msg) + test_openai_v1_chat_completions(MODEL, None, stream=True, messages=msg) + test_openai_v1_chat_completions_n(MODEL, None, stream=False, messages=msg) + test_openai_v1_chat_completions_n(MODEL, None, stream=True, messages=msg) + test_openai_v1_chat_completions_openai_package(MODEL, None, stream=False, messages=msg) + test_openai_v1_chat_completions_openai_package(MODEL, None, stream=True, messages=msg) + test_openai_v1_chat_completions_max_tokens(MODEL, None, stream=False) + test_openai_v1_chat_completions_max_tokens(MODEL, None, stream=True) + test_openai_v1_chat_completions_json(MODEL, None, stream=False) + test_openai_v1_chat_completions_json(MODEL, None, stream=True) + test_openai_v1_chat_completions_ignore_eos(MODEL, None, stream=False) + test_openai_v1_chat_completions_ignore_eos(MODEL, None, stream=True) + test_openai_v1_chat_completions_system_prompt_wrong_pos(MODEL, None, stream=False) + test_openai_v1_chat_completions_system_prompt_wrong_pos(MODEL, None, stream=True) + + test_debug_dump_event_trace(MODEL, None) diff --git a/tests/python/serve/server/test_server_function_call.py b/tests/python/serve/server/test_server_function_call.py new file mode 100644 index 0000000..dfa4f24 --- /dev/null +++ b/tests/python/serve/server/test_server_function_call.py @@ -0,0 +1,208 @@ +""" +Test script for function call in chat completion. To run this script, use the following command: +MLC_SERVE_MODEL_LIB=dist/gorilla-openfunctions-v1-q4f16_1_MLC/gorilla-openfunctions-v1-q4f16_1-cuda.so +MLC_SERVE_MODEL_LIB=${MLC_SERVE_MODEL_LIB} python -m pytest -x tests/python/serve/server/test_server_function_call.py +""" # noqa: E501 + +import json +import os +from typing import Dict, List, Optional, Tuple # noqa: UP035 + +import pytest +import requests + +OPENAI_V1_CHAT_COMPLETION_URL = "http://127.0.0.1:8000/v1/chat/completions" + + +def check_openai_nonstream_response( + response: Dict, # noqa: UP006 + *, + model: str, + object_str: str, + num_choices: int, + finish_reason: List[str], # noqa: UP006 + completion_tokens: Optional[int] = None, +): + print(response) + assert response["model"] == model + assert response["object"] == object_str + + choices = response["choices"] + assert isinstance(choices, list) + assert len(choices) == num_choices + for idx, choice in enumerate(choices): + assert choice["index"] == idx + assert choice["finish_reason"] in finish_reason + + # text: str + message = choice["message"] + assert message["role"] == "assistant" + if choice["finish_reason"] == "tool_calls": + assert message["content"] is None + assert isinstance(message["tool_calls"], list) + else: + assert message["tool_calls"] is None + assert message["content"] is not None + + usage = response["usage"] + assert isinstance(usage, dict) + assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"] + assert usage["prompt_tokens"] > 0 + + if completion_tokens is not None: + assert usage["completion_tokens"] == completion_tokens + + +def check_openai_stream_response( + responses: List[Dict], # noqa: UP006 + *, + model: str, + object_str: str, + num_choices: int, + finish_reason: str, + echo_prompt: Optional[str] = None, + suffix: Optional[str] = None, + stop: Optional[List[str]] = None, # noqa: UP006 + require_substr: Optional[List[str]] = None, # noqa: UP006 +): + assert len(responses) > 0 + + finished = [False for _ in range(num_choices)] + outputs = ["" for _ in range(num_choices)] + for response in responses: + assert response["model"] == model + assert response["object"] == object_str + + choices = response["choices"] + assert isinstance(choices, list) + assert len(choices) == num_choices + for idx, choice in enumerate(choices): + assert choice["index"] == idx + + delta = choice["delta"] + assert delta["role"] == "assistant" + assert isinstance(delta["content"], str) + outputs[idx] += delta["content"] + + if finished[idx]: + assert choice["finish_reason"] == finish_reason + elif choice["finish_reason"] is not None: + assert choice["finish_reason"] == finish_reason + finished[idx] = True + + for output in outputs: + if echo_prompt is not None: + assert output.startswith(echo_prompt) + if suffix is not None: + assert output.endswith(suffix) + if stop is not None: + for stop_str in stop: + assert stop_str not in output + if require_substr is not None: + for substr in require_substr: + assert substr in output + + +tools = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA", + }, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + }, + } +] + + +CHAT_COMPLETION_MESSAGES = [ + # messages #0 + [ + { + "role": "user", + "content": "What is the current weather in Pittsburgh, PA?", + } + ], + # messages #1 + [ + { + "role": "user", + "content": "What is the current weather in Pittsburgh, PA and Tokyo, JP?", + } + ], + # messages #2 + [ + { + "role": "user", + "content": "What is the current weather in Pittsburgh, PA in fahrenheit?", + } + ], +] + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES) +def test_openai_v1_chat_completion_function_call( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, + messages: List[Dict[str, str]], # noqa: UP006 +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + "tools": tools, + } + + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=60) + if not stream: + check_openai_nonstream_response( + response.json(), + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reason=["tool_calls", "error"], + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reason="tool_calls", + ) + + +if __name__ == "__main__": + model_lib = os.environ.get("MLC_SERVE_MODEL_LIB") + if model_lib is None: + raise ValueError( + 'Environment variable "MLC_SERVE_MODEL_LIB" not found. ' + "Please set it to model lib compiled by MLC LLM " + "(e.g., `./dist/gorilla-openfunctions-v1-q4f16_1_MLC/gorilla-openfunctions-v1-q4f16_1-cuda.so`) " # noqa: E501 + "which supports function calls." + ) + MODEL = (os.path.dirname(model_lib), model_lib) + + for msg in CHAT_COMPLETION_MESSAGES: + test_openai_v1_chat_completion_function_call(MODEL, None, stream=False, messages=msg) + test_openai_v1_chat_completion_function_call(MODEL, None, stream=True, messages=msg) diff --git a/tests/python/serve/server/test_server_image.py b/tests/python/serve/server/test_server_image.py new file mode 100644 index 0000000..080e559 --- /dev/null +++ b/tests/python/serve/server/test_server_image.py @@ -0,0 +1,257 @@ +import json +import os +from typing import Dict, List, Optional, Tuple # noqa: UP035 + +import pytest +import regex +import requests + +OPENAI_V1_CHAT_COMPLETION_URL = "http://127.0.0.1:8001/v1/chat/completions" + +JSON_TOKEN_PATTERN = ( + r"((-?(?:0|[1-9]\d*))(\.\d+)?([eE][-+]?\d+)?)|null|true|false|" + r'("((\\["\\\/bfnrt])|(\\u[0-9a-fA-F]{4})|[^"\\\x00-\x1f])*")' +) +JSON_TOKEN_RE = regex.compile(JSON_TOKEN_PATTERN) + + +def is_json_or_json_prefix(s: str) -> bool: + try: + json.loads(s) + return True + except json.JSONDecodeError as e: + # If the JSON decoder reaches the end of s, it is a prefix of a JSON string. + if e.pos == len(s): + return True + # Since json.loads is token-based instead of char-based, there may remain half a token after + # the matching position. + # If the left part is a prefix of a valid JSON token, the output is also valid + regex_match = JSON_TOKEN_RE.fullmatch(s[e.pos :], partial=True) + return regex_match is not None + + +def check_openai_nonstream_response( + response: Dict, # noqa: UP006 + *, + is_chat_completion: bool, + model: str, + object_str: str, + num_choices: int, + finish_reasons: List[str], # noqa: UP006 + completion_tokens: Optional[int] = None, + echo_prompt: Optional[str] = None, + suffix: Optional[str] = None, + stop: Optional[List[str]] = None, # noqa: UP006 + require_substr: Optional[List[str]] = None, # noqa: UP006 + json_mode: bool = False, +): + assert response["model"] == model + assert response["object"] == object_str + + choices = response["choices"] + assert isinstance(choices, list) + assert len(choices) <= num_choices + texts: List[str] = ["" for _ in range(num_choices)] # noqa: UP006 + for choice in choices: + idx = choice["index"] + assert choice["finish_reason"] in finish_reasons + + if not is_chat_completion: + assert isinstance(choice["text"], str) + texts[idx] = choice["text"] + if echo_prompt is not None: + assert texts[idx] + if suffix is not None: + assert texts[idx] + else: + message = choice["message"] + assert message["role"] == "assistant" + assert isinstance(message["content"], str) + texts[idx] = message["content"] + + if stop is not None: + for stop_str in stop: + assert stop_str not in texts[idx] + if require_substr is not None: + for substr in require_substr: + assert substr in texts[idx] + if json_mode: + assert is_json_or_json_prefix(texts[idx]) + + usage = response["usage"] + assert isinstance(usage, dict) + assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"] + assert usage["prompt_tokens"] > 0 + if completion_tokens is not None: + assert usage["completion_tokens"] == completion_tokens + + +def check_openai_stream_response( + responses: List[Dict], # noqa: UP006 + *, + is_chat_completion: bool, + model: str, + object_str: str, + num_choices: int, + finish_reasons: List[str], # noqa: UP006 + completion_tokens: Optional[int] = None, + echo_prompt: Optional[str] = None, + suffix: Optional[str] = None, + stop: Optional[List[str]] = None, # noqa: UP006 + require_substr: Optional[List[str]] = None, # noqa: UP006 + json_mode: bool = False, +): + assert len(responses) > 0 + + finished = [False for _ in range(num_choices)] + outputs = ["" for _ in range(num_choices)] + for response in responses: + assert response["model"] == model + assert response["object"] == object_str + + choices = response["choices"] + assert isinstance(choices, list) + assert len(choices) <= num_choices + for choice in choices: + idx = choice["index"] + + if not is_chat_completion: + assert isinstance(choice["text"], str) + outputs[idx] += choice["text"] + else: + delta = choice["delta"] + assert delta["role"] == "assistant" + assert isinstance(delta["content"], str) + outputs[idx] += delta["content"] + + if finished[idx]: + assert choice["finish_reason"] in finish_reasons + elif choice["finish_reason"] is not None: + assert choice["finish_reason"] in finish_reasons + finished[idx] = True + + if not is_chat_completion: + usage = response["usage"] + assert isinstance(usage, dict) + assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"] + assert usage["prompt_tokens"] > 0 + if completion_tokens is not None: + assert usage["completion_tokens"] <= completion_tokens + + if not is_chat_completion: + if completion_tokens is not None: + assert responses[-1]["usage"]["completion_tokens"] == completion_tokens + + for i, output in enumerate(outputs): + if echo_prompt is not None: + assert output.startswith(echo_prompt) + if suffix is not None: + assert output.endswith(suffix) + if stop is not None: + for stop_str in stop: + assert stop_str not in output + if require_substr is not None: + for substr in require_substr: + assert substr in output + if json_mode: + assert is_json_or_json_prefix(output) + + +CHAT_COMPLETION_MESSAGES = [ + # messages #0 + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "https://llava-vl.github.io/static/images/view.jpg", + }, + {"type": "text", "text": "What does this image represent?"}, + ], + }, + ], + # messages #1 + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "https://llava-vl.github.io/static/images/view.jpg", + }, + {"type": "text", "text": "What does this image represent?"}, + ], + }, + { + "role": "assistant", + "content": "The image represents a serene and peaceful scene of a pier extending over a body of water, such as a lake or a river.er. The pier is made of wood and has a bench on it, providing a place for people to sit and enjoy the view. The pier is situated in a natural environment, surrounded by trees and mountains in the background. This setting creates a tranquil atmosphere, inviting visitors to relax and appreciate the beauty of the landscape.", # noqa: E501 + }, + { + "role": "user", + "content": "What country is the image set in? Give me 10 ranked guesses and reasons why.", # noqa: E501 + }, + ], +] + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES) +def test_openai_v1_chat_completions( + served_model: Tuple[str, str], # noqa: UP006 + launch_server, + stream: bool, + messages: List[Dict[str, str]], # noqa: UP006 +): + # `served_model` and `launch_server` are pytest fixtures + # defined in conftest.py. + + payload = { + "model": served_model[0], + "messages": messages, + "stream": stream, + } + response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=180) + if not stream: + check_openai_nonstream_response( + response.json(), + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion", + num_choices=1, + finish_reasons=["stop"], + ) + else: + responses = [] + for chunk in response.iter_lines(chunk_size=512): + if not chunk or chunk == b"data: [DONE]": + continue + responses.append(json.loads(chunk.decode("utf-8")[6:])) + check_openai_stream_response( + responses, + is_chat_completion=True, + model=served_model[0], + object_str="chat.completion.chunk", + num_choices=1, + finish_reasons=["stop"], + ) + + +if __name__ == "__main__": + model_lib = os.environ.get("MLC_SERVE_MODEL_LIB") + if model_lib is None: + raise ValueError( + 'Environment variable "MLC_SERVE_MODEL_LIB" not found. ' + "Please set it to model lib compiled by MLC LLM " + "(e.g., `dist/Llama-2-7b-chat-hf-q0f16-MLC/Llama-2-7b-chat-hf-q0f16-MLC-cuda.so`)." + ) + + model = os.environ.get("MLC_SERVE_MODEL") + if model is None: + MODEL = (os.path.dirname(model_lib), model_lib) + else: + MODEL = (model, model_lib) + + for msg in CHAT_COMPLETION_MESSAGES: + test_openai_v1_chat_completions(MODEL, None, stream=False, messages=msg) + test_openai_v1_chat_completions(MODEL, None, stream=True, messages=msg) diff --git a/tests/python/serve/test_embedding_engine.py b/tests/python/serve/test_embedding_engine.py new file mode 100644 index 0000000..8599157 --- /dev/null +++ b/tests/python/serve/test_embedding_engine.py @@ -0,0 +1,365 @@ +"""Embedding engine tests in MLC LLM. + +Tests AsyncEmbeddingEngine for both direct (sync) and async embedding inference. +Reuses MLC LLM test infrastructure: markers, require_test_model pattern, +and conventions from test_serve_engine.py. + +Run with real model (requires GPU + compiled embedding model): + MLC_SERVE_EMBEDDING_MODEL_LIB="path/to/model.dylib" \ + pytest -m engine tests/python/serve/test_embedding_engine.py -v + +Environment variables: + MLC_SERVE_EMBEDDING_MODEL_LIB Path to compiled embedding model library (required) + MLC_SERVE_EMBEDDING_MODEL Path to embedding model weight directory + (optional, defaults to dirname of model lib) +""" + +import asyncio +import os + +import numpy as np +import pytest + +# Reuse MLC LLM marker system (registered in tests/python/conftest.py) +pytestmark = [pytest.mark.engine] + +# --------------------------------------------------------------------------- +# Fixtures — follows pattern from serve/server/conftest.py (served_model) +# --------------------------------------------------------------------------- + +EMBEDDING_MODEL_LIB = os.environ.get("MLC_SERVE_EMBEDDING_MODEL_LIB") +EMBEDDING_MODEL_DIR = os.environ.get( + "MLC_SERVE_EMBEDDING_MODEL", + os.path.dirname(EMBEDDING_MODEL_LIB) if EMBEDDING_MODEL_LIB else None, +) + + +def _skip_if_no_model(): + if EMBEDDING_MODEL_LIB is None: + pytest.skip( + 'Environment variable "MLC_SERVE_EMBEDDING_MODEL_LIB" not found. ' + "Set it to a compiled embedding model library " + "(e.g., Qwen3-Embedding-0.6B-q0f32-MLC.dylib)." + ) + if not os.path.isfile(EMBEDDING_MODEL_LIB): + pytest.skip(f"Embedding model library not found at: {EMBEDDING_MODEL_LIB}") + if EMBEDDING_MODEL_DIR is None or not os.path.isdir(EMBEDDING_MODEL_DIR): + pytest.skip(f"Embedding model directory not found at: {EMBEDDING_MODEL_DIR}") + + +@pytest.fixture(scope="module") +def embedding_engine(): + """Module-scoped AsyncEmbeddingEngine — loaded once, shared across tests.""" + _skip_if_no_model() + from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine + + engine = AsyncEmbeddingEngine( + model=EMBEDDING_MODEL_DIR, + model_lib=EMBEDDING_MODEL_LIB, + device="auto", + ) + yield engine + engine.terminate() + + +# --------------------------------------------------------------------------- +# Helpers — reuse cosine_similarity pattern from test_serve_engine.py +# --------------------------------------------------------------------------- + + +def cosine_similarity(a, b): + """Return cosine similarity between two vectors.""" + a, b = np.array(a), np.array(b) + return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))) + + +# =================================================================== +# Engine initialization tests +# =================================================================== + + +def test_engine_model_type(embedding_engine): + """Engine reports a valid model type.""" + assert embedding_engine.model_type in ("encoder", "decoder") + + +def test_engine_pooling_strategy(embedding_engine): + """Engine selects appropriate default pooling strategy.""" + if embedding_engine.model_type == "encoder": + assert embedding_engine.pooling_strategy == "cls" + else: + assert embedding_engine.pooling_strategy == "last" + + +# =================================================================== +# Single-text embedding +# =================================================================== + + +def test_single_text_shape(embedding_engine): + """Single text returns exactly one embedding vector.""" + embeddings, tokens = embedding_engine.embed(["Hello world"]) + assert len(embeddings) == 1 + assert len(embeddings[0]) > 0 + assert tokens > 0 + + +def test_single_text_unit_norm(embedding_engine): + """Embedding output is L2-normalized.""" + embeddings, _ = embedding_engine.embed(["Hello world"]) + norm = float(np.linalg.norm(embeddings[0])) + assert abs(norm - 1.0) < 1e-4, f"Expected unit norm, got {norm}" + + +# =================================================================== +# Batch embedding +# =================================================================== + +BATCH_TEXTS = [ + "Machine learning is fascinating", + "I love pizza", + "Deep learning uses neural networks", +] + + +def test_batch_count(embedding_engine): + """Batch embedding returns one vector per input.""" + embeddings, tokens = embedding_engine.embed(BATCH_TEXTS) + assert len(embeddings) == len(BATCH_TEXTS) + assert tokens > 0 + + +def test_batch_all_normalized(embedding_engine): + """Every vector in a batch is L2-normalized.""" + embeddings, _ = embedding_engine.embed(BATCH_TEXTS) + for i, emb in enumerate(embeddings): + norm = float(np.linalg.norm(emb)) + assert abs(norm - 1.0) < 1e-4, f"Embedding [{i}] norm={norm}" + + +def test_batch_consistent_dimension(embedding_engine): + """All embeddings in a batch have the same dimension.""" + embeddings, _ = embedding_engine.embed(BATCH_TEXTS) + dims = {len(emb) for emb in embeddings} + assert len(dims) == 1, f"Inconsistent dimensions: {dims}" + + +# =================================================================== +# Semantic quality — cosine similarity ranking +# =================================================================== + +SIMILARITY_TEXTS = [ + "What is machine learning?", + "Explain deep learning algorithms", + "I want to order pizza", +] + + +def test_cosine_similarity_ranking(embedding_engine): + """Related texts have higher cosine similarity than unrelated texts.""" + embeddings, _ = embedding_engine.embed(SIMILARITY_TEXTS) + e_ml, e_dl, e_pizza = [np.array(e) for e in embeddings] + sim_related = float(np.dot(e_ml, e_dl)) + sim_unrelated = float(np.dot(e_ml, e_pizza)) + assert sim_related > sim_unrelated, ( + f"Related sim ({sim_related:.4f}) should > unrelated sim ({sim_unrelated:.4f})" + ) + + +# =================================================================== +# Determinism +# =================================================================== + + +def test_deterministic_output(embedding_engine): + """Same input produces identical output across calls.""" + text = ["Deterministic test"] + emb1, _ = embedding_engine.embed(text) + emb2, _ = embedding_engine.embed(text) + cos = cosine_similarity(emb1[0], emb2[0]) + assert cos > 0.9999, f"Expected deterministic output, cosine={cos}" + + +# =================================================================== +# Async embedding +# =================================================================== + + +def test_async_embed(embedding_engine): + """async_embed produces same result as sync embed.""" + text = ["Async test"] + sync_emb, sync_tokens = embedding_engine.embed(text) + + loop = asyncio.new_event_loop() + try: + async_emb, async_tokens = loop.run_until_complete(embedding_engine.async_embed(text)) + finally: + loop.close() + + assert sync_tokens == async_tokens + cos = cosine_similarity(sync_emb[0], async_emb[0]) + assert cos > 0.9999, f"Async vs sync mismatch, cosine={cos}" + + +# =================================================================== +# Edge cases +# =================================================================== + + +def test_empty_string(embedding_engine): + """Empty string should still produce a valid embedding for supported models.""" + embeddings, tokens = embedding_engine.embed([""]) + if embedding_engine.model_type == "encoder": + assert len(embeddings) == 1 + assert len(embeddings[0]) > 0 + assert tokens > 0 + else: + assert len(embeddings) == 1 + assert len(embeddings[0]) > 0 + assert tokens > 0 + + +# =================================================================== +# Long text handling (model-type dependent) +# =================================================================== + + +def test_long_text_decoder_chunked_prefill(embedding_engine): + """[Decoder only] Text >prefill_chunk_size triggers chunked prefill. + ~5000 tokens processed in 3 chunks. Result is unit-norm embedding.""" + if embedding_engine.model_type != "decoder": + pytest.skip("Chunked prefill is decoder-only") + long_text = "word " * 5000 + embeddings, tokens = embedding_engine.embed([long_text]) + assert tokens > 2048, f"Expected >2048 tokens to trigger chunking, got {tokens}" + norm = float(np.linalg.norm(embeddings[0])) + assert abs(norm - 1.0) < 1e-3 + + +def _get_encoder_tokens(embedding_engine, text): + """Replicate encoder preprocessing: tokenize and add [CLS]/[SEP].""" + tokens = list(embedding_engine.tokenizer.encode(text)) + if embedding_engine._cls_token_id is not None and ( + len(tokens) == 0 or tokens[0] != embedding_engine._cls_token_id + ): + tokens = [embedding_engine._cls_token_id, *tokens] + if embedding_engine._sep_token_id is not None and ( + len(tokens) == 0 or tokens[-1] != embedding_engine._sep_token_id + ): + tokens = [*tokens, embedding_engine._sep_token_id] + return tokens + + +def test_long_text_encoder_truncation(embedding_engine): + """[Encoder only] Text exceeding prefill_chunk_size is truncated. + Two texts with the same shared prefix but different suffixes beyond the + limit should produce identical embeddings, since the suffix is truncated + and the retained token prefixes are verified to be identical.""" + if embedding_engine.model_type != "encoder": + pytest.skip("Truncation test is encoder-only") + prefill_chunk = embedding_engine._metadata.get("prefill_chunk_size", 512) + + # Dynamically construct input that exceeds prefill_chunk_size. + unit = "machine learning is great " + suffix_a = " alpha beta gamma " * 200 + suffix_b = " totally different ending " * 200 + unit_tokens = len(list(embedding_engine.tokenizer.encode(unit))) + repeats = max(1, prefill_chunk // max(unit_tokens, 1) + 64) + + # Increase prefix length until both inputs exceed prefill_chunk_size + # and their truncated token prefixes are identical. + while True: + shared_prefix = unit * repeats + full_tokens_a = _get_encoder_tokens(embedding_engine, shared_prefix + suffix_a) + full_tokens_b = _get_encoder_tokens(embedding_engine, shared_prefix + suffix_b) + if ( + len(full_tokens_a) > prefill_chunk + and len(full_tokens_b) > prefill_chunk + and full_tokens_a[:prefill_chunk] == full_tokens_b[:prefill_chunk] + ): + break + repeats += 64 + assert repeats < 200000, "Failed to construct truncation test inputs" + + text_a = shared_prefix + suffix_a + text_b = shared_prefix + suffix_b + + emb_a, tokens_a = embedding_engine.embed([text_a]) + emb_b, tokens_b = embedding_engine.embed([text_b]) + + # Verify truncation happened + assert tokens_a <= prefill_chunk, ( + f"Encoder should truncate to {prefill_chunk}, got {tokens_a} tokens" + ) + assert tokens_b <= prefill_chunk + # Both should be valid unit-norm embeddings + assert abs(float(np.linalg.norm(emb_a[0])) - 1.0) < 1e-3 + assert abs(float(np.linalg.norm(emb_b[0])) - 1.0) < 1e-3 + + # Both truncated to identical token sequences → embeddings must match + cos = cosine_similarity(emb_a[0], emb_b[0]) + assert cos > 0.999, f"Same truncated tokens should match, cosine={cos:.6f}" + + +def test_long_vs_short_semantic_quality(embedding_engine): + """Long text should still capture semantic meaning correctly. + Decoder: chunked prefill preserves full context. + Encoder: truncation keeps most relevant prefix.""" + short_ml = "Machine learning enables systems to learn from data" + long_ml = ( + "Machine learning is a fascinating field of study. " * 200 + + "It enables systems to learn from data." + ) + pizza = "I want to order a pepperoni pizza for dinner" + + embs, _ = embedding_engine.embed([short_ml, long_ml, pizza]) + e_short, e_long, e_pizza = [np.array(e) for e in embs] + + sim_same_topic = float(np.dot(e_short, e_long)) + sim_different = float(np.dot(e_short, e_pizza)) + assert sim_same_topic > sim_different, ( + f"Same topic ({sim_same_topic:.4f}) should > different ({sim_different:.4f})" + ) + + +def test_unicode_text(embedding_engine): + """Unicode input is handled correctly.""" + texts = ["Привет мир", "你好世界", "こんにちは世界"] + embeddings, _ = embedding_engine.embed(texts) + assert len(embeddings) == 3 + for emb in embeddings: + assert abs(float(np.linalg.norm(emb)) - 1.0) < 1e-4 + + +# =================================================================== +# Standalone runner (like test_serve_engine.py) +# =================================================================== + +if __name__ == "__main__": + _skip_if_no_model() + from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine + + engine = AsyncEmbeddingEngine( + model=EMBEDDING_MODEL_DIR, + model_lib=EMBEDDING_MODEL_LIB, + device="auto", + ) + try: + test_engine_model_type(engine) + test_engine_pooling_strategy(engine) + test_single_text_shape(engine) + test_single_text_unit_norm(engine) + test_batch_count(engine) + test_batch_all_normalized(engine) + test_batch_consistent_dimension(engine) + test_cosine_similarity_ranking(engine) + test_deterministic_output(engine) + test_async_embed(engine) + test_empty_string(engine) + test_long_text_decoder_chunked_prefill(engine) + test_long_text_encoder_truncation(engine) + test_long_vs_short_semantic_quality(engine) + test_unicode_text(engine) + print("\nAll embedding engine tests passed!") + finally: + engine.terminate() diff --git a/tests/python/serve/test_event_trace_recorder.py b/tests/python/serve/test_event_trace_recorder.py new file mode 100644 index 0000000..a973867 --- /dev/null +++ b/tests/python/serve/test_event_trace_recorder.py @@ -0,0 +1,48 @@ +import json + +import pytest + +from mlc_llm.serve.event_trace_recorder import EventTraceRecorder + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +def test_event_trace_recorder(): + trace_recorder = EventTraceRecorder() + request_ids = ["x", "y"] + num_decode = 5 + + for request_id in request_ids: + trace_recorder.add_event(request_id, event="start tokenization") + trace_recorder.add_event(request_id, event="finish tokenization") + trace_recorder.add_event(request_id, event="add request") + trace_recorder.add_event(request_id, event="start embed") + trace_recorder.add_event(request_id, event="finish embed") + trace_recorder.add_event(request_id, event="start prefill") + trace_recorder.add_event(request_id, event="finish prefill") + + for _ in range(num_decode): + for request_id in request_ids: + trace_recorder.add_event(request_id, event="start decode") + trace_recorder.add_event(request_id, event="finish decode") + for request_id in request_ids: + trace_recorder.add_event(request_id, event="start detokenization") + trace_recorder.add_event(request_id, event="finish detokenization") + + events = json.loads(trace_recorder.dump_json()) + decode_count = {} + for event in events: + request_id = event["tid"] + if event["name"].startswith("decode"): + if request_id not in decode_count: + decode_count[request_id] = 1 + else: + decode_count[request_id] += 1 + + for _, decode_cnt in decode_count.items(): + assert decode_cnt == num_decode * 2, decode_cnt + + +if __name__ == "__main__": + test_event_trace_recorder() diff --git a/tests/python/serve/test_radix_tree.py b/tests/python/serve/test_radix_tree.py new file mode 100644 index 0000000..67f0cbf --- /dev/null +++ b/tests/python/serve/test_radix_tree.py @@ -0,0 +1,138 @@ +import pytest + +from mlc_llm.serve import PagedRadixTree + +# category "runtime_module" +pytestmark = [pytest.mark.unittest] + + +def test_add(): + prt = PagedRadixTree() + prt.add(0) + assert list(prt.get(0)) == [] + prt.add(1) + assert list(prt.get(1)) == [] + + +def test_remove(): + prt = PagedRadixTree() + capacity = prt.free_capacity() + prt.add(0) + prt.remove(0) + prt.add(0) + prt.extend(0, [1 for _ in range(200)]) + prt.remove(0) + assert prt.free_capacity() == capacity + + prt.add(1) + prt.extend(1, [1 for _ in range(200)]) + capacity = prt.free_capacity() + prt.add(2) + prt.extend(2, [1 for _ in range(100)] + [2 for _ in range(100)]) + prt.remove(2) + assert prt.free_capacity() == capacity + + prt.add(3) + prt.extend(3, [1 for _ in range(200)]) + prt.remove(3) + assert prt.free_capacity() == capacity + + prt.add(4) + prt.add(5) + prt.add(6) + assert prt.free_capacity() == capacity + prt.remove(4) + assert prt.free_capacity() == capacity + prt.remove(5) + assert prt.free_capacity() == capacity + prt.remove(6) + assert prt.free_capacity() == capacity + + +def test_extend(): + prt = PagedRadixTree() + L = prt.free_capacity() // 64 + H = L // 2 + Q = L // 4 + seq_id = 0 + for start_pos in [0, H, L, L + H]: + for length in [Q, L - H, L, 2 * L - H, 2 * L]: + prt.add(seq_id) + if start_pos: + tokens_1 = [seq_id for _ in range(start_pos)] + prt.extend(seq_id, tokens_1) + assert list(prt.get(seq_id)) == tokens_1 + else: + tokens_1 = [] + tokens_2 = [seq_id for _ in range(length)] + prt.extend(seq_id, tokens_2) + assert list(prt.get(seq_id)) == tokens_1 + tokens_2 + seq_id += 1 + + +def test_fork(): + prt = PagedRadixTree() + L = prt.free_capacity() // 64 + H = L // 2 + Q = L // 4 + seq_id = 0 + length_list = [Q, H, L, L + Q, L + H, L * 2] + for p_idx in range(1, len(length_list)): + for c_idx in range(0, p_idx + 1): + prt.add(seq_id) + tokens = [seq_id for _ in range(length_list[p_idx])] + prt.extend(seq_id, tokens) + prt.fork(seq_id + 1, seq_id, length_list[c_idx]) + assert list(prt.get(seq_id + 1)) == tokens[: length_list[c_idx]] + seq_id += 2 + + +def test_fork_2(): + prt = PagedRadixTree() + prt.add(0) + prt.extend(0, [0, 1, 2, 3]) + prt.fork(1, 0, 3) + prt.extend(1, [4]) + prt.fork(2, 0, 3) + prt.extend(2, [5]) + assert prt.match([0, 1, 2, 4]) == (4, (1,)) + assert prt.match([0, 1, 2, 5]) == (4, (2,)) + + +def test_rollback(): + prt = PagedRadixTree() + L = prt.free_capacity() // 64 + H = L // 2 + Q = L // 4 + seq_id = 0 + for start_pos in [H, L, L + H, 2 * L, 3 * L + H]: + for length in [Q, H, L + Q, 2 * L, 2 * L + Q]: + if length > start_pos: + continue + prt.add(seq_id) + tokens = [seq_id for _ in range(start_pos)] + prt.extend(seq_id, tokens) + prt.rollback(seq_id, length) + assert list(prt.get(seq_id)) == tokens[:-length] + seq_id += 1 + + for start_pos in [H, L, L + H, 2 * L, 3 * L + H]: + for length in [Q, H, L + Q, 2 * L, 2 * L + Q]: + if length > start_pos: + continue + prt.add(seq_id) + tokens = [seq_id for _ in range(start_pos)] + prt.extend(seq_id, tokens) + prt.fork(seq_id + 1, seq_id, start_pos) + prt.rollback(seq_id + 1, length) + assert list(prt.get(seq_id + 1)) == tokens[:-length] + seq_id += 2 + + +if __name__ == "__main__": + test_add() + test_remove() + test_extend() + test_fork() + test_fork_2() + test_rollback() diff --git a/tests/python/serve/test_serve_async_engine.py b/tests/python/serve/test_serve_async_engine.py new file mode 100644 index 0000000..617f812 --- /dev/null +++ b/tests/python/serve/test_serve_async_engine.py @@ -0,0 +1,285 @@ +import asyncio +from typing import List # noqa: UP035 + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import AsyncMLCEngine, EngineConfig +from mlc_llm.testing import require_test_model + +prompts = [ + "What is the meaning of life?", + "Introduce the history of Pittsburgh to me. Please elaborate in detail.", + "Write a three-day Seattle travel plan. Please elaborate in detail.", + "What is Alaska famous of? Please elaborate in detail.", + "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501 + "What are the necessary components to assemble a desktop computer? Please elaborate in detail.", + "Why is Vitamin D important to human beings? Please elaborate in detail.", + "Where is milk tea originated from? Please elaborate in detail.", + "Where is the southernmost place in United States? Please elaborate in detail.", + "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501 +] + + +@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") +async def test_engine_generate(model: str): + # Create engine + async_engine = AsyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + + num_requests = 10 + max_tokens = 256 + generation_cfg = GenerationConfig(max_tokens=max_tokens, n=7) + + output_texts: List[List[str]] = [ # noqa: UP006 + ["" for _ in range(generation_cfg.n)] for _ in range(num_requests) + ] + + async def generate_task( + async_engine: AsyncMLCEngine, + prompt: str, + generation_cfg: GenerationConfig, + request_id: str, + ): + print(f"generate task for request {request_id}") + rid = int(request_id) + async for delta_outputs in async_engine._generate( + prompt, generation_cfg, request_id=request_id + ): + if len(delta_outputs) == generation_cfg.n: + for i, delta_output in enumerate(delta_outputs): + output_texts[rid][i] += delta_output.delta_text + else: + assert len(delta_outputs) == 1 + assert len(delta_outputs[0].request_final_usage_json_str) != 0 + + tasks = [ + asyncio.create_task( + generate_task(async_engine, prompts[i], generation_cfg, request_id=str(i)) + ) + for i in range(num_requests) + ] + + await asyncio.gather(*tasks) + + # Print output. + print("All finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + async_engine.terminate() + del async_engine + + +@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") +async def test_chat_completion(model: str): + # Create engine + async_engine = AsyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + + num_requests = 2 + max_tokens = 32 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + async def generate_task(prompt: str, request_id: str): + print(f"generate chat completion task for request {request_id}") + rid = int(request_id) + async for response in await async_engine.chat.completions.create( # noqa: F821 + messages=[{"role": "user", "content": prompt}], + model=model, + max_tokens=max_tokens, + n=n, + request_id=request_id, + stream=True, + ): + for choice in response.choices: + assert choice.delta.role == "assistant" + assert isinstance(choice.delta.content, str) + output_texts[rid][choice.index] += choice.delta.content + + tasks = [ + asyncio.create_task(generate_task(prompts[i], request_id=str(i))) + for i in range(num_requests) + ] + + await asyncio.gather(*tasks) + + # Print output. + print("Chat completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + async_engine.terminate() + del async_engine + + +@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") +async def test_chat_completion_non_stream(model: str): + # Create engine + async_engine = AsyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + + num_requests = 2 + max_tokens = 32 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + async def generate_task(prompt: str, request_id: str): + print(f"generate chat completion task for request {request_id}") + rid = int(request_id) + response = await async_engine.chat.completions.create( # noqa: F821 + messages=[{"role": "user", "content": prompt}], + model=model, + max_tokens=max_tokens, + n=n, + request_id=request_id, + ) + for choice in response.choices: + assert choice.message.role == "assistant" + assert isinstance(choice.message.content, str) + output_texts[rid][choice.index] += choice.message.content + + tasks = [ + asyncio.create_task(generate_task(prompts[i], request_id=str(i))) + for i in range(num_requests) + ] + + await asyncio.gather(*tasks) + + # Print output. + print("Chat completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + async_engine.terminate() + del async_engine + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +async def test_completion(model: str): + # Create engine + async_engine = AsyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + + num_requests = 2 + max_tokens = 128 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + async def generate_task(prompt: str, request_id: str): + print(f"generate completion task for request {request_id}") + rid = int(request_id) + async for response in await async_engine.completions.create( # noqa: F821 + prompt=prompt, + model=model, + max_tokens=max_tokens, + n=n, + request_id=request_id, + stream=True, + extra_body={"debug_config": {"ignore_eos": True}}, + ): + for choice in response.choices: + output_texts[rid][choice.index] += choice.text + + tasks = [ + asyncio.create_task(generate_task(prompts[i], request_id=str(i))) + for i in range(num_requests) + ] + + await asyncio.gather(*tasks) + + # Print output. + print("Completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + async_engine.terminate() + del async_engine + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +async def test_completion_non_stream(model: str): + # Create engine + async_engine = AsyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + + num_requests = 2 + max_tokens = 128 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + async def generate_task(prompt: str, request_id: str): + print(f"generate completion task for request {request_id}") + rid = int(request_id) + response = await async_engine.completions.create( # noqa: F821 + prompt=prompt, + model=model, + max_tokens=max_tokens, + n=n, + request_id=request_id, + extra_body={"debug_config": {"ignore_eos": True}}, + ) + for choice in response.choices: + output_texts[rid][choice.index] += choice.text + + tasks = [ + asyncio.create_task(generate_task(prompts[i], request_id=str(i))) + for i in range(num_requests) + ] + + await asyncio.gather(*tasks) + + # Print output. + print("Completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + async_engine.terminate() + del async_engine + + +if __name__ == "__main__": + asyncio.run(test_engine_generate()) + asyncio.run(test_chat_completion()) + asyncio.run(test_chat_completion_non_stream()) + asyncio.run(test_completion()) + asyncio.run(test_completion_non_stream()) diff --git a/tests/python/serve/test_serve_async_engine_spec.py b/tests/python/serve/test_serve_async_engine_spec.py new file mode 100644 index 0000000..2aaaaac --- /dev/null +++ b/tests/python/serve/test_serve_async_engine_spec.py @@ -0,0 +1,84 @@ +import asyncio +from typing import List # noqa: UP035 + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import AsyncMLCEngine, EngineConfig +from mlc_llm.testing import require_test_model + +prompts = [ + "What is the meaning of life?", + "Introduce the history of Pittsburgh to me. Please elaborate in detail.", + "Write a three-day Seattle travel plan. Please elaborate in detail.", + "What is Alaska famous of? Please elaborate in detail.", + "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501 + "What are the necessary components to assemble a desktop computer? Please elaborate in detail.", + "Why is Vitamin D important to human beings? Please elaborate in detail.", + "Where is milk tea originated from? Please elaborate in detail.", + "Where is the southernmost place in United States? Please elaborate in detail.", + "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501 +] + + +@require_test_model( + "Llama-2-7b-chat-hf-q0f16-MLC", + "Llama-2-7b-chat-hf-q4f16_1-MLC", +) +async def test_engine_generate(model: str, small_model: str): + # Create engine + async_engine = AsyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + additional_models=[small_model], + speculative_mode="small_draft", + ), + ) + + num_requests = 10 + max_tokens = 256 + generation_cfg = GenerationConfig(max_tokens=max_tokens) + + output_texts: List[List[str]] = [ # noqa: UP006 + ["" for _ in range(generation_cfg.n)] for _ in range(num_requests) + ] + + async def generate_task( + async_engine: AsyncMLCEngine, + prompt: str, + generation_cfg: GenerationConfig, + request_id: str, + ): + print(f"generate task for request {request_id}") + rid = int(request_id) + async for delta_outputs in async_engine._generate( + prompt, generation_cfg, request_id=request_id + ): + assert len(delta_outputs) == generation_cfg.n + for i, delta_output in enumerate(delta_outputs): + output_texts[rid][i] += delta_output.delta_text + + tasks = [ + asyncio.create_task( + generate_task(async_engine, prompts[i], generation_cfg, request_id=str(i)) + ) + for i in range(num_requests) + ] + + await asyncio.gather(*tasks) + + # Print output. + print("All finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + async_engine.terminate() + del async_engine + + +if __name__ == "__main__": + asyncio.run(test_engine_generate()) diff --git a/tests/python/serve/test_serve_engine.py b/tests/python/serve/test_serve_engine.py new file mode 100644 index 0000000..d15c410 --- /dev/null +++ b/tests/python/serve/test_serve_engine.py @@ -0,0 +1,241 @@ +from typing import List # noqa: UP035 + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import EngineConfig, MLCEngine +from mlc_llm.testing import require_test_model + +prompts = [ + "What is the meaning of life?", + "Introduce the history of Pittsburgh to me. Please elaborate in detail.", + "Write a three-day Seattle travel plan. Please elaborate in detail.", + "What is Alaska famous of? Please elaborate in detail.", + "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501 + "What are the necessary components to assemble a desktop computer? Please elaborate in detail.", + "Why is Vitamin D important to human beings? Please elaborate in detail.", + "Where is milk tea originated from? Please elaborate in detail.", + "Where is the southernmost place in United States? Please elaborate in detail.", + "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501 +] + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_generate(model: str): + # Create engine + engine = MLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + ), + ) + + num_requests = 10 + max_tokens = 256 + generation_cfg = GenerationConfig(max_tokens=max_tokens, n=7) + + output_texts: List[List[str]] = [ # noqa: UP006 + ["" for _ in range(generation_cfg.n)] for _ in range(num_requests) + ] + for rid in range(num_requests): + print(f"generating for request {rid}") + for delta_outputs in engine._generate(prompts[rid], generation_cfg, request_id=str(rid)): + assert len(delta_outputs) == generation_cfg.n + for i, delta_output in enumerate(delta_outputs): + output_texts[rid][i] += delta_output.delta_text + + # Print output. + print("All finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + engine.terminate() + del engine + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_chat_completion(model: str): + # Create engine + engine = MLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + ), + ) + + num_requests = 2 + max_tokens = 64 + n = 2 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + for rid in range(num_requests): + print(f"chat completion for request {rid}") + for response in engine.chat.completions.create( + messages=[{"role": "user", "content": prompts[rid]}], + model=model, + max_tokens=max_tokens, + n=n, + request_id=str(rid), + stream=True, + ): + for choice in response.choices: + assert choice.delta.role == "assistant" + assert isinstance(choice.delta.content, str) + output_texts[rid][choice.index] += choice.delta.content + + # Print output. + print("Chat completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + engine.terminate() + del engine + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_chat_completion_non_stream(model: str): + # Create engine + engine = MLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + ), + ) + + num_requests = 2 + max_tokens = 64 + n = 2 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + for rid in range(num_requests): + print(f"chat completion for request {rid}") + response = engine.chat.completions.create( + messages=[{"role": "user", "content": prompts[rid]}], + model=model, + max_tokens=max_tokens, + n=n, + request_id=str(rid), + ) + for choice in response.choices: + assert choice.message.role == "assistant" + assert isinstance(choice.message.content, str) + output_texts[rid][choice.index] += choice.message.content + + # Print output. + print("Chat completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + engine.terminate() + del engine + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_completion(model: str): + # Create engine + engine = MLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + ), + ) + + num_requests = 2 + max_tokens = 128 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + for rid in range(num_requests): + print(f"completion for request {rid}") + for response in engine.completions.create( + prompt=prompts[rid], + model=model, + max_tokens=max_tokens, + n=n, + request_id=str(rid), + stream=True, + extra_body={"debug_config": {"ignore_eos": True}}, + ): + for choice in response.choices: + output_texts[rid][choice.index] += choice.text + + # Print output. + print("Completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + engine.terminate() + del engine + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_completion_non_stream(model: str): + # Create engine + engine = MLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + ), + ) + + num_requests = 2 + max_tokens = 128 + n = 1 + output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)] # noqa: UP006 + + for rid in range(num_requests): + print(f"completion for request {rid}") + response = engine.completions.create( + prompt=prompts[rid], + model=model, + max_tokens=max_tokens, + n=n, + request_id=str(rid), + extra_body={"debug_config": {"ignore_eos": True}}, + ) + for choice in response.choices: + output_texts[rid][choice.index] += choice.text + + # Print output. + print("Completion all finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + engine.terminate() + del engine + + +if __name__ == "__main__": + test_engine_generate() + test_chat_completion() + test_chat_completion_non_stream() + test_completion() + test_completion_non_stream() diff --git a/tests/python/serve/test_serve_engine_grammar.py b/tests/python/serve/test_serve_engine_grammar.py new file mode 100644 index 0000000..16e84a0 --- /dev/null +++ b/tests/python/serve/test_serve_engine_grammar.py @@ -0,0 +1,356 @@ +import asyncio +import json +import random +from typing import Dict, List, Literal # noqa: UP035 + +from pydantic import BaseModel + +from mlc_llm.protocol.debug_protocol import DebugConfig +from mlc_llm.protocol.openai_api_protocol import ChatCompletionResponse +from mlc_llm.serve import AsyncMLCEngine, MLCEngine +from mlc_llm.testing import require_test_model + +LLAMA_2_MODEL = "Llama-2-7b-chat-hf-q4f16_1-MLC" +LLAMA_3_MODEL = "Meta-Llama-3-8B-Instruct-q4f16_1-MLC" + + +@require_test_model(LLAMA_3_MODEL) +def test_batch_generation_with_grammar(model: str): + # Engine + engine = MLCEngine(model=model, mode="server") + + # Inputs + system_prompt = "You are a helpful assistant. Always respond only with json." + prompts_list = [ + "Generate a JSON string containing 20 objects:", + "Generate a JSON containing a non-empty list:", + "Generate a JSON with 5 elements:", + "Generate a JSON with a number list, counting from 1 to 20:", + ] + + repeat = 3 + top_p = 0.9 + temperature = 0.6 + max_tokens = 4096 + + # non-json output + responses_text: List[ChatCompletionResponse] = [] # noqa: UP006 + for _ in range(repeat): + for p in prompts_list: + print(f"Start generation task for request {len(responses_text)}") + responses_text.append( + engine.chat.completions.create( + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": p}, + ], + response_format={"type": "text"}, + top_p=top_p, + temperature=temperature, + max_tokens=max_tokens, + seed=random.randint(0, 1 << 30), + extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")}, + ) + ) + + print("Text output") + for req_id, response in enumerate(responses_text): + prompt = prompts_list[req_id % len(prompts_list)] + output = response.choices[0].message.content + print(f"Prompt {req_id}: {prompt}") + print(f"Output {req_id}: {output}\n") + + # json output + responses_json: List[ChatCompletionResponse] = [] # noqa: UP006 + for _ in range(repeat): + for p in prompts_list: + print(f"Start generation task for request {len(responses_json)}") + responses_json.append( + engine.chat.completions.create( + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": p}, + ], + response_format={"type": "json_object"}, + top_p=top_p, + temperature=temperature, + seed=random.randint(0, 1 << 30), + ) + ) + + print("JSON output") + for req_id, response in enumerate(responses_json): + prompt = prompts_list[req_id % len(prompts_list)] + output = str(response.choices[0].message.content) + print(f"Prompt {req_id}: {prompt}") + print(f"Output {req_id}: {output}\n") + json.loads(output) + + print("Engine metrics:", engine.metrics()) + + engine.terminate() + + +@require_test_model(LLAMA_3_MODEL) +def test_batch_generation_with_schema(model: str): + # Create engine + engine = MLCEngine(model=model, mode="server") + + class Product(BaseModel): + product_id: int + is_available: bool + price: float + is_featured: Literal[True] + category: Literal["Electronics", "Clothing", "Food"] + tags: List[str] # noqa: UP006 + stock: Dict[str, int] # noqa: UP006 + + schema_str = json.dumps(Product.model_json_schema()) + + system_prompt = ( + "You are a helpful assistant. Always respond only with JSON based on the " + f"following JSON schema: {schema_str}." + ) + prompt = "Generate a JSON that describes the product according to the given JSON schema." + + repeat = 8 + top_p = 0.9 + temperature = 0.6 + max_tokens = 4096 + + # non-json output + responses_text: List[ChatCompletionResponse] = [] # noqa: UP006 + for i in range(repeat): + print(f"Start generation task for request {i}") + responses_text.append( + engine.chat.completions.create( + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}, + ], + response_format={"type": "text"}, + top_p=top_p, + temperature=temperature, + max_tokens=max_tokens, + seed=random.randint(0, 1 << 30), + extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")}, + ) + ) + + print("Text output") + for req_id, response in enumerate(responses_text): + output = response.choices[0].message.content + print(f"Prompt {req_id}: {prompt}") + print(f"Output {req_id}: {output}\n") + + # json output without schema + responses_json: List[ChatCompletionResponse] = [] # noqa: UP006 + for i in range(repeat): + print(f"Start generation task for request {i}") + responses_json.append( + engine.chat.completions.create( + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}, + ], + response_format={"type": "json_object"}, + top_p=top_p, + temperature=temperature, + max_tokens=max_tokens, + seed=random.randint(0, 1 << 30), + extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")}, + ) + ) + + print("JSON output") + for req_id, response in enumerate(responses_json): + output = response.choices[0].message.content + print(f"Prompt {req_id}: {prompt}") + print(f"Output {req_id}: {output}\n") + + # json output with schema + responses_schema: List[ChatCompletionResponse] = [] # noqa: UP006 + for i in range(repeat): + print(f"Start generation task for request {i}") + responses_schema.append( + engine.chat.completions.create( + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}, + ], + response_format={"type": "json_object", "schema": schema_str}, + top_p=top_p, + temperature=temperature, + max_tokens=max_tokens, + seed=random.randint(0, 1 << 30), + extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")}, + ) + ) + + print("JSON Schema output") + for req_id, response in enumerate(responses_schema): + output = response.choices[0].message.content + print(f"Prompt {req_id}: {prompt}") + print(f"Output {req_id}: {output}\n") + + print("Engine metrics:", engine.metrics()) + + engine.terminate() + + +@require_test_model(LLAMA_3_MODEL) +def test_batch_generation_jump_forward(model: str, jump_forward: bool = True, repeat: int = 1): + # Create engine + engine = MLCEngine(model=model, mode="server") + + class Product(BaseModel): + product_id: int + is_available: bool + price: float + is_featured: Literal[True] + category: Literal["Electronics", "Clothing", "Food"] + tags: List[str] # noqa: UP006 + stock: Dict[str, int] # noqa: UP006 + + schema_str = json.dumps(Product.model_json_schema()) + + system_prompt = ( + "You are a helpful assistant. Always respond only with JSON based on the " + f"following JSON schema: {schema_str}." + ) + prompt = "Generate a JSON that describes the product according to the given JSON schema." + + top_p = 0.9 + temperature = 0.6 + max_tokens = 4096 + grammar_execution_mode = "jump_forward" if jump_forward else "constraint" + + # json output with schema + responses: List[ChatCompletionResponse] = [] # noqa: UP006 + for i in range(repeat): + print(f"Start generation task for request {i}") + responses.append( + engine.chat.completions.create( + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}, + ], + response_format={"type": "json_object", "schema": schema_str}, + top_p=top_p, + temperature=temperature, + max_tokens=max_tokens, + seed=random.randint(0, 1 << 30), + extra_body={ + "debug_config": DebugConfig(grammar_execution_mode=grammar_execution_mode) + }, + ) + ) + + print(f"Jump forward: {jump_forward}, Repeat: {repeat}") + for req_id, response in enumerate(responses): + output = response.choices[0].message.content + print(f"Prompt {req_id}: {prompt}") + print(f"Output {req_id}: {output}\n") + + print("Engine metrics:", engine.metrics()) + + engine.terminate() + + +@require_test_model(LLAMA_3_MODEL) +async def run_async_engine( + model: str, + mode: Literal["text", "json", "schema"] = "schema", + jump_forward: bool = True, + num_requests: int = 8, +): + # Create engine + async_engine = AsyncMLCEngine(model=model, mode="server") + + class Product(BaseModel): + product_id: int + is_available: bool + price: float + is_featured: Literal[True] + category: Literal["Electronics", "Clothing", "Food"] + tags: List[str] # noqa: UP006 + stock: Dict[str, int] # noqa: UP006 + + schema_str = json.dumps(Product.model_json_schema()) + + if mode == "text": + response_format = {"type": "text"} + elif mode == "json": + response_format = {"type": "json_object"} + elif mode == "schema": + response_format = {"type": "json_object", "schema": schema_str} + + system_prompt = ( + "You are a helpful assistant. Always respond only with JSON based on the " + f"following JSON schema: {schema_str}." + ) + prompt = "Generate a JSON that describes the product according to the given JSON schema." + + top_p = 0.9 + temperature = 0.6 + max_tokens = 4096 + grammar_execution_mode = "jump_forward" if jump_forward else "constraint" + + responses = ["" for _ in range(num_requests)] + + async def generate_task(prompt: str, request_id: str): + print(f"Start generation task for request {request_id}") + rid = int(request_id) + async for response in await async_engine.chat.completions.create( # noqa: F821 + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}, + ], + response_format=response_format, + top_p=top_p, + temperature=temperature, + max_tokens=max_tokens, + seed=random.randint(0, 1 << 30), + stream=True, + extra_body={"debug_config": DebugConfig(grammar_execution_mode=grammar_execution_mode)}, + ): + assert len(response.choices) == 1 + choice = response.choices[0] + assert choice.delta.role == "assistant" + assert isinstance(choice.delta.content, str) + responses[rid] += choice.delta.content + + tasks = [ + asyncio.create_task(generate_task(prompt, request_id=str(i))) for i in range(num_requests) + ] + + await asyncio.gather(*tasks) + + print(f"Mode: {mode}, Jump forward: {jump_forward}, Num requests: {num_requests}") + for req_id, output in enumerate(responses): + print(f"Prompt {req_id}: {prompt}") + print(f"Output {req_id}: {output}\n") + + print("Engine metrics:", await async_engine.metrics()) + + async_engine.terminate() + del async_engine + + +def test_async_engine( + mode: Literal["text", "json", "schema"] = "schema", + jump_forward: bool = True, + num_requests: int = 8, +): + asyncio.run(run_async_engine(mode, jump_forward, num_requests)) + + +if __name__ == "__main__": + test_batch_generation_with_grammar() + test_batch_generation_with_schema() + test_batch_generation_jump_forward(False) + test_batch_generation_jump_forward(True) + test_async_engine("schema", False, 1) + test_async_engine("schema", True, 1) + test_async_engine("schema", False, 8) + test_async_engine("schema", True, 8) diff --git a/tests/python/serve/test_serve_engine_image.py b/tests/python/serve/test_serve_engine_image.py new file mode 100644 index 0000000..621dabc --- /dev/null +++ b/tests/python/serve/test_serve_engine_image.py @@ -0,0 +1,56 @@ +import json +from pathlib import Path + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import data +from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine + + +def get_test_image(config) -> data.ImageData: + return data.ImageData.from_url("https://llava-vl.github.io/static/images/view.jpg", config) + + +def test_engine_generate(): + # Create engine + model = "dist/llava-1.5-7b-hf-q4f16_1-MLC/params" + model_lib = "dist/llava-1.5-7b-hf-q4f16_1-MLC/llava-1.5-7b-hf-q4f16_1-MLC.so" + engine = SyncMLCEngine( + model=model, + model_lib=model_lib, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + max_tokens = 256 + + with open(Path(model) / "mlc-chat-config.json", encoding="utf-8") as file: + model_config = json.load(file) + + prompts = [ + [ + data.TextData("USER: "), + get_test_image(model_config), + data.TextData("\nWhat does this image represent? ASSISTANT:"), + ], + [ + data.TextData("USER: "), + get_test_image(model_config), + data.TextData("\nIs there a dog in this image? ASSISTANT:"), + ], + [data.TextData("USER: What is the meaning of life? ASSISTANT:")], + ] + + output_texts, _ = engine.generate( + prompts, GenerationConfig(max_tokens=max_tokens, stop_token_ids=[2]) + ) + + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + +if __name__ == "__main__": + test_engine_generate() diff --git a/tests/python/serve/test_serve_engine_mock.py b/tests/python/serve/test_serve_engine_mock.py new file mode 100644 index 0000000..4ef7842 --- /dev/null +++ b/tests/python/serve/test_serve_engine_mock.py @@ -0,0 +1,39 @@ +"""Mock testing engine I/O conventions + +Mock test only can help checking the overall input +output processing options are passed correctly +""" + +import pytest +import tvm + +from mlc_llm.serve import MLCEngine +from mlc_llm.testing import require_test_model + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +# NOTE: we only need tokenizers in folder +# launch time of mock test is fast so we can put it in unittest +@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC") +def test_completion_api(model: str): + engine = MLCEngine(model, tvm.cpu(), model_lib="mock://echo") + param_dict = { + "top_p": 0.6, + "temperature": 0.9, + "frequency_penalty": 0.1, + "presence_penalty": 0.1, + "n": 2, + } + response = engine.chat.completions.create( + messages=[{"role": "user", "content": "hello"}], + **param_dict, + ) + # echo mock will echo back the generation config + for k, v in param_dict.items(): + assert response.usage.extra[k] == v + + +if __name__ == "__main__": + test_completion_api() diff --git a/tests/python/serve/test_serve_engine_prefix_cache.py b/tests/python/serve/test_serve_engine_prefix_cache.py new file mode 100644 index 0000000..6d8ac98 --- /dev/null +++ b/tests/python/serve/test_serve_engine_prefix_cache.py @@ -0,0 +1,141 @@ +from mlc_llm.protocol.debug_protocol import DebugConfig +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine +from mlc_llm.testing import require_test_model + +prompts = [ + "The meaning of life is", + "According to the history of Pittsburgh,", + "I have a three-day Seattle travel plan. On the first day,", + "Undoubtedly, Alaska is one of the most beautiful places on Earth,", + "Explain difference between Lambda calculus and Turing machine is", + "To assemble a desktop computer, we need the necessary components of", + "Vitamin D is important to human beings, because", + "Refer to history, the milk tea is originated from", + "In the southernmost place in United States,", + "AlphaGo has the capabilities of", +] + + +def test_engine_system_prompt(engine): + system_prompt = "This is a system prompt" + system_prompt_tokens = len(engine.tokenizer.encode(system_prompt)) + max_tokens = 8 + _, _ = engine.generate( + system_prompt, + GenerationConfig( + temperature=0, + max_tokens=max_tokens, + debug_config=DebugConfig(pinned_system_prompt=True), + ), + ) + metrics = engine.metrics() + assert metrics["prefill_tokens_sum"] == system_prompt_tokens + sum_prefill_tokens = system_prompt_tokens + + input_token_lens = [len(engine.tokenizer.encode(prompt)) for prompt in prompts] + + generation_config = GenerationConfig(temperature=0, max_tokens=max_tokens) + _, _ = engine.generate(prompts, generation_config) + metrics = engine.metrics() + assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + sum(input_token_lens) + sum_prefill_tokens = metrics["prefill_tokens_sum"] + + _, _ = engine.generate(system_prompt + " and why ?", generation_config) + metrics = engine.metrics() + # system prompt is reused entirely + assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + 3 + sum_prefill_tokens = metrics["prefill_tokens_sum"] + + _, _ = engine.generate(prompts[:4], generation_config) + metrics = engine.metrics() + # first 4 prompts are removed and need to prefill again + assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + sum(input_token_lens[:4]) + + +def test_engine_multi_round(engine): + num_requests = 10 + max_tokens = 8 + generation_config = GenerationConfig(temperature=0, max_tokens=max_tokens) + input_token_lens = [len(engine.tokenizer.encode(prompt)) for prompt in prompts[:num_requests]] + + output_texts, _ = engine.generate(prompts[:num_requests], generation_config) + metrics = engine.metrics() + assert metrics["prefill_tokens_sum"] == sum(input_token_lens) + sum_prefill_tokens = metrics["prefill_tokens_sum"] + concat_prompt = [] + for i, output in enumerate(output_texts): + concat_prompt.append(prompts[i] + " " + output[0] + " ?") + output_texts, _ = engine.generate(concat_prompt[:num_requests], generation_config) + metrics = engine.metrics() + assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + 2 * num_requests + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_basic_engine_system_prompt(model: str): + # Create engine + engine = SyncMLCEngine( + model=model, + mode="local", + engine_config=EngineConfig( + max_total_sequence_length=4096, + prefix_cache_max_num_recycling_seqs=5, + ), + ) + test_engine_system_prompt(engine) + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_basic_engine_multi_round(model: str): + # Create engine + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + test_engine_multi_round(engine) + + +@require_test_model( + "Llama-2-7b-chat-hf-q0f16-MLC", + "Llama-2-7b-chat-hf-q4f16_1-MLC", +) +def test_engine_spec_multi_round(model: str, small_model: str): + # Create engine + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[small_model], + speculative_mode="small_draft", + ), + ) + + test_engine_multi_round(engine) + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_eagle_multi_round(model: str): + # Create engine + small_model = "dist/Eagle-llama2-7b-chat-q0f16-MLC" + small_model_lib = "dist/Eagle-llama2-7b-chat-q0f16-MLC/Eagle-llama2-7b-chat-q0f16-MLC-cuda.so" + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[(small_model, small_model_lib)], + speculative_mode="eagle", + max_num_sequence=80, + ), + ) + + test_engine_multi_round(engine) + + +if __name__ == "__main__": + test_basic_engine_system_prompt() + test_basic_engine_multi_round() + test_engine_spec_multi_round() + test_engine_eagle_multi_round() diff --git a/tests/python/serve/test_serve_engine_rnn.py b/tests/python/serve/test_serve_engine_rnn.py new file mode 100644 index 0000000..67b120c --- /dev/null +++ b/tests/python/serve/test_serve_engine_rnn.py @@ -0,0 +1,60 @@ +from typing import List # noqa: UP035 + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import EngineConfig, MLCEngine + +prompts = [ + "What is the meaning of life?", + "Introduce the history of Pittsburgh to me. Please elaborate in detail.", + "Write a three-day Seattle travel plan. Please elaborate in detail.", + "What is Alaska famous of? Please elaborate in detail.", + "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501 + "What are the necessary components to assemble a desktop computer? Please elaborate in detail.", + "Why is Vitamin D important to human beings? Please elaborate in detail.", + "Where is milk tea originated from? Please elaborate in detail.", + "Where is the southernmost place in United States? Please elaborate in detail.", + "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501 +] + + +def test_engine_generate() -> None: + engine = MLCEngine( + model="dist/rwkv-6-world-1b6-q0f16-MLC", + model_lib="dist/rwkv-6-world-1b6-q0f16-MLC/rwkv-6-world-1b6-q0f16-MLC-cuda.so", + mode="server", + engine_config=EngineConfig( + max_num_sequence=8, + max_history_size=1, + ), + ) + + num_requests = 10 + max_tokens = 256 + generation_cfg = GenerationConfig(max_tokens=max_tokens, n=7) + + output_texts: List[List[str]] = [ # noqa: UP006 + ["" for _ in range(generation_cfg.n)] for _ in range(num_requests) + ] + for rid in range(num_requests): + print(f"generating for request {rid}") + for delta_outputs in engine._generate(prompts[rid], generation_cfg, request_id=str(rid)): + assert len(delta_outputs) == generation_cfg.n + for i, delta_output in enumerate(delta_outputs): + output_texts[rid][i] += delta_output.delta_text + + # Print output. + print("All finished") + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + engine.terminate() + del engine + + +if __name__ == "__main__": + test_engine_generate() diff --git a/tests/python/serve/test_serve_engine_spec.py b/tests/python/serve/test_serve_engine_spec.py new file mode 100644 index 0000000..61e76e3 --- /dev/null +++ b/tests/python/serve/test_serve_engine_spec.py @@ -0,0 +1,660 @@ +from typing import Callable, List, Optional # noqa: UP035 + +import numpy as np + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import Request, RequestStreamOutput, data +from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine +from mlc_llm.testing import require_test_model + +prompts = [ + "What is the meaning of life?", + "Introduce the history of Pittsburgh to me. Please elaborate in detail.", + "Write a three-day Seattle travel plan. Please elaborate in detail.", + "What is Alaska famous of? Please elaborate in detail.", + "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501 + "What are the necessary components to assemble a desktop computer? Please elaborate in detail.", + "Why is Vitamin D important to human beings? Please elaborate in detail.", + "Where is milk tea originated from? Please elaborate in detail.", + "Where is the southernmost place in United States? Please elaborate in detail.", + "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501 +] + + +def create_requests( + num_requests: int, + stop_token_id: Optional[int] = None, + temperature: float = 0.8, + repetition_penalty: float = 1.0, + max_tokens_low: int = 256, + max_tokens_high: int = 257, +) -> List[Request]: # noqa: UP006 + assert num_requests >= 0 and num_requests <= len(prompts) + + stop_token_ids = [stop_token_id] if stop_token_id is not None else [] + requests = [] + for req_id, prompt in zip(range(num_requests), prompts): + max_tokens = np.random.randint(max_tokens_low, max_tokens_high) + requests.append( + Request( + request_id=str(req_id), + inputs=data.TextData(prompt), + generation_config=GenerationConfig( + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens=max_tokens, + stop_token_ids=stop_token_ids, + ), + ) + ) + return requests + + +@require_test_model( + "Llama-2-7b-chat-hf-q0f16-MLC", + "Llama-2-7b-chat-hf-q4f16_1-MLC", +) +def test_engine_basic(model: str, small_model: str): + """Test engine **without continuous batching**. + + - Add all requests to the engine altogether in the beginning. + - All requests have the same max_tokens. This means all requests + will end together. + - Engine keeps running `step` for estimated number of steps (number of + requests + max_tokens - 1). Then check the output of each request. + """ + + # Hyperparameters for tests (you can try different combinations). + num_requests = len(prompts) # [4, 8, 10] + temperature = 0.9 # [0, 0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.0 # [1.0, 1.01] + max_tokens: int = 256 # [32, 128, 256] + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + + # Define the callback function for request generation results + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + + # Create engine + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[small_model], + speculative_mode="small_draft", + ), + request_stream_callback=fcallback, + ) + + # Create requests + requests = create_requests( + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max_tokens - 1 + # Run steps + for step in range(num_steps): + engine.step() + + for req_id, output in enumerate(outputs): + print(f"Prompt {req_id}: {requests[req_id].inputs[0]}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_eagle_basic(model: str): + """Test engine **without continuous batching**. + + - Add all requests to the engine altogether in the beginning. + - All requests have the same max_tokens. This means all requests + will end together. + - Engine keeps running `step` for estimated number of steps (number of + requests + max_tokens - 1). Then check the output of each request. + - Use Eagle model as speculative model + """ + + # Hyperparameters for tests (you can try different combinations). + num_requests = len(prompts) # [4, 8, 10] + temperature = 0.9 # [0, 0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.0 # [1.0, 1.01] + max_tokens: int = 256 # [32, 128, 256] + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + + # Define the callback function for request generation results + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + + # Create engine + small_model = "dist/Eagle-llama2-7b-chat-q0f16-MLC" + small_model_lib = "dist/Eagle-llama2-7b-chat-q0f16-MLC/Eagle-llama2-7b-chat-q0f16-MLC-cuda.so" + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[(small_model, small_model_lib)], + speculative_mode="eagle", + spec_draft_length=2, + ), + request_stream_callback=fcallback, + ) + + # Create requests + requests = create_requests( + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max_tokens - 1 + # Run steps + for step in range(num_steps): + engine.step() + + for req_id, output in enumerate(outputs): + print(f"Prompt {req_id}: {requests[req_id].inputs[0]}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + + +@require_test_model( + "Llama-2-7b-chat-hf-q0f16-MLC", + "Llama-2-7b-chat-hf-q4f16_1-MLC", +) +def test_engine_continuous_batching_1(model: str, small_model: str): + """Test engine **with continuous batching**. + + - Add all requests to the engine altogether in the beginning. + - All requests have a random maximum generation length. So each + request keeps generating until reaching the maximum length. + - Engine keeps running `step` for estimated number of steps (number of + requests + the maximum max_tokens - 1). Then check the output + of each request. + """ + + # Hyperparameters for tests (you can try different combinations) + num_requests = len(prompts) # [4, 8, 10] + temperature = 0.9 # [0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.00 # [1.0, 1.01] + max_tokens_low = 128 + max_tokens_high = 384 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + finish_time: List[Optional[int]] = [None] * num_requests # noqa: UP006 + + # Define the callback class for request generation results + class CallbackTimer: + timer: int = -1 + + def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]: # noqa: UP006 + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + if stream_outputs[0].finish_reason is not None: + print(f"Request {request_id} finished at step {self.timer}.") + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + finish_time[int(request_id)] = self.timer + + return fcallback + + def step(self) -> None: + self.timer += 1 + + # Create engine + timer = CallbackTimer() + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[small_model], + speculative_mode="small_draft", + ), + request_stream_callback=timer.callback_getter(), + ) + + # Create requests + requests = create_requests( + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens_low, + max_tokens_high=max_tokens_high, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max(request.generation_config.max_tokens for request in requests) - 1 + # Run steps + for step in range(num_steps): + timer.step() + assert timer.timer == step + engine.step() + + for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)): + print(f"Prompt {req_id}: {request.inputs[0]}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + # assert fin_time == request.generation_config.max_tokens - 1 + + +@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_engine_eagle_continuous_batching_1(model: str): + """Test engine **with continuous batching**. + + - Add all requests to the engine altogether in the beginning. + - All requests have a random maximum generation length. So each + request keeps generating until reaching the maximum length. + - Engine keeps running `step` for estimated number of steps (number of + requests + the maximum max_tokens - 1). Then check the output + of each request. + """ + + # Hyperparameters for tests (you can try different combinations) + num_requests = len(prompts) # [4, 8, 10] + temperature = 0.9 # [0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.00 # [1.0, 1.01] + max_tokens_low = 128 + max_tokens_high = 384 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + finish_time: List[Optional[int]] = [None] * num_requests # noqa: UP006 + + # Define the callback class for request generation results + class CallbackTimer: + timer: int = -1 + + def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]: # noqa: UP006 + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + if stream_outputs[0].finish_reason is not None: + print(f"Request {request_id} finished at step {self.timer}.") + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + finish_time[int(request_id)] = self.timer + + return fcallback + + def step(self) -> None: + self.timer += 1 + + # Create engine + small_model = "dist/Eagle-llama2-7b-chat-q4f16_1-MLC" + small_model_lib = ( + "dist/Eagle-llama2-7b-chat-q4f16_1-MLC/Eagle-llama2-7b-chat-q4f16_1-MLC-cuda.so" + ) + timer = CallbackTimer() + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[(small_model, small_model_lib)], + speculative_mode="eagle", + ), + request_stream_callback=timer.callback_getter(), + ) + + # Create requests + requests = create_requests( + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens_low, + max_tokens_high=max_tokens_high, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max(request.generation_config.max_tokens for request in requests) - 1 + # Run steps + for step in range(num_steps): + timer.step() + assert timer.timer == step + engine.step() + + for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)): + print(f"Prompt {req_id}: {request.inputs[0]}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + # assert fin_time == request.generation_config.max_tokens - 1 + + +def compare_output_text(output_text1, output_text2): + if isinstance(output_text1, list) and isinstance(output_text2, list): + for item1, item2 in zip(output_text1, output_text2): + if not compare_output_text(item1, item2): + return False + elif output_text1 != output_text2: + print(output_text1) + print(output_text2) + return False + return True + + +@require_test_model( + "Llama-2-7b-chat-hf-q0f16-MLC", + "Llama-2-7b-chat-hf-q4f16_1-MLC", +) +def test_engine_generate(model: str, small_model: str, compare_precision=False): + # Create engine + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[small_model], + speculative_mode="small_draft", + ), + ) + + num_requests = 10 + max_tokens = 256 + + # Generate output. + if compare_precision: + print("compare precision") + generation_config = GenerationConfig( + temperature=0.0, top_p=0, max_tokens=1024, stop_token_ids=[2], n=1 + ) + engine_single_model = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + ), + ) + output_texts_single_model, _ = engine_single_model.generate( + prompts[:num_requests], generation_config + ) + for req_id, outputs in enumerate(output_texts_single_model): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + # TODO: Add pytorch precision + else: + generation_config = GenerationConfig(max_tokens=max_tokens, n=3) + output_texts, _ = engine.generate(prompts[:num_requests], generation_config) + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + if compare_precision: + precision_flag = compare_output_text(output_texts, output_texts_single_model) + if precision_flag: + print("Accuracy verification succeed\n") + else: + print("Accuracy verification failed\n") + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_eagle_generate(model: str): + # Create engine + small_model = "dist/Eagle-llama2-7b-chat-q4f16_1-MLC" + small_model_lib = ( + "dist/Eagle-llama2-7b-chat-q4f16_1-MLC/Eagle-llama2-7b-chat-q4f16_1-MLC-cuda.so" + ) + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[(small_model, small_model_lib)], + speculative_mode="eagle", + ), + ) + + num_requests = 10 + max_tokens = 256 + + # Generate output. + output_texts, _ = engine.generate( + prompts[:num_requests], GenerationConfig(max_tokens=max_tokens, n=3) + ) + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + +@require_test_model("Llama-2-13b-chat-hf-q4f16_1-MLC") +def test_engine_efficiency(model: str): + """Test engine speculative decoding efficiency.""" + + # Hyperparameters for tests (you can try different combinations). + num_requests = 1 # [4, 8, 10] + temperature = 0.9 # [0, 0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.0 # [1.0, 1.01] + max_tokens: int = 512 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + + # Define the callback function for request generation results + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + + # Create engine + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + request_stream_callback=fcallback, + ) + + # Create requests + requests = create_requests( + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max_tokens - 1 + # Run steps + for step in range(num_steps): + engine.step() + + for eg, name in zip([engine], ["Normal Deconding"]): + metrics = eg.metrics() + print("engine name:", name) + if name == "Speculative Decoding": + print("spec decode metrics:", metrics["spec_decode"]) + print("engine total decode time:", metrics["engine_decode_time_sum"]) + print() + + +@require_test_model( + "Llama-2-13b-chat-hf-q4f16_1-MLC", + "Llama-2-7b-chat-hf-q4f16_1-MLC", +) +def test_engine_spec_efficiency(model: str, small_model: str): + """Test engine speculative decoding efficiency.""" + + # Hyperparameters for tests (you can try different combinations). + num_requests = 1 # [4, 8, 10] + temperature = 0.9 # [0, 0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.0 # [1.0, 1.01] + max_tokens: int = 512 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + + # Define the callback function for request generation results + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + + # Create engine + spec_engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[small_model], + spec_draft_length=6, + speculative_mode="small_draft", + ), + request_stream_callback=fcallback, + ) + + # Create requests + requests = create_requests( + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine + for request in requests: + spec_engine.add_request(request) + + num_steps = num_requests + max_tokens - 1 + # Run steps + for step in range(num_steps): + spec_engine.step() + + for eg, name in zip([spec_engine], ["Speculative Decoding"]): + metrics = eg.metrics() + print("engine name:", name) + if name == "Speculative Decoding": + print("total draft tokens:", metrics["sum_num_draft_tokens"]) + print("total accepted tokens:", metrics["sum_num_accepted_tokens"]) + print( + "Accept rate:", + metrics["sum_num_accepted_tokens"] / (1e-10 + metrics["sum_num_draft_tokens"]), + ) + print("engine total decode time:", metrics["engine_decode_time_sum"]) + print() + + +@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_engine_eagle_spec_efficiency(model: str): + """Test engine speculative decoding efficiency.""" + + # Hyperparameters for tests (you can try different combinations). + num_requests = 1 # [4, 8, 10] + temperature = 0.9 # [0, 0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.0 # [1.0, 1.01] + max_tokens: int = 512 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + + # Define the callback function for request generation results + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + + # Create engine + small_model = "dist/Eagle-llama2-7b-chat-q0f16-MLC" + small_model_lib = "dist/Eagle-llama2-7b-chat-q0f16-MLC/Eagle-llama2-7b-chat-q0f16-MLC-cuda.so" + spec_engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig( + max_total_sequence_length=4096, + additional_models=[(small_model, small_model_lib)], + spec_draft_length=6, + speculative_mode="eagle", + ), + request_stream_callback=fcallback, + ) + + # Create requests + requests = create_requests( + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine + for request in requests: + spec_engine.add_request(request) + + num_steps = num_requests + max_tokens - 1 + # Run steps + for step in range(num_steps): + spec_engine.step() + + for eg, name in zip([spec_engine], ["Speculative Decoding"]): + metrics = eg.metrics() + print("engine name:", name) + if name == "Speculative Decoding": + print("spec decode:", metrics["spec_decode"]) + print("engine total decode time:", metrics["engine_decode_time_sum"]) + print() + + +if __name__ == "__main__": + test_engine_basic() + test_engine_eagle_basic() + test_engine_continuous_batching_1() + test_engine_eagle_continuous_batching_1() + test_engine_generate(compare_precision=True) + test_engine_eagle_generate() + test_engine_efficiency() + test_engine_spec_efficiency() + test_engine_eagle_spec_efficiency() diff --git a/tests/python/serve/test_serve_sync_engine.py b/tests/python/serve/test_serve_sync_engine.py new file mode 100644 index 0000000..9a5807c --- /dev/null +++ b/tests/python/serve/test_serve_sync_engine.py @@ -0,0 +1,474 @@ +from typing import Callable, List, Optional # noqa: UP035 + +import numpy as np + +from mlc_llm.protocol.generation_config import GenerationConfig +from mlc_llm.serve import Request, RequestStreamOutput, data +from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine +from mlc_llm.testing import require_test_model + +prompts = [ + "What is the meaning of life?", + "Introduce the history of Pittsburgh to me. Please elaborate in detail.", + "Write a three-day Seattle travel plan. Please elaborate in detail.", + "What is Alaska famous of? Please elaborate in detail.", + "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.", # noqa: E501 + "What are the necessary components to assemble a desktop computer? Please elaborate in detail.", + "Why is Vitamin D important to human beings? Please elaborate in detail.", + "Where is milk tea originated from? Please elaborate in detail.", + "Where is the southernmost place in United States? Please elaborate in detail.", + "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.", # noqa: E501 +] + + +def create_requests( + engine: SyncMLCEngine, + num_requests: int, + stop_token_id: Optional[int] = None, + temperature: float = 0.8, + repetition_penalty: float = 1.0, + max_tokens_low: int = 256, + max_tokens_high: int = 257, +) -> List[Request]: # noqa: UP006 + assert num_requests >= 0 and num_requests <= len(prompts) + + stop_token_ids = [stop_token_id] if stop_token_id is not None else [] + requests = [] + for req_id, prompt in zip(range(num_requests), prompts): + max_tokens = np.random.randint(max_tokens_low, max_tokens_high) + requests.append( + engine.create_request( + request_id=str(req_id), + inputs=data.TextData(prompt), + generation_config=GenerationConfig( + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens=max_tokens, + stop_token_ids=stop_token_ids, + ), + ) + ) + return requests + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_basic(model: str): + """Test engine **without continuous batching**. + + - Add all requests to the engine altogether in the beginning. + - All requests have the same max_tokens. This means all requests + will end together. + - Engine keeps running `step` for estimated number of steps (number of + requests + max_tokens - 1). Then check the output of each request. + """ + + # Hyperparameters for tests (you can try different combinations). + num_requests = 10 # [4, 8, 10] + temperature = 0.9 # [0, 0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.0 # [1.0, 1.01] + max_tokens: int = 256 # [32, 128, 256] + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + + # Define the callback function for request generation results + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + + # Create engine + engine = SyncMLCEngine( + model=model, + mode="server", + request_stream_callback=fcallback, + ) + + # Create requests + requests = create_requests( + engine, + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max_tokens - 1 + # Run steps + for step in range(num_steps): + engine.step() + + for req_id, output in enumerate(outputs): + print(f"Prompt {req_id}: {requests[req_id].inputs[0]}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_continuous_batching_1(model: str): + """Test engine **with continuous batching**. + + - Add all requests to the engine altogether in the beginning. + - All requests have a random maximum generation length. So each + request keeps generating until reaching the maximum length. + - Engine keeps running `step` for estimated number of steps (number of + requests + the maximum max_tokens - 1). Then check the output + of each request. + """ + + # Hyperparameters for tests (you can try different combinations) + num_requests = 10 # [4, 8, 10] + temperature = 0.9 # [0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.00 # [1.0, 1.01] + max_tokens_low = 128 + max_tokens_high = 384 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + finish_time: List[Optional[int]] = [None] * num_requests # noqa: UP006 + + # Define the callback class for request generation results + class CallbackTimer: + timer: int = -1 + + def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]: # noqa: UP006 + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + if stream_outputs[0].finish_reason is not None: + print(f"Request {request_id} finished at step {self.timer}.") + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + finish_time[int(request_id)] = self.timer + + return fcallback + + def step(self) -> None: + self.timer += 1 + + # Create engine + timer = CallbackTimer() + engine = SyncMLCEngine( + model=model, + mode="server", + request_stream_callback=timer.callback_getter(), + ) + + # Create requests + requests = create_requests( + engine, + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens_low, + max_tokens_high=max_tokens_high, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max(request.generation_config.max_tokens for request in requests) - 1 + # Run steps + for step in range(num_steps): + timer.step() + assert timer.timer == step + engine.step() + + for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)): + print(f"Prompt {req_id}: {request.inputs[0]}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + assert fin_time == request.generation_config.max_tokens - 1, ( + f"finish time = {fin_time}, max tokens = {request.generation_config.max_tokens - 1}" + ) + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_continuous_batching_2(model: str): + """Test engine **with continuous batching**. + + - Add all requests to the engine altogether in the beginning. + - All requests have the stop token. So each request keeps generating + until having the stop token or reaching the maximum length. + - Engine keeps running `step` for estimated number of steps (number of + requests + the maximum max_tokens - 1). Then check the output + of each request. + """ + + # Hyperparameters for tests (you can try different combinations) + num_requests = 10 # [4, 8, 10] + temperature = 0.9 # [0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.00 # [1.0, 1.01] + stop_token_id = 2 + max_tokens = 512 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + finish_time: List[Optional[int]] = [None] * num_requests # noqa: UP006 + + # Define the callback class for request generation results + class CallbackTimer: + timer: int = -1 + + def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]: # noqa: UP006 + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + if stream_outputs[0].finish_reason is not None: + print(f"Request {request_id} finished at step {self.timer}.") + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + finish_time[int(request_id)] = self.timer + + return fcallback + + def step(self) -> None: + self.timer += 1 + + # Create engine + timer = CallbackTimer() + engine = SyncMLCEngine( + model=model, + mode="server", + request_stream_callback=timer.callback_getter(), + ) + + # Create requests + requests = create_requests( + engine, + num_requests, + stop_token_id=stop_token_id, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine + for request in requests: + engine.add_request(request) + + num_steps = num_requests + max_tokens - 1 + # Run steps + for step in range(num_steps): + timer.step() + assert timer.timer == step + engine.step() + + for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)): + print(f"Prompt {req_id}: {request.inputs[0]}") + if fin_time < num_requests + max_tokens - 2: + print(f"Request {req_id} ends early on the stop token") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_continuous_batching_3(model: str): + """Test engine **with continuous batching**. + + - Add requests randomly between time [0, 200). + - All requests have a random maximum generation length. So each + request keeps generating until reaching the maximum length. + - Engine keeps running `step` until all requests finish. + Then check the output of each request. + """ + + # Hyperparameters for tests (you can try different combinations) + num_requests = 10 # [4, 8, 10] + temperature = 0.9 # [0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.00 # [1.0, 1.01] + stop_token_id = 2 + max_tokens_low = 64 + max_tokens_high = 192 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + finish_time: List[Optional[int]] = [None] * num_requests # noqa: UP006 + + # Define the callback class for request generation results + class CallbackTimer: + timer: int = -1 + finished_requests: int = 0 + + def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]: # noqa: UP006 + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + if stream_outputs[0].finish_reason is not None: + print(f"Request {request_id} finished at step {self.timer}.") + self.finished_requests += 1 + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + finish_time[int(request_id)] = self.timer + + return fcallback + + def step(self) -> None: + self.timer += 1 + + def all_finished(self) -> bool: + return self.finished_requests == num_requests + + # Create engine + timer = CallbackTimer() + engine = SyncMLCEngine( + model=model, + mode="server", + request_stream_callback=timer.callback_getter(), + ) + + # Create requests + requests = create_requests( + engine, + num_requests, + stop_token_id=stop_token_id, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens_low, + max_tokens_high=max_tokens_high, + ) + + # Assign the time to add requests to engine + request_add_time = [np.random.randint(0, 200) for _ in range(num_requests)] + + # Run steps + while not timer.all_finished(): + timer.step() + + # Add requests to engine + for req_id, add_time in enumerate(request_add_time): + if add_time == timer.timer: + print(f"add request {req_id} at step {timer.timer}") + engine.add_request(requests[req_id]) + + engine.step() + + for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)): + print(f"Prompt {req_id}: {request.inputs[0]}") + print(f"Finish time: {fin_time}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_generate(model: str): + # Create engine + engine = SyncMLCEngine( + model=model, + mode="server", + engine_config=EngineConfig(max_total_sequence_length=4096), + ) + + num_requests = 10 + max_tokens = 256 + + # Generate output. + output_texts, _ = engine.generate( + prompts[:num_requests], GenerationConfig(max_tokens=max_tokens, n=7) + ) + for req_id, outputs in enumerate(output_texts): + print(f"Prompt {req_id}: {prompts[req_id]}") + if len(outputs) == 1: + print(f"Output {req_id}:{outputs[0]}\n") + else: + for i, output in enumerate(outputs): + print(f"Output {req_id}({i}):{output}\n") + + +@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC") +def test_engine_hybrid_prefill(model: str): + """Test engine **with hybrid prefill**. + + - Add each single request step by step. + - All requests have the same generation length. But due to hybrid prefill, + the earlier request will decode with later request prefill, in single step. + So each request lasts the same steps, and stops generation step by step as well. + - Engine keeps running `step` for the generation length, to finish the last request. + Then check the output of each request. + """ + + # Hyperparameters for tests (you can try different combinations) + num_requests = 10 # [4, 8, 10] + temperature = 0.9 # [0.8, 0.9, 1.0, 1.1] + repetition_penalty = 1.00 # [1.0, 1.01] + max_tokens = 15 + np.random.seed(0) + + # Output list + outputs: List[List[int]] = [[] for _ in range(num_requests)] # noqa: UP006 + finish_time: List[Optional[int]] = [None] * num_requests # noqa: UP006 + + # Define the callback class for request generation results + class CallbackTimer: + timer: int = -1 + + def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]: # noqa: UP006 + def fcallback(delta_outputs: List[RequestStreamOutput]): # noqa: UP006 + for delta_output in delta_outputs: + request_id, stream_outputs = delta_output.unpack() + assert len(stream_outputs) == 1 + if stream_outputs[0].finish_reason is not None: + print(f"Request {request_id} finished at step {self.timer}.") + outputs[int(request_id)] += stream_outputs[0].delta_token_ids + finish_time[int(request_id)] = self.timer + + return fcallback + + def step(self) -> None: + self.timer += 1 + + # Create engine + timer = CallbackTimer() + engine = SyncMLCEngine( + model=model, + mode="server", + request_stream_callback=timer.callback_getter(), + engine_config=EngineConfig(prefill_mode="hybrid"), + ) + + # Create requests + requests = create_requests( + engine, + num_requests, + temperature=temperature, + repetition_penalty=repetition_penalty, + max_tokens_low=max_tokens, + max_tokens_high=max_tokens + 1, + ) + + # Add all requests to engine step by step + for step, request in enumerate(requests): + engine.add_request(request) + timer.step() + assert timer.timer == step + engine.step() + + # Run steps + for step in range(max_tokens): + timer.step() + assert timer.timer == step + num_requests + engine.step() + + for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)): + print(f"Prompt {req_id}: {request.inputs[0]}") + print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n") + assert fin_time == req_id + request.generation_config.max_tokens - 1, ( + f"finish time = {fin_time}, max tokens = {req_id + request.generation_config.max_tokens - 1}" # noqa: E501 + ) + + +if __name__ == "__main__": + test_engine_basic() + test_engine_continuous_batching_1() + test_engine_continuous_batching_2() + test_engine_continuous_batching_3() + test_engine_generate() + test_engine_hybrid_prefill() diff --git a/tests/python/support/test_auto_config.py b/tests/python/support/test_auto_config.py new file mode 100644 index 0000000..276db0b --- /dev/null +++ b/tests/python/support/test_auto_config.py @@ -0,0 +1,43 @@ +import json +import tempfile +from pathlib import Path + +import pytest + +from mlc_llm.support import logging +from mlc_llm.support.auto_config import detect_config + +logging.enable_logging() + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +def _create_json_file(json_path, data): + with open(json_path, "w", encoding="utf-8") as i_f: + json.dump(data, i_f) + + +def test_detect_config(): + with tempfile.TemporaryDirectory() as tmpdir: + base_path = Path(tmpdir) + config_json_path = base_path / "config.json" + _create_json_file(config_json_path, {}) + + assert detect_config(base_path) == config_json_path + assert detect_config(config_json_path) == config_json_path + + +def test_detect_config_fail(): + with pytest.raises(ValueError): + detect_config(Path("do/not/exist")) + + with tempfile.TemporaryDirectory() as tmpdir: + base_path = Path(tmpdir) + with pytest.raises(ValueError): + assert detect_config(base_path) + + +if __name__ == "__main__": + test_detect_config() + test_detect_config_fail() diff --git a/tests/python/support/test_auto_target.py b/tests/python/support/test_auto_target.py new file mode 100644 index 0000000..0386626 --- /dev/null +++ b/tests/python/support/test_auto_target.py @@ -0,0 +1,37 @@ +import pytest +from tvm.target import Target + +from mlc_llm.support.auto_target import _apply_webgpu_subgroups + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +def test_apply_webgpu_subgroups_enables_webgpu_target(): + target = Target("webgpu") + + updated = _apply_webgpu_subgroups(target, True) + + assert updated is not target + assert dict(target.export())["supports_subgroups"] is False + assert dict(updated.export())["supports_subgroups"] is True + + +def test_apply_webgpu_subgroups_non_webgpu_target_is_unchanged(): + target = Target("llvm") + + updated = _apply_webgpu_subgroups(target, True) + + assert updated is target + assert dict(updated.export()) == dict(target.export()) + + +@pytest.mark.parametrize("target_kind", ["webgpu", "llvm"]) +@pytest.mark.parametrize("enable_subgroups", [False, None]) +def test_apply_webgpu_subgroups_disabled_is_unchanged(target_kind, enable_subgroups): + target = Target(target_kind) + + updated = _apply_webgpu_subgroups(target, enable_subgroups) + + assert updated is target + assert dict(updated.export()) == dict(target.export()) diff --git a/tests/python/support/test_auto_weight.py b/tests/python/support/test_auto_weight.py new file mode 100644 index 0000000..879bedc --- /dev/null +++ b/tests/python/support/test_auto_weight.py @@ -0,0 +1,150 @@ +import json +import os +import tempfile +from pathlib import Path + +import pytest + +from mlc_llm.support import logging +from mlc_llm.support.auto_weight import detect_weight + +logging.enable_logging() + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +def _create_json_file(json_path, data): + with open(json_path, "w", encoding="utf-8") as i_f: + json.dump(data, i_f) + + +@pytest.mark.parametrize( + "weight_format, index_filename, result", + [ + ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"), + ( + "huggingface-safetensor", + "model.safetensors.index.json", + "huggingface-safetensor", + ), + ("auto", "pytorch_model.bin.index.json", "huggingface-torch"), + ("auto", "model.safetensors.index.json", "huggingface-safetensor"), + ], +) +def test_detect_weight(weight_format, index_filename, result): + with tempfile.TemporaryDirectory() as tmpdir: + base_path = Path(tmpdir) + if index_filename is not None: + weight_index_file = base_path / index_filename + _create_json_file(weight_index_file, {}) + assert detect_weight(base_path, None, weight_format) == ( + weight_index_file, + result, + ) + + +@pytest.mark.parametrize( + "weight_format, index_filename, result", + [ + ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"), + ( + "huggingface-safetensor", + "model.safetensors.index.json", + "huggingface-safetensor", + ), + ("auto", "pytorch_model.bin.index.json", "huggingface-torch"), + ("auto", "model.safetensors.index.json", "huggingface-safetensor"), + ], +) +def test_detect_weight_in_config_json(weight_format, index_filename, result): + with ( + tempfile.TemporaryDirectory() as config_dir, + tempfile.TemporaryDirectory() as weight_dir, + ): + config_path = Path(config_dir) + weight_path = Path(weight_dir) + config_json_path = config_path / "config.json" + _create_json_file(config_json_path, {"weight_path": weight_dir}) + if index_filename is not None: + weight_index_file = weight_path / index_filename + _create_json_file(weight_index_file, {}) + + assert detect_weight(None, config_json_path, weight_format) == ( + weight_index_file, + result, + ) + + +@pytest.mark.parametrize( + "weight_format, index_filename, result", + [ + ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"), + ( + "huggingface-safetensor", + "model.safetensors.index.json", + "huggingface-safetensor", + ), + ("auto", "pytorch_model.bin.index.json", "huggingface-torch"), + ("auto", "model.safetensors.index.json", "huggingface-safetensor"), + ], +) +def test_detect_weight_same_dir_config_json(weight_format, index_filename, result): + with tempfile.TemporaryDirectory() as tmpdir: + base_path = Path(tmpdir) + config_json_path = base_path / "config.json" + _create_json_file(config_json_path, {}) + if index_filename is not None: + weight_index_file = Path(os.path.join(tmpdir, index_filename)) + _create_json_file(weight_index_file, {}) + assert detect_weight(None, config_json_path, weight_format) == ( + weight_index_file, + result, + ) + + +def test_find_weight_fail(): + with tempfile.TemporaryDirectory() as tmpdir: + base_path = Path(tmpdir) + with pytest.raises(ValueError): + detect_weight(Path("do/not/exist"), base_path, "awq") + with pytest.raises(AssertionError): + detect_weight(None, Path("do/not/exist"), "awq") + + +if __name__ == "__main__": + test_detect_weight("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch") + test_detect_weight( + "huggingface-safetensor", + "model.safetensors.index.json", + "huggingface-safetensor", + ) + test_detect_weight("auto", "pytorch_model.bin.index.json", "huggingface-torch") + test_detect_weight("auto", "model.safetensors.index.json", "huggingface-safetensor") + test_detect_weight_in_config_json( + "huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch" + ) + test_detect_weight_in_config_json( + "huggingface-safetensor", + "model.safetensors.index.json", + "huggingface-safetensor", + ) + test_detect_weight_in_config_json("auto", "pytorch_model.bin.index.json", "huggingface-torch") + test_detect_weight_in_config_json( + "auto", "model.safetensors.index.json", "huggingface-safetensor" + ) + test_detect_weight_same_dir_config_json( + "huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch" + ) + test_detect_weight_same_dir_config_json( + "huggingface-safetensor", + "model.safetensors.index.json", + "huggingface-safetensor", + ) + test_detect_weight_same_dir_config_json( + "auto", "pytorch_model.bin.index.json", "huggingface-torch" + ) + test_detect_weight_same_dir_config_json( + "auto", "model.safetensors.index.json", "huggingface-safetensor" + ) + test_find_weight_fail() diff --git a/tests/python/support/test_cli_convert_weight.py b/tests/python/support/test_cli_convert_weight.py new file mode 100644 index 0000000..d741087 --- /dev/null +++ b/tests/python/support/test_cli_convert_weight.py @@ -0,0 +1,69 @@ +import json +import tempfile +from pathlib import Path + +import pytest + +from mlc_llm.cli import convert_weight as convert_weight_cli + +pytestmark = [pytest.mark.unittest] + + +def test_convert_weight_cli_passes_lora_adapter(monkeypatch): + with tempfile.TemporaryDirectory() as tmp_dir: + temp_path = Path(tmp_dir) + config_path = temp_path / "config.json" + source_dir = temp_path / "source" + source_dir.mkdir(parents=True, exist_ok=True) + source_index = source_dir / "pytorch_model.bin.index.json" + adapter_dir = temp_path / "adapter" + adapter_dir.mkdir(parents=True, exist_ok=True) + output_dir = temp_path / "output" + + config_path.write_text(json.dumps({}), encoding="utf-8") + source_index.write_text(json.dumps({"weight_map": {}}), encoding="utf-8") + + def _fake_detect_device(device): + return device + + def _fake_detect_weight(_weight_path, _config_json_path, _weight_format): + return source_index, "huggingface-torch" + + def _fake_detect_model_type(_model_type, _config): + return "dummy" + + monkeypatch.setattr(convert_weight_cli, "detect_config", Path) + monkeypatch.setattr(convert_weight_cli, "detect_device", _fake_detect_device) + monkeypatch.setattr(convert_weight_cli, "detect_weight", _fake_detect_weight) + monkeypatch.setattr(convert_weight_cli, "detect_model_type", _fake_detect_model_type) + monkeypatch.setattr(convert_weight_cli, "MODELS", {"dummy": object()}) + monkeypatch.setattr(convert_weight_cli, "QUANTIZATION", {"q0f16": object()}) + + call_args = {} + + def _fake_convert_weight(**kwargs): + call_args.update(kwargs) + + monkeypatch.setattr(convert_weight_cli, "convert_weight", _fake_convert_weight) + + convert_weight_cli.main( + [ + str(config_path), + "--quantization", + "q0f16", + "--model-type", + "dummy", + "--source", + str(source_dir), + "--source-format", + "auto", + "--output", + str(output_dir), + "--lora-adapter", + str(adapter_dir), + ] + ) + + assert call_args["lora_adapter"] == adapter_dir + assert call_args["source"] == source_index + assert call_args["source_format"] == "huggingface-torch" diff --git a/tests/python/support/test_convert_weight_lora_merge.py b/tests/python/support/test_convert_weight_lora_merge.py new file mode 100644 index 0000000..e537646 --- /dev/null +++ b/tests/python/support/test_convert_weight_lora_merge.py @@ -0,0 +1,108 @@ +import contextlib +import json +import tempfile +from pathlib import Path + +import pytest + +from mlc_llm.interface import convert_weight as convert_weight_interface + +pytestmark = [pytest.mark.unittest] + + +def test_resolve_base_model_dir(): + with tempfile.TemporaryDirectory() as tmp_dir: + temp_path = Path(tmp_dir) + model_dir = temp_path / "model" + model_dir.mkdir(parents=True, exist_ok=True) + source_file = model_dir / "pytorch_model.bin.index.json" + source_file.write_text(json.dumps({"weight_map": {}}), encoding="utf-8") + + assert convert_weight_interface._resolve_base_model_dir(model_dir) == model_dir + assert convert_weight_interface._resolve_base_model_dir(source_file) == model_dir + + +def test_convert_weight_with_lora_uses_merged_source(monkeypatch): + with tempfile.TemporaryDirectory() as tmp_dir: + temp_path = Path(tmp_dir) + config_path = temp_path / "config.json" + config_path.write_text(json.dumps({}), encoding="utf-8") + + source_dir = temp_path / "source" + source_dir.mkdir(parents=True, exist_ok=True) + source_file = source_dir / "pytorch_model.bin.index.json" + source_file.write_text(json.dumps({"weight_map": {}}), encoding="utf-8") + + adapter_dir = temp_path / "adapter" + adapter_dir.mkdir(parents=True, exist_ok=True) + + merged_dir = temp_path / "merged" + merged_dir.mkdir(parents=True, exist_ok=True) + merged_file = merged_dir / "pytorch_model.bin" + merged_file.write_bytes(b"") + + captured = {} + + @contextlib.contextmanager + def _fake_merge(base_source: Path, lora_adapter: Path): + captured["merge_base_source"] = base_source + captured["merge_lora_adapter"] = lora_adapter + yield merged_dir + + def _fake_detect_weight(weight_path: Path, config_json_path: Path, weight_format: str): + captured["detect_weight_path"] = weight_path + captured["detect_weight_config"] = config_json_path + captured["detect_weight_format"] = weight_format + return merged_file, "huggingface-torch" + + def _fake_convert_args(args): + captured["converted_args"] = args + + monkeypatch.setattr( + convert_weight_interface, "_merge_lora_adapter_with_base_model", _fake_merge + ) + monkeypatch.setattr(convert_weight_interface, "detect_weight", _fake_detect_weight) + monkeypatch.setattr(convert_weight_interface, "_convert_args", _fake_convert_args) + monkeypatch.setattr(convert_weight_interface.ConversionArgs, "display", lambda self: None) + + convert_weight_interface.convert_weight( + config=config_path, + quantization=object(), + model=type("DummyModel", (), {"name": "dummy"})(), + device=object(), + source=source_file, + source_format="huggingface-safetensor", + output=temp_path / "output", + lora_adapter=adapter_dir, + ) + + converted_args = captured["converted_args"] + assert captured["merge_base_source"] == source_file + assert captured["merge_lora_adapter"] == adapter_dir + assert captured["detect_weight_path"] == merged_dir + assert captured["detect_weight_config"] == config_path + assert captured["detect_weight_format"] == "auto" + assert converted_args.source == merged_file + assert converted_args.source_format == "huggingface-torch" + assert converted_args.lora_adapter == adapter_dir + + +def test_convert_weight_with_lora_rejects_awq(): + with tempfile.TemporaryDirectory() as tmp_dir: + temp_path = Path(tmp_dir) + config_path = temp_path / "config.json" + config_path.write_text(json.dumps({}), encoding="utf-8") + adapter_dir = temp_path / "adapter" + adapter_dir.mkdir(parents=True, exist_ok=True) + + with pytest.raises(ValueError, match="only supports source formats"): + convert_weight_interface.convert_weight( + config=config_path, + quantization=object(), + model=type("DummyModel", (), {"name": "dummy"})(), + device=object(), + source=temp_path / "source", + source_format="awq", + output=temp_path / "output", + lora_adapter=adapter_dir, + ) diff --git a/tests/python/tokenizers/test_streamer.py b/tests/python/tokenizers/test_streamer.py new file mode 100644 index 0000000..3bd60a4 --- /dev/null +++ b/tests/python/tokenizers/test_streamer.py @@ -0,0 +1,192 @@ +"""Streamer tests in MLC LLM. + +Please specify the local path to llama2 tokenizer via environment +variable before running this test. +The recommended way to run the tests is to use the following command: + MLC_LLAMA_TOKENIZER_PATH="path/to/llama/tokenizer" \ + pytest -vv tests/python/support/test_text_streamer_stop_handler.py + +Here "MLC_LLAMA_TOKENIZER_PATH" can be chosen from +- a llama2 weight directory (e.g., "path/to/Llama-2-7b-chat-hf"), +- a sentencepiece llama2 tokenizer path + (e.g., "path/to/Llama-2-7b-chat-hf/tokenizer.model"). + +To directly run the Python file (a.k.a., not using pytest), you also need to +specify the tokenizer path via environment variable. +""" + +import time +from typing import List, Tuple # noqa: UP035 + +import pytest + +from mlc_llm.testing import require_test_tokenizers +from mlc_llm.tokenizers import StopStrHandler, TextStreamer, Tokenizer + +# test category "unittest" +pytestmark = [pytest.mark.unittest] + + +# fmt: off +para_input_tokens = [18585, 29892, 1244, 29915, 29879, 263, 3273, 14880, 1048, 953, 29877, 2397, + 29892, 988, 1269, 1734, 338, 5643, 491, 385, 953, 29877, 2397, 29901, 13, 13, + 29950, 1032, 727, 29991, 29871, 243, 162, 148, 142, 306, 29915, 29885, 1244, 304, + 1371, 1234, 738, 5155, 366, 505, 1048, 953, 29877, 2397, 29871, 243, 162, 167, 151, + 29889, 7440, 366, 1073, 393, 953, 29877, 2397, 508, 367, 1304, 304, 27769, 23023, + 1080, 322, 21737, 297, 263, 2090, 322, 1708, 1319, 982, 29973, 29871, 243, 162, 155, + 135, 2688, 508, 884, 367, 1304, 304, 788, 263, 6023, 310, 2022, 2877, 304, 596, 7191, + 322, 11803, 29889, 29871, 243, 162, 149, 152, 1126, 29892, 1258, 366, 1073, 393, 727, + 526, 1584, 953, 29877, 2397, 8090, 322, 14188, 366, 508, 1708, 29973, 29871, 243, 162, + 145, 177, 243, 162, 148, 131, 1105, 29892, 748, 14432, 322, 679, 907, 1230, 411, 953, + 29877, 2397, 29991, 29871, 243, 162, 149, 168, 243, 162, 145, 171] + +DECODED_PARAGRAPH = ( + "Sure, here's a short paragraph about emoji, " + "where each word is followed by an emoji:\n\n" + "Hey there! 👋 I'm here to help answer any questions you have about emoji 🤔. " + "Did you know that emoji can be used to convey emotions and feelings in a " + "fun and playful way? 😄 " + "They can also be used to add a touch of personality to your messages and posts. 💕 " + "And, did you know that there are even emoji games and activities you can play? 🎮👀 " + "So, go ahead and get creative with emoji! 💥🎨" +) +# fmt: on + + +@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_text_streamer(llama_tokenizer_path: str): + text_streamer = TextStreamer(Tokenizer(llama_tokenizer_path)) + total_text = "" + for token in para_input_tokens: + total_text += text_streamer.put([token]) + total_text += text_streamer.finish() + + assert total_text == DECODED_PARAGRAPH + + +def stop_handler_process_tokens( + stop_handler: StopStrHandler, + tokens: List[int], # noqa: UP006 + tokenizer: Tokenizer, +) -> str: + returned_tokens = [] + for token in tokens: + returned_tokens += stop_handler.put(token) + if stop_handler.stop_triggered: + break + + if not stop_handler.stop_triggered: + returned_tokens += stop_handler.finish() + + return tokenizer.decode(returned_tokens) + + +@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_stop_str_handler_stop(llama_tokenizer_path: str): + stop_strs = [" 🤔"] + tokenizer = Tokenizer(llama_tokenizer_path) + stop_handler = StopStrHandler(stop_strs, tokenizer) + + total_text = stop_handler_process_tokens(stop_handler, para_input_tokens, tokenizer) + expected_text = ( + "Sure, here's a short paragraph about emoji, " + "where each word is followed by an emoji:\n\n" + "Hey there! 👋 I'm here to help answer any questions you have about emoji" + ) + + assert total_text == expected_text + + +@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_stop_str_handler_not_stop( + llama_tokenizer_path: str, +): + stop_strs = ["^^"] + tokenizer = Tokenizer(llama_tokenizer_path) + stop_handler = StopStrHandler(stop_strs, tokenizer) + + total_text = stop_handler_process_tokens(stop_handler, para_input_tokens, tokenizer) + assert total_text == DECODED_PARAGRAPH + + +@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_stop_str_handler_return_cached_tokens( + llama_tokenizer_path: str, +): + tokens = para_input_tokens[:26] # until "\n\n" + stop_strs = ["\n\n\n"] + tokenizer = Tokenizer(llama_tokenizer_path) + stop_handler = StopStrHandler(stop_strs, tokenizer) + + total_text = stop_handler_process_tokens(stop_handler, tokens, tokenizer) + expected_text = ( + "Sure, here's a short paragraph about emoji, where each word is followed by an emoji:\n\n" + ) + + assert total_text == expected_text + + +@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_stop_str_handler_throughput( + llama_tokenizer_path: str, +): + stop_strs = ["[INST]"] + tokenizer = Tokenizer(llama_tokenizer_path) + stop_handler = StopStrHandler(stop_strs, tokenizer) + + tokens = para_input_tokens * 20 + returned_tokens = [] + + tbegin = time.perf_counter() + for token in tokens: + returned_tokens += stop_handler.put(token) + assert not stop_handler.stop_triggered + tend = time.perf_counter() + + throughput = len(tokens) / (tend - tbegin) + print( + f"num tokens = {len(tokens)}, " + f"time elapsed = {tend - tbegin:.5f} sec, " + f"throughput = {throughput}" + ) + assert throughput >= 100000 + + +emoji_tokens_expected_result = [ + # HF: "�����", SentencePiece: "�👀" + ([177, 243, 162, 148, 131], ("�����", "�👀")), + # Both: "👀👀" + ([243, 162, 148, 131, 243, 162, 148, 131], ("👀👀",)), + # Both: "👀👀👀" + ([243, 162, 148, 131, 243, 162, 148, 131, 243, 162, 148, 131], ("👀👀👀",)), + # HF: "👀�������", SentencePiece: "👀���👀" + ([243, 162, 148, 131, 162, 148, 131, 243, 162, 148, 131], ("👀�������", "👀���👀")), + # Both: "👀��� have👀" + ([243, 162, 148, 131, 162, 148, 131, 505, 243, 162, 148, 131], ("👀��� have👀",)), +] + + +@pytest.mark.parametrize("tokens_and_results", emoji_tokens_expected_result) +@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC") +def test_text_streamer_emojis( + llama_tokenizer_path: str, + tokens_and_results: Tuple[List[int], Tuple[str]], # noqa: UP006 +): + text_streamer = TextStreamer(Tokenizer(llama_tokenizer_path)) + total_text = "" + tokens, expected_results = tokens_and_results + for token in tokens: + total_text += text_streamer.put([token]) + total_text += text_streamer.finish() + assert total_text in expected_results + + +if __name__ == "__main__": + test_text_streamer() + test_stop_str_handler_stop() + test_stop_str_handler_not_stop() + test_stop_str_handler_return_cached_tokens() + test_stop_str_handler_throughput() + + for tokens_and_res in emoji_tokens_expected_result: + test_text_streamer_emojis(tokens_and_res) diff --git a/version.py b/version.py new file mode 100644 index 0000000..52ed223 --- /dev/null +++ b/version.py @@ -0,0 +1,191 @@ +""" +This is the global script that set the version information of TVM. +This script runs and update all the locations that related to versions + +List of affected files: +- mlc-llm-root/pyproject.toml +""" + +import argparse +import logging +import os +import re +import subprocess + +# Modify the following value during release +# --------------------------------------------------- +# Current version: +# We use the version of the incoming release for code +# that is under development. +# It is also fallback version to be used when --git-describe +# is not invoked, or when the repository does not present the +# git tags in a format that this script can use. +# Two tag formats are supported: +# - vMAJ.MIN.PATCH (e.g. v0.8.0) or +# - vMAJ.MIN.devN (e.g. v0.8.dev0) + +# --------------------------------------------------- + +__version__ = "0.1.dev0" +PROJ_ROOT = os.path.dirname(os.path.abspath(os.path.expanduser(__file__))) + + +def py_str(cstr): + return cstr.decode("utf-8") + + +def git_describe_version(): + """Get PEP-440 compatible public and local version using git describe. + + Returns + ------- + pub_ver: str + Public version. + + local_ver: str + Local version (with additional label appended to pub_ver). + + Notes + ----- + - We follow PEP 440's convention of public version + and local versions. + - Only tags conforming to vMAJOR.MINOR.REV (e.g. "v0.7.0") + are considered in order to generate the version string. + See the use of `--match` in the `git` command below. + + Here are some examples: + + - pub_ver = '0.7.0', local_ver = '0.7.0': + We are at the 0.7.0 release. + - pub_ver = '0.8.dev94', local_ver = '0.8.dev94+g0d07a329e': + We are at the 0.8 development cycle. + The current source contains 94 additional commits + after the most recent tag(v0.7.0), + the git short hash tag of the current commit is 0d07a329e. + """ + cmd = [ + "git", + "describe", + "--tags", + "--match", + "v[0-9]*.[0-9]*.[0-9]*", + "--match", + "v[0-9]*.[0-9]*.dev[0-9]*", + ] + with subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=PROJ_ROOT, + ) as proc: + (out, _) = proc.communicate() + + if proc.returncode != 0: + msg = py_str(out) + logging.warning("git describe: %s", msg) + return None, None + describe = py_str(out).strip() + arr_info = describe.split("-") + + # Remove the v prefix, mainly to be robust + # to the case where v is not presented as well. + if arr_info[0].startswith("v"): + arr_info[0] = arr_info[0][1:] + + # hit the exact tag + if len(arr_info) == 1: + return arr_info[0], arr_info[0] + + if len(arr_info) != 3: + logging.warning("Invalid output from git describe %s", describe) + return None, None + + dev_pos = arr_info[0].find(".dev") + + # Development versions: + # The code will reach this point in case it can't match a full release version, such as v0.7.0. + # 1. in case the last known label looks like vMAJ.MIN.devN e.g. v0.8.dev0, we use + # the current behavior of just using vMAJ.MIN.devNNNN+gGIT_REV + if dev_pos != -1: + dev_version = arr_info[0][: arr_info[0].find(".dev")] + # 2. in case the last known label looks like vMAJ.MIN.PATCH e.g. v0.8.0 + # then we just carry on with a similar version to what git describe provides, which is + # vMAJ.MIN.PATCH.devNNNN+gGIT_REV + else: + dev_version = arr_info[0] + + pub_ver = f"{dev_version}.dev{arr_info[1]}" + local_ver = f"{pub_ver}+{arr_info[2]}" + return pub_ver, local_ver + + +# Implementations +def update(file_name, pattern, repl, dry_run=False): + update = [] + hit_counter = 0 + need_update = False + with open(file_name) as file: + for line in file: + result = re.findall(pattern, line) + if result: + assert len(result) == 1 + hit_counter += 1 + if result[0] != repl: + line = re.sub(pattern, repl, line) + need_update = True + print("%s: %s -> %s" % (file_name, result[0], repl)) + else: + print("%s: version is already %s" % (file_name, repl)) + + update.append(line) + if hit_counter != 1: + raise RuntimeError("Cannot find version in %s" % file_name) + + if need_update and not dry_run: + with open(file_name, "w") as output_file: + for line in update: + output_file.write(line) + + +def sync_version(pub_ver, local_ver, dry_run): + """Synchronize version.""" + # pyproject.toml + update( + os.path.join(PROJ_ROOT, "pyproject.toml"), + r"(?<=version = \")[.0-9a-z\+]+", + pub_ver, + dry_run, + ) + + +def main(): + logging.basicConfig(level=logging.INFO) + parser = argparse.ArgumentParser(description="Detect and synchronize version.") + parser.add_argument( + "--print-version", + action="store_true", + help="Print version to the command line. No changes is applied to files.", + ) + parser.add_argument( + "--git-describe", + action="store_true", + help="Use git describe to generate development version.", + ) + parser.add_argument("--dry-run", action="store_true") + pub_ver, local_ver = git_describe_version() + opt = parser.parse_args() + pub_ver, local_ver = None, None + if opt.git_describe: + pub_ver, local_ver = git_describe_version() + if pub_ver is None: + pub_ver = __version__ + if local_ver is None: + local_ver = __version__ + if opt.print_version: + print(local_ver) + else: + sync_version(pub_ver, local_ver, opt.dry_run) + + +if __name__ == "__main__": + main() diff --git a/web/Makefile b/web/Makefile new file mode 100644 index 0000000..9cc046d --- /dev/null +++ b/web/Makefile @@ -0,0 +1,51 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +TVM_ROOT=$(TVM_SOURCE_DIR) +MLC_LLM_ROOT=$(shell cd ..; pwd) + +INCLUDE_FLAGS = -I$(TVM_ROOT) -I$(TVM_ROOT)/include\ + -I$(TVM_ROOT)/3rdparty/dlpack/include\ + -I$(TVM_ROOT)/3rdparty/compiler-rt\ + -I$(MLC_LLM_ROOT)/3rdparty/tokenizers-cpp\ + -I$(MLC_LLM_ROOT)/3rdparty/tokenizers-cpp/include -I$(MLC_LLM_ROOT)/cpp + +.PHONY: clean all rmtypedep preparetest + +all: dist/wasm/mlc_wasm_runtime.wasm + +EMCC = emcc + +EMCC_CFLAGS = $(INCLUDE_FLAGS) -O3 -std=c++17 -Wno-ignored-attributes + +EMCC_LDFLAGS = --no-entry -s WASM_BIGINT=1 -s ALLOW_MEMORY_GROWTH=1 -s STANDALONE_WASM=1\ + -s ERROR_ON_UNDEFINED_SYMBOLS=0 + +dist/wasm/mlc_wasm_runtime.bc: emcc/mlc_wasm_runtime.cc + @mkdir -p $(@D) + $(EMCC) $(EMCC_CFLAGS) -c -MM -MT dist/wasm/mlc_wasm_runtime.bc emcc/mlc_wasm_runtime.cc >dist/wasm/mlc_wasm_runtime.d + $(EMCC) $(EMCC_CFLAGS) -emit-llvm -c -o dist/wasm/mlc_wasm_runtime.bc emcc/mlc_wasm_runtime.cc + +# Compile to wasm here so that errors can be caught earlier (rather than during export_library) +dist/wasm/mlc_wasm_runtime.wasm: dist/wasm/mlc_wasm_runtime.bc + @mkdir -p $(@D) + $(EMCC) $(EMCC_CFLAGS) -o dist/wasm/mlc_wasm_runtime.wasm $+ $(EMCC_LDFLAGS) + +clean: + @rm -rf dist/wasm lib + +-include dist/wasm/*.d diff --git a/web/README.md b/web/README.md new file mode 100644 index 0000000..9982e14 --- /dev/null +++ b/web/README.md @@ -0,0 +1,28 @@ + + + + + + + + + + + + + + + + + +# MLC-LLM WebAssembly Runtime + +This folder contains MLC-LLM WebAssembly Runtime. + +Please refer to https://llm.mlc.ai/docs/install/emcc.html. + +The main step is running `make` under this folder, a step included in `web/prep_emcc_deps.sh`. + +`make` creates `web/dist/wasm/mlc_wasm_runtime.bc`, which will be included in the model library wasm +when we compile the model. Thus during runtime, runtimes like WebLLM can directly reuse source +code from MLC-LLM. diff --git a/web/emcc/mlc_wasm_runtime.cc b/web/emcc/mlc_wasm_runtime.cc new file mode 100644 index 0000000..3dad33f --- /dev/null +++ b/web/emcc/mlc_wasm_runtime.cc @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/* + * \file mlc_wasm_runtime.cc + * \brief MLC wasm runtime library pack. + */ + +// configurations for tvm logging +#define TVM_LOG_STACK_TRACE 0 +#define TVM_LOG_DEBUG 0 +#define TVM_LOG_CUSTOMIZE 1 + +// Pass in COMPILE_MLC_WASM_RUNTIME so unsupported code would not be compiled in to the .bc file +#define COMPILE_MLC_WASM_RUNTIME 1 +#define __STDC_FORMAT_MACROS 1 +#define PICOJSON_USE_INT64 diff --git a/web/prep_emcc_deps.sh b/web/prep_emcc_deps.sh new file mode 100755 index 0000000..4f0f632 --- /dev/null +++ b/web/prep_emcc_deps.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# This file prepares all the necessary dependencies for the web build. +set -euxo pipefail + +emcc --version +npm --version + +TVM_SOURCE_DIR_SET="${TVM_SOURCE_DIR:-}" + +git submodule update --init --recursive + +CURR_DIR=`pwd` + +if [[ -z "${TVM_SOURCE_DIR_SET}" ]]; then + echo "Do not find TVM_SOURCE_DIR env variable, use 3rdparty/tvm". + echo "Make sure you set TVM_SOURCE_DIR in your env variable to use emcc build correctly" + export TVM_SOURCE_DIR="${TVM_SOURCE_DIR:-${CURR_DIR}/3rdparty/tvm}" +fi + +# Build mlc_wasm_runtime +cd web && make +cd - + +# Build tvm's web runtime +cd ${TVM_SOURCE_DIR}/web && TVM_HOME=${TVM_SOURCE_DIR} make +cd -