{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/squeezeattention-2d-management-of-kv-cache-in","title":"SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget","arxiv_id":"2404.04793","date":"2024-04-07","proceeding":null,"authors":["ZiHao Wang","Bin Cui","Shaoduo Gan"],"abstract":"Optimizing the Key-Value (KV) cache of the Large Language Model (LLM) has been considered critical to saving the cost of inference. Most of the existing KV-cache compression algorithms attempted to sparsify the sequence of tokens by taking advantage of the different importance of tokens. However, most of these methods treat all layers equally, allocating the same KV budget to each layer. This approach is suboptimal, as some layers may be less sensitive to input tokens yet still receive the same budget as others. In this work, we found that by identifying the importance of attention layers, we could optimize the KV-cache jointly from two dimensions, i.e., sequence-wise and layer-wise. Based on our observations regarding layer-wise importance in inference, we propose SqueezeAttention to precisely optimize the allocation of KV-cache budget among layers on-the-fly and then incorporate three representative sequence-wise algorithms to compress the KV-cache for each layer with its very own budget. Specifically, we first measure each layer's importance by calculating the cosine similarity of the input prompt differences before and after the self-attention layers. Based on this similarity, we then categorize the layers into two groups and adjust their KV budgets accordingly. By optimizing the KV-cache from both sequence's and layer's dimensions, SqueezeAttention achieves around 30% to 70% of the memory reductions and up to 2.2 times of throughput improvements in a wide range of LLMs and benchmarks. The code is available at https://github.com/hetailang/SqueezeAttention.","url_abs":"https://arxiv.org/abs/2404.04793v2","url_pdf":"https://arxiv.org/pdf/2404.04793v2.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"squeezeattention-2d-management-of-kv-cache-in","repo_url":"https://github.com/hetailang/squeezeattention","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":0,"framework":"none","reach":{"status":"ok"}}],"tasks":[{"task_slug":"language-modelling","task_name":"Language Modelling"},{"task_slug":"large-language-model","task_name":"Large Language Model"},{"task_slug":"management","task_name":"Management"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":null,"atlas_url":"https://app.syntology.ai/?focus=2404.04793","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2404.04793"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/hetailang/squeezeattention","reach":{"status":"ok"}}],"summary":{"ran_fixture":2,"ran":1,"ran_honours":1,"ran_draft_wrong":3,"unverified":2},"by_repo_kind":{"official":{"samples":9,"ran":7,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":9,"samples":[{"code_sha256_prefix":"30d7eec482ebf6b1","entry":"repeat_kv","repo":"hetailang/squeezeattention","repo_kind":"official","path":"modeling_llama.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/modeling_llama.py","link_basis":"harvester_set","language":"python","status":"ran_fixture","verification_level":2,"contract_check":"RAISES","metamorphic_tier":"invariant","behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"30d7eec482ebf6b1"}},{"code_sha256_prefix":"d61c483a3c2b3156","entry":"apply_rotary_pos_emb","repo":"hetailang/squeezeattention","repo_kind":"official","path":"modeling_llama.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/modeling_llama.py","link_basis":"harvester_set","language":"python","status":"ran_fixture","verification_level":1,"contract_check":"RAISES","metamorphic_tier":"invariant","behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"d61c483a3c2b3156"}},{"code_sha256_prefix":"fed9b132b26ff2cc","entry":"convert_nochange","repo":"hetailang/squeezeattention","repo_kind":"official","path":"pred.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/pred.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"fed9b132b26ff2cc"}},{"code_sha256_prefix":"b349b79d9cc2934b","entry":"count_score","repo":"hetailang/squeezeattention","repo_kind":"official","path":"metrics.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/metrics.py","link_basis":"harvester_set","language":"python","status":"ran_honours","verification_level":1,"contract_check":"HONOURS","metamorphic_tier":"well_formed","behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"b349b79d9cc2934b"}},{"code_sha256_prefix":"e7e75981cb464788","entry":"normalize_answer","repo":"hetailang/squeezeattention","repo_kind":"official","path":"metrics.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/metrics.py","link_basis":"harvester_set","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"e7e75981cb464788"}},{"code_sha256_prefix":"8c5c581f9264c810","entry":"normalize_zh_answer","repo":"hetailang/squeezeattention","repo_kind":"official","path":"metrics.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/metrics.py","link_basis":"harvester_set","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":"deterministic","behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"8c5c581f9264c810"}},{"code_sha256_prefix":"b99eea6376d1e212","entry":"rotate_half","repo":"hetailang/squeezeattention","repo_kind":"official","path":"modeling_llama.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/modeling_llama.py","link_basis":"harvester_set","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":"invariant","behaviour_fingerprint":true,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"b99eea6376d1e212"}},{"code_sha256_prefix":"2e148fc130561394","entry":"build_chat","repo":"hetailang/squeezeattention","repo_kind":"official","path":"pred.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/pred.py","link_basis":"plan_row","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"2e148fc130561394"}},{"code_sha256_prefix":"b138b1dd1b97865f","entry":"parse_args","repo":"hetailang/squeezeattention","repo_kind":"official","path":"pred.py","file_url":"https://github.com/hetailang/squeezeattention/blob/HEAD/pred.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"b138b1dd1b97865f"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}