Optimizing Loops for Softmax and FlashAttention in Triton
Puzzle 8: Long Softmax
This puzzle implements a batched softmax operation. Given a batch of logits of shape [N0, T], compute the softmax for each row. The spec is:
def softmax_spec(x: Float32[4, 200]) -> Float32[4, 200]:
x_max = x.max(1, keepdim=True)[0]
x = x - x_max
x_exp = x.exp()
return x_exp / x_exp.sum(1, keepdim=True)
The challenge requires two solutions: a brrute-force three-loop version and an optimized two-loop version using the online softmax trick.
Brute-Force Approach (Three Loops)
- Loop 1: Find the maximum value in each row.
- Loop 2: Subtract the row maximum from each element, exponentiate, and accumulate the sum.
- Loop 3: Compute the final softmax by dividing each exponentiated value by the row sum.
Implementation:
@triton.jit
def exp_approx(x):
return tl.exp2(1.44269504 * x)
@triton.jit
def softmax_kernel_brute_force(x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr):
block_id_i = tl.program_id(0)
offset_x = block_id_i * B0 + tl.arange(0, B0)
mask_x = offset_x < N0
row_max = tl.zeros([B0, 1], dtype=tl.float32)
# Loop 1: compute row max
for idj in tl.range(0, T, B1):
offset_y = idj + tl.arange(0, B1)
mask_y = offset_y < T
offset_xy = offset_x[:, None] * T + offset_y[None, :]
mask_xy = mask_x[:, None] & mask_y[None, :]
block_value = tl.load(x_ptr + offset_xy, mask_xy, other=float('-inf'))
row_max = tl.maximum(row_max, tl.max(block_value, axis=1, keep_dims=True))
row_sum_exp = tl.zeros([B0, 1], dtype=tl.float32)
# Loop 2: subtract max and compute exp sum
for idj in tl.range(0, T, B1):
offset_y = idj + tl.arange(0, B1)
mask_y = offset_y < T
offset_xy = offset_x[:, None] * T + offset_y[None, :]
mask_xy = mask_x[:, None] & mask_y[None, :]
block_value = tl.load(x_ptr + offset_xy, mask_xy, other=float('-inf'))
block_value -= row_max
row_sum_exp += tl.sum(exp_approx(block_value), axis=1, keep_dims=True)
# Loop 3: compute final softmax
for idj in tl.range(0, T, B1):
offset_y = idj + tl.arange(0, B1)
mask_y = offset_y < T
offset_xy = offset_x[:, None] * T + offset_y[None, :]
mask_xy = mask_x[:, None] & mask_y[None, :]
block_value = tl.load(x_ptr + offset_xy, mask_xy, other=float('-inf'))
softmax_value = exp_approx(block_value - row_max) / row_sum_exp
tl.store(z_ptr + offset_xy, softmax_value, mask_xy)
return
Optimized Approach (Two Loops with On line Softmax)
The online softmax algorithm combines maximum update and sum accumulation in one pass. The recurrence is:
$$\begin{aligned} m_0 &\leftarrow -\infty \ d_0 &\leftarrow 0 \ \text{for } j &\leftarrow 1, V \text{ do} \ &\quad m_j \leftrarow \max(m_{j-1}, x_j) \ &\quad d_j \leftarrow d_{j-1} \times e^{m_{j-1}-m_j} + e^{x_j-m_j} \ \text{end for} \ \text{for } i &\leftarrow 1, V \text{ do} \ &\quad y_i \leftarrow \frac{e^{x_i-m_V}}{d_V} \ \text{end for}\end{aligned}$$
Implementation:
@triton.jit
def exp_approx(x):
return tl.exp2(1.44269504 * x)
@triton.jit
def softmax_kernel(x_ptr, z_ptr, N0, N1, T, B0: tl.constexpr, B1: tl.constexpr):
block_id_i = tl.program_id(0)
offset_x = block_id_i * B0 + tl.arange(0, B0)
mask_x = offset_x < N0
row_max = tl.zeros([B0, 1], dtype=tl.float32)
row_sum_exp = tl.zeros([B0, 1], dtype=tl.float32)
for idj in tl.range(0, T, B1):
offset_y = idj + tl.arange(0, B1)
mask_y = offset_y < T
offset_xy = offset_x[:, None] * T + offset_y[None, :]
mask_xy = mask_x[:, None] & mask_y[None, :]
block_value = tl.load(x_ptr + offset_xy, mask_xy, other=float('-inf'))
new_max = tl.maximum(tl.max(block_value, axis=1, keep_dims=True), row_max)
row_sum_exp = row_sum_exp * exp_approx(row_max - new_max) + tl.sum(exp_approx(block_value - new_max), axis=1, keep_dims=True)
row_max = new_max
for idj in tl.range(0, T, B1):
offset_y = idj + tl.arange(0, B1)
mask_y = offset_y < T
offset_xy = offset_x[:, None] * T + offset_y[None, :]
mask_xy = mask_x[:, None] & mask_y[None, :]
block_value = tl.load(x_ptr + offset_xy, mask_xy, other=float('-inf'))
softmax_value = exp_approx(block_value - row_max) / row_sum_exp
tl.store(z_ptr + offset_xy, softmax_value, mask_xy)
return
Puzzle 9: Simple FlashAttention
This puzzle implements a scalar version of FlashAttention. Given query row q, keys K, values V, compute:
$$z_i = \sum_{j=1}^{T} \text{softmax}(q_i \cdot k_1, \ldots, q_i \cdot k_T)_j \cdot v_j$$
The online softmax trick extends to also accumulate the weighted sum of values in a single loop. The recurrence is:
$$\begin{aligned} x_i &\leftarrow Q[k,:] \cdot K^T[:,i] \ m_i &\leftarrow \max(m_{i-1}, x_i) \ d_i' &\leftarrow d_{i-1}' \cdot e^{m_{i-1} - m_i} + e^{x_i - m_i} \ O_i' &\leftarrow O_{i-1}' \cdot \frac{d_{i-1}'}{d_i'} \cdot e^{m_{i-1} - m_i} + \frac{e^{x_i - m_i}}{d_i'} \cdot V[i,:] \ \end{aligned}$$
Implementation:
@triton.jit
def exp_approx(x):
return tl.exp2(1.44269504 * x)
@triton.jit
def flashatt_kernel(
q_ptr, k_ptr, v_ptr, z_ptr, N0, T, B0: tl.constexpr, B1: tl.constexpr
):
block_id_i = tl.program_id(0)
off_i = block_id_i * B0 + tl.arange(0, B0)
mask_i = off_i < N0
inf = 1.0e6
q = tl.load(q_ptr + off_i, mask=mask_i)
exp_sum = tl.zeros((B0,), dtype=tl.float32)
qk_max = tl.full((B0,), -inf, dtype=tl.float32)
z = tl.zeros((B0,), dtype=tl.float32)
for id_j in tl.range(0, T, B1):
off_j = id_j + tl.arange(0, B1)
mask_j = off_j < T
mask_ij = mask_i[:, None] & mask_j[None, :]
k = tl.load(k_ptr + off_j, mask=mask_j)
qk = q[:, None] * k[None, :] + tl.where(mask_ij, 0, -inf)
new_max = tl.maximum(tl.max(qk, axis=1), qk_max)
qk_exp = exp_approx(qk - new_max[:, None])
factor = exp_approx(qk_max - new_max)
new_exp_sum = exp_sum * factor + tl.sum(qk_exp, axis=1)
v = tl.load(v_ptr + off_j, mask=mask_j, other=0.0)
z = z * factor + tl.sum(qk_exp * v[None, :], axis=1)
qk_max = new_max
exp_sum = new_exp_sum
z = z / exp_sum
tl.store(z_ptr + off_i, z, mask=mask_i)
return
References
- Tri Dao, et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." arXiv:2205.14135.
- Online Softmax: https://arxiv.org/pdf/1805.02867
- Triton Official Tutorials
