"""LuauCoder configuration (``trust_remote_code=True``). Mirrors ``lm.config.ModelConfig`` with transformers' standard names. ``num_attention_heads`` counts attention output heads: with ``diff_v2`` the model runs twice as many query heads against ``num_key_value_heads`` K/V heads. """ from __future__ import annotations import math from transformers import PretrainedConfig __all__ = ["LuauCoderConfig"] class LuauCoderConfig(PretrainedConfig): model_type = "luau_coder" keys_to_ignore_at_inference = ["past_key_values"] def __init__( self, vocab_size: int = 32768, hidden_size: int = 2560, intermediate_size: int = 6144, num_hidden_layers: int = 36, num_attention_heads: int = 20, num_key_value_heads: int = 5, head_dim: int = 128, rms_norm_eps: float = 1e-6, final_logit_softcapping: float | None = None, max_position_embeddings: int = 4096, rope_theta: float = 500000.0, rope_scaling: dict | None = None, lerope_alpha: list[float] | None = None, nope_layers: list[int] | None = None, diff_v2: bool = True, xsa: bool = True, qk_norm: bool = True, use_output_gate: bool | str | None = False, attn_res: bool = False, attn_res_block_size: int = 6, attn_res_delta_block: bool = False, attn_res_partial_source: bool = False, ccgqa: bool = False, ccgqa_seq_kernel: int = 2, ccgqa_head_kernel: int = 2, ccgqa_qk_mean: bool = True, ccgqa_value_shift: bool = True, vha: bool = False, iha: bool = False, iha_pseudo_heads: int = 2, iha_window: int = 0, iha_global_every: int = 5, head_gate_rank: int = 0, initializer_range: float | None = None, use_cache: bool = True, tie_word_embeddings: bool = False, # throughput knobs (no effect on the math) fused_kernels: bool = True, compile_elementwise: bool = True, loss_chunk_size: int = 4096, bos_token_id: int | None = None, eos_token_id: int | list[int] | None = 32744, pad_token_id: int | None = 32745, **kwargs, ): self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.head_dim = head_dim self.rms_norm_eps = rms_norm_eps self.final_logit_softcapping = final_logit_softcapping self.max_position_embeddings = max_position_embeddings self.rope_theta = rope_theta self.rope_scaling = rope_scaling # LeRoPE: learned log-frequency scale per RoPE band, shared by every layer; None = plain RoPE self.lerope_alpha = None if lerope_alpha is None else [float(a) for a in lerope_alpha] self.nope_layers = sorted(nope_layers or []) self.diff_v2 = diff_v2 self.xsa = xsa self.qk_norm = qk_norm # False | "elementwise" | "headwise" self.use_output_gate = ("elementwise" if use_output_gate is True else str(use_output_gate).lower() if use_output_gate else False) self.attn_res = attn_res self.attn_res_block_size = attn_res_block_size self.attn_res_delta_block = attn_res_delta_block self.attn_res_partial_source = attn_res_partial_source self.ccgqa = bool(ccgqa) self.ccgqa_seq_kernel = ccgqa_seq_kernel self.ccgqa_head_kernel = ccgqa_head_kernel self.ccgqa_qk_mean = ccgqa_qk_mean self.ccgqa_value_shift = ccgqa_value_shift self.vha = bool(vha) self.iha = bool(iha) self.iha_pseudo_heads = iha_pseudo_heads self.iha_window = iha_window self.iha_global_every = iha_global_every # gated byte head: [h, sigmoid(u . h + b) * P h] against lm_head = [W, Z]; 0 = plain head self.head_gate_rank = int(head_gate_rank) if self.head_gate_rank < 0: raise ValueError("head_gate_rank must be >= 0") if self.iha and (iha_pseudo_heads < 1 or iha_window < 0 or iha_global_every < 0): raise ValueError("iha: iha_pseudo_heads >= 1, iha_window >= 0 and iha_global_every >= 0") self.initializer_range = hidden_size ** -0.5 if initializer_range is None else initializer_range self.use_cache = use_cache self.fused_kernels = fused_kernels self.compile_elementwise = compile_elementwise self.loss_chunk_size = loss_chunk_size if num_attention_heads % num_key_value_heads: raise ValueError(f"num_attention_heads ({num_attention_heads}) must be a multiple of " f"num_key_value_heads ({num_key_value_heads})") if head_dim % 2: raise ValueError("head_dim must be even (RoPE)") if self.use_output_gate not in (False, "elementwise", "headwise"): raise ValueError(f"use_output_gate: true / 'elementwise', 'headwise' or false, got {use_output_gate!r}") if self.ccgqa and not qk_norm: raise ValueError("ccgqa needs qk_norm") if self.ccgqa and ccgqa_value_shift and (num_key_value_heads * head_dim) % 2: raise ValueError("ccgqa_value_shift needs an even num_key_value_heads * head_dim") if attn_res and math.ceil(num_hidden_layers / attn_res_block_size) > 32: raise ValueError("attn_res: at most 32 blocks (num_hidden_layers / attn_res_block_size)") if any(not 0 <= i < num_hidden_layers for i in self.nope_layers): raise ValueError(f"nope_layers must be in [0, {num_hidden_layers}), got {self.nope_layers}") if rope_scaling is not None: rtype = rope_scaling.get("rope_type", rope_scaling.get("type")) if rtype not in (None, "default", "yarn"): raise ValueError(f"rope_scaling: only 'yarn' is supported, got {rtype!r}") super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs) @property def output_gate(self) -> str | None: """None (no gate), "elementwise" or "headwise".""" return self.use_output_gate or None @property def num_query_heads(self) -> int: """Query heads the attention kernel runs (2x the output heads with DiffAttn V2).""" return self.num_attention_heads * (2 if self.diff_v2 else 1) @property def num_q_latent_heads(self) -> int: """Query heads ``qkv`` produces (VHA: the physical heads every KV group shares).""" return self.num_query_heads // self.num_key_value_heads if self.vha else self.num_query_heads def iha_layer(self, i: int) -> bool: """Whether layer ``i`` runs IHA (the rest are plain global attention layers).""" every = self.iha_global_every return self.iha and not (every and (i + 1) % every == 0)