• background
• softmax - vector form
• gradient of softmax (vector form)
  • Jacobian matrix
  • gradient of \(\frac{\partial L}{\partial \mathbf{x}}\)
  • avoid explicit Jacobian
• softmax - batch form
  • forward pass
  • backward pass
• implementation
  • real forward pass
  • real backward pass
  • a real example
  • native pytorch implementation
  • triton implementation
  • speed comparison
• notions

background

The softmax function is a fundamental operation in deep learning that converts vectors of real numbers into probability distributions. This blog post explores the softmax function’s implementation and optimization using Triton, a programming framework for efficient GPU computations.

The softmax function transforms an input vector into a probability distribution where all elements sum to 1.

softmax - vector form

where:

• \(\mathbf{x} \in \mathbb{R}^d\): input vector.
• \(\mathbf{o} \in \mathbb{R}^d\): output vector, probability distribution.

gradient of softmax (vector form)

We will compute gradients \(\frac{\partial L}{\partial \mathbf{x}}\) given \(\frac{\partial L}{\partial \mathbf{o}}\), where \(L\) is loss function, \(\mathbf{o}\) is softmax output.

Jacobian matrix

softmax is a vector function, the Jacobian matrix is the matrix of all partial derivatives:

For softmax, the derivative has two cases:

1. when \(i = j\), consider \(\mathbf{o}_i = \frac{e^{\mathbf{x}_i}}{\sum_{j=1}^{d} e^{\mathbf{x}_j}}\),

   \[\begin{aligned} \displaystyle \frac{\partial \mathbf{o}_i}{\partial \mathbf{x}_i} &= \frac{ \frac{\partial \left( e^{\mathbf{x}_i} \right)}{\partial \mathbf{x}_i} \cdot \sum_{j=1}^{d} e^{\mathbf{x}_j} - \frac{\partial \left( \sum_{j=1}^{d} e^{\mathbf{x}_j} \right)}{\partial \mathbf{x}_i} \cdot e^{\mathbf{x}_i} }{\left( \sum_{j=1}^{d} e^{\mathbf{x}_j} \right)^2} \\ \displaystyle & = \frac{e^{\mathbf{x}_i} \cdot \sum_{j=1}^{d} e^{\mathbf{x}_j} - e^{\mathbf{x}_i} \cdot e^{\mathbf{x}_i}}{\left( \sum_{j=1}^{d} e^{\mathbf{x}_j} \right)^2} \\ & = \frac{e^{\mathbf{x}_i}}{\sum_{j=1}^{d} e^{\mathbf{x}_j}} \left( 1 - \frac{e^{\mathbf{x}_i}}{\sum_{j=1}^{d} e^{\mathbf{x}_j}} \right) \\ & = \mathbf{o}_i (1 - \mathbf{o}_i) \end{aligned}\]

2. similarly, when \(i \neq j\):

Thus, \((i,j)\)-th element in Jacobian matrix will be:

where \(\mathbf{J}\) has shape \([d \times d]\) and \(\delta_{ij}\) is the Kronecker delta, which is 1 if \(i = j\) and 0 otherwise.

In matrix form, the Jacobian of the softmax is:

where:

• \(\mathbf{o}\) is the output of softmax, the shape is \([d]\).
• \(\mathrm{diag}(\mathbf{o})\) is a diagonal matrix of \(\mathbf{o}\), the shape is \([d \times d]\).
• \(\mathbf{o}\mathbf{o}^\top\) is the outer product of \(\mathbf{o}\) with itself, the shape is \([d \times d]\).

gradient of \(\frac{\partial L}{\partial \mathbf{x}}\)

Given \(\frac{\partial L}{\partial \mathbf{o}}\), we can compute \(\frac{\partial L}{\partial \mathbf{x}}\) using the Jacobian matrix:

where \(\frac{\partial L}{\partial \mathbf{o}}\) has shape \([d]\), \(\mathbf{J}^{\top}\) has shape \([d \times d]\), and \(\frac{\partial L}{\partial \mathbf{x}}\) has shape \([d]\).

avoid explicit Jacobian

Consider the matrix multiplication:

For the \(i\)-th element of \(\frac{\partial L}{\partial \mathbf{x}}\), we can decompose the computation:

This leads to an efficient vector form:

softmax - batch form

\(\mathbf{X}\): A batch of input vectors.

where:

• \(N\) is batch size.
• \(d\) is vector dimension.

forward pass

where \(\mathbf{E} \in \mathbb{R}^{N \times d}\), \(\mathbf{s} \in \mathbb{R}^{N \times 1}\), \(\mathbf{O} \in \mathbb{R}^{N \times d}\).

