"""W4A16 fused-dequant Kimi-Linear hybrid decode.

Layout mirrors the oracle and the baseline. The Model, KDA, MLA, MoE, Block
classes expose the same parameters and buffers (w_q, scales, zeros, etc.) so
the oracle's state_dict loads cleanly. The big projection weights (Q, K, V,
G, O, KV_A, KV_B, all experts) are int4 AWQ/GPTQ-style; we replace every
materializing linear with a custom Triton kernel that streams the int4 bytes
once and dequants on the fly inside the GEMV.

The roofline win comes from the missing-write: the oracle and the baseline
both unpack the int4 weight into a full bf16 matrix before matmul, doubling
the weight traffic. Here the bf16 weight never exists in DRAM. Plus:

* KDA's q/k/v/g projections share the same input, so they are stacked along
  the K dim and computed in one fused kernel call -- 4x the output cols but
  only one launch and one K-loop pass over x.
* MoE batches the 8 routed experts' gate+up projections as a single
  (2*E, in, out) tensor and runs one Triton kernel pass for all of them.
  The down projection is similarly batched over the active experts.
* MLA uses a partial absorb: k_nope is materialized (with the bf16 cast
  needed to match the oracle's precision), but the V projection is
  absorbed -- o = (p^T c_kv) @ W_v_b -- so only half the (L, H, D)
  intermediate is written. The kv_b weight is cached as bf16 since it is
  small and the cast was the only reason to materialize.
* The MLA softmax over (L, H) is done in the (H, L) layout (transpose) --
  the dim=0 path on a tall skinny tensor is 100x slower than dim=1 on the
  transpose; the transpose is a few us and the softmax is a few us.
* QuantLinear / QuantExperts preserve the int4 buffers (w_q, scales, zeros)
  from the oracle's state_dict so the dequant-only-once weight cache for
  MLA's kv_b can read them directly.
"""
from __future__ import annotations

from dataclasses import dataclass, field

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


# --------------------------------------------------------------------------- #
# Triton: fused int4 dequant + GEMV
# --------------------------------------------------------------------------- #
@triton.jit
def _int4_gemv_kernel(
    w_q_ptr, scales_ptr, zeros_ptr, x_ptr, out_ptr,
    K, N,
    stride_wq_k, stride_wq_n,
    stride_sg, stride_sn,
    BLOCK_K: tl.constexpr,
    BLOCK_N: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
):
    """Single linear: y = x @ dequant(w_q, scales, zeros). x is (K,), y is (N,)."""
    pid = tl.program_id(0)
    n_offs = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = n_offs < N

    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)

    n_groups = K // GROUP_SIZE
    for g in range(n_groups):
        k_offs = g * GROUP_SIZE + tl.arange(0, BLOCK_K)
        x = tl.load(x_ptr + k_offs).to(tl.float32)

        wq_row_offs = k_offs // 2
        wq_offs = wq_row_offs[:, None] * stride_wq_k + n_offs[None, :] * stride_wq_n
        wq = tl.load(w_q_ptr + wq_offs, mask=n_mask[None, :], other=0)

        is_even = (k_offs % 2 == 0)
        w_unpacked = tl.where(is_even[:, None], wq & 0xF, (wq >> 4) & 0xF)
        w_unpacked = w_unpacked.to(tl.float32)

        s = tl.load(scales_ptr + g * stride_sg + n_offs * stride_sn,
                    mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(zeros_ptr + g * stride_sg + n_offs * stride_sn,
                    mask=n_mask, other=0.0).to(tl.float32)

        w_bf = (w_unpacked - z[None, :]) * s[None, :]
        acc += tl.sum(x[:, None] * w_bf, axis=0)

    tl.store(out_ptr + n_offs, acc.to(tl.bfloat16), mask=n_mask)


@triton.jit
def _int4_gemv_batched_kernel(
    w_q_ptr, scales_ptr, zeros_ptr, x_ptr, out_ptr,
    K, N, E,
    stride_wq_e, stride_wq_k, stride_wq_n,
    stride_s_e, stride_sg, stride_sn,
    stride_xe, stride_xk,
    BLOCK_K: tl.constexpr,
    BLOCK_N: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
):
    """Batched experts: w_q (E, K//2, N), x (E, K), out (E, N).
    Each expert e uses its own x[e, :] as input."""
    pid_e = tl.program_id(0)
    pid_n = tl.program_id(1)

    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = n_offs < N

    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)

    wq_e_ptr = w_q_ptr + pid_e * stride_wq_e
    s_e_ptr = scales_ptr + pid_e * stride_s_e
    z_e_ptr = zeros_ptr + pid_e * stride_s_e
    x_e_ptr = x_ptr + pid_e * stride_xe

    n_groups = K // GROUP_SIZE
    for g in range(n_groups):
        k_offs = g * GROUP_SIZE + tl.arange(0, BLOCK_K)
        x = tl.load(x_e_ptr + k_offs * stride_xk).to(tl.float32)

        wq_row_offs = k_offs // 2
        wq_offs = wq_row_offs[:, None] * stride_wq_k + n_offs[None, :] * stride_wq_n
        wq = tl.load(wq_e_ptr + wq_offs, mask=n_mask[None, :], other=0)

        is_even = (k_offs % 2 == 0)
        w_unpacked = tl.where(is_even[:, None], wq & 0xF, (wq >> 4) & 0xF)
        w_unpacked = w_unpacked.to(tl.float32)

        s = tl.load(s_e_ptr + g * stride_sg + n_offs * stride_sn,
                    mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_e_ptr + g * stride_sg + n_offs * stride_sn,
                    mask=n_mask, other=0.0).to(tl.float32)

        w_bf = (w_unpacked - z[None, :]) * s[None, :]
        acc += tl.sum(x[:, None] * w_bf, axis=0)

    tl.store(out_ptr + pid_e * N + n_offs, acc.to(tl.bfloat16), mask=n_mask)


