| """Batch-invariant matmul for the one-pass query batch (SPEC.md §12.2). |
| |
| cuBLAS picks a GEMM kernel per shape (tile sizes, split-K), so one row's result changes |
| with the number of rows in the batch. :func:`bi_linear` is a Triton GEMM with ONE fixed |
| tile configuration and no split-K: every output element is the same K-loop over its own |
| row and the weight column, whatever ``M`` is, so a fork's numerics do not depend on what |
| else is in the batch. Rows never mix (each row's dot products read only that row). |
| |
| CPU tensors run ``F.linear`` one row at a time (same shape for every row). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import Tensor |
|
|
| try: |
| import triton |
| import triton.language as tl |
| except ImportError: |
| triton = None |
|
|
| BM, BN, BK = 64, 128, 64 |
| WARPS, STAGES = 4, 3 |
|
|
| if triton is not None: |
|
|
| @triton.jit |
| def _bi_mm(a_ptr, b_ptr, c_ptr, bias_ptr, M, N, K, |
| sam, sak, sbn, sbk, scm, scn, |
| HAS_BIAS: tl.constexpr, IEEE: tl.constexpr, |
| BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr): |
| |
| pid_m = tl.program_id(0) |
| pid_n = tl.program_id(1) |
| rm = pid_m * BM + tl.arange(0, BM) |
| rn = pid_n * BN + tl.arange(0, BN) |
| rk = tl.arange(0, BK) |
| acc = tl.zeros((BM, BN), dtype=tl.float32) |
| for k0 in range(0, K, BK): |
| kk = k0 + rk |
| a = tl.load(a_ptr + rm[:, None] * sam + kk[None, :] * sak, |
| mask=(rm[:, None] < M) & (kk[None, :] < K), other=0.0) |
| b = tl.load(b_ptr + rn[None, :] * sbn + kk[:, None] * sbk, |
| mask=(rn[None, :] < N) & (kk[:, None] < K), other=0.0) |
| if IEEE: |
| acc += tl.dot(a, b, input_precision="ieee") |
| else: |
| acc += tl.dot(a, b) |
| if HAS_BIAS: |
| acc += tl.load(bias_ptr + rn, mask=rn < N, other=0.0).to(tl.float32)[None, :] |
| tl.store(c_ptr + rm[:, None] * scm + rn[None, :] * scn, |
| acc.to(c_ptr.dtype.element_ty), mask=(rm[:, None] < M) & (rn[None, :] < N)) |
|
|
|
|
| if triton is not None: |
|
|
| @triton.jit |
| def _bi_mm_split(a_ptr, b_ptr, p_ptr, M, N, K, KS, |
| sam, sak, sbn, sbk, |
| BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr): |
| |
| pid_m = tl.program_id(0) |
| pid_n = tl.program_id(1) |
| sid = tl.program_id(2) |
| rm = pid_m * BM + tl.arange(0, BM) |
| rn = pid_n * BN + tl.arange(0, BN) |
| rk = tl.arange(0, BK) |
| acc = tl.zeros((BM, BN), dtype=tl.float32) |
| k_lo = sid * KS |
| for k0 in range(k_lo, k_lo + KS, BK): |
| kk = k0 + rk |
| kmask = (kk < K) & (kk < k_lo + KS) |
| a = tl.load(a_ptr + rm[:, None] * sam + kk[None, :] * sak, |
| mask=(rm[:, None] < M) & kmask[None, :], other=0.0) |
| b = tl.load(b_ptr + rn[None, :] * sbn + kk[:, None] * sbk, |
| mask=(rn[None, :] < N) & kmask[:, None], other=0.0) |
| acc += tl.dot(a, b) |
| tl.store(p_ptr + sid * M * N + rm[:, None] * N + rn[None, :], acc, |
| mask=(rm[:, None] < M) & (rn[None, :] < N)) |
|
|
|
|
| @triton.jit |
| def _split_sum(p_ptr, bias_ptr, o_ptr, MN, N, S: tl.constexpr, HAS_BIAS: tl.constexpr, |
| BLOCK: tl.constexpr): |
| |
| offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) |
| mask = offs < MN |
| acc = tl.load(p_ptr + offs, mask=mask, other=0.0) |
| for i in tl.static_range(1, S): |
| acc += tl.load(p_ptr + i * MN + offs, mask=mask, other=0.0) |
| if HAS_BIAS: |
| acc += tl.load(bias_ptr + offs % N, mask=mask, other=0.0).to(tl.float32) |
| tl.store(o_ptr + offs, acc.to(o_ptr.dtype.element_ty), mask=mask) |
|
|
|
|
| SBM, SBN, SBK = 16, 32, 128 |
| SM_TARGET = 256 |
|
|
|
|
| def _split(N: int, K: int) -> int: |
| """Split-K factor from the weight shape only (never from M).""" |
| split = 1 |
| tiles = -(-N // SBN) |
| while tiles * split < SM_TARGET and K // (split * 2) >= 2 * SBK: |
| split *= 2 |
| return split |
|
|
|
|
| def bi_linear_small(x: Tensor, weight: Tensor, bias: Tensor | None = None) -> Tensor: |
| """Batch-invariant GEMM tuned for a few rows (decode steps, LM head): 16-row tiles and |
| a split-K factor fixed by the weight shape; partials are summed in split order. A |
| row's result depends on (row, weight) only, for any number of rows. bf16/fp16 only.""" |
| if not x.is_cuda or triton is None: |
| return bi_linear(x, weight, bias) |
| shape = x.shape |
| a = x.reshape(-1, shape[-1]) |
| w = weight.to(x.dtype) |
| M, K = a.shape |
| N = w.shape[0] |
| split = _split(N, K) |
| ks = -(-K // split) |
| ks = -(-ks // SBK) * SBK |
| part = torch.empty(split, M, N, device=x.device, dtype=torch.float32) |
| if M and N: |
| grid = (triton.cdiv(M, SBM), triton.cdiv(N, SBN), split) |
| _bi_mm_split[grid](a, w, part, M, N, K, ks, a.stride(0), a.stride(1), |
| w.stride(0), w.stride(1), BM=SBM, BN=SBN, BK=SBK, |
| num_warps=4, num_stages=3) |
| out = torch.empty(M, N, device=x.device, dtype=x.dtype) |
| if M and N: |
| block = 1024 |
| _split_sum[(triton.cdiv(M * N, block),)](part, bias if bias is not None else part, out, |
| M * N, N, S=split, HAS_BIAS=bias is not None, |
| BLOCK=block) |
| return out.reshape(*shape[:-1], N) |
|
|
|
|
| def bi_linear(x: Tensor, weight: Tensor, bias: Tensor | None = None) -> Tensor: |
| """``F.linear`` whose per-row result is independent of the other rows (CUDA).""" |
| if not x.is_cuda or triton is None: |
| rows = x.reshape(-1, x.shape[-1]) |
| out = torch.stack([F.linear(r, weight, bias) for r in rows]) if rows.shape[0] else \ |
| rows.new_zeros(0, weight.shape[0]) |
| return out.reshape(*x.shape[:-1], weight.shape[0]) |
| shape = x.shape |
| a = x.reshape(-1, shape[-1]) |
| w = weight.to(x.dtype) |
| M, K = a.shape |
| N = w.shape[0] |
| out = torch.zeros(M, N, device=x.device, dtype=x.dtype) if not K else torch.empty( |
| M, N, device=x.device, dtype=x.dtype) |
| if M and N and K: |
| grid = (triton.cdiv(M, BM), triton.cdiv(N, BN)) |
| _bi_mm[grid](a, w, out, bias if bias is not None else out, M, N, K, |
| a.stride(0), a.stride(1), w.stride(0), w.stride(1), |
| out.stride(0), out.stride(1), |
| HAS_BIAS=bias is not None, IEEE=x.dtype == torch.float32, |
| BM=BM, BN=BN, BK=BK, num_warps=WARPS, num_stages=STAGES) |
| return out.reshape(*shape[:-1], N) |
|
|
|
|
| |
|
|
| if triton is not None: |
|
|
| @triton.jit |
| def _fork_attn(q_ptr, pk_ptr, pv_ptr, ok_ptr, ov_ptr, o_ptr, qstart_ptr, poff_ptr, plen_ptr, |
| L, CAP, scale, WIN, |
| sqb, sql, sqh, spk_h, spk_p, sok_b, sok_h, sok_p, sob, sol, soh, |
| HQ: tl.constexpr, GROUP: tl.constexpr, D: tl.constexpr, |
| BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): |
| |
| |
| |
| |
| |
| |
| |
| pid_m = tl.program_id(0) |
| bh = tl.program_id(1) |
| b = bh // HQ |
| hq = bh % HQ |
| hk = hq // GROUP |
| offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) |
| offs_d = tl.arange(0, D) |
| offs_n = tl.arange(0, BLOCK_N) |
| q0 = tl.load(qstart_ptr + b) |
| P = tl.load(plen_ptr + b) |
| po = tl.load(poff_ptr + b) |
| qmask = offs_m < L |
| q = tl.load(q_ptr + b * sqb + offs_m[:, None] * sql + hq * sqh + offs_d[None, :], |
| mask=qmask[:, None], other=0.0) |
| m_i = tl.full((BLOCK_M,), float("-inf"), tl.float32) |
| l_i = tl.zeros((BLOCK_M,), tl.float32) |
| acc = tl.zeros((BLOCK_M, D), tl.float32) |
| t = q0 + offs_m |
| t_lo = q0 + pid_m * BLOCK_M |
| |
| |
| lo_p = t + P - WIN + 1 |
| lo_o = t - WIN + 1 |
| p0 = tl.maximum(t_lo + P - WIN + 1, 0) // BLOCK_N * BLOCK_N |
| o0 = tl.maximum(t_lo - WIN + 1, 0) // BLOCK_N * BLOCK_N |
| for n0 in range(p0, P, BLOCK_N): |
| n = n0 + offs_n |
| kmask = n < P |
| k = tl.load(pk_ptr + hk * spk_h + (po + n[:, None]) * spk_p + offs_d[None, :], |
| mask=kmask[:, None], other=0.0) |
| s = tl.dot(q, tl.trans(k)) * scale |
| s = tl.where(kmask[None, :] & (n[None, :] >= lo_p[:, None]), s, float("-inf")) |
| m_new = tl.maximum(m_i, tl.max(s, 1)) |
| m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) |
| p = tl.exp(s - m_safe[:, None]) |
| alpha = tl.exp(m_i - m_safe) |
| v = tl.load(pv_ptr + hk * spk_h + (po + n[:, None]) * spk_p + offs_d[None, :], |
| mask=kmask[:, None], other=0.0) |
| acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) |
| l_i = l_i * alpha + tl.sum(p, 1) |
| m_i = m_new |
| hi = tl.minimum(q0 + (pid_m + 1) * BLOCK_M, CAP) |
| for n0 in range(o0, hi, BLOCK_N): |
| n = n0 + offs_n |
| k = tl.load(ok_ptr + b * sok_b + hk * sok_h + n[:, None] * sok_p + offs_d[None, :], |
| mask=(n < hi)[:, None], other=0.0) |
| s = tl.dot(q, tl.trans(k)) * scale |
| s = tl.where((n[None, :] <= t[:, None]) & (n[None, :] >= lo_o[:, None]), s, |
| float("-inf")) |
| m_new = tl.maximum(m_i, tl.max(s, 1)) |
| m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) |
| p = tl.exp(s - m_safe[:, None]) |
| alpha = tl.exp(m_i - m_safe) |
| v = tl.load(ov_ptr + b * sok_b + hk * sok_h + n[:, None] * sok_p + offs_d[None, :], |
| mask=(n < hi)[:, None], other=0.0) |
| acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) |
| l_i = l_i * alpha + tl.sum(p, 1) |
| m_i = m_new |
| out = acc / l_i[:, None] |
| tl.store(o_ptr + b * sob + offs_m[:, None] * sol + hq * soh + offs_d[None, :], |
| out.to(o_ptr.dtype.element_ty), mask=qmask[:, None]) |
|
|
|
|
| def fork_attention(q: Tensor, pk: Tensor, pv: Tensor, ok: Tensor, ov: Tensor, |
| qstart: Tensor, poff: Tensor, plen: Tensor, scale: float, |
| window: int = 0) -> Tensor: |
| """Attention of fork queries over a shared prefix plus each row's own causal keys. |
| |
| q [B, L, Hq, D] (query i of row b sits at own position ``qstart[b] + i``); |
| pk/pv [Hkv, total, D] the prefixes (read by reference, never copied per row): row b |
| reads positions ``poff[b] .. poff[b]+plen[b]-1`` (device tensors, so a CUDA graph can |
| replay over any prefixes); |
| ok/ov [B, Hkv, cap, D] each row's own keys (positions ``0..qstart[b]+L-1`` written). |
| ``window`` > 0 (sliding-window attention, Gemma 4): a query at own position t sees |
| only keys less than ``window`` positions back; the row's prefix slab is then the last |
| ``plen[b]`` positions right before own position 0 (a sliding layer's prefix tail). |
| Returns [B, L, Hq, D]. A query's result depends only on its own row: one program per |
| (row, head, query block), fixed block sizes (by head dim only), keys visited in a fixed |
| order (fully masked blocks are exact no-ops in the online softmax). |
| """ |
| B, L, HQ, D = q.shape |
| hkv = pk.shape[0] |
| if not q.is_cuda or triton is None: |
| return _fork_attention_ref(q, pk, pv, ok, ov, qstart, poff, plen, scale, window) |
| assert pk.stride(-1) == 1 and ok.stride(-1) == 1 and q.stride(-1) == 1 |
| assert pv.stride() == pk.stride() and ov.stride() == ok.stride() |
| out = torch.empty(B, L, HQ, D, device=q.device, dtype=q.dtype) |
| |
| |
| BLOCK_M, BLOCK_N = 16, (32 if D <= 256 else 16) |
| grid = (triton.cdiv(L, BLOCK_M), B * HQ) |
| _fork_attn[grid](q, pk, pv, ok, ov, out, qstart.to(torch.int32), poff.to(torch.int32), |
| plen.to(torch.int32), L, ok.shape[2], |
| scale, int(window) if window > 0 else 1 << 30, |
| q.stride(0), q.stride(1), q.stride(2), pk.stride(0), pk.stride(1), |
| ok.stride(0), ok.stride(1), ok.stride(2), |
| out.stride(0), out.stride(1), out.stride(2), |
| HQ=HQ, GROUP=HQ // hkv, D=D, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, |
| num_warps=4, num_stages=2) |
| return out |
|
|
|
|
| |
|
|
| if triton is not None: |
|
|
| @triton.jit |
| def _fork_gdr(q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr, o_ptr, h0_ptr, h0_idx_ptr, ht_ptr, T, |
| scale, |
| H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, |
| BV: tl.constexpr, STORE: tl.constexpr): |
| |
| |
| |
| |
| i_v, i_nh = tl.program_id(0), tl.program_id(1) |
| i_n, i_hv = i_nh // HV, i_nh % HV |
| i_h = i_hv // (HV // H) |
| o_k = tl.arange(0, K) |
| o_v = i_v * BV + tl.arange(0, BV) |
| p_q = q_ptr + (i_n * T * H + i_h) * K + o_k |
| p_k = k_ptr + (i_n * T * H + i_h) * K + o_k |
| p_v = v_ptr + (i_n * T * HV + i_hv) * V + o_v |
| p_g = g_ptr + i_n * T * HV + i_hv |
| p_b = beta_ptr + i_n * T * HV + i_hv |
| p_o = o_ptr + (i_n * T * HV + i_hv) * V + o_v |
| i_h0 = tl.load(h0_idx_ptr + i_n).to(tl.int64) |
| b_h = tl.load(h0_ptr + i_h0 * HV * K * V + i_hv * K * V + o_k[:, None] * V |
| + o_v[None, :]).to(tl.float32) |
| for _ in range(0, T): |
| b_q = tl.load(p_q).to(tl.float32) |
| b_k = tl.load(p_k).to(tl.float32) |
| b_v = tl.load(p_v).to(tl.float32) |
| b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) |
| b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) |
| b_q = b_q * scale |
| b_beta = tl.load(p_b).to(tl.float32) |
| b_g = tl.load(p_g).to(tl.float32) |
| b_h *= tl.exp(b_g) |
| b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0)) |
| b_h += b_k[:, None] * b_v[None, :] |
| b_o = tl.sum(b_h * b_q[:, None], 0) |
| tl.store(p_o, b_o.to(p_o.dtype.element_ty)) |
| p_q += H * K |
| p_k += H * K |
| p_v += HV * V |
| p_g += HV |
| p_b += HV |
| p_o += HV * V |
| if STORE: |
| tl.store(ht_ptr + i_nh * K * V + o_k[:, None] * V + o_v[None, :], b_h) |
|
|
|
|
| def fork_delta_rule(q: Tensor, k: Tensor, v: Tensor, g: Tensor, beta: Tensor, h0: Tensor, |
| h0_idx: Tensor, store: Tensor | None) -> Tensor: |
| """q/k [B,T,H,K], v [B,T,HV,V], g/beta [B,T,HV], h0 [N,HV,K,V] fp32 contiguous, row b |
| starting from ``h0[h0_idx[b]]``. ``store`` [B,HV,K,V] receives the final state (may |
| alias ``h0`` when ``h0_idx`` is the identity). Returns o [B,T,HV,V].""" |
| B, T, H, K = k.shape |
| HV, V = v.shape[2], v.shape[3] |
| if not q.is_cuda or triton is None: |
| return _fork_delta_rule_ref(q, k, v, g, beta, h0, h0_idx, store) |
| q, k, v = q.contiguous(), k.contiguous(), v.contiguous() |
| g, beta = g.float().contiguous(), beta.contiguous() |
| o = torch.empty_like(v) |
| BV = 8 |
| grid = (triton.cdiv(V, BV), B * HV) |
| _fork_gdr[grid](q, k, v, g, beta, o, h0, h0_idx, store if store is not None else h0, T, |
| K ** -0.5, |
| H=H, HV=HV, K=K, V=V, BV=BV, STORE=store is not None, |
| num_warps=1, num_stages=3) |
| return o |
|
|
|
|
| |
|
|
|
|
| def _fork_attention_ref(q, pk, pv, ok, ov, qstart, poff, plen, scale, window=0): |
| """One query at a time over exactly its visible keys (prefix + own ``0..t``; with a |
| ``window``, only the last ``window`` positions), so the softmax length never depends |
| on the batch width.""" |
| B, L, HQ, D = q.shape |
| group = HQ // pk.shape[0] |
| out = torch.empty_like(q) |
| for b in range(B): |
| t0 = int(qstart[b]) |
| po, P = int(poff[b]), int(plen[b]) |
| for i in range(L): |
| n = t0 + i + 1 |
| p_lo = max(0, P + n - window) if window > 0 else 0 |
| o_lo = max(0, n - window) if window > 0 else 0 |
| keys = torch.cat([pk[:, po + p_lo:po + P], ok[b, :, o_lo:n]], 1) |
| vals = torch.cat([pv[:, po + p_lo:po + P], ov[b, :, o_lo:n]], 1) |
| keys = keys.repeat_interleave(group, 0).float() |
| vals = vals.repeat_interleave(group, 0).float() |
| s = torch.einsum("hd,hnd->hn", q[b, i].float(), keys) * scale |
| out[b, i] = torch.einsum("hn,hnd->hd", s.softmax(-1), vals).to(q.dtype) |
| return out |
|
|
|
|
| def _fork_delta_rule_ref(q, k, v, g, beta, h0, h0_idx, store): |
| from deg.model import delta_rule_reference |
|
|
| rep = v.shape[2] // k.shape[2] |
| if rep > 1: |
| q, k = q.repeat_interleave(rep, 2), k.repeat_interleave(rep, 2) |
| init = h0.index_select(0, h0_idx) |
| o, final = delta_rule_reference(q, k, v, g, beta, init) |
| if store is not None: |
| store.copy_(final) |
| return o |
|
|
|
|
| |
|
|
| if triton is not None: |
|
|
| @triton.jit |
| def _rms_row(x_ptr, w_ptr, o_ptr, C, eps, BLOCK: tl.constexpr): |
| |
| row = tl.program_id(0).to(tl.int64) |
| offs = tl.arange(0, BLOCK) |
| mask = offs < C |
| x = tl.load(x_ptr + row * C + offs, mask=mask, other=0.0).to(tl.float32) |
| w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) |
| r = 1.0 / tl.sqrt(tl.sum(x * x, 0) / C + eps) |
| tl.store(o_ptr + row * C + offs, (x * r * (1.0 + w)).to(o_ptr.dtype.element_ty), |
| mask=mask) |
|
|
| @triton.jit |
| def _gated_rms_row(x_ptr, g_ptr, w_ptr, o_ptr, C, eps, BLOCK: tl.constexpr): |
| |
| row = tl.program_id(0).to(tl.int64) |
| offs = tl.arange(0, BLOCK) |
| mask = offs < C |
| x = tl.load(x_ptr + row * C + offs, mask=mask, other=0.0).to(tl.float32) |
| g = tl.load(g_ptr + row * C + offs, mask=mask, other=0.0).to(tl.float32) |
| w = tl.load(w_ptr + offs, mask=mask, other=0.0) |
| r = 1.0 / tl.sqrt(tl.sum(x * x, 0) / C + eps) |
| xn = (x * r).to(o_ptr.dtype.element_ty) |
| y = (w * xn).to(tl.float32) * (g / (1.0 + tl.exp(-g))) |
| tl.store(o_ptr + row * C + offs, y.to(o_ptr.dtype.element_ty), mask=mask) |
|
|
|
|
| if triton is not None: |
|
|
| @triton.jit |
| def _rms_row_w(x_ptr, w_ptr, o_ptr, C, eps, BLOCK: tl.constexpr): |
| |
| row = tl.program_id(0).to(tl.int64) |
| offs = tl.arange(0, BLOCK) |
| mask = offs < C |
| x = tl.load(x_ptr + row * C + offs, mask=mask, other=0.0).to(tl.float32) |
| w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) |
| r = 1.0 / tl.sqrt(tl.sum(x * x, 0) / C + eps) |
| xn = (x * r).to(o_ptr.dtype.element_ty).to(tl.float32) |
| tl.store(o_ptr + row * C + offs, (w * xn).to(o_ptr.dtype.element_ty), mask=mask) |
|
|
|
|
| if triton is not None: |
|
|
| @triton.jit |
| def _rms_row_g(x_ptr, w_ptr, o_ptr, C, eps, HAS_W: tl.constexpr, BLOCK: tl.constexpr): |
| |
| |
| row = tl.program_id(0).to(tl.int64) |
| offs = tl.arange(0, BLOCK) |
| mask = offs < C |
| x = tl.load(x_ptr + row * C + offs, mask=mask, other=0.0).to(tl.float32) |
| r = 1.0 / tl.sqrt(tl.sum(x * x, 0) / C + eps) |
| y = x * r |
| if HAS_W: |
| y = y * tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) |
| tl.store(o_ptr + row * C + offs, y.to(o_ptr.dtype.element_ty), mask=mask) |
|
|
|
|
| def rms_rows_g(x: Tensor, w: Tensor | None, eps: float) -> Tensor: |
| """Row-invariant Gemma 4 RMSNorm (see :func:`deg.model._rms_g`; ``w`` None: no |
| scale, :func:`deg.model._rms_n`).""" |
| xc = x.contiguous() |
| out = torch.empty_like(xc) |
| C = x.shape[-1] |
| rows = xc.numel() // C |
| if rows: |
| _rms_row_g[(rows,)](xc, w.contiguous() if w is not None else xc, out, C, eps, |
| HAS_W=w is not None, BLOCK=triton.next_power_of_2(C), |
| num_warps=4) |
| return out |
|
|
|
|
| def rms_rows_w(x: Tensor, w: Tensor, eps: float) -> Tensor: |
| """Row-invariant LFM2 RMSNorm (``w * norm(x)``, see :func:`deg.model._rms_w`).""" |
| xc = x.contiguous() |
| out = torch.empty_like(xc) |
| C = x.shape[-1] |
| rows = xc.numel() // C |
| if rows: |
| _rms_row_w[(rows,)](xc, w.contiguous(), out, C, eps, |
| BLOCK=triton.next_power_of_2(C), num_warps=4) |
| return out |
|
|
|
|
| def rms_rows(x: Tensor, w: Tensor, eps: float) -> Tensor: |
| """Row-invariant RMSNorm over the last dim (one program per row, fixed reduction).""" |
| xc = x.contiguous() |
| out = torch.empty_like(xc) |
| C = x.shape[-1] |
| rows = xc.numel() // C |
| if rows: |
| _rms_row[(rows,)](xc, w.contiguous(), out, C, eps, |
| BLOCK=triton.next_power_of_2(C), num_warps=4) |
| return out |
|
|
|
|
| def gated_rms_rows(x: Tensor, gate: Tensor, w: Tensor, eps: float) -> Tensor: |
| """Row-invariant gated RMSNorm (see :func:`deg.model._gated_rms`).""" |
| xc, gc = x.contiguous(), gate.contiguous() |
| out = torch.empty_like(xc) |
| C = x.shape[-1] |
| rows = xc.numel() // C |
| if rows: |
| _gated_rms_row[(rows,)](xc, gc, w.contiguous(), out, C, eps, |
| BLOCK=triton.next_power_of_2(C), num_warps=4) |
| return out |
|
|
