{"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/flexprefill-a-context-aware-sparse-attention","title":"FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference","arxiv_id":"2502.20766","date":"2025-02-28","proceeding":null,"authors":["Xunhao Lai","Jianqiao Lu","Yao Luo","Yiyuan Ma","Xun Zhou"],"abstract":"Large language models (LLMs) encounter computational challenges during long-sequence inference, especially in the attention pre-filling phase, where the complexity grows quadratically with the prompt length. Previous efforts to mitigate these challenges have relied on fixed sparse attention patterns or identifying sparse attention patterns based on limited cases. However, these methods lacked the flexibility to efficiently adapt to varying input demands. In this paper, we introduce FlexPrefill, a Flexible sparse Pre-filling mechanism that dynamically adjusts sparse attention patterns and computational budget in real-time to meet the specific requirements of each input and attention head. The flexibility of our method is demonstrated through two key innovations: 1) Query-Aware Sparse Pattern Determination: By measuring Jensen-Shannon divergence, this component adaptively switches between query-specific diverse attention patterns and predefined attention patterns. 2) Cumulative-Attention Based Index Selection: This component dynamically selects query-key indexes to be computed based on different attention patterns, ensuring the sum of attention scores meets a predefined threshold. FlexPrefill adaptively optimizes the sparse pattern and sparse ratio of each attention head based on the prompt, enhancing efficiency in long-sequence inference tasks. Experimental results show significant improvements in both speed and accuracy over prior methods, providing a more flexible and efficient solution for LLM inference.","url_abs":"https://arxiv.org/abs/2502.20766v1","url_pdf":"https://arxiv.org/pdf/2502.20766v1.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":"flexprefill-a-context-aware-sparse-attention","repo_url":"https://github.com/bytedance/FlexPrefill","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[],"methods":[{"method_slug":"attention","method_name":"Attention"},{"method_slug":"speed","method_name":"SPEED"},{"method_slug":"softmax","method_name":"Softmax"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2502.20766","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2502.20766"}},"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/bytedance/FlexPrefill","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"ran_honours":1,"unverified":10},"by_repo_kind":{"official":{"samples":10,"ran":0,"repositories":1},"community":{"samples":1,"ran":1,"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":"c7bbee93a299753b","entry":"get_num_warps_stages","repo":"ByteDance-Seed/FlexPrefill","repo_kind":"community","path":"flex_prefill/ops/flex_prefill_attention.py","file_url":"https://github.com/ByteDance-Seed/FlexPrefill/blob/HEAD/flex_prefill/ops/flex_prefill_attention.py","link_basis":"first_harvest_node","language":"python","status":"ran_honours","verification_level":1,"contract_check":"HONOURS","metamorphic_tier":"well_formed","behaviour_fingerprint":true,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"c7bbee93a299753b"}},{"code_sha256_prefix":"059a805ecbeca6af","entry":"apply_rotary_pos_emb","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"flex_prefill/modules/glm/glm_self_attention_foward.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/flex_prefill/modules/glm/glm_self_attention_foward.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":"059a805ecbeca6af"}},{"code_sha256_prefix":"78e17a79809e571b","entry":"get_config_example","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"flex_prefill/modules/patch.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/flex_prefill/modules/patch.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":"78e17a79809e571b"}},{"code_sha256_prefix":"6904e1292ef34ff1","entry":"get_dataloader","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"experiments/benchmark/run_ruler.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/experiments/benchmark/run_ruler.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":"6904e1292ef34ff1"}},{"code_sha256_prefix":"0ff1542f9a950e79","entry":"get_tasks","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"experiments/benchmark/run_ruler.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/experiments/benchmark/run_ruler.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":"0ff1542f9a950e79"}},{"code_sha256_prefix":"5f120e3472fb5048","entry":"glm_mlp_forward","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"flex_prefill/modules/glm/glm_mlp_forward.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/flex_prefill/modules/glm/glm_mlp_forward.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":"5f120e3472fb5048"}},{"code_sha256_prefix":"1c80690cb1a8743d","entry":"glm_self_attention_forward","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"flex_prefill/modules/glm/glm_self_attention_foward.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/flex_prefill/modules/glm/glm_self_attention_foward.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":"1c80690cb1a8743d"}},{"code_sha256_prefix":"1bd20ca4f6cfb4eb","entry":"mid_trunc_and_pad","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"experiments/benchmark/utils.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/experiments/benchmark/utils.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":"1bd20ca4f6cfb4eb"}},{"code_sha256_prefix":"2742a6153e7ec429","entry":"remove_duplicates_by_index","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"experiments/benchmark/run_ruler.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/experiments/benchmark/run_ruler.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":"2742a6153e7ec429"}},{"code_sha256_prefix":"08b693a1068fa690","entry":"split_tensor_along_last_dim","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"flex_prefill/modules/glm/glm_self_attention_foward.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/flex_prefill/modules/glm/glm_self_attention_foward.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":"08b693a1068fa690"}},{"code_sha256_prefix":"35b094727218cc43","entry":"str_to_dict","repo":"bytedance/FlexPrefill","repo_kind":"official","path":"experiments/benchmark/utils.py","file_url":"https://github.com/bytedance/FlexPrefill/blob/HEAD/experiments/benchmark/utils.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":"35b094727218cc43"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}