def _int4_gemv(x, w_q, scales, zeros, group_size=128, out=None, block_n=64):
    K = x.shape[0]
    N = w_q.shape[1]
    if out is None:
        out = torch.empty(N, dtype=torch.bfloat16, device=x.device)
    grid = (triton.cdiv(N, block_n),)
    _int4_gemv_kernel[grid](
        w_q, scales, zeros, x, out,
        K, N,
        w_q.stride(0), w_q.stride(1),
        scales.stride(0), scales.stride(1),
        BLOCK_K=group_size, BLOCK_N=block_n, GROUP_SIZE=group_size,
        num_warps=4, num_stages=2,
    )
    return out


@triton.jit
def _int4_gemv_fused_kernel(
    w_q_ptr, scales_ptr, zeros_ptr, x_ptr, out_ptr,
    K, N, N_PROJ: tl.constexpr,
    stride_wq_k, stride_wq_n,
    stride_sg, stride_sn,
    BLOCK_K: tl.constexpr,
    BLOCK_N: tl.constexpr,
    GROUP_SIZE: tl.constexpr,
):
    """N_PROJ projections, each with the same K and N, stacked along the K dim of w_q.
    w_q: (N_PROJ * K//2, N), scales: (N_PROJ * K//g, N), zeros: (N_PROJ * K//g, N).
    Output: (N_PROJ * N,)."""
    pid_n = tl.program_id(0)
    pid_proj = tl.program_id(1)

    n_offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = n_offs < N

    wq_offset = pid_proj * (K // 2) * N
    s_offset = pid_proj * (K // GROUP_SIZE) * N

    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)

    n_groups = K // GROUP_SIZE
    for g in range(n_groups):
        k_offs = g * GROUP_SIZE + tl.arange(0, BLOCK_K)
        x = tl.load(x_ptr + k_offs).to(tl.float32)

        wq_row_offs = k_offs // 2
        wq_offs = wq_offset + wq_row_offs[:, None] * N + n_offs[None, :]
        wq = tl.load(w_q_ptr + wq_offs, mask=n_mask[None, :], other=0)

        is_even = (k_offs % 2 == 0)
        w_unpacked = tl.where(is_even[:, None], wq & 0xF, (wq >> 4) & 0xF)
        w_unpacked = w_unpacked.to(tl.float32)

        s_offs = s_offset + g * N + n_offs
        s = tl.load(scales_ptr + s_offs, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(zeros_ptr + s_offs, mask=n_mask, other=0.0).to(tl.float32)

        w_bf = (w_unpacked - z[None, :]) * s[None, :]
        acc += tl.sum(x[:, None] * w_bf, axis=0)

    out_offset = pid_proj * N
    tl.store(out_ptr + out_offset + n_offs, acc.to(tl.bfloat16), mask=n_mask)


def _int4_gemv_fused(x, w_q_stacked, scales_stacked, zeros_stacked,
                    group_size=128, block_n=64, num_warps=4, num_stages=2, n_proj=4):
    """Stacked: w_q (N_PROJ * K//2, N), scales (N_PROJ * K//g, N), zeros (N_PROJ * K//g, N).
    Output: (N_PROJ * N,)."""
    K_total = w_q_stacked.shape[0] * 2
    N = w_q_stacked.shape[1]
    K = K_total // n_proj
    out = torch.empty(n_proj * N, dtype=torch.bfloat16, device=x.device)
    grid = (triton.cdiv(N, block_n), n_proj)
    _int4_gemv_fused_kernel[grid](
        w_q_stacked, scales_stacked, zeros_stacked, x, out,
        K, N, n_proj,
        w_q_stacked.stride(0), w_q_stacked.stride(1),
        scales_stacked.stride(0), scales_stacked.stride(1),
        BLOCK_K=group_size, BLOCK_N=block_n, GROUP_SIZE=group_size,
        num_warps=num_warps, num_stages=num_stages,
    )
    return out


def _int4_gemv_batched(x, w_q, scales, zeros, group_size=128, out=None, block_n=64):
    """Batched: w_q (E, K//2, N), x (E, K) or (K,), out (E, N)."""
    E, Kh, N = w_q.shape
    K = Kh * 2
    if out is None:
        out = torch.empty(E, N, dtype=torch.bfloat16, device=x.device)
    grid = (E, triton.cdiv(N, block_n))
    # If x is (K,) and E=1, treat as (1, K)
    if x.dim() == 1:
        x = x.unsqueeze(0).expand(E, -1)
    _int4_gemv_batched_kernel[grid](
        w_q, scales, zeros, x, out,
        K, N, E,
        w_q.stride(0), w_q.stride(1), w_q.stride(2),
        scales.stride(0), scales.stride(1), scales.stride(2),
        x.stride(0) if x.dim() >= 2 else K, x.stride(1) if x.dim() >= 2 else 1,
        BLOCK_K=group_size, BLOCK_N=block_n, GROUP_SIZE=group_size,
        num_warps=4, num_stages=2,
    )
    return out


# --------------------------------------------------------------------------- #
# Reference-style W4A16 helpers (for state_dict load compat)
# --------------------------------------------------------------------------- #
@dataclass(frozen=True)
class Config:
    hidden: int = 2304
    kda_heads: int = 32
    kda_head_dim: int = 128
    short_conv: int = 4
    mla_heads: int = 32
    kv_lora: int = 512
    qk_nope: int = 128
    qk_rope: int = 64
    v_head: int = 128
    rope_theta: float = 10000.0
    n_experts: int = 64
    n_active: int = 8
    n_shared: int = 1
    moe_inter: int = 1024
    routed_scaling: float = 2.446
    group: int = 128
    pattern: tuple = ("K", "K", "K", "M")
    dtype: torch.dtype = field(default=torch.bfloat16)


def build_config(shape: dict) -> Config:
    return Config(n_experts=int(shape.get("n_experts", 64)))


def _pack_int4(w_q: torch.Tensor) -> torch.Tensor:
    lo = w_q[0::2] & 0xF
    hi = w_q[1::2] & 0xF
    return (lo | (hi << 4)).contiguous()


def _unpack_int4(w_packed: torch.Tensor, K: int) -> torch.Tensor:
    out = torch.empty((K, w_packed.shape[1]), dtype=torch.uint8, device=w_packed.device)
    out[0::2] = w_packed & 0xF
    out[1::2] = (w_packed >> 4) & 0xF
    return out


def quantize(w_io: torch.Tensor, group: int = 128):
    K, N = w_io.shape
    ng = K // group
    wg = w_io.view(ng, group, N).float()
    wmin = wg.min(dim=1, keepdim=True).values
    wmax = wg.max(dim=1, keepdim=True).values
    scales = (wmax - wmin).clamp_min(1e-8) / 15.0
    zeros = (-wmin / scales).round().clamp(0, 15)
    w_q = ((wg / scales) + zeros).round().clamp(0, 15).to(torch.uint8).view(K, N)
    return _pack_int4(w_q), scales.squeeze(1).to(torch.bfloat16), zeros.squeeze(1).to(torch.bfloat16)


def dequant(w_q: torch.Tensor, scales: torch.Tensor, zeros: torch.Tensor, K: int, group: int) -> torch.Tensor:
    wu = _unpack_int4(w_q, K).to(torch.bfloat16)
    s = scales.repeat_interleave(group, dim=0)
    z = zeros.repeat_interleave(group, dim=0)
    return (wu - z) * s


class QuantLinear(nn.Module):
    def __init__(self, in_f: int, out_f: int, group: int = 128):
        super().__init__()
        assert in_f % group == 0 and in_f % 2 == 0
        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))

    def init_random(self, gen: torch.Generator, std: float = 0.02) -> None:
        w = torch.randn(self.in_f, self.out_f, generator=gen) * std
        wq, s, z = quantize(w, self.group)
        self.w_q.copy_(wq)
        self.scales.copy_(s)
        self.zeros.copy_(z)

    def weight_bf(self) -> torch.Tensor:
        return dequant(self.w_q, self.scales, self.zeros, self.in_f, self.group)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return _int4_gemv(x, self.w_q, self.scales, self.zeros, self.group)


