y = x*w + b, one program per row

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*N
# litmus: buffer w_ptr = N
# litmus: buffer b_ptr = N
# litmus: buffer out_ptr = M*N
# litmus: grid = M

@triton.jit
def elementwise_scale_bias(x_ptr, w_ptr, b_ptr, out_ptr, M, N, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        x = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        w = tl.load(w_ptr + cols, mask=mask, other=0.0)
        b = tl.load(b_ptr + cols, mask=mask, other=0.0)
        tl.store(out_ptr + row * N + cols, x * w + b, mask=mask)


gelu

exact gelu, 0.5x(1+erf(x/sqrt2))

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*N
# litmus: buffer out_ptr = M*N
# litmus: grid = M

@triton.jit
def gelu(x_ptr, out_ptr, M, N, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        x = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        y = 0.5 * x * (1.0 + tl.erf(x * 0.7071067811865476))
        tl.store(out_ptr + row * N + cols, y, mask=mask)


row-wise sum with a strided row base

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_reduce(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros((BLOCK,), dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        acc += v
    tl.store(out_ptr + row, tl.sum(acc, axis=0))


row-wise max with a strided row base

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_max(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.full((BLOCK,), float('-inf'), dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=float('-inf'))
        acc = tl.maximum(acc, v)
    tl.store(out_ptr + row, tl.max(acc, axis=0))


softmax

numerically stable row-wise softmax

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*N
# litmus: buffer out_ptr = M*N
# litmus: grid = M

@triton.jit
def softmax(x_ptr, out_ptr, M, N, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    mx = float('-inf')
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=float('-inf'))
        mx = tl.maximum(mx, tl.max(v, axis=0))
    denom = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=float('-inf'))
        e = tl.where(mask, tl.exp(v - mx), 0.0)
        denom += tl.sum(e, axis=0)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=float('-inf'))
        tl.store(out_ptr + row * N + cols, tl.exp(v - mx) / denom, mask=mask)


layernorm

row-wise layernorm with weight and bias

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer w_ptr = N
# litmus: buffer b_ptr = N
# litmus: buffer out_ptr = M*N
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def layernorm(x_ptr, w_ptr, b_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    s1 = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        s1 += tl.sum(v, axis=0)
    mean = s1 / N
    s2 = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        d = tl.where(mask, v - mean, 0.0)
        s2 += tl.sum(d * d, axis=0)
    var = s2 / N
    rstd = 1.0 / tl.sqrt(var + 1e-5)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        w = tl.load(w_ptr + cols, mask=mask, other=0.0)
        b = tl.load(b_ptr + cols, mask=mask, other=0.0)
        tl.store(out_ptr + row * N + cols, (v - mean) * rstd * w + b, mask=mask)


matmul

tiled fp32 matmul, fp32 accumulator

import triton
import triton.language as tl

# litmus: constexpr BLOCK_M = 32
# litmus: constexpr BLOCK_N = 32
# litmus: constexpr BLOCK_K = 32
# litmus: buffer a_ptr = M*K
# litmus: buffer b_ptr = K*N
# litmus: buffer c_ptr = M*N
# litmus: grid = cdiv(M, BLOCK_M), cdiv(N, BLOCK_N)

@triton.jit
def matmul(a_ptr, b_ptr, c_ptr, M, N, K,
           BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pm = tl.program_id(0)
    pn = tl.program_id(1)
    offs_m = pm * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pn * BLOCK_N + tl.arange(0, BLOCK_N)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K)
        a = tl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :],
                    mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptr + offs_k[:, None] * N + offs_n[None, :],
                    mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
    tl.store(c_ptr + offs_m[:, None] * N + offs_n[None, :], acc,
             mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


fp16 storage with the correct fp32 accumulator


attention

toy single-head attention -- outside the certified fragment

import triton
import triton.language as tl

# litmus: constexpr BLOCK_D = 32
# litmus: buffer q_ptr = M*D
# litmus: buffer k_ptr = L*D
# litmus: buffer v_ptr = L*D
# litmus: buffer out_ptr = M*D
# litmus: grid = M

@triton.jit
def attention(q_ptr, k_ptr, v_ptr, out_ptr, M, L, D, BLOCK_D: tl.constexpr):
    i = tl.program_id(0)
    ds = tl.arange(0, BLOCK_D)
    dmask = ds < D
    q = tl.load(q_ptr + i * D + ds, mask=dmask, other=0.0)
    mx = float('-inf')
    for j in range(0, L):
        kv = tl.load(k_ptr + j * D + ds, mask=dmask, other=0.0)
        mx = tl.maximum(mx, tl.sum(q * kv, axis=0) * 0.17677669529663687)
    denom = 0.0
    acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
    for j in range(0, L):
        kv = tl.load(k_ptr + j * D + ds, mask=dmask, other=0.0)
        s = tl.sum(q * kv, axis=0) * 0.17677669529663687
        p = tl.exp(s - mx)
        denom += p
        vv = tl.load(v_ptr + j * D + ds, mask=dmask, other=0.0)
        acc += p * vv
    tl.store(out_ptr + i * D + ds, acc / denom, mask=dmask)


tail tile written and read without a mask

N=512 is a multiple of BLOCK=32 so no tail exists -- and even at a ragged shape the spilled write lands in a neighbouring allocation, so the kernel's own output is unchanged and no allclose test can see it

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*N
# litmus: buffer w_ptr = N
# litmus: buffer b_ptr = N
# litmus: buffer out_ptr = M*N
# litmus: grid = M

@triton.jit
def elementwise_scale_bias(x_ptr, w_ptr, b_ptr, out_ptr, M, N, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        x = tl.load(x_ptr + row * N + cols)
        w = tl.load(w_ptr + cols)
        b = tl.load(b_ptr + cols)
        tl.store(out_ptr + row * N + cols, x * w + b)


row sum reads past the end of each row

N=512 is a multiple of BLOCK=32, so no tail tile exists

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_reduce(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros((BLOCK,), dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols)
        acc += v
    tl.store(out_ptr + row, tl.sum(acc, axis=0))


softmax normalises over the unmasked tail

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*N
# litmus: buffer out_ptr = M*N
# litmus: grid = M

@triton.jit
def softmax(x_ptr, out_ptr, M, N, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    mx = float('-inf')
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols)
        mx = tl.maximum(mx, tl.max(v, axis=0))
    denom = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols)
        e = tl.where(mask, tl.exp(v - mx), 0.0)
        denom += tl.sum(e, axis=0)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols)
        tl.store(out_ptr + row * N + cols, tl.exp(v - mx) / denom)


last k-tile loaded without the K mask

M, N and K are all multiples of the block sizes

import triton
import triton.language as tl

# litmus: constexpr BLOCK_M = 32
# litmus: constexpr BLOCK_N = 32
# litmus: constexpr BLOCK_K = 32
# litmus: buffer a_ptr = M*K
# litmus: buffer b_ptr = K*N
# litmus: buffer c_ptr = M*N
# litmus: grid = cdiv(M, BLOCK_M), cdiv(N, BLOCK_N)

@triton.jit
def matmul(a_ptr, b_ptr, c_ptr, M, N, K,
           BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pm = tl.program_id(0)
    pn = tl.program_id(1)
    offs_m = pm * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pn * BLOCK_N + tl.arange(0, BLOCK_N)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K)
        a = tl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :],
                    mask=(offs_m[:, None] < M), other=0.0)
        b = tl.load(b_ptr + offs_k[:, None] * N + offs_n[None, :],
                    mask=(offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
    tl.store(c_ptr + offs_m[:, None] * N + offs_n[None, :], acc,
             mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


sum accumulator seeded with 1.0 instead of 0.0

a constant offset of 1.0 on a sum of ~256 uniform values is a 0.4% relative error, inside rtol=1e-2

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_reduce(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.full((BLOCK,), 1.0 / BLOCK, dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        acc += v
    tl.store(out_ptr + row, tl.sum(acc, axis=0))


max accumulator seeded with 0.0 instead of -inf

torch.rand() is non-negative, so max(0, row) == max(row) always -- the identity is never exercised

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_max(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.full((BLOCK,), 0.0, dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=float('-inf'))
        acc = tl.maximum(acc, v)
    tl.store(out_ptr + row, tl.max(acc, axis=0))


matmul accumulator seeded with 1.0

K=512 with uniform inputs gives entries of order 128, so a +1.0 offset is under the 1e-2 relative tolerance

import triton
import triton.language as tl

# litmus: constexpr BLOCK_M = 32
# litmus: constexpr BLOCK_N = 32
# litmus: constexpr BLOCK_K = 32
# litmus: buffer a_ptr = M*K
# litmus: buffer b_ptr = K*N
# litmus: buffer c_ptr = M*N
# litmus: grid = cdiv(M, BLOCK_M), cdiv(N, BLOCK_N)

@triton.jit
def matmul(a_ptr, b_ptr, c_ptr, M, N, K,
           BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pm = tl.program_id(0)
    pn = tl.program_id(1)
    offs_m = pm * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pn * BLOCK_N + tl.arange(0, BLOCK_N)
    acc = tl.full((BLOCK_M, BLOCK_N), 1.0, dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K)
        a = tl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :],
                    mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptr + offs_k[:, None] * N + offs_n[None, :],
                    mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
    tl.store(c_ptr + offs_m[:, None] * N + offs_n[None, :], acc,
             mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


softmax without the max shift: exp overflows

torch.rand() lives in [0,1), where exp(x) never approaches the fp32 range limit

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*N
# litmus: buffer out_ptr = M*N
# litmus: grid = M

@triton.jit
def softmax(x_ptr, out_ptr, M, N, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    mx = float('-inf')
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=float('-inf'))
        mx = tl.maximum(mx, tl.max(v, axis=0))
    denom = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=float('-inf'))
        e = tl.where(mask, tl.exp(v), 0.0)
        denom += tl.sum(e, axis=0)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=float('-inf'))
        tl.store(out_ptr + row * N + cols, tl.exp(v) / denom, mask=mask)


attention scores exponentiated without the max shift

unit-scale q and k keep the scores near zero

import triton
import triton.language as tl

# litmus: constexpr BLOCK_D = 32
# litmus: buffer q_ptr = M*D
# litmus: buffer k_ptr = L*D
# litmus: buffer v_ptr = L*D
# litmus: buffer out_ptr = M*D
# litmus: grid = M

@triton.jit
def attention(q_ptr, k_ptr, v_ptr, out_ptr, M, L, D, BLOCK_D: tl.constexpr):
    i = tl.program_id(0)
    ds = tl.arange(0, BLOCK_D)
    dmask = ds < D
    q = tl.load(q_ptr + i * D + ds, mask=dmask, other=0.0)
    mx = float('-inf')
    for j in range(0, L):
        kv = tl.load(k_ptr + j * D + ds, mask=dmask, other=0.0)
        mx = tl.maximum(mx, tl.sum(q * kv, axis=0) * 0.17677669529663687)
    denom = 0.0
    acc = tl.zeros((BLOCK_D,), dtype=tl.float32)
    for j in range(0, L):
        kv = tl.load(k_ptr + j * D + ds, mask=dmask, other=0.0)
        s = tl.sum(q * kv, axis=0) * 0.17677669529663687
        p = tl.exp(s)
        denom += p
        vv = tl.load(v_ptr + j * D + ds, mask=dmask, other=0.0)
        acc += p * vv
    tl.store(out_ptr + i * D + ds, acc / denom, mask=dmask)


every program read-modify-writes the same output cell

on real hardware the lost update is intermittent; a single allclose run sees it only sometimes

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = 1
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_sum_race(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros((BLOCK,), dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        acc += v
    prev = tl.load(out_ptr)
    tl.store(out_ptr, prev + tl.sum(acc, axis=0))


per-row statistics published through a shared scratch cell

rows of i.i.d. noise have near-identical statistics, so reading another row's mean and variance barely changes the output

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer w_ptr = N
# litmus: buffer b_ptr = N
# litmus: buffer out_ptr = M*N
# litmus: buffer stats_ptr = 2
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def layernorm_race(x_ptr, w_ptr, b_ptr, out_ptr, stats_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    s1 = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        s1 += tl.sum(v, axis=0)
    mean = s1 / N
    s2 = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        d = tl.where(mask, v - mean, 0.0)
        s2 += tl.sum(d * d, axis=0)
    var = s2 / N
    tl.store(stats_ptr + 0, mean)
    tl.store(stats_ptr + 1, var)
    mean = tl.load(stats_ptr + 0)
    var = tl.load(stats_ptr + 1)
    rstd = 1.0 / tl.sqrt(var + 1e-5)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        w = tl.load(w_ptr + cols, mask=mask, other=0.0)
        b = tl.load(b_ptr + cols, mask=mask, other=0.0)
        tl.store(out_ptr + row * N + cols, (v - mean) * rstd * w + b, mask=mask)


tl.debug_barrier() under a lane-dependent condition

divergent barriers produce no numerical error at all until the schedule changes -- they are invisible to any output comparison

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_sum_barrier(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros((BLOCK,), dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        acc += v
    if tl.arange(0, BLOCK)[0] < N:
        tl.debug_barrier()
    tl.store(out_ptr + row, tl.sum(acc, axis=0))


fp16 inputs accumulated in fp16 instead of fp32

at K=512 with values in [0,1) the fp16 accumulator has not yet stagnated; the error stays under rtol=1e-2

import triton
import triton.language as tl

# litmus: constexpr BLOCK_M = 32
# litmus: constexpr BLOCK_N = 32
# litmus: constexpr BLOCK_K = 32
# litmus: buffer a_ptr = M*K
# litmus: buffer b_ptr = K*N
# litmus: buffer c_ptr = M*N
# litmus: grid = cdiv(M, BLOCK_M), cdiv(N, BLOCK_N)

@triton.jit
def matmul(a_ptr, b_ptr, c_ptr, M, N, K,
           BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pm = tl.program_id(0)
    pn = tl.program_id(1)
    offs_m = pm * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pn * BLOCK_N + tl.arange(0, BLOCK_N)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float16)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K)
        a = tl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :],
                    mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptr + offs_k[:, None] * N + offs_n[None, :],
                    mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
    tl.store(c_ptr + offs_m[:, None] * N + offs_n[None, :], acc,
             mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


row sum accumulated in fp16

an fp16 accumulator stagnates above ~2048 and overflows above 65504; with N=512 and values in [0,1) the running total only reaches about 256

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_reduce(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros((BLOCK,), dtype=tl.float16)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * S + cols, mask=mask, other=0.0)
        acc += v
    tl.store(out_ptr + row, tl.sum(acc, axis=0))


row base computed with N instead of the row stride S

benchmarks always pass contiguous tensors, where the row stride equals the row width and the bug is a no-op

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer out_ptr = M
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def row_reduce(x_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros((BLOCK,), dtype=tl.float32)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        acc += v
    tl.store(out_ptr + row, tl.sum(acc, axis=0))


layernorm reads rows with stride N on a padded tensor

import triton
import triton.language as tl

# litmus: constexpr BLOCK = 32
# litmus: buffer x_ptr = M*S
# litmus: buffer w_ptr = N
# litmus: buffer b_ptr = N
# litmus: buffer out_ptr = M*N
# litmus: grid = M
# litmus: subst S = N + P
# litmus: sym P in [0, 8]

@triton.jit
def layernorm(x_ptr, w_ptr, b_ptr, out_ptr, M, N, S, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    s1 = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        s1 += tl.sum(v, axis=0)
    mean = s1 / N
    s2 = 0.0
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        d = tl.where(mask, v - mean, 0.0)
        s2 += tl.sum(d * d, axis=0)
    var = s2 / N
    rstd = 1.0 / tl.sqrt(var + 1e-5)
    for start in range(0, N, BLOCK):
        cols = start + tl.arange(0, BLOCK)
        mask = cols < N
        v = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        w = tl.load(w_ptr + cols, mask=mask, other=0.0)
        b = tl.load(b_ptr + cols, mask=mask, other=0.0)
        tl.store(out_ptr + row * N + cols, (v - mean) * rstd * w + b, mask=mask)