backward pass

We have gradient with respect to softmax output:

we compute the gradient:

where \(\mathbf{O}\) has size \([N \times d]\), and \(\frac{\partial L}{\partial \mathbf{O}}\) has size \([N \times d]\).

where \(\frac{\partial L}{\partial \mathbf{X}} \in \mathbb{R}^{N \times d}\) and \(\mathbf{O} \in \mathbb{R}^{N \times d}\) and \(\mathbf{s}_{grad} \in \mathbb{R}^{N \times 1}\) will be broadcasted to \(\mathbb{R}^{N \times d}\).

implementation

In practice, we subtract the maximum value from each row before applying exp() to prevent numerical overflow:

real forward pass

For input \(\mathbf{X} \in \mathbb{R}^{N \times d}\):

real backward pass

we have \(\frac{\partial L}{\partial \mathbf{O}} \in \mathbb{R}^{N \times d}\) and cached \(\mathbf{O} \in \mathbb{R}^{N \times d}\)

a real example

give a real example to show how to implement softmax and its backward pass in pytorch and triton.

forwards pass is as follows:

backward pass is as follows: \(dO = \begin{bmatrix} 0.1 & 0.2 & 0.7 \\ 0.2 & 0.3 & 0.5 \end{bmatrix}\)

native pytorch implementation

import torch
import torch.nn.functional as F

# Custom Forward Pass (Numerically Stable Softmax)
def softmax_forward(X):
    X_max = torch.max(X, dim=1, keepdim=True)[0]  # Shape: (N, 1)
    E = torch.exp(X - X_max)                     # Shape: (N, d)
    S = torch.sum(E, dim=1, keepdim=True)        # Shape: (N, 1)
    O = E / S                                    # Shape: (N, d)
    return O

# Custom Backward Pass (Gradient Calculation)
def softmax_backward(dL_dO, O):
    s_grad = torch.sum(O * dL_dO, dim=1, keepdim=True)  # Shape: (N, 1)
    dL_dX = O * (dL_dO - s_grad)                        # Shape: (N, d)
    return dL_dX

# Example Inputs
X = torch.tensor([[1.0, 2.0, 3.0], [1.0, 3.0, 5.0]], requires_grad=True)
dL_dO = torch.tensor([[0.1, 0.2, 0.7], [0.2, 0.3, 0.5]])

# Custom Implementation - Forward
O_custom = softmax_forward(X)

# PyTorch Implementation - Forward
O_pytorch = F.softmax(X, dim=1)

# Verify Forward Output
print("Custom Softmax Output:\n", O_custom)
print("PyTorch Softmax Output:\n", O_pytorch)
print("Forward Pass Match:", torch.allclose(O_custom, O_pytorch))

# Custom Implementation - Backward
dL_dX_custom = softmax_backward(dL_dO, O_custom)

# PyTorch Automatic Gradient Calculation
O_pytorch.backward(dL_dO)  # Computes gradient using PyTorch autograd
dL_dX_pytorch = X.grad

# Verify Backward Output
print("\nCustom Gradient w.r.t Input:\n", dL_dX_custom)
print("PyTorch Gradient w.r.t Input:\n", dL_dX_pytorch)
print("Backward Pass Match:", torch.allclose(dL_dX_custom, dL_dX_pytorch))

output:

Custom Softmax Output:
 tensor([[0.0900, 0.2447, 0.6652],
        [0.0159, 0.1173, 0.8668]], grad_fn=<DivBackward0>)
PyTorch Softmax Output:
 tensor([[0.0900, 0.2447, 0.6652],
        [0.0159, 0.1173, 0.8668]], grad_fn=<SoftmaxBackward0>)
Forward Pass Match: True

Custom Gradient w.r.t Input:
 tensor([[-0.0381, -0.0792,  0.1173],
        [-0.0043, -0.0202,  0.0245]], grad_fn=<MulBackward0>)
PyTorch Gradient w.r.t Input:
 tensor([[-0.0381, -0.0792,  0.1173],
        [-0.0043, -0.0202,  0.0245]])
Backward Pass Match: True

triton implementation

from typing import Optional

import torch
import triton
import triton.language as tl


