``` ├── .clang-format ├── .github/ ├── workflows/ ├── main.yml ├── .gitignore ├── .vscode/ ├── settings.json ├── LICENSE ├── README.md ├── book/ ├── .gitignore ├── book.toml ├── sitemap.sh ├── src/ ├── SUMMARY.md ├── copyright.md ├── discord-badge.svg ├── glossary.md ├── preface.md ├── setup.md ├── sitemap.txt ├── sitemap.xml ├── week1-01-attention.md ├── week1-02-positional-encodings.md ├── week1-03-gqa.md ├── week1-04-mlp-transformer.md ├── week1-05-model-1.md ├── week1-06-model-2.md ├── week1-07-generate.md ├── week1-overview.md ├── week2-overview.md ├── week3-overview.md ├── theme/ ├── head.hbs._ ├── build_ext.sh ├── check.py ├── main.py ├── pdm.lock ``` ## /.clang-format ```clang-format path="/.clang-format" BasedOnStyle: Google DerivePointerAlignment: false PointerAlignment: Right ColumnLimit: 120 IndentWidth: 4 AccessModifierOffset: -4 ``` ## /.github/workflows/main.yml ```yml path="/.github/workflows/main.yml" name: CI (main) on: push: branches: - main env: CARGO_TERM_COLOR: always jobs: build: runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v2 - name: setup rust toolchain run: rustup update && rustup toolchain install - uses: taiki-e/install-action@mdbook - name: patch for gh-pages build run: mv book/theme/head.hbs._ book/theme/head.hbs - name: build book run: cd book && mdbook build - uses: actions/upload-pages-artifact@v3 with: path: book/book deploy: needs: build permissions: pages: write id-token: write environment: name: github-pages url: ${{ steps.deployment.outputs.page_url }} runs-on: ubuntu-22.04 if: github.repository == 'skyzh/tiny-llm' steps: - name: Deploy to GitHub Pages id: deployment uses: actions/deploy-pages@v4 ``` ## /.gitignore ```gitignore path="/.gitignore" # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # 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 # 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/ cover/ # 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/ # PyBuilder .pybuilder/ target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: # .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 # UV # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. # This is especially recommended for binary packages to ensure reproducibility, and is more # commonly ignored for libraries. #uv.lock # poetry # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. # This is especially recommended for binary packages to ensure reproducibility, and is more # commonly ignored for libraries. # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control #poetry.lock # pdm # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. #pdm.lock # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it # in version control. # https://pdm.fming.dev/latest/usage/project/#working-with-version-control .pdm.toml .pdm-python .pdm-build/ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm __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 # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ # pytype static type analyzer .pytype/ # Cython debug symbols cython_debug/ # PyCharm # JetBrains specific template is maintained in a separate JetBrains.gitignore that can # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ # Ruff stuff: .ruff_cache/ # PyPI configuration file .pypirc *.dylib *.metallib ``` ## /.vscode/settings.json ```json path="/.vscode/settings.json" { "cmake.ignoreCMakeListsMissing": true } ``` ## /LICENSE ``` path="/LICENSE" 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. ``` ## /README.md # tiny-llm - LLM Serving in a Week [![CI (main)](https://github.com/skyzh/tiny-llm/actions/workflows/main.yml/badge.svg)](https://github.com/skyzh/tiny-llm/actions/workflows/main.yml) Still WIP and in very early stage. A tutorial on LLM serving using MLX for system engineers. The codebase is solely (almost!) based on MLX array/matrix APIs without any high-level neural network APIs, so that we can build the model serving infrastructure from scratch and dig into the optimizations. The goal is to learn the techniques behind efficiently serving an LLM model (i.e., Qwen2 models). ## Book The tiny-llm book is available at [https://skyzh.github.io/tiny-llm/](https://skyzh.github.io/tiny-llm/). You can follow the guide and start building. ## Community You may join skyzh's Discord server and study with the tiny-llm community. [![Join skyzh's Discord Server](book/src/discord-badge.svg)](https://skyzh.dev/join/discord) ## Roadmap | Week + Chapter | Topic | Code | Test | Doc | | -------------- | ----------------------------------------------------------- | ---- | ---- | --- | | 1.1 | Attention | ✅ | ✅ | ✅ | | 1.2 | RoPE | ✅ | ✅ | ✅ | | 1.3 | Grouped Query Attention | ✅ | 🚧 | 🚧 | | 1.4 | RMSNorm and MLP | ✅ | 🚧 | 🚧 | | 1.5 | Transformer Block | ✅ | 🚧 | 🚧 | | 1.6 | Load the Model | ✅ | 🚧 | 🚧 | | 1.7 | Generate Responses (aka Decoding) | ✅ | ✅ | 🚧 | | 2.1 | KV Cache | ✅ | 🚧 | 🚧 | | 2.2 | Quantized Matmul and Linear - CPU | ✅ | 🚧 | 🚧 | | 2.3 | Quantized Matmul and Linear - GPU | ✅ | 🚧 | 🚧 | | 2.4 | Flash Attention and Other Kernels | 🚧 | 🚧 | 🚧 | | 2.5 | Continuous Batching | 🚧 | 🚧 | 🚧 | | 2.6 | Speculative Decoding | 🚧 | 🚧 | 🚧 | | 2.7 | Prompt/Prefix Cache | 🚧 | 🚧 | 🚧 | | 3.1 | Paged Attention - Part 1 | 🚧 | 🚧 | 🚧 | | 3.2 | Paged Attention - Part 2 | 🚧 | 🚧 | 🚧 | | 3.3 | Prefill-Decode Separation | 🚧 | 🚧 | 🚧 | | 3.4 | Scheduler | 🚧 | 🚧 | 🚧 | | 3.5 | Parallelism | 🚧 | 🚧 | 🚧 | | 3.6 | AI Agent | 🚧 | 🚧 | 🚧 | | 3.7 | Streaming API Server | 🚧 | 🚧 | 🚧 | Other topics not covered: quantized/compressed kv cache ## /book/.gitignore ```gitignore path="/book/.gitignore" book ``` ## /book/book.toml ```toml path="/book/book.toml" [book] authors = ["Alex Chi"] language = "en" multilingual = false src = "src" title = "Tiny LLM - LLM Serving in a Week" [preprocessor.toc] command = "mdbook-toc" renderer = ["html"] [output.html] git-repository-url = "https://github.com/skyzh/tiny-llm" ``` ## /book/sitemap.sh ```sh path="/book/sitemap.sh" #!/bin/bash mdbook build static-sitemap-cli -b https://skyzh.github.io/tiny-llm -r book -f xml -o > src/sitemap.xml static-sitemap-cli -b https://skyzh.github.io/tiny-llm -r book -f txt -o > src/sitemap.txt ``` ## /book/src/SUMMARY.md # LLM Serving in a Week [Preface](./preface.md) [Setting Up the Environment](./setup.md) --- - [Week 1: From Matmul to Text](./week1-overview.md) - [Attention and Multi-Head Attention](./week1-01-attention.md) - [Positional Encodings and RoPE](./week1-02-positional-encodings.md) - [Grouped/Multi Query Attention]() - [Multilayer Perceptron Layer and Transformer]() - [Wiring the Qwen2 Model]() - [Loading the Model]() - [Generating the Response]() - [Week 2: Optimizing]() - [Week 3: Serving]() --- [Glossary Index](./glossary.md) ## /book/src/copyright.md

