---
title: 'softdtw-cuda-torch: GPU SoftDTW in PyTorch'
url: https://www.emergentmind.com/topics/softdtw-cuda-torch
type: topic
---

# softdtw-cuda-torch: GPU SoftDTW in PyTorch

Searching arXiv for the specified paper and closely related Soft-DTW background.
softdtw-cuda-torch is an open-source PyTorch library for computing Soft Dynamic Time Warping (SoftDTW) on GPUs. It is presented as a response to three limitations in existing GPU implementations of SoftDTW: a hard sequence-length cap of 1024, numerical instability in the backward pass for small smoothing parameters, and excessive GPU memory consumption from materializing pairwise distance tensors. The implementation introduces tiled anti-diagonal kernel execution, a log-space backward pass, and a fused distance-computation mode that eliminates the $O(B N M)$ intermediate distance tensor, while supporting arbitrary sequence lengths, full PyTorch autograd integration, and Soft-DTW Barycenter computation [2602.17206].

## 1. Soft-DTW formulation

SoftDTW in this library is defined for two real-valued sequences
$x = (x_1,\dots,x_N)$ and $y = (y_1,\dots,y_M)$ with $x_i, y_j \in \mathbb{R}^D$, together with a pointwise cost $d(x_i,y_j)$, usually $\|x_i-y_j\|^2$ [2602.17206]. The method fills an $(N+1)\times(M+1)$ dynamic-programming table $R$ using the boundary conditions
$R_{0,0}=0$ and $R_{i,0}=R_{0,j}=+\infty$ for $i>0$ or $j>0$, and the recurrence
$$
R_{i,j}=d(x_i,y_j)+\operatorname{softmin}_\gamma(R_{i-1,j-1},R_{i-1,j},R_{i,j-1}),
$$
where
$$
\operatorname{softmin}_\gamma(a,b,c)\coloneqq -\gamma \log\left(e^{-a/\gamma}+e^{-b/\gamma}+e^{-c/\gamma}\right).
$$
The Soft-DTW distance is then $sdtw_\gamma(x,y)=R_{N,M}$ [2602.17206].

The smoothing parameter $\gamma>0$ governs the relation between SoftDTW and classical DTW. As $\gamma\to 0$, classic non-differentiable DTW is recovered, whereas for $\gamma>0$ the recurrence is smooth [2602.17206]. This makes SoftDTW suitable for gradient-based optimization and for direct use as a loss in PyTorch training workflows.

The forward pass is evaluated anti-diagonal by anti-diagonal, i.e. across sets of indices with constant $i+j$. For each diagonal index $p=0,\dots,N+M-2$, all valid pairs $(i,j)$ satisfying $i+j=p+1$ are computed. This anti-diagonal structure is central to the GPU implementation because all cells on the same anti-diagonal depend only on the previous two anti-diagonals and can therefore be evaluated in parallel [2602.17206].

## 2. Stable backward pass and gradient computation

The library computes gradients with respect to the input sequences by first evaluating
$$
E_{i,j}\coloneqq \frac{\partial R_{N,M}}{\partial d(x_i,y_j)},
$$
using a reverse dynamic program that starts from $E_{N,M}=1$ and marches backward on the same grid [2602.17206]. In the linear-space formulation,
$E_{i,j}$ depends on neighboring reverse-state values and on derivatives of the soft minimum; for example,
$$
\frac{\partial \operatorname{softmin}}{\partial a}
=
\frac{e^{-a/\gamma}}{e^{-a/\gamma}+e^{-b/\gamma}+e^{-c/\gamma}}.
$$

A central feature of softdtw-cuda-torch is the log-space backward pass for stability. The paper states that the linear-space backward pass overflows when $\gamma$ is small. To prevent this, the implementation computes
$\bar{E}_{i,j}=\log E_{i,j}$ and derives a log-domain recurrence. With
$$
\alpha = (R_{i+1,j} - R_{i,j} - d_{i+1,j})/\gamma,\quad
\beta = (R_{i,j+1} - R_{i,j} - d_{i,j+1})/\gamma,\quad
\delta = (R_{i+1,j+1} - R_{i,j} - d_{i+1,j+1})/\gamma,
$$
the recurrence becomes
$$
\bar{E}_{i,j}
=
\operatorname{logsumexp}(\bar{E}_{i+1,j}+\alpha,\ \bar{E}_{i,j+1}+\beta,\ \bar{E}_{i+1,j+1}+\delta),
$$
with
$$
\operatorname{logsumexp}(a,b,c)=m+\log(e^{a-m}+e^{b-m}+e^{c-m}),\qquad m=\max(a,b,c).
$$
By staying in log-space, the largest exponent is subtracted out, avoiding overflow; after the dynamic program completes, $E=\exp(\bar{E})$ is recovered by a single exponentiation [2602.17206].

