"""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