"""Model architectures shared by the baseline and LongCat training scripts.""" import math import warnings from typing import Optional import torch import torch.nn as nn from transformers import LlamaConfig, LlamaForCausalLM from transformers.models.llama import modeling_llama as llama_modeling try: # Kimi Delta Attention layer (linear attention) from flash-linear-attention. # Optional: only the KDA architecture needs it, so the baseline/longcat scripts # keep importing model.py even when FLA is absent. from fla.layers.kda import KimiDeltaAttention FLA_KDA_IMPORT_ERROR = None except ImportError as exc: # pragma: no cover - exercised only without FLA KimiDeltaAttention = None FLA_KDA_IMPORT_ERROR = exc # Distinct prime hash multipliers ("salts"), one per (order, head) table. Giving # every head a different polynomial base makes the K hash functions genuinely # independent *regardless of table size*, so the historical K->1 collapse — two # heads that shared a table size hashed every n-gram to the identical slot — # cannot recur. Primes larger than the base vocab keep each per-head polynomial # injective over the token range, and being coprime to the table sizes avoids the # base-multiple collision spike the LongCat paper reports (Fig. 3b): the effective # base (multiplier mod table_size) is then a scrambled value rather than the raw # vocab size. All of this is enforced at build time by _validate_ngram_hashing. _DEFAULT_HASH_MULTIPLIERS = ( 40009, 100003, 262147, 524287, 1000003, 2000003, 3000017, 4000037, 5000011, 6000101, 7000127, 8000009, ) # Default minimum pairwise separation between table sizes (fraction of the # smaller size). Near-equal sizes are the condition that silently disabled # multi-head hashing before, so the build refuses to start below this. _DEFAULT_MIN_PAIRWISE_SIZE_GAP = 0.005 def _validate_ngram_hashing( table_sizes: list[int], multipliers: list[int], base_vocab: int, min_pairwise_size_gap: float, ) -> None: """Refuse to build a degenerate n-gram hashing setup. Raises ValueError. Guards, in order of how badly they corrupt the experiment: 1. No two tables may be the *same hash function*. Two heads are identical iff they share both a table size and an effective base (multiplier mod size); that is the exact K->1 collapse. This is the load-bearing invariant. 2. Each multiplier must be >= base vocab, or distinct n-grams alias before the modulus (the base-`m` polynomial stops being injective over token digits). 3. Each multiplier must be coprime to its table size, or one n-gram coordinate collapses into gcd-many classes (the mechanism behind the paper's spike). 4. Table sizes must not be near-duplicates — the config that hid the clone bug. A soft warning also fires when a size sits within 5% of base vocab of an integer multiple of it (paper Fig. 3b), which prime multipliers mitigate but do not fully erase. """ n = len(table_sizes) if len(multipliers) != n: raise ValueError( f"Expected {n} ngram hash multipliers (one per table), got {len(multipliers)}." ) # (1) precise clone check: identical (size, effective base) => identical indices. seen: dict[tuple[int, int], int] = {} for idx, (m, s) in enumerate(zip(multipliers, table_sizes)): key = (s, m % s) if key in seen: raise ValueError( f"N-gram hash tables {seen[key]} and {idx} are the SAME hash function " f"(table size {s}, effective base {m % s}). Two heads that hash " "identically collapse K sub-tables to K=1 — exactly the bug this guard " "exists to prevent. Give them distinct multipliers or distinct sizes." ) seen[key] = idx for m, s in zip(multipliers, table_sizes): # (2) injectivity of the base-`m` polynomial over token digits [0, base_vocab). if m < base_vocab: raise ValueError( f"N-gram hash multiplier {m} must be >= base vocab {base_vocab}; " "a smaller base aliases distinct n-grams before the modulus is applied." ) # (3) coprimality: gcd > 1 collapses a coordinate into gcd-many residues. g = math.gcd(m, s) if g != 1: raise ValueError( f"N-gram hash multiplier {m} shares factor {g} with table size {s}. " "Pick a multiplier coprime to the table size (a prime larger than every " "table size is always safe) so no n-gram coordinate collapses." ) # (4) near-duplicate sizes: the historical trigger for the clone collapse. order = sorted(range(n), key=lambda i: table_sizes[i]) for a, b in zip(order, order[1:]): sa, sb = table_sizes[a], table_sizes[b] rel = abs(sa - sb) / min(sa, sb) if rel < min_pairwise_size_gap: raise ValueError( f"N-gram table sizes {sa} and {sb} differ by only {rel * 100:.3f}% " f"(guard requires >= {min_pairwise_size_gap * 100:.3f}%). Near-equal " "sizes are the condition that silently disabled multi-head hashing " "before; spread the table sizes apart." ) # (5) soft: sizes near an integer multiple of base vocab (paper Fig. 3b). for s in table_sizes: dist = min(s % base_vocab, base_vocab - s % base_vocab) if dist / base_vocab < 0.05: warnings.warn( f"N-gram table size {s} is within {dist} of an integer multiple of base " f"vocab {base_vocab}; the LongCat paper reports collision spikes there. " "The prime multipliers mitigate this, but consider nudging the size.", stacklevel=2, ) class LlamaLongCatNgramConfig(LlamaConfig): """Serializable configuration for :class:`LlamaLongCatNgram`.""" model_type = "llama_longcat_ngram" def __init__( self, ngram_max_n: int = 4, ngram_num_heads: int = 2, ngram_table_vocab_sizes: Optional[list[int]] = None, ngram_embedding_amplification: str = "layer_norm", ngram_hash_multipliers: Optional[list[int]] = None, ngram_min_pairwise_size_gap: float = _DEFAULT_MIN_PAIRWISE_SIZE_GAP, qk_norm: bool = False, qk_norm_eps: Optional[float] = None, **kwargs, ): super().__init__(**kwargs) self.ngram_max_n = ngram_max_n self.ngram_num_heads = ngram_num_heads self.ngram_table_vocab_sizes = ngram_table_vocab_sizes self.ngram_embedding_amplification = ngram_embedding_amplification self.ngram_hash_multipliers = ngram_hash_multipliers self.ngram_min_pairwise_size_gap = ngram_min_pairwise_size_gap self.qk_norm = qk_norm self.qk_norm_eps = qk_norm_eps class LongCatNgramEmbedder(nn.Module): """LongCat N-gram Embedding from Eq. 2 and Eq. 3 of arXiv:2601.21204.""" def __init__(self, config: LlamaLongCatNgramConfig): super().__init__() self.max_n = config.ngram_max_n self.num_heads = config.ngram_num_heads self.base_vocab_size = config.vocab_size self.eos_token_id = config.eos_token_id self.orders = list(range(2, self.max_n + 1)) num_tables = len(self.orders) * self.num_heads table_vocab_sizes = config.ngram_table_vocab_sizes if config.hidden_size % num_tables != 0: raise ValueError( f"hidden_size ({config.hidden_size}) must be divisible by " f"(ngram_max_n-1)*ngram_num_heads ({num_tables})." ) if table_vocab_sizes is None or len(table_vocab_sizes) != num_tables: actual = None if table_vocab_sizes is None else len(table_vocab_sizes) raise ValueError( f"Expected {num_tables} ngram_table_vocab_sizes " f"((ngram_max_n-1)*ngram_num_heads), got {actual}." ) self.sub_dim = config.hidden_size // num_tables # Resolve the per-head hash multipliers (built-in defaults unless the # config overrides them) and refuse to build a degenerate setup. configured = config.ngram_hash_multipliers if configured is None: if num_tables > len(_DEFAULT_HASH_MULTIPLIERS): raise ValueError( f"Need {num_tables} hash multipliers but only " f"{len(_DEFAULT_HASH_MULTIPLIERS)} defaults are defined; pass " "ngram_hash_multipliers explicitly." ) configured = _DEFAULT_HASH_MULTIPLIERS[:num_tables] multipliers: list[int] = [int(m) for m in configured] _validate_ngram_hashing( list(table_vocab_sizes), multipliers, self.base_vocab_size, config.ngram_min_pairwise_size_gap, ) # Persist the resolved list so it is serialized in config.json. config.ngram_hash_multipliers = multipliers self.tables = nn.ModuleDict() self.projections = nn.ModuleDict() self.multipliers: dict[str, int] = {} idx = 0 for n in self.orders: for k in range(self.num_heads): key = f"n{n}_k{k}" self.tables[key] = nn.Embedding( table_vocab_sizes[idx], self.sub_dim ) self.projections[key] = nn.Linear( self.sub_dim, config.hidden_size, bias=False ) self.multipliers[key] = multipliers[idx] idx += 1 amplification = config.ngram_embedding_amplification.strip().lower() if amplification == "layer_norm": self.amplification = nn.LayerNorm(config.hidden_size) self.amplification_scale = 1.0 elif amplification == "sqrt_d": self.amplification = nn.Identity() self.amplification_scale = math.sqrt(config.hidden_size) elif amplification == "none": self.amplification = nn.Identity() self.amplification_scale = 1.0 else: raise ValueError( "ngram_embedding_amplification must be one of " "{'layer_norm', 'sqrt_d', 'none'}, got " f"{amplification!r}." ) def _shift_right(self, x: torch.Tensor, shift: int) -> torch.Tensor: """Causal shift, zeroing context that crosses an EOS boundary.""" if shift == 0: return x pad = x.new_zeros(x.shape[0], shift) shifted = torch.cat([pad, x[:, :-shift]], dim=1) crosses_eos = torch.zeros_like(x, dtype=torch.bool) for offset in range(1, shift + 1): previous = torch.cat( [x.new_zeros(x.shape[0], offset), x[:, :-offset]], dim=1 ) crosses_eos |= previous.eq(self.eos_token_id) return shifted.masked_fill(crosses_eos, 0) def _hash_ngram( self, input_ids: torch.Tensor, n: int, table_size: int, multiplier: int, shifted_tokens: Optional[dict[int, torch.Tensor]] = None, ) -> torch.Tensor: """Eq. 2 with a per-head base: sum_j t[i-j] * multiplier**j mod table_size. `multiplier` is this head's hash salt (a distinct prime >= base vocab), so two heads never compute the same indices even at equal table sizes. The modulus is applied every Horner step, so the result matches the full polynomial mod `table_size` while staying far inside int64. """ h = torch.zeros_like(input_ids) for j in range(n - 1, -1, -1): tok = ( input_ids if j == 0 else shifted_tokens[j] if shifted_tokens is not None else self._shift_right(input_ids, j) ) h = (h * multiplier + tok) % table_size return h def forward( self, input_ids: torch.Tensor, base_embeddings: torch.Tensor ) -> torch.Tensor: """Return amplified Eq. 3 embeddings with shape [B, T, H].""" combined = base_embeddings shifted_tokens = { shift: self._shift_right(input_ids, shift) for shift in range(1, self.max_n) } for n in self.orders: for k in range(self.num_heads): key = f"n{n}_k{k}" table_size = self.tables[key].num_embeddings hash_ids = self._hash_ngram( input_ids, n, table_size, self.multipliers[key], shifted_tokens ) combined = combined + self.projections[key]( self.tables[key](hash_ids) ) combined = combined / (len(self.orders) * self.num_heads + 1) return self.amplification(combined) * self.amplification_scale class LlamaQKNormAttention(llama_modeling.LlamaAttention): """LLaMA attention with per-head RMSNorm on Q and K before RoPE.""" def __init__(self, config: LlamaLongCatNgramConfig, layer_idx: int): super().__init__(config, layer_idx) eps = config.qk_norm_eps if config.qk_norm_eps is not None else config.rms_norm_eps self.q_norm = llama_modeling.LlamaRMSNorm(self.head_dim, eps=eps) self.k_norm = llama_modeling.LlamaRMSNorm(self.head_dim, eps=eps) def forward( self, hidden_states: torch.Tensor, position_embeddings=None, attention_mask=None, past_key_values=None, **kwargs, ): input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = llama_modeling.apply_rotary_pos_emb( query_states, key_states, cos, sin ) if past_key_values is not None: key_states, value_states = past_key_values.update( key_states, value_states, self.layer_idx ) attention_interface = llama_modeling.ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, llama_modeling.eager_attention_forward ) attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs, ) attn_output = self.o_proj(attn_output.reshape(*input_shape, -1).contiguous()) return attn_output, attn_weights class LlamaLongCatNgram(LlamaForCausalLM): """LLaMA using LongCat's standard input N-gram Embedding (NE).""" config_class = LlamaLongCatNgramConfig def __init__(self, config: LlamaLongCatNgramConfig): super().__init__(config) if config.qk_norm: for layer_idx, layer in enumerate(self.model.layers): layer.self_attn = LlamaQKNormAttention(config, layer_idx) self.ngram_embedder = LongCatNgramEmbedder(config) def forward(self, input_ids=None, inputs_embeds=None, **kwargs): if input_ids is not None and inputs_embeds is not None: raise ValueError("Specify exactly one of input_ids or inputs_embeds.") if input_ids is not None: base_embeddings = self.model.embed_tokens(input_ids) inputs_embeds = self.ngram_embedder(input_ids, base_embeddings) input_ids = None return super().forward( input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs ) def prepare_inputs_for_generation(self, input_ids, **kwargs): """Preserve causal n-gram context when HF slices cached decode inputs.""" model_inputs = super().prepare_inputs_for_generation(input_ids, **kwargs) prepared_ids = model_inputs.get("input_ids") if prepared_ids is None: return model_inputs base_embeddings = self.model.embed_tokens(input_ids) full_embeddings = self.ngram_embedder(input_ids, base_embeddings) model_inputs["inputs_embeds"] = full_embeddings[:, -prepared_ids.shape[1] :] model_inputs["input_ids"] = None return model_inputs # Instruct save_pretrained() to package this source file and write AutoClass # metadata. Loading the resulting checkpoint requires trust_remote_code=True. LlamaLongCatNgramConfig.register_for_auto_class() LlamaLongCatNgram.register_for_auto_class("AutoModelForCausalLM") class LlamaKDAConfig(LlamaConfig): """Config for :class:`LlamaKDA` — a hybrid Kimi-Delta / softmax LLaMA. A subset of decoder layers use Kimi Delta Attention (KDA, a gated-delta linear attention); the rest keep standard softmax self-attention (with the same FA backend and optional QK-norm as the baseline). Which layers are which is resolved by :func:`resolve_kda_layer_types`, honoring (in priority order) ``kda_full_attn_layers`` > ``kda_full_attn_range`` > ``kda_full_attn_every``. With none set, every layer is KDA (pure linear attention). """ model_type = "llama_kda" def __init__( self, # --- Hybrid layout: which layers keep softmax (full) attention --- kda_full_attn_layers: Optional[list[int]] = None, kda_full_attn_every: Optional[int] = None, kda_full_attn_range: Optional[list[int]] = None, # Sparse-KDA interleave (the inverse of kda_full_attn_every): one KDA layer # every `kda_every` layers, at indices where i % kda_every == kda_offset; # every other layer is full (GQA) attention. kda_every: Optional[int] = None, kda_offset: int = 0, # --- KDA layer hyperparameters (forwarded to fla KimiDeltaAttention) --- kda_head_dim: int = 128, kda_num_heads: Optional[int] = None, kda_num_v_heads: Optional[int] = None, kda_expand_v: float = 1.0, kda_use_short_conv: bool = True, kda_conv_size: int = 4, kda_conv_bias: bool = False, kda_allow_neg_eigval: bool = False, kda_lower_bound: Optional[float] = None, kda_safe_gate: bool = False, # --- QK-norm applies to the softmax (full-attention) layers only --- qk_norm: bool = False, qk_norm_eps: Optional[float] = None, **kwargs, ): super().__init__(**kwargs) self.kda_full_attn_layers = kda_full_attn_layers self.kda_full_attn_every = kda_full_attn_every self.kda_full_attn_range = kda_full_attn_range self.kda_every = kda_every self.kda_offset = kda_offset self.kda_head_dim = kda_head_dim self.kda_num_heads = kda_num_heads self.kda_num_v_heads = kda_num_v_heads self.kda_expand_v = kda_expand_v self.kda_use_short_conv = kda_use_short_conv self.kda_conv_size = kda_conv_size self.kda_conv_bias = kda_conv_bias self.kda_allow_neg_eigval = kda_allow_neg_eigval self.kda_lower_bound = kda_lower_bound self.kda_safe_gate = kda_safe_gate self.qk_norm = qk_norm self.qk_norm_eps = qk_norm_eps def resolve_kda_layer_types(config: LlamaKDAConfig) -> list[str]: """Return a per-layer list of ``"kda"`` / ``"full"`` (softmax) attention. Priority: explicit ``kda_full_attn_layers`` > contiguous ``kda_full_attn_range`` ``[start, end)`` > interleaved ``kda_full_attn_every`` (the last layer of every block of ``n`` is full attention, e.g. ``4`` -> Kimi/Qwen-style 3:1) > ``kda_every`` (the INVERSE — KDA is the sparse type: one KDA layer every ``kda_every`` layers at ``i % kda_every == kda_offset``, all others full). If none is set, all layers are KDA. """ n = config.num_hidden_layers if config.kda_full_attn_layers is not None: full = set(int(i) for i in config.kda_full_attn_layers) elif config.kda_full_attn_range is not None: start, end = config.kda_full_attn_range full = set(range(int(start), int(end))) elif config.kda_full_attn_every: every = int(config.kda_full_attn_every) if every < 1: raise ValueError(f"kda_full_attn_every must be >= 1, got {every}.") full = {i for i in range(n) if (i + 1) % every == 0} elif getattr(config, "kda_every", None): # Inverse of kda_full_attn_every: KDA is the SPARSE type. One KDA layer every # `kda_every` layers at i % kda_every == kda_offset; every other layer is full. every = int(config.kda_every) if every < 1: raise ValueError(f"kda_every must be >= 1, got {every}.") offset = int(getattr(config, "kda_offset", 0) or 0) % every kda = {i for i in range(n) if i % every == offset} full = set(range(n)) - kda else: full = set() for i in full: if not 0 <= i < n: raise ValueError( f"Full-attention layer index {i} is out of range for " f"num_hidden_layers={n}." ) return ["full" if i in full else "kda" for i in range(n)] class LlamaKDAAttention(nn.Module): """Adapter wrapping fla's :class:`KimiDeltaAttention` for a LLaMA decoder layer. KDA is linear attention: it carries no RoPE and normalizes q/k internally (L2-norm), so ``position_embeddings`` are ignored here. The decoder layer expects a ``(hidden_states, attn_weights)`` pair back; KDA returns a triple, so we drop the cache/weights. A 4-D causal mask (built by ``LlamaModel`` for the softmax layers) is meaningless to KDA — only a 2-D ``[B, T]`` padding mask is forwarded; anything else becomes ``None`` (packed training carries no padding). The module is run **stateless**: it never reads or writes ``past_key_values``. HF's ``LlamaModel`` hands every layer an HF ``DynamicCache`` (incompatible with fla's recurrent-state cache), which is fine for full-sequence LM-loss training and eval but means this wrapper does not support HF incremental ``generate``. """ def __init__(self, config: LlamaKDAConfig, layer_idx: int): super().__init__() if KimiDeltaAttention is None: raise ImportError( "LlamaKDA requires flash-linear-attention (fla) for KimiDeltaAttention." ) from FLA_KDA_IMPORT_ERROR head_dim = config.kda_head_dim num_heads = config.kda_num_heads or (config.hidden_size // head_dim) if num_heads * head_dim != config.hidden_size: # fla supports q/k dim != hidden; Kimi-Linear over-provisions ~1.8x. warnings.warn( f"KDA q/k dim {num_heads * head_dim} != hidden_size " f"{config.hidden_size}; layer params/state will differ from a " "same-width softmax layer.", stacklevel=2, ) self.layer_idx = layer_idx self.kda = KimiDeltaAttention( hidden_size=config.hidden_size, expand_v=config.kda_expand_v, head_dim=head_dim, num_heads=num_heads, num_v_heads=config.kda_num_v_heads, mode="chunk", use_short_conv=config.kda_use_short_conv, conv_size=config.kda_conv_size, conv_bias=config.kda_conv_bias, allow_neg_eigval=config.kda_allow_neg_eigval, safe_gate=config.kda_safe_gate, lower_bound=config.kda_lower_bound, layer_idx=layer_idx, norm_eps=config.rms_norm_eps, ) def forward( self, hidden_states: torch.Tensor, position_embeddings=None, attention_mask=None, past_key_values=None, use_cache=False, **kwargs, ): mask = attention_mask if (attention_mask is not None and attention_mask.dim() == 2) else None forward_kwargs = {k: v for k, v in kwargs.items() if k == "cu_seqlens"} attn_output, _, _ = self.kda( hidden_states=hidden_states, attention_mask=mask, past_key_values=None, use_cache=False, **forward_kwargs, ) return attn_output, None class LlamaKDA(LlamaForCausalLM): """LLaMA whose attention is a config-driven hybrid of KDA and softmax layers.""" config_class = LlamaKDAConfig def __init__(self, config: LlamaKDAConfig): super().__init__(config) layer_types = resolve_kda_layer_types(config) for layer_idx, layer in enumerate(self.model.layers): if layer_types[layer_idx] == "kda": layer.self_attn = LlamaKDAAttention(config, layer_idx) elif config.qk_norm: layer.self_attn = LlamaQKNormAttention(config, layer_idx) # otherwise keep the default softmax LlamaAttention from super().__init__. # Persist the resolved layout so it lands in config.json and can be logged. config.kda_layer_types = layer_types LlamaKDAConfig.register_for_auto_class() LlamaKDA.register_for_auto_class("AutoModelForCausalLM") __all__ = [ "LlamaConfig", "LlamaForCausalLM", "LlamaLongCatNgramConfig", "LongCatNgramEmbedder", "LlamaQKNormAttention", "LlamaLongCatNgram", "LlamaKDAConfig", "LlamaKDAAttention", "LlamaKDA", "resolve_kda_layer_types", "_validate_ngram_hashing", ]