{"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/keyformer-kv-cache-reduction-through-key","title":"Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference","arxiv_id":"2403.09054","date":"2024-03-14","proceeding":null,"authors":["Muhammad Adnan","Akhil Arunkumar","Gaurav Jain","Prashant J. Nair","Ilya Soloveychik","Purushotham Kamath"],"abstract":"Transformers have emerged as the underpinning architecture for Large Language Models (LLMs). In generative language models, the inference process involves two primary phases: prompt processing and token generation. Token generation, which constitutes the majority of the computational workload, primarily entails vector-matrix multiplications and interactions with the Key-Value (KV) Cache. This phase is constrained by memory bandwidth due to the overhead of transferring weights and KV cache values from the memory system to the computing units. This memory bottleneck becomes particularly pronounced in applications that require long-context and extensive text generation, both of which are increasingly crucial for LLMs. This paper introduces \"Keyformer\", an innovative inference-time approach, to mitigate the challenges associated with KV cache size and memory bandwidth utilization. Keyformer leverages the observation that approximately 90% of the attention weight in generative inference focuses on a specific subset of tokens, referred to as \"key\" tokens. Keyformer retains only the key tokens in the KV cache by identifying these crucial tokens using a novel score function. This approach effectively reduces both the KV cache size and memory bandwidth usage without compromising model accuracy. We evaluate Keyformer's performance across three foundational models: GPT-J, Cerebras-GPT, and MPT, which employ various positional embedding algorithms. Our assessment encompasses a variety of tasks, with a particular emphasis on summarization and conversation tasks involving extended contexts. Keyformer's reduction of KV cache reduces inference latency by 2.1x and improves token generation throughput by 2.4x, while preserving the model's accuracy.","url_abs":"https://arxiv.org/abs/2403.09054v2","url_pdf":"https://arxiv.org/pdf/2403.09054v2.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":"keyformer-kv-cache-reduction-through-key","repo_url":"https://github.com/d-matrix-ai/keyformer-llm","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":0,"framework":"pytorch","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[{"task_slug":"text-generation","task_name":"Text Generation"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2403.09054","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2403.09054"}},"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/d-matrix-ai/keyformer-llm","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"ran":3,"unverified":5},"by_repo_kind":{"official":{"samples":8,"ran":3,"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":0,"samples":[{"code_sha256_prefix":"c2217cd556fe5607","entry":"create_sinusoidal_positions","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/gptj-keyformer-lib/modeling_gptj.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/gptj-keyformer-lib/modeling_gptj.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"c2217cd556fe5607"}},{"code_sha256_prefix":"22d927bdb866159f","entry":"get_embed_positions","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/gptj-keyformer-lib/modeling_gptj.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/gptj-keyformer-lib/modeling_gptj.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"22d927bdb866159f"}},{"code_sha256_prefix":"2e3bec9f0c9b1973","entry":"rotate_every_two","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/gptj-keyformer-lib/modeling_gptj.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/gptj-keyformer-lib/modeling_gptj.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"2e3bec9f0c9b1973"}},{"code_sha256_prefix":"69a085d27f83db2a","entry":"flash_attn_fn","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/mpt-keyformer-lib/attention_streaming_llm.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/mpt-keyformer-lib/attention_streaming_llm.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"69a085d27f83db2a"}},{"code_sha256_prefix":"5b88d9d6199110e1","entry":"keyformer_mask","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/mpt-keyformer-lib/attention_lm_eval_harness.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/mpt-keyformer-lib/attention_lm_eval_harness.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"5b88d9d6199110e1"}},{"code_sha256_prefix":"bc9fb74e483fb901","entry":"load_tf_weights_in_gpt2","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/cerebras-keyformer-lib/modeling_gpt2.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/cerebras-keyformer-lib/modeling_gpt2.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"bc9fb74e483fb901"}},{"code_sha256_prefix":"aa98cb8ec8ec9683","entry":"scaled_multihead_dot_product_attention","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/mpt-keyformer-lib/attention_streaming_llm.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/mpt-keyformer-lib/attention_streaming_llm.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"aa98cb8ec8ec9683"}},{"code_sha256_prefix":"a27f1a3818430d86","entry":"triton_flash_attn_fn","repo":"d-matrix-ai/keyformer-llm","repo_kind":"official","path":"models/mpt-keyformer-lib/attention_streaming_llm.py","file_url":"https://github.com/d-matrix-ai/keyformer-llm/blob/HEAD/models/mpt-keyformer-lib/attention_streaming_llm.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"a27f1a3818430d86"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}