For the squared-Euclidean cost, the final gradients with respect to the inputs are obtained via the chain rule:
$$
\frac{\partial sdtw}{\partial x_i}
=
2\left(x_i\cdot \sum_j E_{i,j} - \sum_j E_{i,j}\cdot y_j\right),
$$
$$
\frac{\partial sdtw}{\partial y_j}
=
2\left(y_j\cdot \sum_i E_{i,j} - \sum_i E_{i,j}\cdot x_i\right).
$$
The implementation computes these sums as batched matrix-multiplications in PyTorch [2602.17206]. This ties the dynamic-programming core to the autograd ecosystem without requiring a separate gradient interface.

## 3. GPU execution model

The implementation exploits anti-diagonal parallelism. Because all cells on the same anti-diagonal depend only on diagonals $p-1$ and $p-2$, they can be computed in parallel [2602.17206]. The paper contrasts this with Maghoumi’s original code, which launched one big kernel and used `__syncthreads()`, thereby capping the sequence length at 1024, corresponding to the maximum threads per block [2602.17206].

To remove that limit, softdtw-cuda-torch breaks each diagonal into tiles of at most 256 threads per block and launches one small kernel per diagonal. The host performs the loop over diagonals, while each CUDA kernel computes $R_{i,j}$ for a single tile. Because each anti-diagonal is completed before the next kernel launch, correctness is ensured without intra-block synchronization across diagonals [2602.17206]. In the forward-pass pseudo-code, the number of blocks is given by $\lceil \ell/256 \rceil$, where $\ell$ is the number of cells on the current diagonal.

The memory layout is also specified. The dynamic-programming buffer $R$ is stored as a 3D float tensor of size $(B, N+2, M+2)$, row-major in $(i,j)$. In unfused mode, a precomputed cost tensor $D$ of size $(B,N,M)$ is also stored, so that threads can read $D[b,i-1,j-1]$ in $O(1)$ time. In fused mode, the library omits $D$ entirely; instead, each thread loads $x[b,i]$ and $y[b,j]$ from global memory and computes the local pointwise cost on the fly [2602.17206].

This design directly addresses the sequence-length limitation noted in earlier GPU implementations. A plausible implication is that the principal architectural change is not a modification of the SoftDTW recurrence itself, but a restructuring of the CUDA launch pattern so that the anti-diagonal dynamic program remains valid for arbitrary sequence lengths.

## 4. Memory-efficient fused distance computation

The memory-efficiency contribution is based on an algebraic identity for squared-Euclidean distance:
$$
\|x_i-y_j\|^2 = \|x_i\|^2 - 2\langle x_i,y_j\rangle + \|y_j\|^2.
$$
The implementation precomputes `normX[b,i]=\|x[b,i]\|^2` of size $(B,N)$ and `normY[b,j]=\|y[b,j]\|^2` of size $(B,M)$, and computes the $B\times N\times M$ matrix of dot products $X\cdot Y^\top$ via a single batched GEMM [2602.17206].

In unfused mode, the pairwise distance tensor $D\in\mathbb{R}^{B\times N\times M}$ is materialized once. This gives $O(1)$ kernel-time access per cell to the precomputed cost, at the expense of storage proportional to $4 \text{ bytes}\cdot B\cdot N\cdot M$ floats [2602.17206]. In fused mode, no $D$ tensor is stored. Each kernel thread reads $x[b,i]$ and $y[b,j]$, each of length $D$, and recomputes the cost via the same squared-Euclidean expansion; the backward pass similarly recomputes three neighbor-costs per cell [2602.17206].

The paper summarizes the comparison as follows.

| Mode | Distance-storage complexity | Characterization |
|---|---:|---|
| Unfused | $O(B\cdot N\cdot M)$ | Materializes $D$ |
| Fused | $O(B\cdot (N+M))$ | Eliminates $D$ |

The dynamic-programming tensor $R$ remains $O(B\cdot N\cdot M)$ in both modes; only the storage strategy for the distance tensor changes [2602.17206]. For large $N=M=L$ and batch size $B$, the paper states
$\text{peak\_mem\_unfused} \simeq c\cdot B\cdot L^2$
and
$\text{peak\_mem\_fused} \simeq c\cdot B\cdot 2L$,
from which the claim follows that when $L\gg 1$ the fused mode can save up to approximately 98% of GPU memory [2602.17206].

A common misconception is that the fused mode reduces the total asymptotic storage of SoftDTW. The paper does not make that claim. It states instead that the fused mode eliminates the $O(B\cdot N\cdot M)$ intermediate cost matrix $D$, while the storage for the dynamic-programming buffer $R$ remains identical in both modes [2602.17206].

## 5. PyTorch interface and supported workflows

The library is distributed as a PyTorch package installable with
```bash
pip install softdtw-cuda-torch
```
or from source by cloning `https://github.com/BGU-CS-VIL/sdtw-cuda-torch` and running `pip install -e .` [2602.17206]. Its core API exposes two functions.