class QuantExperts(nn.Module):
    def __init__(self, n: int, in_f: int, out_f: int, group: int = 128):
        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))

    def init_random(self, gen: torch.Generator, std: float = 0.02) -> None:
        for e in range(self.n):
            w = torch.randn(self.in_f, self.out_f, generator=gen) * std
            wq, s, z = quantize(w, self.group)
            self.w_q[e].copy_(wq)
            self.scales[e].copy_(s)
            self.zeros[e].copy_(z)

    def weight_bf(self, e: int) -> torch.Tensor:
        return dequant(self.w_q[e], self.scales[e], self.zeros[e], self.in_f, self.group)


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def _rmsnorm(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    xf = x.float()
    xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6)
    return (xf * w.float()).to(x.dtype)


def _rope_cossin(pos: int, dim: int, theta: float, 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: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    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)


# --------------------------------------------------------------------------- #
# KDA
# --------------------------------------------------------------------------- #
class KDA(nn.Module):
    def __init__(self, cfg: Config):
        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
        # Stacked weights for fused q/k/v/g projection
        # Each has shape (K//2, N) where K=d=2304, N=H*Dk=4096
        # Stacked: (4*K//2, N) = (4608, 4096)
        self._qkv_stacked = False  # lazy init after load_state_dict
        self._cached_qkv = None

    def _build_qkv_stacked(self):
        """Build the stacked q/k/v/g weights from the individual ones."""
        w_q = torch.cat([self.q_proj.w_q, self.k_proj.w_q,
                          self.v_proj.w_q, self.g_proj.w_q], dim=0).contiguous()
        s = torch.cat([self.q_proj.scales, self.k_proj.scales,
                        self.v_proj.scales, self.g_proj.scales], dim=0).contiguous()
        z = torch.cat([self.q_proj.zeros, self.k_proj.zeros,
                        self.v_proj.zeros, self.g_proj.zeros], dim=0).contiguous()
        self._cached_qkv = (w_q, s, z)
        self._qkv_stacked = True

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

    def step(self, x, st):
        H, Dk = self.cfg.kda_heads, self.cfg.kda_head_dim
        if self._cached_qkv is None:
            self._build_qkv_stacked()
        w_q, s, z = self._cached_qkv
        qkvg = _int4_gemv_fused(x, w_q, s, z, self.cfg.group, n_proj=4)
        # qkvg: (4 * H * Dk,) = (16384,)
        qkv = qkvg[:3 * H * Dk].view(3, H * Dk)
        g_flat = qkvg[3 * H * Dk:]
        q, st["cq"] = self._short_conv(qkv[0], st["cq"], 0)
        k, st["ck"] = self._short_conv(qkv[1], st["ck"], 1)
        v, st["cv"] = self._short_conv(qkv[2], st["cv"], 2)
        q = q.view(H, Dk).float() * self.scale
        k = k.view(H, Dk).float()
        v = v.view(H, Dk).float()
        g = (-F.softplus(g_flat.view(H, Dk).float())).view(H, Dk)
        beta = torch.sigmoid(self.beta_proj(x).float())
        S = st["S"] * g.exp()[:, :, None]
        pred = (S * k[:, :, None]).sum(1)
        S = S + beta[:, None, None] * k[:, :, None] * (v - pred)[:, None, :]
        o = (S * q[:, :, None]).sum(1)
        st["S"] = S
        return self.o_proj(o.reshape(H * Dk).to(torch.bfloat16))


