if BACKEND == "triton":
@triton.jit
def _vadd_kernel(a_ptr, b_ptr, c_ptr, n, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
mask = offs < n
a = tl.load(a_ptr + offs, mask=mask)
b = tl.load(b_ptr + offs, mask=mask)
tl.store(c_ptr + offs, a + b, mask=mask)
@triton.jit
def _fused_gelu_kernel(x_ptr, w_ptr, b_ptr, o_ptr, n, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
mask = offs < n
x = tl.load(x_ptr + offs, mask=mask)
w = tl.load(w_ptr + offs, mask=mask)
b = tl.load(b_ptr + offs, mask=mask)
h = x * w + b
c = 0.7978845608028654
z = c * (h + 0.044715 * h * h * h)
e = tl.exp(-2.0 * z)
tanh = (1.0 - e) / (1.0 + e)
g = 0.5 * h * (1.0 + tanh)
tl.store(o_ptr + offs, g, mask=mask)
@triton.jit
def _softmax_kernel(x_ptr, o_ptr, stride, n_cols, BLOCK: tl.constexpr):
row = tl.program_id(0)
cols = tl.arange(0, BLOCK)
mask = cols < n_cols
ptr = x_ptr + row * stride + cols
x = tl.load(ptr, mask=mask, other=-float("inf"))
x = x - tl.max(x, axis=0)
num = tl.exp(x)
den = tl.sum(num, axis=0)
tl.store(o_ptr + row * stride + cols, num / den, mask=mask)
@triton.jit
def _matmul_kernel(A, B, C, M, N, K,
sam, sak, sbk, sbn, scm, scn,
BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
pid_m = tl.program_id(0)
pid_n = tl.program_id(1)
offs_m = pid_m * BM + tl.arange(0, BM)
offs_n = pid_n * BN + tl.arange(0, BN)
offs_k = tl.arange(0, BK)
a_ptr = A + offs_m[:, None] * sam + offs_k[None, :] * sak
b_ptr = B + offs_k[:, None] * sbk + offs_n[None, :] * sbn
acc = tl.zeros((BM, BN), dtype=tl.float32)
for k in range(0, K, BK):
a = tl.load(a_ptr, mask=offs_k[None, :] < K - k, other=0.0)
b = tl.load(b_ptr, mask=offs_k[:, None] < K - k, other=0.0)
acc += tl.dot(a, b)
a_ptr += BK * sak
b_ptr += BK * sbk
c_ptr = C + offs_m[:, None] * scm + offs_n[None, :] * scn
cmask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
tl.store(c_ptr, acc.to(C.dtype.element_ty), mask=cmask)
@triton.jit
def _flash_kernel(Q, K, V, O, sqz, skz, svz, soz,
L, D, scale,
BL: tl.constexpr, BD: tl.constexpr):
pid_l = tl.program_id(0)
z = tl.program_id(1)
offs_l = pid_l * BL + tl.arange(0, BL)
offs_d = tl.arange(0, BD)
q_ptr = Q + z * sqz + offs_l[:, None] * D + offs_d[None, :]
q = tl.load(q_ptr, mask=offs_l[:, None] < L, other=0.0)
m_i = tl.full((BL,), -float("inf"), dtype=tl.float32)
l_i = tl.zeros((BL,), dtype=tl.float32)
acc = tl.zeros((BL, BD), dtype=tl.float32)
for start in range(0, L, BL):
offs_k = start + tl.arange(0, BL)
k_ptr = K + z * skz + offs_k[:, None] * D + offs_d[None, :]
v_ptr = V + z * svz + offs_k[:, None] * D + offs_d[None, :]
k = tl.load(k_ptr, mask=offs_k[:, None] < L, other=0.0)
v = tl.load(v_ptr, mask=offs_k[:, None] < L, other=0.0)
s = tl.dot(q, tl.trans(k)) * scale
s = tl.where(offs_k[None, :] < L, s, -float("inf"))
m_ij = tl.maximum(m_i, tl.max(s, axis=1))
p = tl.exp(s - m_ij[:, None])
alpha = tl.exp(m_i - m_ij)
l_i = l_i * alpha + tl.sum(p, axis=1)
acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v)
m_i = m_ij
acc = acc / l_i[:, None]
o_ptr = O + z * soz + offs_l[:, None] * D + offs_d[None, :]
tl.store(o_ptr, acc.to(O.dtype.element_ty), mask=offs_l[:, None] < L)
def run_vadd(a, b):
c = torch.empty_like(a); n = a.numel()
grid = (triton.cdiv(n, 1024),)
_vadd_kernel[grid](a, b, c, n, BLOCK=1024)
return c
def run_fused_gelu(x, w, b):
o = torch.empty_like(x); n = x.numel()
grid = (triton.cdiv(n, 1024),)
_fused_gelu_kernel[grid](x, w, b, o, n, BLOCK=1024)
return o
def run_softmax(x):
m, ncols = x.shape
o = torch.empty_like(x)
BLOCK = triton.next_power_of_2(ncols)
_softmax_kernel[(m,)](x, o, x.stride(0), ncols, BLOCK=BLOCK)
return o
def run_matmul(a, b):
M, K = a.shape; K2, N = b.shape
c = torch.empty((M, N), device=a.device, dtype=a.dtype)
BM = BN = 64; BK = 32
grid = (triton.cdiv(M, BM), triton.cdiv(N, BN))
_matmul_kernel[grid](a, b, c, M, N, K,
a.stride(0), a.stride(1), b.stride(0), b.stride(1),
c.stride(0), c.stride(1), BM=BM, BN=BN, BK=BK)
return c
def run_flash(q, k, v):
Z, L, D = q.shape
o = torch.empty_like(q)
scale = 1.0 / math.sqrt(D)
BL = 64
grid = (triton.cdiv(L, BL), Z)
_flash_kernel[grid](q, k, v, o,
q.stride(0), k.stride(0), v.stride(0), o.stride(0),
L, D, scale, BL=BL, BD=D)
return o
else:
def run_vadd(a, b): return a + b
def run_fused_gelu(x, w, b): return torch.nn.functional.gelu(x * w + b, approximate="tanh")
def run_softmax(x): return torch.softmax(x, dim=-1)
def run_matmul(a, b): return a @ b
def run_flash(q, k, v): return torch.nn.functional.scaled_dot_product_attention(q, k, v)
