Blame for source/src/deg/kernels.py

"""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:  # pragma: no cover - CUDA only
    import triton
    import triton.language as tl
except ImportError:  # pragma: no cover
    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):
        # C[M, N] = A[M, K] @ B[N, K]^T (+ bias); fp32 accumulation, fixed K order.
        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):
        # partial[s] = A[:, s*KS:(s+1)*KS] @ B[:, s*KS:(s+1)*KS]^T in fp32 (fixed split).
        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):
        # out = partial[0] + partial[1] + ... (fixed order) + bias, cast to the out dtype.
        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  # programs the small-M split aims for (L4: 58 SMs; swept on the decode step)


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:  # CPU reference: one row per call, same shape
        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)


# ------------------------------------------------------------------ fork attention

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):
        # One (row, query head, query block): online softmax over the SHARED prefix keys
        # (read by reference, same for every row) then the row's own keys, causal.
        # WIN (sliding window; 2^30 = none): a query at own position t sees only the keys
        # less than WIN positions back; the prefix holds the P positions right before own
        # position 0.
        # Blocks wholly outside every query's window are skipped (exact no-ops), and a
        # block with no visible key leaves the running softmax unchanged (m_safe).
        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  # own position of each query
        t_lo = q0 + pid_m * BLOCK_M  # smallest own position in this query block
        # window lower bounds: prefix key n is visible iff n >= t + P - WIN + 1, own key n
        # iff n >= t - WIN + 1 (no window: WIN = 2^30, every bound is negative)
        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)
    # D = 512 (Gemma 4 global layers): a [16, 512] fp32 accumulator plus a K/V tile per
    # program; 16-key tiles keep it in registers/shared memory (fixed by D, never by B)
    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


# ------------------------------------------------------- fork gated delta rule (recurrent)

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):
        # FLA's fused recurrent gated delta rule (qk l2norm in kernel, head-wise beta/g),
        # with row n's initial state read at ``h0[h0_idx[n]]`` (a shared prefix state is
        # read by reference by every row forking it). One program per (v block, row, v head),
        # tokens in order: a row's result never depends on other rows.
        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


# ------------------------------------------------------------------ CPU references


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


# ---------------------------------------------------------------- row RMS norms (fixed)

if triton is not None:

    @triton.jit
    def _rms_row(x_ptr, w_ptr, o_ptr, C, eps, BLOCK: tl.constexpr):
        # Qwen3.5 RMSNorm, one program per row: x * rsqrt(mean(x^2) + eps) * (1 + w).
        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):
        # Gated DeltaNet output norm: (w * norm(x).to(dtype)) * silu(gate) -> dtype.
        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):
        # LFM2 RMSNorm, one program per row: w * (x * rsqrt(mean(x^2) + eps)).to(dtype).
        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):
        # Gemma 4 RMSNorm, one program per row: (x * rsqrt(mean(x^2) + eps) * w) in fp32,
        # then cast (HF ``Gemma4RMSNorm``; HAS_W False = ``with_scale=False``, v_norm).
        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