Fused & Furious: Sinkhorn Triton Kernels

In their latest paper, charmingly titled Manifold-Constrained Hyper-Connections¹, DeepSeek improved upon preexisting research on residual stream scaling with two major contributions:

1. an approach that’s more stable than previous papers, with demonstrated performance at scale
2. an efficient infrastructure design that minimizes the overhead of hyper-connections compared to vanilla transformers

My goal in this post is to delve deeper into the second part. More precisely, I want to focus on the Sinkhorn projection step of their architecture: why we need custom fused GPU kernels to implement it, and how to do it.

This post will be in three parts:

1. First, I’ll present increasingly complex residual connection paradigms, culminating in the mHC architecture.
2. Then, I’ll provide some geometric intuition behind Sinkhorn’s algorithm and derive it.
3. Finally, I’ll showcase increasingly fast Triton kernels for the Sinkhorn algorithm.

Figure 1 serves as a teaser of the crazy speedups we achieved. Let’s dive in!

The code for this project can be found here.

Table of Contents

• I. Scaling the residual stream
  • I.A. Residual Matrix Transformers
  • I.B. Hyper-Connections
  • I.C. Manifold-Constrained Hyper-Connections
• II. Sinkhorn’s Algorithm
  • II.A. Projecting under the generalized KL divergence
  • II.B. Deriving the algorithm
• III. Triton kernels for Sinkhorn
  • III.A. Hardware setup
  • III.B. Mini-refresher on memory hierarchy
  • III.C. Implementations
    • Naive PyTorch implementation
    • Basic fused kernel
    • Loading $M$ in registers
    • Block packing
  • III.D. Benchmark
  • III.E. Figures
    • Vanity metric
    • Comparing implementations

I. Scaling the residual stream

A current trend in macro design of transformers is to scale the size of the residual stream. That is, instead of being a vector $\x\in\R^{\D}$, the per-token residual stream becomes a matrix $X\in\R^{n\times\D}$ where $n$ is a new scaling dimension (often a small value e.g. $n=4$). The motivation behind this architectural shift is that in vanilla transformers, model parameters and FLOPs scale quadratically with $\D$ whereas the residual stream scales linearly. Hence, as we keep increasing model size, the residual stream may become a bottleneck.

I.A. Residual Matrix Transformers

My previous post discussed the Residual Matrix Transformer (RMT) architecture, which introduces a matrix representation of the per-token residual stream inspired by outer-product memories. The key idea is that using a matrix instead of a vector increases the effective storage capacity of the residual stream. The implementation of RMT is fairly simple, as the only difference compared to the classic transformer architecture is how we read from and write to the residual stream. Crucially, the micro design remains untouched (i.e. Attention and Feed-Forward blocks do not change), which implies that scaling the residual stream incurs negligible computational overhead, though it does cause memory access overhead if implemented naively (i.e. without fused kernels) as we will see.

In a nutshell, RMT replaces the classic residual connection

with

where $\x_l\in\R^{\D}$ (resp. $X_l\in\R^{n\times\D}$) is the $l$-th layer vector (resp. matrix) residual stream, and $\F_l$ is the $l$-th transformer block.

$\Hpre_l, \Hpost_l \in \R^{1\times n}$ are learned vectors which correspond respectively to READ and WRITE operations on the residual stream memory store.

I.B. Hyper-Connections

Hyper-Connections² (HC) extend the idea of RMT by enabling communication between the $n$ channels of the residual stream.

The motivation is to increase the topological complexity of the residual mapping without increasing the computational complexity. The main improvement over RMT is the introduction of $\Hres_l\in\Rnn$ which acts as a “stream-mixing matrix” allowing the $n$ channels to exchange information instead of simply evolving independently. Also, $\Hpre_l, \Hpost_l, \Hres_l$ are now both static (i.e. learned parameters) and dynamic (i.e. dependent on input $X_l$) to further increase the topological complexity of the connection network (see original paper for more details). Note that we will omit the dependence on $X_l$ to keep notations lightweight (e.g. we’ll write $\Hres_l$ instead of $\Hres_l(X_l)$)

Thus, the residual connection in HC is even richer than in RMT thanks to $\Hres_l$:

I.C. Manifold-Constrained Hyper-Connections

