https://github.com/mit-han-lab/streaming-llm Skip to content Toggle navigation Sign up * Product + Actions Automate any workflow + Packages Host and manage packages + Security Find and fix vulnerabilities + Codespaces Instant dev environments + Copilot Write better code with AI + Code review Manage code changes + Issues Plan and track work + Discussions Collaborate outside of code Explore + All features + Documentation + GitHub Skills + Blog * Solutions For + Enterprise + Teams + Startups + Education By Solution + CI/CD & Automation + DevOps + DevSecOps Resources + Learning Pathways + White papers, Ebooks, Webinars + Customer Stories + Partners * Open Source + GitHub Sponsors Fund open source developers + The ReadME Project GitHub community articles Repositories + Topics + Trending + Collections * Pricing Search or jump to... Search code, repositories, users, issues, pull requests... Search [ ] Clear Search syntax tips Provide feedback We read every piece of feedback, and take your input very seriously. [ ] [ ] Include my email address so I can be contacted Cancel Submit feedback Saved searches Use saved searches to filter your results more quickly Name [ ] Query [ ] To see all available qualifiers, see our documentation. Cancel Create saved search Sign in Sign up You signed in with another tab or window. Reload to refresh your session. You signed out in another tab or window. Reload to refresh your session. You switched accounts on another tab or window. Reload to refresh your session. Dismiss alert {{ message }} mit-han-lab / streaming-llm Public * Notifications * Fork 81 * Star 1.6k Efficient Streaming Language Models with Attention Sinks License MIT license 1.6k stars 81 forks Activity Star Notifications * Code * Issues 6 * Pull requests 1 * Actions * Projects 0 * Security * Insights More * Code * Issues * Pull requests * Actions * Projects * Security * Insights mit-han-lab/streaming-llm This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository. main Switch branches/tags [ ] Branches Tags Could not load branches Nothing to show {{ refName }} default View all branches Could not load tags Nothing to show {{ refName }} default View all tags Name already in use A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch? Cancel Create 1 branch 0 tags Code * Local * Codespaces * Clone HTTPS GitHub CLI [https://github.com/m] Use Git or checkout with SVN using the web URL. [gh repo clone mit-ha] Work fast with our official CLI. Learn more about the CLI. * Open with GitHub Desktop * Download ZIP Sign In Required Please sign in to use Codespaces. Launching GitHub Desktop If nothing happens, download GitHub Desktop and try again. Launching GitHub Desktop If nothing happens, download GitHub Desktop and try again. Launching Xcode If nothing happens, download Xcode and try again. Launching Visual Studio Code Your codespace will open once ready. There was a problem preparing your codespace, please try again. Latest commit @Guangxuan-Xiao Guangxuan-Xiao Update README.md ... 406ed11 Oct 3, 2023 Update README.md 406ed11 Git stats * 21 commits Files Permalink Failed to load latest commit information. Type Name Latest commit message Commit time data upload ppl eval and llama chatbot demo October 2, 2023 06:06 examples upload ppl eval and llama chatbot demo October 2, 2023 06:06 figures update readme October 2, 2023 06:09 streaming_llm core code October 1, 2023 20:51 .gitignore upload ppl eval and llama chatbot demo October 2, 2023 06:06 LICENSE Initial commit September 29, 2023 13:45 README.md Update README.md October 3, 2023 15:45 setup.py core code October 1, 2023 20:51 View code [ ] Efficient Streaming Language Models with Attention Sinks [paper] TL; DR Abstract Usage Environment Setup Run Streaming Llama Chatbot FAQ TODOs Citation README.md Efficient Streaming Language Models with Attention Sinks [paper] schemes streamingllm_demo.mp4 TL;DR We deploy LLMs for infinite-length inputs without sacrificing efficiency and performance. Abstract Deploying Large Language Models (LLMs) in streaming applications such as multi-round dialogue, where long interactions are expected, is urgently needed but poses two major challenges. Firstly, during the decoding stage, caching previous tokens' Key and Value states (KV) consumes extensive memory. Secondly, popular LLMs cannot generalize to longer texts than the training sequence length. Window attention, where only the most recent KVs are cached, is a natural approach --- but we show that it fails when the text length surpasses the cache size. We observe an interesting phenomenon, namely attention sink, that keeping the KV of initial tokens will largely recover the performance of window attention. In this paper, we first demonstrate that the emergence of attention sink is due to the strong attention scores towards initial tokens as a ``sink'' even if they are not semantically important. Based on the above analysis, we introduce StreamingLLM, an efficient framework that enables LLMs trained with a finite length attention window to generalize to infinite sequence length without any fine-tuning. We show that StreamingLLM can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more. In addition, we discover that adding a placeholder token as a dedicated attention sink during pre-training can further improve streaming deployment. In streaming settings, StreamingLLM outperforms the sliding window recomputation baseline by up to 22.2x speedup. Usage Environment Setup conda create -yn streaming python=3.8 conda activate streaming pip install torch torchvision torchaudio pip install transformers accelerate datasets evaluate wandb scikit-learn scipy sentencepiece python setup.py develop Run Streaming Llama Chatbot CUDA_VISIBLE_DEVICES=0 python examples/run_streaming_llama.py --enable_streaming FAQ 1. What does "working on infinite-length inputs" imply for LLMs? Handling infinite-length text with LLMs presents challenges. Notably, storing all previous Key and Value (KV) states demands significant memory, and models might struggle to generate text beyond their training sequence length. StreamingLLM addresses this by retaining only the most recent tokens and attention sinks, discarding intermediate tokens. This enables the model to generate coherent text from recent tokens without a cache reset -- a capability not seen in earlier methods. 2. Is the context window of LLMs expanded? No. The context window remains unchanged. Only the most recent tokens and attention sinks are retained, discarding middle tokens. This means the model can only process the latest tokens. The context window remains constrained by its initial pre-training. For instance, if Llama-2 is pre-trained with a context window of 4096 tokens, then the maximum cache size for StreamingLLM on Llama-2 remains 4096. 3. Can I input an extensive text, like a book, into StreamingLLM for summarization? While you can input a lengthy text, the model will only recognize the latest tokens. Thus, if a book is an input, StreamingLLM might only summarize the concluding paragraphs, which might not be very insightful. As emphasized earlier, we neither expand the LLMs' context window nor enhance their long-term memory. StreamingLLM's strength lies in generating fluent text from recent tokens without needing a cache refresh. 4. What is the ideal use case for StreamingLLM? StreamingLLM is optimized for streaming applications, such as multi-round dialogues. It's ideal for scenarios where a model needs to operate continually without requiring extensive memory or dependency on past data. An example is a daily assistant based on LLMs. StreamingLLM would let the model function continuously, basing its responses on recent conversations without needing to refresh its cache. Earlier methods would either need a cache reset when the conversation length exceeded the training length (losing recent context) or recompute KV states from recent text history, which can be time-consuming. 5. How does StreamingLLM relate to recent works on context extension? StreamingLLM is orthogonal to recent context extension methods and can be integrated with them. In StreamingLLM's context, "context extension" refers to the possibility of using a larger cache size to store more recent tokens. For a practical demonstration, refer to Figure 9 in our paper, where we implement StreamingLLM with models like LongChat-7B-v1.5-32K and Llama-2-7B-32K-Instruct. TODOs We will release the code and data in the following order, please stay tuned! * [*] Release core code of StreamingLLM, including Llama-2, MPT, Falcon, and Pythia. * [*] Release perplexity evaluation code * [*] Release Streaming Llama Chatbot demo. * [ ] Release StreamEval dataset and evaluation code. Citation If you find StreamingLLM useful or relevant to your project and research, please kindly cite our paper: @article{xiao2023streamingllm, title={Efficient Streaming Language Models with Attention Sinks}, author={Xiao, Guangxuan and Tian, Yuandong and Chen, Beidi and Han, Song and Lewis, Mike}, journal={arXiv}, year={2023} } About Efficient Streaming Language Models with Attention Sinks Resources Readme License MIT license Activity Stars 1.6k stars Watchers 22 watching Forks 81 forks Report repository Releases No releases published Packages 0 No packages published Contributors 3 * @Guangxuan-Xiao Guangxuan-Xiao Guangxuan Xiao * @r2d4 r2d4 Matt Rickard * @cosmojg cosmojg Cosmo Languages * Python 100.0% Footer (c) 2023 GitHub, Inc. Footer navigation * Terms * Privacy * Security * Status * Docs * Contact GitHub * Pricing * API * Training * Blog * About You can't perform that action at this time.