kUPS: A Molecular Simulation Engine for the AI Era
Gao, Köhler, Hanke, Ramanan (CuspAI) — Medium Technical Blog, 2026年4月 | OSS: github.com/cusp-ai-oss/kups
🎯 MD/MC/幾何最適化+ML力場をJAX-native・composableに統合。RASPA比49×高速。PyTorch→JAX変換Tojaxも同時公開
① 背景と課題

分子シミュレーションは創薬・材料科学の根幹だが、現状は GROMACS・LAMMPS・RASPA・ASE 等を組み合わせる複雑なワークフローが必要で、各ツールは独自フォーマット・データ構造を持ち変換グルーコードに多大な時間がかかる。

AI 時代の ML ポテンシャル・Boltzmann ジェネレーターを従来 C++ エンジンに統合するには複雑なクロス言語ブリッジが必要で、AI 研究ペースに追いつけない
従来ツールは CPU 最適化設計。GPU サポートは後付けで rigid なインターフェース経由に限られ、batched 実行（複数系並列）を活かせない
新手法を試すたびに別パッケージへの全面移行が必要。1つの新アイデア実装に数週間を要する

→ JAX を基盤とする Python-native エンジンで、MD・MC・ML力場・生成モデルを単一の Propagator 抽象に統一。任意の組み合わせが数行で実現

② kUPS 設計3原則
Propagator: state → next_state（すべての演算子が同一シグネチャ）
  • Composability: 積分器・サーモスタット・MC ムーブ・ML force field が同一 API → 自由に snap together
  • Structured data: columnar Table（structure-of-arrays）→ GPU coalesced memory access を最大化。primary key/foreign key でテーブル間参照
  • Batched execution: leading dimension に system index → jax.vmap で複数系を自動ベクトル化並列
Tojax
任意の PyTorch nn.Module を1行で JAX 関数に変換。jit/vmap/grad 完全互換。実行時オーバーヘッドなし（XLA でネイティブコンパイル）
③ 主な結果 (a) スループット比較
スループット比較 (NVIDIA L4 GPU, MOF系) NVT MD (MACE, 114原子 MOF) ASE × 1.0 (baseline) kUPS × 1.8 ✓ 1.8× µVT GCMC (CO₂吸着, 3078原子 MOF) RASPA (CPU) ×1 kUPS × 49 ✓✓ 49× CPUシリアルに縛られたGCMCがGPU nativeで49倍高速化
③ 主な結果 (b) 対応手法・ポテンシャル一覧
カテゴリ手法
MD アンサンブルNVE, NVT, NPT
Monte CarloNVT-MC, µVT-MC (GCMC)
幾何最適化FIRE, L-BFGS
ML force fieldMACE, UMA (+ Tojax で任意 PyTorch)
古典ポテンシャルLennard-Jones, Coulomb (Ewald), 調和結合/角度, Morse
再現性ビットワイズ再現（JAX random key 明示渡し）

すべて同一 Propagator 抽象 → 任意の組み合わせが自動的に動作

③ 主な結果 (c) コード複雑性の比較
タスク従来ツールkUPS
ML力場でNVT MD数百行 + 複数ファイル数行の Python
ML力場の交換別パッケージに移行1行の変更
MC + MD 組み合わせグルーコード数十行数行追加
GPU バッチ実行別途実装が必要自動（jax.vmap）
PyTorch モデル統合C++ ブリッジ + 数週間tojax(model) 1行

新機能実装: weeks → afternoon に短縮

⑤ 限界点・今後の展開
  • ベンチマークは MOF 系のみ — タンパク質-リガンド・水溶液系での検証は今後
  • JAX の Pure function / 静的形状制約への適応が既存ユーザーの学習コストに
  • MDAnalysis/MDTraj との完全な入出力互換アダプターは利用者側の実装が必要
  • NVIDIA ALCHEMI プラットフォームとの統合を継続開発中
⑥ lib/md への統合シナリオ
MD バックエンド置き換え
RMSDAnalyzer・HBondAnalyzer のバックエンドを kUPS Propagator に変更。同一解析コードで GPU-native batched MD を実行。複数リガンド系の並列 MD が自動化
ML力場の柔軟統合
Tojax で既存 PyTorch 製 MACE/NequIP を1行変換。kUPS 内で jit/vmap 適用済みで追加コード不要
高スループット GCMC
RASPA 比 49× の µVT-MC で溶媒和自由エネルギー・結合ポケット水分子解析を高速化。lib/fep の MMGBSAEngine の前処理に活用
ビットワイズ再現性
JAX の random key 設計により、random_seed 固定での完全再現性をすべての MD/MC で保証。テスト・デバッグコストを大幅削減
⑦ 実装優先度・まとめ
  • Priority: High — lib/md バックエンド高速化への即時適用価値
  • target_module: lib/md（lib/fep にも波及）
  • OSS 公開済み: github.com/cusp-ai-oss/kups + tojax
  • RASPA 比 49×・ASE 比 1.8×（L4 GPU 単体）
X投稿用（120字）
CuspAI がJAX製分子シミュレーションエンジン kUPS をOSS公開。MD/MC/幾何最適化+ML力場を同一APIで統合、RASPA比49倍高速。PyTorch→JAX変換Tojaxも同時公開。 #MolecularDynamics #MLFF
custom:cuspai-2026-kups | CuspAI (Gao, Köhler, Hanke, Ramanan) | 2026-04 | Category: computational_chemistry | lib/md, lib/fep