Squash clean branch history

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,3 @@
+---
+license: apache-2.0
+---

attention_config.py

@@ -0,0 +1,5 @@
+"""Shared attention configuration constants."""
+
+SUPPORTED_SEQUENCE_BOUNDARY_POLICIES = frozenset(
+    {"eos_document", "none", "segment_document"}
+)

config.json

@@ -0,0 +1,67 @@
+{
+  "architectures": [
+    "DatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dat.DatConfig",
+    "AutoModel": "modeling_dat.DatModel",
+    "AutoModelForCausalLM": "modeling_dat.DatForCausalLM",
+    "AutoModelForMaskedLM": "modeling_dat.DatForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_dat.DatForSequenceClassification"
+  },
+  "bos_token_id": 1,
+  "dff_factor": 4,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "ffn_activation": "swiglu",
+  "ffn_hidden_dim_mode": "dff_factor",
+  "hidden_dim": 768,
+  "init_range": 0.15,
+  "init_scheme": "normal_0_02_scaled_projection",
+  "max_position_embeddings": 513,
+  "max_rel_pos": 512,
+  "max_seq_len": 513,
+  "mlm_head_enabled": false,
+  "model_type": "dat",
+  "n_heads_ra": 3,
+  "n_heads_sa": 9,
+  "n_layers": 12,
+  "n_symbols": null,
+  "norm_first": true,
+  "norm_type": "layernorm",
+  "pad_token_id": 3,
+  "pe_type": "rope",
+  "positional_symbols_sinusoidal": false,
+  "ra_n_relations": null,
+  "ra_rel_activation": "identity",
+  "ra_symmetric_rels": false,
+  "ra_type": "rca",
+  "relative_symbols_rope": true,
+  "relsymbolic_dropout": 0.0,
+  "relsymbolic_include_self": false,
+  "relsymbolic_neighborhood_size": 2,
+  "relsymbolic_normalize_rels": true,
+  "relsymbolic_rel_n_heads": 4,
+  "relsymbolic_rel_scale": null,
+  "relsymbolic_symbolic_attn_n_heads": 4,
+  "relsymbolic_symbolic_attn_scale": null,
+  "relsymbolic_trainable_symbols": true,
+  "relsymbolic_use_bias": false,
+  "rope_theta": 10000.0,
+  "segment_boundary_token_id": null,
+  "sequence_boundary_policy": "eos_document",
+  "share_attn_params": false,
+  "shared_symbol_retriever": true,
+  "symbol_dim": null,
+  "symbol_retrieval": "relative",
+  "symbolic_attn_n_heads": null,
+  "symbolic_use_bias": false,
+  "tie_lm_head": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.12.0",
+  "use_bias_ffn": true,
+  "use_bias_out": false,
+  "use_bias_qkv": false,
+  "vocab_size": 16384
+}

configuration_dat.py

@@ -0,0 +1,198 @@
+"""HuggingFace PretrainedConfig for the dual-attention (DAT) decoder LM.
+
+This is the source-of-truth copy. scripts/convert_dat_to_hf.py copies it into a
+generated HF repository (alongside modeling_dat.py and the flattened model
+source) so the model can be loaded with
+``AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True)``.
+
+The config simply carries every field of models.attention.dat.config.DatLMConfig
+so that modeling_dat.py can rebuild the exact DatLMConfig at load time.
+"""
+
+from transformers import PretrainedConfig
+
+
+# Every field needed to rebuild DatLMConfig. New exported fields are appended
+# after existing constructor parameters to preserve positional compatibility.
+# modeling_dat.py rebuilds DatLMConfig as
+# DatLMConfig(**{f: getattr(config, f) for f in DAT_LM_FIELDS}).
+DAT_LM_FIELDS = (
+    "vocab_size",
+    "max_seq_len",
+    "pe_type",
+    "hidden_dim",
+    "n_heads_sa",
+    "n_heads_ra",
+    "n_layers",
+    "dropout",
+    "dff_factor",
+    "ffn_hidden_dim_mode",
+    "ffn_activation",
+    "rope_theta",
+    "max_rel_pos",
+    "init_range",
+    "init_scheme",
+    "norm_type",
+    "norm_first",
+    "use_bias_qkv",
+    "use_bias_out",
+    "use_bias_ffn",
+    "tie_lm_head",
+    "symbol_dim",
+    "n_symbols",
+    "symbolic_attn_n_heads",
+    "symbol_retrieval",
+    "symbolic_use_bias",
+    "shared_symbol_retriever",
+    "share_attn_params",
+    "positional_symbols_sinusoidal",
+    "relative_symbols_rope",
+    "relsymbolic_rel_n_heads",
+    "relsymbolic_symbolic_attn_n_heads",
+    "relsymbolic_neighborhood_size",
+    "relsymbolic_include_self",
+    "relsymbolic_normalize_rels",
+    "relsymbolic_trainable_symbols",
+    "relsymbolic_dropout",
+    "relsymbolic_rel_scale",
+    "relsymbolic_symbolic_attn_scale",
+    "relsymbolic_use_bias",
+    "ra_type",
+    "ra_n_relations",
+    "ra_rel_activation",
+    "ra_symmetric_rels",
+    "sequence_boundary_policy",
+    "segment_boundary_token_id",
+    "pad_token_id",
+    "bos_token_id",
+    "eos_token_id",
+    "mlm_head_enabled",
+)
+
+
+class DatConfig(PretrainedConfig):
+    model_type = "dat"
+    # Expose the model width under HF's conventional name so generic eval
+    # harnesses (e.g. babylm-eval's finetuning head, which reads
+    # config.hidden_size) can size their classification head. The model's own
+    # field is hidden_dim; this maps the alias onto it for both get and set.
+    attribute_map = {"hidden_size": "hidden_dim"}
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        max_seq_len: int = 513,
+        pe_type: str = "rope",
+        hidden_dim: int = 256,
+        n_heads_sa: int = 2,
+        n_heads_ra: int = 2,
+        n_layers: int = 4,
+        dropout: float = 0.0,
+        dff_factor: int = 4,
+        ffn_hidden_dim_mode: str = "dff_factor",
+        ffn_activation: str = "gelu",
+        rope_theta: float = 10000.0,
+        max_rel_pos: int | None = None,
+        init_range: float = 0.15,
+        init_scheme: str = "xavier_uniform",
+        norm_type: str = "rmsnorm",
+        norm_first: bool = True,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+        use_bias_ffn: bool = True,
+        tie_lm_head: bool = True,
+        symbol_dim: int | None = None,
+        n_symbols: int | None = None,
+        symbolic_attn_n_heads: int | None = None,
+        symbol_retrieval: str = "symbolic",
+        symbolic_use_bias: bool = False,
+        shared_symbol_retriever: bool = True,
+        share_attn_params: bool = False,
+        positional_symbols_sinusoidal: bool = False,
+        relative_symbols_rope: bool = False,
+        relsymbolic_rel_n_heads: int = 4,
+        relsymbolic_symbolic_attn_n_heads: int = 4,
+        relsymbolic_neighborhood_size: int = 2,
+        relsymbolic_include_self: bool = False,
+        relsymbolic_normalize_rels: bool = True,
+        relsymbolic_trainable_symbols: bool = True,
+        relsymbolic_dropout: float = 0.0,
+        relsymbolic_rel_scale: float | None = None,
+        relsymbolic_symbolic_attn_scale: float | None = None,
+        relsymbolic_use_bias: bool = False,
+        ra_type: str = "ra",
+        ra_n_relations: int | None = None,
+        ra_rel_activation: str = "identity",
+        ra_symmetric_rels: bool = False,
+        sequence_boundary_policy: str = "eos_document",
+        segment_boundary_token_id: int | None = None,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        mlm_head_enabled: bool = False,
+        **kwargs,
+    ) -> None:
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.pe_type = pe_type
+        self.hidden_dim = hidden_dim
+        self.n_heads_sa = n_heads_sa
+        self.n_heads_ra = n_heads_ra
+        self.n_layers = n_layers
+        self.dropout = dropout
+        self.dff_factor = dff_factor
+        self.ffn_hidden_dim_mode = ffn_hidden_dim_mode
+        self.ffn_activation = ffn_activation
+        self.rope_theta = rope_theta
+        self.max_rel_pos = max_rel_pos
+        self.init_range = init_range
+        self.init_scheme = init_scheme
+        self.norm_type = norm_type
+        self.norm_first = norm_first
+        self.use_bias_qkv = use_bias_qkv
+        self.use_bias_out = use_bias_out
+        self.use_bias_ffn = use_bias_ffn
+        self.tie_lm_head = tie_lm_head
+        self.symbol_dim = symbol_dim
+        self.n_symbols = n_symbols
+        self.symbolic_attn_n_heads = symbolic_attn_n_heads
+        self.symbol_retrieval = symbol_retrieval
+        self.symbolic_use_bias = symbolic_use_bias
+        self.shared_symbol_retriever = shared_symbol_retriever
+        self.share_attn_params = share_attn_params
+        self.positional_symbols_sinusoidal = positional_symbols_sinusoidal
+        self.relative_symbols_rope = relative_symbols_rope
+        self.relsymbolic_rel_n_heads = relsymbolic_rel_n_heads
+        self.relsymbolic_symbolic_attn_n_heads = relsymbolic_symbolic_attn_n_heads
+        self.relsymbolic_neighborhood_size = relsymbolic_neighborhood_size
+        self.relsymbolic_include_self = relsymbolic_include_self
+        self.relsymbolic_normalize_rels = relsymbolic_normalize_rels
+        self.relsymbolic_trainable_symbols = relsymbolic_trainable_symbols
+        self.relsymbolic_dropout = relsymbolic_dropout
+        self.relsymbolic_rel_scale = relsymbolic_rel_scale
+        self.relsymbolic_symbolic_attn_scale = relsymbolic_symbolic_attn_scale
+        self.relsymbolic_use_bias = relsymbolic_use_bias
+        self.ra_type = ra_type
+        self.ra_n_relations = ra_n_relations
+        self.ra_rel_activation = ra_rel_activation
+        self.ra_symmetric_rels = ra_symmetric_rels
+        self.sequence_boundary_policy = sequence_boundary_policy
+        self.segment_boundary_token_id = segment_boundary_token_id
+        self.mlm_head_enabled = mlm_head_enabled
+        # HF base also exposes max_position_embeddings for generation utilities.
+        self.max_position_embeddings = max_seq_len
+        # tie_lm_head is the source of truth for weight tying; keep HF's
+        # tie_word_embeddings in lockstep so from_pretrained never re-ties an
+        # untied head (or fails to tie a tied one). Drop any value coming in via
+        # kwargs (e.g. a serialized config.json) so the two cannot disagree.
+        # NOTE: tie_lm_head controls only lm_head <-> token_embeddings tying.
+        # The mlm_head.linear_out <-> lm_head tying is handled separately by
+        # the get_expanded_tied_weights_keys override in modeling_dat.py.
+        kwargs.pop("tie_word_embeddings", None)
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_lm_head,
+            **kwargs,
+        )

dat_config.py

