{"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/arxiv-2602-06283","title":"SOCKET: SOft Collision Kernel EsTimator for Sparse Attention","arxiv_id":"2602.06283","date":"2026-02-06","proceeding":null,"authors":["Sahil Joshi","Agniva Chowdhury","Wyatt Bellinger","Amar Kanakamedala","Ekam Singh","Hoang Anh Duy Le","Aditya Desai","Anshumali Shrivastava"],"abstract":"Exploiting sparsity during long-context inference is key to scaling large language models, as attention dominates the cost of autoregressive decoding. Sparse attention reduces this cost by restricting computation to a subset of tokens, but its effectiveness depends on efficient scoring and selection at inference time. We revisit Locality-Sensitive Hashing (LSH) and introduce SOCKET, a SOft Collision Kernel EsTimator that replaces hard bucket matches with probabilistic, similarity-aware aggregation. Traditional LSH yields binary collision signals that limit ranking quality and require substantial memory to perform well. In contrast, soft LSH accumulates graded collision evidence across hash tables, preserving top-k ordering with significantly less memory. This reframes LSH from a candidate generator into a principled scoring kernel for sparse attention. Leveraging this property, SOCKET enables efficient token selection without ad hoc voting and matches or surpasses prior sparse attention methods across multiple long-context benchmarks. With a custom CUDA scoring kernel and a Flash Decode Triton backend, SOCKET achieves up to 1.5$\\times$ higher throughput than FlashAttention.","url_abs":"https://arxiv.org/abs/2602.06283","url_pdf":"https://arxiv.org/pdf/2602.06283","source":{"archive":null,"snapshot":"2025-07-28","note":"not in the Papers with Code archive (frozen at the snapshot)","row_kind":"graph","title_abstract_authors_date":"arXiv metadata, CC0 1.0 (https://info.arxiv.org/help/license)"},"code_links":[],"tasks":[],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2602.06283","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2602.06283"}},"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":[{"mentioned_in_github":null,"is_official":null,"provenance":"deterministic:regex_extraction","mentioned_in_paper":null,"url":"https://github.com/amarka8/SOCKET","reach":null}],"summary":{"ran_honours":1,"unverified":19},"by_repo_kind":{"found_in_text":{"samples":20,"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":20,"samples":[{"code_sha256_prefix":"aba86ffbed27ba65","entry":"find_multiple","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/model.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/model.py","link_basis":"plan_row","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":"aba86ffbed27ba65"}},{"code_sha256_prefix":"1102ffb719fa5cde","entry":"apply_rotary_emb","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/model.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/model.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":"1102ffb719fa5cde"}},{"code_sha256_prefix":"a75ad681dca07d95","entry":"attention_mask_to_allowed_prob","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"kernels/socket_triton_kernels.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/kernels/socket_triton_kernels.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":"a75ad681dca07d95"}},{"code_sha256_prefix":"00ff5b090fec69bd","entry":"build_results_row","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"eval/ruler_utils/scorer.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/eval/ruler_utils/scorer.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":"00ff5b090fec69bd"}},{"code_sha256_prefix":"e5bb10b75c0b2e3c","entry":"build_sparse_list_decode","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/kernels/sparse.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/kernels/sparse.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":"e5bb10b75c0b2e3c"}},{"code_sha256_prefix":"55ec8fea60e1d7e7","entry":"dynamically_quantize_per_channel","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/quantize.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/quantize.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":"55ec8fea60e1d7e7"}},{"code_sha256_prefix":"caaea88e7a48e0c2","entry":"expand_block_index_to_heads","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"kernels/block_sparse_flash_attention.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/kernels/block_sparse_flash_attention.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":"caaea88e7a48e0c2"}},{"code_sha256_prefix":"dbabfb89a4b530ac","entry":"get_group_qparams","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/quantize.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/quantize.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":"dbabfb89a4b530ac"}},{"code_sha256_prefix":"75282665af7ff147","entry":"get_kernel_code","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/kernels/soft_hash_collision_loader.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/kernels/soft_hash_collision_loader.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":"75282665af7ff147"}},{"code_sha256_prefix":"9545ad64a8415458","entry":"load_soft_hash_collision","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/kernels/soft_hash_collision_loader.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/kernels/soft_hash_collision_loader.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":"9545ad64a8415458"}},{"code_sha256_prefix":"fd0cfe048b12c6f2","entry":"logits_to_probs","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/generate.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/generate.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":"fd0cfe048b12c6f2"}},{"code_sha256_prefix":"060506d7f037e907","entry":"multinomial_sample_one_no_sync","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/generate.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/generate.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":"060506d7f037e907"}},{"code_sha256_prefix":"146cdcef0c195fd7","entry":"pack_scales_and_zeros","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/quantize.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/quantize.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":"146cdcef0c195fd7"}},{"code_sha256_prefix":"26cf476e99a51b14","entry":"precompute_freqs_cis","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/model.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/model.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":"26cf476e99a51b14"}},{"code_sha256_prefix":"b8b875125635b5fc","entry":"random_walk_indices","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"kernels/block_sparse_flash_attention.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/kernels/block_sparse_flash_attention.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":"b8b875125635b5fc"}},{"code_sha256_prefix":"7ff5f8665dfef840","entry":"random_walk_indices_from_block_probs","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"kernels/block_sparse_flash_attention.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/kernels/block_sparse_flash_attention.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":"7ff5f8665dfef840"}},{"code_sha256_prefix":"89f735317c12c997","entry":"register_args_and_configs","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"pipeline/main_utils.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/pipeline/main_utils.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":"89f735317c12c997"}},{"code_sha256_prefix":"b0856105fbfe6f4f","entry":"sample","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/generate.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/generate.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":"b0856105fbfe6f4f"}},{"code_sha256_prefix":"d86a7d51ce10c6e4","entry":"set_logger","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"pipeline/main_utils.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/pipeline/main_utils.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":"d86a7d51ce10c6e4"}},{"code_sha256_prefix":"42bbf5633ac60efd","entry":"sparse_attention_fwd","repo":"amarka8/SOCKET","repo_kind":"found_in_text","path":"GPT-FAST/kernels/sparse.py","file_url":"https://github.com/amarka8/SOCKET/blob/HEAD/GPT-FAST/kernels/sparse.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":"42bbf5633ac60efd"}}]},"arxiv_metadata":{"licence":"arXiv metadata, CC0 1.0 (https://info.arxiv.org/help/license)","fields":["title","abstract","authors","date"],"primary_category":"cs.LG","source":"arxiv_2026.jsonl"},"syntology_extracted_results":null}