검색 상세

FlashSC: Path-Specialized Product SC Executor for MACE Inference

  • 안성배 (Seongbae An, 아주대학교 일반대학원)

초록/요약

Machine-learning interatomic potentials （MLIPs）, such as MACE, enable high-accuracy molecular-dynamics （MD） simulations, but their practical efficiency depends on the la- tency of repeated force and energy evaluations. Although cuEquivariance accelerates MACE using fused just-in-time GPU kernels, our profiling across nine pre-trained MACE checkpoints shows that the Product block, which performs Symmetric Contraction （SC）, still accounts for 27.4–57.6% of forward latency. This thesis presents FlashSC, a load-time path-specialized CUDA code generator for MACE Symmetric Contraction. FlashSC exploits the fact that the SC path topology of a trained checkpoint is fixed after model loading, while input features and atom indices change at each MD step. Instead of passing generic path metadata to a runtime ex- ecutor at every inference call, FlashSC extracts the polynomial descriptor once, generates checkpoint-specific CUDA code with path indices and coefficients encoded as compile-time literals, and reuses the generated executor during repeated MD evaluations. FlashSC pre- serves the original MACE computation, leaves non-Product blocks unchanged, and falls back to cuEquivariance when specialization is not applicable. We evaluate FlashSC on an NVIDIA RTX 4090 across nine MACE checkpoints and inside LAMMPS ML-IAP. FlashSC accelerates isolated Product SC by 5.1×–9.6×, im- proves full-model forward inference by 1.23×–1.97×, and reduces end-to-end LAMMPS MD-step latency by 1.21×–1.83× while preserving Product SC numerical agreement with cuEquivariance. The benefit is largest in the low-to-moderate local-atom-count regime represented by the evaluated workloads and diminishes as per-atom kernel work domi- nates at large supercell scales. Profiling shows that the speedup is primarily attributable to reduced generic per-call executor overhead, not to a faster GPU compute kernel. These results suggest that load-time specialization is an effective optimization strategy for re- peatedly executed operators with fixed topology, and that MACE Symmetric Contraction is a practical case where this strategy improves MD inference latency.

more

목차

1 Introduction 1
1.1 The Persistent Product Bottleneck 2
1.2 FlashSC: Specializing at Load Time 3
1.3 Thesis Organization 5
2 Background and Related Work 6
2.1 MACE and Symmetric Contraction 6
2.1.1 Computation of Symmetric Contraction 7
2.1.2 Equivariant Feature Layout 8
2.1.3 Body Order and Product SC Cost 8
2.2 Symmetric Contraction Execution in cuEquivariance 9
2.3 Related Work 10
2.3.1 Equivariant MLIP Kernel Optimization 11
2.3.2 Framework-Level and Operator-Level Execution Optimization 11
2.3.3 Positioning of FlashSC 12
3 Motivation 13
3.1 Product Block Dominance in Forward Latency 14
3.2 Non-Kernel Overhead in Product SC 16
3.3 Non-Kernel Overhead under Supercell Replication 19
3.4 Summary 22
4 FlashSC Design and Implementation 23
4.1 Product SC Formulation 23
4.2 Baseline Execution Path 24
4.3 Load-Time Path Specialization 27
4.3.1 Specialization and Reuse Granularity 27
4.4 Path Extraction and Code Generation 28
4.5 Runtime Execution and Fallback 31
4.5.1 Forward Kernel Execution 31
4.5.2 Backward Pass and Fallback 32
4.6 Numerical Agreement 34
5 Experiments and Evaluation 35
5.1 Experiment Setup 35
5.1.1 Testbed 35
5.1.2 Evaluation Checkpoints 37
5.1.3 Compared Backends 39
5.1.4 Protocols and Metrics 39
5.2 Experiment Results 41
5.2.1 Numerical Agreement 41
5.2.2 Isolated Product SC Speedup 43
5.2.3 Mechanism: Kernel and Non-Kernel Time 44
5.2.4 Full-Model Product Block 45
5.2.5 End-to-End Forward and Force-Backward 45
5.2.6 LAMMPS MD-Step Performance 45
5.2.7 Replicate-Scale Behavior 46
5.2.8 Limitations 47
5.2.9 Discussion: When Specialization Pays 48
6 Conclusion and Future Work 61
6.1 Future Work 61
Bibliography 62

more
반출 Meta View 목록