SABLE: GPU-Based Power Flow Accelerator for Sparsity-Aware Batched Learning

Technical Summary: SABLE - GPU-Based Power Flow Accelerator for Sparsity-Aware Batched Learning

Problem Statement
The increasing integration of distributed energy resources (DERs) necessitates high-dimensional, scenario-driven grid security assessments, requiring the solution of thousands to tens of thousands of AC power flow (PF) instances with low latency. While recent studies have successfully accelerated standalone PF problems using GPUs, integrating these approaches into modern differentiable learning frameworks (such as DC3 and DeepLDE) remains challenging. Existing physics-embedded learning systems lack a sparsity-preserving execution path capable of sustaining high throughput during repeated training-time use. Consequently, current methods rely on dense, memory-intensive GPU computations, creating a dominant bottleneck in both runtime and memory consumption.

Three specific limitations hinder progress:

1. Structural Barriers: There is a gap between efficient standalone PF execution and the end-to-end differentiation required for training. Traditional 3D batched sparse representations do not map cleanly to PyTorch's sparse execution model, which lacks robust support for 3D sparse tensors and backward paths.
2. Performance Limits in Repeated Execution: Repeated Jacobian construction and linear equation solving during Newton-Raphson (NR) iterations and training steps incur significant overhead due to redundant sparse-pattern construction, symbolic reanalysis, and memory reallocation.
3. Precision Mismatch: Modern GPUs offer significantly higher throughput for single-precision (SP) and mixed-precision (MP) operations compared to double-precision (DP). However, PF solvers traditionally rely on DP, creating a performance mismatch that is exacerbated when PF layers are embedded in learning loops.

Methodology
The authors present SABLE (Sparsity-aware Accelerator for a Batched Learning Engine), a GPU-based framework designed to execute batched AC power flow within differentiable learning loops while preserving sparsity. The methodology consists of three core components:

1. Block-Diagonal Embedding (BDE) and Zero-Copy Interoperability:

   • SABLE reformulates batched 3D Jacobians into a single, fixed-pattern 2D sparse template using Block-Diagonal Embedding. This transforms $N_{batch}$ independent sparse matrices into one global sparse matrix where each scenario occupies a block on the diagonal.
   • This shared template enables zero-copy interoperability across PyTorch, CuPy, and cuDSS via DLPack and device-pointer binding. Dense state vectors and sparse matrices are shared without redundant memory allocation, keeping the entire pipeline GPU-resident.

2. One-Time Symbolic Analysis and Repeated Numeric Execution:

   • Symbolic Setup: SABLE performs symbolic analysis (sparsity pattern construction, indexing metadata, and LU analysis state) only once for a given topology and batch size.
   • In-Place Value Updates: Custom CUDA kernels update only the numerical values of the Jacobian entries based on cached metadata, avoiding repeated sparse structure reconstruction.
   • Block-Aware Indexing: A custom indexing kernel extracts the specific reduced Jacobians required for forward ($J_1$) and backward ($J_1^T, J_2^T, J_3^T$) passes directly from the base BDE template. This eliminates the need for explicit transpose operations or repeated memory reallocation.
   • Sparse-Direct Solving: The framework utilizes cuDSS for sparse-direct LU refactorization and solving, reusing the unchanged symbolic structure across iterations.

3. Stagewise Mixed-Precision (MP) Strategy:

   • To balance numerical robustness with GPU throughput, SABLE employs a stagewise precision policy. The outer nonlinear logic and residual calculations remain in Double Precision (FP64) to ensure convergence and accuracy.
   • The computationally dominant linear equation solving stages (LU solve) are executed in Single Precision (FP32), leveraging the higher throughput of modern GPU Tensor Cores.

Key Contributions

• Shared Sparse Template: The introduction of BDE allows for a reusable, fixed-pattern 2D sparse template that supports differentiable batched PF, enabling zero-copy interoperability across the software stack.
• Efficient Sparse-Direct Execution: By decoupling symbolic analysis from numeric execution, SABLE avoids redundant memory allocation and sparse-pattern reconstruction, utilizing in-place value updates and block-aware indexing.
• Stagewise MP: The framework implements a hybrid precision strategy that accelerates linear solving without compromising the numerical stability required for PF convergence.

Experimental Results
The authors evaluated SABLE on both standalone PF solving and end-to-end training scenarios using the DC3 and DeepLDE frameworks.

• Standalone Performance: SABLE demonstrated significant speedups over existing baselines. Compared to the CPU-based pandapower, SABLE achieved throughput improvements of up to 253.4$\times$. Compared to the GPU-based ExaPF, SABLE achieved up to 5.7$\times$ speedup. Notably, SABLE remained executable at batch sizes where ExaPF encountered Out-Of-Memory (OOM) errors.
• Component Acceleration: Analysis showed SABLE reduced Jacobian construction time by 8.5$\times$ to 33.2$\times$ and linear solving time by 1.6$\times$ to 3.6$\times$ compared to ExaPF.
• Training Efficiency: In end-to-end training, SABLE expanded the feasible training batch range by up to 64$\times$ (e.g., from batch size 1 to 64 for a 7336-bus system in DC3).
• Throughput: Training throughput improved by up to 206.7$\times$ over the corresponding baseline (DC3) and 64.5$\times$ for DeepLDE.
• Numerical Fidelity: Experiments confirmed that SABLE maintains numerical agreement with reference solutions, with maximum absolute errors below $5.60 \times 10^{-9}$ and convergence behavior identical to double-precision baselines.

Significance
The paper claims that SABLE addresses the critical bottlenecks of memory consumption and runtime in physics-embedded learning for power systems. By providing a sparsity-preserving, GPU-resident execution path, SABLE enables the training of large-scale models that were previously infeasible due to memory constraints or computational cost. The framework is presented as a practical solution for accelerating optimization and learning loops that require repeated batched power flow solving and differentiation.