Add transformers configuration_vortex.py + modeling_vortex.py and auto_map

Registers the architecture with AutoConfig/AutoModelForCausalLM so the model loads via trust_remote_code. Adds the KV cache (generate() was previously impossible), a bottom-right-aligned attention mask for cached decoding, and real ModelOutput types. Weights are unchanged - model.safetensors is byte-identical and not re-uploaded.

config.json

@@ -2,7 +2,13 @@
   "architectures": [
     "VortexForCausalLM"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_vortex.VortexConfig",
+    "AutoModelForCausalLM": "modeling_vortex.VortexForCausalLM"
+  },
+  "bos_token_id": 1,
   "dtype": "float32",
+  "eos_token_id": 2,
   "hidden_size": 512,
   "initializer_range": 0.02,
   "intermediate_size": 1072,
@@ -11,6 +17,7 @@
   "num_attention_heads": 8,
   "num_hidden_layers": 18,
   "num_key_value_heads": 2,
+  "pad_token_id": 0,
   "rms_norm_eps": 1e-06,
   "rope_interleaved": true,
   "rope_theta": 10000.0,

configuration_vortex.py

@@ -0,0 +1,205 @@
+"""Vortex configuration — Hugging Face `PretrainedConfig` subclass.
+
+Self-contained on purpose. When `trust_remote_code=True` is used,
+`transformers` copies `configuration_vortex.py` and `modeling_vortex.py` into
+`~/.cache/huggingface/modules/transformers_modules/<repo>/` and imports them as
+*top-level* modules. Any import of a sibling file in this repository (e.g.
+`from config import VortexArch`) would fail at that point, so this file may only
+depend on the standard library and `transformers`.
+
+Registering with the auto classes is what makes the checkpoint loadable with a
+plain `AutoModelForCausalLM.from_pretrained(...)`:
+
+    AutoConfig.register("vortex", VortexConfig)
+    AutoModelForCausalLM.register(VortexConfig, VortexForCausalLM)
+
+`src/export_hf.py` writes the equivalent `auto_map` block into `config.json`,
+which is the serialised form of those two calls.
+"""
+
+from __future__ import annotations
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class VortexConfig(PretrainedConfig):
+    """Configuration for the Vortex decoder-only Transformer.
+
+    The defaults are the `vortex-50m-16k` preset: a 512d x 18L model with a
+    16,384-token tied embedding table and 8Q/2KV grouped-query attention.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 16384):
+            Size of the token embedding table. With `tie_word_embeddings=True`
+            this is also the size of the output head, and it is the single
+            biggest lever on the parameter budget at this scale.
+        hidden_size (`int`, *optional*, defaults to 512):
+            Model dimension. Must be divisible by `num_attention_heads`.
+        num_hidden_layers (`int`, *optional*, defaults to 18):
+            Number of decoder blocks.
+        num_attention_heads (`int`, *optional*, defaults to 8):
+            Number of query heads. `hidden_size // num_attention_heads` is the
+            head dimension and must be even for RoPE.
+        num_key_value_heads (`int`, *optional*, defaults to 2):
+            Number of key/value heads. Fewer than `num_attention_heads` selects
+            grouped-query attention (GQA); must divide `num_attention_heads`.
+        intermediate_size (`int`, *optional*, defaults to 1072):
+            SwiGLU feed-forward width, ~2.09x `hidden_size`.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-6):
+            Epsilon inside every RMSNorm.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            RoPE base. Higher values stretch the wavelength of the
+            high-frequency rotary components.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            Maximum context length. The RoPE tables are built to this size and
+            grow on demand if a longer sequence is actually seen.
+        use_qk_norm (`bool`, *optional*, defaults to `True`):
+            Per-head RMSNorm on queries and keys before the attention matmul.
+            The main defence against attention entropy collapse in small
+            models; costs 2 * head_dim parameters per layer.
+        tie_word_embeddings (`bool`, *optional*, defaults to `True`):
+            Share the `lm_head` weight with the input embedding. Halves the
+            vocabulary-sized parameter cost.
+        zero_init_residual (`bool`, *optional*, defaults to `True`):
+            Initialise `o_proj` and `down_proj` to exactly zero so every block is
+            an identity at step 0. Only affects fresh initialisation — it has no
+            effect on loading trained weights.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            Standard deviation of the normal init for linear and embedding
+            weights.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Return a key/value `Cache` from `forward` so `generate` runs in
+            O(1) per token instead of re-running the full prefix.
+        scale_residual (`bool`, *optional*, defaults to `False`):
+            Scale residual branch outputs by `1/sqrt(2 * num_hidden_layers)`
+            (GPT-2 style). Redundant next to zero-init residuals, so off.
+        rope_interleaved (`bool`, *optional*, defaults to `True`):
+            `True` uses the GPT-NeoX split-half pairing (`x1, x2 = x.chunk(2)`);
+            `False` uses the interleaved-even/odd pairing. Recorded for
+            provenance; the split-half layout is what the released weights were
+            trained with.
+    """
+
+    model_type = "vortex"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Defaults mirror the `vortex-50m-16k` preset (src/config.py::VortexArch).
+    # They are duplicated rather than imported so this file stays standalone.
+    def __init__(
+        self,
+        vocab_size: int = 16_384,
+        hidden_size: int = 512,
+        num_hidden_layers: int = 18,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 2,
+        intermediate_size: int = 1_072,
+        rms_norm_eps: float = 1e-6,
+        rope_theta: float = 10_000.0,
+        max_position_embeddings: int = 2_048,
+        use_qk_norm: bool = True,
+        tie_word_embeddings: bool = True,
+        zero_init_residual: bool = True,
+        initializer_range: float = 0.02,
+        use_cache: bool = True,
+        scale_residual: bool = False,
+        rope_interleaved: bool = True,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        pad_token_id: int = 0,
+        **kwargs,
+    ):
+        self.vocab_size = int(vocab_size)
+        self.hidden_size = int(hidden_size)
+        self.num_hidden_layers = int(num_hidden_layers)
+        self.num_attention_heads = int(num_attention_heads)
+        self.num_key_value_heads = int(num_key_value_heads)
+        self.intermediate_size = int(intermediate_size)
+        self.rms_norm_eps = float(rms_norm_eps)
+        self.rope_theta = float(rope_theta)
+        self.max_position_embeddings = int(max_position_embeddings)
+        self.use_qk_norm = bool(use_qk_norm)
+        self.zero_init_residual = bool(zero_init_residual)
+        self.initializer_range = float(initializer_range)
+        self.scale_residual = bool(scale_residual)
+        self.rope_interleaved = bool(rope_interleaved)
+        self.name_or_path = kwargs.pop("name_or_path", "")
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=bool(tie_word_embeddings),
+            **kwargs,
+        )
+
+        # `use_cache` is a model-level flag, not a base `PretrainedConfig`
+        # attribute — transformers 5 dropped it from the base class, so setting
+        # it here is what makes `config.use_cache` readable on a loaded config.
+        self.use_cache = bool(use_cache)
+
+        self.validate()
+
+    # ── derived ──────────────────────────────────────────────────────
+    # `hidden_size` and `num_attention_heads` are also the names of the two
+    # outermost `__init__` parameters, so these are read from the instance
+    # rather than the caller's arguments. A `head_dim` passed in the config JSON
+    # is a *derived* value: recomputing it keeps the model and its config from
+    # disagreeing if someone edits one and not the other.
+
+    @property
+    def head_dim(self) -> int:
+        """Query/key/value head dimension."""
+        return self.hidden_size // self.num_attention_heads
+
+    @property
+    def num_query_groups(self) -> int:
+        """Query heads served by each KV head under GQA."""
+        return self.num_attention_heads // self.num_key_value_heads
+
+    # ── validation ───────────────────────────────────────────────────
+    def validate(self) -> None:
+        """Reject an illegal shape at construction time.
+
+        Without this, a bad GQA split surfaces as an opaque SDPA error
+        ("heads in key and value must divide the number of heads in query")
+        layers deep inside a forward pass.
+        """
+        if self.hidden_size <= 0:
+            raise ValueError(f"hidden_size must be positive, got {self.hidden_size}")
+        if self.num_attention_heads <= 0:
+            raise ValueError(
+                f"num_attention_heads must be positive, got {self.num_attention_heads}"
+            )
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size {self.hidden_size} is not divisible by "
+                f"num_attention_heads {self.num_attention_heads}"
+            )
+        if self.num_key_value_heads < 1:
+            raise ValueError(
+                f"num_key_value_heads must be >= 1, got {self.num_key_value_heads}"
+            )
+        if self.num_key_value_heads > self.num_attention_heads:
+            raise ValueError(
+                f"num_key_value_heads {self.num_key_value_heads} exceeds "
+                f"num_attention_heads {self.num_attention_heads}"
+            )
+        if self.num_attention_heads % self.num_key_value_heads != 0:
+            raise ValueError(
+                f"num_attention_heads {self.num_attention_heads} is not divisible by "
+                f"num_key_value_heads {self.num_key_value_heads}; GQA needs whole "
+                f"query groups"
+            )
+        if self.head_dim % 2 != 0:
+            raise ValueError(
+                f"head_dim {self.head_dim} must be even for RoPE; got "
+                f"hidden_size {self.hidden_size} / {self.num_attention_heads} heads"
+            )
+        if self.vocab_size <= 0:
+            raise ValueError(f"vocab_size must be positive, got {self.vocab_size}")
+
+
+__all__ = ["VortexConfig"]

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "5.17.0",
+  "use_cache": true
+}

