"""Matched backbone with native FLA GLA/MLP/norm and budgeted softmax.

The block follows fla/models/{transformer,gla}/modeling_*.py at the revision
in sources.json. GLA kernels and layer internals are imported unchanged.
Our SDPA/FlexAttention adapter replaces only the softmax attention backend.
"""
from __future__ import annotations

import hashlib
import os
from dataclasses import asdict, dataclass
from pathlib import Path

os.environ.setdefault('TORCHINDUCTOR_CACHE_DIR', str(Path.home() / '.cache/torchinductor'))
os.environ.setdefault('TRITON_CACHE_DIR', str(Path.home() / '.triton/cache'))

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention.flex_attention import create_block_mask, flex_attention

from fla.layers.gla import GatedLinearAttention
from fla.modules import FusedLinearCrossEntropyLoss, GatedMLP, RMSNorm, RotaryEmbedding
from fla.modules.rotary.ops import rotary_embedding


@dataclass
class ModelConfig:
    variant: str = "dense"
    layers: int = 24
    hidden: int = 1024
    intermediate: int = 2816
    vocab: int = 32000
    heads: int = 16
    gla_heads: int = 4
    context: int = 2048
    dense_ends: int = 2
    seed: int = 1
    backend: str = "flex"
    recent_block: int = 0
    history_policy: str = "prefix"
    minimum_history: int = 1
    mamba_outer: bool = False
    mamba_d_state: int = 128
    mamba_expand: int = 2
    mamba_headdim: int = 64
    mamba_d_conv: int = 4
    mamba_state_dtype: str = "float32"

    def mamba_layer(self, i):
        return self.variant == 'mamba2' or (self.mamba_outer and
                (i < self.dense_ends or i >= self.layers - self.dense_ends))

    def dense_layer(self, i):
        if self.mamba_outer or self.variant == 'mamba2':
            return False
        return self.variant == "dense" or (
            self.variant in ("hybrid_gla", "taildrop")
            and (i < self.dense_ends or i >= self.layers - self.dense_ends)
        )

    def to_dict(self):
        result = asdict(self)
        # Preserve every old min1 checkpoint/config field verbatim.
        if self.minimum_history == 1:
            result.pop('minimum_history')
        # Old checkpoint configurations and numerical paths remain unchanged.
        if not self.mamba_outer and self.variant != 'mamba2':
            for name in tuple(result):
                if name.startswith('mamba_'):
                    result.pop(name)
        return result


def sample_prefixes(batch, length, *, generator, device):
    """Each (sequence, query) uniformly chooses from unique {1,2,4,...,N}."""
    n = torch.arange(1, length + 1, device=device, dtype=torch.int64)
    # Floating log2 can round exact powers of two downward on CUDA.
    # Integer bucket boundaries preserve every candidate, including full N.
    powers = torch.tensor([1 << p for p in range(length.bit_length())], device=device)
    k = torch.bucketize(n, powers, right=True) - 1
    count = k + 1 + ((n & (n - 1)) != 0).long()
    choice = (torch.rand((batch, length), device=device, generator=generator) * count).long()
    return torch.where(choice <= k, 1 << choice, n)


compiled_flex = torch.compile(flex_attention, dynamic=False)
compiled_block_mask = torch.compile(create_block_mask, dynamic=False)