First, `soft_dtw` has signature
```python
soft_dtw(
    X: Tensor[B,N,D],
    Y: Tensor[B,M,D],
    gamma: float = 1.0,
    device: torch.device = None,
    fused: bool = True
) -> Tensor[B]
```
and computes Soft-DTW distances for each pair `(X[b], Y[b])` [2602.17206].

Second, `soft_dtw_barycenter` has signature
```python
soft_dtw_barycenter(
    sequences: Tensor[K,N,D],
    gamma: float = 1.0,
    n_iter: int = 1000,
    lr: float = 0.1,
    device: torch.device = None
) -> Tensor[N,D]
```
and performs gradient-based Soft-DTW barycenter computation for $K$ time series [2602.17206].

The functions support full autograd, so `.backward()` can be called on the returned `Tensor[B]` [2602.17206]. The example training loop provided in the paper uses a model whose output has shape `[B,N,D]`, computes
```python
loss = soft_dtw(pred, Y, gamma=0.5, fused=True).mean()
```
and then performs `zero_grad`, `backward`, and `step` within a standard `torch.optim.Adam` workflow [2602.17206]. The loss returned by `soft_dtw` is a batch vector and is averaged to obtain a scalar.

This API design places softdtw-cuda-torch in the category of differentiable sequence-alignment losses directly usable in end-to-end training. A plausible implication is that the library targets scenarios where exact DTW-style alignment behavior is desired without sacrificing the differentiability required by contemporary PyTorch optimization pipelines.

## 6. Benchmark results and operational trade-offs

The performance evaluation compares Maghoumi’s `pytorch-softdtw-cuda` with the new implementation in two modes, unfused and fused. Benchmarks were measured on an NVIDIA GTX1080 and averaged over 5 runs [2602.17206]. The reported forward-plus-backward peak GPU memory figures are:

| $B$ | $L$ | Maghoumi (MB) | Ours unfused (MB) | Ours fused (MB) |
|---:|---:|---:|---:|---:|
| 16 | 128 | 275 | 26 | 23 |
| 16 | 512 | 4,136 | 137 | 89 |
| 32 | 512 | 8,256 | 257 | 161 |

The paper further reports memory saving versus Maghoumi of 91.6% at $(B,L)=(16,128)$, 97.9% at $(16,512)$, and 98.1% at $(32,512)$ [2602.17206]. It also states that Maghoumi runs out of memory above $L=512$ for $B=32$, whereas the fused mode of softdtw-cuda-torch still works up to $L=2048$ [2602.17206].

For wall-clock runtime, again measuring forward plus backward:

| $B$ | $L$ | Maghoumi (ms) | Ours unfused (ms) | Ours fused (ms) |
|---:|---:|---:|---:|---:|
| 16 | 128 | 7.7 | 1.8 | 47.0 |
| 16 | 512 | 83.2 | 16.0 | 200.8 |
| 32 | 512 | 2,791 | 41.8 | 429.5 |

The paper summarizes these results by stating that the unfused path is up to approximately $50\times$ faster than Maghoumi’s original, while the fused path is 10–15× slower than the library’s unfused mode but reduces memory by 40–98% [2602.17206]. It also reports a maximum speedup of approximately $50\times$ for unfused versus Maghoumi at $L=512$, $B=32$, and a maximum memory saving of up to 98% for fused versus Maghoumi [2602.17206].

These figures delineate the principal deployment trade-off. Unfused mode is the throughput-oriented configuration, whereas fused mode is the memory-oriented configuration. The paper presents both as complementary rather than mutually exclusive choices, allowing the practitioner to select between precomputed distances and on-the-fly recomputation depending on hardware constraints and problem scale.

## 7. Scope, significance, and limitations addressed

The paper’s conclusion identifies four main properties: removal of the 1024-length cap via tiled anti-diagonals, a numerically stable backward pass via log-space dynamic programming, a fused mode that eliminates the $O(B\cdot N\cdot M)$ cost matrix and saves up to 98% memory, and seamless integration with PyTorch autograd together with support for barycenter computation [2602.17206].

Within that framing, softdtw-cuda-torch is best understood as an implementation-focused contribution rather than a reformulation of SoftDTW itself. Its novelty lies in CUDA execution strategy, numerical stabilization of the backward recursion, and memory-management choices for distance computation. This suggests that the library occupies the intersection of differentiable dynamic programming, GPU kernel design, and practical PyTorch systems engineering.

The implementation also sharpens a recurrent tension in differentiable sequence alignment: numerical robustness, GPU occupancy, and memory footprint often pull in different directions. The paper resolves this tension by separating concerns into distinct mechanisms—tiled anti-diagonal execution for scalability, log-space recurrence for stability, and fused distance computation for memory reduction—while preserving a single user-facing API [2602.17206].

Source: https://www.emergentmind.com/topics/softdtw-cuda-torch