softdtw-cuda-torch: GPU SoftDTW in PyTorch

• softdtw-cuda-torch is an open-source PyTorch library for computing differentiable SoftDTW on GPUs, addressing sequence-length limits, numerical instability, and memory consumption.
• It employs tiled anti-diagonal kernel execution and a log-space backward pass to ensure efficient parallel computation and numerical stability in gradient calculations.
• The library offers unfused and fused modes, balancing speed and memory efficiency, and supports full PyTorch autograd integration along with Soft-DTW barycenter computation.

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 (Weber et al., 19 Feb 2026).

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$ (Weber et al., 19 Feb 2026). 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 $x = (x_1,\dots,x_N)$0 or $x = (x_1,\dots,x_N)$1, and the recurrence

$x = (x_1,\dots,x_N)$2

where

$x = (x_1,\dots,x_N)$3

The Soft-DTW distance is then $x = (x_1,\dots,x_N)$4 (Weber et al., 19 Feb 2026).

The smoothing parameter $x = (x_1,\dots,x_N)$5 governs the relation between SoftDTW and classical DTW. As $x = (x_1,\dots,x_N)$6, classic non-differentiable DTW is recovered, whereas for $x = (x_1,\dots,x_N)$7 the recurrence is smooth (Weber et al., 19 Feb 2026). 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 $x = (x_1,\dots,x_N)$8. For each diagonal index $x = (x_1,\dots,x_N)$9, all valid pairs $y = (y_1,\dots,y_M)$0 satisfying $y = (y_1,\dots,y_M)$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 (Weber et al., 19 Feb 2026).

2. Stable backward pass and gradient computation

The library computes gradients with respect to the input sequences by first evaluating

$y = (y_1,\dots,y_M)$2

using a reverse dynamic program that starts from $y = (y_1,\dots,y_M)$3 and marches backward on the same grid (Weber et al., 19 Feb 2026). In the linear-space formulation, $y = (y_1,\dots,y_M)$4 depends on neighboring reverse-state values and on derivatives of the soft minimum; for example,

$y = (y_1,\dots,y_M)$5

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 $y = (y_1,\dots,y_M)$6 is small. To prevent this, the implementation computes $y = (y_1,\dots,y_M)$7 and derives a log-domain recurrence. With

$y = (y_1,\dots,y_M)$8

the recurrence becomes

$y = (y_1,\dots,y_M)$9

with

$x_i, y_j \in \mathbb{R}^D$0

By staying in log-space, the largest exponent is subtracted out, avoiding overflow; after the dynamic program completes, $x_i, y_j \in \mathbb{R}^D$1 is recovered by a single exponentiation (Weber et al., 19 Feb 2026).

For the squared-Euclidean cost, the final gradients with respect to the inputs are obtained via the chain rule:

$x_i, y_j \in \mathbb{R}^D$2

$x_i, y_j \in \mathbb{R}^D$3

The implementation computes these sums as batched matrix-multiplications in PyTorch (Weber et al., 19 Feb 2026). 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 $x_i, y_j \in \mathbb{R}^D$4 and $x_i, y_j \in \mathbb{R}^D$5, they can be computed in parallel (Weber et al., 19 Feb 2026). 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 (Weber et al., 19 Feb 2026).

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 $x_i, y_j \in \mathbb{R}^D$6 for a single tile. Because each anti-diagonal is completed before the next kernel launch, correctness is ensured without intra-block synchronization across diagonals (Weber et al., 19 Feb 2026). In the forward-pass pseudo-code, the number of blocks is given by $x_i, y_j \in \mathbb{R}^D$7, where $x_i, y_j \in \mathbb{R}^D$8 is the number of cells on the current diagonal.

The memory layout is also specified. The dynamic-programming buffer $x_i, y_j \in \mathbb{R}^D$9 is stored as a 3D float tensor of size $d(x_i,y_j)$0, row-major in $d(x_i,y_j)$1. In unfused mode, a precomputed cost tensor $d(x_i,y_j)$2 of size $d(x_i,y_j)$3 is also stored, so that threads can read $d(x_i,y_j)$4 in $d(x_i,y_j)$5 time. In fused mode, the library omits $d(x_i,y_j)$6 entirely; instead, each thread loads $d(x_i,y_j)$7 and $d(x_i,y_j)$8 from global memory and computes the local pointwise cost on the fly (Weber et al., 19 Feb 2026).

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:

$d(x_i,y_j)$9