def prefix_block_mask(lengths, kv_length, recent_block=0):
    batch, q_length = lengths.shape

    def visible(b, h, q, k):
        # create_block_mask may inspect padding beyond the actual query extent.
        prefix = k < lengths[b, q.clamp(max=q_length - 1)]
        if recent_block:
            start = (q // recent_block - 1) * recent_block
            return (q < q_length) & (k <= q) & (prefix | (k >= start))
        return (q < q_length) & prefix

    return compiled_block_mask(visible, batch, None, q_length, kv_length,
                               device=lengths.device, BLOCK_SIZE=128)


def suffix_block_mask(lengths, kv_length):
    """Latest K visible positions, including the current token within K."""
    batch, q_length = lengths.shape

    def visible(b, h, q, k):
        retained = lengths[b, q.clamp(max=q_length - 1)]
        return ((q < q_length) & (k < kv_length) & (k <= q)
                & (k >= q - retained + 1))

    return compiled_block_mask(visible, batch, None, q_length, kv_length,
                               device=lengths.device, BLOCK_SIZE=128)


class ResearchCache:
    """The small cache interface used by unchanged native GLA layers.

    Softmax stores the selected history in compact tensors. Suffix caches use
    bounded rings; their physical slot order need not be chronological. Seen
    token counts are separate, so RoPE positions never become cache positions.
    """
    def __init__(self, layers, max_length=None):
        if max_length is not None and max_length < 1:
            raise ValueError('Static cache capacity must be positive')
        self.states = [None] * layers
        self.seen = [0] * layers
        self.max_length = max_length

    def __len__(self):
        return len(self.states)

    def __getitem__(self, layer_idx):
        return self.states[layer_idx]

    def get_seq_length(self, layer_idx=0):
        return self.seen[layer_idx]

    def update(self, *, layer_idx, offset=1, **states):
        old = self.states[layer_idx] or {}
        old.update(states)
        self.states[layer_idx] = old
        self.seen[layer_idx] += offset
        return old

    def tensor_bytes(self):
        allocations = {}

        def visit(x):
            if isinstance(x, torch.Tensor):
                storage = x.untyped_storage()
                allocations[(str(x.device), storage.data_ptr())] = storage.nbytes()
            elif isinstance(x, dict):
                for v in x.values():
                    visit(v)
            elif isinstance(x, (list, tuple)):
                for v in x:
                    visit(v)

        visit(self.states)
        return sum(allocations.values())

    def metadata_bytes(self):
        """Host Python object sizes, deduplicated; GPU storage is counted above."""
        import sys
        seen = set()

        def size(x):
            if id(x) in seen:
                return 0
            seen.add(id(x))
            total = sys.getsizeof(x)
            if isinstance(x, dict):
                total += sum(size(k) + size(v) for k, v in x.items())
            elif isinstance(x, (list, tuple)):
                total += sum(size(v) for v in x)
            return total

        return sys.getsizeof(self) + size(self.__dict__)


class SoftmaxAttention(nn.Module):
    def __init__(self, cfg, layer_idx, taildrop):
        super().__init__()
        self.layer_idx = layer_idx
        self.heads, self.dim = cfg.heads, cfg.hidden // cfg.heads
        self.taildrop, self.backend = taildrop, cfg.backend
        self.context = cfg.context
        self.recent_block = cfg.recent_block if taildrop else 0
        self.history_policy = cfg.history_policy if taildrop else "prefix"
        self.q_proj = nn.Linear(cfg.hidden, cfg.hidden, bias=False)
        self.k_proj = nn.Linear(cfg.hidden, cfg.hidden, bias=False)
        self.v_proj = nn.Linear(cfg.hidden, cfg.hidden, bias=False)
        self.o_proj = nn.Linear(cfg.hidden, cfg.hidden, bias=False)
        self.rotary = RotaryEmbedding(self.dim, base=10000.0)

    def forward(self, x, *, lengths=None, block_mask=None, cap=None, cache=None):
        b, t, _ = x.shape
        if cache is not None and self.history_policy == "suffix" and cap is not None:
            if lengths is not None:
                raise ValueError('A persistent cache requires a fixed suffix policy')
            return self.suffix_cached_forward(x, cap, cache)
        if cache is not None and self.recent_block and cap is not None:
            if lengths is not None:
                raise ValueError('A persistent cache requires a fixed prefix policy')
            return self.protected_cached_forward(x, cap, cache)
        # Once a chronological prefix is full, no later token contributes KV.
        # Keep every query at its original position and reuse the complete KV
        # prefix. The uncached/training path below remains unchanged.
        if (cache is not None and self.taildrop and self.history_policy == "prefix"
                and cap is not None and lengths is None):
            offset = cache.get_seq_length(self.layer_idx)
            state = cache.states[self.layer_idx]
            if state is not None and offset >= cap:
                if self.training:
                    raise ValueError("Persistent cache is for inference only")
                if (state.get('policy_cap') != cap
                        or state.get('policy_history', 'prefix') != 'prefix'):
                    raise ValueError('A populated prefix cache cannot change its retained-history policy')
                if cache.max_length is not None and offset + t > cache.max_length:
                    raise ValueError('Decode would exceed the declared cache capacity')
                q = self.q_proj(x).view(b, t, self.heads, self.dim)
                rotary = self.rotary
                rotary._update_cos_sin_cache(max(self.context, offset + t), device=q.device, dtype=q.dtype)
                q = rotary_embedding(x=q, cos=rotary._cos_cached, sin=rotary._sin_cached,
                                     interleaved=rotary.interleaved, seqlen_offsets=offset).transpose(1, 2)
                k, v = state['k'], state['v']
                cache.update(layer_idx=self.layer_idx, offset=t, k=k, v=v, policy_cap=cap)
                # All prefix keys precede all queries in this branch, including
                # when a caller decodes several tokens in one call.
                out = F.scaled_dot_product_attention(q, k, v)
                return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))
        q, k, v = [p(x).view(b, t, self.heads, self.dim)
                   for p in (self.q_proj, self.k_proj, self.v_proj)]
        offset = 0 if cache is None else cache.get_seq_length(self.layer_idx)
        q, k = self.rotary(q, k, seqlen_offset=offset,
                           max_seqlen=max(self.context, offset + t))
        q, k, v = (z.transpose(1, 2) for z in (q, k, v))
        cap = cap if self.taildrop else None
        if cache is not None:
            if self.training:
                raise ValueError("Persistent cache is for inference only")
            state = cache.states[self.layer_idx]
            if state is not None and (state.get('policy_cap') != cap
                                      or state.get('policy_history', 'prefix') != 'prefix'):
                raise ValueError('A populated prefix cache cannot change its retained-history policy')
            keep = t if cap is None else min(t, max(0, cap - offset))
            if cache.max_length is not None:
                if offset + t > cache.max_length:
                    raise ValueError('Decode would exceed the declared cache capacity')
                capacity = cache.max_length if cap is None else min(cap, cache.max_length)
                if state is None:
                    shape = (b, self.heads, capacity, self.dim)
                    kb, vb = k.new_empty(shape), v.new_empty(shape)
                else:
                    kb, vb = state['k_buffer'], state['v_buffer']
                if keep:
                    kb[:, :, offset:offset+keep].copy_(k[:, :, :keep])
                    vb[:, :, offset:offset+keep].copy_(v[:, :, :keep])
                valid = min(offset + t, capacity)
                k, v = kb[:, :, :valid], vb[:, :, :valid]
                # Count full allocated capacity, including unused slots; never
                # treat a view's numel as its physical persistent storage.
                cache.update(layer_idx=self.layer_idx, offset=t, k=k, v=v,
                             k_buffer=kb, v_buffer=vb, policy_cap=cap)
            else:
                # clone ensures a prefix view cannot retain discarded backing storage.
                new_k, new_v = k[:, :, :keep].clone(), v[:, :, :keep].clone()
                if state is not None:
                    if keep:
                        k = torch.cat((state['k'], new_k), dim=2)
                        v = torch.cat((state['v'], new_v), dim=2)
                    else:
                        k, v = state['k'], state['v']
                else:
                    k, v = new_k, new_v
                cache.update(layer_idx=self.layer_idx, offset=t, k=k, v=v, policy_cap=cap)
        elif cap is not None and not self.recent_block and self.history_policy == "prefix":
            k, v = k[:, :, :cap], v[:, :, :cap]

        if self.taildrop and lengths is not None:
            if self.backend == "flex":
                out = compiled_flex(q, k, v, block_mask=block_mask)
            else:
                ki = torch.arange(k.shape[2], device=x.device)[None, None, :]
                if self.history_policy == "suffix":
                    qi = torch.arange(t, device=x.device)[None, :, None]
                    mask = (ki <= qi) & (ki >= qi - lengths[:, :, None] + 1)
                else:
                    mask = ki < lengths[:, :, None]
                if self.recent_block:
                    qi = torch.arange(t, device=x.device)[None, :, None]
                    start = (qi // self.recent_block - 1) * self.recent_block
                    mask = (ki <= qi) & (mask | (ki >= start))
                out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask[:, None])
        elif self.history_policy == "suffix" and cap is not None:
            # Prompt queries obey the same rolling window as continuation.
            # Keeping a dense prefill here would change the inference policy.
            if cap >= t:
                out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
            else:
                qi = torch.arange(t, device=x.device)[:, None]
                ki = torch.arange(t, device=x.device)[None, :]
                mask = (ki <= qi) & (ki >= qi - cap + 1)
                out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        elif self.recent_block and cap is not None:
            qi = torch.arange(t, device=x.device)[:, None]
            ki = torch.arange(k.shape[2], device=x.device)[None, :]
            start = (qi // self.recent_block - 1) * self.recent_block
            mask = (ki <= qi) & ((ki < cap) | (ki >= start))
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        elif offset == 0 and k.shape[2] == t:
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        elif t == 1:
            out = F.scaled_dot_product_attention(q, k, v)
        else:
            qi = offset + torch.arange(t, device=x.device)
            ki = torch.arange(k.shape[2], device=x.device)
            out = F.scaled_dot_product_attention(q, k, v, attn_mask=ki[None, :] <= qi[:, None])
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))

    def suffix_cached_forward(self, x, cap, cache):
        """Exact rolling suffix with a physically bounded C-slot KV ring.

        All current K/V projections are computed. A one-token call overwrites
        its expired slot before unmasked attention over valid ring slots;
        attention is invariant to their physical permutation. Chunked calls
        reconstruct chronological old keys temporarily, apply the per-query
        causal suffix mask, and then write only the final suffix to the ring.
        RoPE always uses absolute stream positions, never ring indices.
        """
        if self.training:
            raise ValueError('Persistent cache is for inference only')
        if cap < 1:
            raise ValueError('A suffix must retain at least one token')
        b, t, _ = x.shape
        if t < 1:
            raise ValueError('Cached input must contain at least one token')
        offset = cache.get_seq_length(self.layer_idx)
        end = offset + t
        if cache.max_length is not None and end > cache.max_length:
            raise ValueError('Decode would exceed the declared cache capacity')
        state = cache.states[self.layer_idx]
        if state is not None and (state.get('policy_cap') != cap
                                  or state.get('policy_history') != 'suffix'):
            raise ValueError('A populated suffix cache cannot change its retained-history policy')
        if state is None and offset:
            raise ValueError('Suffix cache state is missing for previously seen tokens')

        q, k, v = [p(x).view(b, t, self.heads, self.dim)
                   for p in (self.q_proj, self.k_proj, self.v_proj)]
        q, k = self.rotary(q, k, seqlen_offset=offset,
                           max_seqlen=max(self.context, end))
        q, k, v = (z.transpose(1, 2) for z in (q, k, v))
        capacity = cap if cache.max_length is None else min(cap, cache.max_length)
        if state is None:
            shape = (b, self.heads, capacity, self.dim)
            kb, vb = k.new_empty(shape), v.new_empty(shape)
        else:
            kb, vb = state['k_buffer'], state['v_buffer']
        valid = min(end, capacity)

        if t == 1:
            slot = offset % capacity
            kb[:, :, slot:slot+1].copy_(k)
            vb[:, :, slot:slot+1].copy_(v)
            out = F.scaled_dot_product_attention(q, kb[:, :, :valid], vb[:, :, :valid])
        else:
            old_valid = min(offset, capacity)
            if old_valid:
                # After wrap, the oldest slot is the next slot to overwrite.
                oldest = (offset - old_valid) % capacity
                if oldest:
                    old_k = torch.cat((kb[:, :, oldest:], kb[:, :, :oldest]), dim=2)
                    old_v = torch.cat((vb[:, :, oldest:], vb[:, :, :oldest]), dim=2)
                else:
                    old_k, old_v = kb[:, :, :old_valid], vb[:, :, :old_valid]
                keys = torch.cat((old_k, k), dim=2)
                values = torch.cat((old_v, v), dim=2)
            else:
                keys, values = k, v
            qi = torch.arange(offset, end, device=x.device)[:, None]
            ki = torch.arange(offset - old_valid, end, device=x.device)[None, :]
            mask = (ki <= qi) & (ki >= qi - cap + 1)
            out = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask)

            # Temporary chunk tensors do not become persistent backing storage.
            # Lay the final chronological suffix into its absolute-position ring
            # slots so the next single-token call needs only one overwrite.
            kept_k, kept_v = keys[:, :, -valid:], values[:, :, -valid:]
            start = (end - valid) % capacity
            first = min(valid, capacity - start)
            kb[:, :, start:start+first].copy_(kept_k[:, :, :first])
            vb[:, :, start:start+first].copy_(kept_v[:, :, :first])
            if first < valid:
                kb[:, :, :valid-first].copy_(kept_k[:, :, first:])
                vb[:, :, :valid-first].copy_(kept_v[:, :, first:])

        cache.update(layer_idx=self.layer_idx, offset=t, k=kb[:, :, :valid],
                     v=vb[:, :, :valid], k_buffer=kb, v_buffer=vb,
                     policy_cap=cap, policy_history='suffix')
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))

    def protected_cached_forward(self, x, cap, cache):
        """Exact union of prefix and previous/current blocks, with physical eviction.

        Prefix slots precede a contiguous recent suffix. Their absolute positions
        are implicit in Python counters. Single-token decoding calls unmasked
        SDPA and shifts the suffix only when the protected block changes.
        Multi-token calls use the identical per-query policy before eviction.
        """
        if self.training:
            raise ValueError('Persistent cache is for inference only')
        b, t, _ = x.shape
        offset = cache.get_seq_length(self.layer_idx)
        end = offset + t
        if cache.max_length is not None and end > cache.max_length:
            raise ValueError('Decode would exceed the declared cache capacity')
        state = cache.states[self.layer_idx]
        if state is not None and (state.get('policy_cap') != cap or
                                  state.get('policy_recent_block') != self.recent_block):
            raise ValueError('A populated prefix cache cannot change its retained-history policy')
        q, k, v = [p(x).view(b, t, self.heads, self.dim)
                   for p in (self.q_proj, self.k_proj, self.v_proj)]
        q, k = self.rotary(q, k, seqlen_offset=offset,
                           max_seqlen=max(self.context, end))
        q, k, v = (z.transpose(1, 2) for z in (q, k, v))
        capacity = cap + 2 * self.recent_block
        if cache.max_length is not None:
            capacity = min(capacity, cache.max_length)
        if state is None:
            shape = (b, self.heads, capacity, self.dim)
            kb, vb = k.new_empty(shape), v.new_empty(shape)
        else:
            kb, vb = state['k_buffer'], state['v_buffer']
        prefix_count = min(cap, end)
        tail_start = max(cap, ((end - 1) // self.recent_block - 1) * self.recent_block)
        tail_count = max(0, end - tail_start)
        valid = prefix_count + tail_count
        if t == 1:
            if offset < cap:
                kb[:, :, offset:offset+1].copy_(k)
                vb[:, :, offset:offset+1].copy_(v)
            else:
                old_start = state['tail_start'] if state is not None else tail_start
                survivors = max(0, offset - tail_start)
                shift = tail_start - old_start
                if shift > 0 and survivors:
                    # copy_ with overlapping source/destination is undefined.
                    kb[:, :, cap:cap+survivors].copy_(kb[:, :, cap+shift:cap+shift+survivors].clone())
                    vb[:, :, cap:cap+survivors].copy_(vb[:, :, cap+shift:cap+shift+survivors].clone())
                kb[:, :, valid-1:valid].copy_(k)
                vb[:, :, valid-1:valid].copy_(v)
            out = F.scaled_dot_product_attention(q, kb[:, :, :valid], vb[:, :, :valid])
        else:
            new_positions = torch.arange(offset, end, device=x.device)
            if state is None:
                keys, values, positions = k, v, new_positions
            else:
                keys = torch.cat((state['k'], k), dim=2)
                values = torch.cat((state['v'], v), dim=2)
                old_positions = torch.cat((
                    torch.arange(min(cap, offset), device=x.device),
                    torch.arange(min(offset, state['tail_start']), offset, device=x.device)))
                positions = torch.cat((old_positions, new_positions))
            qi, ki = new_positions[:, None], positions[None, :]
            start = (qi // self.recent_block - 1) * self.recent_block
            mask = (ki <= qi) & ((ki < cap) | (ki >= start))
            out = F.scaled_dot_product_attention(q, keys, values, attn_mask=mask)
            kb[:, :, :prefix_count].copy_(keys[:, :, :prefix_count])
            vb[:, :, :prefix_count].copy_(values[:, :, :prefix_count])
            if tail_count:
                kb[:, :, prefix_count:valid].copy_(keys[:, :, -tail_count:])
                vb[:, :, prefix_count:valid].copy_(values[:, :, -tail_count:])
        cache.update(layer_idx=self.layer_idx, offset=t, k=kb[:, :, :valid],
                     v=vb[:, :, :valid], k_buffer=kb, v_buffer=vb,
                     policy_cap=cap, policy_recent_block=self.recent_block,
                     tail_start=tail_start)
        return self.o_proj(out.transpose(1, 2).reshape(b, t, -1))


class Block(nn.Module):
    def __init__(self, cfg, i):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.hidden, eps=1e-6)
        self.mlp_norm = RMSNorm(cfg.hidden, eps=1e-6)
        self.linear = cfg.variant in ("gla", "hybrid_gla") and not cfg.dense_layer(i)
        if self.linear:
            self.attn = GatedLinearAttention(
                mode='chunk', hidden_size=cfg.hidden, expand_k=0.5, expand_v=1.0,
                num_heads=cfg.gla_heads, use_short_conv=False, use_output_gate=True,
                gate_fn='swish', norm_eps=1e-6, fuse_norm=True, layer_idx=i,
            )
        else:
            self.attn = SoftmaxAttention(cfg, i, cfg.variant == "taildrop" and not cfg.dense_layer(i))
        self.mlp = GatedMLP(cfg.hidden, intermediate_size=cfg.intermediate,
                            hidden_act='swish', fuse_swiglu=True)

    def forward(self, x, *, lengths=None, block_mask=None, cap=None, cache=None):
        residual = x
        x = self.attn_norm(x)
        if self.linear:
            x = self.attn(x, past_key_values=cache, use_cache=cache is not None)[0]
        else:
            x = self.attn(x, lengths=lengths, block_mask=block_mask, cap=cap, cache=cache)
        x, residual = self.mlp_norm(x, residual, True)
        return residual + self.mlp(x)


class ResearchLM(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        if cfg.variant not in ('dense', 'gla', 'hybrid_gla', 'taildrop', 'mamba2'):
            raise ValueError(cfg.variant)
        if cfg.recent_block < 0 or (cfg.recent_block and cfg.variant != 'taildrop'):
            raise ValueError('Recent protection is a nonnegative tail-drop policy')
        if cfg.history_policy not in ('prefix', 'suffix'):
            raise ValueError('History policy must be prefix or suffix')
        if cfg.history_policy == 'suffix' and (cfg.variant != 'taildrop' or cfg.recent_block):
            raise ValueError('Suffix history requires taildrop with no extra recent protection')
        from pretraining_suffix_min.sampling import validate_minimum
        validate_minimum(cfg.minimum_history)
        if cfg.minimum_history != 1 and (cfg.history_policy != 'suffix' or cfg.variant != 'taildrop'
                                         or cfg.recent_block):
            raise ValueError('Minimum history is defined only for unprotected suffix training')
        self.cfg = cfg
        if cfg.variant == 'mamba2':
            from pretraining_pure_mamba2.layers import PureMambaBlock as MambaOuterBlock
            from pretraining_pure_mamba2.spec import validate_config
            validate_config(cfg)
        elif cfg.mamba_outer:
            from pretraining_mamba2_outer.layers import validate_config, MambaOuterBlock
            validate_config(cfg)
        self.embeddings = nn.Embedding(cfg.vocab, cfg.hidden)
        self.layers = nn.ModuleList(
            MambaOuterBlock(cfg, i) if cfg.mamba_layer(i) else Block(cfg, i)
            for i in range(cfg.layers))
        # All softmax layers use identical RoPE dimensions/base/positions. Share
        # their non-trainable tables, as model-wide positional metadata.
        shared_rotary = None
        for layer in self.layers:
            if isinstance(layer.attn, SoftmaxAttention):
                if shared_rotary is None:
                    if cfg.mamba_outer:
                        from pretraining_mamba2_outer.layers import BoundedRotary
                        shared_rotary = BoundedRotary(cfg.hidden // cfg.heads, max_cached=cfg.context)
                        layer.attn.rotary = shared_rotary
                    else:
                        shared_rotary = layer.attn.rotary
                else:
                    layer.attn.rotary = shared_rotary
        self.norm = RMSNorm(cfg.hidden, eps=1e-6)
        self.lm_head = nn.Linear(cfg.hidden, cfg.vocab, bias=False)
        self.lm_head.weight = self.embeddings.weight
        self.loss_fn = FusedLinearCrossEntropyLoss(accumulate_grad_in_fp32=True)
        self.initialize_shared()

    @torch.no_grad()
    def initialize_shared(self):
        # Logical-name + shape seeding makes compatible tensors identical across
        # variants without depending on module construction order or device count.
        for name, param in self.named_parameters():
            logical_name = name
            if (self.cfg.mamba_outer or self.cfg.variant == 'mamba2') and name.startswith('layers.'):
                index = int(name.split('.')[1])
                if self.cfg.mamba_layer(index):
                    # Native dt_bias/A_log/D and convolution/projection
                    # initializers were seeded independently in the constructor.
                    continue
                from pretraining_mamba2_outer.spec import shared_parameter_name
                logical_name = shared_parameter_name(name, self.cfg.dense_ends)
            if param.ndim < 2:
                param.fill_(0 if name.endswith('bias') else 1)
            else:
                key = f"{self.cfg.seed}:{logical_name}:{tuple(param.shape)}".encode()
                seed = int.from_bytes(hashlib.sha256(key).digest()[:8], 'little') % (2**63 - 1)
                generator = torch.Generator(device=param.device).manual_seed(seed)
                param.normal_(mean=0, std=0.02, generator=generator)

    def forward(self, tokens, targets=None, *, lengths=None, cap=None, cache=None,
                generator=None, last_only=False):
        if self.cfg.variant == 'mamba2':
            if lengths is not None or cap is not None:
                raise ValueError('Pure Mamba-2 has recurrent state, not an attention-window policy')
            if cache is not None:
                from pretraining_pure_mamba2.layers import validate_cache_call
                validate_cache_call(self, tokens, cache)
        elif self.cfg.mamba_outer and cache is not None:
            from pretraining_mamba2_outer.layers import validate_cache_call
            validate_cache_call(self, tokens, lengths, cap, cache)
        if cap is not None and cap < 1:
            raise ValueError("An attention history must retain at least one token")
        if self.cfg.variant == 'taildrop' and self.training and lengths is None:
            if generator is None:
                raise ValueError("Training tail-drop requires its checkpointed RNG")
            if self.cfg.minimum_history == 1:
                lengths = sample_prefixes(*tokens.shape, generator=generator, device=tokens.device)
            else:
                from pretraining_suffix_min.sampling import sample_suffix_lengths
                lengths = sample_suffix_lengths(*tokens.shape, minimum=self.cfg.minimum_history,
                                                generator=generator, device=tokens.device)
        block_mask = None
        if lengths is not None and self.cfg.variant == 'taildrop' and self.cfg.backend == 'flex':
            if self.cfg.history_policy == 'suffix':
                block_mask = suffix_block_mask(lengths, tokens.shape[1])
            else:
                block_mask = prefix_block_mask(lengths, tokens.shape[1], self.cfg.recent_block)
        x = self.embeddings(tokens)
        for layer in self.layers:
            x = layer(x, lengths=lengths, block_mask=block_mask, cap=cap, cache=cache)
        x = self.norm(x)
        if targets is not None:
            # Targets are pre-shifted in the common packed stream. All 2048
            # inputs predict a target, including the final position.
            return self.loss_fn(x, targets, self.lm_head.weight)
        if last_only:
            x = x[:, -1:]
        return self.lm_head(x)

    def parameter_count(self):
        return sum(p.numel() for p in self.parameters())