AMPED: Multi-GPU MTTKRP Acceleration

• AMPED is a multi-GPU algorithm for accelerating the Matricized Tensor Times Khatri-Rao Product, a key operation in sparse tensor decomposition.
• It leverages mode-wise tensor partitioning, greedy load balancing, and ring-based peer-to-peer exchanges to overcome single-GPU memory and performance limits.
• The method achieves nearly linear scaling, with a 5.1x speedup over state-of-the-art single-GPU baselines, demonstrating efficient multi-GPU resource utilization.

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 (Wijeratne et al., 20 Jul 2025).

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 $N$0 contributes

$N$1

In CP-ALS decomposition, MTTKRP is repeated for each mode $N$2 every iteration. AMPED is therefore targeted at the dominant kernel rather than at a peripheral stage of the decomposition pipeline (Wijeratne et al., 20 Jul 2025).

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

The motivating workload regime is billion-scale sparsity, specifically $N$3-$N$4 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 $N$5 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 $N$6 GPUs, with $N$7, 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 (Wijeratne et al., 20 Jul 2025).

3. Partitioning, data placement, and communication structure

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

At the inter-device level, the output-mode indices $N$9 are split into $A_m\in\mathbb{R}^{I_m\times R}$0 contiguous blocks of approximately equal size:

$A_m\in\mathbb{R}^{I_m\times R}$1

The corresponding inter-device partition is

$A_m\in\mathbb{R}^{I_m\times R}$2

Because no two GPUs update the same output row, there are no inter-GPU $A_m\in\mathbb{R}^{I_m\times R}$3 dependencies during the per-mode compute phase. This eliminates inter-GPU race dependencies during execution.

At the inter-SM level, each $A_m\in\mathbb{R}^{I_m\times R}$4 on GPU $A_m\in\mathbb{R}^{I_m\times R}$5 is subdivided into

$A_m\in\mathbb{R}^{I_m\times R}$6

equally sized chunks, where $A_m\in\mathbb{R}^{I_m\times R}$7 is the number of SMs per GPU. Each chunk feeds one threadblock. The total number of inter-SM partitions in mode $A_m\in\mathbb{R}^{I_m\times R}$8 is

$A_m\in\mathbb{R}^{I_m\times R}$9

The data-movement pattern is correspondingly structured. For each mode $R$0, the host transfers $R$1 nonzeros to the GPUs exactly once. At the end of the mode, each GPU sends its updated $R$2 rows of size approximately $R$3 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 $R$4, and its local fragment of $R$5 (Wijeratne et al., 20 Jul 2025).

4. Dynamic load balancing and execution flow

Although the inter-device partitions are based on equal index ranges, the actual nonzero counts $R$6 may differ substantially because sparsity is irregular. AMPED addresses this with a greedy Longest-Processing-Time assignment:

1. sort $R$7 in descending $R$8,
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 $R$9, where $n$0 is perfect balance. In practice, the reported compute-time skew across GPUs is less than $n$1 (Wijeratne et al., 20 Jul 2025).

The per-mode execution flow follows this load-balanced partitioning. For each mode $n$2, the host selects $n$3 partitions, each GPU loads its $n$4 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 $n$5 and then atomically accumulates into $n$6. After a global GPU barrier, the updated $n$7 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$8.

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

• $n$9,
• $n$0,
• $n$1,
• $n$2.

Then each GPU must accommodate simultaneously

• $n$3 for coordinates and values of the largest partition,
• $n$4 for the full factor-matrix copies,
• $n$5 for the local $n$6 fragment.

AMPED requires choosing $n$7 so that

$n$8

In the reported experiments, using $n$9 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 (Wijeratne et al., 20 Jul 2025).

6. Experimental evaluation

The evaluation platform consists of an AMD EPYC 9654 host with 192 threads and 1.5 TB DDR4, together with $$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$0 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	$$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$1 M $$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$2 M $$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$3 M	1.7
Patents	$$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$4 K $$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$5 K	3.6
Reddit-15	$$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$6 M $$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$7 K $$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$8 M	4.7
Twitch	5-mode, $$\mathcal{X}_{(n)}\in\mathbb{R}^{I_n\times (I_0\cdots I_{n-1}I_{n+1}\cdots I_{N-1})}$$9 M $$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),$$0 M $$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),$$1 M $$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),$$2 K $$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),$$3 K	0.5

The baseline set includes BLCO, MM-CSF, HiCOO-GPU, and FLYCOO-GPU. For a single ALS iteration at $$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),$$4, 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 $$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),$$5 GPUs, respectively, which the paper characterizes as nearly linear scaling. The sustained per-GPU GFLOP/s also grows linearly with $$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),$$6 (Wijeratne et al., 20 Jul 2025).

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 $$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),$$7. 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 $$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),$$8 of total time (Wijeratne et al., 20 Jul 2025).

7. Limitations, future directions, and acronym ambiguity

The reported limitations are primarily communication- and heterogeneity-related. Communication overhead remains approximately $$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),$$9 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 (Wijeratne et al., 20 Jul 2025).

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 (Cho et al., 6 Jun 2025) and "Adaptive Multi-stage Non-edge Pruning for Edge Detection" in Transformer-based edge detection (Gao et al., 29 Mar 2026). A related but distinct acronym, AMP, refers to Approximate Message Passing in high-dimensional inference (Feng et al., 2021). 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 (Wijeratne et al., 20 Jul 2025).