# --------------------------------------------------------------------------- #
# MLA — absorb pattern (kv_b precomputed as bf16 since it's small)
# --------------------------------------------------------------------------- #
class MLA(nn.Module):
    def __init__(self, cfg: Config):
        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
        self._cached_kv_b = None  # bf16 cache
        self._cached_W_k_b_flat = None  # bf16 weight for the k_nope matmul
        self._cached_W_v_b = None  # bf16 weight for the v matmul (used in absorb)

    def _kv_b_bf16(self):
        if self._cached_kv_b is None or self._cached_kv_b.device != self.kv_b.w_q.device:
            # Precompute bf16 weight (kv_lora, H*(Hn+Hv))
            self._cached_kv_b = self.kv_b.weight_bf()
        return self._cached_kv_b

    def step(self, x, st):
        cfg = self.cfg
        H = cfg.mla_heads
        Hn, Hr, Hv = cfg.qk_nope, cfg.qk_rope, cfg.v_head
        pos = st["c_kv"].shape[0]

        q = self.q_proj(x).view(H, Hn + Hr)
        q_nope = q[:, :Hn]
        q_rope = q[:, Hn:]

        kv = self.kv_a(x)
        c_kv = kv[: cfg.kv_lora]
        k_rope = kv[cfg.kv_lora:]

        cos, sin = _rope_cossin(pos, Hr, cfg.rope_theta, x.device)
        q_rope = _apply_rope(q_rope, cos, sin)
        k_rope = _apply_rope(k_rope, cos, sin)
        st["c_kv"] = torch.cat([st["c_kv"], c_kv[None]], 0)
        st["k_rope"] = torch.cat([st["k_rope"], k_rope[None]], 0)

        # Compute only the k_nope part of kvb (we'll absorb the v part)
        # Materialize k_nope with bf16 cast to match reference precision
        if self._cached_W_k_b_flat is None:
            kv_b = self._kv_b_bf16()  # (kv_lora, H*(Hn+Hv))
            kv_lora = cfg.kv_lora
            kv_b_3d = kv_b.view(kv_lora, H, Hn + Hv)  # (kv_lora, H, Hn+Hv)
            W_k_b_3d = kv_b_3d[..., :Hn].permute(1, 0, 2).contiguous()  # (H, kv_lora, Hn)
            W_v_b_3d = kv_b_3d[..., Hn:].permute(1, 0, 2).contiguous()  # (H, kv_lora, Hv)
            # Cache as bf16 for matmul speed
            self._cached_W_k_b_flat = W_k_b_3d.view(H * Hn, kv_lora).t().contiguous().to(torch.bfloat16)
            self._cached_W_v_b = W_v_b_3d
        W_k_b_flat = self._cached_W_k_b_flat
        W_v_b_3d = self._cached_W_v_b
        kv_lora = cfg.kv_lora

        # bf16 matmul: c_kv (L+1, kv_lora) @ W_k_b_flat (kv_lora, H*Hn) = k_nope (L+1, H*Hn)
        k_nope = (st["c_kv"] @ W_k_b_flat).view(-1, H, Hn).float()
        # scores
        q_nope_f = q_nope.float()
        q_rope_f = q_rope.float()
        scores = (torch.einsum("hd,lhd->lh", q_nope_f, k_nope)
                  + torch.einsum("hd,ld->lh", q_rope_f, st["k_rope"].float())) * self.scale
        # softmax over dim=0 is slow on (L+1, H); transpose to (H, L+1), softmax dim=1, transpose back
        p = torch.softmax(scores.T, dim=1).T.to(x.dtype)
        # Absorb v: attn_out = p @ c_kv @ W_v_b
        attn_out_latent = (p.float().T @ st["c_kv"].float())  # (H, kv_lora)
        o = torch.einsum("hc,hcd->hd", attn_out_latent, W_v_b_3d.float()).reshape(H * Hv)
        return self.o_proj(o.to(torch.bfloat16))


