When Particles Touch: A Story of Fast Tensor Contractions on AI Accelerators

Introduction

When a spacecraft plunges back to Earth at twenty times the speed of sound, the air around it stops acting like a gentle breeze and becomes a chaotic, fiery storm of hot particles.

At the edge of space, the Earth's atmosphere is incredibly thin. Because of this, engineers cannot rely on the standard, everyday rules of aerodynamics to design a safe heat shield. Up there, the air isn't a smooth, continuous fluid. Instead, to understand how the intense heat and pressure will affect the ship, you have to zoom all the way in and look at the gas as a collection of individual atoms and molecules smashing into the spacecraft at ultra-high speeds.

The description of a gas under these extreme conditions requires a delicate accounting of particle collisions. Pioneered by the Austrian physicist Ludwig Boltzmann in the 1870s, it constitutes the most complete model to predict gas flows for high-tech applications.

The underlying mathematical framework and tensor generation are open for academic validation as described in our preprint. At Simkinetic we have developed a highly optimized, proprietary GPU implementation to achieve near-optimal TFLOP/s for enterprise scale.

Reduction of Collision Physics

As an external observer of the gas, the colliding particles can come from any direction in three-dimensional space at various speeds. This poses a huge computational challenge as we need to resolve the different directions and magnitudes for all particle pairs. However, if we put ourselves into the shoes of one of the particles, the second one, our collision target, seems to stand still. Once they touch and collide, they do not scatter arbitrarily in three-dimensional space but are constrained to a two-dimensional plane owing to the conservation of angular momentum.

Comparing this to the naive outside view of the gas, the collision process simplifies from an eight-dimensional problem (two particles, each 3D, and a 2D scattering direction) to a five-dimensional problem! This gives a massive simulation speedup as computational effort typically increases by an order of magnitude with every added dimension. Applied to our problem, this means we can construct the collision physics in fractions of a second that would previously take hours or even days to assemble.

Massive Improvements by Tensor Factorizations

We have removed all redundancy of the global frame and condensed the relevant physics from the collision process without approximations. To apply this not just to one pair of collisions but to every possible collision, we perform two steps. First, we leave our position as a global observer and rotate our view to align it with one particle. Next, we apply our condensed collision physics. Mathematically, this describes a factorization of a tensor,

$$C_{q_1, q_2, q_3}^{k_1, k_2, k_3} = \mathcal R^{(c)}_{k_1, k_2, k_3} \times \mathcal G_{q_1, q_2, q_3}^{(c)}.$$

Here, we follow the notation of Equation (22) in our preprint. The indices $k_1, k_2, k_3$ represent the condensed collision physics in the plane, while the indices $q_1, q_2, q_3$ track the rotations into the single-particle picture. As the geometry of the collision process is highly constrained, only a vanishingly small number of combinations of $(q_1, q_2, q_3)$ have a nonzero contribution. While the condensed collision tensor $\mathcal R$ is densely populated and small, the rotation tensor $\mathcal G$ is large but extremely sparsely populated. Compared to the naive storage of the left-hand side, the collision tensor, this yields a reduction by a factor of one thousand in memory consumption!

Factorizations are an important pillar of modern simulation technology. Their ultimate goal is to break the complicated nature of the tensor into easier sub-problems, often revealing the true nature of the tensor. Applying this to our collision process, we split the nine-dimensional collision tensor into a rotation into the particle frame and the condensed five-dimensional collision physics.

Specialized Tensor Contraction on AI Accelerators

Compared to the full size of the collision tensor, the latter is tiny and we are able to employ established fast techniques on modern hardware to compute its action on the particles. The situation is more challenging for the first part, the tensor that rotates the global frame into the particle frame. It acts as a filter that strips all the redundancies of the global observer frame. This means that only a tiny fraction of its values are actually nonzero. To truly profit from the massive dimensional reduction, we have designed tailored data structures and algorithms to store and apply the rotation filter with high throughput on modern hardware. Owing to the sparse structure of the tensor factorization, AI accelerators with their high memory bandwidth are the perfect platform for our implementation. In the following, we explore the performance of our technology on a multicore CPU and subsequent improvements on AI accelerators. We report the speed of different implementations in trillions of floating-point operations per second (TFLOP/s) together with the percentage of the theoretical peak performance on the corresponding device.

The CPU reference implementation, benchmarked on 128 cores of an AMD EPYC 9754 pinned at 2.25 GHz, achieves 600 GFLOP/s, which amounts to 13% of the peak performance in double precision using AVX-512. As the problem size grows, throughput climbs steeply until the tensor is large enough to saturate the memory bandwidth. From that point on, performance plateaus regardless of how much further the problem grows.

The performance drastically improves with our optimized in-house GPU kernels that are designed for efficient streaming of the tensor factorization, early caching and maximum register use per thread on modern AI accelerators. In addition, we implemented dynamic dispatch to account for different cache sizes and available intrinsics on tensor cores. On the Nvidia H200, a state-of-the-art AI accelerator, we achieve a blazing 12 TFLOP/s, 36% of its peak performance. This is a massive improvement over the CPU implementation: a twenty-fold increase in throughput that even clears the CPU's theoretical peak by a factor of more than two, as the chart below makes vivid.