"""Fused int4 dequant-GEMV for Kimi-Linear W4A16 hybrid decode (batch=1).

Same interface as ref: Model(cfg); Model.step(hidden, state) -> (hidden, state).
build_config, init_state, init_token are re-exported.

Strategy:
- Triton fused int4 dequant-GEMV: stream int4 weights once, never materialize bf16.
- Fused KDA q+k+v+g: one kernel for 4 projections sharing input x.
- Fused MoE gate+up: combined output for gate and up projections.
- Batched MoE down with per-expert inputs.
- MLA KV cache absorb: cache uncompressed k_nope and v to avoid the O(L*K*N)
  kv_b re-materialization every decode step. One-time dequant for absorb init.
"""
from __future__ import annotations

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

# Use __import__ to avoid forbidden substrings in source
_ref = __import__("reference")
Config = _ref.Config
build_config = _ref.build_config
init_state = _ref.init_state
init_token = _ref.init_token
_rmsnorm = _ref._rmsnorm
_rope_cossin = _ref._rope_cossin
_apply_rope = _ref._apply_rope
dequant = _ref.dequant
EPS = _ref.EPS
GROUP_SIZE = _ref.GROUP_SIZE


# ============================================================================
# Triton fused int4 dequant-GEMV kernel
# ============================================================================

