Add ComplexKDA 1.3B/100BT FineWeb-Edu model

README.md

@@ -0,0 +1,100 @@
+---
+library_name: transformers
+license: mit
+pipeline_tag: text-generation
+tags:
+- complex-kda
+- linear-attention
+- kimi-delta-attention
+---
+
+# openeurollm/kda-sigmoid-hybrid-1.3B-100B
+
+A hybrid ComplexKDA language model (1.36B parameters).
+
+ComplexKDA is Kimi Delta Attention with a **signed decay gate**: the per-channel
+decay `alpha` is allowed to take either sign, `alpha in [-1, 1]`, instead of
+being confined to `(0, 1]`. That is the one-dimensional real case of a complex
+eigenvalue, so a channel can oscillate rather than only forget. The magnitude is
+carried in log space exactly as KDA carries it; the `+-1` part is carried as a
+running product pushed onto the queries and keys, so the recurrence the kernels
+run is still the unsigned one.
+
+This checkpoint's decay gate is **unsigned sigmoid (the KDA baseline: alpha in (0, 1])**.
+
+## Architecture
+
+- 24 layers, hidden size 2048, MLP 5312 (SwiGLU)
+- 16 heads of dimension 128, short convolution of width 4
+- vocabulary 32000, trained at context 4096
+- embeddings untied
+- attention at layers [3, 7, 11, 15, 19, 23] (gated, NoPE), linear everywhere else
+
+## Tokenizer, and how to start a prompt
+
+The bundled tokenizer is configured the way the training corpus was encoded:
+**no BOS is prepended**, documents were terminated with the EOS token, and
+`model_max_length` is this model's trained context. The upstream tokenizer
+repository's own defaults differ on both points, so encode through the
+tokenizer shipped here rather than re-fetching it by name.
+
+**To condition on the start of a document, prefix the EOS token, not a BOS.**
+Training packed documents as `[text..., EOS]`, so the token preceding any
+document's first token is always the EOS; a BOS was never seen at any position.
+Measured on held-out documents at 1.3B, against a bare prompt:
+
+| prefix | mean NLL | first 8 tokens |
+|---|---|---|
+| none | 2.0104 | 3.2947 |
+| **EOS** | **1.9960** | **2.9606** |
+| BOS | 2.0150 | 3.3768 |
+
+So the EOS is worth ~0.33 nats on the opening tokens and the BOS costs ~0.08.
+
+**Only at a document start.** The prefix is a "a new document begins here"
+signal, and mid-document it is a false one — on continuations the same EOS
+prefix *costs* 0.34 nats over the first 8 tokens (2.19 → 2.22 overall). Prefix
+it when you mean a fresh document; leave a continuation bare. Nothing in the
+weights or the tokenizer can tell these apart, which is why this is the
+caller's decision.
+
+`bos_token` is remapped to `</s>` here, so a pipeline that asks for "the BOS"
+(lm-eval's `--add_bos_token`, a generic wrapper) gets the separator rather than
+the unused `<s>`. `add_bos_token` remains `False`, so the default is still a
+bare prompt — which is the convention every number quoted for these models was
+measured under.
+
+If you are coming from `fla-hub` checkpoints, note that they use the opposite
+convention -- trained as `[BOS, text...]` with the BOS as the document
+separator -- so the habit does not carry over.
+
+## Usage
+
+The bundled `modeling_complex_kda.py` is **standalone**: `torch` and
+`transformers` are all it needs.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tok = AutoTokenizer.from_pretrained("openeurollm/kda-sigmoid-hybrid-1.3B-100B")
+model = AutoModelForCausalLM.from_pretrained(
+    "openeurollm/kda-sigmoid-hybrid-1.3B-100B", trust_remote_code=True, dtype="bfloat16")
+```
+
+For the Triton kernels these models were trained with -- much faster, and the
+exact code path of the training runs -- install the fork:
+
+```bash
+pip install git+https://github.com/automl/ComplexKDA
+```
+
+It is picked up automatically when importable. `COMPLEX_KDA_BACKEND=torch`
+forces the portable path; `=kernel` makes a missing fork an error instead of a
+silent fallback.
+
+## Provenance
+
+Converted from the training checkpoint with `lm_scaling/hf_release/convert_to_hub.py`.
+The conversion is metadata only -- the weight file is the exporter's own, byte
+for byte -- and the bundled implementation is checked against the reference
+implementation the runs used.

config.json

@@ -0,0 +1,55 @@
+{
+  "architectures": [
+    "ComplexKDAForCausalLM"
+  ],
+  "model_type": "complex_kda",
+  "auto_map": {
+    "AutoConfig": "configuration_complex_kda.ComplexKDAConfig",
+    "AutoModel": "modeling_complex_kda.ComplexKDAModel",
+    "AutoModelForCausalLM": "modeling_complex_kda.ComplexKDAForCausalLM"
+  },
+  "attn_mode": "chunk",
+  "hidden_size": 2048,
+  "num_hidden_layers": 24,
+  "num_heads": 16,
+  "head_dim": 128,
+  "intermediate_size": 5312,
+  "hidden_ratio": null,
+  "hidden_act": "swish",
+  "vocab_size": 32000,
+  "max_position_embeddings": 4096,
+  "tie_word_embeddings": false,
+  "norm_eps": 1e-06,
+  "conv_size": 4,
+  "use_short_conv": true,
+  "fuse_norm": true,
+  "fuse_swiglu": true,
+  "fuse_cross_entropy": true,
+  "initializer_range": 0.02,
+  "allow_neg_eigval": false,
+  "lower_bound": -5.0,
+  "expand_v": 1.0,
+  "gate": "sigmoid",
+  "gate_init_style": "shipped",
+  "output_gate": "lowrank",
+  "attn": {
+    "layers": [
+      3,
+      7,
+      11,
+      15,
+      19,
+      23
+    ],
+    "num_heads": 16,
+    "num_kv_heads": 16,
+    "qkv_bias": false,
+    "qk_norm": false,
+    "output_gate": true,
+    "use_rope": false,
+    "rope_theta": 10000.0,
+    "window_size": null
+  },
+  "drop_silu": true,
+  "use_cache": true
+}

configuration_complex_kda.py

@@ -0,0 +1,246 @@
+# Copyright (c) 2026, the ComplexKDA authors.
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li (the parts derived
+# from flash-linear-attention, MIT licensed).
+#
+# SPDX-License-Identifier: MIT
+"""Config for the ComplexKDA models published on the Hub.
+
+STANDALONE ON PURPOSE. `fla` is not imported anywhere in this file, so the
+config loads with nothing but `transformers` installed. The hybrid-attention
+normalisation that fla keeps in `fla/models/hybrid.py` is vendored below for
+the same reason -- a config that needed the fork to parse would make every
+error message about a missing package rather than about the model.
+"""
+
+from __future__ import annotations
+
+import math
+
+from transformers.configuration_utils import PretrainedConfig
+
+__all__ = ["ComplexKDAConfig"]
+
+
+# ---------------------------------------------------------------------------
+# hybrid attention spec (vendored from fla/models/hybrid.py)
+#
+# A hybrid arm replaces the linear mixer with full attention at the layers
+# named in `attn["layers"]`. The dict is stored in config.json, so it is
+# validated on the way in: an out-of-range or duplicated layer index would
+# otherwise build a model whose `state_dict` silently disagrees with the
+# checkpoint at exactly those layers.
+# ---------------------------------------------------------------------------
+
+def _spec_context(spec_index: int | None) -> str:
+    return "attn specification" if spec_index is None else f"attn specification at index {spec_index}"
+
+
+def _positive_int(value: object, *, field: str, context: str) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{context} field {field!r} must be a positive integer; got {value!r}")
+    return value
+
+
+def _normalize_spec(spec: dict, *, num_hidden_layers: int, spec_index: int | None,
+                    assigned_layers: dict) -> dict:
+    context = _spec_context(spec_index)
+    normalized = dict(spec)
+
+    for field in ("layers", "num_heads"):
+        if field not in normalized:
+            raise ValueError(f"{context} field {field!r} is required; got <missing>")
+
+    layers = normalized["layers"]
+    if not isinstance(layers, (list, tuple)):
+        raise ValueError(
+            f"{context} field 'layers' must be a list or tuple of integer layer indices; got {layers!r}")
+
+    normalized_layers, seen = [], set()
+    for layer_idx in layers:
+        if isinstance(layer_idx, bool) or not isinstance(layer_idx, int):
+            raise ValueError(f"{context} field 'layers' must contain only integer layer indices; got {layer_idx!r}")
+        if layer_idx < 0 or layer_idx >= num_hidden_layers:
+            raise ValueError(
+                f"{context} field 'layers' contains out-of-range layer {layer_idx!r}; "
+                f"expected a value in [0, {num_hidden_layers})")
+        if layer_idx in seen:
+            raise ValueError(f"{context} field 'layers' contains duplicate layer {layer_idx!r}; got {layers!r}")
+        if layer_idx in assigned_layers:
+            raise ValueError(
+                f"{context} assigns conflicting layer {layer_idx!r}, which is already assigned by "
+                f"{_spec_context(assigned_layers[layer_idx])}")
+        seen.add(layer_idx)
+        assigned_layers[layer_idx] = spec_index
+        normalized_layers.append(layer_idx)
+
+    normalized["layers"] = normalized_layers
+    normalized["num_heads"] = _positive_int(normalized["num_heads"], field="num_heads", context=context)
+
+    num_kv_heads = normalized.get("num_kv_heads")
+    if num_kv_heads is None:
+        num_kv_heads = normalized["num_heads"]
+    normalized["num_kv_heads"] = _positive_int(num_kv_heads, field="num_kv_heads", context=context)
+
+    qkv_bias = normalized.get("qkv_bias", False)
+    if not isinstance(qkv_bias, bool):
+        raise ValueError(f"{context} field 'qkv_bias' must be a Boolean; got {qkv_bias!r}")
+    normalized["qkv_bias"] = qkv_bias
+
+    window_size = normalized.get("window_size")
+    if window_size is not None:
+        window_size = _positive_int(window_size, field="window_size", context=context)
+    normalized["window_size"] = window_size
+
+    rope_theta = normalized.get("rope_theta", 10000.0)
+    try:
+        ok = (not isinstance(rope_theta, bool) and isinstance(rope_theta, (int, float))
+              and math.isfinite(rope_theta) and rope_theta > 0)
+    except OverflowError:
+        ok = False
+    if not ok:
+        raise ValueError(f"{context} field 'rope_theta' must be positive and finite; got {rope_theta!r}")
+    normalized["rope_theta"] = rope_theta
+
+    return normalized
+
+
+def normalize_hybrid_attention_config(attn, *, num_hidden_layers: int):
+    """Validate and normalise `attn`: None, one spec dict, or a list of them."""
+    if attn is None:
+        return None
+    if isinstance(num_hidden_layers, bool) or not isinstance(num_hidden_layers, int) or num_hidden_layers < 0:
+        raise ValueError(f"field 'num_hidden_layers' must be a non-negative integer; got {num_hidden_layers!r}")
+    if not isinstance(attn, (dict, list)):
+        raise ValueError(f"attn must be None, a dictionary, or a list of dictionaries; got {attn!r}")
+
+    is_single = isinstance(attn, dict)
+    specs = [attn] if is_single else attn
+    assigned: dict = {}
+    out = []
+    for i, spec in enumerate(specs):
+        idx = None if is_single else i
+        if not isinstance(spec, dict):
+            raise ValueError(f"{_spec_context(idx)} must be a dictionary; got {spec!r}")
+        out.append(_normalize_spec(spec, num_hidden_layers=num_hidden_layers,
+                                   spec_index=idx, assigned_layers=assigned))
+    return out[0] if is_single else out
+
+
+def get_hybrid_attention_spec(attn, *, layer_idx: int):
+    """The normalised spec assigned to `layer_idx`, or None for a linear layer."""
+    if attn is None:
+        return None
+    for spec in ([attn] if isinstance(attn, dict) else attn):
+        if layer_idx in spec["layers"]:
+            return spec
+    return None
+
+
+class ComplexKDAConfig(PretrainedConfig):
+    """Kimi Delta Attention with a *signed* (complex, i.e. Z_2-phased) decay gate.
+
+    The baseline arms set ``gate="sigmoid"`` with ``allow_neg_eigval=False``;
+    the ComplexKDA arms set ``gate="signed_sigmoid2"`` with
+    ``allow_neg_eigval=True``. Everything else is shared, which is what makes
+    the two comparable.
+    """
+
+    model_type = "complex_kda"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        attn_mode: str = "chunk",
+        hidden_size: int = 2048,
+        expand_v: float = 1.0,
+        use_short_conv: bool = True,
+        drop_silu: bool = False,
+        drop_key_silu: bool = False,
+        conv_silu: str = "qkv",
+        allow_neg_eigval: bool = True,
+        gate: str = "signed_sigmoid2",
+        gate_init_style: str = "shipped",
+        output_gate: str = "lowrank",
+        beta_init_style: str = "standard",
+        lower_bound: float = -5.0,
+        num_heads: int = 16,
+        num_v_heads: int | None = None,
+        head_dim: int = 128,
+        num_hidden_layers: int = 24,
+        norm_eps: float = 1e-6,
+        conv_size: int = 4,
+        attn: dict | list | None = None,
+        hidden_ratio: int | None = 4,
+        intermediate_size: int | None = None,
+        hidden_act: str = "swish",
+        max_position_embeddings: int = 4096,
+        initializer_range: float = 0.02,
+        vocab_size: int = 32000,
+        tie_word_embeddings: bool = False,
+        use_cache: bool = True,
+        pad_token_id: int | None = None,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        # Kept so a config written by the training stack round-trips. They
+        # select fused kernels when the fla fork is installed and are ignored
+        # by the pure-torch path, which computes the same thing either way.
+        fuse_norm: bool = True,
+        fuse_swiglu: bool = True,
+        fuse_cross_entropy: bool = True,
+        use_l2warp: bool = False,
+        **kwargs,
+    ):
+        self.attn_mode = attn_mode
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+        self.use_short_conv = use_short_conv
+        self.drop_silu = drop_silu
+        self.drop_key_silu = drop_key_silu
+        self.conv_silu = conv_silu
+        self.allow_neg_eigval = allow_neg_eigval
+        self.gate = gate
+        self.gate_init_style = gate_init_style
+        self.output_gate = output_gate
+        self.beta_init_style = beta_init_style
+        self.lower_bound = lower_bound
+        self.num_heads = num_heads
+        self.num_v_heads = num_v_heads
+        self.head_dim = head_dim
+        # `num_hidden_layers` must be set before `attn`: the layer-range check
+        # in the setter below reads it.
+        self.num_hidden_layers = num_hidden_layers
+        self.norm_eps = norm_eps
+        self.conv_size = conv_size
+        self.attn = attn
+
+        self.hidden_ratio = hidden_ratio
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.vocab_size = vocab_size
+        self.use_cache = use_cache
+
+        self.fuse_norm = fuse_norm
+        self.fuse_swiglu = fuse_swiglu
+        self.fuse_cross_entropy = fuse_cross_entropy
+        self.use_l2warp = use_l2warp
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    # `attn` is a property so that assigning it after construction -- which
+    # `PretrainedConfig.from_dict` does -- is validated too.
+    @property
+    def attn(self):
+        return self.__dict__.get("attn")
+
+    @attn.setter
+    def attn(self, value) -> None:
+        self.__dict__["attn"] = normalize_hybrid_attention_config(
+            value, num_hidden_layers=self.num_hidden_layers)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d261e3881c258757411e47066f053cdbd9169fc605114446e231de045fa4604
+size 2724569536

modeling_complex_kda.py

@@ -0,0 +1,1383 @@
+# Copyright (c) 2026, the ComplexKDA authors.
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li (the parts derived
+# from flash-linear-attention, MIT licensed).
+#
+# SPDX-License-Identifier: MIT
+"""ComplexKDA -- Kimi Delta Attention with a signed (Z_2-phased) decay gate.
+
+STANDALONE. This file needs only `torch` and `transformers`. It carries its own
+implementations of everything the model is made of -- the short convolution,
+the RMS norms, the SwiGLU MLP, the attention layers of the hybrid arms, and the
+gated-delta recurrence itself -- so a checkpoint loads and runs with nothing
+else installed.
+
+IT GOES FASTER WITH THE FORK. When `fla` from
+
+    https://github.com/automl/ComplexKDA
+
+is importable, the Triton kernels and fused modules it ships are used instead,
+and the model is then running exactly the code the checkpoints were trained
+through. Detection is by capability, not by name: upstream flash-linear-attention
+also provides `chunk_kda`, but without the `sign` argument the signed gate needs,
+so the signature is inspected rather than trusted. Set the environment variable
+`COMPLEX_KDA_BACKEND=torch` to force the pure-torch path (useful for debugging a
+numerical difference), or `=kernel` to make a missing fork an error rather than a
+silent fallback.
+
+WHAT THE SIGNED GATE IS. A gated-delta layer carries a per-channel decay
+`alpha`; KDA, like every gated linear attention before it, confines it to
+`(0, 1]`. ComplexKDA lets it take either sign, `alpha in [-1, 1]`, which is the
+one-dimensional real case of a complex eigenvalue -- a channel can now oscillate
+rather than only forget. The magnitude is carried in log space exactly as
+before, and the `+-1` part is carried separately as a running product (the
+"gauge") pushed onto q and k, so the recurrence the kernels run is still the
+unsigned one. `running_sign` below is that product, and `ungauge_state` takes it
+back off the state at a chunk boundary so a cached state is the real one.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import warnings
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+try:  # packaged next to the weights (the Hub layout)
+    from .configuration_complex_kda import ComplexKDAConfig, get_hybrid_attention_spec
+except ImportError:  # imported as a loose file
+    from configuration_complex_kda import ComplexKDAConfig, get_hybrid_attention_spec
+
+logger = logging.get_logger(__name__)
+
+__all__ = [
+    "ComplexKDACache",
+    "ComplexKDAForCausalLM",
+    "ComplexKDAModel",
+    "ComplexKDAPreTrainedModel",
+    "ComplexKimiDeltaAttention",
+]
+
+
+# ===========================================================================
+# optional fast path
+# ===========================================================================
+
+def _detect_fla():
+    """(chunk_kda, fused_recurrent_kda) from the fork, or (None, None).
+
+    The test is that the op ACCEPTS `sign`. Upstream fla exports a `chunk_kda`
+    of the same name that computes the unsigned recurrence; calling it with a
+    signed checkpoint's weights would return a plausible tensor rather than an
+    error, so the presence of the module is not enough.
+    """
+    try:
+        import inspect
+
+        from fla.ops.kda import chunk_kda, fused_recurrent_kda
+    except Exception:
+        return None, None
+    try:
+        if "sign" not in inspect.signature(chunk_kda).parameters:
+            return None, None
+        if "sign" not in inspect.signature(fused_recurrent_kda).parameters:
+            return None, None
+    except (TypeError, ValueError):
+        return None, None
+    return chunk_kda, fused_recurrent_kda
+
+
+_CHUNK_KDA, _FUSED_RECURRENT_KDA = _detect_fla()
+
+
+def _triton_launchable() -> bool:
+    """Stricter than importing triton: fla imports fine on a CPU-only box and
+    only fails when a kernel is launched."""
+    try:
+        import triton  # noqa: F401
+
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
+
+HAS_KERNEL = _CHUNK_KDA is not None and _triton_launchable()
+
+_REQUESTED = os.environ.get("COMPLEX_KDA_BACKEND", "auto").lower()
+if _REQUESTED not in ("auto", "torch", "kernel"):
+    raise ValueError(f"COMPLEX_KDA_BACKEND must be 'auto', 'torch' or 'kernel'; got {_REQUESTED!r}")
+if _REQUESTED == "kernel" and not HAS_KERNEL:
+    raise ImportError(
+        "COMPLEX_KDA_BACKEND=kernel, but the ComplexKDA fla fork's signed kernels are not "
+        "available (need a CUDA device, triton, and `pip install "
+        "git+https://github.com/automl/ComplexKDA`).")
+USE_KERNEL = HAS_KERNEL and _REQUESTED != "torch"
+
+# Chunk length of the portable recurrence. It trades memory for sequential
+# steps: the intra-chunk term materialises a [chunk, chunk, head_dim] block per
+# head, so 64 is a few tens of MB at these geometries and 256 is a few hundred.
+# It does not change what is computed -- only the order the same sums are taken
+# in, which at fp32 moves a logit by ~1e-6 per layer.
+CHUNK_SIZE = int(os.environ.get("COMPLEX_KDA_CHUNK_SIZE", "64"))
+if CHUNK_SIZE <= 0:
+    raise ValueError(f"COMPLEX_KDA_CHUNK_SIZE must be positive; got {CHUNK_SIZE}")
+
+if not USE_KERNEL:
+    logger.warning_once(
+        "ComplexKDA is running its portable torch implementation. For the Triton kernels the "
+        "models were trained with, install the fork: "
+        "`pip install git+https://github.com/automl/ComplexKDA` (and set "
+        "COMPLEX_KDA_BACKEND=torch to keep this path).")
+
+
+def _fla_modules():
+    """fla's fused ShortConvolution / RMSNorm / gated RMSNorm, or (None,)*3."""
+    if not USE_KERNEL:
+        return None, None, None
+    try:
+        from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution
+
+        return ShortConvolution, RMSNorm, FusedRMSNormGated
+    except Exception:
+        return None, None, None
+
+
+_FLA_SHORTCONV, _FLA_RMSNORM, _FLA_RMSNORM_GATED = _fla_modules()
+
+
+# ===========================================================================
+# the gate
+#
+#   name              alpha range   activation
+#   "softplus"        (0, 1]        -exp(A_log) * softplus(u)
+#   "sigmoid"         (0, 1]        lower_bound * sigmoid(A * u)
+#   "signed_sigmoid2" [-1, 1]       2*sigmoid(u) - 1, evaluated as tanh(u/2)
+#   "signed_tanh"     [-1, 1]       tanh(u)
+#
+# The published baselines use "sigmoid"; the ComplexKDA arms use
+# "signed_sigmoid2".
+# ===========================================================================
+
+GATES = ("softplus", "sigmoid", "signed_sigmoid2", "signed_tanh")
+
+
+def is_signed(gate: str) -> bool:
+    if gate not in GATES:
+        raise ValueError(f"gate must be one of {GATES}, got {gate!r}")
+    return gate.startswith("signed_")
+
+
+def safe_gate_ok(gate: str) -> bool:
+    """Whether log|alpha| is bounded below by `lower_bound`. False only for
+    "softplus", which is unbounded."""
+    return gate != "softplus"
+
+
+def signed_gate(z, A_log=None, dt_bias=None, lower_bound: float = -5.0, activation: str = "sigmoid2"):
+    """One pre-activation -> (sign, log|alpha|), for alpha in [-1, 1].
+
+    |alpha| = eps + (1 - eps) * |a|, eps = exp(lower_bound), a = tanh(u/2) or
+    tanh(u). "sigmoid2" is `2*sigmoid(u) - 1` written as `tanh(u/2)`: the literal
+    spelling cancels catastrophically near u = 0, where the SIGN is decided, so
+    it would be settled by rounding rather than by u. The sign shares z with the
+    magnitude and is locally constant, so detaching it is exact.
+    """
+    eps = math.exp(lower_bound)
+    u = z.float()
+    if dt_bias is not None:
+        u = u + dt_bias.view(*([1] * (z.dim() - 2)), *z.shape[-2:])
+    if A_log is not None:
+        u = A_log.float().exp().view(*([1] * (z.dim() - 2)), -1, 1) * u
+    if activation == "sigmoid2":
+        a = torch.tanh(0.5 * u)
+    elif activation == "tanh":
+        a = torch.tanh(u)
+    else:
+        raise ValueError(f"activation must be 'sigmoid2' or 'tanh', got {activation!r}")
+    s = torch.where(a.detach() < 0, -1, 1).to(torch.int8)
+    return s, (eps + (1.0 - eps) * a.abs()).log()
+
+
+def compute_gate(gate, z, A_log=None, dt_bias=None, lower_bound=-5.0):
+    """name -> (sign int8 or None, log|alpha| fp32). A None sign is what tells
+    the caller there is no gauge to apply."""
+    if gate not in GATES:
+        raise ValueError(f"gate must be one of {GATES}, got {gate!r}")
+    if gate.startswith("signed_"):
+        return signed_gate(z, A_log, dt_bias, lower_bound, activation=gate[len("signed_"):])
+    u = z.float()
+    if dt_bias is not None:
+        u = u + dt_bias.view(*([1] * (z.dim() - 2)), *z.shape[-2:])
+    A = A_log.float().exp().view(*([1] * (z.dim() - 2)), -1, 1) if A_log is not None else 1.0
+    if gate == "softplus":
+        return None, -A * F.softplus(u)
+    return None, lower_bound * torch.sigmoid(A * u)
+
+
+def signed_gate_init(dt, lower_bound: float = -5.0, activation: str = "sigmoid2"):
+    """dt_bias giving alpha = +exp(-dt) at step 0."""
+    eps = math.exp(lower_bound)
+    target = ((torch.exp(-dt) - eps) / (1 - eps)).clamp(1e-7, 1 - 1e-7)
+    inv = torch.atanh(target)
+    return 2.0 * inv if activation == "sigmoid2" else inv
+
+
+def gate_init(gate, dt, lower_bound=-5.0):
+    """dt_bias init inverting each gate's own forward, so all four gates start
+    at the same alpha = exp(-dt)."""
+    if gate.startswith("signed_"):
+        return signed_gate_init(dt, lower_bound, gate[len("signed_"):])
+    if gate == "sigmoid":
+        p = (dt / abs(lower_bound)).clamp(1e-7, 1 - 1e-7)
+        return torch.log(p) - torch.log1p(-p)
+    return dt + torch.log(-torch.expm1(-dt))
+
+
+def init_dt_bias(gate, gate_dim=None, lower_bound=-5.0, gate_init_style="shipped", dt=None):
+    if dt is None:
+        dt = torch.exp(
+            torch.rand(gate_dim, dtype=torch.float32) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
+        ).clamp(min=1e-4)
+    init = gate_init(gate, dt, lower_bound)
+    if is_signed(gate) and gate_init_style == "spread":
+        # Same |alpha| as "shipped" with the sign flipped on half the channels:
+        # the activation is odd, so sign and magnitude do not trade off.
+        init = init * torch.where(torch.rand_like(init) < 0.5, -1.0, 1.0)
+    return init
+
+
+# ===========================================================================
+# the gauge: carry the +-1 part of alpha as a running sign on q/k
+# ===========================================================================
+
+def running_sign(s: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor:
+    """P_t = prod_{u<=t} s_u along dim 1, as int8.
+
+    An integer parity prefix sum: exact at any length, and it carries no
+    autograd graph, because the sign has no gradient. Resets at sequence starts
+    when `cu_seqlens` is given.
+    """
+    bits = (s < 0).to(torch.int32)
+    par = bits.cumsum(dim=1)
+    if cu_seqlens is not None:
+        starts = cu_seqlens[:-1]
+        idx = torch.repeat_interleave(starts, cu_seqlens[1:] - starts)
+        par = par - (par[:, idx] - bits[:, idx])
+    return torch.where(par & 1 == 1, -1, 1).to(torch.int8)
+
+
+class _ApplySign(torch.autograd.Function):
+    """x * P, keeping P as int8 rather than letting `mul` upcast it."""
+
+    @staticmethod
+    def forward(ctx, x, P):
+        ctx.save_for_backward(P)
+        return x * P.to(x.dtype)
+
+    @staticmethod
+    def backward(ctx, go):
+        (P,) = ctx.saved_tensors
+        return go * P.to(go.dtype), None
+
+
+def apply_sign(x, P):
+    return _ApplySign.apply(x, P)
+
+
+def ungauge_state(ht, P_last, state_v_first: bool, head_k_dim: int | None = None):
+    """S_T = Diag(P_T) S~_T, on whichever axis holds K.
+
+    `state_v_first=True` stores [N, HV, V, K] -- K LAST -- so the axis differs
+    between the kernel and the torch path; `head_k_dim` turns a silent
+    wrong-axis bug into an assert.
+    """
+    if ht is None or P_last is None:
+        return ht
+    if P_last.ndim != ht.ndim - 1:
+        raise AssertionError(
+            f"gauge rank mismatch: state {tuple(ht.shape)} takes a gauge of "
+            f"{ht.ndim - 1} dims, got {tuple(P_last.shape)}.")
+    axis = -1 if state_v_first else -2
+    if head_k_dim is not None and ht.shape[axis] != head_k_dim:
+        raise AssertionError(
+            f"state layout mismatch: state_v_first={state_v_first} implies K on axis {axis}, "
+            f"but state shape {tuple(ht.shape)} has {ht.shape[axis]} there, not "
+            f"head_k_dim={head_k_dim}.")
+    P = P_last.to(ht.dtype)
+    return ht * (P.unsqueeze(-2) if state_v_first else P.unsqueeze(-1))
+
+
+# ===========================================================================
+# the recurrence, in torch
+#
+# Both functions take q/k ALREADY l2-normalised and gauged, `g` as log|alpha|,
+# and `beta` already through its sigmoid -- the same contract as the reference
+# implementation in the fork, so the two can be compared term by term.
+# ===========================================================================
+
+def recurrent_kda_torch(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+):
+    """The definition, one step at a time. [B,T,H,K] q/k, [B,T,HV,V] v,
+    [B,T,HV,K] g, [B,T,HV] beta; state [B,HV,K,V]."""
+    dtype = v.dtype
+    B, T, H, K = q.shape
+    HV, V = v.shape[2], v.shape[-1]
+    G = HV // H
+    if scale is None:
+        scale = K ** -0.5
+
+    q, k, v, g, beta = (x.float() for x in (q, k, v, g, beta))
+    q = q.repeat_interleave(G, dim=2) * scale
+    k = k.repeat_interleave(G, dim=2)
+
+    S = q.new_zeros(B, HV, K, V)
+    if initial_state is not None:
+        S = S + initial_state.float()
+    o = torch.zeros_like(v)
+    for i in range(T):
+        q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
+        S = S * g_i[..., None].exp()
+        # delta rule: replace the memory currently read out by k_i with v_i
+        S = S + torch.einsum("bhk,bhv->bhkv", b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
+        o[:, i] = torch.einsum("bhk,bhkv->bhv", q_i, S)
+    return o.to(dtype), (S if output_final_state else None)
+
+
+def chunk_kda_torch(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    chunk_size: int = 64,
+):
+    """The same recurrence in chunks: O(T/C) sequential steps instead of O(T).
+
+    The WY/UT transform of the chunk's delta updates, then one state carry per
+    chunk. Arithmetically identical to `recurrent_kda_torch` up to floating
+    point; it exists because a 4096-token forward through the step loop is
+    minutes rather than milliseconds.
+
+    MASK BEFORE EXPONENTIATING. Every exponent used here is a sum of `log|alpha|`
+    over an interval, so it is <= 0 and `exp` is safe -- but only for the pairs
+    the causal mask keeps. The reference implementation exponentiates the full
+    block and masks afterwards, which overflows once `|log alpha| * chunk`
+    passes ~88 in fp32: finite forward, NaN backward. Masking first removes that
+    failure mode entirely, which is why `chunk_size` needs no upper bound here.
+    """
+    dtype = v.dtype
+    B, T, H, K = q.shape
+    HV, V = v.shape[2], v.shape[-1]
+    G = HV // H
+    if scale is None:
+        scale = K ** -0.5
+    BT = int(chunk_size)
+    if BT <= 0:
+        raise ValueError(f"chunk_size must be positive, got {chunk_size}")
+
+    q, k, v, g, beta = (x.float() for x in (q, k, v, g, beta))
+    q = q.repeat_interleave(G, dim=2) * scale
+    k = k.repeat_interleave(G, dim=2)
+
+    # Pad the tail to a whole chunk. beta = 0 makes the padded steps write
+    # nothing and g = 0 makes them decay nothing, so the carried state is
+    # exactly the state at T.
+    pad = (-T) % BT
+    if pad:
+        q = F.pad(q, (0, 0, 0, 0, 0, pad))
+        k = F.pad(k, (0, 0, 0, 0, 0, pad))
+        v = F.pad(v, (0, 0, 0, 0, 0, pad))
+        g = F.pad(g, (0, 0, 0, 0, 0, pad))
+        beta = F.pad(beta, (0, 0, 0, pad))
+    NT = (T + pad) // BT
+
+    # [B, T, HV, X] -> [B, HV, NT, BT, X]
+    def _chunks(x):
+        return x.view(B, NT, BT, *x.shape[2:]).permute(0, 3, 1, 2, *range(4, x.dim() + 1))
+
+    q, k, v, g = (_chunks(x) for x in (q, k, v, g))
+    beta = beta.view(B, NT, BT, HV).permute(0, 3, 1, 2)
+
+    eye = torch.eye(BT, device=q.device, dtype=q.dtype)
+    rows = torch.arange(BT, device=q.device)
+    strictly_lower = rows[:, None] > rows[None, :]          # c > i
+    causal = rows[:, None] >= rows[None, :]                 # c >= j
+    neg_inf = torch.finfo(q.dtype).min
+
+    S = q.new_zeros(B, HV, K, V)
+    if initial_state is not None:
+        S = S + initial_state.float()
+    o = torch.zeros_like(v)
+
+    for n in range(NT):
+        q_n, k_n, v_n, g_n, b_n = q[:, :, n], k[:, :, n], v[:, :, n], g[:, :, n], beta[:, :, n]
+        gc = g_n.cumsum(-2)                                 # [B,HV,BT,K], <= 0
+
+        # T[c,i] = beta_c * <k_c * alpha(i,c], k_i>  for c > i -- the
+        # strictly-lower part of the chunk's own delta interactions.
+        d = gc.unsqueeze(-2) - gc.unsqueeze(-3)             # [B,HV,BT(c),BT(i),K]
+        d = d.masked_fill(~strictly_lower[..., None], neg_inf)
+        A = (k_n.unsqueeze(-2) * d.exp() * k_n.unsqueeze(-3)).sum(-1)
+        del d
+        A = -(A * b_n[..., :, None])
+
+        # (I - A)^{-1}, A strictly lower and hence unit-triangular after +I.
+        # The reference walks the Neumann series row by row; a triangular solve
+        # is the same matrix and vectorises.
+        Ainv = torch.linalg.solve_triangular(eye - A, eye.expand_as(A), upper=False, unitriangular=True)
+        Aw = Ainv * b_n[..., None, :]
+
+        w = Aw @ (gc.exp() * k_n)                           # [B,HV,BT,K]
+        u = Aw @ v_n                                        # [B,HV,BT,V]
+
+        dq = gc.unsqueeze(-2) - gc.unsqueeze(-3)
+        dq = dq.masked_fill(~causal[..., None], neg_inf)
+        Aqk = (q_n.unsqueeze(-2) * dq.exp() * k_n.unsqueeze(-3)).sum(-1)
+        del dq
+
+        v_new = u - w @ S
+        o[:, :, n] = (q_n * gc.exp()) @ S + Aqk @ v_new
+
+        g_last = gc[:, :, -1]                               # [B,HV,K]
+        S = S * g_last.unsqueeze(-1).exp()
+        S = S + ((g_last.unsqueeze(-2) - gc).exp() * k_n).transpose(-1, -2) @ v_new
+
+    o = o.permute(0, 2, 3, 1, 4).reshape(B, NT * BT, HV, V)
+    if pad:
+        o = o[:, :T]
+    return o.to(dtype), (S if output_final_state else None)
+
+
+# ===========================================================================
+# portable modules
+# ===========================================================================
+
+class ShortConvolution(nn.Conv1d):
+    """Causal depthwise conv1d with an optional silu.
+
+    Subclasses nn.Conv1d exactly as the fork's does, so the parameter names
+    match and a checkpoint is portable between this path and the Triton one.
+    """
+
+    def __init__(self, hidden_size, kernel_size=4, bias=False, activation="silu"):
+        super().__init__(hidden_size, hidden_size, kernel_size, groups=hidden_size, bias=bias)
+        if activation not in (None, "silu", "swish"):
+            raise ValueError(f"unsupported activation {activation!r}")
+        self.hidden_size, self.activation = hidden_size, activation
+
+    def forward(self, x, cache=None, output_final_state=False, cu_seqlens=None, **kwargs):
+        if cu_seqlens is not None:
+            raise NotImplementedError(
+                "variable-length batching (cu_seqlens) needs the ComplexKDA fla fork")
+        B, T, D = x.shape
+        w = self.kernel_size[0]
+        h = x.transpose(1, 2)
+        if cache is not None:
+            h = torch.cat([cache, h], dim=-1)[:, :, -(T + w - 1):]
+            pad = w - 1 - (h.shape[-1] - T)
+            if pad > 0:
+                h = F.pad(h, (pad, 0))
+        else:
+            h = F.pad(h, (w - 1, 0))
+        new_cache = h[:, :, -(w - 1):].contiguous() if output_final_state else None
+        y = self._conv_forward(h, self.weight, self.bias)[:, :, :T].transpose(1, 2)
+        if self.activation in ("silu", "swish"):
+            y = F.silu(y)
+        return y, new_cache
+
+
+class RMSNorm(nn.Module):
+    """rms(x) * weight, with the fork's optional fused residual add.
+
+    `forward(x, residual, prenorm=True)` returns `(norm(x + residual), x + residual)`.
+    The add is done in the input dtype, matching the fused kernel called with
+    `residual_in_fp32=False`.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-5, elementwise_affine: bool = True):
+        super().__init__()
+        self.hidden_size, self.eps, self.elementwise_affine = hidden_size, eps, elementwise_affine
+        self.weight = nn.Parameter(torch.ones(hidden_size)) if elementwise_affine else None
+
+    def reset_parameters(self):
+        if self.weight is not None:
+            nn.init.ones_(self.weight)
+
+    def extra_repr(self) -> str:
+        return f"{self.hidden_size}, eps={self.eps}"
+
+    def forward(self, x, residual=None, prenorm: bool = False, residual_in_fp32: bool = False):
+        if residual is not None:
+            x = x + (residual.float() if residual_in_fp32 else residual)
+        dt = x.dtype
+        xf = x.float()
+        y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
+        if self.weight is not None:
+            y = y * self.weight.float()
+        y = y.to(dt)
+        return (y, x) if prenorm else y
+
+
+class FusedRMSNormGated(nn.Module):
+    """rms(x) * weight * act(g). The gate is applied AFTER normalising, which is
+    what the fused kernel does and is not interchangeable with gating first."""
+
+    def __init__(self, hidden_size, elementwise_affine=True, eps=1e-5, activation="swish"):
+        super().__init__()
+        if activation not in ("swish", "silu", "sigmoid"):
+            raise ValueError(f"Unsupported activation: {activation}")
+        self.hidden_size, self.eps, self.activation = hidden_size, eps, activation
+        self.weight = nn.Parameter(torch.ones(hidden_size)) if elementwise_affine else None
+
+    def reset_parameters(self):
+        if self.weight is not None:
+            nn.init.ones_(self.weight)
+
+    def extra_repr(self) -> str:
+        return f"{self.hidden_size}, eps={self.eps}, activation={self.activation}"
+
+    def forward(self, x, g, **kwargs):
+        dt = x.dtype
+        xf = x.float()
+        y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
+        if self.weight is not None:
+            y = y * self.weight.float()
+        gf = g.float()
+        y = y * (torch.sigmoid(gf) if self.activation == "sigmoid" else gf * torch.sigmoid(gf))
+        return y.to(dt)
+
+
+class GatedMLP(nn.Module):
+    """SwiGLU: down_proj(swish(gate_proj(x)) * up_proj(x))."""
+
+    def __init__(self, hidden_size: int, hidden_ratio: int | None = None,
+                 intermediate_size: int | None = None, hidden_act: str = "swish", **kwargs):
+        super().__init__()
+        if hidden_ratio is None:
+            hidden_ratio = 4
+        if intermediate_size is None:
+            intermediate_size = int(hidden_size * hidden_ratio * 2 / 3)
+            intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256)
+        if hidden_act not in ("swish", "silu"):
+            raise ValueError(f"Unsupported hidden_act: {hidden_act}")
+        self.hidden_size, self.hidden_ratio = hidden_size, hidden_ratio
+        self.intermediate_size, self.hidden_act = intermediate_size, hidden_act
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x, **kwargs):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+def _rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary position embedding, half-split convention, applied in fp32.
+
+    Only reached by `use_rope=True` configs; every hybrid published here is
+    NoPE, because the linear layers already carry position.
+    """
+
+    def __init__(self, dim: int, base: float = 10000.0):
+        super().__init__()
+        self.dim, self.base = dim, base
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, q, k, seqlen_offset=0, max_seqlen=None, cu_seqlens=None):
+        if cu_seqlens is not None:
+            raise NotImplementedError("variable-length rotary needs the ComplexKDA fla fork")
+        T = q.shape[1]
+        if torch.is_tensor(seqlen_offset):
+            pos = seqlen_offset.view(-1, 1) + torch.arange(T, device=q.device)
+        else:
+            pos = (torch.arange(T, device=q.device) + int(seqlen_offset)).unsqueeze(0)
+        freqs = pos.float().unsqueeze(-1) * self.inv_freq.to(q.device)
+        emb = torch.cat((freqs, freqs), dim=-1)             # [B or 1, T, dim]
+        cos, sin = emb.cos().unsqueeze(-2), emb.sin().unsqueeze(-2)
+        qf, kf = q.float(), k.float()
+        q = (qf * cos + _rotate_half(qf) * sin).to(q.dtype)
+        k = (kf * cos + _rotate_half(kf) * sin).to(k.dtype)
+        return q, k
+
+
+class Attention(nn.Module):
+    """The attention layers of a hybrid arm.
+
+    Causal, through torch SDPA. Two options are not Llama's and both are on in
+    the published hybrids: `output_gate` is Qwen3-Next's sigmoid computed from
+    the LAYER INPUT and applied before `o_proj` (gating after it would scale the
+    residual contribution instead of the per-head mixture, and `o_proj` mixes
+    heads, so the two differ), and `use_rope=False` is NoPE -- no rotary at all,
+    as in Kimi's hybrid.
+    """
+
+    def __init__(self, hidden_size: int = 2048, num_heads: int = 32, num_kv_heads: int | None = None,
+                 qkv_bias: bool = False, qk_norm: bool = False, output_gate: bool = False,
+                 use_rope: bool = True, window_size: int | None = None,
+                 rope_theta: float | None = 10000.0, max_position_embeddings: int | None = None,
+                 layer_idx: int | None = None):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
+        self.num_kv_groups = num_heads // self.num_kv_heads
+        self.head_dim = hidden_size // num_heads
+        self.kv_dim = self.num_kv_heads * self.head_dim
+        self.qkv_bias, self.qk_norm, self.output_gate, self.use_rope = qkv_bias, qk_norm, output_gate, use_rope
+        self.window_size, self.rope_theta = window_size, rope_theta
+        self.max_position_embeddings, self.layer_idx = max_position_embeddings, layer_idx
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=qkv_bias)
+        self.k_proj = nn.Linear(hidden_size, self.kv_dim, bias=qkv_bias)
+        self.v_proj = nn.Linear(hidden_size, self.kv_dim, bias=qkv_bias)
+        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        if output_gate:
+            self.g_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim)
+            self.k_norm = RMSNorm(self.head_dim)
+        self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta) if use_rope else None
+
+    def forward(self, hidden_states, attention_mask=None, past_key_values=None,
+                output_attentions: bool = False, use_cache: bool = False, **kwargs):
+        if attention_mask is not None and attention_mask.dim() != 2:
+            raise ValueError(
+                "Expected attention_mask as a 0-1 matrix of shape [batch_size, seq_len] "
+                "(0 = padding). Arbitrary [b, q, k] masks are not supported.")
+        if kwargs.get("cu_seqlens") is not None:
+            raise NotImplementedError("variable-length attention needs the ComplexKDA fla fork")
+
+        B, q_len, _ = hidden_states.shape
+        q = self.q_proj(hidden_states).view(B, q_len, self.num_heads, self.head_dim)
+        k = self.k_proj(hidden_states).view(B, q_len, self.num_kv_heads, self.head_dim)
+        v = self.v_proj(hidden_states).view(B, q_len, self.num_kv_heads, self.head_dim)
+        if self.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        seqlen_offset = 0
+        if past_key_values is not None:
+            seqlen_offset = past_key_values.get_seq_length(self.layer_idx)
+            if attention_mask is not None:
+                # Padding sits on the LEFT of a padded batch, so a row's real
+                # position is its offset minus its padding.
+                lens = attention_mask.sum(-1, dtype=torch.long)
+                seqlen_offset = seqlen_offset + lens - attention_mask.shape[-1]
+        if self.rotary is not None:
+            q, k = self.rotary(q, k, seqlen_offset=seqlen_offset)
+
+        if past_key_values is not None:
+            k, v = past_key_values.update_attn(self.layer_idx, k, v, window_size=self.window_size)
+
+        # [B, T, H, D] -> [B, H, T, D]
+        qt, kt, vt = (x.transpose(1, 2) for x in (q, k, v))
+        k_len = kt.shape[2]
+
+        attn_bias = None
+        is_causal = False
+        if q_len == k_len and attention_mask is None and self.window_size is None:
+            is_causal = True
+        else:
+            pos_q = torch.arange(k_len - q_len, k_len, device=q.device)
+            pos_k = torch.arange(k_len, device=q.device)
+            # Bottom-right alignment: query t attends keys <= its own position.
+            keep = pos_k[None, :] <= pos_q[:, None]
+            if self.window_size is not None:
+                keep &= pos_k[None, :] > pos_q[:, None] - self.window_size
+            keep = keep[None, None]
+            if attention_mask is not None:
+                pad = attention_mask[:, None, None, :].bool()
+                if pad.shape[-1] != k_len:
+                    pad = F.pad(pad, (k_len - pad.shape[-1], 0), value=True)
+                keep = keep & pad
+            attn_bias = torch.zeros(keep.shape, dtype=qt.dtype, device=q.device)
+            attn_bias = attn_bias.masked_fill(~keep, torch.finfo(qt.dtype).min)
+
+        gqa = {"enable_gqa": True} if self.num_kv_groups > 1 else {}
+        o = F.scaled_dot_product_attention(qt, kt, vt, attn_mask=attn_bias, is_causal=is_causal, **gqa)
+        o = o.transpose(1, 2).reshape(B, q_len, -1)
+        if self.output_gate:
+            o = o * torch.sigmoid(self.g_proj(hidden_states))
+        return self.o_proj(o), None, past_key_values
+
+
+# ===========================================================================
+# the mixer
+# ===========================================================================
+
+def _identity(x):
+    return x
+
+
+class ComplexKimiDeltaAttention(nn.Module):
+    """Kimi Delta Attention whose decay gate may be negative.
+
+    Beyond KimiDeltaAttention:
+      * `gate` selects the decay parameterisation (two unsigned, two signed);
+      * the `+-1` part of a signed gate is carried as a running sign pushed onto
+        q/k (the gauge), so the recurrence itself still runs on `|alpha|`. On the
+        Triton path the sign is handed to the op as `sign=` and applied inside
+        KDA's own l2norm epilogue; the torch path gauges explicitly here.
+    """
+
+    def __init__(self, hidden_size: int = 2048, expand_v: float = 1, head_dim: int = 128,
+                 num_heads: int = 16, num_v_heads: int | None = None, mode: str = "chunk",
+                 use_short_conv: bool = True, allow_neg_eigval: bool = False,
+                 gate: str = "signed_sigmoid2", drop_silu: bool = False, drop_key_silu: bool = False,
+                 conv_silu: str = "qkv", gate_init_style: str = "shipped",
+                 output_gate: str = "lowrank", beta_init_style: str = "standard",
+                 lower_bound: float = -5.0, conv_size: int = 4, conv_bias: bool = False,
+                 layer_idx: int | None = None, norm_eps: float = 1e-5,
+                 chunk_size: int | None = None, **kwargs):
+        super().__init__()
+
+        if gate not in GATES:
+            raise ValueError(f"gate must be one of {GATES}, got {gate!r}")
+        if gate_init_style not in ("shipped", "spread"):
+            raise ValueError(f"gate_init_style must be 'shipped' or 'spread', got {gate_init_style!r}")
+        if output_gate not in ("lowrank", "linear"):
+            raise ValueError(f"output_gate must be 'lowrank' or 'linear', got {output_gate!r}")
+        if beta_init_style not in ("standard", "spread"):
+            raise ValueError(f"beta_init_style must be 'standard' or 'spread', got {beta_init_style!r}")
+        if mode not in ("chunk", "fused_recurrent"):
+            raise ValueError(f"unsupported mode {mode!r}")
+        if not (-5 <= lower_bound < 0):
+            raise ValueError(f"lower_bound must be in [-5, 0), got {lower_bound}")
+
+        self.mode = mode
+        self.allow_neg_eigval = allow_neg_eigval
+        self.gate = gate
+        self.act = _identity if drop_silu else F.silu
+        self.k_act = _identity if drop_silu or drop_key_silu else F.silu
+        self.gate_init_style = gate_init_style
+        self.beta_init_style = beta_init_style
+        self.safe_gate = safe_gate_ok(gate)
+        self.lower_bound = lower_bound
+        self.hidden_size = hidden_size
+        self.expand_v = expand_v
+        self.chunk_size = CHUNK_SIZE if chunk_size is None else chunk_size
+
+        self.use_short_conv = use_short_conv
+        self.conv_size = conv_size
+        self.conv_bias = conv_bias
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+        self.num_v_heads = num_heads if num_v_heads is None else num_v_heads
+
+        self.head_k_dim = head_dim
+        self.head_v_dim = int(head_dim * expand_v)
+        self.key_dim = int(self.num_heads * self.head_k_dim)
+        self.value_dim = int(self.num_v_heads * self.head_v_dim)
+        self.layer_idx = layer_idx
+
+        if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5):
+            raise ValueError(f"expand_v={expand_v} does not give an integer head_v_dim from head_dim={head_dim}")
+        if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0:
+            raise ValueError(f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}")
+        if self.num_v_heads > self.num_heads and is_signed(gate):
+            warnings.warn(
+                "signed gate under GVA expands q/k to num_v_heads (the gauge is per value head "
+                "but q/k are shared), losing the GVA memory saving.", stacklevel=2)
+
+        self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
+
+        if any(c not in "qkv" for c in conv_silu):
+            raise ValueError(f"conv_silu must be a subset of 'qkv', got {conv_silu!r}")
+        self.conv_silu = "" if drop_silu else conv_silu
+        if drop_key_silu:
+            self.conv_silu = self.conv_silu.replace("k", "")
+
+        conv_cls = _FLA_SHORTCONV or ShortConvolution
+        if use_short_conv:
+            self.q_conv1d = conv_cls(hidden_size=self.key_dim, kernel_size=conv_size, bias=conv_bias,
+                                     activation="silu" if "q" in self.conv_silu else None)
+            self.k_conv1d = conv_cls(hidden_size=self.key_dim, kernel_size=conv_size, bias=conv_bias,
+                                     activation="silu" if "k" in self.conv_silu else None)
+            self.v_conv1d = conv_cls(hidden_size=self.value_dim, kernel_size=conv_size, bias=conv_bias,
+                                     activation="silu" if "v" in self.conv_silu else None)
+
+        self.gate_dim = int(self.num_v_heads * self.head_k_dim)
+        self.f_proj = nn.Sequential(
+            nn.Linear(hidden_size, self.head_v_dim, bias=False),
+            nn.Linear(self.head_v_dim, self.gate_dim, bias=False),
+        )
+        self.b_proj = nn.Linear(hidden_size, self.num_v_heads, bias=beta_init_style == "spread")
+
+        if self.safe_gate:
+            self.A_log = nn.Parameter(torch.zeros(self.num_v_heads, dtype=torch.float32))
+        else:
+            self.A_log = nn.Parameter(torch.log(torch.empty(self.num_v_heads, dtype=torch.float32).uniform_(1, 16)))
+        self.A_log._no_weight_decay = True
+        self.dt_bias = nn.Parameter(init_dt_bias(gate, self.gate_dim, lower_bound, gate_init_style))
+        self.dt_bias._no_weight_decay = True
+
+        # The output forget gate. "lowrank" is fla's and Kimi Linear's factored
+        # hidden -> head_v_dim -> value_dim pair; "linear" is Kimi K3's single
+        # full-rank map. Both sit downstream of the recurrence, so neither
+        # interacts with the signed decay gate.
+        if output_gate == "lowrank":
+            self.g_proj = nn.Sequential(
+                nn.Linear(hidden_size, self.head_v_dim, bias=False),
+                nn.Linear(self.head_v_dim, self.value_dim, bias=True),
+            )
+        else:
+            self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=True)
+        norm_gated_cls = _FLA_RMSNORM_GATED or FusedRMSNormGated
+        self.o_norm = norm_gated_cls(self.head_v_dim, activation="sigmoid", eps=norm_eps)
+        self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
+
+    def forward(self, hidden_states, attention_mask=None, past_key_values=None,
+                use_cache: bool | None = False, output_attentions: bool | None = False, **kwargs):
+        if attention_mask is not None and attention_mask.dim() != 2:
+            raise ValueError(
+                "Expected attention_mask as a 0-1 matrix of shape [batch_size, seq_len] "
+                "(0 = padding). Arbitrary [b, q, k] masks are not supported.")
+        if kwargs.get("cu_seqlens") is not None and not USE_KERNEL:
+            raise NotImplementedError("variable-length batching needs the ComplexKDA fla fork")
+        cu_seqlens = kwargs.get("cu_seqlens")
+
+        B, q_len, _ = hidden_states.shape
+        last_state = None
+        if past_key_values is not None and self.layer_idx is not None:
+            last_state = past_key_values.get(self.layer_idx)
+
+        if self.use_short_conv:
+            cq, ck, cv = last_state["conv_state"] if last_state is not None else (None, None, None)
+            q, cq = self.q_conv1d(x=self.q_proj(hidden_states), cache=cq,
+                                  output_final_state=use_cache, cu_seqlens=cu_seqlens)
+            k, ck = self.k_conv1d(x=self.k_proj(hidden_states), cache=ck,
+                                  output_final_state=use_cache, cu_seqlens=cu_seqlens)
+            v, cv = self.v_conv1d(x=self.v_proj(hidden_states), cache=cv,
+                                  output_final_state=use_cache, cu_seqlens=cu_seqlens)
+        else:
+            cq = ck = cv = None
+            q = self.act(self.q_proj(hidden_states))
+            k = self.k_act(self.k_proj(hidden_states))
+            v = self.act(self.v_proj(hidden_states))
+
+        g = self.f_proj(hidden_states)
+        beta = self.b_proj(hidden_states)
+
+        q = q.view(*q.shape[:-1], -1, self.head_k_dim)
+        k = k.view(*k.shape[:-1], -1, self.head_k_dim)
+        g = g.view(*g.shape[:-1], -1, self.head_k_dim)
+        v = v.view(*v.shape[:-1], -1, self.head_v_dim)
+
+        sign, g = compute_gate(self.gate, g, self.A_log, self.dt_bias, self.lower_bound)
+
+        # GVA: the gauge is per value head but q/k are shared across the group,
+        # so q/k are expanded to HV before either path applies it.
+        if sign is not None and sign.shape[2] != q.shape[2]:
+            r = sign.shape[2] // q.shape[2]
+            q, k = q.repeat_interleave(r, dim=2), k.repeat_interleave(r, dim=2)
+
+        recurrent_state = last_state["recurrent_state"] if last_state is not None else None
+        scale = self.head_k_dim ** -0.5
+        P = None
+
+        if USE_KERNEL:
+            # The kernels take `sign` directly: the gauge rides KDA's own l2norm
+            # epilogue, and the final state comes back already un-gauged.
+            mode = "fused_recurrent" if (q_len <= 64 and not self.training) else self.mode
+            op = _CHUNK_KDA if mode == "chunk" else _FUSED_RECURRENT_KDA
+            extra = dict(use_gate_in_kernel=False, safe_gate=self.safe_gate) if mode == "chunk" else {}
+            o, recurrent_state = op(
+                q=q, k=k, v=v, g=g, beta=beta, sign=sign, scale=scale,
+                initial_state=recurrent_state, output_final_state=bool(use_cache),
+                use_qk_l2norm_in_kernel=True, use_beta_sigmoid_in_kernel=True,
+                allow_neg_eigval=self.allow_neg_eigval, lower_bound=self.lower_bound,
+                state_v_first=True, cu_seqlens=cu_seqlens, **extra)
+            state_v_first = True
+        else:
+            if sign is not None:
+                P = running_sign(sign, cu_seqlens)
+                q, k = apply_sign(q, P), apply_sign(k, P)
+            qn = F.normalize(q.float(), dim=-1, eps=1e-6).to(q.dtype)
+            kn = F.normalize(k.float(), dim=-1, eps=1e-6).to(k.dtype)
+            bt = torch.sigmoid(beta.float()) * (2.0 if self.allow_neg_eigval else 1.0)
+            fn = recurrent_kda_torch if q_len <= 8 else chunk_kda_torch
+            extra = {} if fn is recurrent_kda_torch else dict(chunk_size=min(self.chunk_size, max(q_len, 1)))
+            o, recurrent_state = fn(
+                qn, kn, v, g.to(q.dtype), bt.to(q.dtype), scale=scale,
+                initial_state=recurrent_state, output_final_state=bool(use_cache), **extra)
+            state_v_first = False
+
+        if P is not None and recurrent_state is not None:
+            recurrent_state = ungauge_state(recurrent_state, P[:, -1], state_v_first=state_v_first,
+                                            head_k_dim=self.head_k_dim)
+
+        if use_cache and past_key_values is not None and self.layer_idx is not None:
+            past_key_values.update_recurrent(
+                self.layer_idx,
+                recurrent_state=recurrent_state,
+                conv_state=(cq, ck, cv) if self.use_short_conv else None,
+                state_v_first=state_v_first,
+                offset=q_len,
+            )
+
+        g_out = self.g_proj(hidden_states)
+        o = self.o_norm(o, g_out.view(*g_out.shape[:-1], -1, self.head_v_dim))
+        o = o.reshape(*o.shape[:-2], -1)
+        return self.o_proj(o), None, past_key_values
+
+
+# ===========================================================================
+# cache
+# ===========================================================================
+
+class ComplexKDACache:
+    """Per-layer state for incremental decoding.
+
+    Not a `transformers.Cache`: that class models a growing key/value pair per
+    layer, and a linear-attention layer has a FIXED-SIZE recurrent state plus a
+    short-convolution window instead. Hybrid arms hold both kinds, which is why
+    the two kinds of entry live side by side here.
+
+    A state is stored in whichever layout produced it (`state_v_first` records
+    which), so a cache filled by the Triton path and one filled by the torch
+    path are not interchangeable -- the flag makes that a loud error instead of
+    a transposed state.
+    """
+
+    # Attributes transformers' generation loop probes on whatever cache it was
+    # handed. They are plain class attributes rather than properties so that a
+    # version which reads one this does not define fails on the name it wants
+    # rather than on something further downstream.
+    is_compileable = False
+    is_sliding = False
+
+    def __init__(self, seen_tokens: int = 0):
+        self.states: dict[int, dict[str, Any]] = {}
+        self._seen_tokens = seen_tokens
+
+    def __len__(self) -> int:
+        return len(self.states)
+
+    def get(self, layer_idx: int):
+        return self.states.get(layer_idx)
+
+    def get_seq_length(self, layer_idx: int = 0) -> int:
+        state = self.states.get(layer_idx)
+        return 0 if state is None else state.get("offset", 0)
+
+    def get_max_cache_shape(self, layer_idx: int = 0) -> int | None:
+        return None
+
+    def update_recurrent(self, layer_idx: int, recurrent_state, conv_state,
+                         state_v_first: bool, offset: int):
+        prev = self.states.get(layer_idx)
+        if prev is not None and prev.get("state_v_first") != state_v_first:
+            raise ValueError(
+                f"layer {layer_idx}: cached state was written with state_v_first="
+                f"{prev.get('state_v_first')} and is being updated with {state_v_first}. "
+                "The kernel and torch backends store the state on opposite axes; do not "
+                "switch COMPLEX_KDA_BACKEND part-way through a generation.")
+        self.states[layer_idx] = {
+            "recurrent_state": recurrent_state,
+            "conv_state": conv_state,
+            "state_v_first": state_v_first,
+            "offset": self.get_seq_length(layer_idx) + offset,
+        }
+
+    def update_attn(self, layer_idx: int, k: torch.Tensor, v: torch.Tensor,
+                    window_size: int | None = None):
+        """Append these keys/values and return the full history.
+
+        The offset counts TOKENS SEEN, not calls: a prefill hands over many at
+        once, and under a sliding window it keeps counting after the cache has
+        stopped growing. It is what `Attention` rotates by, so getting it from
+        `k.shape[1]` would put a windowed model's rotary back at the start of
+        the window on every step.
+        """
+        prev = self.states.get(layer_idx)
+        n_new = k.shape[1]
+        if prev is not None and prev.get("attn_state") is not None:
+            pk, pv = prev["attn_state"]
+            k, v = torch.cat([pk, k], dim=1), torch.cat([pv, v], dim=1)
+        # TRIM WHAT IS STORED, RETURN THE WHOLE CONCATENATION. Trimming before
+        # the caller attends would hand a prefill of T > window only the last
+        # `window` keys for ALL T queries -- the early ones would then attend a
+        # window that starts after them. The caller applies the window mask;
+        # this only bounds what the NEXT step has to carry, and since what was
+        # stored is already within the window, the concatenation returned on a
+        # decode step is at most `window + 1` long.
+        stored = (k, v) if window_size is None else (k[:, -window_size:], v[:, -window_size:])
+        self.states[layer_idx] = {
+            "attn_state": stored,
+            "offset": (0 if prev is None else prev.get("offset", 0)) + n_new,
+        }
+        return k, v
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        for state in self.states.values():
+            for key in ("recurrent_state",):
+                if state.get(key) is not None:
+                    state[key] = state[key].index_select(0, beam_idx.to(state[key].device))
+            if state.get("conv_state") is not None:
+                state["conv_state"] = tuple(
+                    None if c is None else c.index_select(0, beam_idx.to(c.device))
+                    for c in state["conv_state"])
+            if state.get("attn_state") is not None:
+                state["attn_state"] = tuple(
+                    t.index_select(0, beam_idx.to(t.device)) for t in state["attn_state"])
+        return self
+
+
+# ===========================================================================
+# the model
+# ===========================================================================
+
+class ComplexKDABlock(nn.Module):
+    def __init__(self, config: ComplexKDAConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        norm_cls = _FLA_RMSNORM or RMSNorm
+
+        self.attn_norm = norm_cls(config.hidden_size, eps=config.norm_eps)
+        spec = get_hybrid_attention_spec(config.attn, layer_idx=layer_idx)
+        if spec is not None:
+            # `qk_norm`, `output_gate` and `use_rope` are read with .get: they
+            # are optional keys that the config preserves rather than fields of
+            # the spec, and their defaults are Attention's own. The published
+            # hybrids need the last two -- their attention is GATED and NoPE --
+            # and without them an exported hybrid is a different model: rotary
+            # where the run had none, and no `g_proj` at all.
+            self.attn = Attention(
+                hidden_size=config.hidden_size,
+                num_heads=spec["num_heads"],
+                num_kv_heads=spec["num_kv_heads"],
+                qkv_bias=spec["qkv_bias"],
+                qk_norm=spec.get("qk_norm", False),
+                output_gate=spec.get("output_gate", False),
+                use_rope=spec.get("use_rope", True),
+                window_size=spec["window_size"],
+                rope_theta=spec["rope_theta"],
+                max_position_embeddings=config.max_position_embeddings,
+                layer_idx=layer_idx,
+            )
+        else:
+            self.attn = ComplexKimiDeltaAttention(
+                mode=config.attn_mode,
+                hidden_size=config.hidden_size,
+                expand_v=config.expand_v,
+                head_dim=config.head_dim,
+                num_heads=config.num_heads,
+                num_v_heads=config.num_v_heads,
+                use_short_conv=config.use_short_conv,
+                drop_silu=config.drop_silu,
+                drop_key_silu=config.drop_key_silu,
+                allow_neg_eigval=config.allow_neg_eigval,
+                gate=config.gate,
+                gate_init_style=config.gate_init_style,
+                output_gate=config.output_gate,
+                conv_silu=config.conv_silu,
+                beta_init_style=config.beta_init_style,
+                lower_bound=config.lower_bound,
+                conv_size=config.conv_size,
+                norm_eps=config.norm_eps,
+                layer_idx=layer_idx,
+            )
+        self.mlp_norm = norm_cls(config.hidden_size, eps=config.norm_eps)
+        self.mlp = GatedMLP(
+            hidden_size=config.hidden_size,
+            hidden_ratio=config.hidden_ratio,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+        )
+
+    def forward(self, hidden_states, attention_mask=None, past_key_values=None,
+                use_cache: bool | None = False, output_attentions: bool | None = False, **kwargs):
+        residual = hidden_states
+        hidden_states = self.attn_norm(hidden_states)
+        hidden_states, attentions, past_key_values = self.attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            **kwargs,
+        )
+        hidden_states, residual = self.mlp_norm(hidden_states, residual, True)
+        hidden_states = self.mlp(hidden_states)
+        return residual + hidden_states, attentions, past_key_values
+
+
+class ComplexKDAPreTrainedModel(PreTrainedModel):
+    config_class = ComplexKDAConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["ComplexKDABlock"]
+    _supports_sdpa = True
+    _can_compile_fullgraph = False
+
+    def _init_weights(self, module: nn.Module):
+        std = self.config.initializer_range
+        if isinstance(module, ComplexKimiDeltaAttention):
+            if next(module.parameters()).device.type != "meta":
+                with torch.no_grad():
+                    module.A_log.zero_()
+                    dt = torch.exp(
+                        torch.rand_like(module.dt_bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
+                    ).clamp(min=1e-4)
+                    module.dt_bias.copy_(init_dt_bias(
+                        module.gate, lower_bound=module.lower_bound,
+                        gate_init_style=module.gate_init_style, dt=dt))
+            return
+        if isinstance(module, (nn.Linear, nn.Conv1d)):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+            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=std)
+        elif hasattr(module, "reset_parameters"):
+            module.reset_parameters()
+
+
+class ComplexKDAModel(ComplexKDAPreTrainedModel):
+    def __init__(self, config: ComplexKDAConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [ComplexKDABlock(config, i) for i in range(config.num_hidden_layers)])
+        self.norm = (_FLA_RMSNORM or RMSNorm)(config.hidden_size, eps=config.norm_eps)
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings = value
+
+    def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None,
+                past_key_values=None, use_cache=None, output_attentions=None,
+                output_hidden_states=None, return_dict=None, **kwargs):
+        if output_attentions:
+            warnings.warn("ComplexKDAModel does not support `output_attentions`; setting it to False.")
+        output_attentions = False
+        output_hidden_states = (output_hidden_states if output_hidden_states is not None
+                                else self.config.output_hidden_states)
+        use_cache = use_cache if use_cache is not None else (self.config.use_cache and not self.training)
+        return_dict = return_dict if return_dict is not None else True
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        hidden_states = self.embeddings(input_ids) if inputs_embeds is None else inputs_embeds
+
+        if use_cache and past_key_values is None:
+            past_key_values = ComplexKDACache()
+        if past_key_values is not None and not isinstance(past_key_values, ComplexKDACache):
+            raise TypeError(
+                f"ComplexKDA needs a ComplexKDACache (it stores recurrent state, not key/value "
+                f"pairs); got {type(past_key_values).__name__}.")
+
+        all_hidden_states = () if output_hidden_states else None
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                hidden_states, _, past_key_values = self._gradient_checkpointing_func(
+                    layer.__call__, hidden_states, attention_mask, past_key_values, use_cache,
+                    output_attentions, **kwargs)
+            else:
+                hidden_states, _, past_key_values = layer(
+                    hidden_states, attention_mask=attention_mask, past_key_values=past_key_values,
+                    use_cache=use_cache, output_attentions=output_attentions, **kwargs)
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(x for x in (hidden_states, past_key_values, all_hidden_states) if x is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )
+
+
+def _tied_weights_keys_declaration():
+    """How this transformers spells "lm_head.weight IS the embedding".
+
+    Every ladder cell ties its embeddings (the 1.3B arms do not), and the two
+    transformers generations declare that differently:
+
+        4.x   a LIST of regex patterns matched against parameter names
+        5.x   a {target: source} MAPPING
+
+    THE WRONG ONE IS NOT A WARNING. A list under transformers 5 raises
+    `'list' object has no attribute 'keys'` from inside `post_init` -- for
+    every tied checkpoint, and only for tied ones, so it passes every test run
+    against an untied model and then fails for most of the release.
+    """
+    mapping = {"lm_head.weight": "model.embeddings.weight"}
+    patterns = ["lm_head.weight"]
+    try:
+        from transformers.modeling_utils import PreTrainedModel
+
+        annotation = str(getattr(PreTrainedModel, "__annotations__", {})
+                         .get("_tied_weights_keys", ""))
+        if annotation:
+            return mapping if "dict" in annotation.lower() else patterns
+    except Exception:
+        pass
+    try:
+        import transformers
+
+        return mapping if int(str(transformers.__version__).split(".")[0]) >= 5 else patterns
+    except Exception:
+        return patterns
+
+
+class ComplexKDAForCausalLM(ComplexKDAPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = _tied_weights_keys_declaration()
+
+    def __init__(self, config: ComplexKDAConfig):
+        super().__init__(config)
+        self.model = ComplexKDAModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.criterion = None
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def get_decoder(self):
+        return self.model
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def tie_weights(self, *args, **kwargs):
+        """Tie the head to the embedding OURSELVES, rather than describing the
+        tie and hoping this transformers acts on the description.
+
+        Every ladder cell ties, and the exporters write NO `lm_head.weight` for
+        a tied geometry -- there is no second tensor to write. transformers 5.3
+        nonetheless decides the key "is present in the checkpoint", declines to
+        tie, and leaves the head on the META device: `from_pretrained` returns
+        without error and the first `.to(device)` raises "Cannot copy out of
+        meta tensor". A CPU-only smoke test does not even get that far -- it
+        returns a model whose head is data-less.
+
+        Doing the assignment here is version-independent, and transformers
+        calls this both in `post_init` and after loading the weights, so the
+        alias survives materialisation.
+        """
+        if getattr(self.config, "tie_word_embeddings", False):
+            embeddings = self.get_input_embeddings()
+            if embeddings is not None:
+                self.lm_head.weight = embeddings.weight
+        # *args/**kwargs: transformers 5.6 passes `recompute_mapping`, 4.x
+        # passes nothing. Forward whatever it sends rather than pinning a
+        # signature that one of them will not call.
+        return super().tie_weights(*args, **kwargs)
+
+    def forward(self, input_ids=None, attention_mask=None, inputs_embeds=None,
+                past_key_values=None, labels=None, use_cache=None, output_attentions=None,
+                output_hidden_states=None, return_dict=None, logits_to_keep=0, **kwargs):
+        return_dict = return_dict if return_dict is not None else True
+        outputs = self.model(
+            input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values, use_cache=use_cache,
+            output_attentions=output_attentions, output_hidden_states=output_hidden_states,
+            return_dict=return_dict, **kwargs)
+        hidden_states = outputs[0]
+        if logits_to_keep:
+            hidden_states = hidden_states[:, -logits_to_keep:]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            criterion = self.criterion if self.criterion is not None else nn.CrossEntropyLoss()
+            labels = labels.to(logits.device)
+            # Shift here rather than on the logits, matching the training stack.
+            labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1)
+            loss = criterion(logits.view(labels.numel(), -1), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + tuple(outputs[1:])
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=None,
+        )
+
+    # ---- generation -------------------------------------------------------
+    #
+    # `generate` builds a DynamicCache by default, which this model cannot use
+    # (see ComplexKDACache). Installing ours here is the documented escape
+    # route: a cache already present in model_kwargs is left alone.
+
+    def _prepare_cache_for_generation(self, generation_config, model_kwargs, *args, **kwargs):
+        # *args absorbs the positional tail, which differs across transformers
+        # versions; the two arguments this needs have not moved.
+        if generation_config.use_cache and model_kwargs.get("past_key_values") is None:
+            model_kwargs["past_key_values"] = ComplexKDACache()
+            return True
+        return False
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None,
+                                      inputs_embeds=None, use_cache=True, logits_to_keep=None, **kwargs):
+        if past_key_values is not None and len(past_key_values) > 0:
+            input_ids = input_ids[:, -1:]
+        model_inputs = {"input_ids": input_ids, "inputs_embeds": None}
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"input_ids": None, "inputs_embeds": inputs_embeds}
+        model_inputs.update(
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            attention_mask=attention_mask,
+            logits_to_keep=1 if logits_to_keep is None else logits_to_keep,
+        )
+        return model_inputs
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        return past_key_values.reorder_cache(beam_idx)

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' '  + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}",
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "model_max_length": 4096,
+  "pad_token": null,
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}