# --------------------------------------------------------------------------- #
# MoE
# --------------------------------------------------------------------------- #
class MoE(nn.Module):
    def __init__(self, cfg: Config):
        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 _ffn(self, x, idx, weights, gate, up, down):
        """Run a single FFN for the given experts (batched)."""
        cfg = self.cfg
        E = idx.numel()
        # Fuse gate + up: (2*E, in, out) stacked
        wq_g = gate.w_q[idx]
        wq_u = up.w_q[idx]
        wq_gu = torch.cat([wq_g, wq_u], dim=0)  # (2*E, in/2, out)
        s_g = gate.scales[idx]
        s_u = up.scales[idx]
        s_gu = torch.cat([s_g, s_u], dim=0)  # (2*E, ng, out)
        z_g = gate.zeros[idx]
        z_u = up.zeros[idx]
        z_gu = torch.cat([z_g, z_u], dim=0)  # (2*E, ng, out)
        # Batched call: (2*E, K) input (x broadcast), (2*E, m) output
        if x.dim() == 1:
            x_b = x.unsqueeze(0).expand(2 * E, -1)
        else:
            x_b = x.expand(2 * E, -1)
        gu = _int4_gemv_batched(x_b, wq_gu, s_gu, z_gu, cfg.group)
        g = gu[:E].float()
        u = gu[E:].float()
        h = F.silu(g) * u  # (E, m)
        # For the down projection, each expert has a different input h[e, :]
        wq_d = down.w_q[idx]
        s_d = down.scales[idx]
        z_d = down.zeros[idx]
        d = _int4_gemv_batched(h.to(torch.bfloat16), wq_d, s_d, z_d, cfg.group).float()
        return (weights[:, None] * d).sum(0) if weights is not None else d.sum(0)

    def step(self, x):
        cfg = self.cfg
        probs = torch.softmax(self.router(x).float(), dim=-1)
        w, idx = torch.topk(probs, cfg.n_active)
        w = (w / (w.sum() + 1e-9) * cfg.routed_scaling)
        out = self._ffn(x, idx, w, self.gate, self.up, self.down)
        sidx = torch.arange(cfg.n_shared, device=x.device)
        out = out + self._ffn(x, sidx, None, self.s_gate, self.s_up, self.s_down)
        return out.to(torch.bfloat16)