@@ -0,0 +1,283 @@
+"""Configuration for the owned dual-attention decoder LM."""
+
+from dataclasses import dataclass
+
+from .attention_config import SUPPORTED_SEQUENCE_BOUNDARY_POLICIES
+
+
+SUPPORTED_PE_TYPES = {"sinusoidal", "learned", "relative", "rope", "none"}
+SUPPORTED_SYMBOL_RETRIEVAL = {"symbolic", "positional", "relative", "relsymbolic"}
+SUPPORTED_RA_TYPES = {"ra", "rca", "disrca"}
+SUPPORTED_RA_ACTIVATIONS = {"softmax", "identity", "relu", "tanh", "sigmoid", "gelu"}
+SUPPORTED_FFN_ACTIVATIONS = {"gelu", "relu", "swiglu", "identity"}
+SUPPORTED_FFN_HIDDEN_DIM_MODES = {"dff_factor", "swiglu_parameter_matched"}
+SUPPORTED_INIT_SCHEMES = {"xavier_uniform", "normal_0_02_scaled_projection"}
+
+
+@dataclass(frozen=True)
+class DatLMConfig:
+    vocab_size: int
+    max_seq_len: int
+    pe_type: str = "rope"
+    hidden_dim: int = 256
+    n_heads_sa: int = 2
+    n_heads_ra: int = 2
+    n_layers: int = 4
+    dropout: float = 0.0
+    dff_factor: int = 4
+    ffn_hidden_dim_mode: str = "dff_factor"
+    ffn_activation: str = "gelu"
+    rope_theta: float = 10000.0
+    max_rel_pos: int | None = None
+    sequence_boundary_policy: str = "eos_document"
+    segment_boundary_token_id: int | None = None
+    init_range: float = 0.15
+    init_scheme: str = "xavier_uniform"
+    norm_type: str = "rmsnorm"
+    norm_first: bool = True
+    use_bias_qkv: bool = False
+    use_bias_out: bool = True
+    use_bias_ffn: bool = True
+    tie_lm_head: bool = True
+    symbol_dim: int | None = None
+    n_symbols: int | None = None
+    symbolic_attn_n_heads: int | None = None
+    symbol_retrieval: str = "symbolic"
+    symbolic_use_bias: bool = False
+    shared_symbol_retriever: bool = True
+    share_attn_params: bool = False
+    positional_symbols_sinusoidal: bool = False
+    relative_symbols_rope: bool = False
+    relsymbolic_rel_n_heads: int = 4
+    relsymbolic_symbolic_attn_n_heads: int = 4
+    relsymbolic_neighborhood_size: int = 2
+    relsymbolic_include_self: bool = False
+    relsymbolic_normalize_rels: bool = True
+    relsymbolic_trainable_symbols: bool = True
+    relsymbolic_dropout: float = 0.0
+    relsymbolic_rel_scale: float | None = None
+    relsymbolic_symbolic_attn_scale: float | None = None
+    relsymbolic_use_bias: bool = False
+    ra_type: str = "ra"
+    ra_n_relations: int | None = None
+    ra_rel_activation: str = "identity"
+    ra_symmetric_rels: bool = False
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    mlm_head_enabled: bool = False
+
+    def __post_init__(self) -> None:
+        if self.vocab_size <= 0:
+            raise ValueError(f"vocab_size must be positive, got {self.vocab_size}")
+        if self.max_seq_len <= 0:
+            raise ValueError(f"max_seq_len must be positive, got {self.max_seq_len}")
+        if self.pe_type not in SUPPORTED_PE_TYPES:
+            raise ValueError(f"Unsupported pe_type: {self.pe_type}")
+        if self.hidden_dim <= 0:
+            raise ValueError(f"hidden_dim must be positive, got {self.hidden_dim}")
+        if self.n_heads_sa <= 0:
+            raise ValueError(f"n_heads_sa must be positive for DAT, got {self.n_heads_sa}")
+        if self.n_heads_ra <= 0:
+            raise ValueError(f"n_heads_ra must be positive for DAT, got {self.n_heads_ra}")
+        total_heads = self.total_n_heads
+        if self.hidden_dim % total_heads != 0:
+            raise ValueError(
+                f"hidden_dim ({self.hidden_dim}) must be divisible by total DAT heads "
+                f"({total_heads} = {self.n_heads_sa} SA + {self.n_heads_ra} RA)"
+            )
+        if self.n_layers <= 0:
+            raise ValueError(f"n_layers must be positive, got {self.n_layers}")
+        if not 0.0 <= self.dropout < 1.0:
+            raise ValueError(f"dropout must be in [0.0, 1.0), got {self.dropout}")
+        if self.dff_factor <= 0:
+            raise ValueError(f"dff_factor must be positive, got {self.dff_factor}")
+        if self.ffn_hidden_dim_mode not in SUPPORTED_FFN_HIDDEN_DIM_MODES:
+            raise ValueError(f"Unsupported ffn_hidden_dim_mode for DAT: {self.ffn_hidden_dim_mode}")
+        if self.ffn_activation not in SUPPORTED_FFN_ACTIVATIONS:
+            raise ValueError(f"Unsupported ffn_activation for DAT: {self.ffn_activation}")
+        if self.ffn_hidden_dim_mode == "swiglu_parameter_matched" and self.ffn_activation != "swiglu":
+            raise ValueError(
+                "ffn_hidden_dim_mode='swiglu_parameter_matched' requires "
+                f"ffn_activation='swiglu', got {self.ffn_activation}"
+            )
+        if self.rope_theta <= 0.0:
+            raise ValueError(f"rope_theta must be positive, got {self.rope_theta}")
+        if self.max_rel_pos is not None and self.max_rel_pos <= 0:
+            raise ValueError(f"max_rel_pos must be positive when provided, got {self.max_rel_pos}")
+        if self.sequence_boundary_policy not in SUPPORTED_SEQUENCE_BOUNDARY_POLICIES:
+            raise ValueError(
+                "Unsupported sequence_boundary_policy for DAT: "
+                f"{self.sequence_boundary_policy}"
+            )
+        if self.sequence_boundary_policy == "segment_document":
+            if self.segment_boundary_token_id is None:
+                raise ValueError(
+                    "segment_boundary_token_id is required when "
+                    "sequence_boundary_policy='segment_document'"
+                )
+        if self.init_range <= 0.0:
+            raise ValueError(f"init_range must be positive, got {self.init_range}")
+        if self.init_scheme not in SUPPORTED_INIT_SCHEMES:
+            raise ValueError(f"Unsupported init_scheme for DAT: {self.init_scheme}")
+        if self.norm_type not in {"layernorm", "rmsnorm"}:
+            raise ValueError(
+                f"norm_type must be 'layernorm' or 'rmsnorm', got {self.norm_type}"
+            )
+        if self.pe_type == "rope" and self.head_dim % 2 != 0:
+            raise ValueError(
+                "RoPE requires even DAT head_dim, got "
+                f"{self.head_dim} from hidden_dim={self.hidden_dim}, total_heads={total_heads}"
+            )
+        if self.pe_type == "sinusoidal" and self.hidden_dim % 2 != 0:
+            raise ValueError(f"Sinusoidal encoding requires even hidden_dim, got {self.hidden_dim}")
+        if self.share_attn_params and self.n_heads_sa != self.n_heads_ra:
+            raise ValueError(
+                "share_attn_params=True requires n_heads_sa == n_heads_ra, "
+                f"got {self.n_heads_sa} and {self.n_heads_ra}"
+            )
+        if self.symbol_dim is not None and self.symbol_dim <= 0:
+            raise ValueError(f"symbol_dim must be positive when provided, got {self.symbol_dim}")
+        if self.symbolic_attn_n_heads is not None and self.symbolic_attn_n_heads <= 0:
+            raise ValueError(
+                "symbolic_attn_n_heads must be positive when provided, "
+                f"got {self.symbolic_attn_n_heads}"
+            )
+        if self.symbol_retrieval not in SUPPORTED_SYMBOL_RETRIEVAL:
+            raise ValueError(f"Unsupported symbol_retrieval for DAT: {self.symbol_retrieval}")
+        if self.symbol_retrieval == "symbolic":
+            symbolic_heads = self.resolved_symbolic_attn_n_heads
+            if self.hidden_dim % symbolic_heads != 0:
+                raise ValueError(
+                    f"hidden_dim ({self.hidden_dim}) must be divisible by symbolic_attn_n_heads "
+                    f"({symbolic_heads}) for symbolic retrieval"
+                )
+            if self.resolved_symbol_dim % symbolic_heads != 0:
+                raise ValueError(
+                    f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by symbolic_attn_n_heads "
+                    f"({symbolic_heads}) for symbolic retrieval"
+                )
+        if (
+            self.symbol_retrieval == "positional"
+            and self.positional_symbols_sinusoidal
+            and self.resolved_symbol_dim % 2 != 0
+        ):
+            raise ValueError(
+                "Sinusoidal positional symbols require even symbol_dim, "
+                f"got {self.resolved_symbol_dim}"
+            )
+        if self.relative_symbols_rope:
+            if self.symbol_retrieval != "relative":
+                raise ValueError(
+                    "relative_symbols_rope=True requires symbol_retrieval='relative', "
+                    f"got {self.symbol_retrieval!r}"
+                )
+            if self.resolved_symbol_dim % 2 != 0:
+                raise ValueError(
+                    "RoPE relative symbols require even symbol_dim, "
+                    f"got {self.resolved_symbol_dim}"
+                )
+        if self.positional_symbols_sinusoidal and self.symbol_retrieval != "positional":
+            raise ValueError(
+                "positional_symbols_sinusoidal=True requires symbol_retrieval='positional', "
+                f"got {self.symbol_retrieval!r}"
+            )
+        if self.resolved_n_symbols <= 0:
+            raise ValueError(f"resolved_n_symbols must be positive, got {self.resolved_n_symbols}")
+        if self.symbol_retrieval == "relsymbolic":
+            if self.relsymbolic_rel_n_heads <= 0:
+                raise ValueError(
+                    f"relsymbolic_rel_n_heads must be positive, got {self.relsymbolic_rel_n_heads}"
+                )
+            if self.hidden_dim % self.relsymbolic_rel_n_heads != 0:
+                raise ValueError(
+                    f"hidden_dim ({self.hidden_dim}) must be divisible by "
+                    f"relsymbolic_rel_n_heads ({self.relsymbolic_rel_n_heads})"
+                )
+            if self.relsymbolic_symbolic_attn_n_heads <= 0:
+                raise ValueError(
+                    "relsymbolic_symbolic_attn_n_heads must be positive, "
+                    f"got {self.relsymbolic_symbolic_attn_n_heads}"
+                )
+            if self.hidden_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
+                raise ValueError(
+                    f"hidden_dim ({self.hidden_dim}) must be divisible by "
+                    f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
+                )
+            if self.resolved_symbol_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
+                raise ValueError(
+                    f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by "
+                    f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
+                )
+            if self.relsymbolic_neighborhood_size <= 0:
+                raise ValueError(
+                    "relsymbolic_neighborhood_size must be positive, "
+                    f"got {self.relsymbolic_neighborhood_size}"
+                )
+            if not 0.0 <= self.relsymbolic_dropout < 1.0:
+                raise ValueError(
+                    f"relsymbolic_dropout must be in [0.0, 1.0), got {self.relsymbolic_dropout}"
+                )
+            if self.relsymbolic_rel_scale is not None and self.relsymbolic_rel_scale <= 0.0:
+                raise ValueError(
+                    "relsymbolic_rel_scale must be positive when provided, "
+                    f"got {self.relsymbolic_rel_scale}"
+                )
+            if (
+                self.relsymbolic_symbolic_attn_scale is not None
+                and self.relsymbolic_symbolic_attn_scale <= 0.0
+            ):
+                raise ValueError(
+                    "relsymbolic_symbolic_attn_scale must be positive when provided, "
+                    f"got {self.relsymbolic_symbolic_attn_scale}"
+                )
+        if self.ra_type not in SUPPORTED_RA_TYPES:
+            raise ValueError(f"Unsupported ra_type for DAT: {self.ra_type}")
+        if self.ra_n_relations is not None and self.ra_n_relations <= 0:
+            raise ValueError(
+                f"ra_n_relations must be positive when provided, got {self.ra_n_relations}"
+            )
+        if self.ra_type != "ra" and self.ra_n_relations is not None:
+            raise ValueError(f"ra_n_relations applies only to ra_type='ra', got {self.ra_type}")
+        if self.ra_type != "ra" and self.ra_symmetric_rels:
+            raise ValueError(f"ra_symmetric_rels applies only to ra_type='ra', got {self.ra_type}")
+        n_relations = self.resolved_ra_n_relations
+        if self.ra_type == "ra" and (self.head_dim * self.n_heads_ra) % n_relations != 0:
+            raise ValueError(
+                f"head_dim * n_heads_ra ({self.head_dim * self.n_heads_ra}) must be "
+                f"divisible by ra_n_relations ({n_relations})"
+            )
+        if self.ra_rel_activation not in SUPPORTED_RA_ACTIVATIONS:
+            raise ValueError(
+                f"Unsupported ra_rel_activation for DAT: {self.ra_rel_activation}"
+            )
+
+    @property
+    def total_n_heads(self) -> int:
+        return self.n_heads_sa + self.n_heads_ra
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_dim // self.total_n_heads
+
+    @property
+    def resolved_symbol_dim(self) -> int:
+        return self.hidden_dim if self.symbol_dim is None else self.symbol_dim
+
+    @property
+    def resolved_symbolic_attn_n_heads(self) -> int:
+        return self.total_n_heads if self.symbolic_attn_n_heads is None else self.symbolic_attn_n_heads
+
+    @property
+    def resolved_ffn_hidden_dim(self) -> int:
+        if self.ffn_hidden_dim_mode == "swiglu_parameter_matched":
+            return int(8 / 3 * self.hidden_dim)
+        return self.hidden_dim * self.dff_factor
+
+    @property
+    def resolved_n_symbols(self) -> int:
+        return self.max_seq_len if self.n_symbols is None else self.n_symbols
+
+    @property
+    def resolved_ra_n_relations(self) -> int:
+        return self.n_heads_ra if self.ra_n_relations is None else self.ra_n_relations

