## 1. Project Overview & Quickstart (mlc-ai/web-llm) ## File: README.md ## Overview WebLLM is a high-performance in-browser LLM inference engine that brings language model inference directly onto web browsers with hardware acceleration. Everything runs inside the browser with no server support and is accelerated with WebGPU. WebLLM is **fully compatible with [OpenAI API](https://platform.openai.com/docs/api-reference/chat).** That is, you can use the same OpenAI API on **any open source models** locally, with functionalities including streaming, JSON-mode, function-calling (WIP), etc. We can bring a lot of fun opportunities to build AI assistants for everyone and enable privacy while enjoying GPU acceleration. You can use WebLLM as a base [npm package](https://www.npmjs.com/package/@mlc-ai/web-llm) and build your own web application on top of it by following the examples below. This project is a companion project of [MLC LLM](https://github.com/mlc-ai/mlc-llm), which enables universal deployment of LLM across hardware environments. ## Key Features - **In-Browser Inference**: WebLLM is a high-performance, in-browser language model inference engine that leverages WebGPU for hardware acceleration, enabling powerful LLM operations directly within web browsers without server-side processing. - [**Full OpenAI API Compatibility**](#full-openai-compatibility): Seamlessly integrate your app with WebLLM using OpenAI API with functionalities such as streaming, JSON-mode, logit-level control, seeding, and more. - **Structured JSON Generation**: WebLLM supports state-of-the-art JSON mode structured generation, implemented in the WebAssembly portion of the model library for optimal performance. Check [WebLLM JSON Playground](https://huggingface.co/spaces/mlc-ai/WebLLM-JSON-Playground) on HuggingFace to try generating JSON output with custom JSON schema. - [**Extensive Model Support**](#built-in-models): WebLLM natively supports a range of models including Llama 3, Phi 3, Gemma, Mistral, Qwen(通义千问), and many others, making it versatile for various AI tasks. For the complete supported model list, check [MLC Models](https://mlc.ai/models). - [**Custom Model Integration**](#custom-models): Easily integrate and deploy custom models in MLC format, allowing you to adapt WebLLM to specific needs and scenarios, enhancing flexibility in model deployment. - **Plug-and-Play Integration**: Easily integrate WebLLM into your projects using package managers like NPM and Yarn, or directly via CDN, complete with comprehensive [examples](./examples/) and a modular design for connecting with UI components. - **Streaming & Real-Time Interactions**: Supports streaming chat completions, allowing real-time output generation which enhances interactive applications like chatbots and virtual assistants. - **Web Worker & Service Worker Support**: Optimize UI performance and manage the lifecycle of models efficiently by offloading computations to separate worker threads or service workers. - **Chrome Extension Support**: Extend the functionality of web browsers through custom Chrome extensions using WebLLM, with examples available for building both basic and advanced extensions. ## Built-in Models Check the complete list of available models on [MLC Models](https://mlc.ai/models). WebLLM supports a subset of these available models and the list can be accessed at [`prebuiltAppConfig.model_list`](https://github.com/mlc-ai/web-llm/blob/main/src/config.ts#L293). Here are the primary families of models currently supported: - **Llama**: Llama 3, Llama 2, Hermes-2-Pro-Llama-3 - **Phi**: Phi 3, Phi 2, Phi 1.5 - **Gemma**: Gemma-2B - **Mistral**: Mistral-7B-v0.3, Hermes-2-Pro-Mistral-7B, NeuralHermes-2.5-Mistral-7B, OpenHermes-2.5-Mistral-7B - **Qwen (通义千问)**: Qwen2 0.5B, 1.5B, 7B If you need more models, [request a new model via opening an issue](https://github.com/mlc-ai/web-llm/issues/new/choose) or check [Custom Models](#custom-models) for how to compile and use your own models with WebLLM. ## Jumpstart with Examples Learn how to use WebLLM to integrate large language models into your application and generate chat completions through this simple Chatbot example: [](https://jsfiddle.net/neetnestor/4nmgvsa2/) [](https://codepen.io/neetnestor/pen/vYwgZaG) For an advanced example of a larger, more complicated project, check [WebLLM Chat](https://github.com/mlc-ai/web-llm-chat/blob/main/app/client/webllm.ts). More examples for different use cases are available in the [examples](./examples/) folder. ## Get Started WebLLM offers a minimalist and modular interface to access the chatbot in the browser. The package is designed in a modular way to hook to any of the UI components. ### Installation #### Package Manager ```sh # npm npm install @mlc-ai/web-llm # yarn yarn add @mlc-ai/web-llm # or pnpm pnpm install @mlc-ai/web-llm ``` Then import the module in your code. ```typescript // Import everything import * as webllm from "@mlc-ai/web-llm"; // Or only import what you need import { CreateMLCEngine } from "@mlc-ai/web-llm"; ``` #### CDN Delivery Thanks to [jsdelivr.com](https://www.jsdelivr.com/package/npm/@mlc-ai/web-llm), WebLLM can be imported directly through URL and work out-of-the-box on cloud development platforms like [jsfiddle.net](https://jsfiddle.net/), [Codepen.io](https://codepen.io/), and [Scribbler](https://scribbler.live): ```javascript import * as webllm from "https://esm.run/@mlc-ai/web-llm"; ``` It can also be dynamically imported as: ```javascript const webllm = await import("https://esm.run/@mlc-ai/web-llm"); ``` ### Create MLCEngine Most operations in WebLLM are invoked through the `MLCEngine` interface. You can create an `MLCEngine` instance and loading the model by calling the `CreateMLCEngine()` factory function. (Note that loading models requires downloading and it can take a significant amount of time for the very first run without caching previously. You should properly handle this asynchronous call.) ```typescript import { CreateMLCEngine } from "@mlc-ai/web-llm"; // Callback function to update model loading progress const initProgressCallback = (initProgress) => { console.log(initProgress); }; const selectedModel = "Llama-3.1-8B-Instruct-q4f32_1-MLC"; const engine = await CreateMLCEngine( selectedModel, { initProgressCallback: initProgressCallback }, // engineConfig ); ``` Under the hood, this factory function does the following steps for first creating an engine instance (synchronous) and then loading the model (asynchronous). You can also do them separately in your application. ```typescript import { MLCEngine } from "@mlc-ai/web-llm"; // This is a synchronous call that returns immediately const engine = new MLCEngine({ initProgressCallback: initProgressCallback, }); // This is an asynchronous call and can take a long time to finish await engine.reload(selectedModel); ``` ### Cache Backend Policy WebLLM supports four cache backends through `AppConfig.cacheBackend`: - `"cache"`: browser [Cache API](https://developer.mozilla.org/en-US/docs/Web/API/Cache) (default). - `"indexeddb"`: browser [IndexedDB](https://developer.mozilla.org/en-US/docs/Web/API/IndexedDB_API). - `"opfs"`: browser [Origin Private File System (OPFS)](https://developer.mozilla.org/en-US/docs/Web/API/File_System_API/Origin_private_file_system). - `"cross-origin"`: experimental Chrome [Cross-Origin Storage API](https://github.com/WICG/cross-origin-storage) extension backend. Install the [Cross-Origin Storage extension](https://chromewebstore.google.com/detail/cross-origin-storage/denpnpcgjgikjpoglpjefakmdcbmlgih) to use it. (If the extension isn't installed, WebLLM falls back to the default cache automatically.) Example: ```typescript import { CreateMLCEngine, prebuiltAppConfig } from "@mlc-ai/web-llm"; const appConfig = { ...prebuiltAppConfig, cacheBackend: "cross-origin" }; const engine = await CreateMLCEngine("Llama-3.1-8B-Instruct-q4f32_1-MLC", { appConfig, }); ``` Notes: - If `"opfs"` is selected in an environment without OPFS support, cache operations fail with an OPFS availability error. - When using `"opfs"`, `appConfig.opfsAccessMode` can be set to `"auto"` to use OPFS sync access handles where supported, or `"sync"` to require sync access handles. The default is `"async"`. - The `"cross-origin"` backend requires installing and enabling a compatible browser extension. - Cross-origin backend currently does not support programmatic tensor-cache deletion; clearing is extension-managed. ### Chat Completion After successfully initializing the engine, you can now invoke chat completions using OpenAI style chat APIs through the `engine.chat.completions` interface. For the full list of parameters and their descriptions, check [section below](#full-openai-compatibility) and [OpenAI API reference](https://platform.openai.com/docs/api-reference/chat/create). (Note: The `model` parameter is not supported and will be ignored here. Instead, call `CreateMLCEngine(model)` or `engine.reload(model)` instead as shown in the [Create MLCEngine](#create-mlcengine) above.) ```typescript const messages = [ { role: "system", content: "You are a helpful AI assistant." }, { role: "user", content: "Hello!" }, ]; const reply = await engine.chat.completions.create({ messages, }); console.log(reply.choices[0].message); console.log(reply.usage); ``` ### Streaming WebLLM also supports streaming chat completion generating. To use it, simply pass `stream: true` to the `engine.chat.completions.create` call. ```typescript const messages = [ { role: "system", content: "You are a helpful AI assistant." }, { role: "user", content: "Hello!" }, ]; // Chunks is an AsyncGenerator object const chunks = await engine.chat.completions.create({ messages, temperature: 1, stream: true, // <-- Enable streaming stream_options: { include_usage: true }, }); let reply = ""; for await (const chunk of chunks) { reply += chunk.choices[0]?.delta.content || ""; console.log(reply); if (chunk.usage) { console.log(chunk.usage); // only last chunk has usage } } const fullReply = await engine.getMessage(); console.log(fullReply); ``` ## Advanced Usage ### Using Workers You can put the heavy computation in a worker script to optimize your application performance. To do so, you need to: 1. Create a handler in the worker thread that communicates with the frontend while handling the requests. 2. Create a Worker Engine in your main application, which under the hood sends messages to the handler in the worker thread. For detailed implementations of different kinds of Workers, check the following sections. #### Dedicated Web Worker WebLLM comes with API support for WebWorker so you can hook the generation process into a separate worker thread so that the computing in the worker thread won't disrupt the UI. We create a handler in the worker thread that communicates with the frontend while handling the requests. ```typescript // worker.ts import { WebWorkerMLCEngineHandler } from "@mlc-ai/web-llm"; // A handler that resides in the worker thread const handler = new WebWorkerMLCEngineHandler(); self.onmessage = (msg: MessageEvent) => { handler.onmessage(msg); }; ``` In the main logic, we create a `WebWorkerMLCEngine` that implements the same `MLCEngineInterface`. The rest of the logic remains the same. ```typescript // main.ts import { CreateWebWorkerMLCEngine } from "@mlc-ai/web-llm"; async function main() { // Use a WebWorkerMLCEngine instead of MLCEngine here const engine = await CreateWebWorkerMLCEngine( new Worker(new URL("./worker.ts", import.meta.url), { type: "module", }), selectedModel, { initProgressCallback }, // engineConfig ); // everything else remains the same } ``` ### Use Service Worker WebLLM comes with API support for ServiceWorker so you can hook the generation process into a service worker to avoid reloading the model in every page visit and optimize your application's offline experience. (Note, Service Worker's life cycle is managed by the browser and can be killed any time without notifying the webapp. `ServiceWorkerMLCEngine` will try to keep the service worker thread alive by periodically sending heartbeat events, but your application should also include proper error handling. Check `keepAliveMs` and `missedHeatbeat` in [`ServiceWorkerMLCEngine`](https://github.com/mlc-ai/web-llm/blob/main/src/service_worker.ts#L234) for more details.) We create a handler in the worker thread that communicates with the frontend while handling the requests. ```typescript // sw.ts import { ServiceWorkerMLCEngineHandler } from "@mlc-ai/web-llm"; let handler: ServiceWorkerMLCEngineHandler; self.addEventListener("activate", function (event) { handler = new ServiceWorkerMLCEngineHandler(); console.log("Service Worker is ready"); }); ``` Then in the main logic, we register the service worker and create the engine using `CreateServiceWorkerMLCEngine` function. The rest of the logic remains the same. ```typescript // main.ts import { MLCEngineInterface, CreateServiceWorkerMLCEngine, } from "@mlc-ai/web-llm"; if ("serviceWorker" in navigator) { navigator.serviceWorker.register( new URL("sw.ts", import.meta.url), // worker script { type: "module" }, ); } const engine: MLCEngineInterface = await CreateServiceWorkerMLCEngine( selectedModel, { initProgressCallback }, // engineConfig ); ``` You can find a complete example on how to run WebLLM in service worker in [examples/service-worker](examples/service-worker/). ### Chrome Extension You can also find examples of building Chrome extension with WebLLM in [examples/chrome-extension](examples/chrome-extension/) and [examples/chrome-extension-webgpu-service-worker](examples/chrome-extension-webgpu-service-worker/). The latter one leverages service worker, so the extension is persistent in the background. Additionally, you can explore another full project of a Chrome extension, WebLLM Assistant, which leverages WebLLM [here](https://github.com/mlc-ai/web-llm-assistant). ## Full OpenAI Compatibility WebLLM is designed to be fully compatible with [OpenAI API](https://platform.openai.com/docs/api-reference/chat). Thus, besides building a simple chatbot, you can also have the following functionalities with WebLLM: - [streaming](examples/streaming): return output as chunks in real-time in the form of an AsyncGenerator - [json-mode](examples/json-mode): efficiently ensure output is in JSON format, see [OpenAI Reference](https://platform.openai.com/docs/guides/text-generation/chat-completions-api) for more. - [seed-to-reproduce](examples/seed-to-reproduce): use seeding to ensure a reproducible output with fields `seed`. - [function-calling](examples/function-calling) (WIP): function calling with fields `tools` and `tool_choice` (with preliminary support); or manual function calling without `tools` or `tool_choice` (keeps the most flexibility). ## Integrity Verification WebLLM supports optional integrity verification for model artifacts using [SRI (Subresource Integrity)](https://developer.mozilla.org/en-US/docs/Web/Security/Subresource_Integrity) hashes. When the `integrity` field is set on a `ModelRecord`, WebLLM will verify the downloaded config, WASM, and tokenizer files against the provided hashes before loading. ```typescript import { CreateMLCEngine } from "@mlc-ai/web-llm"; const appConfig = { model_list: [ { model: "https://huggingface.co/mlc-ai/Llama-3.2-1B-Instruct-q4f16_1-MLC", model_id: "Llama-3.2-1B-Instruct-q4f16_1-MLC", model_lib: "https://raw.githubusercontent.com/user/model-libs/main/model.wasm", integrity: { config: "sha256-", model_lib: "sha256-", tokenizer: { "tokenizer.json": "sha256-", }, onFailure: "error", // "error" (default) throws IntegrityError, "warn" logs and continues }, }, ], }; const engine = await CreateMLCEngine("Llama-3.2-1B-Instruct-q4f16_1-MLC", { appConfig, }); ``` You can generate SRI hashes for model files with: ```bash # SHA-256 openssl dgst -sha256 -binary | openssl base64 -A | sed 's/^/sha256-/' # SHA-384 openssl dgst -sha384 -binary | openssl base64 -A | sed 's/^/sha384-/' # SHA-512 openssl dgst -sha512 -binary | openssl base64 -A | sed 's/^/sha512-/' ``` > The `openssl` commands require a Unix-like shell (macOS/Linux). On Windows, run `openssl` via [Git Bash](https://gitforwindows.org/) or [WSL](https://learn.microsoft.com/en-us/windows/wsl/). If a hash does not match, an `IntegrityError` is thrown (or a warning is logged when `onFailure: "warn"`). All fields in `integrity` are optional — only specified artifacts will be verified. When the `integrity` field is omitted entirely, WebLLM behaves exactly as before (no verification). See the [integrity-verification example](examples/integrity-verification/) for a complete working demo. ## Custom Models WebLLM works as a companion project of [MLC LLM](https://github.com/mlc-ai/mlc-llm) and it supports custom models in MLC format. It reuses the model artifact and builds the flow of MLC LLM. To compile and use your own models with WebLLM, please check out [MLC LLM document](https://llm.mlc.ai/docs/deploy/webllm.html) on how to compile and deploy new model weights and libraries to WebLLM. Here, we go over the high-level idea. There are two elements of the WebLLM package that enable new models and weight variants. - `model`: Contains a URL to model artifacts, such as weights and meta-data. - `model_lib`: A URL to the web assembly library (i.e. wasm file) that contains the executables to accelerate the model computations. Both are customizable in the WebLLM. ```typescript import { CreateMLCEngine } from "@mlc-ai/web-llm"; async main() { const appConfig = { "model_list": [ { "model": "/url/to/my/llama", "model_id": "MyLlama-3b-v1-q4f32_0", "model_lib": "/url/to/myllama3b.wasm", } ], }; // override default const chatOpts = { "repetition_penalty": 1.01 }; // load a prebuilt model // with a chat option override and app config // under the hood, it will load the model from myLlamaUrl // and cache it in the browser cache // The chat will also load the model library from "/url/to/myllama3b.wasm", // assuming that it is compatible to the model in myLlamaUrl. const engine = await CreateMLCEngine( "MyLlama-3b-v1-q4f32_0", { appConfig }, // engineConfig chatOpts, ); } ``` In many cases, we only want to supply the model weight variant, but not necessarily a new model (e.g. `NeuralHermes-Mistral` can reuse `Mistral`'s model library). For examples of how a model library can be shared by different model variants, see `webllm.prebuiltAppConfig`. ## Build WebLLM Package From Source NOTE: you don't need to build from source unless you would like to modify the WebLLM package. To use the npm, simply follow [Get Started](#get-started) or any of the [examples](examples) instead. To build from source, simply run: ```bash npm install npm run build ``` Then, to test the effects of your code change in an example, inside `examples/get-started/package.json`, change from `"@mlc-ai/web-llm": "^0.2.84"` to `"@mlc-ai/web-llm": ../..`. Then run: ```bash cd examples/get-started npm install npm start ``` Note that sometimes you would need to switch between `file:../..` and `../..` to trigger npm to recognize new changes. In the worst case, you can run: ```bash cd examples/get-started rm -rf node_modules dist package-lock.json .parcel-cache npm install npm start ``` ### In case you need to build TVMjs from source WebLLM's runtime largely depends on TVMjs: https://github.com/apache/tvm/tree/main/web While it is also available as an npm package: https://www.npmjs.com/package/@mlc-ai/web-runtime, you can build it from source if needed by following the steps below. 1. Install [emscripten](https://emscripten.org). It is an LLVM-based compiler that compiles C/C++ source code to WebAssembly. - Follow the [installation instruction](https://emscripten.org/docs/getting_started/downloads.html#installation-instructions-using-the-emsdk-recommended) to install the latest emsdk. - Source `emsdk_env.sh` by `source path/to/emsdk_env.sh`, so that `emcc` is reachable from PATH and the command `emcc` works. We can verify the successful installation by trying out `emcc` terminal. 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 ``` Init error, LinkError: WebAssembly.instantiate(): Import #6 module="wasi_snapshot_preview1" function="proc_exit": function import requires a callable ``` 2. In `./package.json`, change from `"@mlc-ai/web-runtime": "0.18.0-dev2",` to `"@mlc-ai/web-runtime": "file:./tvm_home/web",`. 3. Setup necessary environment Prepare all the necessary dependencies for web build: ```shell ./scripts/prep_deps.sh ``` In this step, if `$TVM_SOURCE_DIR` is not defined in the environment, we will execute the following line to build `tvmjs` dependency: ```shell git clone https://github.com/mlc-ai/relax 3rdparty/tvm-unity --recursive ``` This clones the current HEAD of `mlc-ai/relax`. However, it may not always be the correct branch or commit to clone. To build a specific npm version from source, refer to the version bump PR, which states which branch (i.e. `mlc-ai/relax` or `apache/tvm`) and which commit the current WebLLM version depends on. For instance, version 0.2.52, according to its version bump PR https://github.com/mlc-ai/web-llm/pull/521, is built by checking out the following commit https://github.com/apache/tvm/commit/e6476847753c80e054719ac47bc2091c888418b6 in `apache/tvm`, rather than the HEAD of `mlc-ai/relax`. Besides, `--recursive` is necessary and important. Otherwise, you may encounter errors like `fatal error: 'dlpack/dlpack.h' file not found`. 4. Build WebLLM Package ```shell npm run build ``` 5. Validate some of the sub-packages You can then go to the subfolders in [examples](examples) to validate some of the sub-packages. We use Parcelv2 for bundling. Although Parcel is not very good at tracking parent directory changes sometimes. When you make a change in the WebLLM package, try to edit the `package.json` of the subfolder and save it, which will trigger Parcel to rebuild. ## Links - [Demo App: WebLLM Chat](https://chat.webllm.ai/) - If you want to run LLM on native runtime, check out [MLC-LLM](https://github.com/mlc-ai/mlc-llm) - You might also be interested in [Web Stable Diffusion](https://github.com/mlc-ai/web-stable-diffusion/). ## Acknowledgement This project is initiated by members from CMU Catalyst, UW SAMPL, SJTU, OctoML, and the MLC community. We would love to continue developing and supporting the open-source ML community. This project is only possible thanks to the shoulders open-source ecosystems that we stand on. We want to thank the Apache TVM community and developers of the TVM Unity effort. The open-source ML community members made these models publicly available. PyTorch and Hugging Face communities make these models accessible. We would like to thank the teams behind Vicuna, SentencePiece, LLaMA, and Alpaca. We also would like to thank the WebAssembly, Emscripten, and WebGPU communities. Finally, thanks to Dawn and WebGPU developers. ## Citation If you find this project to be useful, please cite: ``` @misc{ruan2026webllmhighperformanceinbrowserllm, title={WebLLM: A High-Performance In-Browser LLM Inference Engine}, author={Charlie F. Ruan and Yucheng Qin and Akaash R. Parthasarathy and Xun Zhou and Ruihang Lai and Hongyi Jin and Yixin Dong and Bohan Hou and Meng-Shiun Yu and Yiyan Zhai and Sudeep Agarwal and Hangrui Cao and Siyuan Feng and Tianqi Chen}, year={2026}, eprint={2412.15803}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2412.15803}, } ``` ## Contributors [⬆ Back to Top ⬆](#top) --- ## File: docs/developer/add_models.rst Adding Models ============= WebLLM allows you to compile custom language models using `MLC-LLM `_ and then serve the compiled model through WebLLM. For instructions on how to compile and add custom models to WebLLM, please refer to the `MLC-LLM documentation `_. --- ## File: docs/developer/building_from_source.rst Building From Source ==================== Clone the Repository --------------------- .. code-block:: bash git clone https://github.com/mlc-ai/web-llm.git cd web-llm Install Dependencies --------------------- .. code-block:: bash npm install Build the Project ----------------- .. code-block:: bash npm run build Test Changes ------------ To test your changes, you can reuse an existing example or create a new example that specifically tests the new functionality you wish to provide. To test the effects of your code change in an example, inside ``examples//package.json``, change ``"@mlc-ai/web-llm": "^0.2.xx"`` to ``"@mlc-ai/web-llm": "../.."`` to let it reference your local code. Note that sometimes you may need to switch between ``"file:../.."`` and ``"../.."`` to trigger npm to recognize new changes. .. code-block:: bash cd examples/ # Modify package.json as described npm install npm start --- ## File: docs/user/advanced_usage.rst Advanced Use Cases ================== Using Workers ------------- You can put the heavy computation in a worker script to optimize your application performance. To do so, you need to: Create a handler in the worker thread that communicates with the frontend while handling the requests. Create a worker engine in your main application that sends messages to the handler in the worker thread under the hood. For detailed implementations of different kinds of workers, look at the following sections. Using Web Workers ^^^^^^^^^^^^^^^^^ WebLLM comes with API support for `Web Workers `_ so you can offload the computation-heavy generation work into a separate worker thread. WebLLM has implemented cross-thread communication through messages under the hood, so manual implementation is not required. In the worker script, import and instantiate a ``WebWorkerMLCEngineHandler``, which handles communication with other scripts and processes incoming requests. .. code-block:: typescript // worker.ts import { WebWorkerMLCEngineHandler } from "@mlc-ai/web-llm"; const handler = new WebWorkerMLCEngineHandler(); self.onmessage = (msg: MessageEvent) => { handler.onmessage(msg); }; In the main script, import and instantiate a ``WebWorkerMLCEngine`` that implements the same ``MLCEngineInterface`` and exposes the same APIs. Then, simply use it as you would a normal ``MLCEngine``. .. code-block:: typescript import { CreateWebWorkerMLCEngine } from "@mlc-ai/web-llm"; async function runWorker() { const engine = await CreateWebWorkerMLCEngine( new Worker(new URL("./worker.ts", import.meta.url), { type: "module" }), "Llama-3.1-8B-Instruct" ); const messages = [{ role: "user", content: "How does WebLLM use workers?" }]; const reply = await engine.chat.completions.create({ messages }); console.log(reply.choices[0].message.content); } runWorker(); Under the hood, ``WebWorkerMLCEngine`` does **not** perform any computation. It translates all calls into messages and sends them to the ``WebWorkerMLCEngineHandler`` for processing. The worker thread receives these messages and processes the actual computation using a hidden engine, and returns the result to the main thread using messages. Service Workers ^^^^^^^^^^^^^^^ WebLLM also supports offloading computation using `Service Workers `_. This allows you to avoid reloading the model between page refreshes and optimize your application's offline experience. (Note, the lifecycle of a Service Worker is managed by the browser and can be killed any time without notifying the web application. WebLLM's ``ServiceWorkerMLCEngine`` attempts to keep the service worker thread alive by periodically sending heartbeat events. However, the script could still be killed at any time by Chrome, and your application should include proper error handling. Check `keepAliveMs` and `missedHeartbeat` in `ServiceWorkerMLCEngine `_ for more details.) In the worker script, import and instantiate ``ServiceWorkerMLCEngineHandler``, which handles communication with page scripts and processes incoming requests. .. code-block:: typescript // sw.ts import { ServiceWorkerMLCEngineHandler } from "@mlc-ai/web-llm"; self.addEventListener("activate", () => { const handler = new ServiceWorkerMLCEngineHandler(); console.log("Service Worker activated!"); }); Then, in the main page script, register the service worker and instantiate the engine using the ``CreateServiceWorkerMLCEngine`` factory function that implements the same ``MLCEngineInterface`` and exposes the same APIs. Then, simply use it as you would a normal ``MLCEngine``. .. code-block:: typescript // main.ts import { MLCEngineInterface, CreateServiceWorkerMLCEngine } from "@mlc-ai/web-llm"; if ("serviceWorker" in navigator) { navigator.serviceWorker.register( new URL("sw.ts", import.meta.url), // worker script { type: "module" }, ); } const engine: MLCEngineInterface = await CreateServiceWorkerMLCEngine( selectedModel, { initProgressCallback }, // engineConfig ); Similar to the ``WebWorkerMLCEngine`` above, the ``ServiceWorkerMLCEngine`` is also a proxy and does not perform any actual computation. Instead, it forwards all calls to the service worker thread and receives the result through messages. Chrome Extension ---------------- WebLLM can be used in Chrome extensions to empower local LLM inference. You can find examples of building Chrome extension using WebLLM in `examples/chrome-extension `_ and `examples/chrome-extension-webgpu-service-worker `_. The latter leverages Service Worker, so the extension is persistent in the background. Additionally, we have a full Chrome extension project, `WebLLM Assistant `_, which leverages WebLLM to provide a personal web browsing copilot assistant experience. Feel free to check it out and contribute if you are interested. Additional Customization ------------------------ Using IndexedDB Cache ^^^^^^^^^^^^^^^^^^^^^ By default, WebLLM caches model artifacts using the `Cache API `_ for faster subsequent model loads. You can alternatively use `IndexedDB caching `_ by setting ``appConfig.cacheBackend = "indexeddb"``. When changing only the cache backend, preserve the prebuilt model list by spreading ``prebuiltAppConfig``. .. code-block:: typescript import { AppConfig, CreateMLCEngine, prebuiltAppConfig } from "@mlc-ai/web-llm"; const appConfig: AppConfig = { ...prebuiltAppConfig, cacheBackend: "indexeddb", }; const engine = await CreateMLCEngine("Llama-3.1-8B-Instruct-q4f32_1-MLC", { appConfig, }); Using Cross-Origin Storage Cache ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ WebLLM also supports caching model artifacts across different origins using the experimental Cross-Origin Storage API. You can enable this cache backend by setting ``appConfig.cacheBackend = "cross-origin"``. For users with the `Cross-Origin Storage browser extension `_ installed, resources will then be cached and shared across origins. This means two independent apps opted into this cache backend using the same AI model will download and cache the required resources only once. See the `cache usage example `_ for more details. If Cross-Origin Storage isn't available, WebLLM will automatically fall back to using the default cache. .. code-block:: typescript import { AppConfig, CreateMLCEngine, prebuiltAppConfig } from "@mlc-ai/web-llm"; const appConfig: AppConfig = { ...prebuiltAppConfig, cacheBackend: "cross-origin", }; const engine = await CreateMLCEngine("Llama-3.1-8B-Instruct-q4f32_1-MLC", { appConfig, }); Customizing Token Behavior ^^^^^^^^^^^^^^^^^^^^^^^^^^ You can modify `logit_bias` in `GenerationConfig` to control token likelihood. Setting a token's bias to a positive value increases its likelihood of being generated, while a negative value decreases it. A large negative value (e.g., -100) can effectively prevent the token from being generated. .. code-block:: typescript const messages = [ { role: "user", content: "Describe WebLLM in detail." }, ]; const response = await engine.chatCompletion({ messages, logit_bias: { "50256": -100 }, // Example: Prevent specific token generation }); --- ## File: docs/user/api_reference.rst .. _api-reference: WebLLM API Reference ==================== The ``MLCEngine`` class is the core interface of WebLLM. It enables model loading, chat completions, embeddings, and other operations. Below, we document its methods, along with the associated configuration interfaces. Interfaces ---------- The following interfaces are used as parameters or configurations within ``MLCEngine`` methods. They are linked to their respective methods for reference. MLCEngineConfig ^^^^^^^^^^^^^^^ Optional configurations for ``CreateMLCEngine()`` and ``CreateWebWorkerMLCEngine()``. - **Fields**: - ``appConfig``: Configure the app, including the list of models and whether to use IndexedDB cache. - ``initProgressCallback``: A callback for showing model loading progress. - ``logitProcessorRegistry``: A registry for stateful logit processors (see ``webllm.LogitProcessor``). - **Usage**: - ``appConfig``: Contains application-specific settings, including: - Model configurations. - IndexedDB caching preferences. - ``initProgressCallback``: Allows developers to visualize model loading progress by implementing a callback. - ``logitProcessorRegistry``: A ``Map`` object for registering custom logit processors. Only applies to ``MLCEngine``. .. note:: All fields are optional, and ``logitProcessorRegistry`` is only used in ``MLCEngine``. Example: .. code-block:: typescript const engine = await CreateMLCEngine("Llama-3.1-8B-Instruct", { appConfig: { /* app-specific config */ }, initProgressCallback: (progress) => console.log(progress), }); GenerationConfig ^^^^^^^^^^^^^^^^ Configurations for a single generation task, primarily used in chat completions. - **Fields**: - ``repetition_penalty``, ``ignore_eos``: Parameters specific to MLC models. - ``top_p``, ``temperature``, ``max_tokens``, ``stop``: Common parameters shared with OpenAI APIs. - ``frequency_penalty``, ``presence_penalty``: Tune repetition behavior following OpenAI semantics. - ``logit_bias``, ``n``, ``logprobs``, ``top_logprobs``: Advanced sampling controls. - ``response_format``, ``enable_thinking``, ``enable_latency_breakdown``: Additional OpenAI-style request features. - **Usage**: - Fields like ``repetition_penalty`` and ``ignore_eos`` give explicit control over repetition handling and whether the model stops at the EOS token, respectively. - Common parameters shared with OpenAI APIs (e.g., ``temperature``, ``top_p``) ensure compatibility while still falling back to the values configured during ``MLCEngine.reload()`` when omitted. - ``frequency_penalty`` and ``presence_penalty`` mirror OpenAI's bounds ``[-2, 2]``; providing only one will default the other to ``0``. - ``response_format`` (for JSON or other schema outputs), ``enable_thinking``, and ``enable_latency_breakdown`` pass through directly to the engine and surface enhanced telemetry or structured responses when the underlying model supports them. Example: .. code-block:: typescript const messages = [ { role: "system", content: "You are a helpful assistant." }, { role: "user", content: "Explain WebLLM." }, ]; const response = await engine.chatCompletion({ messages, top_p: 0.9, temperature: 0.8, max_tokens: 150, }); ChatConfig ^^^^^^^^^^ Model's baseline configuration loaded from ``mlc-chat-config.json`` when ``MLCEngine.reload()`` runs. ``ChatOptions`` (and therefore the ``chatOpts`` argument to ``reload``) can override any subset of these fields. - **Fields** (subset): - ``tokenizer_files``, ``tokenizer_info``: Files and parameters required to initialize the tokenizer. - ``conv_template``, ``conv_config``: Conversation templates that define prompts, separators, and role formatting. - ``context_window_size``, ``sliding_window_size``, ``attention_sink_size``: KV-cache and memory settings. - Default generation knobs such as ``repetition_penalty``, ``frequency_penalty``, ``presence_penalty``, ``top_p``, and ``temperature``. - **Usage**: - Loaded automatically for each model; provides defaults that ``GenerationConfig`` falls back to when fields are omitted. - Override selected values per model load by supplying ``chatOpts`` (``Partial``) to ``MLCEngine.reload()``. Example: .. code-block:: typescript await engine.reload("Llama-3.1-8B-Instruct", { temperature: 0.7, repetition_penalty: 1.1, context_window_size: 4096, }); ChatCompletionRequest ^^^^^^^^^^^^^^^^^^^^^ Defines the structure for chat completion requests. - **Base Interface**: ``ChatCompletionRequestBase`` - Contains parameters such as ``messages``, ``stream``, ``frequency_penalty``, and ``presence_penalty``. - **Sub-interfaces**: - ``ChatCompletionRequestNonStreaming``: For non-streaming completions. - ``ChatCompletionRequestStreaming``: For streaming completions. - **Usage**: - Combines settings from ``GenerationConfig`` and ``ChatCompletionRequestBase`` to provide complete control over chat behavior. - The ``stream`` parameter enables streaming responses, improving interactivity in conversational agents. - The ``logit_bias`` feature allows controlling token generation probabilities, providing a mechanism to restrict or encourage specific outputs. Example: .. code-block:: typescript const response = await engine.chatCompletion({ messages: [ { role: "user", content: "Tell me about WebLLM." }, ], stream: true, }); Model Loading ------------- ``MLCEngine.reload(modelId: string | string[], chatOpts?: ChatOptions | ChatOptions[]): Promise`` Loads the specified model(s) into the engine. Uses ``MLCEngineConfig`` during initialization. - Parameters: - ``modelId``: Identifier(s) for the model(s) to load. - ``chatOpts``: Configuration for generation (see ``ChatConfig``). Example: .. code-block:: typescript await engine.reload(["Llama-3.1-8B", "Gemma-2B"], [ { temperature: 0.7 }, { top_p: 0.9 }, ]); ``MLCEngine.unload(): Promise`` Unloads all loaded models and clears their associated configurations. Example: .. code-block:: typescript await engine.unload(); --- Chat Completions ---------------- ``MLCEngine.chat.completions.create(request: ChatCompletionRequest): Promise>`` Generates chat-based completions using a specified request configuration. - Parameters: - ``request``: A ``ChatCompletionRequest`` instance. Example: .. code-block:: typescript const response = await engine.chat.completions.create({ messages: [ { role: "system", content: "You are a helpful AI assistant." }, { role: "user", content: "What is WebLLM?" }, ], temperature: 0.8, stream: false, }); --- Utility Methods ^^^^^^^^^^^^^^^ ``MLCEngine.getMessage(modelId?: string): Promise`` Retrieves the current output message from the specified model. - Parameters: - ``modelId``: (Optional) Identifier of model to query. Omitting modelId only works when the engine currently has a single model loaded. ``MLCEngine.resetChat(keepStats?: boolean, modelId?: string): Promise`` Resets the chat history and optionally retains usage statistics. - Parameters: - ``keepStats``: (Optional) If true, retains usage statistics. - ``modelId``: (Optional) Identifier of the model to reset. Omitting modelId only works when the engine currently has a single model loaded. GPU Information ---------------- The following methods provide detailed information about the GPU used for WebLLM computations. ``MLCEngine.getGPUVendor(): Promise`` Retrieves the vendor name of the GPU used for computations. This is useful for understanding hardware capabilities during inference. - **Returns**: A string indicating the GPU vendor (e.g., "Intel", "NVIDIA"). Example: .. code-block:: typescript const gpuVendor = await engine.getGPUVendor(); console.log(``GPU Vendor: ${gpuVendor}``); ``MLCEngine.getMaxStorageBufferBindingSize(): Promise`` Returns the maximum storage buffer size supported by the GPU. This is important when working with larger models that require significant memory for processing. - **Returns**: A number representing the maximum size in bytes. Example: .. code-block:: typescript const maxBufferSize = await engine.getMaxStorageBufferBindingSize(); console.log(``Max Storage Buffer Binding Size: ${maxBufferSize}``); --- ## File: docs/user/basic_usage.rst Basic Usage ================ Model Records in WebLLM ----------------------- Each of the model available WebLLM is registered as an instance of ``ModelRecord`` and can be accessed at `webllm.prebuiltAppConfig.model_list `__. Creating an MLCEngine --------------------- WebLLM APIs are exposed through the ``MLCEngine`` interface. You can create an ``MLCEngine`` instance and load the model by calling the CreateMLCEngine() factory function. (Note that loading models requires downloading and it can take a significant amount of time for the very first run without previous caching. You should properly handle this asynchronous call.) ``MLCEngine`` can be instantiated in two ways: 1. Using the factory function ``CreateMLCEngine``. 2. Instantiating the ``MLCEngine`` class directly and using ``reload()`` to load models. .. code-block:: typescript import { CreateMLCEngine, MLCEngine } from "@mlc-ai/web-llm"; // Initialize with a progress callback const initProgressCallback = (progress) => { console.log("Model loading progress:", progress); }; // Using CreateMLCEngine const engine = await CreateMLCEngine("Llama-3.1-8B-Instruct", { initProgressCallback }); // Direct instantiation const engineInstance = new MLCEngine({ initProgressCallback }); await engineInstance.reload("Llama-3.1-8B-Instruct"); Under the hood, this factory function ``CreateMLCEngine`` does the following steps for first creating an engine instance (synchronous) and then loading the model (asynchronous). You can also do them separately in your application. .. code-block:: typescript import { MLCEngine } from "@mlc-ai/web-llm"; // This is a synchronous call that returns immediately const engine = new MLCEngine({ initProgressCallback: initProgressCallback }); // This is an asynchronous call and can take a long time to finish await engine.reload(selectedModel); Chat Completion --------------- Chat completions can be invoked using OpenAI style chat APIs through the ``engine.chat.completions`` interface of an initialized ``MLCEngine``. For the full list of parameters and their descriptions, check :ref:`api-reference` for full list of parameters. (Note: Since the model is determined during ``MLCEngine`` instantiation, the ``model`` parameter is not supported and will be **ignored**. Instead, call ``CreateMLCEngine(model)`` or ``engine.reload(model)`` to reinitialize the engine to use a specific model.) .. code-block:: typescript const messages = [ { role: "system", content: "You are a helpful AI assistant." }, { role: "user", content: "Hello!" } ]; const reply = await engine.chat.completions.create({ messages, }); console.log(reply.choices[0].message); console.log(reply.usage); Streaming Chat Completion ------------------------- Streaming chat completion could be enabled by passsing ``stream: true`` parameter to the `engine.chat.completions.create` call configuration. Check :ref:`api-reference` for full list of parameters. .. code-block:: typescript const messages = [ { role: "system", content: "You are a helpful AI assistant." }, { role: "user", content: "Hello!" }, ] // chunks is an AsyncGenerator object const chunks = await engine.chat.completions.create({ messages, temperature: 1, stream: true, // <-- Enable streaming stream_options: { include_usage: true }, }); let reply = ""; for await (const chunk of chunks) { reply += chunk.choices[0]?.delta.content || ""; console.log(reply); if (chunk.usage) { console.log(chunk.usage); // only last chunk has usage } } const fullReply = await engine.getMessage(); console.log(fullReply); Chatbot Examples ---------------- Learn how to use WebLLM to integrate large language models into your applications and generate chat completions through this simple Chatbot example: - `Example in JSFiddle `_ - `Example in CodePen `_ For an advanced example of a larger, more complicated project, look at `WebLLM Chat `_. More examples for different use cases are available in the `WebLLM examples folder `_. --- ## File: docs/user/get_started.rst Getting Started with WebLLM =========================== This guide will help you set up WebLLM in your project, install necessary dependencies, and verify your setup. WebLLM Chat ----------- If you want to experience AI Chat supported by local LLM inference and understand how WebLLM works, 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. Installation ------------ WebLLM offers a minimalist and modular interface to access the chatbot in the browser. The package is designed in a modular way to hook to any of the UI components. WebLLM is available as an `npm package `_ and is also CDN-delivered. Therefore, you can install WebLLM using Node.js package managers like npm, yarn, or pnpm, or directly import the pacakge via CDN. Using Package Managers ^^^^^^^^^^^^^^^^^^^^^^ Install WebLLM via your preferred package manager: .. code-block:: bash # npm npm install @mlc-ai/web-llm # yarn yarn add @mlc-ai/web-llm # pnpm pnpm install @mlc-ai/web-llm Import WebLLM into your project: .. code-block:: javascript // Import everything import * as webllm from "@mlc-ai/web-llm"; // Or only import what you need import { CreateMLCEngine } from "@mlc-ai/web-llm"; Using CDN ^^^^^^^^^ Thanks to `jsdelivr.com `_, WebLLM can be imported directly through URL and work out-of-the-box on cloud development platforms like `jsfiddle.net `_, `Codepen.io `_, and `Scribbler `_: .. code-block:: javascript import * as webllm from "https://esm.run/@mlc-ai/web-llm"; This method is especially useful for online environments like CodePen, JSFiddle, or local experiments. Verifying Installation ^^^^^^^^^^^^^^^^^^^^^^ Run the following script to verify the installation: .. code-block:: javascript import { CreateMLCEngine } from "@mlc-ai/web-llm"; console.log("WebLLM loaded successfully!"); Online IDE Sandbox ------------------ Instead of setting WebLLM locally, you can also try it on online Javascript IDE sandboxes like: - `Example in JSFiddle `_ - `Example in CodePen `_ --- ## File: docs/index.rst 👋 Welcome to WebLLM ==================== `GitHub `_ | `WebLLM Chat `_ | `NPM `_ | `Discord `_ WebLLM is a high-performance in-browser language model inference engine that brings large language models (LLMs) to web browsers with hardware acceleration. With WebGPU support, it allows developers to build AI-powered applications directly within the browser environment, removing the need for server-side processing and ensuring privacy. 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. Key Features ------------ - 🌐 In-Browser Inference: Run LLMs directly in the browser - 🚀 WebGPU Acceleration: Leverage hardware acceleration for optimal performance - 🔄 OpenAI API Compatibility: Seamless integration with standard AI workflows - 📦 Multiple Model Support: Works with Llama, Phi, Gemma, Mistral, and more Start exploring WebLLM by `chatting with WebLLM Chat `_, and start building webapps with high-performance local LLM inference with the following guides and tutorials. .. toctree:: :maxdepth: 2 :caption: User Guide user/get_started.rst user/basic_usage.rst user/advanced_usage.rst user/api_reference.rst .. toctree:: :maxdepth: 2 :caption: Developer Guide developer/building_from_source.rst developer/add_models.rst --- ## File: docs/README.md # WebLLM Documentation The documentation was built upon [Sphinx](https://www.sphinx-doc.org/en/master/). ## Dependencies Run the following command in this directory to install dependencies first: ```bash pip3 install -r requirements.txt ``` ## Build the Documentation Then you can build the documentation by running: ```bash make html ``` ## View the Documentation Run the following command to start a simple HTTP server: ```bash cd _build/html python3 -m http.server ``` Then you can view the documentation in your browser at `http://localhost:8000` (the port can be customized by appending ` -p PORT_NUMBER` in the python command above). --- ## File: examples/abort-reload/README.md # WebLLM Get Started App This folder provides a demo for cancelling model fetching after calling `engine.reload()`. ```bash npm install npm start ``` Note if you would like to hack WebLLM core package. You can change web-llm dependencies as `"file:../.."`, and follow the build from source instruction in the project to build webllm locally. This option is only recommended if you would like to hack WebLLM core package. ## 2. Official Technical Reference & Guides (mlc-ai/docs) ## File: README.md # TVM Unity Documentation The documentation was built upon [Sphinx](https://www.sphinx-doc.org/en/master/). ## Dependencies Run the following command in this directory to install dependencies first: ```bash pip install -r requirements.txt ``` ## Build the Documentation Then you can build the documentation by running: ```bash make html ``` ## View the Documentation Run the following command to start a simple HTTP server: ```bash python -m http.server -d _build/html 8000 ``` Then you can view the documentation in your browser at `http://localhost:8000` (the port can be customized by changing the port number). --- ## File: docs/deep_dive/tensor_ir/abstraction.rst .. 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. .. _tir-abstraction: Tensor Program Abstraction -------------------------- Before we dive into the details of TensorIR, let's first introduce what is a primitive tensor function. Primitive tensor functions are functions that correspond to a single "unit" of computational operation. For example, a convolution operation can be a primitive tensor function, and a fused convolution + relu operation can also be a primitive tensor function. Usually, a typical abstraction for primitive tensor function implementation contains the following elements: multi-dimensional buffers, loop nests that drive the tensor computations, and finally, the compute statements themselves. .. code:: python from tvm.script import tir as T @T.prim_func def main( A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32"), C: T.Buffer((128,), "float32"), ) -> None: for i in range(128): with T.block("C"): vi = T.axis.spatial(128, i) C[vi] = A[vi] + B[vi] Key Elements of Tensor Programs ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ The demonstrated primitive tensor function calculates the element-wise sum of two vectors. The function: - Accepts three **multi-dimensional buffers** as parameters, and generates one **multi-dimensional buffer** as output. - Incorporates a solitary **loop nest** ``i`` that facilitates the computation. - Features a singular **compute statement** that calculates the element-wise sum of the two vectors. Extra Structure in TensorIR ~~~~~~~~~~~~~~~~~~~~~~~~~~~ Crucially, we are unable to execute arbitrary transformations on the program, as certain computations rely on the loop's sequence. Fortunately, the majority of primitive tensor functions we focus on possess favorable properties, such as independence among loop iterations. For instance, the aforementioned program includes block and iteration annotations: - The **block annotation** ``with T.block("C")`` signifies that the block is the fundamental computation unit designated for scheduling. A block may encompass a single computation statement, multiple computation statements with loops, or opaque intrinsics such as Tensor Core instructions. - The **iteration annotation** ``T.axis.spatial``, indicating that variable ``vi`` is mapped to ``i``, and all iterations are independent. While this information isn't crucial for *executing* the specific program, it proves useful when transforming the program. Consequently, we can confidently parallelize or reorder loops associated with ``vi``, provided we traverse all the index elements from 0 to 128. --- ## File: docs/deep_dive/tensor_ir/index.rst .. 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. .. _tensor-ir: TensorIR ======== TensorIR is one of the core abstraction in Apache TVM Unity stack, which is used to represent and optimize the primitive tensor functions. .. toctree:: :maxdepth: 2 abstraction learning tutorials/creation tutorials/transformation --- ## File: docs/deep_dive/tensor_ir/learning.rst .. 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. .. _tir-learning: Understand TensorIR Abstraction =============================== TensorIR is the tensor program abstraction in Apache TVM, which is one of the standard machine learning compilation frameworks. The principal objective of tensor program abstraction is to depict loops and associated hardware acceleration options, including threading, the application of specialized hardware instructions, and memory access. To help our explanations, let us use the following sequence of tensor computations as a motivating example. Specifically, for two :math:`128 \times 128` matrices ``A`` and ``B``, let us perform the following two steps of tensor computations. .. math:: Y_{i, j} &= \sum_k A_{i, k} \times B_{k, j} \\ C_{i, j} &= \mathbb{relu}(Y_{i, j}) = \mathbb{max}(Y_{i, j}, 0) The above computations resemble a typical primitive tensor function commonly seen in neural networks, a linear layer with relu activation. We use TensorIR to depict the above computations as follows. Before we invoke TensorIR, let's use native Python codes with NumPy to show the computation: .. code:: python def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray): Y = np.empty((128, 128), dtype="float32") for i in range(128): for j in range(128): for k in range(128): if k == 0: Y[i, j] = 0 Y[i, j] = Y[i, j] + A[i, k] * B[k, j] for i in range(128): for j in range(128): C[i, j] = max(Y[i, j], 0) With the low-level NumPy example in mind, now we are ready to introduce TensorIR. The code block below shows a TensorIR implementation of ``mm_relu``. The particular code is implemented in a language called TVMScript, which is a domain-specific dialect embedded in python AST. .. code:: python @tvm.script.ir_module class MyModule: @T.prim_func def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")): Y = T.alloc_buffer((128, 128), dtype="float32") for i, j, k in T.grid(128, 128, 128): with T.block("Y"): vi = T.axis.spatial(128, i) vj = T.axis.spatial(128, j) vk = T.axis.reduce(128, k) with T.init(): Y[vi, vj] = T.float32(0) Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] for i, j in T.grid(128, 128): with T.block("C"): vi = T.axis.spatial(128, i) vj = T.axis.spatial(128, j) C[vi, vj] = T.max(Y[vi, vj], T.float32(0)) Next, let's invest the elements in the above TensorIR program. Function Parameters and Buffers ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ **The function parameters correspond to the same set of parameters on the numpy function.** .. code:: python # TensorIR def mm_relu(A: T.Buffer[(128, 128), "float32"], B: T.Buffer[(128, 128), "float32"], C: T.Buffer[(128, 128), "float32"]): ... # NumPy def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray): ... Here ``A``, ``B``, and ``C`` takes a type named ``T.Buffer``, which with shape argument ``(128, 128)`` and data type ``float32``. This additional information helps possible MLC process to generate code that specializes in the shape and data type. **Similarly, TensorIR also uses a buffer type in intermediate result allocation.** .. code:: python # TensorIR Y = T.alloc_buffer((128, 128), dtype="float32") # NumPy Y = np.empty((128, 128), dtype="float32") Loop Iterations ~~~~~~~~~~~~~~~ **There are also direct correspondence of loop iterations.** ``T.grid`` is a syntactic sugar in TensorIR for us to write multiple nested iterators. .. code:: python # TensorIR with `T.grid` for i, j, k in T.grid(128, 128, 128): ... # TensorIR with `range` for i in range(128): for j in range(128): for k in range(128): ... # NumPy for i in range(128): for j in range(128): for k in range(128): ... Computational Block ~~~~~~~~~~~~~~~~~~~ A significant distinction lies in computational statements: **TensorIR incorporates an additional construct termed** ``T.block``. .. code:: python # TensorIR with T.block("Y"): vi = T.axis.spatial(128, i) vj = T.axis.spatial(128, j) vk = T.axis.reduce(128, k) with T.init(): Y[vi, vj] = T.float32(0) Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] # NumPy vi, vj, vk = i, j, k if vk == 0: Y[vi, vj] = 0 Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] A **block** represents a fundamental computation unit within TensorIR. Importantly, a block encompasses more information than standard NumPy code. It comprises a set of block axes ``(vi, vj, vk)`` and the computations delineated around them. .. code:: python vi = T.axis.spatial(128, i) vj = T.axis.spatial(128, j) vk = T.axis.reduce(128, k) The above three lines declare the **key properties** about block axes in the following syntax. .. code:: python [block_axis] = T.axis.[axis_type]([axis_range], [mapped_value]) These three lines convey the following details: - They specify the binding of ``vi``, ``vj``, ``vk`` (in this instance, to ``i``, ``j``, ``k``). - They declare the original range intended for ``vi``, ``vj``, ``vk`` (the 128 in ``T.axis.spatial(128, i)``). - They announce the properties of the iterators (spatial, reduce). Block Axis Properties ~~~~~~~~~~~~~~~~~~~~~ Let's delve deeper into the properties of the block axis. These properties signify the axis's relationship to the computation in progress. The block comprises three axes ``vi``, ``vj``, and ``vk``, meanwhile the block reads the buffer ``A[vi, vk]``, ``B[vk, vj]`` and writs the buffer ``Y[vi, vj]``. Strictly speaking, the block performs (reduction) updates to Y, which we label as write for the time being, as we don't require the value of Y from another block. Significantly, for a fixed value of ``vi`` and ``vj``, the computation block yields a point value at a spatial location of ``Y`` (``Y[vi, vj]``) that is independent of other locations in ``Y`` (with different ``vi``, ``vj`` values). We can refer to ``vi``, ``vj`` as **spatial axes** since they directly correspond to the start of a spatial region of buffers that the block writes to. The axes involved in reduction (``vk``) are designated as **reduce axes**. Why Extra Information in Block ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ One crucial observation is that the additional information (block axis range and their properties) makes the block to be **self-contained** when it comes to the iterations that it is supposed to carry out independent from the external loop-nest ``i, j, k``. The block axis information also provides additional properties that help us to validate the correctness of the external loops that are used to carry out the computation. For example, the above code block will result in an error because the loop expects an iterator of size 128, but we only bound it to a for loop of size 127. .. code:: python # wrong program due to loop and block iteration mismatch for i in range(127): with T.block("C"): vi = T.axis.spatial(128, i) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ error here due to iterator size mismatch ... Sugars for Block Axes Binding ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ In situations where each of the block axes is directly mapped to an outer loop iterator, we can use ``T.axis.remap`` to declare the block axis in a single line. .. code:: python # SSR means the properties of each axes are "spatial", "spatial", "reduce" vi, vj, vk = T.axis.remap("SSR", [i, j, k]) which is equivalent to .. code:: python vi = T.axis.spatial(range_of_i, i) vj = T.axis.spatial(range_of_j, j) vk = T.axis.reduce (range_of_k, k) So we can also write the programs as follows. .. code:: python @tvm.script.ir_module class MyModuleWithAxisRemapSugar: @T.prim_func def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")): Y = T.alloc_buffer((128, 128), dtype="float32") for i, j, k in T.grid(128, 128, 128): with T.block("Y"): vi, vj, vk = T.axis.remap("SSR", [i, j, k]) with T.init(): Y[vi, vj] = T.float32(0) Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj] for i, j in T.grid(128, 128): with T.block("C"): vi, vj = T.axis.remap("SS", [i, j]) C[vi, vj] = T.max(Y[vi, vj], T.float32(0)) --- ## File: docs/get_started/install.rst .. 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. .. _install: Installing Apache TVM ===================== .. contents:: Table of Contents :local: :depth: 2 Option 1. Prebuilt Package -------------------------- To help our community to use Apache TVM, a nightly prebuilt developer package is provided by `MLC community `_ Please visit the installation page for installation instructions: https://mlc.ai/package/. 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. Step 1. Install Dependencies ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ TVM requires the following dependencies: - CMake (>= 3.24.0) - LLVM (recommended >= 15) - Git - A recent C++ compiler supporting C++ 17, at the minimum - GCC 7.1 - Clang 5.0 - Apple Clang 9.3 - Visual Studio 2019 (v16.7) - Python (>= 3.8) - (Optional) Conda (Strongly Recommended) For Ubuntu/Debian users, the following APT Repository may help: - CMake: https://apt.kitware.com - LLVM: https://apt.llvm.org Step 2. Get Source from Github ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ First, You can also choose to clone the source repo from Github. The code of Apache TVM is hosted under the `Apache TVM `_ .. code:: bash git clone https://github.com/apache/tvm --recursive .. note:: It's important to use the ``--recursive`` flag when cloning the TVM repository, which will automatically clone the submodules. If you forget to use this flag, you can manually clone the submodules by running ``git submodule update --init --recursive`` in the root directory of the TVM repository. Step 3. Configure and Build ~~~~~~~~~~~~~~~~~~~~~~~~~~~ Create a build directory and run CMake to configure the build. The following example shows how to build .. code:: bash cd tvm rm -rf build && mkdir build && cd build # Specify the build configuration via CMake options cp ../cmake/config.cmake . Then you can configure the build by running CMake: .. code:: bash # controls default compilation flags (Candidates: Release, Debug, RelWithDebInfo) 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_METAL OFF)" >> config.cmake echo "set(USE_VULKAN OFF)" >> config.cmake echo "set(USE_OPENCL OFF)" >> config.cmake # cuBLAS, cuDNN, cutlass support, turn on if needed echo "set(USE_CUBLAS OFF)" >> config.cmake echo "set(USE_CUDNN OFF)" >> config.cmake echo "set(USE_CUTLASS OFF)" >> 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. Once config.cmake is edited accordingly, kick off build with the commands below: .. code:: bash cmake .. && cmake --build . --parallel $(nproc) A success build should produce ``libtvm`` and ``libtvm_runtime`` under ``build/`` directory. .. tabs :: .. code-tab :: bash Install via environment variable export TVM_HOME=/path/to/tvm export PYTHONPATH=$TVM_HOME/python:$PYTHONPATH .. code-tab :: bash Install via pip local project cd /path-to-tvm/python pip install -e . Step 4. Validate Installation ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Please the following code to validate the TVM installation: .. code:: bash python -c "import tvm; print(tvm.__file__)" If the installation is successful, you should see the path to the TVM Python package printed out. Also, please verify you installed the versions .. code:: bash python -c "import tvm.relax; print(\"OK\")" If the installation is successful, you should see ``OK`` printed out. --- ## File: docs/get_started/overview.md # Overview Apache TVM is a machine learning compilation framework, following the principle of **Python-first development** and **universal deployment**. It takes in pre-trained machine learning models, compiles and generates deployable modules that can be embedded and run everywhere. Apache TVM also enables customizing optimization processes to introduce new optimizations, libraries, codegen and more. ## Key Principle - **Python-first**: the optimization process is fully customizable in Python. It is easy to customize the optimization pipeline without recompiling the TVM stack. - **Composable**: the optimization process is composable. It is easy to compose new optimization passes, libraries and codegen to the existing pipeline. ## Key Goals - **Optimize** performance of ML workloads, composing libraries and codegen. - **Deploy** ML workloads to a diverse set of new environments, including new runtime and new hardware. - **Continuously improve and customize** ML deployment pipeline in Python by quickly customizing library dispatching, bringing in customized operators and code generation. ## Key Flow Here is a typical flow of using TVM to deploy a machine learning model. For a runnable example, please refer to [Quick Start](#quick_start) 1. **Import/construct an ML model** TVM supports importing models from various frameworks, such as PyTorch, TensorFlow for generic ML models. Meanwhile, we can create models directly using Relax frontend for scenarios of large language models. 2. **Perform composable optimization** transformations via `pipelines` The pipeline encapsulates a collection of transformations to achieve two goals: - **Graph Optimizations**: such as operator fusion, and layout rewrites. - **Tensor Program Optimization**: Map the operators to low-level implementations (both library or codegen) :::{note} The twos are goals but not the stages of the pipeline. The two optimizations are performed **at the same level**, or separately in two stages. ::: 3. **Build and universal deploy** Apache TVM aims to provide a universal deployment solution to bring machine learning everywhere with every language with minimum runtime support. TVM runtime can work in non-Python environments, so it works on mobile, edge devices or even bare metal devices. Additionally, TVM runtime comes with native data structures, and can also have zero copy exchange with the existing ecosystem (PyTorch, TensorFlow, TensorRT, etc.) using DLPack support. --- ## File: docs/reference/api/relax/analysis.rst .. 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.relax.analysis ------------------ .. automodule:: tvm.relax.analysis :members: :imported-members: --- ## File: docs/reference/api/relax/block_builder.rst .. 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.relax.block_builder ----------------------- .. automodule:: tvm.relax.block_builder :members: