Parallelizing Large-Scale Tensor Network Contraction on Multiple GPUs

Abstract: Exact tensor network contraction underpins quantum circuit simulation, quantum error correction, combinatorial optimization, and many-body dynamics. The dominant parallelization strategy, slicing, scales exponentially and incurs redundant computation. We present a multi-GPU framework that instead distributes intermediate tensors across devices with explicit communication, converting a fixed contraction path into a communication-efficient schedule via GEMM-oriented mode reordering and communication-aware mode distribution planning. Within a single DGX H100 node (8 GPUs, NVLink), distribution delivers $7$--$173\times$ extra speedup beyond embarrassingly parallel slicing, capturing nearly all of the available compute reduction (87--101%) because NVLink's high bandwidth keeps communication small relative to compute. Scaling the same four workloads to 1024 H100 GPUs over InfiniBand, the extra speedup beyond slicing ranges from $42\times$ to $67{,}869\times$, demonstrating that communication-aware distributed contraction far surpasses slicing-based scaling limits for frontier tensor networks.

• The paper introduces a framework that transforms a fixed contraction path into an optimal multi-GPU execution schedule for large tensor networks.
• The methodology combines GEMM-oriented mode reordering with dynamic programming to minimize communication and maximize computational efficiency on modern clusters.
• Experimental results show unprecedented speedups and scalability, validating the effectiveness of communication-aware distributed contraction over traditional slicing methods.

Parallelizing Large-Scale Tensor Network Contraction on Multiple GPUs

Introduction and Motivation

The contraction of large tensor networks (TNs) is central to modern computational approaches in quantum circuit simulation, quantum error correction (QEC), combinatorial optimization, and quantum many-body dynamics. The system bottleneck traditionally arises not solely from floating-point throughput, but from memory capacity constraints: intermediates in the contraction tree often exceed the limits of a single GPU, rendering standard approaches infeasible for large-scale workloads. Slicing exposes parallelism by dividing the computation into independent sub-tasks, but incurs an exponential scaling in redundant computation and cannot circumvent underlying complexity limitations related to the treewidth of the network.

This work introduces a framework for translating a fixed contraction path into an optimal multi-GPU execution schedule. The framework leverages GEMM-oriented mode reordering and a dynamic-programming-based communication-aware mode distribution strategy. The combined effects are to minimize communication relative to computation, optimally exploit available memory and compute across devices, and enable contractions that are infeasible on any single device.

Figure 1: Theoretical complexity reduction from distributing intermediate tensors across GPUs for six workloads, demonstrating rapid computation cost decrease as device count increases.

Technical Approach

GEMM-Oriented Mode Reordering

A fixed contraction path is described as a sequence of pairwise contractions. To maximally exploit BLAS throughput and minimize memory transpositions, the authors employ a deterministic backward pass that propagates mode lifetime constraints. Each tensor is reordered such that its longest-lived indices reside in the leading memory dimensions, mapping directly to the $[\text{retained} \| \text{reduced}]$ layout required by GEMM routines. This eliminates the need for runtime transpositions and produces a globally consistent layout throughout the contraction tree.

Figure 2: Method for GEMM-oriented mode reordering on a subtree, yielding transposition-free operand layouts aligned with contraction requirements.

Communication-Aware Mode Distribution and Redistribution

Given a memory threshold for single-device feasibility, large intermediates are distributed by selecting the minimal prefix of leading modes that can be evenly partitioned among all GPUs. Prefix selection is natural given the lifetime-ordered mode arrangement, ensuring data contiguity and block granularity for inter-GPU communication, which is essential for efficient data movement on both NVLink (intra-node) and InfiniBand (inter-node) fabrics.

A dynamic-programming search, parameterized by hardware-aware cost models for both computation and communication, determines where elective redistributions (all-to-all shuffles of distributed modes) should be performed. The DP algorithm exploits structural features of contraction trees and input mode lifetime to localize communication at size-minimal points in the tree, avoiding latency-bound transfers, and preventing pathological redistribution in large stride scenarios.

Figure 3: Dynamic programming-based redistribution point selection, distributing communication in favorable regions of the contraction tree on the Zuchongzhi n60m24 benchmark.

The distinction between slicing and distribution is rendered explicit: slicing trades memory for redundant computation, while distribution partitions memory and coordinates explicit communication, fundamentally trading off arithmetic cost for network boundedness.

Figure 4: Conceptual comparison between slicing (independent sub-tasks, redundant computation) and distribution (partitioned intermediates, explicit communication).

The schedule produced by the planner is executed using cuTENSORMp. This library, built upon the NCCL communication substrate, automatically overlaps compute and communication, decomposes large contractions into tiles, and manages per-device and collective communication buffers, all transparent to the user.

Figure 5: End-to-end workflow, illustrating the offline planning (mode reordering, distribution planning) and runtime execution pipeline (cuTENSORMp-based overlapped compute and communication).

Experimental Results

Benchmarking focuses on four application domains: (1) quantum circuit simulation (Zuchongzhi n60m24), (2) QEC (distance-7 rotated surface code), (3) combinatorial problems (King’s graph independent set enumeration), and (4) quantum many-body dynamics (on rectangular, hexagonal, and triangular lattices).

Single-node (8×H100, NVLink) experiments yield extra speedups (beyond ideal slicing) from $7$–$173\times$, capturing 87–101% of the maximum FLOP reduction predicted by compute-only analyses. These results reflect near-saturation of H100 FLOP rates (28–33 TFLOP/s/GPU), and validate that, on high-bandwidth NVLink, communication is not the limiting factor for distributed contraction.

Scaling is extended to 1024 GPUs over InfiniBand. Despite lower interconnect bandwidth, extra speedups persist, with values ranging from $41.8\times$ (circuit simulation) to $67{,}869\times$ (many-body hexagonal). The very large speedups reflect dramatic reductions in the number of sliced indices and, consequently, contractions that would otherwise be infeasible in memory. The results confirm that communication-aware distributed contraction extends scalability far past the exponential wall faced by slicing.

Figure 6: Projected full-contraction speedup vs. embarrassingly parallel slicing baseline across 1–1024 H100 GPUs; extra speedups sustained even over low-bandwidth inter-node networks.

Non-monotonicities in speedup curves across device count are attributed to contraction-path quality variation (the optimizer being NP-hard and run on a fixed time budget) and network topology transitions (e.g., moving from single-node NVLink to multi-node InfiniBand).

Implications and Outlook

This work systematizes the conversion of high-quality contraction paths into hardware-efficient, communication-minimizing execution schedules suitable for modern accelerator clusters. The strong numerical results reinforce that communication, not compute, is the principal barrier to scaling TN contraction beyond slicing-based paradigms. By overlapping pipelined GEMM operations and coordinated communication, the proposed method essentially saturates available hardware capabilities.

On next-generation platforms—such as NVIDIA GB200 NVL72 and future MNNVL fabrics—where single-domain memory and bandwidth expand, a higher fraction of theoretical complexity reduction from distribution will be achievable in practice. This positions communication-optimal execution as a critical enabler for large-scale classical simulation in quantum information science and beyond.

The method’s generality is worth emphasizing: all planning decisions are reduced to contraction-tree structure, mode lifetimes, and hardware cost parameters; no application-specific assumptions are made. The dynamic-programming algorithm for redistribution selection and mode tracking is combinatorially efficient, and empirical performance indicates near-optimal utilization of both device and communication resources.

Conclusion

This research establishes a robust, general framework for multi-GPU TN contraction that supersedes the limitations of slicing via GEMM-optimized mode arrangement and communication-aware data distribution. The magnitude and persistence of extra speedup, especially as device counts scale into the thousands and workloads transition off a single node, highlight the practical effectiveness of the approach. Future advances in interconnect technology will further bridge the gap between theoretical and realized complexity reduction, solidifying communication-aware distributed contraction as the new standard for high-dimensional tensor network simulation.