FlashSchNet: Fast and Accurate Coarse-Grained Neural Network Molecular Dynamics

Motivation

• Known: GNN 기반 포텐셜(SchNet 등)은 분자동역학의 정확도와 전이성을 향상시키지만, 단편화된 커널과 메모리 병목으로 인해 고전적 포스필드보다 느리다. PyTorch와 JAX 같은 고수준 프레임워크에서 SchNet 구현은 중간 텐서를 GPU HBM에 반복 저장하고 scatter 집계에서 원자적 경쟁으로 인한 오버헤드를 겪는다.
• Gap: 기존 GNN-MD 구현은 GPU HBM과 on-chip SRAM 간의 IO 효율성을 고려하지 않아 메모리 대역폭 활용이 불충분하고, CGSchNet은 단 2.5% 모델 FLOPs 활용도만 달성한다. 또한 기존 GNN-MD는 고전적 coarse-grained 포스필드(MARTINI)보다 실제 벽시간 시뮬레이션 속도에서 뒤떨어진다.
• Why: 분자동역학 시뮬레이션은 신약개발, 재료과학 등에 필수적이나, 접근 가능한 시간 규모와 시스템 크기가 포스 계산 병목에 의해 제한된다. GNN 포텐셜이 학습 기반의 정확도 이점을 실제 워크플로우에서 활용하려면 고전적 포스필드와 경쟁력 있는 속도가 필수적이다.
• Approach: IO 바운드 문제를 root cause로 진단하고, radial basis 확장, message passing, scatter 집계, MLP 가중치 양자화를 각각 IO 인식 커널로 재설계하여 중간 텐서 저장을 제거하고 메모리 트래픽을 감축한다.

Achievement

Figure 1: Left: Memory-throughput trade-off for SchNet-style GNN-MD. FlashSchNet achieves 5×

• Flash radial basis: pairwise 거리 계산, Gaussian basis 확장, cosine cutoff을 단일 tiled pass로 융합하여 각 거리를 한 번만 계산하고 on-chip SRAM에서 모든 basis 함수에 재사용
• Flash message passing: cutoff masking, neighbor gather, filter 곱셈, reduction을 하나의 커널로 융합하여 edge 텐서 O(E×F)를 HBM에 저장하지 않음
• Flash aggregation: scatter-add를 CSR segment reduce로 재공식화하여 원자적 쓰기를 feature dimension만큼 감소시키고 forward/backward 모두에서 contention-free 누적
• Channel-wise 16-bit 양자화: SchNet MLP 가중치의 낮은 per-channel 동적 범위를 활용하여 처리량 개선, 정확도 손실 무시할 수준
• 성능: RTX PRO 6000 단일 GPU에서 269-bead 단백질 64개 병렬 복제본에 대해 1000 ns/day 달성 (CGSchNet 대비 6.5배, MARTINI와 동등 수준), 메모리 피크 80% 감축

How

• Radial basis: 거리 $r_{ij}$ 계산과 Gaussian basis $\\exp(-\\alpha_k r_{ij}^2)$, cosine envelope을 tiling 전략으로 동일 SRAM 블록에서 처리
• Message passing: 동적 neighbor list를 기반으로 cutoff mask, gather 인덱스, filter MLP 적용, 결과 reduction을 통합 커널에서 순차 처리
• Aggregation: CSR(Compressed Sparse Row) 형식으로 neighbor graph를 인코딩하여 segment reduce를 원자적 연산 대신 사용, 각 atom별 feature별로 직렬 누적
• 양자화: SchNet 필터 네트워크의 각 채널 가중치를 per-channel 통계에 따라 FP16으로 양자화, 역전파 시에도 축적 오차 최소화
• 병렬화: 64개 독립적 분자 복제본을 배치 차원으로 처리하여 GPU 활용도 극대화, 각 복제본의 neighbor graph는 동적으로 업데이트

Originality

• IO 인식(IO-aware) 원칙을 GNN-MD에 처음 도입하여 메모리 계층 구조(HBM↔SRAM) 최적화를 체계적으로 적용
• Flash attention 등 기존 fusing 기법을 molecular dynamics의 동적 neighbor graph 특성에 맞게 재설계 (scatter-add → CSR segment reduce 변환, per-channel 양자화)
• SchNet 아키텍처의 구조적 특성(sparse interactions, low per-channel dynamic range)을 활용한 algorithmic 수준의 메모리 트래픽 감축
• GNN 포텐셜이 고전적 포스필드를 속도에서 추월한 첫 사례 (1000 ns/day는 MARTINI 대비 경쟁력 있는 수준)

Limitation & Further Study

• 일반화: 현재 구현은 SchNet 스타일 CFConv 아키텍처에 특화되어 있으며, NequIP, Allegro, MACE 같은 더 정확한 E(3)-equivariant 모델로의 확장은 명확하지 않음
• 모델 크기: 단일 RTX PRO 6000의 메모리 제약으로 인해 더 큰 단백질 시스템이나 높은 feature dimension에서의 확장성 미검증
• 정확도-속도 트레이드오프: 16-bit 양자화는 일부 모델에서 정확도 저하 가능성 있으며, 매우 높은 정확도가 필요한 응용에서는 부적합할 수 있음
• 동적 neighbor list 갱신: 매 timestep마다 neighbor list를 재계산하는 오버헤드가 분석되지 않았으며, 이것이 전체 성능에 미치는 영향 불명확
• 후속 연구: (1) E(3)-equivariant 모델에 대한 IO-aware 구현, (2) 멀티-GPU 확장성 검증, (3) 더 공격적인 양자화 기법(INT8, dynamic quantization) 탐색

Evaluation

총평: FlashSchNet은 GNN 기반 분자동역학의 실용성을 획기적으로 향상시키는 체계적이고 정교한 IO-aware 최적화 연구로, 세부 커널 설계부터 end-to-end 성능 검증까지 우수한 완성도를 보인다. 기존 GNN 포텐셜을 고전적 포스필드 수준의 속도로 끌어올린 첫 사례로서 계산 생화학 분야에 즉시 영향을 미칠 것으로 기대된다.