DeepSeek’s Manifold-Constrained Hyper-Connections (mHC) address one critical flaw in HC: instability as the network depth $L$ scales.

Indeed, HC abandons the identity mapping of residual connections and replaces it with $\Hres_l$. Yet, the identity mapping is crucial to the training stability of deep neural network architectures³.

It’s easy to see why replacing $\mathbb{I}_n$ with $\Hres_l$ is problematic: if we recursively extend the residual connection across multiple layers, we have:

1. \(\x_{l+m} = \x_l + \sum_{k=0}^{m-1}\F_{k+l}(\x_{l+k})\) for the standard residual connection
2. \(X_{l+m} = X_l + \sum_{k=0}^{m-1}{\Hpost_{l+k}}^T \F_{l+k}(\Hpre_{l+k} X_{l+k})\) for RMT
3. \(X_{l+m} = \big(\overleftarrow{\prod}_{k=0}^{m-1}\Hres_{l+k}\big) X_l + \sum_{k=0}^{m-1} \big(\overleftarrow{\prod}_{j=k+1}^{m-1}\Hres_{l+j}\big) {\Hpost_{l+k}}^T \F_{l+k}(\Hpre_{l+k} X_{l+k})\) for HC

where $\overleftarrow{\prod}$ denotes the reverse product i.e. $\overleftarrow{\prod} _ {k=0}^{m-1}\Hres_{l+k} = \Hres_{l+m-1}\Hres_{l+m-2}\cdots\Hres_{l+1}\Hres_{l}$.

The matrix product $\Pi=\overleftarrow{\prod} _ {k=0}^{m-1}\Hres_{l+k}$ has no reason to behave nicely i.e. keep a spectral norm close to 1. In fact, mHC shows that instead of preserving the signal strength, $\Pi$ tends to amplify or attenuate it as depth increases, resulting in signal rescaling across several orders of magnitudes, which is problematic for both the forward pass (exploding/vanishing activations) and the backward pass (exploding/vanishing gradients). The propagation of instability as depth increases is illustrated in Figure 2.

DeepSeek’s solution to this instability problem is to constrain $\Hres_l$ to be a bistochastic matrix i.e. a square matrix with non-negative entries where each row and column sums to 1.

Bistochastic matrices have three nice properties that make them ideal for stabilizing the residual stream:

1. they have a spectral norm of 1, meaning that they preserve signal strength across layers
2. they are closed under matrix multiplication, meaning that the product of multiple bistochastic matrices is also bistochastic
3. they can be given a probabilistic / optimal transport interpretation as soft permutations, meaning that they allow each channel to attend to and exchange with every other channel, without collapsing to a single channel

To enforce the bistochastic constraint, we must first compute $\Hres_l$ as in HC, and then project it onto the manifold of bistochastic matrices $\B$, also known as Birkhoff’s polytope. The next section goes into the details of this projection, which uses Sinkhorn’s algorithm.

II. Sinkhorn’s Algorithm

Quite plainly, the goal of Sinkhorn’s algorithm is to project a given square matrix $M\in\Rnn$⁴ onto the Birkhoff polytope

\(\B=\lbrace P\in\Rnn \mid P\one=\one, \one^T P=\one^T, P\geq 0\rbrace\).

Note that $\B$ is convex as we will use this property later. In fact it is the convex hull generated by the $n!$ permutation matrices in $\Rnn$, which are the extreme points of $\B$.⁵

Let’s now present Sinkhorn’s algorithm and derive its formulation intuitively.

II.A. Projecting under the generalized KL divergence

We’ve been talking about projections, but to define a projection, one needs a metric. In other words, we need a sensible metric $d: \Rnn \times \Rnn \to \Rnn$ and then we can try to find \(\text{Proj}_d^{\B}(M) = \arg\min_{P \in \B} d(P,M)\) for a given matrix $M \in \Rnn$.

One classic metric on $\Rnn$ is the one induced by the Frobenius norm ${\Vert M \Vert} _ F = \sqrt{\T (MM^T)} = \sqrt{ \sum_{1 \leq i,j \leq n} M_{ij}^2}$.

However, using this metric would mean we’re doing a Euclidean ($L^2$) projection, which intuitively doesn’t feel right here given the probabilistic / optimal transport interpretation of bistochastic matrices. It would also require solving a linear programming problem in $O(n^3)$, which isn’t ideal.

