I/O 인지형 레이어 구현을 통한 GNN의 효율적인 스케일링에 관하여

그래프 신경망 (Graph Neural Networks, GNNs)은 희소하고 불규칙한 메모리 접근으로 인해 병목 현상이 발생합니다. DGL 및 PyTorch Geometric과 같은 대중적인 프레임워크는 일반적인 메시지 패싱 (message passing)을 지원하지만, 복잡한 레이어들은 종종 에지 단위 (edge-wise)의 중간 결과물을 실체화하여 메모리 트래픽을 증가시키고 대규모 그래프에서의 확장성을 제한합니다. 우리는 I/O 및 연산 강도 (arithmetic-intensity) 중심의 관점을 취하며, 널리 사용되는 레이어들이 세 가지 커널 제품군으로 분류됨을 보여줍니다: SpMM 기반 컨볼루션 (convolutions), 리덕션 (reduction) 기반 집계 (aggregations), 그리고 어텐션 (attention) 기반 레이어 (GATv2/Graph Transformer). 각 제품군에 대해, 우리는 데이터 이동을 줄이고, 지역성 (locality)을 개선하며, 실제적인 그래프 환경에서도 견고하게 작동하는 GPU 커널을 개발합니다. 또한 그래프 재정렬 (graph reordering)을 연구하였으며, 그 영향이 커널 매핑에 따라 달라진다는 것을 발견했습니다: 즉, 재정렬은 피처 병렬 (feature-parallel) 설계보다 이웃 병렬 (neighbor-parallel, gather-dominated) 커널에 더 일관되게 이점을 제공합니다. 실증적으로, 우리의 퓨즈드 어텐션 (fused attention) 커널은 Graph Transformer에 대해 최대 $\textbf{3.9}\times$의 속도 향상 (중앙값 $\textbf{1.6}\times$)을 달성하였으며, 텐서 코어 (Tensor Core, block-sparse) 변형 모델은 국소적으로 밀집된 그래프에서 최대 $\textbf{7.3}\times$의 성능 향상을 보였습니다. GATv2의 경우, 피크 메모리 (peak memory)를 최대 $\textbf{76}\times$ (중앙값 $\textbf{6}\times$)까지 줄이면서 최대 $\textbf{8.5}\times$의 속도 향상 (중앙값 $\textbf{2.0}\times$)을 달성했습니다. 우리의 차수 인지형 (degree-aware) 리덕션 커널은 최대 $\textbf{10}\times$의 속도 향상 (중앙값 $\textbf{2.6}\times$)을 달성합니다. SpMM 기반 레이어의 경우, 적절히 캐싱된 cuSPARSE는 DGL 대비 최대 $\textbf{8}\times$의 속도 향상을 달성하였으며, 대부분의 평가에서 검증된 커스텀 베이스라인 (custom baselines)보다 우수한 성능을 보였습니다. 우리는 재현 가능하고 하드웨어 인지적인 (hardware-aware) GNN 가속을 지원하기 위해, 우리의 구현체를 즉시 교체 가능한 (drop-in replacements) 형태로 공개합니다.