"""Fused W4A16 Kimi-Linear hybrid decode (batch=1).

Beats baseline.py by fusing the int4 unpack + per-group asymmetric dequant
directly into the GEMV (the weights are streamed once as int4, never
materialized to bf16), plus the MLA absorb reformulation so the latent KV cache
is never re-expanded through kv_b.
"""
from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
import triton
import triton.language as tl

EPS = 1.0e-6
GROUP = 128


# --------------------------------------------------------------------------- #
# Fused int4 dequant GEMV (batch=1), split-K, optional expert gather.
# y[e,n] = sum_k x[e_or_0,k] * (unpack(wq)[e,k,n] - zeros[e,k//G,n]) * scales[..]
# Using the group identity:
#   y = sum_g scales_g * ( sum_{k in g} x_k*w_k  -  zeros_g * sum_{k in g} x_k )
# --------------------------------------------------------------------------- #
def _gemv_configs():
    cfgs = []
    for BN in (32, 64, 128):
        for SK in (4, 8, 16):
            for nw in (1, 2):
                for ns in (2, 3, 4):
                    cfgs.append(triton.Config(
                        {"BLOCK_N": BN, "SPLIT_K": SK}, num_warps=nw, num_stages=ns))
    return cfgs


@triton.autotune(configs=_gemv_configs(), key=["N", "NG", "Eact"], reset_to_zero=["y_ptr"])
@triton.jit
def _gemv_kernel(
    x_ptr, wq_ptr, s_ptr, z_ptr, idx_ptr, y_ptr,
    K, N, NG, Eact,
    stride_xe, stride_xk,
    stride_we, stride_wk, stride_wn,
    stride_se, stride_sg, stride_sn,
    stride_ye, stride_yn,
    GROUP: tl.constexpr, USE_IDX: tl.constexpr,
    BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr,
):
    pid_n = tl.program_id(0)
    pid_e = tl.program_id(1)
    pid_k = tl.program_id(2)
    if USE_IDX:
        e = tl.load(idx_ptr + pid_e).to(tl.int64)
    else:
        e = pid_e
    n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    nmask = n < N
    wq_e = wq_ptr + e * stride_we
    s_e = s_ptr + e * stride_se
    z_e = z_ptr + e * stride_se
    x_e = x_ptr + pid_e * stride_xe
    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
    HG: tl.constexpr = GROUP // 2
    rows = tl.arange(0, HG)
    for g in range(pid_k, NG, SPLIT_K):
        kp = g * HG + rows
        wq = tl.load(wq_e + kp[:, None] * stride_wk + n[None, :] * stride_wn,
                     mask=nmask[None, :], other=0)
        lo = (wq & 0xF).to(tl.float32)
        hi = ((wq >> 4) & 0xF).to(tl.float32)
        xe_ = tl.load(x_e + (2 * kp) * stride_xk)
        xo_ = tl.load(x_e + (2 * kp + 1) * stride_xk)
        raw = tl.sum(xe_[:, None] * lo + xo_[:, None] * hi, axis=0)
        xs = tl.sum(xe_ + xo_)
        s = tl.load(s_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
        z = tl.load(z_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
        acc += s * (raw - z * xs)
    if SPLIT_K == 1:
        tl.store(y_ptr + pid_e * stride_ye + n * stride_yn, acc, mask=nmask)
    else:
        tl.atomic_add(y_ptr + pid_e * stride_ye + n * stride_yn, acc, mask=nmask)


def gemv(x, wq, scales, zeros, idx=None, out=None):
    """x:[K] (shared) or [Eact,K]; wq:[E,K//2,N] or [K//2,N]; -> [Eact,N] fp32."""
    if wq.dim() == 2:
        wq = wq.unsqueeze(0)
        scales = scales.unsqueeze(0)
        zeros = zeros.unsqueeze(0)
    E, Kp, N = wq.shape
    K = Kp * 2
    NG = scales.shape[1]
    use_idx = idx is not None
    Eact = idx.numel() if use_idx else E
    if x.dim() == 1:
        stride_xe, stride_xk = 0, x.stride(0)
    else:
        stride_xe, stride_xk = x.stride(0), x.stride(1)
    if out is None:
        out = torch.empty((Eact, N), device=x.device, dtype=torch.float32)
    out.zero_()
    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), Eact, meta["SPLIT_K"])
    _gemv_kernel[grid](
        x, wq, scales, zeros, idx if use_idx else x, out,
        K, N, NG, Eact,
        stride_xe, stride_xk,
        wq.stride(0), wq.stride(1), wq.stride(2),
        scales.stride(0), scales.stride(1), scales.stride(2),
        out.stride(0), out.stride(1),
        GROUP=GROUP, USE_IDX=use_idx,
    )
    return out