Your feedback is greatly appreciated. Welcome to join our Discord Community.
Found an issue? Create an issue / pull request on github.com/skyzh/tiny-llm.
tiny-llm-book © 2025 by Alex Chi Z is licensed under CC BY-NC-SA 4.0.

## /book/src/discord-badge.svg ```svg path="/book/src/discord-badge.svg" MDbot#9808skyzh's server ``` ## /book/src/glossary.md # Glossary Index - [Scaled Dot Product Attention](./week1-01-attention.md) - [Multi Head Attention](./week1-01-attention.md) - [Linear](./week1-01-attention.md) - [Rotary Positional Encoding](./week1-02-positional-encodings.md) {{#include copyright.md}} ## /book/src/preface.md # Preface This course is designed for systems engineers who want to understand how LLMs work. As a system engineer, I always wonder how things work internally and how to optimize them. I had a hard time figuring out the LLM stuff. Most of the open source projects that serve LLMs are highly optimized with CUDA kernels and other low-level optimizations. It is not easy to understand the whole picture by looking at a codebase of 100k lines of code. Therefore, I decided to implement an LLM serving project from scratch -- with only matrix manipulations APIs, so that I can understand what it takes to load those LLM model parameters and do the math magic to generate text. You can think of this course as an LLM version of CMU Deep Learning Systems course's [needle](https://github.com/dlsyscourse/hw1/tree/main/python/needle) project. ## Prerequisites You should have some experience with the basics of deep learning and have some idea of how PyTorch works. Some recommended resources are: - CMU [Intro to Machine Learning](https://www.cs.cmu.edu/~mgormley/courses/10601/) -- this course teaches you the basics of machine learning - CMU [Deep Learning Systems](https://dlsyscourse.org) -- this course teaches you how to build PyTorch from scratch ## Environment Setup This course uses [MLX](https://github.com/ml-explore/mlx), an array/machine learning library for Apple Silicon. Nowaways it's much easier to get an Apple Silicon device than NVIDIA GPUs. In theory you can also do this course with PyTorch or numpy, but we just don't have the test infra to support them. We test your implementation against PyTorch's CPU implementation and MLX's implementation to ensure correctness. ## Course Structure This course is divided into 3 weeks. We will serve the Qwen2-7B-Instruct model and optimize it throughout the course. * Week 1: serve Qwen2 with purely matrix manipulation APIs. Just Python. * Week 2: optimizations, implement C++/Metal custom kernels to make the model run faster. * Week 3: more optimizations, batch the requests to serve the model with high throughput. ## How to Use This Book The thing you are reading right now is the tiny-llm book. It is designed more like a guidebook instead of a textbook that explains everything from scratch. In this course, we provide the materials that we find useful on the Internet when the author(s) implemented the tiny-llm project. The Internet does a better job of explaining the concepts and I do not think it is necessary to repeat everything here. Think of this as a guide (of a list of tasks) and some hints! We will also unify the language of the Internet materials so that it is easier to correspond them to the codebase. For example, we will have a unified dimension symbols for the tensors. You do not need to figure out what `H`, `L`, `E` stands for and what dimension of the matrixes are passed into the function. ## Community You may join skyzh's Discord server and study with the tiny-llm community. [![Join skyzh's Discord Server](discord-badge.svg)](https://skyzh.dev/join/discord) ## Get Started Now, you can start to set up the environment following the instructions in [Setting Up the Environment](./setup.md) and begin your journey to build tiny-llm! {{#include copyright.md}} ## /book/src/setup.md # Setting Up the Environment To follow along this course, you will need a mactonish device with Apple Silicon. We manage the codebase with pdm. ## Install pdm Please follow the [offcial guide](https://pdm-project.org/en/latest/) to install pdm. ## Clone the Repository ```bash git clone https://github.com/skyzh/tiny-llm ``` The repository is organized as follows: ``` src/tiny_llm -- your implementation src/tiny_llm_week1_ref -- reference implementation of week 1 tests/ -- unit tests for your implementation tests_ref_impl_week1/ -- unit tests for the reference implementation of week 1 book/ -- the book ``` We provide all reference implementations and you can refer to them if you get stuck in the course. ## Install Dependencies ```bash cd tiny-llm pdm install -v # this will automatically create a virtual environment and install all dependencies ``` ## Check the Installation ```bash pdm run python check.py # The reference solution should pass all the tests pdm run test_ref_impl_week1 ``` ## Run Unit Tests Your code is in `src/tiny_llm`. You can run the unit tests with: ```bash pdm run test ``` ## Download the Model Parameters We will use the Qwen2-7B-Instruct model for this course. It takes ~20GB of memory in week 1 to load the model parameters. If you do not have enough memory, you can consider using the smaller 0.5B model. (We will make the course compatible with it in the future; meanwhile, you have to figure out things on your own if you use the 0.5B model. Likely, this only matters after week 1 day 6 when you start to load the model parameters.) Follow the guide of [this page](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli) to install the huggingface cli. You should install it in your user directory/globally instead of in the tiny-llm virtual environment created by pdm. The model parameters are hosted on Hugging Face. Once you authenticated your cli with the credentials, you can download them with: ```bash # do not do this in the virtual environment created by pdm huggingface-cli login huggingface-cli download Qwen/Qwen2-7B-Instruct-MLX ``` Then, you can run: ```bash pdm run main --solution week1_ref ``` It should load the model and print some text. In week 2, we will write some kernels in C++/Metal, and we will need to set up additional tools for that. We will cover it later. {{#include copyright.md}} ## /book/src/sitemap.txt https://skyzh.github.io/tiny-llm https://skyzh.github.io/tiny-llm/glossary https://skyzh.github.io/tiny-llm/preface https://skyzh.github.io/tiny-llm/setup https://skyzh.github.io/tiny-llm/week1-overview ## /book/src/sitemap.xml ```xml path="/book/src/sitemap.xml" https://skyzh.github.io/tiny-llm 2025-04-26T03:38:02.503Z https://skyzh.github.io/tiny-llm/glossary 2025-04-26T03:38:02.504Z https://skyzh.github.io/tiny-llm/preface 2025-04-26T03:38:02.503Z https://skyzh.github.io/tiny-llm/setup 2025-04-26T03:38:02.504Z https://skyzh.github.io/tiny-llm/week1-overview 2025-04-26T03:38:02.504Z ``` ## /book/src/week1-01-attention.md # Week 1 Day 1: Attention and Multi-Head Attention In day 1, we will implement the basic attention layer and the multi-head attention layer. Attention layers take a input sequence and focus on different parts of the sequence when generating the output. Attention layers are the key building blocks of the Transformer models. [📚 Reading: Transformer Architecture](https://huggingface.co/learn/llm-course/chapter1/6) We use the Qwen2 model for text generation. The model is a decoder-only model. The input of the model is a sequence of token embeddings. The output of the model is the most likely next token ID. [📚 Reading: LLM Inference, the Decode Phase](https://huggingface.co/learn/llm-course/chapter1/8) Back to the attention layer. The attention layer takes a query, a key, and a value. In a classic implementation, all of them are of the same shape: `N.. x L x D`. `N..` is zero or some number of dimensions for batches. Within each of the batch, `L` is the sequence length and `D` is the dimension of the embedding for a given head in the sequence. So, for example, if we have a sequence of 1024 tokens, where each of the token has a 512-dimensional embedding (head_dim), we will pass a tensor of the shape `N.. x 1024 x 512` to the attention layer. ## Task 1: Implement `scaled_dot_product_attention` In this task, we will implement the scaled dot product attention function. ``` pdm run test -k week_1_day_1_task_1 -v ``` **📚 Readings** * [Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/) * [PyTorch Scaled Dot Product Attention API](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) (assume `enable_gqa=False`, assume dim_k=dim_v=dim_q and H_k=H_v=H_q) * [MLX Scaled Dot Product Attention API](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.fast.scaled_dot_product_attention.html) (assume dim_k=dim_v=dim_q and H_k=H_v=H_q) * [Attention is All You Need](https://arxiv.org/abs/1706.03762) Implement `scaled_dot_product_attention`. The function takes key, value, and query of the same dimensions. ``` L is seq_len, in PyTorch API it's S (source len) D is head_dim K: N.. x L x D V: N.. x L x D Q: N.. x L x D output: N.. x L x D ``` You may use `softmax` provided by mlx and implement it later in week 2. Because we are always using the attention layer within the multi-head attention layer, the actual tensor shape when serving the model will be: ``` K: 1 x H x L x D V: 1 x H x L x D Q: 1 x H x L x D output: 1 x H x L x D mask: 1 x H x L x L ``` .. though the attention layer only cares about the last two dimensions. The test case will test any shape of the batching dimension. At the end of this task, you should be able to pass the following tests: ``` pdm run test -k test_attention_simple pdm run test -k test_attention_with_mask ``` ## Task 2: Implement `MultiHeadAttention` In this task, we will implement the multi-head attention layer. ``` src/tiny_llm/attention.py ``` **📚 Readings** * [Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/) * [PyTorch MultiHeadAttention API](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html) (assume dim_k=dim_v=dim_q and H_k=H_v=H_q) * [MLX MultiHeadAttention API](https://ml-explore.github.io/mlx/build/html/python/nn/_autosummary/mlx.nn.MultiHeadAttention.html) (assume dim_k=dim_v=dim_q and H_k=H_v=H_q) * [The Illustrated GPT-2 (Visualizing Transformer Language Models)](https://jalammar.github.io/illustrated-gpt2) helps you better understand what key, value, and query are. Implement `MultiHeadAttention`. The layer takes a batch of vectors `x`, maps it through the K, V, Q weight matrixes, and use the attention function we implemented in day 1 to compute the result. The output needs to be mapped using the O weight matrix. You will also need to implement the `linear` function first. For `linear`, it takes a tensor of the shape `N.. x I`, a weight matrix of the shape `O x I`, and a bias vector of the shape `O`. The output is of the shape `N.. x O`. `I` is the input dimension and `O` is the output dimension. For the `MultiHeadAttention` layer, the input tensor `x` has the shape `N x L x E`, where `E` is the dimension of the embedding for a given token in the sequence. The `K/Q/V` weight matrixes will map the tensor into key, value, and query separately, where the dimension `E` will be mapped into a dimension of size `H x D`, which means that the token embedding gets mapped into `H` heads, each with a dimension of `D`. You can directly reshape the tensor to split the `H x D` dimension into two dimensions of `H` and `D` to get `H` heads for the token. Then, apply the attention function to each of the head (this requires a transpose, using `swapaxes` in mlx). The attention function takes `N.. x H x L x D` as input so that it produces an output for each of the head of the token. Then, you can transpose it into `N.. x L x H x D` and reshape it so that all heads get merged back together with a shape of `N.. x L x (H x D)`. Map it through the output weight matrix to get the final output. ``` E is hidden_size or embed_dim or dims or model_dim H is num_heads D is head_dim L is seq_len, in PyTorch API it's S (source len) W_q/k/v: E x (H x D) output/x: N x L x E W_o: (H x D) x E ``` At the end of the day, you should be able to pass the following tests: ``` pdm run test -k week_1_day_1_task_2 -v ``` {{#include copyright.md}} ## /book/src/week1-02-positional-encodings.md # Week 1 Day 2: Positional Encodings and RoPE In day 2, we will implement the positional embedding used in the Qwen2 model: Rotary Postional Encoding. In a transformer model, we need a way to embed the information of the position of a token into the input of the attention layers. In Qwen2, positional embedding is applied within the multi head attention layer on the query and key vectors. **📚 Readings** - [You could have designed state of the art positional encoding](https://huggingface.co/blog/designing-positional-encoding) - [Roformer: Enhanced Transformer with Rotary Positional Encoding](https://arxiv.org/pdf/2104.09864) ## Task 1: Implement Rotary Postional Encoding "RoPE" You will need to modify the following file: ``` src/tiny_llm/positional_encoding.py ``` In traditional RoPE (as described in the readings), the positional encoding is applied to each head of the query and key vectors. You can pre-compute the frequencies when initializing the `RoPE` class. If `offset` is not provided, the positional encoding will be applied to the entire sequence: 0th frequency applied to the 0th token, up to the (L-1)-th token. Otherwise, the positional encoding will be applied to the sequence according to the offset slice. If the offset slice is 5..10, then the sequence length provided to the layer would be 5, and the 0th token will be applied with the 5th frequency. ``` x: (N, L, H, D) cos/sin_freqs: (MAX_SEQ_LEN, D // 2) ``` In the traditional form of RoPE, each head on the dimension of `D` is viewed as consequtive complex pairs. That is to say, if D = 8, then, x[0] and x[1] are a pair, x[2] and x[3] are another pair, and so on. A pair gets the same frequency from `cos/sin_freqs`. ``` output[0] = x[0] * cos_freqs[0] + x[1] * sin_freqs[0] output[1] = x[0] * -sin_freqs[0] + x[1] * cos_freqs[0] output[2] = x[2] * cos_freqs[1] + x[3] * sin_freqs[1] output[3] = x[2] * -sin_freqs[1] + x[3] * cos_freqs[1] ...and so on ``` You can do this by reshaping `x` to (N, L, H, D // 2, 2) and then applying the above formula to each pair. **📚 Readings** - [PyTorch RotaryPositionalEmbeddings API](https://pytorch.org/torchtune/stable/generated/torchtune.modules.RotaryPositionalEmbeddings.html) - [MLX Implementation of RoPE before the custom metal kernel implementation](https://github.com/ml-explore/mlx/pull/676/files) You can test your implementation by running the following command: ``` pdm run test -k week_1_day_2_task_1 -v ``` ## Task 2: Implement `RoPE` in the non-traditional form The Qwen2 model uses a non-traditional form of RoPE. In this form, the head embedding dimension is split into two halves, and the two halves are applied with different frequencies. Let's say `x1 = x[.., :HALF_DIM]` and `x2 = x[.., HALF_DIM:]`. ``` output[0] = x1[0] * cos_freqs[0] + x2[0] * sin_freqs[0] output[HALF_DIM] = x1[0] * -sin_freqs[0] + x2[0] * cos_freqs[0] output[1] = x1[1] * cos_freqs[1] + x2[1] * sin_freqs[1] output[HALF_DIM + 1] = x1[1] * -sin_freqs[1] + x2[1] * cos_freqs[1] ...and so on ``` You can do this by directly getting the first half / second half of the embedding dimension of `x` and applying the frequencies to each half separately. You can test your implementation by running the following command: ``` pdm run test -k week_1_day_2_task_2 -v ``` **📚 Readings** - [vLLM implementation of RoPE](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding.py) {{#include copyright.md}} ## /book/src/week1-03-gqa.md # Grouped/Multi Query Attention ## /book/src/week1-04-mlp-transformer.md # Multilayer Perceptron Layer and Transformer ## /book/src/week1-05-model-1.md # Wiring the Qwen2 Model ## /book/src/week1-06-model-2.md # Loading the Model ## /book/src/week1-07-generate.md # Generating the Response ## /book/src/week1-overview.md # Week 1: From Matmul to Text In this week, we will start from the basic matrix operations and see how those these matrix manipulations can turn the Qwen2 model parameters into a model that generates text. We will implement the neural network layers used in the Qwen2 model using mlx's matrix APIs. We will use the Qwen2-7B-Instruct model for this week. As we need to dequantize the model parameters, the 4GB model needs 20GB of memory in week 1. If you do not have enough memory, you can consider using the smaller 0.5B model (we do not have infra to test it so you need to figure out things on your own unfortunately). The MLX version of the Qwen2-7B-Instruct model we downloaded in the setup is an int4 quantized version of the original bfloat16 model. ## What We will Cover * Attention, Multi-Head Attention, and Grouped/Multi Query Attention * Positional Embeddings and RoPE * Put the attention layers together and implement the whole Transformer block * Implement the MLP layer and the whole Transformer model * Load the Qwen2 model parameters and generate text ## What We will Not Cover To make the journey as interesting as possible, we will skip a few things for now: * How to quantize/dequantize a model -- that will be part of week 2. The Qwen2 model is quantized so we will need to dequantize them before we can use them in our layer implementations. * Actually we still used some APIs other than matrix manipulations -- like softmax, exp, log, etc. But they are simple and not implementing them would not affect the learning experience. * Tokenizer -- we will not implement the tokenizer from scratch. We will use the `mlx_lm` tokenizer to tokenize the input. * Loading the model weights -- I don't think it's an interesting thing to learn how to decode those tensor dump files, so we will use the `mlx_lm` to load the model and steal the weights from the loaded model into our layer implementations. ## Qwen2 Models You can try the Qwen2 model with MLX/vLLM. You can read the blog post below to have some idea of what we will build within this course. At the end of this week, we will be able to chat with the model -- that is to say, use Qwen2 to generate text, as a casual language model. The reference implementation of the Qwen2 model can be found in huggingface transformers, vLLM, and mlx-lm. You may utilize these resources to better understand the internals of the model and what we will implement in this week. **📚 Readings** - [Qwen2.5: A Party of Foundation Models!](https://qwenlm.github.io/blog/qwen2.5/) - [Key Concepts of the Qwen2 Model](https://qwen.readthedocs.io/en/latest/getting_started/concepts.html) - [Huggingface Transformers - Qwen2](https://github.com/huggingface/transformers/tree/main/src/transformers/models/qwen2) - [vLLM Qwen2](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/qwen2.py) - [mlx-lm Qwen2](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/qwen2.py) - [Qwen2 Technical Report](https://arxiv.org/pdf/2407.10671) - [Qwen2.5 Technical Report](https://arxiv.org/pdf/2412.15115) {{#include copyright.md}} ## /book/src/week2-overview.md https://github.com/ml-explore/mlx/blob/main/mlx/backend/cpu/quantized.cpp https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/linear.py MLX uses INT4 W4A16 https://ml-explore.github.io/mlx/build/html/dev/extensions.html https://github.com/ggml-org/llama.cpp/blob/master/ggml/src/ggml-metal/ggml-metal.metal https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/kernels/quantized.h#L962 pdm run ./build_ext.sh speculative decoding prefill and decode separation quantized kv cache Assert return data type ## /book/src/week3-overview.md ## /book/theme/head.hbs._ ```_ path="/book/theme/head.hbs._" ``` ## /build_ext.sh ```sh path="/build_ext.sh" #!/bin/bash set -e pdm run build-ext-ref cp src/extensions_ref/build/lib/tiny_llm_ext_ref/tiny_llm_ext_ref.metallib .venv/lib/python3.12/site-packages/mlx/lib/ pdm run test-week2-ref -k 'week_2_day_2' ``` ## /check.py ```py path="/check.py" import mlx.core as mx import torch with mx.stream(mx.cpu): a = mx.array([1, 2, 3]) b = mx.array([4, 5, 6]) c = mx.add(a, b) print(c) with mx.stream(mx.gpu): a = mx.array([1, 2, 3]) b = mx.array([4, 5, 6]) c = mx.add(a, b) print(c) print( torch.add( torch.tensor([1, 2, 3], device="cpu"), torch.tensor([4, 5, 6], device="cpu") ) ) ``` ## /main.py ```py path="/main.py" from mlx_lm import load import mlx.core as mx import argparse parser = argparse.ArgumentParser() parser.add_argument("--model", type=str, default="Qwen/Qwen2-7B-Instruct-MLX") parser.add_argument( "--prompt", type=str, default="Give me a short introduction to large language model.", ) parser.add_argument("--solution", type=str, default="tiny_llm") parser.add_argument("--device", type=str, default="gpu") args = parser.parse_args() use_mlx = False if args.solution == "tiny_llm": from tiny_llm import Qwen2Model, simple_generate print("Using your tiny_llm solution") elif args.solution == "tiny_llm_week1_ref" or args.solution == "week1_ref": from tiny_llm_week1_ref import Qwen2Model, simple_generate print("Using tiny_llm_week1_ref solution") elif args.solution == "tiny_llm_week2_ref" or args.solution == "week2_ref": from tiny_llm_week2_ref import Qwen2Model, simple_generate print("Using tiny_llm_week2_ref solution") elif args.solution == "mlx": use_mlx = True from mlx_lm.generate import stream_generate print("Using the original mlx model") else: raise ValueError(f"Solution {args.solution} not supported") mlx_model, tokenizer = load( args.model, tokenizer_config={"eos_token": "<|im_end|>"}, model_config={"tie_word_embeddings": False, "rope_traditional": False}, ) with mx.stream(mx.gpu if args.device == "gpu" else mx.cpu): if use_mlx: tiny_llm_model = mlx_model else: tiny_llm_model = Qwen2Model(mlx_model) messages = [ {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": args.prompt}, ] prompt = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) if not use_mlx: simple_generate(tiny_llm_model, tokenizer, prompt) else: for resp in stream_generate(tiny_llm_model, tokenizer, prompt): print(resp.text, end="", flush=True) ``` The content has been capped at 50000 tokens, and files over NaN bytes have been omitted. The user could consider applying other filters to refine the result. The better and more specific the context, the better the LLM can follow instructions. If the context seems verbose, the user can refine the filter using uithub. Thank you for using https://uithub.com - Perfect LLM context for any GitHub repo.