---
title: 'AMPED: Multi-GPU MTTKRP Acceleration'
url: https://www.emergentmind.com/topics/amped
type: topic
---

# AMPED: Multi-GPU MTTKRP Acceleration

AMPED is a multi-GPU parallel algorithm for accelerating Matricized Tensor Times Khatri-Rao Product (MTTKRP), the computational bottleneck in sparse tensor decomposition and a repeated kernel in CP-ALS. It was introduced in "AMPED: Accelerating MTTKRP for Billion-Scale Sparse Tensor Decomposition on Multiple GPUs" as a method for billion-scale sparse tensors whose memory footprint and irregular access patterns exceed the practical limits of a single GPU. The design combines mode-wise tensor partitioning, greedy load balancing, and ring-based peer-to-peer exchange of updated factor-matrix rows, and on real-world billion-scale tensors it achieves a 5.1x geometric mean speedup in total execution time over state-of-the-art GPU baselines using 4 GPUs on a single CPU node [2507.15121].

## 1. Mathematical setting and role of MTTKRP

Let $\mathcal{X}\in\mathbb{R}^{I_0\times I_1\times\cdots\times I_{N-1}}$ be an $N$-mode sparse tensor, and let $A_m\in\mathbb{R}^{I_m\times R}$ denote the factor matrices of rank $R$. For mode $n$, the mode-$n$ unfolding is $\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$, and the MTTKRP is

$$
Y_{(n)}=\mathcal{X}_{(n)}\cdot\left(A_{N-1}\odot A_{N-2}\odot\cdots\odot A_{n+1}\odot A_{n-1}\odot\cdots\odot A_0\right),
$$

with $Y_{(n)}\in\mathbb{R}^{I_n\times R}$ and $\odot$ denoting the Khatri-Rao product. In elementwise form, each nonzero tensor entry $x=\mathcal{X}(i_0,\ldots,i_{N-1})$ contributes

$$
\forall r\in[0,R),\quad Y_{(n)}(i_n,r)\leftarrow Y_{(n)}(i_n,r)+x\cdot\prod_{m\neq n}A_m(i_m,r).
$$

In CP-ALS decomposition, MTTKRP is repeated for each mode $n$ every iteration. AMPED is therefore targeted at the dominant kernel rather than at a peripheral stage of the decomposition pipeline [2507.15121].

## 2. Single-GPU limits and the motivation for multi-GPU execution

The motivating workload regime is billion-scale sparsity, specifically $10^9$-$10^{10}$ nonzeros. At that scale, sparse tensors can exceed a typical GPU's 48 GB-80 GB of DRAM. Even when the tensor fits, highly irregular accesses to factor rows and atomic updates on $Y$ lead to poor SM utilization and load imbalance. The paper also identifies a cost in out-of-memory, streaming-based single-GPU approaches such as BLCO: heavy PCIe streaming overhead.

These constraints motivate distributing the workload across $d$ GPUs, with $d>1$, to aggregate both memory capacity and compute throughput. AMPED is explicitly constructed to scale beyond the memory and performance limits of a single device while preserving a mode-wise update structure suitable for CP-ALS [2507.15121].

## 3. Partitioning, data placement, and communication structure

AMPED uses a two-level static partitioning scheme for each mode $n$.

At the inter-device level, the output-mode indices $[0..I_n)$ are split into $d$ contiguous blocks of approximately equal size:

$$
I_n=\bigcup_{j=0}^{d-1} I_{n,j},\qquad |I_{n,j}|\approx I_n/d.
$$

The corresponding inter-device partition is

$$
\mathrm{iDp}_{n,j}=\{\text{all nonzeros whose mode-}n\text{ index}\in I_{n,j}\}.
$$

Because no two GPUs update the same output row, there are no inter-GPU $Y$ dependencies during the per-mode compute phase. This eliminates inter-GPU race dependencies during execution.

At the inter-SM level, each $\mathrm{iDp}_{n,j}$ on GPU $j$ is subdivided into

$$
t_{n,j}=\left\lceil |\mathrm{iDp}_{n,j}|/g\right\rceil
$$

equally sized chunks, where $g$ is the number of SMs per GPU. Each chunk feeds one threadblock. The total number of inter-SM partitions in mode $n$ is

$$
\tau_n=\sum_{j=0}^{d-1} t_{n,j}\approx |\mathcal{X}.\mathrm{nnz}|/(d\cdot g).
$$

The data-movement pattern is correspondingly structured. For each mode $n$, the host transfers $\sum_j |\mathrm{iDp}_{n,j}|$ nonzeros to the GPUs exactly once. At the end of the mode, each GPU sends its updated $Y$ rows of size approximately $(I_n/d)\times R$ floats to all other GPUs via ring P2P. Per device, GPU DRAM holds three categories of data simultaneously: its local nonzero partition, a full local copy of all factor matrices $\{A_0,\ldots,A_{N-1}\}$, and its local fragment of $Y_{(n)}$ [2507.15121].

## 4. Dynamic load balancing and execution flow

Although the inter-device partitions are based on equal index ranges, the actual nonzero counts $\eta_{n,j}=|\mathrm{iDp}_{n,j}|$ may differ substantially because sparsity is irregular. AMPED addresses this with a greedy Longest-Processing-Time assignment:

1. sort $\{\mathrm{iDp}_{n,j}\}$ in descending $\eta_{n,j}$,
2. iteratively assign the largest remaining partition to the GPU whose current assigned load is smallest.

By Graham '69, this assignment guarantees that no GPU's assigned load exceeds $(4/3)\cdot \mathrm{OPT}$, where $\mathrm{OPT}$ is perfect balance. In practice, the reported compute-time skew across GPUs is less than $1\%$ [2507.15121].

