"""Fused Triton kernels shared by the non-causal engine.

_resid_norm_kernel     x[r] += y[src[r]] * scale (optional), then RMSNorm(x[r]) * w -> bf16
_resid_ln_split_kernel x[r] += y[r] * scale, LayerNorm(x[r]) * w + b, written as [hi | lo | hi] bf16 for an fp32-accurate head GEMM
_interleave            reorder SwiGLU rows [gate | up] -> gate0, up0, gate1, up1, ... (quack's gated epilogue layout)
"""
import torch
import triton
import triton.language as tl


@triton.jit
def _resid_norm_kernel(x, y, src, scale, w, hout, D: tl.constexpr, EPS: tl.constexpr, ADD: tl.constexpr):
    r = tl.program_id(0)
    c = tl.arange(0, D)
    xr = x + r.to(tl.int64) * D
    xv = tl.load(xr + c)
    if ADD:
        sr = tl.load(src + r).to(tl.int64)
        yv = tl.load(y + sr * D + c).to(tl.float32)
        xv = xv + yv * tl.load(scale + c)
        tl.store(xr + c, xv)
    rs = tl.rsqrt(tl.sum(xv * xv, 0) / D + EPS)
    tl.store(hout + r.to(tl.int64) * D + c, (xv * rs * tl.load(w + c)).to(tl.bfloat16))


@triton.jit
def _resid_ln_split_kernel(x, y, scale, w, b, aout, D: tl.constexpr, EPS: tl.constexpr):
    r = tl.program_id(0)
    c = tl.arange(0, D)
    xv = tl.load(x + r.to(tl.int64) * D + c)
    yv = tl.load(y + r.to(tl.int64) * D + c).to(tl.float32)
    xv = xv + yv * tl.load(scale + c)
    mean = tl.sum(xv, 0) / D
    xc = xv - mean
    rs = tl.rsqrt(tl.sum(xc * xc, 0) / D + EPS)
    o = xc * rs * tl.load(w + c) + tl.load(b + c)
    hi = o.to(tl.bfloat16)
    lo = (o - hi.to(tl.float32)).to(tl.bfloat16)
    ar = aout + r.to(tl.int64) * (3 * D)
    tl.store(ar + c, hi)
    tl.store(ar + D + c, lo)
    tl.store(ar + 2 * D + c, hi)


def _interleave(w):
    """rows [gate | up] -> gate0, up0, gate1, up1, ... (quack's gated-epilogue layout)."""
    I = w.shape[0] // 2
    return torch.stack([w[:I], w[I:]], 1).reshape(2 * I, *w.shape[1:]).contiguous()