@triton.autotune(configs=_gemv_configs(), key=["N", "NG", "Eact"], reset_to_zero=["y_ptr"])
@triton.jit
def _gemv_acc_kernel(
    x_ptr, wq_ptr, s_ptr, z_ptr, idx_ptr, sc_ptr, y_ptr,
    K, N, NG, Eact,
    stride_xe, stride_xk,
    stride_we, stride_wk, stride_wn,
    stride_se, stride_sg, stride_sn,
    GROUP: tl.constexpr, BLOCK_N: tl.constexpr, SPLIT_K: tl.constexpr,
):
    """Weighted accumulate GEMV: y[n] += scale[e] * sum_k x[e,k]*w[e,k,n], over all e."""
    pid_n = tl.program_id(0)
    pid_e = tl.program_id(1)
    pid_k = tl.program_id(2)
    e = tl.load(idx_ptr + pid_e).to(tl.int64)
    sc = tl.load(sc_ptr + pid_e)
    n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    nmask = n < N
    wq_e = wq_ptr + e * stride_we
    s_e = s_ptr + e * stride_se
    z_e = z_ptr + e * stride_se
    x_e = x_ptr + pid_e * stride_xe
    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
    HG: tl.constexpr = GROUP // 2
    rows = tl.arange(0, HG)
    for g in range(pid_k, NG, SPLIT_K):
        kp = g * HG + rows
        wq = tl.load(wq_e + kp[:, None] * stride_wk + n[None, :] * stride_wn,
                     mask=nmask[None, :], other=0)
        lo = (wq & 0xF).to(tl.float32)
        hi = ((wq >> 4) & 0xF).to(tl.float32)
        xe_ = tl.load(x_e + (2 * kp) * stride_xk)
        xo_ = tl.load(x_e + (2 * kp + 1) * stride_xk)
        raw = tl.sum(xe_[:, None] * lo + xo_[:, None] * hi, axis=0)
        xs = tl.sum(xe_ + xo_)
        s = tl.load(s_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
        z = tl.load(z_e + g * stride_sg + n * stride_sn, mask=nmask, other=0.0).to(tl.float32)
        acc += s * (raw - z * xs)
    tl.atomic_add(y_ptr + n, sc * acc, mask=nmask)


def gemv_acc(x, wq, scales, zeros, idx, scale, out):
    """x:[Eact,K]; accumulate scale[e]*(x[e]@dequant(wq[idx[e]])) into out[N]."""
    E, Kp, N = wq.shape
    K = Kp * 2
    NG = scales.shape[1]
    Eact = idx.numel()
    out.zero_()
    grid = lambda meta: (triton.cdiv(N, meta["BLOCK_N"]), Eact, meta["SPLIT_K"])
    _gemv_acc_kernel[grid](
        x, wq, scales, zeros, idx, scale, out,
        K, N, NG, Eact,
        x.stride(0), x.stride(1),
        wq.stride(0), wq.stride(1), wq.stride(2),
        scales.stride(0), scales.stride(1), scales.stride(2),
        GROUP=GROUP,
    )
    return out


@triton.jit
def _router_logits_kernel(x_ptr, w_ptr, out_ptr, D, BLOCK_D: tl.constexpr):
    """One program per expert: logit[e] = sum_d x[d] * W[e,d].
    x rounded to bf16 to match the reference's bf16 router (keeps top-k selection identical)."""
    e = tl.program_id(0)
    acc = tl.zeros([BLOCK_D], dtype=tl.float32)
    for d0 in range(0, D, BLOCK_D):
        dd = d0 + tl.arange(0, BLOCK_D)
        dm = dd < D
        xv = tl.load(x_ptr + dd, mask=dm, other=0.0).to(tl.bfloat16).to(tl.float32)
        acc += xv * tl.load(w_ptr + e * D + dd, mask=dm, other=0.0).to(tl.float32)
    tl.store(out_ptr + e, tl.sum(acc, 0))


@triton.jit
def _router_topk_kernel(
    logits_ptr, idxg_ptr, idxgu_ptr, scale_ptr,
    E, scaling, GATE_OFF,
    NACT: tl.constexpr, NS: tl.constexpr,
    BLOCK_E: tl.constexpr, BLOCK_J: tl.constexpr,
):
    """top-k(logits) + softmax + gather-index build (one program).
    Emits idx_g[NACT+NS], idx_gu[2*(NACT+NS)] (gate then up), scale[NACT+NS]."""
    e = tl.arange(0, BLOCK_E)
    emask = e < E
    logits = tl.where(emask, tl.load(logits_ptr + e, mask=emask, other=0.0), -float("inf"))
    jr = tl.arange(0, BLOCK_J)
    idx_vec = tl.zeros([BLOCK_J], dtype=tl.int32)
    val_vec = tl.full([BLOCK_J], -float("inf"), dtype=tl.float32)
    for j in range(NACT):
        mx = tl.max(logits, 0)
        am = tl.argmax(logits, 0).to(tl.int32)
        idx_vec = tl.where(jr == j, am, idx_vec)
        val_vec = tl.where(jr == j, mx, val_vec)
        logits = tl.where(e == am, -float("inf"), logits)
    idx_vec = tl.where(jr == NACT, E, idx_vec)               # shared expert at slot NACT
    routed = jr < NACT
    vmax = tl.max(tl.where(routed, val_vec, -float("inf")), 0)
    ev = tl.where(routed, tl.exp(val_vec - vmax), 0.0)
    sm = ev / tl.sum(ev, 0) * scaling
    scale_vec = tl.where(routed, sm, tl.where(jr == NACT, 1.0, 0.0))
    ns_mask = jr < (NACT + NS)
    tl.store(idxg_ptr + jr, idx_vec, mask=ns_mask)
    tl.store(scale_ptr + jr, scale_vec, mask=ns_mask)
    tl.store(idxgu_ptr + jr, idx_vec, mask=ns_mask)
    tl.store(idxgu_ptr + (NACT + NS) + jr, idx_vec + GATE_OFF, mask=ns_mask)


def router_topk(x, w, n_active, n_shared, gate_off, scaling, E, logits_buf):
    dev = x.device
    ns = n_active + n_shared
    idx_g = torch.empty(ns, device=dev, dtype=torch.int32)
    idx_gu = torch.empty(2 * ns, device=dev, dtype=torch.int32)
    scale = torch.empty(ns, device=dev, dtype=torch.float32)
    D = x.shape[-1]
    _router_logits_kernel[(E,)](x, w, logits_buf, D, BLOCK_D=triton.next_power_of_2(D))
    _router_topk_kernel[(1,)](
        logits_buf, idx_g, idx_gu, scale, E, scaling, gate_off,
        NACT=n_active, NS=n_shared,
        BLOCK_E=triton.next_power_of_2(E), BLOCK_J=triton.next_power_of_2(ns),
    )
    return idx_g, idx_gu, scale


@triton.jit
def _silu_mul_kernel(g_ptr, u_ptr, o_ptr, n_elem, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    off = pid * BLOCK + tl.arange(0, BLOCK)
    m = off < n_elem
    g = tl.load(g_ptr + off, mask=m, other=0.0)
    u = tl.load(u_ptr + off, mask=m, other=0.0)
    tl.store(o_ptr + off, (g * tl.sigmoid(g)) * u, mask=m)


def silu_mul(g, u):
    o = torch.empty_like(g)
    n = g.numel()
    _silu_mul_kernel[(triton.cdiv(n, 1024),)](g, u, o, n, BLOCK=1024)
    return o


@triton.jit
def _kda_conv_kernel(
    qkv_ptr, cw_ptr, st_ptr, out_ptr, C,
    scw_t, scw_n, scw_k, BLOCK: tl.constexpr,
):
    """Depthwise causal conv (kernel 4) + SiLU on q/k/v, with in-place ring shift.
    qkv_ptr: [>=3, C] (rows q,k,v); st_ptr: [3,3,C] window state; out_ptr: [3,C]."""
    pid_t = tl.program_id(0)
    pid_n = tl.program_id(1)
    n = pid_n * BLOCK + tl.arange(0, BLOCK)
    m = n < C
    val = tl.load(qkv_ptr + pid_t * C + n, mask=m, other=0.0).to(tl.float32)
    cw = cw_ptr + pid_t * scw_t + n * scw_n
    w0 = tl.load(cw + 0 * scw_k, mask=m, other=0.0).to(tl.float32)
    w1 = tl.load(cw + 1 * scw_k, mask=m, other=0.0).to(tl.float32)
    w2 = tl.load(cw + 2 * scw_k, mask=m, other=0.0).to(tl.float32)
    w3 = tl.load(cw + 3 * scw_k, mask=m, other=0.0).to(tl.float32)
    sb = st_ptr + pid_t * 3 * C + n
    s0 = tl.load(sb + 0 * C, mask=m, other=0.0).to(tl.float32)
    s1 = tl.load(sb + 1 * C, mask=m, other=0.0).to(tl.float32)
    s2 = tl.load(sb + 2 * C, mask=m, other=0.0).to(tl.float32)
    acc = s0 * w0 + s1 * w1 + s2 * w2 + val * w3
    tl.store(out_ptr + pid_t * C + n, acc * tl.sigmoid(acc), mask=m)
    tl.store(sb + 0 * C, s1.to(tl.bfloat16), mask=m)
    tl.store(sb + 1 * C, s2.to(tl.bfloat16), mask=m)
    tl.store(sb + 2 * C, val.to(tl.bfloat16), mask=m)


def kda_conv(qkv, conv_w, state):
    """qkv:[>=3,C]; conv_w:[3,C,4]; state:[3,3,C] bf16 (in place). -> out:[3,C] fp32."""
    C = qkv.shape[1]
    out = torch.empty((3, C), device=qkv.device, dtype=torch.float32)
    grid = (3, triton.cdiv(C, 1024))
    _kda_conv_kernel[grid](
        qkv, conv_w, state, out, C,
        conv_w.stride(0), conv_w.stride(1), conv_w.stride(2), BLOCK=1024,
    )
    return out


@triton.jit
def _kda_recur_kernel(
    S_ptr, q_ptr, k_ptr, v_ptr, g_ptr, beta_ptr, o_ptr,
    stride_sh, stride_sk, stride_sv, scale,
    DK: tl.constexpr, DV: tl.constexpr, BLOCK_V: tl.constexpr,
):
    pid_h = tl.program_id(0)
    pid_v = tl.program_id(1)
    koff = tl.arange(0, DK)
    voff = pid_v * BLOCK_V + tl.arange(0, BLOCK_V)
    key = tl.load(k_ptr + pid_h * DK + koff)
    qry = tl.load(q_ptr + pid_h * DK + koff) * scale
    graw = tl.load(g_ptr + pid_h * DK + koff)
    decay = tl.sigmoid(-graw)                              # = exp(-softplus(graw))
    beta = tl.load(beta_ptr + pid_h)
    val = tl.load(v_ptr + pid_h * DV + voff)
    sbase = S_ptr + pid_h * stride_sh + koff[:, None] * stride_sk + voff[None, :] * stride_sv
    S = tl.load(sbase)
    Sd = S * decay[:, None]
    pred = tl.sum(Sd * key[:, None], axis=0)
    Snew = Sd + (beta * key)[:, None] * (val - pred)[None, :]
    o = tl.sum(Snew * qry[:, None], axis=0)
    tl.store(sbase, Snew)
    tl.store(o_ptr + pid_h * DV + voff, o)


def kda_recur(S, q, k, v, g_raw, beta, scale, BLOCK_V=64):
    """S:[H,DK,DV] fp32 in place; q,k,v,g_raw:[H,DK]; beta:[H]. q scaled, decay=sigmoid(-g_raw)."""
    H, DK, DV = S.shape
    o = torch.empty((H, DV), device=S.device, dtype=torch.float32)
    grid = (H, triton.cdiv(DV, BLOCK_V))
    _kda_recur_kernel[grid](
        S, q, k, v, g_raw, beta, o,
        S.stride(0), S.stride(1), S.stride(2), scale,
        DK=DK, DV=DV, BLOCK_V=BLOCK_V,
    )
    return o


def _unpack_dequant(wq, scales, zeros, group=GROUP):
    """Dequant a single [K//2,N] int4 weight to [K,N] bf16 (only used for kv_b)."""
    Kp, N = wq.shape
    wu = torch.empty((2 * Kp, N), dtype=torch.uint8, device=wq.device)
    wu[0::2] = wq & 0xF
    wu[1::2] = (wq >> 4) & 0xF
    s = scales.repeat_interleave(group, dim=0)
    z = zeros.repeat_interleave(group, dim=0)
    return (wu.to(torch.bfloat16) - z) * s


# --------------------------------------------------------------------------- #
# modules (buffer/param names identical to reference for state_dict load)
# --------------------------------------------------------------------------- #
class QuantLinear(nn.Module):
    def __init__(self, in_f, out_f, group=GROUP):
        super().__init__()
        self.in_f, self.out_f, self.group = in_f, out_f, group
        ng = in_f // group
        self.register_buffer("w_q", torch.zeros(in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(ng, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(ng, out_f, dtype=torch.bfloat16))


class QuantExperts(nn.Module):
    def __init__(self, n, in_f, out_f, group=GROUP):
        super().__init__()
        self.n, self.in_f, self.out_f, self.group = n, in_f, out_f, group
        ng = in_f // group
        self.register_buffer("w_q", torch.zeros(n, in_f // 2, out_f, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(n, ng, out_f, dtype=torch.bfloat16))
        self.register_buffer("zeros", torch.zeros(n, ng, out_f, dtype=torch.bfloat16))


@triton.jit
def _rmsnorm_kernel(x_ptr, w_ptr, o_ptr, N: tl.constexpr, eps, BLOCK: tl.constexpr):
    off = tl.arange(0, BLOCK)
    m = off < N
    x = tl.load(x_ptr + off, mask=m, other=0.0).to(tl.float32)
    w = tl.load(w_ptr + off, mask=m, other=0.0).to(tl.float32)
    r = tl.rsqrt(tl.sum(x * x) / N + eps)
    tl.store(o_ptr + off, x * r * w, mask=m)


def _rmsnorm(x, w):
    """Fused RMSNorm; returns fp32 (ready for the dequant-GEMV)."""
    N = x.shape[-1]
    o = torch.empty(N, device=x.device, dtype=torch.float32)
    BLOCK = triton.next_power_of_2(N)
    _rmsnorm_kernel[(1,)](x, w, o, N, EPS, BLOCK=BLOCK)
    return o


@triton.jit
def _rmsnorm_add_kernel(x_ptr, y_ptr, w_ptr, resid_ptr, o_ptr, N: tl.constexpr, eps, BLOCK: tl.constexpr):
    off = tl.arange(0, BLOCK)
    m = off < N
    # round y to bf16 before the add to match reference's bf16 residual exactly
    yb = tl.load(y_ptr + off, mask=m, other=0.0).to(tl.bfloat16).to(tl.float32)
    s = tl.load(x_ptr + off, mask=m, other=0.0).to(tl.float32) + yb
    sb = s.to(tl.bfloat16)
    tl.store(resid_ptr + off, sb, mask=m)
    sf = sb.to(tl.float32)                          # match reference: norm the bf16 residual
    r = tl.rsqrt(tl.sum(sf * sf) / N + eps)
    w = tl.load(w_ptr + off, mask=m, other=0.0).to(tl.float32)
    tl.store(o_ptr + off, sf * r * w, mask=m)


def _rmsnorm_add(x, y, w):
    """resid = x + y (bf16); returns (resid, rmsnorm(resid) fp32)."""
    N = x.shape[-1]
    resid = torch.empty(N, device=x.device, dtype=torch.bfloat16)
    o = torch.empty(N, device=x.device, dtype=torch.float32)
    _rmsnorm_add_kernel[(1,)](x, y, w, resid, o, N, EPS, BLOCK=triton.next_power_of_2(N))
    return resid, o


def _rope_cossin(pos, dim, theta, device):
    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
    ang = pos * inv
    return torch.cos(ang), torch.sin(ang)


def _apply_rope(x, cos, sin):
    xf = x.float()
    even, odd = xf[..., 0::2], xf[..., 1::2]
    out = torch.empty_like(xf)
    out[..., 0::2] = even * cos - odd * sin
    out[..., 1::2] = odd * cos + even * sin
    return out.to(x.dtype)


class KDA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        H, Dk, d = cfg.kda_heads, cfg.kda_head_dim, cfg.hidden
        self.q_proj = QuantLinear(d, H * Dk, cfg.group)
        self.k_proj = QuantLinear(d, H * Dk, cfg.group)
        self.v_proj = QuantLinear(d, H * Dk, cfg.group)
        self.g_proj = QuantLinear(d, H * Dk, cfg.group)
        self.beta_proj = nn.Linear(d, H, bias=False, dtype=cfg.dtype)
        self.conv_w = nn.Parameter(torch.empty(3, H * Dk, cfg.short_conv, dtype=cfg.dtype))
        self.o_proj = QuantLinear(H * Dk, d, cfg.group)
        self.scale = Dk ** -0.5

    def _stack(self):
        self._wq = torch.stack([self.q_proj.w_q, self.k_proj.w_q,
                                self.v_proj.w_q, self.g_proj.w_q])
        self._s = torch.stack([self.q_proj.scales, self.k_proj.scales,
                               self.v_proj.scales, self.g_proj.scales])
        self._z = torch.stack([self.q_proj.zeros, self.k_proj.zeros,
                               self.v_proj.zeros, self.g_proj.zeros])

    def _conv(self, val, prev, idx):
        # val:[C] fp32, prev:[3,C] bf16, conv_w[idx]:[C,4]
        win = torch.cat([prev.float(), val[None]], dim=0)      # [4,C]
        w = self.conv_w[idx].float().transpose(0, 1)            # [4,C]
        out = F.silu((win * w).sum(0))                          # [C]
        return out, win[1:].to(prev.dtype)

    def step(self, x, st):
        H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim
        if not hasattr(self, "_wq"):
            self._stack()
        xf = x.float()
        qkvg = gemv(xf, self._wq, self._s, self._z)            # [4, H*Dk]
        q, st["cq"] = self._conv(qkvg[0], st["cq"], 0)
        k, st["ck"] = self._conv(qkvg[1], st["ck"], 1)
        v, st["cv"] = self._conv(qkvg[2], st["cv"], 2)
        q = q.view(H, Dk).contiguous()
        k = k.view(H, Dk).contiguous()
        v = v.view(H, Dk).contiguous()
        g_raw = qkvg[3].view(H, Dk).contiguous()
        beta = torch.sigmoid(self.beta_proj(x.to(torch.bfloat16)).float()).contiguous()
        o = kda_recur(st["S"], q, k, v, g_raw, beta, self.scale)   # updates S in place
        of = o.reshape(H * Dk)
        return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)

    def step_g(self, x, S_buf, conv_state):
        H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim
        if not hasattr(self, "_wq"):
            self._stack()
        xf = x.float()
        qkvg = gemv(xf, self._wq, self._s, self._z)                # [4, C]
        conv_out = kda_conv(qkvg, self.conv_w, conv_state)         # [3, C], state updated
        beta = torch.sigmoid(self.beta_proj(x.to(torch.bfloat16)).float())
        o = kda_recur(S_buf, conv_out[0], conv_out[1], conv_out[2], qkvg[3], beta, self.scale)
        of = o.reshape(H * Dk)
        return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)


class MLA(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        H, d = cfg.mla_heads, cfg.hidden
        self.q_proj = QuantLinear(d, H * (cfg.qk_nope + cfg.qk_rope), cfg.group)
        self.kv_a = QuantLinear(d, cfg.kv_lora + cfg.qk_rope, cfg.group)
        self.kv_b = QuantLinear(cfg.kv_lora, H * (cfg.qk_nope + cfg.v_head), cfg.group)
        self.o_proj = QuantLinear(H * cfg.v_head, d, cfg.group)
        self.scale = (cfg.qk_nope + cfg.qk_rope) ** -0.5

    def _prep(self):
        cfg = self.cfg
        H, nope, vh, lora = cfg.mla_heads, cfg.qk_nope, cfg.v_head, cfg.kv_lora
        W = _unpack_dequant(self.kv_b.w_q, self.kv_b.scales, self.kv_b.zeros).view(lora, H, nope + vh)
        # per-head absorb matrices, bf16 for tensor-core matmuls
        self._Wkn = W[:, :, :nope].permute(1, 0, 2).contiguous()   # [H,512,nope]
        self._Wvn = W[:, :, nope:].permute(1, 0, 2).contiguous()   # [H,512,vh]
        dev = self.kv_b.w_q.device
        self._inv = 1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.qk_rope, 2, device=dev, dtype=torch.float32) / cfg.qk_rope))
        # combined q_proj || kv_a (same input x); split output at qsplit
        self._qkva_wq = torch.cat([self.q_proj.w_q, self.kv_a.w_q], dim=1)
        self._qkva_s = torch.cat([self.q_proj.scales, self.kv_a.scales], dim=1)
        self._qkva_z = torch.cat([self.q_proj.zeros, self.kv_a.zeros], dim=1)
        self._qsplit = H * (nope + cfg.qk_rope)

    def step(self, x, st):
        cfg = self.cfg
        H = cfg.mla_heads
        nope, rope, vh, lora = cfg.qk_nope, cfg.qk_rope, cfg.v_head, cfg.kv_lora
        if not hasattr(self, "_Wkn"):
            self._prep()
        pos = st["c_kv"].shape[0]
        xf = x.float()
        qkva = gemv(xf, self._qkva_wq, self._qkva_s, self._qkva_z).squeeze(0)
        q = qkva[:self._qsplit].view(H, nope + rope)
        q_nope = q[:, :nope]
        q_rope = q[:, nope:]
        kv = qkva[self._qsplit:]
        c_kv = kv[:lora].to(torch.bfloat16)
        k_rope = kv[lora:]
        cos, sin = _rope_cossin(pos, rope, cfg.rope_theta, x.device)
        q_rope = _apply_rope(q_rope, cos, sin).to(torch.bfloat16)
        k_rope = _apply_rope(k_rope, cos, sin).to(torch.bfloat16)
        st["c_kv"] = torch.cat([st["c_kv"], c_kv[None]], 0)
        st["k_rope"] = torch.cat([st["k_rope"], k_rope[None]], 0)
        ckv = st["c_kv"]                                           # [L,512] bf16
        krope = st["k_rope"]                                       # [L,64] bf16
        # absorb q into latent space (per-head [512,nope] @ [nope] -> [512])
        q_absorb = torch.einsum("hdc,hc->hd", self._Wkn, q_nope.to(torch.bfloat16))  # [H,512] bf16
        scores = (q_absorb @ ckv.t() + q_rope @ krope.t()).float() * self.scale       # [H,L]
        p = torch.softmax(scores, dim=-1).to(torch.bfloat16)                          # [H,L]
        pc = p @ ckv                                                                   # [H,512] bf16
        o = torch.einsum("hd,hdc->hc", pc, self._Wvn)                                  # [H,vh] bf16
        of = o.reshape(H * vh).float()
        return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)

    def step_g(self, x, c_kv_buf, k_rope_buf, pos, maxlen):
        """Graph variant: pos is a [1] int64 GPU scalar; cache written in place at pos."""
        cfg = self.cfg
        H = cfg.mla_heads
        nope, rope, vh, lora = cfg.qk_nope, cfg.qk_rope, cfg.v_head, cfg.kv_lora
        if not hasattr(self, "_Wkn"):
            self._prep()
        xf = x.float()
        qkva = gemv(xf, self._qkva_wq, self._qkva_s, self._qkva_z).squeeze(0)
        q = qkva[:self._qsplit].view(H, nope + rope)
        q_nope = q[:, :nope]
        q_rope = q[:, nope:]
        kv = qkva[self._qsplit:]
        c_kv_new = kv[:lora].to(torch.bfloat16)
        k_rope_new = kv[lora:]
        inv = self._inv
        ang = pos.float() * inv                                  # [rope/2]
        cos, sin = torch.cos(ang), torch.sin(ang)
        q_rope = _apply_rope(q_rope, cos, sin).to(torch.bfloat16)
        k_rope_new = _apply_rope(k_rope_new, cos, sin).to(torch.bfloat16)
        c_kv_buf.index_copy_(0, pos, c_kv_new[None])
        k_rope_buf.index_copy_(0, pos, k_rope_new[None])
        q_absorb = torch.einsum("hdc,hc->hd", self._Wkn, q_nope.to(torch.bfloat16))    # [H,512]
        scores = (q_absorb @ c_kv_buf.t() + q_rope @ k_rope_buf.t()).float() * self.scale  # [H,maxlen]
        mask = torch.arange(maxlen, device=x.device) > pos
        scores = scores.masked_fill(mask[None, :], float("-inf"))
        p = torch.softmax(scores, dim=-1).to(torch.bfloat16)
        pc = p @ c_kv_buf
        o = torch.einsum("hd,hdc->hc", pc, self._Wvn)
        of = o.reshape(H * vh).float()
        return gemv(of, self.o_proj.w_q, self.o_proj.scales, self.o_proj.zeros).squeeze(0)


class MoE(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        d, m, E = cfg.hidden, cfg.moe_inter, cfg.n_experts
        self.router = nn.Linear(d, E, bias=False, dtype=cfg.dtype)
        self.gate = QuantExperts(E, d, m, cfg.group)
        self.up = QuantExperts(E, d, m, cfg.group)
        self.down = QuantExperts(E, m, d, cfg.group)
        self.s_gate = QuantExperts(cfg.n_shared, d, m, cfg.group)
        self.s_up = QuantExperts(cfg.n_shared, d, m, cfg.group)
        self.s_down = QuantExperts(cfg.n_shared, m, d, cfg.group)

    def _prep(self):
        dev = self.gate.w_q.device
        # combined routed+shared expert weights: index E (=n_experts) is the shared expert
        E = self.cfg.n_experts
        # combined gate||up weights: rows [gate(E), s_gate(1), up(E), s_up(1)]
        self._gu_wq = torch.cat([self.gate.w_q, self.s_gate.w_q, self.up.w_q, self.s_up.w_q], 0)
        self._gu_s = torch.cat([self.gate.scales, self.s_gate.scales, self.up.scales, self.s_up.scales], 0)
        self._gu_z = torch.cat([self.gate.zeros, self.s_gate.zeros, self.up.zeros, self.s_up.zeros], 0)
        self._d_wq = torch.cat([self.down.w_q, self.s_down.w_q], 0)
        self._d_s = torch.cat([self.down.scales, self.s_down.scales], 0)
        self._d_z = torch.cat([self.down.zeros, self.s_down.zeros], 0)
        self._gate_off = E + 1                                            # up block offset in gu
        self._out = torch.empty(self.cfg.hidden, device=dev, dtype=torch.float32)
        self._logits = torch.empty(E, device=dev, dtype=torch.float32)

    def step(self, x):
        cfg = self.cfg
        if not hasattr(self, "_gu_wq"):
            self._prep()
        ns = cfg.n_active + cfg.n_shared
        xf = x.float()
        idx_g, idx_gu, scale = router_topk(
            xf, self.router.weight, cfg.n_active, cfg.n_shared,
            self._gate_off, cfg.routed_scaling, cfg.n_experts, self._logits)
        gu = gemv(xf, self._gu_wq, self._gu_s, self._gu_z, idx=idx_gu)   # [2*ns, m]
        h = silu_mul(gu[:ns], gu[ns:])                                    # [ns, m]
        gemv_acc(h, self._d_wq, self._d_s, self._d_z, idx_g, scale, self._out)
        return self._out


class Block(nn.Module):
    def __init__(self, cfg, kind):
        super().__init__()
        self.kind = kind
        self.attn_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.moe_norm = nn.Parameter(torch.ones(cfg.hidden, dtype=cfg.dtype))
        self.attn = KDA(cfg) if kind == "K" else MLA(cfg)
        self.moe = MoE(cfg)

    def step(self, x, st):
        h = x + self.attn.step(_rmsnorm(x, self.attn_norm), st)
        return h + self.moe.step(_rmsnorm(h, self.moe_norm))


HEADROOM = 256


class _Bundle:
    """Persistent buffers + captured CUDA graph for one context length."""
    __slots__ = ("maxlen", "h_in", "h_out", "pos", "cpos", "S", "conv",
                 "c_kv", "k_rope", "graph", "mla_idx")


class Model(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self._graphs = {}
        self._use_graph = True

    # ---- eager reference path (always correct) ---------------------------- #
    def _step_eager(self, hidden, state):
        for i, blk in enumerate(self.blocks):
            hidden = blk.step(hidden, state[i])
        return hidden, state

    # ---- graph machinery -------------------------------------------------- #
    def _forward_graph(self, b: _Bundle):
        ki = 0
        resid, add = b.h_in, None         # deferred residual: actual stream = resid + add
        for i, blk in enumerate(self.blocks):
            if add is None:
                xn = _rmsnorm(resid, blk.attn_norm)
            else:
                resid, xn = _rmsnorm_add(resid, add, blk.attn_norm)   # fold prev moe-residual add
            if blk.kind == "K":
                a = blk.attn.step_g(xn, b.S[ki], b.conv[ki])
                ki += 1
            else:
                a = blk.attn.step_g(xn, b.c_kv, b.k_rope, b.pos, b.maxlen)
            resid, xn2 = _rmsnorm_add(resid, a, blk.moe_norm)         # fold attn-residual add
            add = blk.moe.step(xn2)
        b.h_out.copy_(resid + add.to(torch.bfloat16))
        b.pos.add_(1)

    def _make_bundle(self, ctx, device):
        cfg = self.cfg
        H, Dk = cfg.kda_heads, cfg.kda_head_dim
        C = H * Dk
        b = _Bundle()
        b.maxlen = ctx + HEADROOM
        b.h_in = torch.zeros(cfg.hidden, device=device, dtype=torch.bfloat16)
        b.h_out = torch.zeros(cfg.hidden, device=device, dtype=torch.bfloat16)
        b.pos = torch.zeros(1, device=device, dtype=torch.int64)
        nk = sum(1 for k in cfg.pattern if k == "K")
        b.S = [torch.zeros(H, Dk, Dk, device=device, dtype=torch.float32) for _ in range(nk)]
        b.conv = [torch.zeros(3, 3, C, device=device, dtype=torch.bfloat16) for _ in range(nk)]
        b.c_kv = torch.zeros(b.maxlen, cfg.kv_lora, device=device, dtype=torch.bfloat16)
        b.k_rope = torch.zeros(b.maxlen, cfg.qk_rope, device=device, dtype=torch.bfloat16)
        b.mla_idx = cfg.pattern.index("M")
        return b

    def _capture(self, ctx, device):
        b = self._make_bundle(ctx, device)
        b.pos.fill_(ctx)
        s = torch.cuda.Stream()
        s.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(s):
            for _ in range(3):
                self._forward_graph(b)
        torch.cuda.current_stream().wait_stream(s)
        torch.cuda.synchronize()
        b.graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(b.graph):
            self._forward_graph(b)
        self._graphs[ctx] = b
        return b

    def _load_state(self, b, state, ctx):
        """Copy a fresh dict state into bundle buffers; rebind dict to buffer views."""
        ki = 0
        for i, blk in enumerate(self.blocks):
            st = state[i]
            if blk.kind == "K":
                b.S[ki].copy_(st["S"])
                b.conv[ki][0].copy_(st["cq"]); b.conv[ki][1].copy_(st["ck"]); b.conv[ki][2].copy_(st["cv"])
                st["S"] = b.S[ki]
                st["cq"] = b.conv[ki][0]; st["ck"] = b.conv[ki][1]; st["cv"] = b.conv[ki][2]
                ki += 1
            else:
                b.c_kv[:ctx].copy_(st["c_kv"]); b.k_rope[:ctx].copy_(st["k_rope"])
                st["c_kv"] = b.c_kv[:ctx]; st["k_rope"] = b.k_rope[:ctx]
        b.pos.fill_(ctx)
        b.cpos = ctx

    def step(self, hidden, state):
        if not self._use_graph or not hidden.is_cuda:
            return self._step_eager(hidden, state)
        mla_idx = self.cfg.pattern.index("M")
        ckv = state[mla_idx]["c_kv"]
        try:
            b = self._cur if getattr(self, "_cur", None) is not None else None
            is_cont = (b is not None and ckv.data_ptr() == b.c_kv.data_ptr())
            if is_cont and b.cpos + 1 >= b.maxlen:        # cache overflow: drop to eager
                self._use_graph = False
                return self._step_eager(hidden, state)
            if not is_cont:
                ctx = ckv.shape[0]
                b = self._graphs.get(ctx) or self._capture(ctx, hidden.device)
                self._load_state(b, state, ctx)
                self._cur = b
            b.h_in.copy_(hidden)
            b.graph.replay()
            b.cpos += 1
            p = b.cpos
            state[mla_idx]["c_kv"] = b.c_kv[:p]
            state[mla_idx]["k_rope"] = b.k_rope[:p]
            return b.h_out, state
        except Exception:
            self._use_graph = False
            self._cur = None
            return self._step_eager(hidden, state)