# --------------------------------------------------------------------------- #
# Block / Model
# --------------------------------------------------------------------------- #
class Block(nn.Module):
    def __init__(self, cfg: Config, kind: str):
        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))


class Model(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList(Block(cfg, k) for k in cfg.pattern)
        self.reset_parameters()

    def reset_parameters(self):
        g = torch.Generator(device="cpu").manual_seed(1234)
        for mod in self.modules():
            if isinstance(mod, (QuantLinear, QuantExperts)):
                mod.init_random(g)
            elif isinstance(mod, nn.Linear):
                nn.init.normal_(mod.weight, 0.0, 0.02, generator=g)
            elif isinstance(mod, KDA):
                nn.init.normal_(mod.conv_w, 0.0, 0.1, generator=g)

    def step(self, hidden, state):
        for i, blk in enumerate(self.blocks):
            hidden = blk.step(hidden, state[i])
        return hidden, state


# --------------------------------------------------------------------------- #
# state / inputs
# --------------------------------------------------------------------------- #
def init_state(cfg: Config, context_len: int, seed: int) -> list:
    dev = torch.device("cuda:0")
    g = torch.Generator(device=dev).manual_seed(seed)
    H, Dk = cfg.kda_heads, cfg.kda_head_dim
    C = H * Dk
    state = []
    for kind in cfg.pattern:
        if kind == "K":
            state.append({
                "S": torch.randn(H, Dk, Dk, device=dev, generator=g) * 0.05,
                "cq": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
                "ck": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
                "cv": torch.randn(cfg.short_conv - 1, C, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
            })
        else:
            state.append({
                "c_kv": torch.randn(context_len, cfg.kv_lora, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
                "k_rope": torch.randn(context_len, cfg.qk_rope, device=dev, generator=g, dtype=cfg.dtype) * 0.1,
            })
    return state


def init_token(cfg: Config, seed: int) -> torch.Tensor:
    dev = torch.device("cuda:0")
    g = torch.Generator(device=dev).manual_seed(seed + 1)
    return torch.randn(cfg.hidden, device=dev, generator=g, dtype=cfg.dtype) * 0.25


if __name__ == "__main__":
    cfg = build_config({"n_experts": 64})
    m = Model(cfg).cuda().eval()
    st = init_state(cfg, context_len=2048, seed=0)
    h = init_token(cfg, seed=0)
    with torch.no_grad():
        for _ in range(4):
            h, st = m.step(h, st)
    torch.cuda.synchronize()
    print(f"ok: out {tuple(h.shape)} finite {torch.isfinite(h).all().item()}")