FlashEKGR: Fast Embedding-Based Knowledge Graph Reasoning Models Training

Embedding-based Knowledge Graph Reasoning (EKGR) models embed entities, relations, and queries into a unified vector space for logical reasoning. However, their training suffers from severe inefficiencies due to redundant computations, I/O bottlenecks, and unoptimized GPU kernels. To address this, we propose FlashEKGR, a multi-level co-optimization framework designed to unleash the full potential of modern hardware for EKGR training. At the logical level, we propose Logical Operator Parallelism (LOP) to eliminate redundant computations and I/O by merging common logic operator prefixes in query structures. At the dataflow level, we design a Staleness-Free Pipeline (SFP) to effectively overlap the latency of data preparation with model computation. At the GPU execution level, we introduce two key techniques: Layout-Aware Kernel Design (LAK), which resolves the memory-access inefficiency caused by the asymmetric data layouts inherent in negative samples, and Dynamic CUDA Graph Caching (DGC), which significantly reduces kernel launch overhead for dynamic query structures. Extensive experiments on public datasets show that FlashEKGR achieves end-to-end speedups ranging from 2.06× to 5.67× over state-of-the-art frameworks, while maintaining reasoning accuracy. Our code is released at https://github.com/BUPT-Reasoning-Lab/FlashEKGR.