The per-mode execution flow follows this load-balanced partitioning. For each mode $n$, the host selects $T[n]=\mathcal{X}_{(n)}$ partitions, each GPU loads its $\mathrm{iDp}_{n,j}$ into device memory, and a grid is launched in which each block corresponds to one inter-SM partition. Within the threadblock, OFMRU performs the per-nonzero update: it computes a temporary vector $\mathrm{tmp}[0..R-1]=\mathrm{val}\cdot\prod_{m\neq n}A_m(i_m,r)$ and then atomically accumulates into $Y_{(n)}(i_n,0..R-1)$. After a global GPU barrier, the updated $Y$ fragments are exchanged by ring-based peer-to-peer communication; after a second barrier, each GPU holds the full updated factor matrix for mode $n$.

A plausible implication is that AMPED separates irregular sparse accumulation from dense factor synchronization in a way that makes the communication phase explicit and analyzable, rather than hidden behind repeated fine-grained coherence events.

## 5. Memory model and feasibility bounds

The per-GPU memory model is given in terms of the largest assigned partition and the size of the factor matrices. Let

- $\eta_{\max}=\max_j |\mathrm{iDp}_{n,j}|$,
- $S_{\mathrm{val}}=\mathrm{sizeof(value)}$,
- $S_{\mathrm{idx}}=4\ \mathrm{B}$,
- $M_A=\sum_m I_m\cdot R\cdot S_{\mathrm{val}}$.

Then each GPU must accommodate simultaneously

- $\eta_{\max}\cdot (S_{\mathrm{val}}+N\cdot S_{\mathrm{idx}})$ for coordinates and values of the largest partition,
- $M_A$ for the full factor-matrix copies,
- $(I_n/d)\cdot R\cdot S_{\mathrm{val}}$ for the local $Y_{(n)}$ fragment.

AMPED requires choosing $d$ so that

$$
\eta_{\max}(S_{\mathrm{val}}+N\cdot S_{\mathrm{idx}})+M_A+(I_n/d)\cdot R\cdot S_{\mathrm{val}}
\le \mathrm{GPU\_global\_DRAM}.
$$

In the reported experiments, using $d=4$ on 48 GB GPUs, all four billion-scale tensors fit comfortably. This bound is central to the algorithm's claim of scaling beyond a single GPU: feasibility is achieved not by streaming away the tensor-matrix interaction, but by partitioning the tensor while replicating factor matrices and restricting output ownership per device [2507.15121].

## 6. Experimental evaluation

The evaluation platform consists of an AMD EPYC 9654 host with 192 threads and 1.5 TB DDR4, together with $4\times$ NVIDIA RTX 6000 Ada GPUs with 48 GB GDDR6 and PCIe+GPUDirect P2P. The tensors are drawn from FROSTT and LiveRec.

| Tensor | Shape | nnz (B) |
|---|---|---:|
| Amazon | $4.8$ M $\times\ 1.8$ M $\times\ 1.8$ M | 1.7 |
| Patents | $46 \times 239.2$ K $\times 239.2$ K | 3.6 |
| Reddit-15 | $8.2$ M $\times 177$ K $\times 8.1$ M | 4.7 |
| Twitch | 5-mode, $15.5$ M $\times 6.2$ M $\times 0.8$ M $\times 6.1$ K $\times 6.1$ K | 0.5 |

The baseline set includes BLCO, MM-CSF, HiCOO-GPU, and FLYCOO-GPU. For a single ALS iteration at $R=32$, AMPED achieves a geometric-mean 5.1x speedup over the best single-GPU baseline. Relative to its own 1-GPU execution, it is 1.9x, 2.3x, and 3.3x faster with $d=2,3,4$ GPUs, respectively, which the paper characterizes as nearly linear scaling. The sustained per-GPU GFLOP/s also grows linearly with $d$ [2507.15121].

The load-balancing study is particularly sharp. Comparing equal-nnz splits against the greedy assignment, the unbalanced case incurs a 5-10x slowdown, while AMPED maintains compute-time skew below $1\%$. The communication/computation breakdown is workload dependent: for very sparse tensors with huge nnz, specifically Patents and Reddit, host-to-GPU streaming dominates communication; for large-index tensors, specifically Amazon and Twitch, inter-GPU factor exchange is approximately $30\%$ of total time [2507.15121].

## 7. Limitations, future directions, and acronym ambiguity

The reported limitations are primarily communication- and heterogeneity-related. Communication overhead remains approximately $30\%$ on modes with very large output-index sets, and the paper identifies NVLink or RDMA interconnects as possible remedies. Static partitioning also cannot react to dynamic run-time imbalances due to hot-rows, so hybrid static/dynamic schemes are proposed as a future improvement. Additional stated directions are extending AMPED to heterogeneous nodes mixing CPU, GPU, and FPGA, with device-specific cost models, and overlapping host-to-GPU streaming, GPU compute, and P2P communication with CUDA-graphs [2507.15121].

The acronym itself is overloaded in recent arXiv literature. "AMPED" also denotes "Adaptive Multi-objective Projection for balancing Exploration and skill Diversification" in skill-based reinforcement learning [2506.05980] and "Adaptive Multi-stage Non-edge Pruning for Edge Detection" in Transformer-based edge detection [2603.27661]. A related but distinct acronym, AMP, refers to Approximate Message Passing in high-dimensional inference [2105.02180]. Within sparse tensor decomposition, however, AMPED refers specifically to the multi-GPU MTTKRP algorithm whose core innovations are task-independent partitioning of billion-scale sparse tensors across GPUs, provably good greedy load balancing of nonzeros, elimination of inter-GPU race dependencies during per-mode compute, and ring-based P2P factor exchange with minimal barrier overhead [2507.15121].

Source: https://www.emergentmind.com/topics/amped