The implementation precomputes normX[b,i]=\|x[b,i]\|^2 of size $\|x_i-y_j\|^2$0 and normY[b,j]=\|y[b,j]\|^2 of size $\|x_i-y_j\|^2$1, and computes the $\|x_i-y_j\|^2$2 matrix of dot products $\|x_i-y_j\|^2$3 via a single batched GEMM (Weber et al., 19 Feb 2026).

In unfused mode, the pairwise distance tensor $\|x_i-y_j\|^2$4 is materialized once. This gives $\|x_i-y_j\|^2$5 kernel-time access per cell to the precomputed cost, at the expense of storage proportional to $\|x_i-y_j\|^2$6 floats (Weber et al., 19 Feb 2026). In fused mode, no $\|x_i-y_j\|^2$7 tensor is stored. Each kernel thread reads $\|x_i-y_j\|^2$8 and $\|x_i-y_j\|^2$9, each of length $(N+1)\times(M+1)$0, and recomputes the cost via the same squared-Euclidean expansion; the backward pass similarly recomputes three neighbor-costs per cell (Weber et al., 19 Feb 2026).

The paper summarizes the comparison as follows.

Mode	Distance-storage complexity	Characterization
Unfused	$(N+1)\times(M+1)$1	Materializes $(N+1)\times(M+1)$2
Fused	$(N+1)\times(M+1)$3	Eliminates $(N+1)\times(M+1)$4

The dynamic-programming tensor $(N+1)\times(M+1)$5 remains $(N+1)\times(M+1)$6 in both modes; only the storage strategy for the distance tensor changes (Weber et al., 19 Feb 2026). For large $(N+1)\times(M+1)$7 and batch size $(N+1)\times(M+1)$8, the paper states

$(N+1)\times(M+1)$9

and $R$0, from which the claim follows that when $R$1 the fused mode can save up to approximately 98% of GPU memory (Weber et al., 19 Feb 2026).

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 $R$2 intermediate cost matrix $R$3, while the storage for the dynamic-programming buffer $R$4 remains identical in both modes (Weber et al., 19 Feb 2026).

5. PyTorch interface and supported workflows

The library is distributed as a PyTorch package installable with $R_{i,0}=R_{0,j}=+\infty$1 or from source by cloning https://github.com/BGU-CS-VIL/sdtw-cuda-torch and running pip install -e . (Weber et al., 19 Feb 2026). Its core API exposes two functions.

First, soft_dtw has signature $R_{i,0}=R_{0,j}=+\infty$2 and computes Soft-DTW distances for each pair (X[b], Y[b]) (Weber et al., 19 Feb 2026).

Second, soft_dtw_barycenter has signature $R_{i,0}=R_{0,j}=+\infty$3 and performs gradient-based Soft-DTW barycenter computation for $R$5 time series (Weber et al., 19 Feb 2026).

The functions support full autograd, so .backward() can be called on the returned Tensor[B] (Weber et al., 19 Feb 2026). The example training loop provided in the paper uses a model whose output has shape [B,N,D], computes $R_{i,0}=R_{0,j}=+\infty$4 and then performs zero_grad, backward, and step within a standard torch.optim.Adam workflow (Weber et al., 19 Feb 2026). 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 (Weber et al., 19 Feb 2026). The reported forward-plus-backward peak GPU memory figures are:

$R$6	$R$7	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 $R$8, 97.9% at $R$9, and 98.1% at $R_{0,0}=0$0 (Weber et al., 19 Feb 2026). It also states that Maghoumi runs out of memory above $R_{0,0}=0$1 for $R_{0,0}=0$2, whereas the fused mode of softdtw-cuda-torch still works up to $R_{0,0}=0$3 (Weber et al., 19 Feb 2026).

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

$R_{0,0}=0$4	$R_{0,0}=0$5	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 $R_{0,0}=0$6 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% (Weber et al., 19 Feb 2026). It also reports a maximum speedup of approximately $R_{0,0}=0$7 for unfused versus Maghoumi at $R_{0,0}=0$8, $R_{0,0}=0$9, and a maximum memory saving of up to 98% for fused versus Maghoumi (Weber et al., 19 Feb 2026).

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 $R_{i,0}=R_{0,j}=+\infty$0 cost matrix and saves up to 98% memory, and seamless integration with PyTorch autograd together with support for barycenter computation (Weber et al., 19 Feb 2026).

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 (Weber et al., 19 Feb 2026).