FuseSampleAgg：面向小批量 GNN 的融合邻居采样与聚合
机器学习
2025-11-18 v1
摘要
我们提出了 FuseSampleAgg，这是一个 CUDA 算子，它将邻居采样和均值聚合融合到单次传递中，用于一跳和两跳的 GraphSAGE。通过消除块物化和额外的内核启动，FuseSampleAgg 减少了内存流量和开销，同时通过保存的索引重放保留了 GraphSAGE 的均值语义。在 Reddit、ogbn-arxiv 和 ogbn-products 基准测试（批量大小 1024，启用自动混合精度）中，我们观察到在 ogbn-products 上步时间加速高达 51 倍，在 Reddit 上使用扇出 10-10 和 15-10 时加速约 4 倍，在 ogbn-arxiv 上使用较大扇出时加速约 3.3 倍，同时峰值 GPU 内存分别减少了高达 100 倍、36 倍和约 3.5 倍。该算子是确定性的，可与标准 PyTorch 优化器集成，并附带可从 CSV 日志重现所有表格和图形的脚本。代码和脚本可在 https://github.com/SV25-22/FuseSampleAgg 获取。
引用
@article{arxiv.2511.13645,
title = {FuseSampleAgg: Fused Neighbor Sampling and Aggregation for Mini-batch GNNs},
author = {Aleksandar Stanković},
journal= {arXiv preprint arXiv:2511.13645},
year = {2025}
}
备注
15 pages. Code and reproducibility scripts: https://github.com/SV25-22/FuseSampleAgg