training.py

"""
Addressed State Attention (ASA) - Training Harness

Efficient implementation optimized for language model training.
For mechanistic analysis and interventions, use asm_analysis.py instead.

Repository: https://github.com/DigitalDaimyo/AddressedStateAttention
Paper:  https://github.com/DigitalDaimyo/AddressedStateAttention/paper_drafts
"""

import math
from dataclasses import dataclass
from typing import Optional, Dict, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

__all__ = [
    'AddressedStateAttention',
    'ASMBlock',
    'ASMLanguageModel', 
    'ASMTrainConfig',
    'build_model_from_cfg',
]

# -------------------------
# RoPE helper (rotate-half)
# -------------------------
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)

class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, base: float = 10000.0):
        super().__init__()
        assert dim % 2 == 0, "RoPE requires even dim"
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._cos_cached = None
        self._sin_cached = None
        self._t_cached = None
        self._device_cached = None

    def get_cos_sin(self, T: int, device, dtype):
        if self._t_cached == T and self._cos_cached is not None and self._device_cached == device:
            return self._cos_cached.to(dtype=dtype), self._sin_cached.to(dtype=dtype)

        t = torch.arange(T, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum("t,f->tf", t, self.inv_freq)  # [T, d/2]
        emb = torch.cat([freqs, freqs], dim=-1)            # [T, d]
        cos = emb.cos()[None, None, :, :]                  # [1,1,T,d]
        sin = emb.sin()[None, None, :, :]                  # [1,1,T,d]

        self._t_cached = T
        self._device_cached = device
        self._cos_cached = cos
        self._sin_cached = sin
        return cos.to(dtype=dtype), sin.to(dtype=dtype)

def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    return (x * cos) + (_rotate_half(x) * sin)

# -------------------------
# ALiBi slopes helper
# -------------------------
def alibi_slopes(num_heads: int, device=None, dtype=torch.float32) -> torch.Tensor:
    def get_slopes(n):
        def power_of_2_slopes(n):
            start = 2.0 ** (-(2.0 ** -(math.log2(n) - 3)))
            ratio = start
            return [start * (ratio ** i) for i in range(n)]
        if math.log2(n).is_integer():
            return power_of_2_slopes(n)
        closest = 2 ** math.floor(math.log2(n))
        return power_of_2_slopes(closest) + get_slopes(2 * closest)[0::2][: n - closest]
    return torch.tensor(get_slopes(num_heads), device=device, dtype=dtype)

def _inv_softplus(y: torch.Tensor) -> torch.Tensor:
    return torch.log(torch.expm1(y))

class AddressedStateAttention(nn.Module):
    """
    ASA with integral slotspace refine fused into the compiled chunk kernel.
    Fixes included:
      (1) pad slotspace RoPE cos/sin to CH (identity on padded positions)
      (2) build valid_mask_c even when attention_mask is None (padding-only)
      (3) pad write logits with -inf (so padded positions contribute zero to scan)
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int = 12,
        num_slots: int = 16,
        dropout: float = 0.1,

        # temps / numerics
        read_temperature: float = 1.0,
        write_temperature: float = 1.0,
        state_fp32: bool = True,
        slot_dropout: float = 0.0,
        normalize_k: bool = False,

        # write geometry
        use_rope_keys: bool = True,
        rope_base: float = 10000.0,

        # write bias
        use_alibi_write: bool = True,
        alibi_strength_init: float = 0.1,
        learn_alibi_strength: bool = True,
        min_strength: float = 0.0,

        # content read gamma
        use_content_read: bool = True,
        content_read_init: float = -4.0,
        content_read_max_gamma: float = 3.0,

        # slotspace refine (INTEGRAL)
        use_slotspace_refine: bool = True,  # compat only
        slotspace_dim: int = 8,
        slotspace_gate_init: float = -4.0,
        slotspace_dropout: float = 0.05,
        slotspace_signed_weights: bool = True,

        # slotspace RoPE (Q/K only)
        use_rope_slotspace: bool = True,
        rope_base_slotspace: float = 100000.0,

        # perf
        write_chunk_size: int = 1024,
        enable_compiled: bool = True,
    ):
        super().__init__()
        assert embed_dim % num_heads == 0
        assert (slotspace_dim % 2) == 0, "slotspace_dim must be even if RoPE enabled"

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.num_slots = num_slots
        self.head_dim = embed_dim // num_heads

        self.dropout = nn.Dropout(dropout)

        self.read_temperature = float(read_temperature)
        self.write_temperature = float(write_temperature)
        self.state_fp32 = bool(state_fp32)
        self.slot_dropout = float(slot_dropout)
        self.normalize_k = bool(normalize_k)

        self.use_rope_keys = bool(use_rope_keys)
        self.use_alibi_write = bool(use_alibi_write)
        self.learn_alibi_strength = bool(learn_alibi_strength)
        self.min_strength = float(min_strength)

        self.use_content_read = bool(use_content_read)
        self.content_read_max_gamma = float(content_read_max_gamma)

        self.slotspace_dim = int(slotspace_dim)
        self.slotspace_dropout = nn.Dropout(float(slotspace_dropout))
        self.slotspace_signed_weights = bool(slotspace_signed_weights)

        self.use_rope_slotspace = bool(use_rope_slotspace)
        self.write_chunk_size = int(write_chunk_size)

        H, K, d = self.num_heads, self.num_slots, self.head_dim
        M = self.slotspace_dim

        self.slot_keys = nn.Parameter(torch.randn(H, K, d) / math.sqrt(d))

        self.Wk_write = nn.Linear(embed_dim, embed_dim, bias=False)
        self.Wv_write = nn.Linear(embed_dim, embed_dim, bias=False)
        self.Wq_read  = nn.Linear(embed_dim, embed_dim, bias=False)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)

        self.rope = RotaryEmbedding(d, base=rope_base) if self.use_rope_keys else None

        if self.use_alibi_write:
            self.register_buffer("_alibi_slopes", alibi_slopes(H), persistent=False)
        else:
            self.register_buffer("_alibi_slopes", torch.zeros(H), persistent=False)

        if self.use_alibi_write and self.learn_alibi_strength:
            init = torch.tensor(float(alibi_strength_init) - self.min_strength).clamp_min(1e-8)
            self._alibi_strength_param = nn.Parameter(_inv_softplus(init))
        else:
            self._alibi_strength_param = None
            self.alibi_strength = float(alibi_strength_init)

        if self.use_content_read:
            self._content_read_gamma_raw = nn.Parameter(torch.tensor(float(content_read_init)))
        else:
            self._content_read_gamma_raw = None

        self.slot_in  = nn.Linear(K, M, bias=False)
        self.slot_q   = nn.Linear(M, M, bias=False)
        self.slot_k   = nn.Linear(M, M, bias=False)
        self.slot_v   = nn.Linear(M, M, bias=False)
        self.slot_out = nn.Linear(M, K, bias=False)

        self._slotspace_gate_raw = nn.Parameter(torch.tensor(float(slotspace_gate_init)))

        self.rope_slotspace = RotaryEmbedding(M, base=float(rope_base_slotspace)) if self.use_rope_slotspace else None

        self._compiled = None
        if enable_compiled:
            self.enable_compiled_kernel()

    def enable_compiled_kernel(self):
        if self._compiled is None:
            self._compiled = torch.compile(self._asa_chunk_fused, dynamic=False, fullgraph=False)

    def _alibi_strength(self, dtype, device) -> torch.Tensor:
        if not (self.use_alibi_write and self.learn_alibi_strength):
            return torch.tensor(getattr(self, "alibi_strength", 0.0), dtype=dtype, device=device)
        return (F.softplus(self._alibi_strength_param) + self.min_strength).to(dtype=dtype, device=device)

    def _content_read_gamma(self, dtype, device) -> torch.Tensor:
        if not self.use_content_read:
            return torch.tensor(0.0, dtype=dtype, device=device)
        g = F.softplus(self._content_read_gamma_raw)
        if self.content_read_max_gamma is not None and self.content_read_max_gamma > 0:
            g = g.clamp(max=self.content_read_max_gamma)
        return g.to(dtype=dtype, device=device)

    def _slotspace_gate(self, dtype, device) -> torch.Tensor:
        return F.softplus(self._slotspace_gate_raw).to(dtype=dtype, device=device)

    @staticmethod
    def _safe_exp_sub_max(s: torch.Tensor, m: torch.Tensor) -> torch.Tensor:
        diff = s - m
        diff = diff.masked_fill(~torch.isfinite(m), float("-inf"))
        return torch.exp(diff)

    @staticmethod
    def _phi(x: torch.Tensor) -> torch.Tensor:
        return F.elu(x) + 1.0

    @staticmethod
    def _pad_time_slice(x: torch.Tensor, t0: int, L: int, CH: int, dim: int):
        sl = x.narrow(dim, t0, L)
        if L == CH:
            return sl, None
        pad_shape = list(sl.shape)
        pad_shape[dim] = CH - L
        pad = torch.zeros(pad_shape, device=sl.device, dtype=sl.dtype)
        xpad = torch.cat([sl, pad], dim=dim)
        mask = torch.zeros((CH,), device=sl.device, dtype=torch.bool)
        mask[:L] = True
        return xpad, mask

    def _asa_chunk_fused(
        self,
        wlog_c: torch.Tensor,         # [B,H,K,CH]
        v_c: torch.Tensor,            # [B,H,CH,d]
        q_c: torch.Tensor,            # [B,H,CH,d]
        slot_keys_dk: torch.Tensor,   # [1,H,d,K]
        pos_cos_s: Optional[torch.Tensor],  # [1,1,CH,M] or None
        pos_sin_s: Optional[torch.Tensor],  # [1,1,CH,M] or None
        content_gamma: torch.Tensor,
        rtemp_t: torch.Tensor,
        gate_t: torch.Tensor,
        m_state: torch.Tensor,        # [B,H,K]
        denom_state: torch.Tensor,    # [B,H,K]
        numer_state: torch.Tensor,    # [B,H,K,d]
        S_state: torch.Tensor,        # [B,H,M,M]
        Z_state: torch.Tensor,        # [B,H,M]
        valid_mask_c: Optional[torch.Tensor],  # [B,1,CH,1] or None
        do_dropout: bool,
        dropout_p: float,
        signed_slot_w: bool,
    ):
        B, H, K, CH = wlog_c.shape
        d = numer_state.shape[-1]
        M = S_state.shape[-1]
        inv_sqrt_d = 1.0 / math.sqrt(d)

        # ----- WRITE prefix-softmax scan -----
        m_c, _ = torch.cummax(wlog_c, dim=-1)              # [B,H,K,CH]
        m_new = torch.maximum(m_state.unsqueeze(-1), m_c)  # [B,H,K,CH]
        scale = torch.exp(m_state.unsqueeze(-1) - m_new)   # [B,H,K,CH]

        denom_c = denom_state.unsqueeze(-1) * scale                          # [B,H,K,CH]
        numer_c = numer_state.unsqueeze(-2) * scale.unsqueeze(-1)            # [B,H,K,CH,d]

        w_new = self._safe_exp_sub_max(wlog_c, m_new)                        # [B,H,K,CH]
        denom_c = denom_c + torch.cumsum(w_new, dim=-1)                      # [B,H,K,CH]
        numer_c = numer_c + torch.cumsum(w_new.unsqueeze(-1) * v_c.unsqueeze(2), dim=-2)  # [B,H,K,CH,d]

        # ----- Routing logits -----
        read_logits_key = torch.matmul(q_c, slot_keys_dk) * inv_sqrt_d        # [B,H,CH,K]

        if self.use_content_read:
            numer_for_dot = numer_c.to(q_c.dtype).permute(0, 1, 3, 2, 4)      # [B,H,CH,K,d]
            denom_for_div = denom_c.to(q_c.dtype).permute(0, 1, 3, 2)         # [B,H,CH,K]
            read_logits_content = (q_c.unsqueeze(-2) * numer_for_dot).sum(dim=-1) * inv_sqrt_d
            read_logits_content = read_logits_content / denom_for_div.clamp_min(1e-8)
            read_logits = read_logits_key + content_gamma.to(read_logits_key.dtype) * read_logits_content
        else:
            read_logits = read_logits_key

        read_w = torch.softmax(read_logits / rtemp_t, dim=-1)                # [B,H,CH,K]

        # ----- EXACT base output -----
        inv_denom = (1.0 / denom_c.clamp_min(1e-8)).to(numer_c.dtype)         # [B,H,K,CH]
        w_scaled = read_w.to(numer_c.dtype).permute(0, 1, 3, 2) * inv_denom   # [B,H,K,CH]
        out_base = (w_scaled.unsqueeze(-1) * numer_c).sum(dim=2)              # [B,H,CH,d]

        # ----- Slotspace refine -----
        u = self.slot_in(read_w.to(out_base.dtype))                           # [B,H,CH,M]
        q_s = self.slot_q(u)
        k_s = self.slot_k(u)
        v_s = self.slot_v(u)

        if self.use_rope_slotspace and (pos_cos_s is not None) and (pos_sin_s is not None):
            q_s = apply_rope(q_s, pos_cos_s, pos_sin_s)
            k_s = apply_rope(k_s, pos_cos_s, pos_sin_s)

        if valid_mask_c is not None:
            q_s = q_s * valid_mask_c
            k_s = k_s * valid_mask_c
            v_s = v_s * valid_mask_c

        qf = self._phi(q_s)
        kf = self._phi(k_s)

        kv = kf.unsqueeze(-1) * v_s.unsqueeze(-2)                             # [B,H,CH,M,M]
        S_c = torch.cumsum(kv, dim=2) + S_state.unsqueeze(2)                  # [B,H,CH,M,M]
        Z_c = torch.cumsum(kf, dim=2) + Z_state.unsqueeze(2)                  # [B,H,CH,M]
        Z_c = Z_c.clamp_min(1e-8)

        num = torch.matmul(qf.unsqueeze(-2), S_c).squeeze(-2)                 # [B,H,CH,M]
        den = (qf * Z_c).sum(dim=-1, keepdim=True).clamp_min(1e-8)            # [B,H,CH,1]
        u2 = num / den                                                        # [B,H,CH,M]

        S_state_new = S_c[:, :, -1, :, :]
        Z_state_new = Z_c[:, :, -1, :]

        if do_dropout and dropout_p > 0.0:
            keep = (torch.rand_like(u2) > dropout_p).to(u2.dtype) / (1.0 - dropout_p)
            u2 = u2 * keep

        slot_w = self.slot_out(u2)                                            # [B,H,CH,K]
        if signed_slot_w:
            slot_w = torch.tanh(slot_w)
        else:
            slot_w = torch.softmax(slot_w, dim=-1)

        slot_w_scaled = slot_w.to(numer_c.dtype).permute(0, 1, 3, 2) * inv_denom
        delta = (slot_w_scaled.unsqueeze(-1) * numer_c).sum(dim=2)             # [B,H,CH,d]

        out = out_base + gate_t.to(out_base.dtype) * delta

        m_state_new = m_new[:, :, :, -1]
        denom_state_new = denom_c[:, :, :, -1]
        numer_state_new = numer_c[:, :, :, -1, :]

        return out, read_w, m_state_new, denom_state_new, numer_state_new, S_state_new, Z_state_new

    def forward(
        self,
        x: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        return_info: bool = False,
        return_light_stats: bool = False,
    ) -> Tuple[torch.Tensor, Optional[Dict[str, torch.Tensor]]]:

        B, T, C = x.shape
        H, K, d = self.num_heads, self.num_slots, self.head_dim
        M = self.slotspace_dim

        k_write = self.Wk_write(x).reshape(B, T, H, d).transpose(1, 2)  # [B,H,T,d]
        v_write = self.Wv_write(x).reshape(B, T, H, d).transpose(1, 2)  # [B,H,T,d]
        q_read  = self.Wq_read(x).reshape(B, T, H, d).transpose(1, 2)   # [B,H,T,d]

        if self.normalize_k:
            k_write = F.normalize(k_write, dim=-1, eps=1e-8)

        if self.use_rope_keys:
            cos, sin = self.rope.get_cos_sin(T, device=x.device, dtype=k_write.dtype)
            k_write = apply_rope(k_write, cos, sin)

        slot_keys = self.slot_keys
        if self.training and self.slot_dropout > 0.0:
            drop = (torch.rand((H, K), device=x.device) < self.slot_dropout)
            slot_keys = slot_keys * (~drop).to(slot_keys.dtype).unsqueeze(-1)

        slot_keys_dk = slot_keys.transpose(-1, -2).unsqueeze(0).to(q_read.dtype)  # [1,H,d,K]

        write_logits_raw = torch.matmul(k_write.to(q_read.dtype), slot_keys_dk).permute(0, 1, 3, 2) / math.sqrt(d)

        state_dtype = torch.float32 if (self.state_fp32 and x.dtype != torch.float32) else x.dtype
        write_logits = write_logits_raw.to(state_dtype)

        wtemp = max(1e-6, self.write_temperature)
        write_logits = write_logits / wtemp

        if self.use_alibi_write:
            strength = self._alibi_strength(dtype=state_dtype, device=x.device)
            slopes = self._alibi_slopes.to(device=x.device, dtype=state_dtype) * strength
            pos = torch.arange(T, device=x.device, dtype=state_dtype)
            write_logits = write_logits + slopes.view(1, H, 1, 1) * pos.view(1, 1, 1, T)

        valid = None
        if attention_mask is not None:
            valid = attention_mask.to(dtype=torch.bool)
            write_logits = write_logits.masked_fill(~valid.view(B, 1, 1, T), float("-inf"))

        content_gamma = self._content_read_gamma(dtype=q_read.dtype, device=x.device)
        rtemp_t = torch.tensor(max(1e-6, self.read_temperature), device=x.device, dtype=q_read.dtype)
        gate_t = self._slotspace_gate(dtype=state_dtype, device=x.device)

        denom_state = torch.zeros((B, H, K), device=x.device, dtype=state_dtype)
        numer_state = torch.zeros((B, H, K, d), device=x.device, dtype=state_dtype)
        m_state = torch.full((B, H, K), float("-inf"), device=x.device, dtype=state_dtype)

        S_state = torch.zeros((B, H, M, M), device=x.device, dtype=state_dtype)
        Z_state = torch.zeros((B, H, M), device=x.device, dtype=state_dtype)

        out_h = torch.empty((B, H, T, d), device=x.device, dtype=state_dtype)

        if self.use_rope_slotspace:
            cos_s_full, sin_s_full = self.rope_slotspace.get_cos_sin(T, device=x.device, dtype=state_dtype)
        else:
            cos_s_full = sin_s_full = None

        CH = self.write_chunk_size
        kernel = self._compiled if self._compiled is not None else self._asa_chunk_fused

        do_dropout = bool(self.training and self.slotspace_dropout.p > 0.0)
        dropout_p = float(self.slotspace_dropout.p)
        signed_slot_w = bool(self.slotspace_signed_weights)

        for t0 in range(0, T, CH):
            t1 = min(T, t0 + CH)
            L = t1 - t0

            wlog_c, mask = self._pad_time_slice(write_logits, t0, L, CH, dim=3)              # [B,H,K,CH]
            v_c,   _     = self._pad_time_slice(v_write.to(state_dtype), t0, L, CH, dim=2)  # [B,H,CH,d]
            q_c,   _     = self._pad_time_slice(q_read, t0, L, CH, dim=2)                   # [B,H,CH,d]

            # (3) ensure padded write logits contribute zero mass
            if mask is not None:
                wlog_c = wlog_c.clone()
                wlog_c[:, :, :, L:] = float("-inf")

            # (2) build valid_mask_c even when attention_mask is None (padding-only)
            valid_mask_c = None
            if (valid is not None) or (mask is not None):
                if valid is None:
                    vm_pad = mask.view(1, CH).expand(B, CH)  # [B,CH]
                else:
                    if mask is None:
                        vm_pad = valid[:, t0:t1]
                    else:
                        vm = valid[:, t0:t1]
                        vm_pad = torch.zeros((B, CH), device=x.device, dtype=torch.bool)
                        vm_pad[:, :L] = vm
                valid_mask_c = vm_pad.view(B, 1, CH, 1).to(state_dtype)

            # (1) slotspace RoPE slice PADDED TO CH (identity on padded positions)
            if self.use_rope_slotspace:
                cos_slice = cos_s_full[:, :, t0:t1, :]  # [1,1,L,M]
                sin_slice = sin_s_full[:, :, t0:t1, :]  # [1,1,L,M]
                if L == CH:
                    cos_s, sin_s = cos_slice, sin_slice
                else:
                    cos_s = torch.ones((1, 1, CH, M), device=x.device, dtype=state_dtype)
                    sin_s = torch.zeros((1, 1, CH, M), device=x.device, dtype=state_dtype)
                    cos_s[:, :, :L, :] = cos_slice
                    sin_s[:, :, :L, :] = sin_slice
            else:
                cos_s = sin_s = None

            out_c, read_w_c, m_state, denom_state, numer_state, S_state, Z_state = kernel(
                wlog_c, v_c, q_c, slot_keys_dk,
                cos_s, sin_s,
                content_gamma, rtemp_t, gate_t,
                m_state, denom_state, numer_state,
                S_state, Z_state,
                valid_mask_c,
                do_dropout, dropout_p,
                signed_slot_w,
            )

            if mask is not None:
                out_c = out_c * mask.view(1, 1, CH, 1).to(out_c.dtype)

            out_h[:, :, t0:t1, :] = out_c[:, :, :L, :]

        out = out_h.transpose(1, 2).reshape(B, T, C)
        out = self.out_proj(out)
        out = self.dropout(out)

        info = None
        if return_info or return_light_stats:
            info = {
                "content_read_gamma": content_gamma.detach().to(torch.float32).cpu(),
                "slotspace_gate": gate_t.detach().to(torch.float32).cpu(),
            }
        return out, info


# ============================================================================
# Addressed State Models (ASM): Config + Block + LM
# - Naming aligned with paper: slots, read/write, slot-space refinement
# - No compatibility layer (fresh public tooling)
# ============================================================================


# ============================================================================
# Config
# ============================================================================
@dataclass
class ASMTrainConfig:
    # Data
    dataset_name: str = "wikitext"
    dataset_config: str = "wikitext-103-raw-v1"
    tokenizer_name: str = "gpt2"

    max_seq_len: int = 256
    stride_frac_val: float = 0.50
    seed: int = 1337
    micro_batch_size: int = 2
    grad_accum_steps: int = 8
    # Sample budgets
    train_samples_target: int = 100_000_000
    val_samples_target: int = 25_000

    # Training
    batch_size: int = 64
    learning_rate: float = 3e-4
    weight_decay: float = 0.01
    betas: Tuple[float, float] = (0.9, 0.95)
    grad_clip: float = 1.0
    warmup_steps: int = 1_000
    total_steps: int = 75_000
    eval_interval: int = 1_000
    log_interval: int = 100

    # Model
    vocab_size: int = 50257
    embed_dim: int = 384
    num_layers: int = 23
    num_heads: int = 8
    num_slots: int = 32
    mlp_ratio: float = 4.0
    dropout: float = 0.1
    tie_weights: bool = True

    # Addressed State Attention (ASA) / numerics
    read_temperature: float = 1.0
    write_temperature: float = 1.0
    slot_dropout: float = 0.05
    state_fp32: bool = True
    normalize_k: bool = False

    # Positions
    use_abs_pos: bool = False
    use_rope_keys: bool = True
    rope_base: float = 10000.0
    use_alibi_write: bool = True
    alibi_strength_init: float = 0.1
    learn_alibi_strength: bool = True
    min_strength: float = 0.0

    # Content-conditioned read term (gamma)
    use_content_read: bool = True
    content_read_init: float = -4.0
    content_read_max_gamma: float = 3.0

    # Optional slot-space refinement (formerly "k-space")
    use_slotspace_refine: bool = True
    slotspace_dim: int = 64
    slotspace_gate_init: float = -4.0
    slotspace_dropout: float = 0.05
    slotspace_signed_weights: bool = True

    # RoPE inside slot-space matcher (Q/K only)
    use_rope_slotspace: bool = True
    rope_base_slotspace: float = 100000.0

    # Perf knobs (behavior-identical)
    write_chunk_size: int = 128
    enable_compiled: bool = True

    # Analytics
    eval_max_batches: int = 150
    analytics_last_k: int = 4

    # IO / caches
    output_dir: str = "./drive/MyDrive/asm_outputs"
    tag: str = "asm_wikitext"
    cache_dir: str = "./drive/MyDrive/asm_caches/fineweb/1B"
    val_windows_cache: str = "./drive/MyDrive/asm_val_cache_windows_1024.pkl"


# ============================================================================
# Block
# ============================================================================
class ASMBlock(nn.Module):
    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        num_slots: int,
        mlp_ratio: float = 4.0,
        dropout: float = 0.1,

        # temperatures / numerics
        read_temperature: float = 1.0,
        write_temperature: float = 1.0,
        state_fp32: bool = True,
        slot_dropout: float = 0.0,
        normalize_k: bool = False,

        # positions
        use_rope_keys: bool = True,
        rope_base: float = 10000.0,
        use_alibi_write: bool = True,

        # ALiBi params
        alibi_strength_init: float = 0.1,
        learn_alibi_strength: bool = True,
        min_strength: float = 0.0,

        # content-conditioned read (gamma)
        use_content_read: bool = True,
        content_read_init: float = -4.0,
        content_read_max_gamma: float = 3.0,

        # optional slot-space refinement
        use_slotspace_refine: bool = True,
        slotspace_dim: int = 32,
        slotspace_gate_init: float = -10.0,
        slotspace_dropout: float = 0.0,
        slotspace_signed_weights: bool = True,

        # RoPE inside slot-space matcher
        use_rope_slotspace: bool = True,
        rope_base_slotspace: float = 100000.0,

        # chunk sizes
        write_chunk_size: int = 128,
        enable_compiled: bool = False,
    ):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)

        self.asa = AddressedStateAttention(
            embed_dim=embed_dim,
            num_heads=num_heads,
            num_slots=num_slots,
            dropout=dropout,

            read_temperature=read_temperature,
            write_temperature=write_temperature,
            state_fp32=state_fp32,
            slot_dropout=slot_dropout,
            normalize_k=normalize_k,

            use_rope_keys=use_rope_keys,
            rope_base=rope_base,
            use_alibi_write=use_alibi_write,
            alibi_strength_init=alibi_strength_init,
            learn_alibi_strength=learn_alibi_strength,
            min_strength=min_strength,

            use_content_read=use_content_read,
            content_read_init=content_read_init,
            content_read_max_gamma=content_read_max_gamma,

            use_slotspace_refine=use_slotspace_refine,
            slotspace_dim=slotspace_dim,
            slotspace_gate_init=slotspace_gate_init,
            slotspace_dropout=slotspace_dropout,
            slotspace_signed_weights=slotspace_signed_weights,

            use_rope_slotspace=use_rope_slotspace,
            rope_base_slotspace=rope_base_slotspace,

            write_chunk_size=write_chunk_size,
            enable_compiled=enable_compiled,

        )

        self.norm2 = nn.LayerNorm(embed_dim)
        hidden = int(embed_dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, hidden, bias=False),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, embed_dim, bias=False),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, return_info: bool = False, return_light_stats: Optional[bool] = None):
        a, info = self.asa(self.norm1(x), attention_mask=attention_mask, return_info=return_info, return_light_stats=return_light_stats)
        x = x + a
        x = x + self.mlp(self.norm2(x))
        return x, info


# ============================================================================
# LM
# ============================================================================
class ASMLanguageModel(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        embed_dim: int = 384,
        num_layers: int = 6,
        num_heads: int = 8,
        num_slots: int = 8,
        max_seq_len: int = 1024,
        mlp_ratio: float = 4.0,
        dropout: float = 0.1,

        # temperatures / numerics
        read_temperature: float = 1.0,
        write_temperature: float = 1.0,
        state_fp32: bool = True,
        slot_dropout: float = 0.05,
        normalize_k: bool = False,

        tie_weights: bool = True,

        # LM-level abs pos
        use_abs_pos: bool = False,

        # positions
        use_rope_keys: bool = True,
        rope_base: float = 10000.0,
        use_alibi_write: bool = True,

        # ALiBi
        alibi_strength_init: float = 0.1,
        learn_alibi_strength: bool = True,
        min_strength: float = 0.0,

        # content-conditioned read (gamma)
        use_content_read: bool = True,
        content_read_init: float = -4.0,
        content_read_max_gamma: float = 3.0,

        # optional slot-space refinement
        use_slotspace_refine: bool = True,
        slotspace_dim: int = 32,
        slotspace_gate_init: float = -10.0,
        slotspace_dropout: float = 0.0,
        slotspace_signed_weights: bool = True,

        # RoPE inside slot-space matcher
        use_rope_slotspace: bool = True,
        rope_base_slotspace: float = 100000.0,

        # chunk sizes
        write_chunk_size: int = 128,
        enable_compiled: bool = False,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.max_seq_len = max_seq_len
        self.use_abs_pos = bool(use_abs_pos)

        self.tok = nn.Embedding(vocab_size, embed_dim)
        self.pos = nn.Embedding(max_seq_len, embed_dim) if self.use_abs_pos else None
        self.drop = nn.Dropout(dropout)

        self.blocks = nn.ModuleList([
            ASMBlock(
                embed_dim=embed_dim,
                num_heads=num_heads,
                num_slots=num_slots,
                mlp_ratio=mlp_ratio,
                dropout=dropout,

                read_temperature=read_temperature,
                write_temperature=write_temperature,
                state_fp32=state_fp32,
                slot_dropout=slot_dropout,
                normalize_k=normalize_k,

                use_rope_keys=use_rope_keys,
                rope_base=rope_base,
                use_alibi_write=use_alibi_write,

                alibi_strength_init=alibi_strength_init,
                learn_alibi_strength=learn_alibi_strength,
                min_strength=min_strength,

                use_content_read=use_content_read,
                content_read_init=content_read_init,
                content_read_max_gamma=content_read_max_gamma,

                use_slotspace_refine=use_slotspace_refine,
                slotspace_dim=slotspace_dim,            slotspace_gate_init=slotspace_gate_init,
                slotspace_dropout=slotspace_dropout,
                slotspace_signed_weights=slotspace_signed_weights,
                use_rope_slotspace=use_rope_slotspace,
                rope_base_slotspace=rope_base_slotspace,

                write_chunk_size=write_chunk_size,
                enable_compiled=enable_compiled,
            )
            for _ in range(num_layers)
        ])

        self.norm = nn.LayerNorm(embed_dim)
        self.lm_head = nn.Linear(embed_dim, vocab_size, bias=False)
        if tie_weights:
            self.lm_head.weight = self.tok.weight

        self.apply(self._init)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.LayerNorm):
            nn.init.ones_(m.weight)
            nn.init.zeros_(m.bias)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        return_info: bool = False,
        return_light_stats: Optional[bool] = None,
    ):
        B, T = input_ids.shape
        assert T <= self.max_seq_len, f"T={T} exceeds max_seq_len={self.max_seq_len}"

        x = self.tok(input_ids)
        if self.use_abs_pos:
            pos = torch.arange(T, device=input_ids.device).unsqueeze(0).expand(B, -1)
            x = x + self.pos(pos)

        x = self.drop(x)

        infos = []
        for blk in self.blocks:
            x, info = blk(x, attention_mask=attention_mask, return_info=return_info, return_light_stats=return_light_stats)
            if return_info:
                infos.append(info)

        x = self.norm(x)
        logits = self.lm_head(x)
        return (logits, infos) if return_info else logits


# ============================================================================
# Convenience: build model from config
# ============================================================================
def build_model_from_cfg(cfg: ASMTrainConfig) -> ASMLanguageModel:
    return ASMLanguageModel(
        vocab_size=cfg.vocab_size,
        embed_dim=cfg.embed_dim,
        num_layers=cfg.num_layers,
        num_heads=cfg.num_heads,
        num_slots=cfg.num_slots,
        max_seq_len=cfg.max_seq_len,
        mlp_ratio=cfg.mlp_ratio,
        dropout=cfg.dropout,

        read_temperature=cfg.read_temperature,
        write_temperature=cfg.write_temperature,
        state_fp32=cfg.state_fp32,
        slot_dropout=cfg.slot_dropout,
        normalize_k=cfg.normalize_k,

        tie_weights=cfg.tie_weights,

        use_abs_pos=cfg.use_abs_pos,
        use_rope_keys=cfg.use_rope_keys,
        rope_base=cfg.rope_base,
        use_alibi_write=cfg.use_alibi_write,

        alibi_strength_init=cfg.alibi_strength_init,
        learn_alibi_strength=cfg.learn_alibi_strength,
        min_strength=cfg.min_strength,

        use_content_read=cfg.use_content_read,
        content_read_init=cfg.content_read_init,
        content_read_max_gamma=cfg.content_read_max_gamma,

        use_slotspace_refine=cfg.use_slotspace_refine,
        slotspace_dim=cfg.slotspace_dim,
        slotspace_gate_init=cfg.slotspace_gate_init,
        slotspace_dropout=cfg.slotspace_dropout,
        slotspace_signed_weights=cfg.slotspace_signed_weights,
        use_rope_slotspace=cfg.use_rope_slotspace,
        rope_base_slotspace=cfg.rope_base_slotspace,

        write_chunk_size=cfg.write_chunk_size,
        enable_compiled=cfg.enable_compiled,
    )