Triton
Python 風の DSL で GPU カーネルを記述し、ブロック単位のタイル抽象で自動的に並列化・メモリ管理する OpenAI 由来のコンパイラ。cuBLAS/cuDNN の固定実装に縛られず、行列サイズに合わせたタイル戦略や任意のカーネル融合を書ける。
Realtime-VLA での役割
realtime-vla / vla-inference-graph-optimization では、pi0 推論を24の GEMM に分解し各々を Triton で実装。依存は PyTorch + Triton のみ。
- タイルチューニング:
BLOCK_SIZE_M/N/Kを手動探索し、cuBLAS が最適配分しないサイズで約1.5 ms 改善。 - カーネル融合:
rms_matmul_n_2048_2560_qkv_rope（RMSNorm+QKV+RoPE）、matmul_small_res_gate（matmul+残差+gate）等、複数演算を1カーネルに。LeRobot の ~21 kernels/層 を ~5 kernels/層 に削減。 - Partial Split-k: 512×1152×1152 を均等分配可能な2つの matmul に分割し1カーネル化。
- ソフトウェアバリア:
tl.atomic_add/tl.atomic_orで SM 数ぶんのブロックを起動し全ブロック同期を実現（roofline-model の同期コスト下限推定に使用）。