{"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-2605-31500","title":"On Efficient Scaling of GNNs via IO-Aware Layers Implementations","arxiv_id":"2605.31500","date":"2026-05-29","proceeding":"ICML","authors":["Daria Fomina","Daniil Krasylnikov","Alexey Boykov","Andrey Dolgovyazov","Vyacheslav Zhdanovskiy","Fedor Velikonivtsev"],"abstract":"Graph Neural Networks (GNNs) are bottlenecked by sparse, irregular memory access. Popular frameworks such as DGL and PyTorch Geometric support general message passing, but complex layers often materialize edge-wise intermediates, increasing memory traffic and limiting scalability on large graphs. We take an I/O- and arithmetic-intensity--centric view and show that widely used layers fall into three kernel families: SpMM-based convolutions, reduction-based aggregations, and attention-based layers (GATv2/Graph Transformer). For each family, we develop GPU kernels that reduce data movement, improve locality, and remain robust across realistic graphs. We also study graph reordering and find that its impact depends on the kernel mapping: it benefits neighbor-parallel (gather-dominated) kernels more consistently than feature-parallel designs. Empirically, our fused attention kernels reach up to $\\textbf{3.9}\\times$ speedup for Graph Transformer (median $\\textbf{1.6}\\times$), with Tensor Core (block-sparse) variants up to $\\textbf{7.3}\\times$ on locally dense graphs; for GATv2 we reach up to $\\textbf{8.5}\\times$ speedup (median $\\textbf{2.0}\\times$) while reducing peak memory by up to $\\textbf{76}\\times$ (median $\\textbf{6}\\times$). Our degree-aware reduction kernels achieve up to $\\textbf{10}\\times$ speedup (median $\\textbf{2.6}\\times$). For SpMM-based layers, properly cached cuSPARSE achieves up to $\\textbf{8}\\times$ speedup over DGL and outperforms evaluated custom baselines in the majority of evaluations. We release our implementations as drop-in replacements to support reproducible, hardware-aware GNN acceleration.","url_abs":"https://arxiv.org/abs/2605.31500","url_pdf":"https://arxiv.org/pdf/2605.31500","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=2605.31500","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2605.31500"}},"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/yandex-research/On-Efficient-Scaling-Of-GNNs","reach":null}],"summary":{"ran":10,"unverified":3},"by_repo_kind":{"found_in_text":{"samples":13,"ran":10,"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":"0d4f54fe2894720f","entry":"activation_factory","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/models/base.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/models/base.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"0d4f54fe2894720f"}},{"code_sha256_prefix":"0f2a52d87e993417","entry":"compute_accuracy","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/training/metrics.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/training/metrics.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"0f2a52d87e993417"}},{"code_sha256_prefix":"5c551f60830ac131","entry":"compute_f1","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/training/metrics.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/training/metrics.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"5c551f60830ac131"}},{"code_sha256_prefix":"cd6f55a7a1817bde","entry":"compute_precision","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/training/metrics.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/training/metrics.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"cd6f55a7a1817bde"}},{"code_sha256_prefix":"6889b571f6b4480a","entry":"human_bytes","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/benchmarking/memory.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/benchmarking/memory.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"6889b571f6b4480a"}},{"code_sha256_prefix":"0b3c340b374bd042","entry":"make_schedule","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/benchmarking/profiler.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/benchmarking/profiler.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"0b3c340b374bd042"}},{"code_sha256_prefix":"123fd36fe2eb70d5","entry":"model_param_and_buffer_bytes","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/benchmarking/memory.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/benchmarking/memory.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"123fd36fe2eb70d5"}},{"code_sha256_prefix":"22d3a8460297800d","entry":"norm_factory","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/models/base.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/models/base.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"22d3a8460297800d"}},{"code_sha256_prefix":"964dcee138e3fb6a","entry":"register","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/models/registry.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/models/registry.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"964dcee138e3fb6a"}},{"code_sha256_prefix":"33866dae05556ec2","entry":"tensor_nbytes","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/benchmarking/memory.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/benchmarking/memory.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":true,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"33866dae05556ec2"}},{"code_sha256_prefix":"b484b587dc7496d6","entry":"build","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/models/registry.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/models/registry.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"b484b587dc7496d6"}},{"code_sha256_prefix":"4c6a28957b22dbb4","entry":"get_gpu_info","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/benchmarking/microbench.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/benchmarking/microbench.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"4c6a28957b22dbb4"}},{"code_sha256_prefix":"eb879070473f1596","entry":"measure_memory","repo":"yandex-research/On-Efficient-Scaling-Of-GNNs","repo_kind":"found_in_text","path":"src/benchmarking/microbench.py","file_url":"https://github.com/yandex-research/On-Efficient-Scaling-Of-GNNs/blob/HEAD/src/benchmarking/microbench.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"BSD-3-Clause","inline_ok":true,"mcp_get_code":{"code_sha256":"eb879070473f1596"}}]},"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}