Instead, one much more sensible metric⁶ is the KL divergence. Since we’re dealing with bistochastic matrices, we use the generalized KL divergence:

The extra terms serve a purpose: $P_{ij}$ ensures $\KL(P\Vert M)$ is minimized iff $P=M$, and $M_{ij}$ guarantees non-negativity.

Thus, the problem we want to solve is:

Crucially, $P \mapsto \KL(P\Vert M)$ is strictly convex and $\B$ is convex, meaning that $(S)$ has a unique solution!

II.B. Deriving the algorithm

Now that we’ve defined the projection as solving $(S)$, let’s derive Sinkhorn’s algorithm. We’ll tackle this problem just like any optimization problem, using Lagrange’s multipliers.

First, we introduce the Lagrangian:

where $\mathbf{f}, \mathbf{g} \in \R^n$ are Lagrange multipliers.

Then, solving for $\frac{\partial \mathcal{L}}{\partial P_{ij}}=0$ yields $P_{ij}=e^{-f_i} M_{ij} e^{-g_j}=u_i M_{ij} v_j$ where we introduced $\mathbf{u}=\exp{(-\mathbf{f})}$ and $\mathbf{v}=\exp{(-\mathbf{g})}$.

Next, plugging this into $\frac{\partial \mathcal{L}}{\partial f_i}=0$ yields $u_i = 1 / (\sum_j P_{ij}v_j)$, i.e. $\mathbf{u}=1 \oslash (M\mathbf{v})$.

Likewise, we get $\mathbf{v}=1 \oslash (M^T\mathbf{u})$.

Finally, we know that the (unique!) solution of $(S)$ is of the form

where $\mathbf{u}, \mathbf{v}$ satisfy the above equations.

This leads to the following iterative algorithm, known as Sinkhorn’s algorithm⁷:

Note that $n_{iter}$ is a hyperparameter that controls the accuracy of the projection: the larger it is, the closer $P$ is to the true projection. In practice, $n_{iter}$ is often set between 10 and 20 ($n_{iter}=20$ in mHC).

III. Triton kernels for Sinkhorn

What’s remarkable (and indeed remarked in the age of GPU scientific computing⁸) about Sinkhorn’s algorithm is its simplicity: it only requires matrix-vector multiplications and element-wise operations! This makes it extremely well-suited for GPU implementations, as these operations can be efficiently parallelized between the (tens of) thousands of cores of modern GPUs.

However, the looping nature of the algorithm introduces severe memory-boundedness if implemented naively, as each iteration requires reading and writing the entire matrix $M$ as well as the vectors $\mathbf{u}, \mathbf{v}$. This can be mitigated by fusing the algorithm into a single kernel, which helps dramatically reduce memory access overhead and improve performance.

Indeed, if we look at the memory access pattern of the Sinkhorn algorithm, we have:

Thus, excluding the final write, memory access scales with $n_{iter}$, which is very inefficient.

Also, if we look at FLOPS, we see that each iteration requires only $4n^2+2n$ FLOPS for $2n^2+6n$ memory accesses. With 4 bytes per word since we’re using FP32 precision, this yields an arithmetic intensity of roughly 0.5 FLOPS/byte, which is terrible. Hence, Sinkhorn’s algorithm is completely memory-bound.

Sinkhorn’s algorithm memory-boundedness can be addressed with kernel fusion. We will now explore increasingly complex implementations based on Triton kernels⁹.

III.A. Hardware setup

We’re benchmarking on a NVIDIA RTX 4000 Ada Generation (released Aug 9th, 2023) with:

• 20GB of VRAM (GDDR6, not HBM)
• 360 GB/s memory bandwidth and 160-bit memory bus
• 48 Streaming Multiprocessors (SMs)
• 192 Tensor Cores (4th gen, 4 per SM)
• 6144 CUDA cores (Ada architecture, 128 per SM)

Delivering a peak performance of:

• 327 TFLOPS for Tensor Cores (Achtung: that’s in FP8 and “with sparsity”¹⁰)
• 26 TFLOPS for CUDA Cores (FP32)

This isn’t some fancy dual-die GB300 or whatever, but it’s still enough to do some serious benchmarking. Let’s get to it!

III.B. Mini-refresher on memory hierarchy

