from __future__ import annotations import json from dataclasses import asdict, dataclass, fields from pathlib import Path from typing import Any LINEAR_LAYER_TYPES = frozenset({"gdn2", "kda", "diag_kdn", "bigdn2"}) ATTENTION_LAYER_TYPES = frozenset({"full_attention", "gated_mla", "parallax", "bidirectional"}) # Gated-residual token-mixer pairs: Senryu-2's causal one, and Fugue's bidirectional one. GATED_RESIDUAL_MIXERS = frozenset({("gdn2", "parallax"), ("bigdn2", "bidirectional")}) FUGUE_MODEL_TYPE = "fugue" RESIDUAL_MODES = frozenset({"pre_norm", "gated_residual"}) DIAG_KDN_HEAD_DIMS = (64, 128) @dataclass(frozen=True) class TinyGDNConfig: architecture: str = "TinyGDNForCausalLM" model_type: str = "tiny_gdn" vocab_size: int = 49_152 # Deep-thin sizing is deliberate: controlled sub-billion studies find # depth materially more valuable than width around the 125M-150M scale. hidden_size: int = 512 intermediate_size: int = 1_472 num_hidden_layers: int = 32 num_attention_heads: int = 4 num_key_value_heads: int = 1 attention_head_dim: int = 128 full_attention_interval: int = 4 attention_dropout: float = 0.0 partial_rotary_factor: float = 0.5 rope_theta: float = 1_000_000.0 linear_num_heads: int = 4 linear_num_value_heads: int = 4 linear_head_dim: int = 128 linear_expand_v: float = 1.0 linear_conv_kernel_dim: int = 4 allow_negative_eigenvalues: bool = False linear_layer_type: str = "gdn2" attention_layer_type: str = "full_attention" mlp_activation: str = "swiglu" use_block_attn_res: bool = False kda_safe_gate: bool = True kda_lower_bound: float = -5.0 mla_q_lora_rank: int | None = 512 mla_kv_lora_rank: int = 256 mla_qk_nope_head_dim: int = 128 mla_v_head_dim: int = 128 situ_gate_cap: float = 4.0 situ_up_cap: float = 25.0 # "gated_residual" replaces pre-norm with Qwen's multi-branch gated # residual stream: each sublayer reads a gated mix of `residual_branches` # normalized branches and writes its output back with per-branch scales. residual_mode: str = "pre_norm" residual_branches: int = 4 residual_gate_rank: int = 128 # GDN-2 output gate as one full-rank projection (Kimi K3) instead of fla's # low-rank pair through head_v_dim. linear_full_rank_gate: bool = False # Hashed n-gram memory (Engram): `ngram_heads_per_order` tables per order, # each with a distinct prime row count near `ngram_table_rows`, injected # before block `ngram_layer_index`. An empty `ngram_orders` disables it. ngram_orders: tuple[int, ...] = () ngram_heads_per_order: int = 4 ngram_table_rows: int = 2_097_152 ngram_embedding_dim: int = 128 ngram_layer_index: int = 1 ngram_conv_kernel: int = 4 max_position_embeddings: int = 32_768 training_sequence_length: int = 2_048 rms_norm_eps: float = 1e-6 initializer_range: float = 0.02 tie_word_embeddings: bool = True shared_layer_indices: tuple[int, ...] = () # MTP is an opt-in ablation at this scale; static MTP is not assumed to # improve a 150M model without a controlled pilot. mtp_num_heads: int = 0 mtp_adapter_rank: int = 128 mtp_loss_weight: float = 0.0 bos_token_id: int = 0 eos_token_id: int = 1 pad_token_id: int = 2 unk_token_id: int = 3 # Masked-diffusion (Fugue) models replace noised tokens with this id. mask_token_id: int | None = None # Fugue only: 0 denoises whole sequences bidirectionally (Rondo); G > 0 # generates G-token blocks left to right over a causal clean stream # (Fugue-1, docs/architecture.md "Fugue-1"). fugue_block_size: int = 0 def __post_init__(self) -> None: for name in ("shared_layer_indices", "ngram_orders"): object.__setattr__(self, name, tuple(getattr(self, name))) if self.vocab_size <= 0 or self.vocab_size > 65_536: raise ValueError("vocab_size must fit the uint16 token dataset") if self.hidden_size != self.num_attention_heads * self.attention_head_dim: raise ValueError("hidden_size must equal num_attention_heads * attention_head_dim") if self.hidden_size != self.linear_num_heads * self.linear_head_dim: raise ValueError("hidden_size must equal linear_num_heads * linear_head_dim") if self.linear_num_value_heads < self.linear_num_heads: raise ValueError("linear_num_value_heads must be at least linear_num_heads") if self.linear_num_value_heads % self.linear_num_heads != 0: raise ValueError("linear_num_value_heads must be divisible by linear_num_heads") if self.num_attention_heads % self.num_key_value_heads != 0: raise ValueError("num_attention_heads must be divisible by num_key_value_heads") if self.num_hidden_layers % self.full_attention_interval != 0: raise ValueError("num_hidden_layers must be divisible by full_attention_interval") if not 0.0 < self.partial_rotary_factor <= 1.0: raise ValueError("partial_rotary_factor must be in (0, 1]") rotary_dim = int(self.attention_head_dim * self.partial_rotary_factor) if rotary_dim <= 0 or rotary_dim % 2: raise ValueError("The partial rotary dimension must be positive and even") if self.training_sequence_length > self.max_position_embeddings: raise ValueError("training_sequence_length exceeds max_position_embeddings") if len(set(self.shared_layer_indices)) != len(self.shared_layer_indices): raise ValueError("shared_layer_indices must be unique") if any( index < 0 or index >= self.num_hidden_layers for index in self.shared_layer_indices ): raise ValueError("shared_layer_indices contains an invalid layer") if self.mtp_num_heads < 0: raise ValueError("mtp_num_heads cannot be negative") if self.mtp_num_heads and self.mtp_adapter_rank <= 0: raise ValueError("mtp_adapter_rank must be positive when MTP is enabled") if not 0.0 <= self.mtp_loss_weight <= 1.0: raise ValueError("mtp_loss_weight must be between zero and one") for token_id in ( self.bos_token_id, self.eos_token_id, self.pad_token_id, self.unk_token_id, ): if not 0 <= token_id < self.vocab_size: raise ValueError(f"Special token ID {token_id} is outside the vocabulary") if self.linear_layer_type not in LINEAR_LAYER_TYPES: raise ValueError( f"linear_layer_type must be one of {sorted(LINEAR_LAYER_TYPES)}" ) if self.linear_layer_type == "diag_kdn": if self.linear_head_dim not in DIAG_KDN_HEAD_DIMS: raise ValueError( f"diag_kdn needs linear_head_dim in {DIAG_KDN_HEAD_DIMS}" ) if self.linear_num_value_heads != self.linear_num_heads: raise ValueError("diag_kdn needs linear_num_value_heads == linear_num_heads") if self.linear_expand_v != 1.0: raise ValueError("diag_kdn needs linear_expand_v == 1.0") if self.allow_negative_eigenvalues: raise ValueError("diag_kdn derives its gain from the covariance; negative eigenvalues are unsupported") if self.attention_layer_type not in ATTENTION_LAYER_TYPES: raise ValueError( f"attention_layer_type must be one of {sorted(ATTENTION_LAYER_TYPES)}" ) if self.residual_mode not in RESIDUAL_MODES: raise ValueError(f"residual_mode must be one of {sorted(RESIDUAL_MODES)}") gated = self.residual_mode == "gated_residual" if gated: if self.residual_branches < 2: raise ValueError("gated_residual needs at least two residual_branches") if self.residual_gate_rank <= 0: raise ValueError("residual_gate_rank must be positive") if self.use_block_attn_res: raise ValueError("gated_residual replaces block attention residuals") if self.mtp_num_heads: raise ValueError("gated_residual models do not support MTP heads") if self.shared_layer_indices: raise ValueError("gated_residual models do not support shared layers") if (self.linear_layer_type, self.attention_layer_type) not in GATED_RESIDUAL_MIXERS: raise ValueError( "gated_residual models pair gdn2 with parallax attention, or bigdn2 with bidirectional" ) elif ( self.attention_layer_type in {"parallax", "bidirectional"} or self.linear_layer_type == "bigdn2" or self.ngram_orders ): raise ValueError( "parallax and bidirectional mixers and n-gram memory need residual_mode 'gated_residual'" ) fugue = self.model_type == FUGUE_MODEL_TYPE if fugue != (self.linear_layer_type == "bigdn2"): raise ValueError("Fugue models, and only Fugue models, use bidirectional bigdn2 mixers") if fugue: if self.mask_token_id is None or not 0 <= self.mask_token_id < self.vocab_size: raise ValueError("Fugue models need a mask_token_id inside the vocabulary") if self.ngram_orders: raise ValueError("Fugue models have no n-gram memory") if self.fugue_block_size < 0: raise ValueError("fugue_block_size cannot be negative") if self.fugue_block_size and self.training_sequence_length % self.fugue_block_size: raise ValueError("training_sequence_length must be a multiple of fugue_block_size") elif self.fugue_block_size: raise ValueError("fugue_block_size applies to Fugue models only") if self.linear_full_rank_gate and self.linear_layer_type not in {"gdn2", "bigdn2"}: raise ValueError("linear_full_rank_gate applies to gdn2 and bigdn2 layers") if self.ngram_orders: if any(order < 2 for order in self.ngram_orders): raise ValueError("ngram_orders must be at least 2") if tuple(sorted(set(self.ngram_orders))) != tuple(self.ngram_orders): raise ValueError("ngram_orders must be unique and ascending") if self.ngram_heads_per_order <= 0 or self.ngram_embedding_dim <= 0: raise ValueError("n-gram heads and embedding width must be positive") if self.ngram_table_rows <= self.vocab_size: raise ValueError("ngram_table_rows must exceed the vocabulary") if not 0 <= self.ngram_layer_index < self.num_hidden_layers: raise ValueError("ngram_layer_index must name a block") if self.ngram_conv_kernel <= 0: raise ValueError("ngram_conv_kernel must be positive") if self.mlp_activation not in {"swiglu", "situ_glu"}: raise ValueError("mlp_activation must be 'swiglu' or 'situ_glu'") if self.mla_kv_lora_rank <= 0: raise ValueError("mla_kv_lora_rank must be positive") if self.mla_q_lora_rank is not None and self.mla_q_lora_rank <= 0: raise ValueError("mla_q_lora_rank must be positive or None") if self.mla_qk_nope_head_dim <= 0 or self.mla_v_head_dim <= 0: raise ValueError("MLA head dimensions must be positive") if self.kda_lower_bound >= 0.0: raise ValueError("kda_lower_bound must be negative log-decay") if self.situ_gate_cap <= 0.0 or self.situ_up_cap <= 0.0: raise ValueError("SiTU caps must be positive") @property def layer_types(self) -> tuple[str, ...]: return tuple( self.attention_layer_type if (index + 1) % self.full_attention_interval == 0 else self.linear_layer_type for index in range(self.num_hidden_layers) ) @property def rotary_dim(self) -> int: return int(self.attention_head_dim * self.partial_rotary_factor) @property def effective_num_layers(self) -> int: return self.num_hidden_layers + len(self.shared_layer_indices) @property def ngram_table_count(self) -> int: return len(self.ngram_orders) * self.ngram_heads_per_order @property def ngram_memory_dim(self) -> int: return self.ngram_table_count * self.ngram_embedding_dim def to_dict(self) -> dict[str, Any]: payload = asdict(self) payload["layer_types"] = list(self.layer_types) return payload def save_json(self, path: Path) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text( json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8", ) @classmethod def from_json(cls, path: Path) -> TinyGDNConfig: payload = json.loads(path.read_text(encoding="utf-8")) payload.pop("layer_types", None) for name in ("shared_layer_indices", "ngram_orders"): if name in payload: payload[name] = tuple(payload[name]) allowed = {item.name for item in fields(cls)} return cls(**{key: value for key, value in payload.items() if key in allowed}) def haiku_config(**overrides: Any) -> TinyGDNConfig: """650M Haiku: 3×KDA + 1×Gated-MLA, AttnRes, SiTU-GLU, MTP.""" payload = { "architecture": "HaikuForCausalLM", "model_type": "haiku", "vocab_size": 65_536, "hidden_size": 1024, "intermediate_size": 3840, "num_hidden_layers": 36, "num_attention_heads": 8, "num_key_value_heads": 8, "attention_head_dim": 128, "full_attention_interval": 4, "linear_num_heads": 8, "linear_num_value_heads": 8, "linear_head_dim": 128, "linear_layer_type": "kda", "attention_layer_type": "gated_mla", "mlp_activation": "situ_glu", "use_block_attn_res": True, "mtp_num_heads": 2, "mtp_adapter_rank": 128, "mtp_loss_weight": 0.2, "max_position_embeddings": 32_768, "training_sequence_length": 2048, } payload.update(overrides) return TinyGDNConfig(**payload) def haiku_kdn_config(**overrides: Any) -> TinyGDNConfig: """Haiku with Diagonal KDN in place of KDA in the 27 linear layers. Opt-in: at equal tokens it trailed KDA in Haiku pilots and trains ~23% slower at full depth (docs/architecture.md, "Haiku v2"). """ return haiku_config(**{"linear_layer_type": "diag_kdn", **overrides}) def senryu_config(**overrides: Any) -> TinyGDNConfig: """800M Senryu: 3×Gated DeltaNet-2 + 1×Gated-MLA over 44 layers, AttnRes, SiTU-GLU. Haiku's width, heads, tokenizer, and residual design with GDN-2's channel-wise erase and write gates in the 33 linear layers. A rank-512 query LoRA costs as many parameters as a full 1024×1024 query, so queries are full rank, and the KV latent doubles to 512. The feed-forward narrows to 3392 to pay for the depth. Next-token only: MTP has no pilot win at this scale (docs/architecture.md, "Senryu"). """ return haiku_config( **{ "architecture": "SenryuForCausalLM", "model_type": "senryu", "intermediate_size": 3392, "num_hidden_layers": 44, "linear_layer_type": "gdn2", "mla_q_lora_rank": None, "mla_kv_lora_rank": 512, "mtp_num_heads": 0, "mtp_loss_weight": 0.0, **overrides, } ) def senryu2_1b_config(**overrides: Any) -> TinyGDNConfig: """Senryu-2-1B: a 1.09B Senryu-2 backbone plus a 2.2B host-side n-gram memory. 28 layers of 3×Gated DeltaNet-2 (full-rank output gate) + 1×Parallax attention, width 1536, SiTU-GLU 3584, a 4-branch gated residual stream, and orders-2/3 hashed n-gram tables injected before block 1 (docs/architecture.md, "Senryu-2-1B"). """ payload: dict[str, Any] = { "architecture": "Senryu2ForCausalLM", "model_type": "senryu2", "vocab_size": 65_536, "hidden_size": 1536, "intermediate_size": 3584, "num_hidden_layers": 28, "num_attention_heads": 12, "num_key_value_heads": 4, "attention_head_dim": 128, "full_attention_interval": 4, "partial_rotary_factor": 0.5, "rope_theta": 1_000_000.0, "linear_num_heads": 12, "linear_num_value_heads": 12, "linear_head_dim": 128, "linear_layer_type": "gdn2", "attention_layer_type": "parallax", "mlp_activation": "situ_glu", "residual_mode": "gated_residual", "residual_branches": 4, "residual_gate_rank": 128, "linear_full_rank_gate": True, "ngram_orders": (2, 3), "ngram_heads_per_order": 4, "ngram_table_rows": 2_097_152, "ngram_embedding_dim": 128, "ngram_layer_index": 1, "ngram_conv_kernel": 4, "max_position_embeddings": 32_768, "training_sequence_length": 32_768, "mtp_num_heads": 0, "mtp_loss_weight": 0.0, } payload.update(overrides) return TinyGDNConfig(**payload) def monody_config(**overrides: Any) -> TinyGDNConfig: """Monody: the autoregressive control in the Rondo/Monody data-efficiency experiment. Senryu-2's backbone at width 1024 with no n-gram memory: 16 layers of 3×Gated DeltaNet-2 + 1×Parallax attention, SiTU-GLU 2816, 8 query and 4 KV heads, trained on 2,048-token windows (docs/architecture.md, "Fugue"). """ return senryu2_1b_config( **{ "hidden_size": 1024, "intermediate_size": 2816, "num_hidden_layers": 16, "num_attention_heads": 8, "num_key_value_heads": 4, "linear_num_heads": 8, "linear_num_value_heads": 8, "ngram_orders": (), "max_position_embeddings": 8192, "training_sequence_length": 2048, **overrides, } ) def rondo_config(**overrides: Any) -> TinyGDNConfig: """Rondo: the first Fugue model, a masked-diffusion LM on Monody's backbone. Monody's widths, depth and heads with every token mixer bidirectional: 12 bidirectional Gated DeltaNet-2 layers and 4 bidirectional softmax attention layers, denoising `<|mask|>` (tokenizer v2, id 22). The backward scans' own convolutions and decay projections leave it within 0.3% of Monody's parameter count (docs/architecture.md, "Fugue"). """ return monody_config( **{ "architecture": "FugueForMaskedDiffusion", "model_type": FUGUE_MODEL_TYPE, "linear_layer_type": "bigdn2", "attention_layer_type": "bidirectional", "mask_token_id": 22, **overrides, } ) def fugue1_1b_config(**overrides: Any) -> TinyGDNConfig: """Fugue-1-1B: a 1.1B block-diffusion LM on Senryu-2-1B's backbone. Senryu-2-1B's widths, depth and heads without n-gram memory, as a Fugue model that writes 32-token blocks left to right: a causal clean stream (next-token loss, and the exact cache every block reads) and a noisy stream denoising each block bidirectionally against that cache, trained on 4,096-token windows (docs/architecture.md, "Fugue-1"). """ return senryu2_1b_config( **{ "architecture": "FugueForMaskedDiffusion", "model_type": FUGUE_MODEL_TYPE, "linear_layer_type": "bigdn2", "attention_layer_type": "bidirectional", "mask_token_id": 22, "fugue_block_size": 32, "ngram_orders": (), "max_position_embeddings": 32_768, "training_sequence_length": 4096, **overrides, } ) def prompttercet_config(**overrides: Any) -> TinyGDNConfig: """PromptTercet: a 501M Senryu-2 prompter, which writes the user's side of a chat. Senryu-2-1B's layout at width 1024: 24 layers of 3×Gated DeltaNet-2 + 1×Parallax attention, SiTU-GLU 2816, 8 query and 4 KV heads, and n-gram tables of 1,048,576 rows, 1.09B host-side parameters, half of Senryu-2-1B's (docs/architecture.md, "PromptTercet"). """ return senryu2_1b_config( **{ "hidden_size": 1024, "intermediate_size": 2816, "num_hidden_layers": 24, "num_attention_heads": 8, "num_key_value_heads": 4, "linear_num_heads": 8, "linear_num_value_heads": 8, "ngram_table_rows": 1_048_576, **overrides, } )