dat_core.py

@@ -0,0 +1,328 @@
+"""Dual-attention primitives for decoder language models."""
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .transformer_components import PositionalInfo
+from .transformer_core import MultiHeadAttentionBase
+
+
+def _activate_scores(scores: torch.Tensor, activation: str) -> torch.Tensor:
+    if activation == "softmax":
+        return torch.softmax(scores, dim=-1)
+    if activation == "identity":
+        return scores
+    if activation == "relu":
+        return torch.relu(scores)
+    if activation == "tanh":
+        return torch.tanh(scores)
+    if activation == "sigmoid":
+        return torch.sigmoid(scores)
+    if activation == "gelu":
+        return torch.nn.functional.gelu(scores)
+    raise ValueError(f"Unsupported attention activation: {activation}")
+
+
+class RelationalAttentionBase(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        dropout: float = 0.0,
+        n_relations: Optional[int] = None,
+        rel_activation: str = "identity",
+        symmetric_rels: bool = False,
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        head_dim = hidden_dim // total_n_heads
+        output_dim = n_heads * head_dim
+        super().__init__(
+            query_dim=hidden_dim,
+            output_dim=output_dim,
+            key_dim=hidden_dim,
+            value_dim=symbol_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            total_n_heads=total_n_heads,
+            activation="softmax",
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.symbol_dim = symbol_dim
+        self.rel_activation = rel_activation
+        self.symmetric_rels = symmetric_rels
+        self.use_relative_positional_symbols = use_relative_positional_symbols
+        self.n_relations = n_heads if n_relations is None else n_relations
+        total_rel_dim = self.head_dim * n_heads
+        if total_rel_dim % self.n_relations != 0:
+            raise ValueError(
+                f"head_dim * n_heads ({total_rel_dim}) must be divisible by n_relations "
+                f"({self.n_relations})"
+            )
+        self.rel_proj_dim = total_rel_dim // self.n_relations
+        self.rel_scale = 1.0 / math.sqrt(self.rel_proj_dim)
+        rel_total_dim = self.n_relations * self.rel_proj_dim
+        self.wq_rel = nn.Linear(hidden_dim, rel_total_dim, bias=False)
+        self.wk_rel = self.wq_rel if symmetric_rels else nn.Linear(
+            hidden_dim,
+            rel_total_dim,
+            bias=False,
+        )
+        nn.init.xavier_uniform_(self.wq_rel.weight)
+        if self.wk_rel is not self.wq_rel:
+            nn.init.xavier_uniform_(self.wk_rel.weight)
+
+    def _compute_base_attention(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        q = self._reshape_for_multihead(self.q_proj(x), batch_size, seq_len)
+        k = self._reshape_for_multihead(self.k_proj(x), batch_size, seq_len)
+        scores = self._compute_attn_scores(q, k, pos_info)
+        return self._apply_activation_and_mask(scores, mask)
+
+    def compute_relational_scores(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        return_as: str,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        q_rel = self.wq_rel(x).view(batch_size, seq_len, self.n_relations, self.rel_proj_dim)
+        k_rel = self.wk_rel(x).view(batch_size, seq_len, self.n_relations, self.rel_proj_dim)
+        q_rel = q_rel.transpose(1, 2)
+        k_rel = k_rel.transpose(1, 2)
+        relations = torch.matmul(q_rel, k_rel.transpose(-2, -1)) * self.rel_scale
+        processed_mask = self._process_mask(mask)
+        if self.rel_activation == "softmax" and processed_mask is not None:
+            relations.masked_fill_(~processed_mask, torch.finfo(relations.dtype).min)
+        relations = _activate_scores(relations, self.rel_activation)
+        if processed_mask is not None:
+            relations = relations.masked_fill(~processed_mask, 0.0)
+        if return_as == "vectors":
+            return relations.permute(0, 2, 3, 1)
+        if return_as == "scores":
+            return relations
+        raise ValueError(f"return_as must be 'vectors' or 'scores', got {return_as}")
+
+    def combine_attention_and_relations(
+        self,
+        attn_weights: torch.Tensor,
+        relation_scores: torch.Tensor,
+    ) -> torch.Tensor:
+        return attn_weights * relation_scores
+
+    def _process_symbols_with_attention(
+        self,
+        symbols: torch.Tensor,
+        attn_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size, _, seq_len_q, seq_len_k = attn_weights.shape
+        values = self.v_proj(symbols)
+        if self.use_relative_positional_symbols:
+            values = values.view(seq_len_q, seq_len_k, self.n_heads, self.head_dim)
+            return torch.einsum("bhij,ijhd->bihd", attn_weights, values)
+
+        values = values.view(batch_size, seq_len_k, self.n_heads, self.head_dim)
+        values = values.transpose(1, 2)
+        output = torch.matmul(attn_weights, values)
+        return output.transpose(1, 2)
+
+    def _apply_output_projection(self, output: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len = output.shape[:2]
+        output = output.contiguous().view(batch_size, seq_len, self.output_dim)
+        output = self.o_proj(output)
+        output = self.dropout(output)
+        return output
+
+    def _validate_symbols(self, x: torch.Tensor, symbols: torch.Tensor) -> None:
+        if self.use_relative_positional_symbols:
+            seq_len = x.shape[1]
+            expected_shape = (seq_len, seq_len, self.symbol_dim)
+            if tuple(symbols.shape) != expected_shape:
+                raise ValueError(
+                    f"Relative symbols must have shape {expected_shape}, got {tuple(symbols.shape)}"
+                )
+
+
+class RelationalAttention(RelationalAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        n_relations: int,
+        dropout: float = 0.0,
+        rel_activation: str = "identity",
+        symmetric_rels: bool = False,
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        super().__init__(
+            hidden_dim=hidden_dim,
+            symbol_dim=symbol_dim,
+            n_heads=n_heads,
+            total_n_heads=total_n_heads,
+            dropout=dropout,
+            n_relations=n_relations,
+            rel_activation=rel_activation,
+            symmetric_rels=symmetric_rels,
+            use_relative_positional_symbols=use_relative_positional_symbols,
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.wr_proj = nn.Parameter(torch.empty(n_heads, self.head_dim, n_relations))
+        nn.init.xavier_uniform_(self.wr_proj)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbols: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        self._validate_symbols(x, symbols)
+        attn_weights = self._compute_base_attention(x, mask, pos_info)
+        attn_weights = self.attn_dropout(attn_weights)
+        relation_vectors = self.compute_relational_scores(x, mask, return_as="vectors")
+        attended_symbols = self._process_symbols_with_attention(symbols, attn_weights)
+        projected_relations = torch.einsum(
+            "bhij,bijr,hdr->bihd",
+            attn_weights,
+            relation_vectors,
+            self.wr_proj,
+        )
+        output = self._apply_output_projection(attended_symbols + projected_relations)
+        self.last_attn_weights = attn_weights.detach()
+        return output, {"attention": attn_weights, "relations": relation_vectors}
+
+
+class RelationalCrossAttention(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        dropout: float = 0.0,
+        activation: str = "identity",
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        head_dim = hidden_dim // total_n_heads
+        super().__init__(
+            query_dim=hidden_dim,
+            output_dim=n_heads * head_dim,
+            key_dim=hidden_dim,
+            value_dim=symbol_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            total_n_heads=total_n_heads,
+            activation=activation,
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.symbol_dim = symbol_dim
+        self.use_relative_positional_symbols = use_relative_positional_symbols
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbols: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        if self.use_relative_positional_symbols:
+            batch_size, seq_len, _ = x.shape
+            expected_shape = (seq_len, seq_len, self.symbol_dim)
+            if tuple(symbols.shape) != expected_shape:
+                raise ValueError(
+                    f"Relative symbols must have shape {expected_shape}, got {tuple(symbols.shape)}"
+                )
+
+            q = self._reshape_for_multihead(self.q_proj(x), batch_size, seq_len)
+            k = self._reshape_for_multihead(self.k_proj(x), batch_size, seq_len)
+            values = self.v_proj(symbols).view(seq_len, seq_len, self.n_heads, self.head_dim)
+            scores = self._compute_attn_scores(q, k, pos_info)
+            weights = self._apply_activation_and_mask(scores, mask)
+            weights = self.attn_dropout(weights)
+            output = torch.einsum("bhij,ijhd->bihd", weights, values)
+            output = output.contiguous().view(batch_size, seq_len, self.output_dim)
+            output = self.o_proj(output)
+            output = self.dropout(output)
+            self.last_attn_weights = weights.detach()
+            return output, {"attention": weights}
+
+        output, weights = super().forward(
+            query=x,
+            key=x,
+            value=symbols,
+            mask=mask,
+            pos_info=pos_info,
+        )
+        return output, {"attention": weights}
+
+
+class DisentangledRelationalCrossAttention(RelationalAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        dropout: float = 0.0,
+        rel_activation: str = "identity",
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        super().__init__(
+            hidden_dim=hidden_dim,
+            symbol_dim=symbol_dim,
+            n_heads=n_heads,
+            total_n_heads=total_n_heads,
+            dropout=dropout,
+            n_relations=None,
+            rel_activation=rel_activation,
+            symmetric_rels=False,
+            use_relative_positional_symbols=use_relative_positional_symbols,
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbols: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        self._validate_symbols(x, symbols)
+        attn_weights = self._compute_base_attention(x, mask, pos_info)
+        relation_scores = self.compute_relational_scores(x, mask, return_as="scores")
+        combined_weights = self.attn_dropout(
+            self.combine_attention_and_relations(attn_weights, relation_scores)
+        )
+        output = self._process_symbols_with_attention(symbols, combined_weights)
+        output = self._apply_output_projection(output)
+        self.last_attn_weights = attn_weights.detach()
+        return output, {
+            "attention": attn_weights,
+            "relations": relation_scores,
+            "combined": combined_weights,
+        }

dat_lm.py

@@ -0,0 +1,475 @@
+"""Decoder-only LM built from dual-attention blocks."""
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from .dat_config import DatLMConfig
+from .dat_core import (
+    DisentangledRelationalCrossAttention,
+    RelationalAttention,
+    RelationalCrossAttention,
+)
+from .transformer_core import SelfAttention
+from .dat_symbols import (
+    PositionalSymbolRetriever,
+    RelationalSymbolicAttentionRetriever,
+    RelativePositionalSymbolRetriever,
+    SymbolicAttentionRetriever,
+)
+from .masks import build_decoder_attention_mask
+from .mlm import MaskClassifier
+from .transformer_components import FeedForward, PositionalEncoding, PositionalInfo, RMSNorm
+
+
+class DatDecoderBlock(nn.Module):
+    def __init__(self, config: DatLMConfig):
+        super().__init__()
+        self.norm_first = config.norm_first
+        # input_dim is the block state width; hidden_dim is the total DAT width
+        # used with total_n_heads to derive the shared SA/RA head_dim.
+        self.sensory_attention = SelfAttention(
+            input_dim=config.hidden_dim,
+            n_heads=config.n_heads_sa,
+            hidden_dim=config.hidden_dim,
+            total_n_heads=config.total_n_heads,
+            dropout=config.dropout,
+            supports_relative=config.pe_type == "relative",
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+        self.relational_attention = _build_relational_attention(config)
+        if config.share_attn_params:
+            self.sensory_attention.q_proj = self.relational_attention.q_proj
+            self.sensory_attention.k_proj = self.relational_attention.k_proj
+        self.feed_forward = FeedForward(
+            input_dim=config.hidden_dim,
+            hidden_dim=config.resolved_ffn_hidden_dim,
+            dropout=config.dropout,
+            activation=config.ffn_activation,
+            use_bias=config.use_bias_ffn,
+        )
+        self.norm1 = _create_norm(config.norm_type, config.hidden_dim)
+        self.norm2 = _create_norm(config.norm_type, config.hidden_dim)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbol_retriever: nn.Module,
+        mask: torch.Tensor,
+        pos_info: Optional[PositionalInfo],
+    ) -> torch.Tensor:
+        if self.norm_first:
+            normed_x = self.norm1(x)
+            # Retrieve symbols at the block boundary so pre-norm attention and
+            # input-dependent symbol retrieval use the same representation.
+            symbols = symbol_retriever(normed_x)
+            sensory_output, _ = self.sensory_attention(normed_x, mask=mask, pos_info=pos_info)
+            relational_output, _ = self.relational_attention(
+                normed_x,
+                symbols,
+                mask=mask,
+                pos_info=pos_info,
+            )
+            x = x + self.dropout(torch.cat((sensory_output, relational_output), dim=-1))
+            x = x + self.dropout(self.feed_forward(self.norm2(x)))
+            return x
+
+        symbols = symbol_retriever(x)
+        sensory_output, _ = self.sensory_attention(x, mask=mask, pos_info=pos_info)
+        relational_output, _ = self.relational_attention(
+            x,
+            symbols,
+            mask=mask,
+            pos_info=pos_info,
+        )
+        x = self.norm1(x + self.dropout(torch.cat((sensory_output, relational_output), dim=-1)))
+        x = self.norm2(x + self.dropout(self.feed_forward(x)))
+        return x
+
+
+class DatDecoderLM(nn.Module):
+    def __init__(self, config: DatLMConfig):
+        super().__init__()
+        self.config = config
+        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_dim)
+        self.embedding_dropout = nn.Dropout(config.dropout)
+        position_dim = (
+            config.hidden_dim if config.pe_type in {"sinusoidal", "learned", "none"}
+            else config.head_dim
+        )
+        self.position_encoder = PositionalEncoding(
+            embedding_dim=position_dim,
+            pe_type=config.pe_type,
+            max_len=config.max_seq_len,
+            theta=config.rope_theta,
+            max_rel_pos=config.max_rel_pos,
+            init_range=config.init_range,
+        )
+        self.symbol_retrievers = _build_symbol_retrievers(config)
+        self.layers = nn.ModuleList(DatDecoderBlock(config) for _ in range(config.n_layers))
+        self.final_norm = _create_norm(config.norm_type, config.hidden_dim)
+        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+        if config.mlm_head_enabled:
+            self.mlm_head = MaskClassifier(
+                hidden_dim=config.hidden_dim,
+                vocab_size=config.vocab_size,
+                norm_type=config.norm_type,
+                ffn_activation=config.ffn_activation,
+                use_bias_ffn=config.use_bias_ffn,
+                dropout=config.dropout,
+                word_embedding=self.lm_head.weight if not config.tie_lm_head else self.token_embeddings.weight,
+            )
+        self._init_weights()
+
+    @property
+    def device(self) -> torch.device:
+        return self.token_embeddings.weight.device
+
+    def _init_weights(self) -> None:
+        if self.config.init_scheme == "xavier_uniform":
+            nn.init.xavier_uniform_(
+                self.token_embeddings.weight,
+                gain=nn.init.calculate_gain("linear"),
+            )
+            if self.config.tie_lm_head:
+                self.lm_head.weight = self.token_embeddings.weight
+            else:
+                nn.init.xavier_uniform_(
+                    self.lm_head.weight,
+                    gain=nn.init.calculate_gain("linear"),
+                )
+            return
+
+        if self.config.init_scheme == "normal_0_02_scaled_projection":
+            self._init_normal_0_02_scaled_projection()
+            if self.config.tie_lm_head:
+                self.lm_head.weight = self.token_embeddings.weight
+            return
+
+        raise ValueError(f"Unsupported init_scheme: {self.config.init_scheme}")
+
+    def _init_normal_0_02_scaled_projection(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+        for module in self.modules():
+            if isinstance(module, RelationalAttention):
+                nn.init.normal_(module.wr_proj, mean=0.0, std=0.02)
+            elif isinstance(module, SymbolicAttentionRetriever):
+                nn.init.normal_(module.template_features, mean=0.0, std=1.0)
+                nn.init.normal_(module.symbol_library, mean=0.0, std=1.0)
+            elif isinstance(module, RelativePositionalSymbolRetriever) and not module.rope:
+                nn.init.xavier_uniform_(module.position_encoder.rel_pos_embeddings_table.weight)
+
+        scaled_std = 0.02 / math.sqrt(2 * self.config.n_layers)
+        ffn_scaled_suffix = (
+            "feed_forward.w_up.weight"
+            if self.config.ffn_activation == "swiglu"
+            else "feed_forward.linear2.weight"
+        )
+        for name, parameter in self.named_parameters():
+            # Exclude the MLM head from depth-scaled init: its feed_forward
+            # down-projection shares the same suffix as decoder layers but
+            # should use the regular 0.02 initialization.
+            if name.startswith("mlm_head"):
+                continue
+            if name.endswith("o_proj.weight") or name.endswith(ffn_scaled_suffix):
+                nn.init.normal_(parameter, mean=0.0, std=scaled_std)
+
+    def _build_attention_mask(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        bidirectional: bool = False,
+    ) -> torch.Tensor:
+        return build_decoder_attention_mask(
+            input_ids=input_ids,
+            pad_token_id=self.config.pad_token_id,
+            eos_token_id=self.config.eos_token_id,
+            sequence_boundary_policy=self.config.sequence_boundary_policy,
+            attention_mask=attention_mask,
+            segment_boundary_token_id=self.config.segment_boundary_token_id,
+            bidirectional=bidirectional,
+        )
+
+    def encode_for_objective(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        bidirectional: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if input_ids.dim() != 2:
+            raise ValueError(f"input_ids must be rank-2 [batch, seq], got shape {tuple(input_ids.shape)}")
+        if input_ids.shape[1] > self.config.max_seq_len:
+            raise ValueError(
+                f"Sequence length {input_ids.shape[1]} exceeds max_seq_len {self.config.max_seq_len}"
+            )
+
+
+        mask = self._build_attention_mask(input_ids, attention_mask, bidirectional=bidirectional)
+        token_embeddings = self.token_embeddings(input_ids)
+        hidden_states = token_embeddings
+        pos_info = self.position_encoder.get_positional_info(input_ids.shape[1], input_ids.device)
+        if pos_info.apply_to_embeddings:
+            if pos_info.embeddings is None:
+                raise ValueError(f"Embedding-level pe_type {pos_info.pe_type} did not provide embeddings.")
+            hidden_states = hidden_states + pos_info.embeddings.unsqueeze(0)
+        hidden_states = self.embedding_dropout(hidden_states)
+
+        for symbol_retriever, layer in zip(self.symbol_retrievers, self.layers):
+            hidden_states = layer(
+                hidden_states,
+                symbol_retriever=symbol_retriever,
+                mask=mask,
+                pos_info=pos_info,
+            )
+
+        hidden_states = self.final_norm(hidden_states)
+        return token_embeddings, hidden_states
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        targets: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        _, hidden_states = self.encode_for_objective(input_ids, attention_mask=attention_mask)
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if targets is not None:
+            if targets.shape != input_ids.shape:
+                raise ValueError(
+                    f"targets shape must match input_ids shape, got {tuple(targets.shape)} "
+                    f"vs {tuple(input_ids.shape)}"
+                )
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                targets.reshape(-1),
+                ignore_index=-1,
+            )
+
+        return logits, loss
+
+    def forward_mlm(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode bidirectionally and apply the MLM head.
+
+        Args:
+            input_ids: Masked input token ids [batch, seq].
+            attention_mask: Validity mask [batch, seq].
+
+        Returns:
+            MLM logits at every position [batch, seq, vocab].
+        """
+        if not self.config.mlm_head_enabled:
+            raise ValueError(
+                "forward_mlm requires mlm_head_enabled=True; "
+                "the MLM head is not instantiated."
+            )
+        _, hidden_states = self.encode_for_objective(
+            input_ids, attention_mask=attention_mask, bidirectional=True
+        )
+        return self.mlm_head(hidden_states)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: int = 100,
+        temperature: float = 1.0,
+        top_k: Optional[int] = None,
+        do_sample: bool = True,
+        eos_token_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        if temperature <= 0.0:
+            raise ValueError(f"temperature must be positive, got {temperature}")
+        if max_new_tokens < 0:
+            raise ValueError(f"max_new_tokens must be non-negative, got {max_new_tokens}")
+        if top_k is not None and top_k <= 0:
+            raise ValueError(f"top_k must be positive when provided, got {top_k}")
+        if eos_token_id is None:
+            eos_token_id = self.config.eos_token_id
+
+        if attention_mask is None:
+            generated_sequences = [row.clone() for row in input_ids]
+        else:
+            current_attention_mask = attention_mask.bool()
+            generated_sequences = []
+            for row, row_mask in zip(input_ids, current_attention_mask):
+                generated_sequence = row[row_mask]
+                if generated_sequence.numel() == 0:
+                    raise ValueError("Each input row must contain at least one unmasked token for generation.")
+                generated_sequences.append(generated_sequence)
+
+        finished = torch.tensor(
+            [sequence[-1].item() == eos_token_id for sequence in generated_sequences],
+            dtype=torch.bool,
+            device=input_ids.device,
+        )
+
+        for _ in range(max_new_tokens):
+            context_sequences = [sequence[-self.config.max_seq_len :] for sequence in generated_sequences]
+            context_ids = pad_sequence(
+                context_sequences,
+                batch_first=True,
+                padding_value=self.config.pad_token_id,
+            )
+            context_mask = pad_sequence(
+                [
+                    torch.ones(sequence.shape[0], dtype=torch.bool, device=input_ids.device)
+                    for sequence in context_sequences
+                ],
+                batch_first=True,
+                padding_value=False,
+            )
+            logits, _ = self(context_ids, attention_mask=context_mask)
+            last_positions = context_mask.long().sum(dim=1) - 1
+            batch_indices = torch.arange(logits.shape[0], device=logits.device)
+            next_token_logits = logits[batch_indices, last_positions, :] / temperature
+
+            if top_k is not None:
+                k = min(top_k, next_token_logits.size(-1))
+                top_values, _ = torch.topk(next_token_logits, k=k)
+                cutoff = top_values[:, -1].unsqueeze(-1)
+                next_token_logits = next_token_logits.masked_fill(next_token_logits < cutoff, float("-inf"))
+
+            if do_sample:
+                probs = F.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            for row_index in range(next_token.shape[0]):
+                if finished[row_index]:
+                    continue
+                generated_sequences[row_index] = torch.cat(
+                    [generated_sequences[row_index], next_token[row_index]],
+                )
+                if next_token[row_index, 0].item() == eos_token_id:
+                    finished[row_index] = True
+
+            if torch.all(finished):
+                break
+
+        return pad_sequence(
+            generated_sequences,
+            batch_first=True,
+            padding_value=self.config.pad_token_id,
+        )
+
+
+def _build_symbol_retrievers(config: DatLMConfig) -> nn.ModuleList:
+    retriever = _build_symbol_retriever(config)
+    if config.shared_symbol_retriever:
+        return nn.ModuleList([retriever] * config.n_layers)
+    return nn.ModuleList(_build_symbol_retriever(config) for _ in range(config.n_layers))
+
+
+def _build_symbol_retriever(config: DatLMConfig) -> nn.Module:
+    if config.symbol_retrieval == "symbolic":
+        return SymbolicAttentionRetriever(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_symbols=config.resolved_n_symbols,
+            n_heads=config.resolved_symbolic_attn_n_heads,
+            dropout=config.dropout,
+            use_bias=config.symbolic_use_bias,
+        )
+    if config.symbol_retrieval == "positional":
+        return PositionalSymbolRetriever(
+            symbol_dim=config.resolved_symbol_dim,
+            max_len=config.max_seq_len,
+            sinusoidal=config.positional_symbols_sinusoidal,
+        )
+    if config.symbol_retrieval == "relative":
+        max_rel_distance = config.max_rel_pos if config.max_rel_pos is not None else max(1, config.max_seq_len // 2)
+        return RelativePositionalSymbolRetriever(
+            symbol_dim=config.resolved_symbol_dim,
+            max_rel_distance=max_rel_distance,
+            rope=config.relative_symbols_rope,
+            theta=config.rope_theta,
+        )
+    if config.symbol_retrieval == "relsymbolic":
+        return RelationalSymbolicAttentionRetriever(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            rel_n_heads=config.relsymbolic_rel_n_heads,
+            symbolic_attn_n_heads=config.relsymbolic_symbolic_attn_n_heads,
+            n_symbols=config.resolved_n_symbols,
+            neighborhood_size=config.relsymbolic_neighborhood_size,
+            include_self=config.relsymbolic_include_self,
+            normalize_rels=config.relsymbolic_normalize_rels,
+            dropout=config.relsymbolic_dropout,
+            trainable_symbols=config.relsymbolic_trainable_symbols,
+            rel_scale=config.relsymbolic_rel_scale,
+            symbolic_attn_scale=config.relsymbolic_symbolic_attn_scale,
+            use_bias=config.relsymbolic_use_bias,
+        )
+    raise ValueError(f"Unsupported symbol_retrieval: {config.symbol_retrieval}")
+
+
+def _build_relational_attention(config: DatLMConfig) -> nn.Module:
+    use_relative_symbols = config.symbol_retrieval == "relative"
+    if config.ra_type == "ra":
+        return RelationalAttention(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_heads=config.n_heads_ra,
+            total_n_heads=config.total_n_heads,
+            n_relations=config.resolved_ra_n_relations,
+            dropout=config.dropout,
+            rel_activation=config.ra_rel_activation,
+            symmetric_rels=config.ra_symmetric_rels,
+            use_relative_positional_symbols=use_relative_symbols,
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+    if config.ra_type == "rca":
+        return RelationalCrossAttention(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_heads=config.n_heads_ra,
+            total_n_heads=config.total_n_heads,
+            dropout=config.dropout,
+            activation=config.ra_rel_activation,
+            use_relative_positional_symbols=use_relative_symbols,
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+    if config.ra_type == "disrca":
+        return DisentangledRelationalCrossAttention(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_heads=config.n_heads_ra,
+            total_n_heads=config.total_n_heads,
+            dropout=config.dropout,
+            rel_activation=config.ra_rel_activation,
+            use_relative_positional_symbols=use_relative_symbols,
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+    raise ValueError(f"Unsupported ra_type: {config.ra_type}")
+
+
+def _create_norm(norm_type: str, hidden_dim: int) -> nn.Module:
+    if norm_type == "layernorm":
+        return nn.LayerNorm(hidden_dim)
+    if norm_type == "rmsnorm":
+        return RMSNorm(hidden_dim)
+    raise ValueError(f"Unsupported norm_type: {norm_type}")

dat_symbols.py

@@ -0,0 +1,263 @@
+"""Symbol retrieval modules for the owned DAT backend."""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .transformer_components import PositionalEncoding
+
+
+class SymbolicAttentionRetriever(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_symbols: int,
+        n_heads: int,
+        dropout: float = 0.0,
+        trainable_symbols: bool = True,
+        scale: float | None = None,
+        use_bias: bool = False,
+    ):
+        super().__init__()
+        if hidden_dim % n_heads != 0:
+            raise ValueError(f"hidden_dim ({hidden_dim}) must be divisible by n_heads ({n_heads})")
+        if symbol_dim % n_heads != 0:
+            raise ValueError(f"symbol_dim ({symbol_dim}) must be divisible by n_heads ({n_heads})")
+        if scale is not None and scale <= 0.0:
+            raise ValueError(f"scale must be positive when provided, got {scale}")
+
+        self.hidden_dim = hidden_dim
+        self.symbol_dim = symbol_dim
+        self.n_symbols = n_symbols
+        self.n_heads = n_heads
+        self.head_dim = hidden_dim // n_heads
+        self.symbol_head_dim = symbol_dim // n_heads
+        self.scale = 1.0 / math.sqrt(self.head_dim) if scale is None else scale
+
+        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias)
+        self.template_features = nn.Parameter(torch.empty(n_symbols, hidden_dim))
+        self.symbol_library = nn.Parameter(
+            torch.empty(n_symbols, symbol_dim),
+            requires_grad=trainable_symbols,
+        )
+        self.dropout = nn.Dropout(dropout)
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        nn.init.normal_(self.template_features, mean=0.0, std=0.9)
+        nn.init.normal_(self.symbol_library, mean=0.0, std=0.9)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+
+        queries = self.q_proj(x)
+        queries = queries.view(batch_size, seq_len, self.n_heads, self.head_dim)
+        queries = queries.transpose(1, 2)
+
+        keys = self.template_features.view(self.n_symbols, self.n_heads, self.head_dim)
+        keys = keys.transpose(0, 1).unsqueeze(0).expand(batch_size, -1, -1, -1)
+
+        values = self.symbol_library.view(self.n_symbols, self.n_heads, self.symbol_head_dim)
+        values = values.transpose(0, 1).unsqueeze(0).expand(batch_size, -1, -1, -1)
+
+        scores = torch.matmul(queries, keys.transpose(-2, -1)) * self.scale
+        weights = F.softmax(scores, dim=-1)
+        weights = self.dropout(weights)
+        retrieved = torch.matmul(weights, values)
+        retrieved = retrieved.transpose(1, 2).contiguous()
+        return retrieved.view(batch_size, seq_len, self.symbol_dim)
+
+
+class PositionalSymbolRetriever(nn.Module):
+    def __init__(
+        self,
+        symbol_dim: int,
+        max_len: int,
+        sinusoidal: bool = False,
+    ):
+        super().__init__()
+        self.symbol_dim = symbol_dim
+        self.max_len = max_len
+        self.sinusoidal = sinusoidal
+        pe_type = "sinusoidal" if sinusoidal else "learned"
+        self.position_encoder = PositionalEncoding(
+            embedding_dim=symbol_dim,
+            pe_type=pe_type,
+            max_len=max_len,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        pos_info = self.position_encoder.get_positional_info(seq_len, x.device)
+        if pos_info.embeddings is None:
+            raise ValueError(f"Positional symbol retrieval produced no embeddings for seq_len={seq_len}")
+        return pos_info.embeddings.unsqueeze(0).expand(batch_size, -1, -1)
+
+
+class RelativePositionalSymbolRetriever(nn.Module):
+    def __init__(
+        self,
+        symbol_dim: int,
+        max_rel_distance: int,
+        rope: bool = False,
+        theta: float = 10000.0,
+    ):
+        super().__init__()
+        if rope and symbol_dim % 2 != 0:
+            raise ValueError(f"RoPE relative symbols require even symbol_dim, got {symbol_dim}")
+        if theta <= 0.0:
+            raise ValueError(f"theta must be positive, got {theta}")
+        self.symbol_dim = symbol_dim
+        self.max_rel_distance = max_rel_distance
+        self.rope = rope
+        self.theta = theta
+        if rope:
+            self.register_buffer("_rope_relative_cache", torch.empty(0), persistent=False)
+        else:
+            self.position_encoder = PositionalEncoding(
+                embedding_dim=symbol_dim,
+                pe_type="relative",
+                max_len=max_rel_distance * 2,
+                max_rel_pos=max_rel_distance,
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        seq_len = x.shape[1]
+        if self.rope:
+            return self._rope_relative_symbols(seq_len, x.device, x.dtype)
+
+        pos_info = self.position_encoder.get_positional_info(seq_len, x.device)
+        if pos_info.rel_embeddings is None:
+            raise ValueError(f"Relative symbol retrieval produced no embeddings for seq_len={seq_len}")
+        return pos_info.rel_embeddings
+
+    def _rope_relative_symbols(
+        self,
+        seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        cached = self._rope_relative_cache
+        if (
+            cached.shape[0] >= seq_len
+            and cached.device == device
+            and cached.dtype == dtype
+        ):
+            return cached[:seq_len, :seq_len, :]
+
+        positions = torch.arange(seq_len, device=device)
+        distances = positions[None, :] - positions[:, None]
+        if self.max_rel_distance is not None:
+            distances = torch.clamp(distances, -self.max_rel_distance, self.max_rel_distance)
+        inv_freq = PositionalEncoding._rope_inv_freq(self.symbol_dim, self.theta, device=device)
+        phases = distances.to(torch.float32).unsqueeze(-1) * inv_freq
+        symbols = torch.empty(seq_len, seq_len, self.symbol_dim, device=device)
+        symbols[..., 0::2] = torch.cos(phases)
+        symbols[..., 1::2] = torch.sin(phases)
+        self._rope_relative_cache = symbols.to(dtype=dtype)
+        return self._rope_relative_cache
+
+
+class RelationalSymbolicAttentionRetriever(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        rel_n_heads: int,
+        symbolic_attn_n_heads: int,
+        n_symbols: int,
+        neighborhood_size: int = 2,
+        include_self: bool = False,
+        normalize_rels: bool = True,
+        dropout: float = 0.0,
+        trainable_symbols: bool = True,
+        rel_scale: float | None = None,
+        symbolic_attn_scale: float | None = None,
+        use_bias: bool = False,
+    ):
+        super().__init__()
+        if hidden_dim % rel_n_heads != 0:
+            raise ValueError(
+                f"hidden_dim ({hidden_dim}) must be divisible by rel_n_heads ({rel_n_heads})"
+            )
+        if rel_scale is not None and rel_scale <= 0.0:
+            raise ValueError(f"rel_scale must be positive when provided, got {rel_scale}")
+        if symbolic_attn_scale is not None and symbolic_attn_scale <= 0.0:
+            raise ValueError(
+                "symbolic_attn_scale must be positive when provided, "
+                f"got {symbolic_attn_scale}"
+            )
+
+        self.hidden_dim = hidden_dim
+        self.symbol_dim = symbol_dim
+        self.rel_n_heads = rel_n_heads
+        self.neighborhood_size = neighborhood_size
+        self.include_self = include_self
+        self.normalize_rels = normalize_rels
+        self.neighborhood_dim = neighborhood_size + (1 if include_self else 0)
+        self.rel_feature_dim = rel_n_heads * self.neighborhood_dim
+        rel_head_dim = hidden_dim // rel_n_heads
+        self.rel_scale = 1.0 / math.sqrt(rel_head_dim) if rel_scale is None else rel_scale
+
+        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias)
+        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias)
+        self.rel_to_hidden = nn.Linear(self.rel_feature_dim, hidden_dim, bias=True)
+        self.symbolic_attention = SymbolicAttentionRetriever(
+            hidden_dim=hidden_dim,
+            symbol_dim=symbol_dim,
+            n_symbols=n_symbols,
+            n_heads=symbolic_attn_n_heads,
+            dropout=dropout,
+            trainable_symbols=trainable_symbols,
+            scale=symbolic_attn_scale,
+            use_bias=use_bias,
+        )
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        nn.init.xavier_uniform_(self.q_proj.weight)
+        nn.init.xavier_uniform_(self.k_proj.weight)
+        if self.q_proj.bias is not None:
+            nn.init.zeros_(self.q_proj.bias)
+        if self.k_proj.bias is not None:
+            nn.init.zeros_(self.k_proj.bias)
+        nn.init.xavier_uniform_(self.rel_to_hidden.weight)
+        nn.init.zeros_(self.rel_to_hidden.bias)
+
+    def _compute_neighborhood_indices(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        positions = torch.arange(seq_len, device=device).unsqueeze(1)
+        # Decoder-LM adaptation: DSSL uses bidirectional neighborhoods, but this
+        # owned backend must keep symbol retrieval causal. Keep the old causal
+        # neighborhood order: current/immediate-past positions before older ones.
+        if self.include_self:
+            offsets = torch.arange(0, self.neighborhood_size + 1, device=device).unsqueeze(0)
+        else:
+            offsets = torch.arange(1, self.neighborhood_size + 1, device=device).unsqueeze(0)
+        return (positions - offsets).clamp(0, seq_len - 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        head_dim = self.hidden_dim // self.rel_n_heads
+
+        queries = self.q_proj(x).view(batch_size, seq_len, self.rel_n_heads, head_dim)
+        keys = self.k_proj(x).view(batch_size, seq_len, self.rel_n_heads, head_dim)
+        queries = queries.transpose(1, 2)
+        keys = keys.transpose(1, 2)
+
+        neighbor_indices = self._compute_neighborhood_indices(seq_len, x.device)
+        neighborhood_keys = keys[:, :, neighbor_indices]
+        neighborhood_relations = torch.einsum("bhid,bhijd->bhij", queries, neighborhood_keys)
+        if self.normalize_rels:
+            neighborhood_relations = F.softmax(neighborhood_relations * self.rel_scale, dim=-1)
+
+        neighborhood_relations = neighborhood_relations.permute(0, 2, 3, 1)
+        neighborhood_relations = neighborhood_relations.contiguous().view(
+            batch_size,
+            seq_len,
+            self.rel_feature_dim,
+        )
+        relational_features = self.rel_to_hidden(neighborhood_relations)
+        return self.symbolic_attention(relational_features)

masks.py

@@ -0,0 +1,61 @@
+"""Attention-mask helpers shared by repo-owned decoder LMs."""
+
+from typing import Optional
+
+import torch
+
+
+def build_decoder_attention_mask(
+    input_ids: torch.Tensor,
+    pad_token_id: int,
+    eos_token_id: int,
+    sequence_boundary_policy: str,
+    attention_mask: Optional[torch.Tensor] = None,
+    segment_boundary_token_id: Optional[int] = None,
+    bidirectional: bool = False,
+) -> torch.Tensor:
+    if input_ids.ndim != 2:
+        raise ValueError(
+            "input_ids must be rank-2 [batch, seq], "
+            f"got shape {tuple(input_ids.shape)}"
+        )
+    _, seq_len = input_ids.shape
+    if attention_mask is None:
+        valid_tokens = input_ids != pad_token_id
+    else:
+        valid_tokens = attention_mask.bool()
+
+    query_mask = valid_tokens.unsqueeze(2)
+    key_mask = valid_tokens.unsqueeze(1)
+    if bidirectional:
+        mask = query_mask & key_mask
+    else:
+        directionality = torch.tril(
+            torch.ones(seq_len, seq_len, dtype=torch.bool, device=input_ids.device)
+        ).unsqueeze(0)
+        mask = directionality & query_mask & key_mask
+
+    if sequence_boundary_policy == "none":
+        return mask
+    if sequence_boundary_policy == "segment_document":
+        if segment_boundary_token_id is None:
+            raise ValueError(
+                "segment_boundary_token_id is required when "
+                "sequence_boundary_policy='segment_document'"
+            )
+        # The boundary token starts the next segment: cumsum increments on the
+        # boundary position, so the marker attends with the following tokens.
+        segment_ids = torch.cumsum(
+            input_ids == segment_boundary_token_id, dim=1
+        )
+        same_segment = segment_ids.unsqueeze(1) == segment_ids.unsqueeze(2)
+        return mask & same_segment
+    if sequence_boundary_policy != "eos_document":
+        raise ValueError(f"Unsupported sequence_boundary_policy: {sequence_boundary_policy}")
+
+    # Next-token training predicts EOS from the preceding document token.
+    # Once EOS is present as an input token, it starts the next segment so the
+    # following document is not predicted with prior-document context.
+    document_ids = torch.cumsum(input_ids == eos_token_id, dim=1)
+    same_document = document_ids.unsqueeze(1) == document_ids.unsqueeze(2)
+    return mask & same_document

mlm.py

@@ -0,0 +1,95 @@
+"""MLM head for the hybrid MLM/CLM objective.
+
+Ported from temp/gpt-bert-main/pretraining/model.py MaskClassifier, but
+configurable to match the trunk's customizability:
+- norm_type: "layernorm" or "rmsnorm" (reuses the existing field).
+- ffn_activation: "gelu", "swiglu", "relu", or "identity" (reuses the
+  existing field). The head's nonlinearity follows this choice.
+- use_bias_ffn: whether the linear layers in the head use bias.
+- weight tying: the final linear is tied to the word embedding, matching
+  gpt-bert's MaskClassifier and our lm_head.
+
+The head is controlled by model config (mlm_head_enabled) so checkpoint
+loading and HF export remain strict and honest.
+"""
+
+import torch
+import torch.nn as nn
+
+from .transformer_components import RMSNorm, FeedForward
+
+
+def _create_norm(norm_type: str, hidden_dim: int) -> nn.Module:
+    if norm_type == "layernorm":
+        return nn.LayerNorm(hidden_dim)
+    if norm_type == "rmsnorm":
+        return RMSNorm(hidden_dim)
+    raise ValueError(f"Unsupported norm_type: {norm_type}")
+
+
+class MaskClassifier(nn.Module):
+    """MLM prediction head with configurable norm and activation.
+
+    Architecture (matching gpt-bert's MaskClassifier):
+        norm -> FeedForward (activation, bias configurable) -> norm -> dropout -> linear (tied)
+
+    The final linear is tied to the word embedding weight, matching
+    gpt-bert's MaskClassifier and our lm_head.
+
+    Args:
+        hidden_dim: Model hidden dimension.
+        vocab_size: Vocabulary size for the output projection.
+        norm_type: "layernorm" or "rmsnorm".
+        ffn_activation: "gelu", "swiglu", "relu", or "identity".
+        use_bias_ffn: Whether the intermediate linear layers use bias.
+        dropout: Dropout probability.
+        word_embedding: The word embedding weight to tie the final linear to.
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        vocab_size: int,
+        norm_type: str,
+        ffn_activation: str,
+        use_bias_ffn: bool,
+        dropout: float,
+        word_embedding: nn.Parameter,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.vocab_size = vocab_size
+        self.norm_type = norm_type
+        self.ffn_activation = ffn_activation
+        self.use_bias_ffn = use_bias_ffn
+        self.dropout_p = dropout
+
+        self.norm1 = _create_norm(norm_type, hidden_dim)
+        self.feed_forward = FeedForward(
+            input_dim=hidden_dim,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            activation=ffn_activation,
+            use_bias=use_bias_ffn,
+        )
+        self.norm2 = _create_norm(norm_type, hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.linear_out = nn.Linear(hidden_dim, vocab_size, bias=False)
+        self.linear_out.weight = word_embedding
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the MLM head to contextualized embeddings.
+
+        Args:
+            x: Contextualized embeddings [batch, seq_len, hidden_dim] or
+               flattened [num_masked, hidden_dim].
+
+        Returns:
+            Logits [batch, seq_len, vocab_size] or [num_masked, vocab_size].
+        """
+        x = self.norm1(x)
+        x = self.feed_forward(x)
+        x = self.norm2(x)
+        x = self.dropout(x)
+        x = self.linear_out(x)
+        return x

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37693c16dbb081d2adff7f0de2db3ecca7c22f6fa710007ee1caaed7b7072f4b
+size 493209616

modeling_dat.py

@@ -0,0 +1,290 @@
+"""HuggingFace PreTrainedModel wrapper for the dual-attention (DAT) decoder LM.
+
+Source-of-truth copy. scripts/convert_dat_to_hf.py copies it into a generated HF
+repository together with configuration_dat.py and the flattened model source
+(dat_config.py, dat_core.py, dat_symbols.py, dat_lm.py, transformer_core.py,
+transformer_components.py).
+
+Why every model module is imported directly here: with trust_remote_code on a
+local directory, transformers 4.46.3 only copies the entry module's *direct*
+relative imports into its dynamic-module cache (it does not recurse). Importing
+all flattened modules here forces every one of them to be copied, after which
+their own relative imports resolve because the files are co-located.
+"""
+
+import collections
+
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput, MaskedLMOutput, SequenceClassifierOutput
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+# Force every flattened module to be copied into the dynamic-module cache.
+from .attention_config import SUPPORTED_SEQUENCE_BOUNDARY_POLICIES  # noqa: F401
+from .masks import build_decoder_attention_mask  # noqa: F401
+from .mlm import MaskClassifier  # noqa: F401
+from .transformer_components import FeedForward, PositionalEncoding  # noqa: F401
+from .transformer_core import SelfAttention  # noqa: F401
+from .dat_symbols import SymbolicAttentionRetriever  # noqa: F401
+from .dat_core import RelationalAttention  # noqa: F401
+from .dat_config import DatLMConfig
+from .dat_lm import DatDecoderLM
+from .configuration_dat import DatConfig, DAT_LM_FIELDS
+
+
+def _refresh_rope_buffers(model: DatDecoderLM, rope_theta: float) -> None:
+    pe = model.position_encoder
+    if hasattr(pe, "pe_type") and pe.pe_type == "rope":
+        cos, sin = pe._precompute_rope_freqs(pe.embedding_dim, pe.max_len, rope_theta)
+        pe.register_buffer("rope_cos", cos, persistent=False)
+        pe.register_buffer("rope_sin", sin, persistent=False)
+        pe._rope_uninitialized = True
+
+
+def _collect_shared_weight_keys(module: torch.nn.Module) -> dict[str, str]:
+    # remove_duplicate=False is essential: by default named_parameters yields
+    # each shared tensor only once, which would hide exactly the duplicates we
+    # need to declare (and which state_dict still emits under every name).
+    names_by_storage: dict[int, list[str]] = collections.defaultdict(list)
+    for name, parameter in module.named_parameters(remove_duplicate=False):
+        names_by_storage[id(parameter)].append(name)
+    shared_keys: dict[str, str] = {}
+    for names in names_by_storage.values():
+        if len(names) > 1:
+            names = sorted(names)
+            # Prefer token_embeddings.weight as the canonical source
+            # (HF convention: output weights tied to input embeddings).
+            source = names[0]
+            for n in names:
+                if n.endswith("token_embeddings.weight"):
+                    source = n
+                    break
+            for target in names:
+                if target != source:
+                    shared_keys[target] = source
+    return shared_keys
+
+
+class _ExpandedTiedWeightsMixin:
+    def get_expanded_tied_weights_keys(self, all_submodels: bool = False) -> dict:
+        # HF's default returns {} when tie_word_embeddings=False, which would
+        # discard the mlm_head.linear_out <-> lm_head tie for NextLat models
+        # (tie_lm_head=False). Return the dynamically computed mapping directly
+        # so all shared weights are properly tied on load.
+        return dict(self._tied_weights_keys)
+
+
+class DatModel(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = ""
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        # Work around accelerate/transformers leaving persistent=False buffers uninitialized
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> BaseModelOutput:
+        _, hidden_states = self.model.encode_for_objective(input_ids, attention_mask=attention_mask)
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class DatForCausalLM(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = {"model.lm_head.weight": "model.token_embeddings.weight"}
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        # Work around accelerate/transformers leaving persistent=False buffers uninitialized
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def get_output_embeddings(self) -> torch.nn.Module:
+        return self.model.lm_head
+
+    def set_output_embeddings(self, new_embeddings: torch.nn.Module) -> None:
+        self.model.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ) -> CausalLMOutput:
+        logits, _ = self.model(input_ids, attention_mask=attention_mask)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = torch.nn.functional.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits)
+
+
+class DatForMaskedLM(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def get_output_embeddings(self) -> torch.nn.Module:
+        return self.model.mlm_head.linear_out
+
+    def set_output_embeddings(self, new_embeddings: torch.nn.Module) -> None:
+        self.model.mlm_head.linear_out = new_embeddings
+        # Keep lm_head in sync with mlm_head.linear_out so that
+        # resize_token_embeddings updates both tied heads.
+        self.model.lm_head.weight = new_embeddings.weight
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ) -> MaskedLMOutput | tuple:
+        return_dict = kwargs.pop("return_dict", None)
+        if return_dict is None:
+            return_dict = self.config.return_dict
+        if kwargs:
+            raise TypeError(f"Unexpected keyword argument(s): {list(kwargs)}")
+        logits = self.model.forward_mlm(input_ids, attention_mask=attention_mask)
+
+        loss = None
+        if labels is not None:
+            loss = CrossEntropyLoss(ignore_index=-100)(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1),
+            )
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(loss=loss, logits=logits)
+
+class DatForSequenceClassification(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        self.num_labels = getattr(config, "num_labels", 2)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self.score = torch.nn.Linear(config.hidden_dim, self.num_labels, bias=False)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ) -> SequenceClassifierOutput:
+        _, hidden_states = self.model.encode_for_objective(input_ids, attention_mask=attention_mask)
+        logits = self.score(hidden_states)
+
+        batch_size = input_ids.shape[0]
+
+        if self.config.pad_token_id is None:
+            last_non_pad_token = -1
+        else:
+            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+            if (non_pad_mask.sum(-1) == 0).any().item():
+                raise ValueError("Cannot pool sequence-classification logits for all-padding sequences")
+            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+        loss = None
+        if labels is not None:
+            if getattr(self.config, "problem_type", None) is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=pooled_logits,
+            hidden_states=hidden_states,
+        )

special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": "<s>", "eos_token": "</s>", "unk_token": "<unk>", "sep_token": "</s>", "pad_token": "<pad>", "cls_token": "<s>", "mask_token": "<mask>"}

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+    "tokenizer_class": "PreTrainedTokenizerFast",
+    "bos_token": "<s>",
+    "eos_token": "</s>",
+    "unk_token": "<unk>",
+    "sep_token": "</s>",
+    "pad_token": "<pad>",
+    "cls_token": "<s>",
+    "mask_token": "<mask>"
+}

transformer_components.py

@@ -0,0 +1,282 @@
+"""Minimal transformer components for decoder LM training."""
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+@dataclass(frozen=True)
+class PositionalInfo:
+    pe_type: str
+    embeddings: Optional[Tensor] = None
+    rope_freqs: Optional[Tuple[Tensor, Tensor]] = None
+    rel_embeddings: Optional[Tensor] = None
+    apply_to_embeddings: bool = False
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x: Tensor) -> Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.dtype in (torch.float16, torch.bfloat16):
+            output = self._norm(x.float()).type_as(x)
+        else:
+            output = self._norm(x)
+        return output * self.weight
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        use_bias: bool = True,
+    ):
+        super().__init__()
+        self.activation = activation
+        self.use_swiglu = activation == "swiglu"
+        self.dropout = nn.Dropout(dropout)
+        if self.use_swiglu:
+            self.w_gate = nn.Linear(input_dim, hidden_dim, bias=use_bias)
+            self.w_up = nn.Linear(input_dim, hidden_dim, bias=use_bias)
+            self.w_down = nn.Linear(hidden_dim, input_dim, bias=use_bias)
+            self.activation_fn = nn.SiLU()
+        else:
+            self.linear1 = nn.Linear(input_dim, hidden_dim, bias=use_bias)
+            self.linear2 = nn.Linear(hidden_dim, input_dim, bias=use_bias)
+            if activation == "gelu":
+                self.activation_fn = nn.GELU()
+            elif activation == "relu":
+                self.activation_fn = nn.ReLU()
+            elif activation == "identity":
+                self.activation_fn = nn.Identity()
+            else:
+                raise ValueError(f"Unsupported ffn activation: {activation}")
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight, gain=1.0)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.use_swiglu:
+            gate = self.activation_fn(self.w_gate(x))
+            up = self.w_up(x)
+            x = gate * up
+            x = self.dropout(x)
+            return self.w_down(x)
+
+        x = self.linear1(x)
+        x = self.activation_fn(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        pe_type: str,
+        max_len: int,
+        theta: float = 10000.0,
+        max_rel_pos: Optional[int] = None,
+        init_range: float = 0.15,
+    ):
+        super().__init__()
+        if embedding_dim <= 0:
+            raise ValueError(f"embedding_dim must be positive, got {embedding_dim}")
+        if max_len <= 0:
+            raise ValueError(f"max_len must be positive, got {max_len}")
+        if pe_type not in {"sinusoidal", "learned", "relative", "rope", "none"}:
+            raise ValueError(f"Unsupported pe_type: {pe_type}")
+        if theta <= 0.0:
+            raise ValueError(f"theta must be positive, got {theta}")
+        if init_range <= 0.0:
+            raise ValueError(f"init_range must be positive, got {init_range}")
+
+        self.embedding_dim = embedding_dim
+        self.pe_type = pe_type
+        self.max_len = max_len
+
+        if pe_type == "sinusoidal":
+            if embedding_dim % 2 != 0:
+                raise ValueError(f"Sinusoidal encoding requires even embedding_dim, got {embedding_dim}")
+            pe = self._precompute_sinusoidal_embeddings(
+                embedding_dim,
+                max_len,
+                device=torch.empty(0).device,
+            )
+            self.register_buffer("sinusoidal_embeddings", pe, persistent=False)
+            self._sinusoidal_uninitialized = True
+        elif pe_type == "learned":
+            self.position_embeddings = nn.Embedding(max_len, embedding_dim)
+            nn.init.uniform_(self.position_embeddings.weight, -init_range, init_range)
+        elif pe_type == "relative":
+            self.max_relative_position = max_rel_pos if max_rel_pos is not None else max_len // 2
+            if self.max_relative_position <= 0:
+                raise ValueError(
+                    f"max_relative_position must be positive, got {self.max_relative_position}"
+                )
+            num_embeddings = 2 * self.max_relative_position + 1
+            self.rel_pos_embeddings_table = nn.Embedding(num_embeddings, embedding_dim)
+            nn.init.uniform_(self.rel_pos_embeddings_table.weight, -init_range, init_range)
+        elif pe_type == "rope":
+            if embedding_dim % 2 != 0:
+                raise ValueError(f"RoPE requires even embedding_dim, got {embedding_dim}")
+            self.theta = theta
+            cos, sin = self._precompute_rope_freqs(embedding_dim, max_len, theta)
+            self.register_buffer("rope_cos", cos, persistent=False)
+            self.register_buffer("rope_sin", sin, persistent=False)
+            self._rope_uninitialized = True
+
+    def get_positional_info(self, seq_len: int, device: torch.device) -> PositionalInfo:
+        if self.pe_type == "none":
+            return PositionalInfo(pe_type="none")
+        if self.pe_type == "sinusoidal":
+            if seq_len > self.sinusoidal_embeddings.size(0):
+                raise ValueError(
+                    f"seq_len {seq_len} exceeds precomputed sinusoidal length {self.sinusoidal_embeddings.size(0)}"
+                )
+            if (
+                getattr(self, "_sinusoidal_uninitialized", False)
+                or self.sinusoidal_embeddings.device.type == "meta"
+            ):
+                embeddings = self._precompute_sinusoidal_embeddings(
+                    self.embedding_dim,
+                    self.max_len,
+                    device=device,
+                )
+                self.register_buffer("sinusoidal_embeddings", embeddings, persistent=False)
+                self._sinusoidal_uninitialized = False
+            return PositionalInfo(
+                pe_type="sinusoidal",
+                embeddings=self.sinusoidal_embeddings[:seq_len].to(device),
+                apply_to_embeddings=True,
+            )
+        if self.pe_type == "learned":
+            if seq_len > self.max_len:
+                raise ValueError(
+                    f"seq_len {seq_len} exceeds maximum learned position length {self.max_len}"
+                )
+            positions = torch.arange(seq_len, dtype=torch.long, device=device)
+            return PositionalInfo(
+                pe_type="learned",
+                embeddings=self.position_embeddings(positions),
+                apply_to_embeddings=True,
+            )
+        if self.pe_type == "relative":
+            return PositionalInfo(
+                pe_type="relative",
+                rel_embeddings=self._generate_relative_embeddings(seq_len, device),
+                apply_to_embeddings=False,
+            )
+        if seq_len > self.rope_cos.size(0):
+            # Dynamically extend RoPE buffer (standard practice for extrapolation)
+            cos, sin = self._precompute_rope_freqs(self.embedding_dim, seq_len, self.theta)
+            self.register_buffer("rope_cos", cos, persistent=False)
+            self.register_buffer("rope_sin", sin, persistent=False)
+            self._rope_uninitialized = False
+        if getattr(self, "_rope_uninitialized", True) or torch.isnan(self.rope_cos[0, 0]):
+            cos, sin = self._precompute_rope_freqs(self.embedding_dim, self.max_len, self.theta)
+            self.rope_cos.copy_(cos)
+            self.rope_sin.copy_(sin)
+            self._rope_uninitialized = False
+        return PositionalInfo(
+            pe_type="rope",
+            rope_freqs=(
+                self.rope_cos[:seq_len].to(device),
+                self.rope_sin[:seq_len].to(device),
+            ),
+            apply_to_embeddings=False,
+        )
+
+    @staticmethod
+    def _rope_inv_freq(dim: int, theta: float, device: torch.device | str) -> Tensor:
+        return 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
+
+    @staticmethod
+    def _precompute_sinusoidal_embeddings(
+        dim: int,
+        end: int,
+        device: torch.device | str,
+    ) -> Tensor:
+        pe = torch.zeros(end, dim, device=device)
+        position = torch.arange(0, end, dtype=torch.float, device=device).unsqueeze(1)
+        div_term = PositionalEncoding._rope_inv_freq(dim, 10000.0, device=device)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe
+
+    @staticmethod
+    def _precompute_rope_freqs(dim: int, end: int, theta: float) -> Tuple[Tensor, Tensor]:
+        freqs = PositionalEncoding._rope_inv_freq(dim, theta, device="cpu")
+        positions = torch.arange(end, device="cpu")
+        freqs = torch.outer(positions, freqs).float()
+        return torch.cos(freqs), torch.sin(freqs)
+
+    def _generate_relative_embeddings(self, seq_len: int, device: torch.device) -> Tensor:
+        range_q = torch.arange(seq_len, device=device)
+        range_k = torch.arange(seq_len, device=device)
+        distance_mat = range_k[None, :] - range_q[:, None]
+        distance_mat_clipped = torch.clamp(
+            distance_mat,
+            -self.max_relative_position,
+            self.max_relative_position,
+        )
+        final_indices = distance_mat_clipped + self.max_relative_position
+        return self.rel_pos_embeddings_table(final_indices.long())
+
+
+def reshape_for_broadcast(freqs_cis: Tensor, x: Tensor) -> Tensor:
+    ndim = x.ndim
+    if ndim < 2:
+        raise ValueError(f"Input tensor x must have at least 2 dimensions, got {ndim}")
+    if x.shape[-2] != freqs_cis.shape[0] or x.shape[-1] != freqs_cis.shape[-1]:
+        raise ValueError(
+            f"Shape mismatch for RoPE broadcasting: freqs_cis {freqs_cis.shape}, x {x.shape}"
+        )
+    shape = [1] * ndim
+    shape[-2] = x.shape[-2]
+    shape[-1] = x.shape[-1]
+    return freqs_cis.view(shape)
+
+
+def apply_rotary_emb(xq: Tensor, xk: Tensor, freqs_cos: Tensor, freqs_sin: Tensor) -> Tuple[Tensor, Tensor]:
+    if xq.shape[-1] % 2 != 0:
+        raise ValueError(f"Query feature dimension must be even for RoPE, got {xq.shape[-1]}")
+    if xk.shape[-1] % 2 != 0:
+        raise ValueError(f"Key feature dimension must be even for RoPE, got {xk.shape[-1]}")
+
+    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
+    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
+
+    freqs_cos_q = reshape_for_broadcast(freqs_cos[: xq.shape[-2]], xq_r)
+    freqs_sin_q = reshape_for_broadcast(freqs_sin[: xq.shape[-2]], xq_r)
+    freqs_cos_k = reshape_for_broadcast(freqs_cos[: xk.shape[-2]], xk_r)
+    freqs_sin_k = reshape_for_broadcast(freqs_sin[: xk.shape[-2]], xk_r)
+
+    xq_out_r = xq_r * freqs_cos_q - xq_i * freqs_sin_q
+    xq_out_i = xq_r * freqs_sin_q + xq_i * freqs_cos_q
+    xk_out_r = xk_r * freqs_cos_k - xk_i * freqs_sin_k
+    xk_out_i = xk_r * freqs_sin_k + xk_i * freqs_cos_k
+
+    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2)
+    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2)
+
+    return xq_out.type_as(xq), xk_out.type_as(xk)

transformer_core.py

@@ -0,0 +1,327 @@
+"""Minimal self-attention core for decoder language models."""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from .transformer_components import PositionalInfo, apply_rotary_emb
+
+
+def _activate_scores(scores: torch.Tensor, activation: str) -> torch.Tensor:
+    if activation == "softmax":
+        return torch.softmax(scores, dim=-1)
+    if activation == "identity":
+        return scores
+    if activation == "relu":
+        return torch.relu(scores)
+    if activation == "tanh":
+        return torch.tanh(scores)
+    if activation == "sigmoid":
+        return torch.sigmoid(scores)
+    if activation == "gelu":
+        return torch.nn.functional.gelu(scores)
+    raise ValueError(f"Unsupported attention activation: {activation}")
+
+
+class MultiHeadAttentionBase(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        output_dim: int,
+        key_dim: Optional[int] = None,
+        value_dim: Optional[int] = None,
+        n_heads: int = 4,
+        hidden_dim: int = 64,
+        dropout: float = 0.0,
+        total_n_heads: Optional[int] = None,
+        activation: str = "softmax",
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        super().__init__()
+        self.query_dim = query_dim
+        self.key_dim = query_dim if key_dim is None else key_dim
+        self.value_dim = self.key_dim if value_dim is None else value_dim
+        self.output_dim = output_dim
+        self.n_heads = n_heads
+        self.hidden_dim = hidden_dim
+        self.total_n_heads = n_heads if total_n_heads is None else total_n_heads
+        self.activation = activation
+        self.use_bias_qkv = use_bias_qkv
+        self.use_bias_out = use_bias_out
+
+        if self.total_n_heads <= 0:
+            raise ValueError(f"total_n_heads must be positive, got {self.total_n_heads}")
+        if hidden_dim % self.total_n_heads != 0:
+            raise ValueError(
+                f"hidden_dim ({hidden_dim}) must be divisible by total_n_heads ({self.total_n_heads})"
+            )
+
+        self.head_dim = hidden_dim // self.total_n_heads
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+        projected_dim = n_heads * self.head_dim
+        self.q_proj = nn.Linear(self.query_dim, projected_dim, bias=use_bias_qkv)
+        self.k_proj = nn.Linear(self.key_dim, projected_dim, bias=use_bias_qkv)
+        self.v_proj = nn.Linear(self.value_dim, projected_dim, bias=use_bias_qkv)
+        self.o_proj = nn.Linear(projected_dim, output_dim, bias=use_bias_out)
+        self.dropout = nn.Dropout(dropout)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.last_attn_weights = None
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        gain = 1.0 / math.sqrt(2.0)
+        for projection in (self.q_proj, self.k_proj, self.v_proj, self.o_proj):
+            nn.init.xavier_uniform_(projection.weight, gain=gain)
+            if projection.bias is not None:
+                nn.init.zeros_(projection.bias)
+
+    def _reshape_for_multihead(self, x: torch.Tensor, batch_size: int, seq_len: int) -> torch.Tensor:
+        x = x.view(batch_size, seq_len, self.n_heads, self.head_dim)
+        return x.transpose(1, 2)
+
+    def _process_mask(self, mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        if mask is None:
+            return None
+        if mask.dim() == 2:
+            if mask.shape[0] != mask.shape[1]:
+                raise ValueError(
+                    "2D decoder attention masks must be square [seq, seq]; "
+                    f"got {tuple(mask.shape)}. Expand padding masks at the LM boundary."
+                )
+            return mask.bool().unsqueeze(0).unsqueeze(0)
+        if mask.dim() == 3:
+            return mask.bool().unsqueeze(1)
+        raise ValueError(f"Mask must be 2D or 3D, got {mask.dim()}D")
+
+    def _compute_attn_scores(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> torch.Tensor:
+        if pos_info is not None and pos_info.rope_freqs is not None:
+            freqs_cos, freqs_sin = pos_info.rope_freqs
+            q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin)
+        return torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+    def _apply_activation_and_mask(self, scores: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:
+        processed_mask = self._process_mask(mask)
+        if processed_mask is not None and self.activation in {"softmax", "sigmoid", "tanh"}:
+            scores.masked_fill_(~processed_mask, torch.finfo(scores.dtype).min)
+        weights = _activate_scores(scores, self.activation)
+        if processed_mask is not None:
+            weights = weights.masked_fill(~processed_mask, 0.0)
+        return weights
+
+    def _apply_mask(self, scores: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:
+        return self._apply_activation_and_mask(scores, mask)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: Optional[torch.Tensor] = None,
+        value: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, _ = query.shape
+        key = query if key is None else key
+        value = key if value is None else value
+        key_len = key.shape[1]
+
+        if key.shape[0] != batch_size or value.shape[0] != batch_size:
+            raise ValueError(
+                f"Batch size mismatch: query={batch_size}, key={key.shape[0]}, value={value.shape[0]}"
+            )
+        if key.shape[1] != value.shape[1]:
+            raise ValueError(
+                f"Key and value sequence length mismatch: {key.shape[1]} vs {value.shape[1]}"
+            )
+
+        q = self.q_proj(query)
+        k = self.k_proj(key)
+        v = self.v_proj(value)
+
+        q = self._reshape_for_multihead(q, batch_size, seq_len)
+        k = self._reshape_for_multihead(k, batch_size, key_len)
+        v = self._reshape_for_multihead(v, batch_size, key_len)
+
+        attn_scores = self._compute_attn_scores(q, k, pos_info)
+        attn_weights = self._apply_activation_and_mask(attn_scores, mask)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, self.n_heads * self.head_dim)
+        attn_output = self.o_proj(attn_output)
+        attn_output = self.dropout(attn_output)
+
+        self.last_attn_weights = attn_weights.detach()
+        return attn_output, attn_weights
+
+
+class SelfAttention(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        input_dim: int,
+        n_heads: int = 4,
+        hidden_dim: int = 64,
+        dropout: float = 0.0,
+        supports_relative: bool = False,
+        total_n_heads: Optional[int] = None,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        head_dim = hidden_dim // (n_heads if total_n_heads is None else total_n_heads)
+        super().__init__(
+            query_dim=input_dim,
+            output_dim=n_heads * head_dim,
+            key_dim=input_dim,
+            value_dim=input_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            total_n_heads=total_n_heads,
+            activation="softmax",
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.supports_relative = supports_relative
+        if self.supports_relative:
+            self.rel_k_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            self.rel_v_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            nn.init.xavier_uniform_(self.rel_k_proj.weight)
+            nn.init.xavier_uniform_(self.rel_v_proj.weight)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.supports_relative and pos_info is not None and pos_info.rel_embeddings is not None:
+            batch_size, seq_len, _ = x.shape
+            rel_embeddings = pos_info.rel_embeddings
+            if rel_embeddings.shape[:2] != (seq_len, seq_len):
+                raise ValueError(
+                    f"Relative embeddings shape {rel_embeddings.shape} does not match sequence length {seq_len}"
+                )
+            if rel_embeddings.shape[2] != self.head_dim:
+                raise ValueError(
+                    f"Relative embeddings dim {rel_embeddings.shape[2]} does not match head_dim {self.head_dim}"
+                )
+
+            rel_k = self.rel_k_proj(rel_embeddings)
+            rel_v = self.rel_v_proj(rel_embeddings)
+
+            q = self.q_proj(x)
+            k = self.k_proj(x)
+            v = self.v_proj(x)
+
+            q = self._reshape_for_multihead(q, batch_size, seq_len)
+            k = self._reshape_for_multihead(k, batch_size, seq_len)
+            v = self._reshape_for_multihead(v, batch_size, seq_len)
+
+            attn_scores = self._compute_attn_scores(q, k, pos_info)
+            rel_scores = torch.einsum("bhid,ijd->bhij", q, rel_k) * self.scale
+            attn_scores = attn_scores + rel_scores
+
+            attn_weights = self._apply_mask(attn_scores, mask)
+            attn_weights = self.attn_dropout(attn_weights)
+
+            attn_output = torch.matmul(attn_weights, v)
+            rel_output = torch.einsum("bhij,ijd->bhid", attn_weights, rel_v)
+            attn_output = attn_output + rel_output
+
+            attn_output = attn_output.transpose(1, 2).contiguous()
+            attn_output = attn_output.view(batch_size, seq_len, self.n_heads * self.head_dim)
+            attn_output = self.o_proj(attn_output)
+            attn_output = self.dropout(attn_output)
+            self.last_attn_weights = attn_weights.detach()
+            return attn_output, attn_weights
+
+        return super().forward(query=x, mask=mask, pos_info=pos_info)
+
+
+class CrossAttention(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        context_dim: Optional[int] = None,
+        n_heads: int = 8,
+        hidden_dim: int = 64,
+        dropout: float = 0.0,
+        supports_relative: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        resolved_context_dim = input_dim if context_dim is None else context_dim
+        super().__init__(
+            query_dim=input_dim,
+            key_dim=resolved_context_dim,
+            value_dim=resolved_context_dim,
+            output_dim=output_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            activation="softmax",
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.supports_relative = supports_relative
+        if supports_relative:
+            self.rel_k_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            self.rel_v_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            nn.init.xavier_uniform_(self.rel_k_proj.weight)
+            nn.init.xavier_uniform_(self.rel_v_proj.weight)
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        context: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.supports_relative and pos_info is not None and pos_info.rel_embeddings is not None:
+            batch_size, seq_len_q, _ = inputs.shape
+            _, seq_len_k, _ = context.shape
+            rel_embeddings = pos_info.rel_embeddings
+            if rel_embeddings.shape[:2] != (seq_len_q, seq_len_k):
+                raise ValueError(
+                    f"Relative embeddings shape {rel_embeddings.shape} does not match "
+                    f"sequence lengths {(seq_len_q, seq_len_k)}"
+                )
+            if rel_embeddings.shape[2] != self.head_dim:
+                raise ValueError(
+                    f"Relative embeddings dim {rel_embeddings.shape[2]} does not match head_dim {self.head_dim}"
+                )
+
+            rel_k = self.rel_k_proj(rel_embeddings)
+            rel_v = self.rel_v_proj(rel_embeddings)
+            q = self._reshape_for_multihead(self.q_proj(inputs), batch_size, seq_len_q)
+            k = self._reshape_for_multihead(self.k_proj(context), batch_size, seq_len_k)
+            v = self._reshape_for_multihead(self.v_proj(context), batch_size, seq_len_k)
+            attn_scores = self._compute_attn_scores(q, k, pos_info)
+            attn_scores = attn_scores + torch.einsum("bhid,ijd->bhij", q, rel_k) * self.scale
+            attn_weights = self._apply_activation_and_mask(attn_scores, mask)
+            attn_weights = self.attn_dropout(attn_weights)
+            self.last_attn_weights = attn_weights.detach()
+            attn_output = torch.matmul(attn_weights, v)
+            attn_output = attn_output + torch.einsum("bhij,ijd->bhid", attn_weights, rel_v)
+            attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, -1)
+            attn_output = self.o_proj(attn_output)
+            attn_output = self.dropout(attn_output)
+            return attn_output, attn_weights
+
+        return super().forward(
+            query=inputs,
+            key=context,
+            value=context,
+            mask=mask,
+            pos_info=pos_info,
+        )