I’ve said that Sinkhorn’s algorithm is memory-bound. This may seem unclear if you’re not familiar with “memory hierarchy”. I’ll recap it very briefly here, focusing on GPUs.

Basically, GPUs have a hierarchical memory system with different types of memory at different levels of the hierarchy. While the full theory is quite complicated, a simple binary model of GPU memory is sufficient to grasp the main challenges. Thus, you can think of GPU memory as being split in two components:

1. SRAM where “S” stands for “Static”, in reference to the static nature of the stored information (uses a 6-transistor cell).
2. VRAM where “V” stands for “Video” since GPUs were mostly for video games & animation historically. It’s also called DRAM because it is DRAM i.e. it uses dynamic memory cells (based on capacitors which must be refreshed frequently). Whenever we talk about HBM, we’re talking about this memory¹¹.

The key is that SRAM is fast but small, whereas VRAM is slow but large. One goal of GPU kernels is to maximize the amount of computation done in SRAM (i.e. inside the registers) and minimize the amount of data exchanged with VRAM. That’s because moving data between VRAM and SRAM is extremely slow (orders of magnitude) compared to computation inside the registers.

Below is a table summarizing the key differences between SRAM and VRAM. Figure 5 gives a more detailed view of the memory hierarchy of the NVIDIA H100 GPU.

Note: despite its name, “local memory” actually lives in the GPU’s VRAM. This memory is accessed when registers are full but we need more storage capacity. Register spill over into local memory is something we really want to avoid as it completely kills performance due to the slow access time of VRAM.

The key thing to remember is that by default, data lives in VRAM until you want to do stuff with it and then it gets loaded into SRAM for computation, then the result is written back to VRAM.

This means that if you implement things naively, each operation costs you a back and forth between VRAM and SRAM. For example, imagine you have 2 matrices $A$ and $B$ and want to compute their product $C$. Three things happen sequentially:

1. $A$ and $B$ are loaded into SRAM. (VRAM → SRAM)
2. The computation is done in SRAM. (no memory transfer)
3. The result $C$ is written back to VRAM. (SRAM → VRAM)

This seems fair. But now imagine you have a matrix $A$ and want to do 2 things on it: first square it, then exponentiate it, both element-wise. If you do it naively, you’ll have to do 4 back and forths between VRAM and SRAM:

1. $A$ is loaded into SRAM. (VRAM → SRAM)
2. $A$ is squared element-wise in SRAM. (no memory transfer)
3. The result $A^2$ is written back to VRAM. (SRAM → VRAM)
4. $A^2$ is loaded into SRAM. (VRAM → SRAM)
5. $A^2$ is exponentiated element-wise in SRAM. (no memory transfer)
6. The result $\exp(A^2)$ is written back to VRAM. (SRAM → VRAM)

This is clearly suboptimal. What if we could do both operations in one go? That’s where fused kernels come in. A fused kernel is a kernel that does multiple sequential operations in the registers, hence minimizing the number of back and forths between VRAM and SRAM. This is exactly what we’re going to do in the next section.¹²

Also, whenever we refer to registers, think of it as a more granular way to refer to SRAM. In practice, the SRAM is more than just registers (it also includes shared memory & L1/L2 cache) but for the purpose of this blog post, you can equate these two concepts in your mind.

There’s a lot more to be said about memory hierarchy but we’ll stop here for now. If you want to learn more, I recommend Aleksa Gordic’s deep dive.

III.C. Implementations

Now that the hardware is out of the way, let’s dive into the implementations. We’ll start with a naive PyTorch implementation, then move on to the Triton kernels!

Naive PyTorch implementation

We begin with a simple PyTorch implementation for reference.

import torch

torch.set_float32_matmul_precision("high")  # we want to leverage Tensor Cores!

def sinkhorn_pytorch(
    log_M: torch.Tensor,  # logits
    n_iter: int,  # increase for better convergence
    epsilon: float,  # numerical stability
) -> torch.Tensor:
    """PyTorch baseline for comparison with Triton kernels."""
    M = torch.exp(log_M)
    M_T = M.transpose(-1, -2)  # O(1) stride manipulation
    B, N, _ = M.shape

    # initialize scalers
    u = torch.ones(B, N, 1, device=M.device)
    v = torch.ones(B, N, 1, device=M.device)

    # loop
    for _ in range(n_iter):
        u = 1.0 / (M @ v + epsilon)  # row normalization
        v = 1.0 / (M_T @ u + epsilon)  # column normalization

    # final scaled matrix
    return u * M * v.transpose(-1, -2)