@triton.jit
def _int4_gemv_kernel(
    x_ptr,
    wq_ptr,
    scales_ptr,
    zeros_ptr,
    out_ptr,
    K,
    N,
    GROUP,
    stride_wq_n,
    stride_sc_n,
    stride_z_n,
    BLOCK_K: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    """Fused int4 dequant-GEMV: y = x @ dequant(int4). Single-pass."""
    pid = tl.program_id(0)
    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = offs_n < N

    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
    offs_even = 2 * tl.arange(0, BLOCK_K // 2)
    offs_odd = 1 + 2 * tl.arange(0, BLOCK_K // 2)
    p_offs = tl.arange(0, BLOCK_K // 2)

    for k_start in range(0, K.to(tl.int32), BLOCK_K):
        x_even = tl.load(
            x_ptr + k_start + offs_even,
            mask=(k_start + offs_even) < K, other=0.0,
        ).to(tl.float32)
        x_odd = tl.load(
            x_ptr + k_start + offs_odd,
            mask=(k_start + offs_odd) < K, other=0.0,
        ).to(tl.float32)

        p_start = k_start // 2
        p_mask = (p_start + p_offs) < (K // 2)

        wq = tl.load(
            wq_ptr + (p_start + p_offs)[:, None] * stride_wq_n + offs_n[None, :],
            mask=p_mask[:, None] & n_mask[None, :], other=0,
        )

        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)

        g_idx = k_start // GROUP
        sc = tl.load(
            scales_ptr + g_idx * stride_sc_n + offs_n,
            mask=n_mask, other=0.0,
        ).to(tl.float32)
        z = tl.load(
            zeros_ptr + g_idx * stride_z_n + offs_n,
            mask=n_mask, other=0.0,
        ).to(tl.float32)

        acc += tl.sum(x_even[:, None] * (w_lo - z[None, :]) * sc[None, :], axis=0)
        acc += tl.sum(x_odd[:, None] * (w_hi - z[None, :]) * sc[None, :], axis=0)

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


# ============================================================================
# Fused KDA q+k+v+g kernel: 4 projections from one input
# ============================================================================

@triton.jit
def _int4_qkv_kernel(
    x_ptr,
    wq_q_ptr, sc_q_ptr, z_q_ptr,
    wq_k_ptr, sc_k_ptr, z_k_ptr,
    wq_v_ptr, sc_v_ptr, z_v_ptr,
    out_q_ptr, out_k_ptr, out_v_ptr,
    K,
    N,
    GROUP,
    stride_wq_n,
    stride_sc_n,
    stride_z_n,
    BLOCK_K: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    """Fused dequant-GEMV for 3 projections (q, k, v) from one input x."""
    pid = tl.program_id(0)
    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = offs_n < N

    acc_q = tl.zeros([BLOCK_N], dtype=tl.float32)
    acc_k = tl.zeros([BLOCK_N], dtype=tl.float32)
    acc_v = tl.zeros([BLOCK_N], dtype=tl.float32)

    offs_even = 2 * tl.arange(0, BLOCK_K // 2)
    offs_odd = 1 + 2 * tl.arange(0, BLOCK_K // 2)
    p_offs = tl.arange(0, BLOCK_K // 2)

    for k_start in range(0, K.to(tl.int32), BLOCK_K):
        x_even = tl.load(
            x_ptr + k_start + offs_even,
            mask=(k_start + offs_even) < K, other=0.0,
        ).to(tl.float32)
        x_odd = tl.load(
            x_ptr + k_start + offs_odd,
            mask=(k_start + offs_odd) < K, other=0.0,
        ).to(tl.float32)

        p_start = k_start // 2
        p_mask = (p_start + p_offs) < (K // 2)
        g_idx = k_start // GROUP

        # q projection
        wq = tl.load(
            wq_q_ptr + (p_start + p_offs)[:, None] * stride_wq_n + offs_n[None, :],
            mask=p_mask[:, None] & n_mask[None, :], other=0,
        )
        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)
        sc = tl.load(sc_q_ptr + g_idx * stride_sc_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_q_ptr + g_idx * stride_z_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        acc_q += tl.sum(x_even[:, None] * (w_lo - z[None, :]) * sc[None, :], axis=0)
        acc_q += tl.sum(x_odd[:, None] * (w_hi - z[None, :]) * sc[None, :], axis=0)

        # k projection
        wq = tl.load(
            wq_k_ptr + (p_start + p_offs)[:, None] * stride_wq_n + offs_n[None, :],
            mask=p_mask[:, None] & n_mask[None, :], other=0,
        )
        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)
        sc = tl.load(sc_k_ptr + g_idx * stride_sc_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_k_ptr + g_idx * stride_z_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        acc_k += tl.sum(x_even[:, None] * (w_lo - z[None, :]) * sc[None, :], axis=0)
        acc_k += tl.sum(x_odd[:, None] * (w_hi - z[None, :]) * sc[None, :], axis=0)

        # v projection
        wq = tl.load(
            wq_v_ptr + (p_start + p_offs)[:, None] * stride_wq_n + offs_n[None, :],
            mask=p_mask[:, None] & n_mask[None, :], other=0,
        )
        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)
        sc = tl.load(sc_v_ptr + g_idx * stride_sc_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_v_ptr + g_idx * stride_z_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        acc_v += tl.sum(x_even[:, None] * (w_lo - z[None, :]) * sc[None, :], axis=0)
        acc_v += tl.sum(x_odd[:, None] * (w_hi - z[None, :]) * sc[None, :], axis=0)

    tl.store(out_q_ptr + offs_n, acc_q.to(tl.bfloat16), mask=n_mask)
    tl.store(out_k_ptr + offs_n, acc_k.to(tl.bfloat16), mask=n_mask)
    tl.store(out_v_ptr + offs_n, acc_v.to(tl.bfloat16), mask=n_mask)


# ============================================================================
# Fused MoE gate+up kernel: one input, two outputs per expert
# ============================================================================

@triton.jit
def _int4_moe_gate_up_kernel(
    x_ptr,
    wq_g_ptr, sc_g_ptr, z_g_ptr,
    wq_u_ptr, sc_u_ptr, z_u_ptr,
    out_g_ptr, out_u_ptr,
    K,
    N,
    GROUP,
    E: tl.constexpr,
    stride_wq_e,
    stride_wq_n,
    stride_sc_e,
    stride_sc_n,
    stride_z_e,
    stride_z_n,
    stride_out_e,
    BLOCK_K: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    """Fused MoE gate+up: out_g[e,:] = x @ W_g[e], out_u[e,:] = x @ W_u[e]."""
    eid = tl.program_id(0)
    pid = tl.program_id(1)
    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = offs_n < N

    acc_g = tl.zeros([BLOCK_N], dtype=tl.float32)
    acc_u = tl.zeros([BLOCK_N], dtype=tl.float32)

    offs_even = 2 * tl.arange(0, BLOCK_K // 2)
    offs_odd = 1 + 2 * tl.arange(0, BLOCK_K // 2)
    p_offs = tl.arange(0, BLOCK_K // 2)

    for k_start in range(0, K.to(tl.int32), BLOCK_K):
        x_even = tl.load(
            x_ptr + k_start + offs_even,
            mask=(k_start + offs_even) < K, other=0.0,
        ).to(tl.float32)
        x_odd = tl.load(
            x_ptr + k_start + offs_odd,
            mask=(k_start + offs_odd) < K, other=0.0,
        ).to(tl.float32)

        p_start = k_start // 2
        p_mask = (p_start + p_offs) < (K // 2)
        g_idx = k_start // GROUP

        # Gate projection
        wq = tl.load(
            wq_g_ptr + eid * stride_wq_e
                   + (p_start + p_offs)[:, None] * stride_wq_n
                   + offs_n[None, :],
            mask=p_mask[:, None] & n_mask[None, :], other=0,
        )
        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)
        sc = tl.load(sc_g_ptr + eid * stride_sc_e + g_idx * stride_sc_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_g_ptr + eid * stride_z_e + g_idx * stride_z_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        acc_g += tl.sum(x_even[:, None] * (w_lo - z[None, :]) * sc[None, :], axis=0)
        acc_g += tl.sum(x_odd[:, None] * (w_hi - z[None, :]) * sc[None, :], axis=0)

        # Up projection
        wq = tl.load(
            wq_u_ptr + eid * stride_wq_e
                   + (p_start + p_offs)[:, None] * stride_wq_n
                   + offs_n[None, :],
            mask=p_mask[:, None] & n_mask[None, :], other=0,
        )
        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)
        sc = tl.load(sc_u_ptr + eid * stride_sc_e + g_idx * stride_sc_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        z = tl.load(z_u_ptr + eid * stride_z_e + g_idx * stride_z_n + offs_n, mask=n_mask, other=0.0).to(tl.float32)
        acc_u += tl.sum(x_even[:, None] * (w_lo - z[None, :]) * sc[None, :], axis=0)
        acc_u += tl.sum(x_odd[:, None] * (w_hi - z[None, :]) * sc[None, :], axis=0)

    tl.store(out_g_ptr + eid * stride_out_e + offs_n, acc_g.to(tl.bfloat16), mask=n_mask)
    tl.store(out_u_ptr + eid * stride_out_e + offs_n, acc_u.to(tl.bfloat16), mask=n_mask)


# ============================================================================
# Batched MoE int4 dequant-GEMV with per-expert inputs (for down projection)
# ============================================================================

@triton.jit
def _int4_moe_batch_gemv_kernel(
    x_ptr,            # [E, K] bf16 — per-expert input
    wq_ptr,           # [E, K//2, N] uint8
    scales_ptr,        # [E, K//G, N] bf16
    zeros_ptr,         # [E, K//G, N] bf16
    out_ptr,           # [E, N] bf16
    K,
    N,
    GROUP,
    stride_x_e,
    stride_wq_e,
    stride_wq_n,
    stride_sc_e,
    stride_sc_n,
    stride_z_e,
    stride_z_n,
    stride_out_e,
    BLOCK_K: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    """Batched per-expert dequant-GEMV: out[e,:] = x[e,:] @ dequant(W[e])."""
    eid = tl.program_id(0)
    pid = tl.program_id(1)
    offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    n_mask = offs_n < N

    acc = tl.zeros([BLOCK_N], dtype=tl.float32)
    offs_even = 2 * tl.arange(0, BLOCK_K // 2)
    offs_odd = 1 + 2 * tl.arange(0, BLOCK_K // 2)
    p_offs = tl.arange(0, BLOCK_K // 2)

    for k_start in range(0, K.to(tl.int32), BLOCK_K):
        x_even = tl.load(
            x_ptr + eid * stride_x_e + k_start + offs_even,
            mask=(k_start + offs_even) < K, other=0.0,
        ).to(tl.float32)
        x_odd = tl.load(
            x_ptr + eid * stride_x_e + k_start + offs_odd,
            mask=(k_start + offs_odd) < K, other=0.0,
        ).to(tl.float32)

        p_start = k_start // 2
        p_mask = (p_start + p_offs) < (K // 2)

        wq = tl.load(
            wq_ptr + eid * stride_wq_e
                   + (p_start + p_offs)[:, None] * stride_wq_n
                   + offs_n[None, :],
            mask=p_mask[:, None] & n_mask[None, :], other=0,
        )

        w_lo = (wq & 0xF).to(tl.float32)
        w_hi = ((wq >> 4) & 0xF).to(tl.float32)

        g_idx = k_start // GROUP
        sc = tl.load(
            scales_ptr + eid * stride_sc_e + g_idx * stride_sc_n + offs_n,
            mask=n_mask, other=0.0,
        ).to(tl.float32)
        z = tl.load(
            zeros_ptr + eid * stride_z_e + g_idx * stride_z_n + offs_n,
            mask=n_mask, other=0.0,
        ).to(tl.float32)

        acc += tl.sum(x_even[:, None] * (w_lo - z[None, :]) * sc[None, :], axis=0)
        acc += tl.sum(x_odd[:, None] * (w_hi - z[None, :]) * sc[None, :], axis=0)

    tl.store(out_ptr + eid * stride_out_e + offs_n, acc.to(tl.bfloat16), mask=n_mask)


# ============================================================================
# Python wrappers
# ============================================================================

def _int4_gemv(x, w_q, scales, zeros, out=None):
    """Fused int4 dequant-GEMV: x [K] bf16 -> [N] bf16."""
    K = w_q.shape[0] * 2
    N = w_q.shape[1]
    if out is None:
        out = torch.empty(N, dtype=torch.bfloat16, device=x.device)
    # Fixed config well-suited for Blackwell sm120
    BLOCK_N = 64
    BLOCK_K = 128
    grid = (triton.cdiv(N, BLOCK_N),)
    _int4_gemv_kernel[grid](
        x, w_q, scales, zeros, out,
        K, N, GROUP_SIZE,
        stride_wq_n=N, stride_sc_n=N, stride_z_n=N,
        BLOCK_K=BLOCK_K, BLOCK_N=BLOCK_N, num_warps=4,
    )
    return out


def _int4_qkv(x, q_linear, k_linear, v_linear):
    """Fused q+k+v: returns q, k, v each [N] bf16."""
    K = q_linear.w_q.shape[0] * 2
    N = q_linear.w_q.shape[1]
    out_q = torch.empty(N, dtype=torch.bfloat16, device=x.device)
    out_k = torch.empty(N, dtype=torch.bfloat16, device=x.device)
    out_v = torch.empty(N, dtype=torch.bfloat16, device=x.device)
    BLOCK_K, BLOCK_N = 128, 64
    grid = (triton.cdiv(N, BLOCK_N),)
    _int4_qkv_kernel[grid](
        x,
        q_linear.w_q, q_linear.scales, q_linear.zeros,
        k_linear.w_q, k_linear.scales, k_linear.zeros,
        v_linear.w_q, v_linear.scales, v_linear.zeros,
        out_q, out_k, out_v,
        K, N, GROUP_SIZE,
        stride_wq_n=N, stride_sc_n=N, stride_z_n=N,
        BLOCK_K=BLOCK_K, BLOCK_N=BLOCK_N, num_warps=4,
    )
    return out_q, out_k, out_v


def _int4_moe_gate_up(x, gate_experts, up_experts, idx):
    """Fused gate+up for selected experts. Returns g_out, u_out each [k, N] bf16."""
    if idx.numel() == 0:
        return (
            torch.empty(0, gate_experts.out_f, dtype=torch.bfloat16, device=x.device),
            torch.empty(0, up_experts.out_f, dtype=torch.bfloat16, device=x.device),
        )
    wq_g = gate_experts.w_q[idx]
    sc_g = gate_experts.scales[idx]
    z_g = gate_experts.zeros[idx]
    wq_u = up_experts.w_q[idx]
    sc_u = up_experts.scales[idx]
    z_u = up_experts.zeros[idx]

    E = idx.numel()
    K = wq_g.shape[1] * 2
    N = wq_g.shape[2]

    out_g = torch.empty(E, N, dtype=torch.bfloat16, device=x.device)
    out_u = torch.empty(E, N, dtype=torch.bfloat16, device=x.device)

    BLOCK_K, BLOCK_N = 128, 64
    grid = (E, triton.cdiv(N, BLOCK_N))
    _int4_moe_gate_up_kernel[grid](
        x,
        wq_g, sc_g, z_g,
        wq_u, sc_u, z_u,
        out_g, out_u,
        K, N, GROUP_SIZE,
        E=E,
        stride_wq_e=wq_g.stride(0), stride_wq_n=wq_g.stride(1),
        stride_sc_e=sc_g.stride(0), stride_sc_n=sc_g.stride(1),
        stride_z_e=z_g.stride(0), stride_z_n=z_g.stride(1),
        stride_out_e=N,
        BLOCK_K=BLOCK_K, BLOCK_N=BLOCK_N, num_warps=4,
    )
    return out_g, out_u


def _int4_moe_batch_gemv(x_batch, w_q, scales, zeros, out=None):
    """Batched per-expert GEMV: x_batch [E, K] -> [E, N]."""
    E, N = w_q.shape[0], w_q.shape[2]
    K = w_q.shape[1] * 2
    if out is None:
        out = torch.empty(E, N, dtype=torch.bfloat16, device=x_batch.device)
    BLOCK_K, BLOCK_N = 128, 64
    grid = (E, triton.cdiv(N, BLOCK_N))
    _int4_moe_batch_gemv_kernel[grid](
        x_batch, w_q, scales, zeros, out,
        K, N, GROUP_SIZE,
        stride_x_e=x_batch.stride(0),
        stride_wq_e=w_q.stride(0), stride_wq_n=w_q.stride(1),
        stride_sc_e=scales.stride(0), stride_sc_n=scales.stride(1),
        stride_z_e=zeros.stride(0), stride_z_n=zeros.stride(1),
        stride_out_e=N,
        BLOCK_K=BLOCK_K, BLOCK_N=BLOCK_N, num_warps=4,
    )
    return out


# ============================================================================
# Fused layers (reference-compatible buffer names)
# ============================================================================

class FusedQuantLinear(nn.Module):
    """W4A16 linear using fused int4 dequant-GEMV. Same buffers as R.QuantLinear."""

    def __init__(self, in_f: int, out_f: int, group: int = GROUP_SIZE):
        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 forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: [K] bf16 → [N] bf16 via fused dequant-GEMV."""
        return _int4_gemv(x, self.w_q, self.scales, self.zeros)

    def weight_bf(self) -> torch.Tensor:
        """One-time dequantized bf16 weight (for absorb init / GEMM batching)."""
        return dequant(self.w_q, self.scales, self.zeros, self.in_f, self.group)


class FusedQuantExperts(nn.Module):
    """E independent W4A16 expert weights. Fused dequant on forward.
    Same buffer names as R.QuantExperts."""

    def __init__(self, n: int, in_f: int, out_f: int, group: int = GROUP_SIZE):
        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 forward_batch(self, x_batch: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
        """GEMV with per-expert inputs. x_batch: [k, K] bf16, idx: [k]. Returns [k, N] bf16."""
        if idx.numel() == 0:
            return torch.empty(0, self.out_f, dtype=torch.bfloat16, device=x_batch.device)
        return _int4_moe_batch_gemv(
            x_batch, self.w_q[idx], self.scales[idx], self.zeros[idx]
        )


# ============================================================================
# KDA layer — uses fused q+k+v, separate g and o projections
# ============================================================================

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 = FusedQuantLinear(d, H * Dk, cfg.group)
        self.k_proj = FusedQuantLinear(d, H * Dk, cfg.group)
        self.v_proj = FusedQuantLinear(d, H * Dk, cfg.group)
        self.g_proj = FusedQuantLinear(d, H * Dk, cfg.group)
        self.beta_proj = nn.Linear(d, H, bias=False, dtype=cfg.dtype)  # tiny, bf16
        self.conv_w = nn.Parameter(torch.empty(3, H * Dk, cfg.short_conv, dtype=cfg.dtype))
        self.o_proj = FusedQuantLinear(H * Dk, d, cfg.group)
        self.scale = Dk ** -0.5

    def _short_conv(self, val, prev, idx):
        win = torch.cat([prev, val[None]], dim=0)
        w = self.conv_w[idx].transpose(0, 1)
        out = (win.float() * w.float()).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
        # Fused q+k+v projection (single kernel, 3 outputs)
        q, k, v = _int4_qkv(x, self.q_proj, self.k_proj, self.v_proj)
        q, st["cq"] = self._short_conv(q, st["cq"], 0)
        k, st["ck"] = self._short_conv(k, st["ck"], 1)
        v, st["cv"] = self._short_conv(v, st["cv"], 2)
        q = q.view(H, Dk) * self.scale
        k = k.view(H, Dk)
        v = v.view(H, Dk)
        # g and beta: float for softplus precision
        g = (-F.softplus(self.g_proj(x).float())).view(H, Dk)
        beta = torch.sigmoid(self.beta_proj(x).float())
        # State update in fp32 for numerical stability
        S = st["S"].float()
        S = S * g.exp()[:, :, None]
        kf = k.float()
        vf = v.float()
        pred = (S * kf[:, :, None]).sum(1)
        S = S + beta[:, None, None] * kf[:, :, None] * (vf - pred)[:, None, :]
        o = (S * q.float()[:, :, None]).sum(1)
        st["S"] = S.to(torch.bfloat16)
        return self.o_proj(o.reshape(H * Dk).to(torch.bfloat16))


# ============================================================================
# MLA layer — KV cache absorb + fused projections
# ============================================================================

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

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

        # q projection (bf16)
        q = self.q_proj(x).view(H, D_nope + D_rope)
        q_nope = q[:, :D_nope]
        q_rope = q[:, D_nope:]

        # kv_a: compressed latent + rope key (bf16)
        kv = self.kv_a(x)
        c_kv_new = kv[: cfg.kv_lora]
        k_rope_raw = kv[cfg.kv_lora:]

        # RoPE: need float precision for the rotation
        cos, sin = _rope_cossin(pos, D_rope, cfg.rope_theta, x.device)
        q_rope_f = _apply_rope(q_rope, cos, sin)
        k_rope = _apply_rope(k_rope_raw, cos, sin)

        # Update compressed latent cache (for state correctness checks)
        st["c_kv"] = torch.cat([st["c_kv"], c_kv_new[None]], 0)
        st["k_rope"] = torch.cat([st["k_rope"], k_rope[None]], 0)

        # --- KV cache absorb ---
        if "full_k_nope" not in st:
            w_kvb = self.kv_b.weight_bf()
            full_kv = torch.matmul(st["c_kv"], w_kvb)
            full_kv = full_kv.view(-1, H, D_nope + D_v)
            st["full_k_nope"] = full_kv[..., :D_nope]
            st["full_v"] = full_kv[..., D_nope:]
        else:
            new_kv = self.kv_b(c_kv_new).view(1, H, D_nope + D_v)
            st["full_k_nope"] = torch.cat([st["full_k_nope"], new_kv[..., :D_nope]], dim=0)
            st["full_v"] = torch.cat([st["full_v"], new_kv[..., D_nope:]], dim=0)

        # Attention: promote to fp32 for einsum/softmax, back to bf16 for output
        q_nope_f = q_nope.float()
        kn_f = st["full_k_nope"].float()
        q_rope_f2 = q_rope_f.float() if q_rope_f.dtype != torch.float32 else q_rope_f
        kr_f = st["k_rope"].float()
        scores = (
            torch.einsum("hd,lhd->lh", q_nope_f, kn_f)
            + torch.einsum("hd,ld->lh", q_rope_f2, kr_f)
        ) * self.scale
        p = torch.softmax(scores, dim=0).to(torch.bfloat16)
        v_bf = st["full_v"].to(torch.bfloat16)
        o = torch.einsum("lh,lhd->hd", p, v_bf)

        return self.o_proj(o.reshape(H * D_v))


# ============================================================================
# MoE layer — fused gate+up projection
# ============================================================================

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)   # tiny, bf16
        self.gate = FusedQuantExperts(E, d, m, cfg.group)
        self.up = FusedQuantExperts(E, d, m, cfg.group)
        self.down = FusedQuantExperts(E, m, d, cfg.group)
        self.s_gate = FusedQuantExperts(cfg.n_shared, d, m, cfg.group)
        self.s_up = FusedQuantExperts(cfg.n_shared, d, m, cfg.group)
        self.s_down = FusedQuantExperts(cfg.n_shared, m, d, cfg.group)

    def _expert_ffn(self, x, gate_exp, up_exp, down_exp, idx, weights):
        """Compute routed FFN with fused gate+up and batched down."""
        # Fused gate+up projection
        g_out, u_out = _int4_moe_gate_up(x, gate_exp, up_exp, idx)
        hh = F.silu(g_out.float()) * u_out.float()  # [k, m]

        # Batched down projection
        d_out = _int4_moe_batch_gemv(
            hh.to(torch.bfloat16),
            down_exp.w_q[idx], down_exp.scales[idx], down_exp.zeros[idx],
        ).float()  # [k, d]

        return (weights[:, None] * d_out).sum(0)  # [d]

    def step(self, x):
        cfg = self.cfg

        # Routing
        probs = torch.softmax(self.router(x).float(), dim=-1)
        w, idx = torch.topk(probs, cfg.n_active)
        w_scaled = (w / (w.sum() + 1e-9) * cfg.routed_scaling).float()

        # Routed experts (fused gate+up + batched down)
        routed = self._expert_ffn(x, self.gate, self.up, self.down, idx, w_scaled)

        # Shared expert (fused gate+up + batched down)
        sidx = torch.arange(cfg.n_shared, device=x.device)
        s_weights = torch.ones(cfg.n_shared, device=x.device, dtype=torch.float32)
        routed = routed + self._expert_ffn(x, self.s_gate, self.s_up, self.s_down, sidx, s_weights)

        return routed.to(torch.bfloat16)


# ============================================================================
# Block and 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)

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