@triton.jit
def softmax_fwd_kernel(
    X,
    O,
    D: tl.constexpr,
    B: tl.constexpr
):
    i_n = tl.program_id(0)
    o_d = tl.arange(0, B)
    m_d = o_d < D

    X_max = tl.max(tl.load(X + i_n * D + o_d, mask=m_d, other=-float('inf')), 0)
    E = tl.exp(tl.load(X + i_n * D + o_d, mask=m_d, other=-float('inf')) - X_max)
    S = tl.sum(E, 0)
    P = E / S

    tl.store(O + i_n * D + o_d, P.to(O.dtype.element_ty), mask=m_d)


@triton.jit
def softmax_bwd_kernel(
    O,
    dO,
    dX,
    D: tl.constexpr,
    B: tl.constexpr
):
    i_n = tl.program_id(0)
    o_d = tl.arange(0, B)
    m_d = o_d < D

    P = tl.load(O + i_n * D + o_d, mask=m_d, other=0.)
    dP = tl.load(dO + i_n * D + o_d, mask=m_d, other=0.)
    s_grad = tl.sum(P * dP, 0)
    dX_row = P * (dP - s_grad)

    tl.store(dX + i_n * D + o_d, dX_row.to(dX.dtype.element_ty), mask=m_d)


def softmax_fwd(
    X: torch.Tensor,
    dtype: Optional[torch.dtype] = torch.float
) -> torch.Tensor:
    shape = X.shape
    X = X.view(-1, X.shape[-1])

    N, D = X.shape
    B = triton.next_power_of_2(D)

    O = torch.empty_like(X, dtype=dtype)
    softmax_fwd_kernel[(N,)](
        X=X,
        O=O,
        D=D,
        B=B
    )
    return O.view(*shape)


def softmax_bwd(
    O: torch.Tensor,
    dO: torch.Tensor,
    dtype: Optional[torch.dtype] = torch.float
) -> torch.Tensor:
    shape = O.shape
    O = O.view(-1, O.shape[-1])
    dX = torch.empty_like(O, dtype=dtype)

    N, D = O.shape
    B = triton.next_power_of_2(D)
    softmax_bwd_kernel[(N,)](
        O=O,
        dO=dO,
        dX=dX,
        D=D,
        B=B
    )
    return dX.view(*shape)




# Test code to verify correctness
import torch.nn.functional as F

# Example inputs
X = torch.tensor([[1.0, 2.0, 3.0], [1.0, 3.0, 5.0]], requires_grad=True, device='cuda')
dP = torch.tensor([[0.1, 0.2, 0.7], [0.2, 0.3, 0.5]], device='cuda')

# Forward pass
P_triton = softmax_fwd(X)
P_torch = F.softmax(X, dim=1)

# Verify forward pass
print( "P_triton:\n", P_triton)
print( "P_torch:\n", P_torch)
print("Forward Pass Match:", torch.allclose(P_triton, P_torch))

# Backward pass

dX_triton = softmax_bwd(P_triton, dP)
P_torch.backward(dP)
dX_torch = X.grad

# Verify backward pass
print( "dX_triton:\n", dX_triton)
print( "dX_torch:\n", dX_torch)
print("Backward Pass Match:", torch.allclose(dX_triton, dX_torch))

output:

P_triton:
 tensor([[0.0900, 0.2447, 0.6652],
        [0.0159, 0.1173, 0.8668]], device='cuda:0')
P_torch:
 tensor([[0.0900, 0.2447, 0.6652],
        [0.0159, 0.1173, 0.8668]], device='cuda:0', grad_fn=<SoftmaxBackward0>)
Forward Pass Match: True
dX_triton:
 tensor([[-0.0381, -0.0792,  0.1173],
        [-0.0043, -0.0202,  0.0245]], device='cuda:0')
dX_torch:
 tensor([[-0.0381, -0.0792,  0.1173],
        [-0.0043, -0.0202,  0.0245]], device='cuda:0')
Backward Pass Match: True

speed comparison

The performance comparison between PyTorch and Triton implementations reveals:

Results show

• forward pass: triton implementation is stable, while the PyTorch implementation is faster for most batch sizes but shows fluctuations for a few.
• backward pass: triton implementation outperforms the pytorch implementation across most batch sizes. (the comparison may not be entirely fair, as triton caches the output \(O\), whereas pytorch’s handling intermediate values is unclear.)

notions

Note:

• Symbols like \(x\), \(\mathbf{x}\), \(\mathbf{X}\) represent scalars, vectors, or matrices, where uppercase denotes batch forms.
• \(\mathbf{X}_{:,i}\) denotes a column vector, \(\mathbf{X}_{i,:}\) denotes a row vector, \(\mathbf{X}_{i,j}\) and denote the \((i,j)\)-th element
• \(\mathbf{x}_i\) denote the \(i\)-th element.