Software, and making SME useful

The versions come from the arXiv paper on interatomic potentials. NSCS has described the STAR libraries only by their role.

Both teams that ran on the full machine used the same layout: 16 MPI ranks per node, one per 38-core cluster, with OpenMP inside. It follows from the memory page: a rank that stays on its cluster reads 450 GB/s from its own HBM stack, and one that strays gets half or less.

Inside the 38 cores, the interatomic potential team split the work:

The Earth-observation team reports 37 cores per cluster available to the application, one left to the OS.

On a GPU, PyTorch overlaps communication and compute through CUDA streams. The CPU backend has no such thing: a collective blocks the thread that calls it. Both teams wrote their way around that.

The Earth-observation team wrote two SME micro-kernels, 64 × K by K × 16 and 16 × K by K × 64, chosen at run time from the matrix shape. K comes from the L2 capacity, so a panel stays in the 28.5 MB of the cluster. The streamed operand is prefetched into HBM; partial sums in FP32 also stay in HBM when K is large. Bias and scaling are fused into the write-back. Single-node step time on their 6-billion-parameter model went from 51.3 s to 5.0 s, a 10.3× gain, once all of this was in place.

NSCS's own slide on the same problem: SME for the outer-product part of stencils, square GEMM tiles and the tall-skinny products of attention; SVE for the row updates, boundaries and softmax; buffers pre-allocated in HBM.

The interatomic-potential team summarised their whole stack in one figure; panel (g) is the SDMA and HBM pipeline described above.

One number every user of a big machine should look at: on the full system, MPI initialisation took 923.5 s and Python imports plus HBM and PyTorch initialisation another 531.8 s. Twenty-four minutes before the first step. On a 22,000-node job, start-up is part of the cost of a run.

What these papers suggest, if you bring a code to LineShine: