"""TileLang kernels for the Laya (ModernBERT + decision head) encoder. All kernels take bf16 activations, accumulate in fp32. Row count M is a runtime symbol so one compiled kernel serves every batch/sequence bucket; M must be a multiple of 16 (the caller pads); out-of-bounds rows are predicated by TileLang. """ import tilelang import tilelang.language as T DT, ACC = "bfloat16", "float" FAST = {tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True} def _act(x, kind): if kind == "gelu": # exact erf-GELU, what HF "gelu" means return 0.5 * x * (1.0 + T.erf(x * 0.7071067811865476)) if kind == "relu": return T.max(x, 0.0) return x # ----------------------------------------------------------------------------- GEMM @tilelang.jit(pass_configs=FAST) def gemm_kernel(N, K, bias=False, act="none", bm=64, bn=128, bk=64, stages=3, threads=128): """C[M,N] = act(A[M,K] @ W[N,K]^T + b).""" M = T.dynamic("M") @T.prim_func def main(A: T.Tensor((M, K), DT), W: T.Tensor((N, K), DT), Bv: T.Tensor((N,), ACC), C: T.Tensor((M, N), DT)): with T.Kernel(T.ceildiv(N, bn), T.ceildiv(M, bm), threads=threads) as (bx, by): A_s = T.alloc_shared((bm, bk), DT) W_s = T.alloc_shared((bn, bk), DT) C_l = T.alloc_fragment((bm, bn), ACC) T.clear(C_l) for k in T.Pipelined(T.ceildiv(K, bk), num_stages=stages): T.copy(A[by * bm, k * bk], A_s) T.copy(W[bx * bn, k * bk], W_s) T.gemm(A_s, W_s, C_l, transpose_B=True) for i, j in T.Parallel(bm, bn): v = C_l[i, j] if bias: v = v + Bv[bx * bn + j] C_l[i, j] = _act(v, act) T.copy(C_l, C[by * bm, bx * bn]) return main @tilelang.jit(pass_configs=FAST) def gemm_geglu_kernel(F, K, bm=64, bn=64, bk=64, stages=3, threads=128): """ModernBERT GLU MLP up-projection, fused: C[M,F] = gelu(A @ Wi[:F]^T) * (A @ Wi[F:]^T).""" M = T.dynamic("M") @T.prim_func def main(A: T.Tensor((M, K), DT), W: T.Tensor((2 * F, K), DT), C: T.Tensor((M, F), DT)): with T.Kernel(T.ceildiv(F, bn), T.ceildiv(M, bm), threads=threads) as (bx, by): A_s = T.alloc_shared((bm, bk), DT) Wi_s = T.alloc_shared((bn, bk), DT) Wg_s = T.alloc_shared((bn, bk), DT) Ci = T.alloc_fragment((bm, bn), ACC) Cg = T.alloc_fragment((bm, bn), ACC) T.clear(Ci); T.clear(Cg) for k in T.Pipelined(T.ceildiv(K, bk), num_stages=stages): T.copy(A[by * bm, k * bk], A_s) T.copy(W[bx * bn, k * bk], Wi_s) T.copy(W[F + bx * bn, k * bk], Wg_s) T.gemm(A_s, Wi_s, Ci, transpose_B=True) T.gemm(A_s, Wg_s, Cg, transpose_B=True) for i, j in T.Parallel(bm, bn): Ci[i, j] = _act(Ci[i, j], "gelu") * Cg[i, j] T.copy(Ci, C[by * bm, bx * bn]) return main # ----------------------------------------------------------------------------- LayerNorm (+residual) @tilelang.jit(pass_configs=FAST) def add_ln_kernel(D, residual=True, bias=False, eps=1e-5, bm=4, threads=32): """X (fp32 residual stream) += R (bf16 branch output, if residual); Y (bf16) = LN(X) * w (+ b). The residual stream stays in fp32 exactly like the stock autocast path: ModernBERT-large's residual activations reach ~3e4, where bf16's 8-bit mantissa would lose ~100 units per add and drift layer by layer.""" M = T.dynamic("M") @T.prim_func def main(X: T.Tensor((M, D), ACC), R: T.Tensor((M, D), DT), Wv: T.Tensor((D,), ACC), Bv: T.Tensor((D,), ACC), Y: T.Tensor((M, D), DT)): with T.Kernel(T.ceildiv(M, bm), threads=threads) as bx: x = T.alloc_fragment((bm, D), ACC) xs = T.alloc_fragment((bm, D), ACC) mean = T.alloc_fragment((bm,), ACC) var = T.alloc_fragment((bm,), ACC) Xb = T.alloc_shared((bm, D), ACC) Rb = T.alloc_shared((bm, D), DT) Yb = T.alloc_shared((bm, D), DT) T.copy(X[bx * bm, 0], Xb) T.copy(Xb, x) if residual: T.copy(R[bx * bm, 0], Rb) T.copy(Rb, xs) for i, j in T.Parallel(bm, D): x[i, j] = x[i, j] + xs[i, j] T.copy(x, Xb) T.copy(Xb, X[bx * bm, 0]) T.reduce_sum(x, mean, dim=1) for i in T.Parallel(bm): mean[i] = mean[i] / D for i, j in T.Parallel(bm, D): xs[i, j] = (x[i, j] - mean[i]) * (x[i, j] - mean[i]) T.reduce_sum(xs, var, dim=1) for i in T.Parallel(bm): var[i] = T.rsqrt(var[i] / D + eps) for i, j in T.Parallel(bm, D): v = (x[i, j] - mean[i]) * var[i] * Wv[j] if bias: v = v + Bv[j] xs[i, j] = v T.copy(xs, Yb) T.copy(Yb, Y[bx * bm, 0]) return main # ----------------------------------------------------------------------------- RoPE (in place on packed qkv) @tilelang.jit(pass_configs=FAST) def rope_kernel(H, Dh, bm=32, threads=128): """QKV[M, 3*H*Dh] packed as (q|k|v)(h)(d). Rotates q and k in place (rotate-half convention, fp32 math). cos/sin: [L, Dh/2]. Row r has position r % L. M and L are runtime symbols.""" M, L = T.dynamic("M"), T.dynamic("L") half = Dh // 2 W = 2 * H * Dh # q and k columns @T.prim_func def main(QKV: T.Tensor((M, 3 * H * Dh), DT), Cos: T.Tensor((L, half), ACC), Sin: T.Tensor((L, half), ACC)): with T.Kernel(T.ceildiv(M, bm), threads=threads) as bx: for i, c in T.Parallel(bm, W // 2): r = bx * bm + i pos = r % L hh = c // half # which (q|k, head) d = c % half c0 = hh * Dh + d c1 = c0 + half x0 = T.cast(QKV[r, c0], ACC) x1 = T.cast(QKV[r, c1], ACC) cs = Cos[pos, d] sn = Sin[pos, d] QKV[r, c0] = T.cast(x0 * cs - x1 * sn, DT) QKV[r, c1] = T.cast(x1 * cs + x0 * sn, DT) return main # ----------------------------------------------------------------------------- flash attention (padding mask + sliding window) @tilelang.jit(pass_configs=FAST) def attn_kernel(B, L, H, Dh, window=0, bm=64, bn=64, stages=1, threads=128): """QKV: [B, L, 3, H, Dh] bf16 (a view of the packed [M, 3*H*Dh] buffer). Lens: [B] int32 valid length. O: [B, L, H*Dh]. window>0 => bidirectional sliding window |i-j| <= window. Masked scores use a large finite negative so fully-masked (padding) rows stay finite. B and/or L may be None: they then become runtime symbols (one compile serves every shape, at the cost of predicated loads -- ~4x slower for full attention at L=1024, free for short inputs).""" scale = (1.0 / Dh) ** 0.5 * 1.44269504 # log2(e) if B is None: B = T.dynamic("B") if L is None: L = T.dynamic("L") NEG = -1e9 @T.prim_func def main(QKV: T.Tensor((B, L, 3, H, Dh), DT), Lens: T.Tensor((B,), "int32"), O: T.Tensor((B, L, H * Dh), DT)): with T.Kernel(T.ceildiv(L, bm), H, B, threads=threads) as (bx, by, bz): Q_s = T.alloc_shared((bm, Dh), DT) K_s = T.alloc_shared((bn, Dh), DT) V_s = T.alloc_shared((bn, Dh), DT) O_s = T.alloc_shared((bm, Dh), DT) s = T.alloc_fragment((bm, bn), ACC) s_c = T.alloc_fragment((bm, bn), DT) o = T.alloc_fragment((bm, Dh), ACC) m = T.alloc_fragment((bm,), ACC) m_prev = T.alloc_fragment((bm,), ACC) sc = T.alloc_fragment((bm,), ACC) rs = T.alloc_fragment((bm,), ACC) l = T.alloc_fragment((bm,), ACC) T.annotate_layout({Q_s: tilelang.layout.make_swizzled_layout(Q_s)}) T.copy(QKV[bz, bx * bm:(bx + 1) * bm, 0, by, :], Q_s) T.fill(o, 0); T.fill(l, 0); T.fill(m, NEG) n = Lens[bz] if window > 0: k_lo = T.max(0, (bx * bm - window) // bn) k_hi = T.min(T.ceildiv(L, bn), T.ceildiv(T.min(n, (bx + 1) * bm + window), bn)) else: k_lo = 0 k_hi = T.ceildiv(n, bn) for k in T.Pipelined(k_lo, k_hi, num_stages=stages): T.copy(QKV[bz, k * bn:(k + 1) * bn, 1, by, :], K_s) for i, j in T.Parallel(bm, bn): qi = bx * bm + i kj = k * bn + j if window > 0: ok = (kj < n) & (qi - kj <= window) & (kj - qi <= window) else: ok = kj < n s[i, j] = T.if_then_else(ok, 0.0, NEG) T.gemm(Q_s, K_s, s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) T.copy(QKV[bz, k * bn:(k + 1) * bn, 2, by, :], V_s) T.copy(m, m_prev) T.reduce_max(s, m, dim=1, clear=False) for i in T.Parallel(bm): sc[i] = T.exp2(m_prev[i] * scale - m[i] * scale) for i, j in T.Parallel(bm, bn): s[i, j] = T.exp2(s[i, j] * scale - m[i] * scale) T.reduce_sum(s, rs, dim=1) for i in T.Parallel(bm): l[i] = l[i] * sc[i] + rs[i] T.copy(s, s_c) for i, j in T.Parallel(bm, Dh): o[i, j] = o[i, j] * sc[i] T.gemm(s_c, V_s, o, policy=T.GemmWarpPolicy.FullRow) for i, j in T.Parallel(bm, Dh): o[i, j] = o[i, j] / T.max(l[i], 1e-30) T.copy(o, O_s) T.copy(O_s, O[bz, bx * bm:(bx + 1) * bm, by * Dh:(by + 1) * Dh]) return main