modeling_vortex.py

@@ -0,0 +1,970 @@
+"""Vortex modeling — Hugging Face `PreTrainedModel` implementation.
+
+Self-contained on purpose. With `trust_remote_code=True`, `transformers` copies
+`configuration_vortex.py` and `modeling_vortex.py` into
+`~/.cache/huggingface/modules/transformers_modules/<repo>/` and imports them as
+a package, so nothing here may import a sibling file from this repository.
+`configuration_vortex` is the only dependency and it travels with this module, so
+the pair is always copied together — see `_load_config_class` for why the import
+is written the way it is.
+
+What this adds over a bare `nn.Module` port, and why each piece is needed for
+`AutoModelForCausalLM` / `generate` / `Trainer` to work:
+
+* **Key/value cache.** `use_cache` was a config field with no implementation —
+  every `forward` recomputed the whole prefix. `VortexAttention` now consumes a
+  `transformers` `Cache`, which is what makes `model.generate()` viable.
+* **Position offsets under a cache.** RoPE was sliced `cos[:T]`, i.e. positions
+  were always 0-based. With a cache the query block starts at `past_len`; the
+  rotary tables are now sliced `[offset : offset + T]`. RoPE is relative, so this
+  leaves the pretraining fast path bit-identical.
+* **A correct attention mask on the cached path.** SDPA's `is_causal=True`
+  assumes top-left alignment and is only right when the cache is empty. Cached
+  steps with left padding need an explicit bottom-right-aligned mask, which is
+  what `VortexModel._build_causal_mask` builds. The empty-cache/no-padding case
+  still takes the `is_causal=True` fast path, so training numerics and memory are
+  unchanged.
+* **Real `ModelOutput`s.** The previous `CausalLMOutput` was a plain object, so
+  `output.logits` worked but nothing HF-side (generation, `Trainer`, tensor
+  logging) recognised it.
+* **Standard input plumbing** — `attention_mask`, `position_ids`,
+  `inputs_embeds`, `num_items_in_batch`, `logits_to_keep`.
+
+Two deliberate deviations from HF naming conventions:
+
+* The decoder submodules keep their original names (`attn`, `ln_attn`, `ln_mlp`)
+  rather than `self_attn`, `input_layernorm`, `post_attention_layernorm`. HF's
+  `self_attn` means *cross*-attention, which this architecture does not have.
+  More importantly, the released checkpoints on the Hub use the current names,
+  and `from_pretrained` matches `state_dict` keys literally — renaming would
+  break every one of them unless a key-remapping table were threaded through
+  `from_pretrained`, which is a per-version API in transformers 5.x. The outer
+  names (`model.*`, `embed_tokens`, `lm_head`, `norm`) already match HF.
+* `logits_to_keep` is not decoration. Computing `(B, T, vocab_size)` logits for a
+  full 2048-token batch is the largest single memory term in a training step, and
+  the whole point of the chunked loss path is to never materialise it. That is
+  why `labels=` returns `logits=None` unless logits are explicitly asked for.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import math
+import os
+import sys
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+
+def _load_config_class():
+    """Import the sibling `configuration_vortex` module.
+
+    Two different import mechanics have to be satisfied:
+
+    * Under `trust_remote_code`, `transformers` copies both files into its
+      dynamic-module package and imports them as a package, so a *relative*
+      import is the one that resolves.
+    * Running the repo's own tests imports this file as a top-level module from
+      `src/`, where there is no package and no `__package__`.
+
+    A plain top-level `from configuration_vortex import ...` is not an option:
+    `dynamic_module_utils.check_imports` runs `importlib.import_module` on every
+    statically-detected import *before* the sibling has been copied next to this
+    file, so it fails with "No module named 'configuration_vortex'" and a
+    misleading `pip install configuration_vortex`. Loading by file path keeps the
+    statement out of the AST the checker inspects.
+    """
+    if __package__:
+        from .configuration_vortex import VortexConfig
+
+        return VortexConfig
+
+    spec = importlib.util.spec_from_file_location(
+        "configuration_vortex",
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "configuration_vortex.py"),
+    )
+    module = importlib.util.module_from_spec(spec)
+    # Registered before exec so the dataclass-free class object survives even if
+    # something inside the module re-enters this lookup.
+    sys.modules["configuration_vortex"] = module
+    spec.loader.exec_module(module)
+    return module.VortexConfig
+
+
+VortexConfig = _load_config_class()
+
+logger = logging.get_logger(__name__)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Norm
+# ──────────────────────────────────────────────────────────────────────
+class VortexRMSNorm(nn.Module):
+    """RMSNorm with the reduction and the norm-weight multiply in fp32.
+
+    Upcasting is the point: with 18 pre-norm blocks in bf16 autocast, a bf16
+    reduction over the residual stream loses enough precision to stall training.
+    The output is cast back so the residual add stays in the activation dtype.
+    """
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.normalized_shape = (hidden_size,)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return (self.weight.float() * hidden_states).to(input_dtype)
+
+    def extra_repr(self) -> str:
+        return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Rotary position embedding
+# ──────────────────────────────────────────────────────────────────────
+def build_rope_cache(
+    head_dim: int,
+    max_seq_len: int,
+    base: float = 10_000.0,
+    device=None,
+    dtype: torch.dtype = torch.float32,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Build the `(max_seq_len, head_dim / 2)` cos/sin tables for RoPE."""
+    if head_dim % 2 != 0:
+        raise ValueError(f"head_dim must be even, got {head_dim}")
+    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    position_ids = torch.arange(max_seq_len, device=device, dtype=torch.float32)
+    freqs = torch.outer(position_ids, inv_freq)
+    return freqs.cos().to(dtype), freqs.sin().to(dtype)
+
+
+def apply_rope(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    offset: int = 0,
+) -> torch.Tensor:
+    """Rotate the last dim of `x` (GPT-NeoX split-half pairing).
+
+    `x` is `(batch, heads, seq, head_dim)`. `cos`/`sin` are `(seq, head_dim / 2)`
+    *absolute* position tables; `offset` selects the starting position, which is
+    what puts a cached query block on the right rotary phase.
+
+    Pre-sliced tables with the default `offset=0` are still accepted, so the
+    direct-call form used by the verification suite keeps working.
+    """
+    if x.shape[-2] != cos.shape[0] or offset != 0:
+        T = x.shape[-2]
+        cos = cos[offset : offset + T]
+        sin = sin[offset : offset + T]
+    cos = cos.unsqueeze(0).unsqueeze(0)
+    sin = sin.unsqueeze(0).unsqueeze(0)
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class VortexRotaryEmbedding(nn.Module):
+    """Per-model RoPE table, built once and shared by every attention layer.
+
+    Held in non-persistent state so it never becomes a checkpoint tensor — it is
+    fully determined by `head_dim`, `rope_theta` and the current device/dtype.
+    """
+
+    def __init__(self, config: VortexConfig, device=None):
+        super().__init__()
+        self.config = config
+        self.max_seq_len_cached = config.max_position_embeddings
+
+        # Plain attributes, deliberately not buffers. `from_pretrained` builds
+        # the model on a meta device and materialises only the tensors it finds
+        # in the checkpoint, so a *non-persistent* buffer is left as
+        # uninitialised memory: the model loads without error and every RoPE
+        # application is garbage. Keeping this out of `state_dict` also means the
+        # key layout stays identical to the released training checkpoints, which
+        # is what lets `load_state_dict(strict=True)` accept them.
+        self._inv_freq: Optional[torch.Tensor] = None
+        self._inv_freq_device: Optional[torch.device] = device
+        self._cos: Optional[torch.Tensor] = None
+        self._sin: Optional[torch.Tensor] = None
+        self._cached_len = 0
+        self._cached_dtype: Optional[torch.dtype] = None
+
+    def _get_inv_freq(self, device: torch.device) -> torch.Tensor:
+        head_dim = self.config.head_dim
+        if self._inv_freq is None or self._inv_freq_device != device:
+            self._inv_freq = 1.0 / (
+                self.config.rope_theta
+                ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim)
+            )
+            self._inv_freq_device = device
+            # Invalidate the cos/sin tables; they were built from the old one.
+            self._cos = self._sin = None
+        return self._inv_freq
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return cos/sin tables covering at least `seq_len` positions."""
+        device, dtype = x.device, x.dtype
+        if (
+            self._cos is None
+            or self._cached_len < seq_len
+            or self._cos.device != device
+            or self._cached_dtype != dtype
+        ):
+            inv_freq = self._get_inv_freq(device)
+            self._cached_len = max(seq_len, self.config.max_position_embeddings)
+            position_ids = torch.arange(self._cached_len, device=device, dtype=torch.float32)
+            freqs = torch.outer(position_ids, inv_freq)
+            self._cos = freqs.cos().to(dtype)
+            self._sin = freqs.sin().to(dtype)
+            self._cached_dtype = dtype
+        return self._cos, self._sin
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Attention
+# ──────────────────────────────────────────────────────────────────────
+class VortexAttention(nn.Module):
+    """Causal grouped-query attention with optional QK-Norm.
+
+    Goes through `F.scaled_dot_product_attention` with no hand-written softmax,
+    which lets PyTorch dispatch to FlashAttention-2 on Ampere and later and to
+    the math backend everywhere else. `enable_gqa` avoids materialising repeated
+    KV heads; the `repeat_interleave` branch only runs on torch < 2.5.
+    """
+
+    def __init__(self, config: VortexConfig, layer_idx: int = 0):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.n_heads = config.num_attention_heads
+        self.n_kv = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.n_groups = self.n_heads // self.n_kv
+        self.rope_theta = config.rope_theta
+        self.use_qk_norm = bool(config.use_qk_norm)
+        self.scale = self.head_dim**-0.5
+
+        hidden_size = config.hidden_size
+        self.q_proj = nn.Linear(hidden_size, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, self.n_kv * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, self.n_kv * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, hidden_size, bias=False)
+
+        if self.use_qk_norm:
+            # Per-head RMS over head_dim, applied before the attention matmul.
+            # Without it, small models hit attention entropy collapse early: a
+            # few heads saturate, their softmax goes one-hot, and those heads
+            # are dead for the rest of the run. Costs 2 * head_dim params/layer.
+            self.q_norm = VortexRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = VortexRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        else:
+            self.q_norm = self.k_norm = nn.Identity()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        position_offset: int = 0,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Cache] = None,
+    ) -> torch.Tensor:
+        B, T, C = x.shape
+
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_kv, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_kv, self.head_dim).transpose(1, 2)
+
+        # QK-Norm: bound the pre-softmax logits before RoPE mixes them.
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if position_embeddings is None:
+            position_embeddings = build_rope_cache(
+                self.head_dim, position_offset + T, self.rope_theta, x.device, x.dtype
+            )
+        cos, sin = position_embeddings
+        q = apply_rope(q, cos, sin, offset=position_offset)
+        k = apply_rope(k, cos, sin, offset=position_offset)
+
+        if past_key_value is not None:
+            k, v = past_key_value.update(k, v, self.layer_idx)
+
+        # `is_causal=True` is only correct when the cache is empty: SDPA assumes
+        # top-left alignment, and a cached block queries a suffix of the key
+        # sequence. `VortexModel` hands over an explicit mask whenever that is
+        # the case and leaves it `None` for the prefill fast path.
+        is_causal = attention_mask is None and T > 1
+
+        try:
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                is_causal=is_causal,
+                scale=self.scale,
+                enable_gqa=self.n_groups > 1,
+            )
+        except TypeError:  # torch < 2.5 has no `enable_gqa`
+            if self.n_groups > 1:
+                k = k.repeat_interleave(self.n_groups, dim=1)
+                v = v.repeat_interleave(self.n_groups, dim=1)
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attention_mask,
+                dropout_p=0.0,
+                is_causal=is_causal,
+                scale=self.scale,
+            )
+
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.o_proj(out)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# MLP
+# ──────────────────────────────────────────────────────────────────────
+class VortexMLP(nn.Module):
+    """SwiGLU feed-forward: `down(silu(gate(x)) * up(x))`."""
+
+    def __init__(self, config: VortexConfig):
+        super().__init__()
+        intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = ACT2FN["silu"]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Block
+# ──────────────────────────────────────────────────────────────────────
+class VortexBlock(nn.Module):
+    """Pre-norm block: attention and MLP each add onto the residual stream."""
+
+    def __init__(self, config: VortexConfig, layer_idx: int = 0):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.attn = VortexAttention(config, layer_idx=layer_idx)
+        self.mlp = VortexMLP(config)
+        self.ln_attn = VortexRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.ln_mlp = VortexRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # GPT-2 style 1/sqrt(2L) branch scaling. Off by default: it is redundant
+        # next to zero-initialised residual outputs, which already make every
+        # block an exact identity at init.
+        self.resid_scale = (
+            1.0 / math.sqrt(2.0 * config.num_hidden_layers) if config.scale_residual else 1.0
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        position_offset: int = 0,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Cache] = None,
+    ) -> torch.Tensor:
+        x = x + self.resid_scale * self.attn(
+            self.ln_attn(x),
+            position_embeddings=position_embeddings,
+            position_offset=position_offset,
+            attention_mask=attention_mask,
+            past_key_value=past_key_value,
+        )
+        x = x + self.resid_scale * self.mlp(self.ln_mlp(x))
+        return x
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Base
+# ──────────────────────────────────────────────────────────────────────
+class VortexPreTrainedModel(PreTrainedModel):
+    """Weight init, tied-embedding bookkeeping and tokenizer plumbing."""
+
+    config_class = VortexConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["VortexBlock"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_sdpa = True
+    # SDPA already dispatches to FlashAttention-2 kernels on Ampere+, but the
+    # `attn_implementation="flash_attention_2"` HF interface is not implemented
+    # here. Claiming support would let `from_pretrained` pick a code path that
+    # does not exist for this architecture.
+    _supports_flash_attn = False
+    _supports_attention_backend = False
+    _supports_cache_class = True
+    _supports_static_cache = True
+    _can_record_outputs = {"hidden_states": VortexBlock, "attentions": VortexAttention}
+
+    def _init_weights(self, module: nn.Module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            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):
+            # A small vocab (16K) is far more tolerant than a 151K one, but
+            # scaling down keeps initial logits O(1) rather than O(10).
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, VortexRMSNorm):
+            nn.init.ones_(module.weight)
+        # Dispatched per-submodule by `PreTrainedModel.post_init`, which applies
+        # this over the whole tree. Overriding it here is what makes a freshly
+        # constructed model an identity passthrough without a separate traversal.
+        self._zero_init_residuals(module)
+
+    def _zero_init_residuals(self, module: Optional[nn.Module] = None) -> None:
+        """Zero `o_proj` and `down_proj` so every block starts as an identity.
+
+        With 18 stacked pre-norm blocks, default init compounds the residual
+        variance and saturates the stream before step 0. Zeroing the two branch
+        outputs makes the untrained network a clean passthrough, so the initial
+        loss is ln(vocab_size) = 9.7 rather than the hundreds default init gives.
+
+        Dispatched per-module by `PreTrainedModel.post_init` via
+        `_init_weights`; the recursion is over `self.modules()` so it also works
+        when called with no argument.
+        """
+        if not getattr(self.config, "zero_init_residual", True):
+            return
+        if module is not None:
+            if isinstance(module, VortexAttention):
+                nn.init.zeros_(module.o_proj.weight)
+            elif isinstance(module, VortexMLP):
+                nn.init.zeros_(module.down_proj.weight)
+            return
+        for block in self.model.layers:
+            nn.init.zeros_(block.attn.o_proj.weight)
+            nn.init.zeros_(block.mlp.down_proj.weight)
+
+    def resize_token_embeddings(
+        self,
+        new_num_tokens: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = None,
+        mean_resizing: bool = True,
+    ) -> nn.Embedding:
+        """Grow the embedding table, never shrink it.
+
+        Growing pads with fresh normal noise. Shrinking is refused rather than
+        silently truncating: rows that have been trained keep meaning something,
+        and a truncated table yields a model that evaluates fine and answers
+        with the wrong tokens.
+        """
+        old_embeddings = self.get_input_embeddings()
+        if old_embeddings is None:
+            raise ValueError("cannot resize embeddings on a model with no input embeddings")
+
+        old_num_tokens, embedding_dim = old_embeddings.weight.shape
+        if new_num_tokens is None:
+            new_num_tokens = old_num_tokens
+        if pad_to_multiple_of is not None:
+            new_num_tokens = math.ceil(new_num_tokens / pad_to_multiple_of) * pad_to_multiple_of
+        new_num_tokens = int(new_num_tokens)
+
+        if new_num_tokens < old_num_tokens:
+            raise ValueError(
+                f"cannot shrink token embeddings {old_num_tokens} -> {new_num_tokens}; "
+                f"the vocabulary must only be extended"
+            )
+        if new_num_tokens == old_num_tokens:
+            return old_embeddings
+
+        new_embeddings = nn.Embedding(
+            new_num_tokens, embedding_dim, device=old_embeddings.weight.device
+        )
+        with torch.no_grad():
+            new_embeddings.weight.normal_(mean=0.0, std=self.config.initializer_range)
+            new_embeddings.weight[:old_num_tokens].copy_(old_embeddings.weight)
+        self.set_input_embeddings(new_embeddings)
+        self.config.vocab_size = new_num_tokens
+
+        # Keep the head in step with a tied table.
+        if self.config.tie_word_embeddings and self.get_output_embeddings() is not None:
+            self.tie_weights()
+        return new_embeddings
+
+    def tie_weights(self, recompute_mapping: bool = False, missing_keys=None):
+        """Alias the `lm_head` weight onto the embedding table.
+
+        Overridden rather than inherited because `missing_keys` has two
+        incompatible shapes across transformers versions: a `set` in 4.x and a
+        mapping in 4.56+. Only `lm_head` is tied here, so it is dropped from the
+        "missing" report under either shape.
+        """
+        if getattr(self.config, "tie_word_embeddings", False):
+            output_embeddings = self.get_output_embeddings()
+            input_embeddings = self.get_input_embeddings()
+            if output_embeddings is not None and input_embeddings is not None:
+                output_embeddings.weight = input_embeddings.weight
+        if missing_keys is None:
+            return
+        discard = getattr(missing_keys, "discard", None)
+        if callable(discard):
+            discard("lm_head.weight")
+            return
+        if hasattr(missing_keys, "pop"):
+            try:
+                missing_keys.pop("lm_head.weight")
+            except TypeError:  # mapping-style pop(key, default)
+                missing_keys.pop("lm_head.weight", None)
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Model
+# ──────────────────────────────────────────────────────────────────────
+class VortexModel(VortexPreTrainedModel):
+    """Embedding + decoder stack + final norm."""
+
+    def __init__(self, config: VortexConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [VortexBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = VortexRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = VortexRotaryEmbedding(config)
+
+        # Set by `PreTrainedModel.gradient_checkpointing_enable`, which targets
+        # any submodule carrying this attribute.
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed_tokens = value
+
+    # ── attention mask ───────────────────────────────────────────────
+    @staticmethod
+    def _build_causal_mask(
+        q_len: int,
+        kv_len: int,
+        attention_mask_2d: Optional[torch.Tensor],
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Bottom-right-aligned boolean mask, `True` = attend.
+
+        Two things SDPA's `is_causal=True` cannot express:
+
+        1. **Alignment.** With `past_len` cached keys, query `i` sits at absolute
+           position `past_len + i`, so it may attend to keys `0 .. past_len + i`.
+           Top-left alignment would bar the cached keys from every query.
+        2. **Padding.** Left-padded batches need the pad columns removed.
+
+        The self-diagonal is force-enabled on top of the mask so no query row is
+        ever fully masked. A fully-masked row makes softmax return `NaN`, and
+        those `NaN`s then ride in the padded key/value vectors into the next
+        layer, where a `0 * NaN` in the weighted sum spreads them. Letting a
+        padded query attend to itself is harmless — that position is masked out
+        for every other query, so it cannot leak.
+        """
+        key_positions = torch.arange(kv_len, device=device)
+        query_positions = torch.arange(q_len, device=device) + (kv_len - q_len)
+        mask = (key_positions[None, :] <= query_positions[:, None])[None, None, :, :]
+
+        if attention_mask_2d is not None:
+            padding = attention_mask_2d.to(device=device)[:, None, None, :].bool()
+            mask = mask & padding
+
+        self_attends = (key_positions[None, :] == query_positions[:, None])[None, None, :, :]
+        return (mask | self_attends).expand(-1, 1, -1, -1).contiguous()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = bool(output_attentions)
+        output_hidden_states = bool(output_hidden_states)
+        return_dict = True if return_dict is None else bool(return_dict)
+
+        checkpointing = bool(getattr(self, "gradient_checkpointing", False)) and self.training
+        use_cache = self.config.use_cache if use_cache is None else bool(use_cache)
+
+        # Checkpointing recomputes each block during the backward pass, and
+        # `Cache.update` mutates in place -- so a cached forward would append the
+        # same keys a second time and corrupt every downstream layer's mask
+        # (observed: the cache silently doubling from 24 to 48 entries).
+        # Training never needs the cache anyway, so it is dropped here. Inference
+        # is unaffected because `self.training` is False.
+        if checkpointing:
+            use_cache = False
+            past_key_values = None
+
+        if output_attentions:
+            raise NotImplementedError(
+                "`output_attentions=True` is not supported: each block returns only "
+                "hidden states, because attention runs fused inside SDPA."
+            )
+
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("provide exactly one of `input_ids` or `inputs_embeds`")
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if past_key_values is None and use_cache:
+            past_key_values = DynamicCache(config=self.config)
+
+        batch_size, seq_len, _ = inputs_embeds.shape
+        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
+        kv_len = past_len + seq_len
+
+        if position_ids is None:
+            # RoPE is relative, so shifting every position by the same constant
+            # leaves every attention score unchanged. Deriving absolute
+            # positions from the cache length is therefore correct even for the
+            # left-padded batches `generate` builds.
+            position_ids = torch.arange(past_len, past_len + seq_len, device=inputs_embeds.device)
+            position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        elif position_ids.shape[-1] == kv_len and past_len > 0:
+            position_ids = position_ids[:, past_len:]
+
+        # A 2D `(batch, kv_len)` padding mask is what `generate` passes; a 4D
+        # mask is taken as already built. Anything else is ignored rather than
+        # guessed at.
+        padding_mask_2d = None
+        if attention_mask is not None and attention_mask.dim() == 2:
+            padding_mask_2d = attention_mask
+            attention_mask = None
+
+        if attention_mask is None and (past_len > 0 or padding_mask_2d is not None):
+            attention_mask = self._build_causal_mask(
+                q_len=seq_len,
+                kv_len=kv_len,
+                attention_mask_2d=padding_mask_2d,
+                device=inputs_embeds.device,
+            )
+
+        # One RoPE table for the whole stack rather than one per layer.
+        position_embeddings = self.rotary_emb(inputs_embeds, kv_len)
+
+        hidden_states = inputs_embeds
+        all_hidden_states = () if output_hidden_states else None
+
+        checkpoint_fn = getattr(self, "_gradient_checkpointing_func", None)
+        if checkpointing and checkpoint_fn is None:
+            checkpoint_fn = lambda fn, *args: checkpoint(fn, *args, use_reentrant=False)
+
+        for block in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            if checkpointing:
+                hidden_states = checkpoint_fn(
+                    block,
+                    hidden_states,
+                    position_embeddings,
+                    past_len,
+                    attention_mask,
+                    past_key_values,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    position_embeddings=position_embeddings,
+                    position_offset=past_len,
+                    attention_mask=attention_mask,
+                    past_key_value=past_key_values,
+                )
+
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return (hidden_states, past_key_values if use_cache else None, all_hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )
+
+
+# ──────────────────────────────────────────────────────────────────────
+# Causal LM
+# ──────────────────────────────────────────────────────────────────────
+class VortexForCausalLM(VortexPreTrainedModel, GenerationMixin):
+    """Vortex with a tied language-modelling head.
+
+    `VortexModel` is the base model (`base_model_prefix = "model"`), so
+    `save_pretrained` writes `model.embed_tokens.weight`, `model.layers.N.*` and
+    `model.norm.weight` — the same key layout as the training checkpoints, which
+    is what lets this class load them unchanged.
+
+    `GenerationMixin` is inherited explicitly. From transformers v4.50 onward
+    `PreTrainedModel` no longer provides it, so without this second base the
+    model silently loses `generate`, `generate_from_model` and sampling helpers.
+    It must come *after* `PreTrainedModel` in the MRO.
+    """
+
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head.weight": "model.embed_tokens.weight"}
+    _pp_plan = {"embed_tokens": ["model.embed_tokens"], "layers": ["model.layers"]}
+
+    def __init__(self, config: VortexConfig):
+        super().__init__(config)
+        self.model = VortexModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Without this a directly-constructed model keeps PyTorch's default init
+        # — no `initializer_range`, no zeroed residual outputs. `from_pretrained`
+        # calls it too, but construction has to be self-sufficient or
+        # `VortexForCausalLM(config).to(device)` silently trains a broken model.
+        self.post_init()
+        if config.tie_word_embeddings:
+            self.tie_weights()
+
+    def post_init(self) -> None:
+        """Initialise weights, then apply the zero-init residual scheme.
+
+        `PreTrainedModel.post_init` is what registers `all_tied_weights_keys`,
+        parallel plans and device-map hints, and it is also what drives
+        `_init_weights` over every submodule. It must be delegated to rather than
+        shadowed, but on its own it leaves `o_proj` and `down_proj` at their
+        normal init, so the second pass below is what actually makes an untrained
+        model a passthrough.
+        """
+        super().post_init()
+        self._zero_init_residuals()
+        if getattr(self.config, "tie_word_embeddings", False):
+            self.tie_weights()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.lm_head = new_embeddings
+
+    def get_decoder(self) -> VortexModel:
+        return self.model
+
+    # ── loss ─────────────────────────────────────────────────────────
+    def _chunked_cross_entropy(
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor,
+        chunk_size: int,
+        num_items_in_batch: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Cross-entropy without materialising `(N, vocab_size)` logits.
+
+        This is the single largest avoidable memory term in a training step: at
+        batch 32 x 2048 tokens x 16384 vocab in fp32 the logits alone are 4.3GB.
+        Accumulating in time-chunks holds the peak at `chunk_size` rows instead.
+        """
+        n_valid = (labels != -100).sum()
+        if n_valid.item() == 0:
+            return hidden_states.sum() * 0.0  # keep the graph connected
+
+        total = hidden_states.new_zeros((), dtype=torch.float32)
+        for start in range(0, hidden_states.shape[0], chunk_size):
+            logits = self.lm_head(hidden_states[start : start + chunk_size]).float()
+            total = total + F.cross_entropy(
+                logits,
+                labels[start : start + chunk_size],
+                ignore_index=-100,
+                reduction="sum",
+            )
+            del logits
+
+        if num_items_in_batch is not None:
+            # `Trainer` normalises by a token count accumulated across
+            # gradient-accumulation steps. Matching it is what keeps the loss it
+            # reports comparable to the standalone training loop's.
+            return total / num_items_in_batch.to(total.device)
+        return total / n_valid.clamp_min(1).float()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        chunk_size: int = 0,
+        num_items_in_batch: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""Causal language modelling.
+
+        Args:
+            labels (`torch.LongTensor`, *optional*):
+                Targets for next-token prediction. When given, `logits` comes back
+                `None` unless `logits_to_keep` asks for it — the loss accumulates
+                in chunks precisely so the full `(batch, seq, vocab)` tensor never
+                has to exist.
+            logits_to_keep (`int`, *optional*, defaults to 0):
+                Return logits for only the last `n` positions. 0 means all of
+                them when no loss is being computed, and none when one is.
+                `transformers` sets this to 1 during `generate`; passing any value
+                alongside `labels` is how to ask for a loss *and* logits.
+            chunk_size (`int`, *optional*, defaults to 0):
+                Rows per cross-entropy chunk. 0 selects 1024.
+
+        Returns:
+            [`CausalLMOutputWithPast`]: `logits`, `loss`, and `past_key_values`
+            when `use_cache` is set.
+        """
+        return_dict = True if return_dict is None else bool(return_dict)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        past_key_values = outputs.past_key_values
+
+        # Keep only the tail when asked. During generation this is the single new
+        # position, so the vocab-sized projection runs on one row instead of the
+        # whole sequence.
+        keep = int(logits_to_keep.item()) if isinstance(logits_to_keep, torch.Tensor) else int(logits_to_keep)
+
+        loss = None
+        if labels is not None:
+            # Next-token alignment: predict token t+1 from position t.
+            shift_hidden = hidden_states[..., :-1, :].reshape(-1, hidden_states.shape[-1])
+            shift_labels = labels[..., 1:].reshape(-1)
+            loss = self._chunked_cross_entropy(
+                shift_hidden, shift_labels, chunk_size or 1024, num_items_in_batch
+            )
+
+        if labels is not None and keep == 0:
+            # Keep the memory win. Ask with `logits_to_keep=1` if you need logits
+            # alongside a loss.
+            logits = None
+        else:
+            tail = hidden_states[:, -keep:, :] if keep > 0 else hidden_states
+            logits = self.lm_head(tail)
+
+        if not return_dict:
+            return (logits, loss) if loss is None else (logits, loss, past_key_values)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    # ── generation ───────────────────────────────────────────────────
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs,
+    ) -> dict:
+        """Trim model inputs to the block `generate` is about to run.
+
+        Handled entirely by `GenerationMixin`: it slices `input_ids` down to the
+        tokens not yet in the cache. That slice is load-bearing rather than an
+        optimisation — resending the full prefix would recompute it and corrupt
+        the cache. Overridden only to keep the signature aligned with
+        `transformers` 5.x and to forward `position_ids`, which the base
+        implementation pops and re-slices.
+        """
+        return super().prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+    # ── gradient checkpointing ───────────────────────────────────────
+    def gradient_checkpointing_enable(
+        self,
+        gradient_checkpointing_kwargs: Optional[dict] = None,
+        **kwargs,
+    ) -> None:
+        """Recompute decoder activations in the backward pass instead of storing them.
+
+        Trades roughly 20-30% step time for most of the activation memory, which
+        is what lets one 40GB card hold a large batch at 2K context. Only active
+        in training mode — `VortexModel.forward` gates on `self.training`.
+        """
+        super().gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs=gradient_checkpointing_kwargs, **kwargs
+        )
+        self.model.gradient_checkpointing = True
+        if not hasattr(self.model, "_gradient_checkpointing_func"):
+            self.model._gradient_checkpointing_func = lambda fn, *args: checkpoint(
+                fn, *args, use_reentrant=False
+            )
+
+    def gradient_checkpointing_disable(self, **kwargs) -> None:
+        super().gradient_checkpointing_disable(**kwargs)
+        self.model.gradient_checkpointing = False
+
+
+__all__ = [
+    "VortexConfig",
+    "VortexPreTrainedModel",
+    "VortexModel",
+    "VortexForCausalLM",
+    "VortexBlock",
+    "VortexAttention",
+    "VortexMLP",
+    "VortexRMSNorm",
+    "VortexRotaryEmbedding",
+    "build_rope_cache",
+    "apply_rope",
+]