This version is highly inefficient because $M, \mathbf{u}, \mathbf{v}$ are each read from and written to global memory at each iteration, effectively scaling memory access with $n_{iter}$. We can be smarter!

In addition, we define a slightly more optimized baseline, which technically is a fused kernel but uses poor auto-fusion heuristics provided by PyTorch’s torch.compile. We also use torch.inference_mode to drop the computational graph and tensor version tracking since we’re only interested in the forward pass. Thus, this second baseline has no unfair disadvantage compared to the Triton kernels coming next.

Formally, we construct this second baseline as a wrapper of the first one:

sinkhorn_pytorch_compiled = torch.inference_mode()(torch.compile(sinkhorn_pytorch))

This second version is still inefficient, but it’s a good sanity check to ensure that our Triton kernels are actually faster.

Basic fused kernel

Currently we have two types of objects doing back and forths between global memory and registers:

1. vector scalers $\mathbf{u}$, $\mathbf{v}$ of size $n$
2. matrix $M$ of size $(n,n)$

In my first kernel attempt, I decided to have $\mathbf{u}$ and $\mathbf{v}$ live in registers but keep $M$ in global memory for now. This gave a kernel looking something like the pseudo-code below:

Note that we use an accumulator $\mathbf{t}$ for the column normalization step. This is to avoid reading $M$ column-by-column, which is inefficient due to strided memory access¹³ (whereas reading $M$ row-by-row offers efficient coalesced memory access).

This first kernel is a good start as we’ve reduced the I/O bandwidth caused by $\mathbf{u}$ and $\mathbf{v}$’s back-and-forths in global memory, but $M$ is still read twice per iteration, meaning that memory access still scales with $n_{iter}$. We can do much better!

Loading $M$ in registers

First of all, we need a bit of context: remember that we’re using Sinkhorn’s algorithm to project $\Hres_l$ onto Birkhoff’s polytope. But $\Hres_l$ is a tiny matrix since the scaling factor $n$ of the residual stream is a small integer value. For the sake of simplicity, we’ll stick with $n=4$ like in mHC. This means that $M$ is effectively a $4\times 4$ matrix, which can easily fit in registers! Thus, we can update the kernel to have $M$ live in the registers, which will save us a lot more I/O bandwidth!

This gives us exactly the same algorithm as above, except that $M$ is loaded in the registers at the beginning of the kernel, so we don’t need to read it from memory at every row / column normalization step! One other difference: the exp operation on $M$ is now done on-the-fly inside the registers, which saves one back-and-forth between global memory and registers.

Block packing

The previous solution may seem optimal, but we can in fact still do much better!

I realized this by experimenting on my own and observing the unoptimized PyTorch solution beating my kernel for batch sizes of 2048 and above. The reason is simple: PyTorch is smart enough to pack small matrices together whenever it can, helping the GPU better saturate its cores. More precisely, PyTorch (via cuBLAS) uses highly optimized BMM (batch matrix multiplication) kernels that process multiple matrices per thread block to saturate the SMs. In contrast, our naive kernel assigns one block per $4\times4$ matrix, leaving the vast majority of threads in the warp idle.

So, we too need to process our matrices in batches instead of one-by-one. To do so, we can pack our $B$ matrices into

tensors of shape $(\text{BLOCK_SIZE},n,n)$, and then process a full tensor per block. This technique is known as block packing¹⁴, it’s useful when the data objects you’re processing are too small to saturate the GPU’s threads. See Figure 5 for an illustration.

This is much more efficient than using one block per matrix, which is what the two previous kernels are doing. The reason for that is that a block has significant resources: 4 warps by default, each with 32 threads, meaning a total of 128 threads. Thus, assigning a single $4\times 4$ matrix per block means we’re effectively dedicating 8 threads per matrix element, which is way too much and leads to underutilization i.e. idle threads. This in turn greatly slows down the kernel as it hinders parallelism (since the idle threads cannot be used to process other matrices).

Choosing the right $\text{BLOCK_SIZE}$ i.e. how many matrices to pack in a single block is a bit of an art. One easy method is to use on triton.autotune to benchmark different values. I did so for $n=4$ and found that $\text{BLOCK_SIZE}=64$ yielded the best results.¹⁵

Our last kernel thus reuses the same layout as the previous one, except that we now process a full tensor of shape $(\text{BLOCK_SIZE},n,n)$ per block. For $n=4$, we thus end up using tensors of dimension $(64, 4, 4)$.

III.D. Benchmark

We have presented in total 5 implementations of Sinkhorn’s algorithm:

1. naive PyTorch (sinkhorn_pytorch)
2. naive Pytorch but auto-fused and with inference mode (sinkhorn_pytorch_compiled)
3. Triton kernel with scalers $\mathbf{u}, \mathbf{v}$ in registers but $M$ in global memory (sinkhorn_A_in_global_memory)
4. Triton kernel with $\mathbf{u}, \mathbf{v}, M$ in registers (sinkhorn_A_in_registers)
5. Triton kernel with $\mathbf{u}, \mathbf{v}, M$ in registers and block packing (sinkhorn_A_in_registers_block_packing)

Note: A refers to the input matrix $M$, as I happened to have used different naming conventions in the code.

We will now benchmark these implementations on our NVIDIA RTX 4000 Ada Generation.

We’ll stick to $n=4$ to have simple 2D plots instead of heatmaps / 3D plots.

As for the batch size $B$, we will sweep from $1$ to $2^{24}\sim 16M$. The reason for sweeping to such large batch sizes is twofold:

1. because $\Hres_l$ is defined at the token-level, it means that for a given network layer $l$, we need as many Sinkhorn projections as we have tokens in our batch. Thus, the relevant batch size $B$ to benchmark on is the microbatch size $mbs$. Today’s pretraining pipelines commonly use microbatch sizes in the range $mbs\in[4096, 16384]$¹⁶, which justifies benchmarking at least up to $B=16k$.
2. it allows us to escape the uninteresting latency-bound regime and get to the much cooler memory-bound regime (that’s why we push $B$ much further than the $16k$ required for pretraining)

III.E. Figures

Vanity metric

Figure 8 shows the speedup of the most optimized Triton kernel (sinkhorn_A_in_registers_block_packing) over the compiled PyTorch implementation (sinkhorn_pytorch_compiled). We can see that for batch sizes below 1k, both implementations operate in the latency-bound regime, hence the speedup is roughly constant (and scales with $n_{iter}$). Past this threshold, we enter the memory-bound regime where the Triton kernel’s I/O awareness shines!

In any case, the speedup ranges from 20x to 139x, which is pretty cool!

Comparing implementations

The next three figures compare the implementations across three performance metrics:

1. Execution time (lower is better) on Figure 8
2. Memory bandwidth (higher is better) on Figure 9
3. Compute throughput (higher is better) on Figure 10

Note that the execution time uses the median (instead of the mean) to avoid being skewed by outliers, since kernel execution times tend to exhibit positive skew.

We also compute the 98% confidence interval (q01 to q99) to ensure the statistical significance of the results.

Also, the memory bandwidth and compute throughput are computed as follows:

• Memory bandwidth = $\frac{\text{total bytes read from/written to global memory}}{\text{median execution time}}$
• Compute throughput = $\frac{\text{total FLOPS}}{\text{median execution time}}$

Across all three metrics, we observe an order relation: sinkhorn_pytorch < sinkhorn_pytorch_compiled < sinkhorn_A_in_global_memory < sinkhorn_A_in_registers < sinkhorn_A_in_registers_block_packing, as expected!

Also, we consistently witness two regimes:

1. $B\ll 1k$ is the latency-bound regime: the amount of data is not sufficient to keep all the GPU’s SMs busy, and as such execution time is roughly constant. In this regime, the GPU’s throughput scales linearly with batch size as more SMs are utilized, meaning we can effectively scale memory bandwidth and compute throughput “for free” thanks to increased parallelism.
2. $B\gg 1k$ is the memory-bound regime: there’s now enough data to use all the GPU’s SMs, and as such execution time now scales linearly with batch size i.e. no more “free” performance from increased parallelism. That’s why global throughput mostly plateaus beyond $B=16k$.

References: