Jet-Long dynamic bifocal RoPE: 4K->10K zero-shot context extension (arXiv:2607.07740)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Escarda-LLMs.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,313 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- text-generation
+- small-models
+- base-model
+- mla
+- jepa
+- experimental
+pipeline_tag: text-generation
+library_name: transformers
+datasets:
+- HuggingFaceFW/fineweb-edu
+---
+
+<img src="Escarda-LLMs.png" width="50%">
+
+
+# Escarda-86M-Base
+
+## 🔭 Jet-Long context extension (native 4K → 10K)
+
+This is the **Jet-Long** edition of [`Escarda-86M-Base`](https://huggingface.co/Quazim0t0/Escarda-86M-Base).
+It extends the usable context from the native **4,096**-token training window to
+**10,240** tokens with **no fine-tuning** and **no change to short-context behaviour**,
+by adding *dynamic bifocal RoPE* from **Jet-Long** ([arXiv:2607.07740](https://arxiv.org/abs/2607.07740), NVIDIA).
+
+### What was applied
+
+Jet-Long pairs a **local window** (`w0 = 2048`, classic RoPE) with a **remote window**
+whose position map aliases far-apart tokens back onto the pretrained rotation grid:
+
+```
+f(x) = floor(x / G),   G = max(1, ceil(L / 4096))
+```
+
+`G` adapts to the **current** sequence length `L`, so:
+
+- **`L ≤ 4096` → `G = 1` → f is the identity → the model is bit-for-bit the base model.**
+  (Verified: max |Δlogit| between Jet-Long on/off within the window is `0.000e+00`.)
+- **`L > 4096`** → the remote window keeps every rotation in-distribution, so the model
+  extrapolates instead of collapsing.
+
+Implementation notes specific to this SpikeWhaleLM build:
+- Only the **decoupled RoPE partition** (16 of 64 head dims) is aliased; the NoPE partition
+  is untouched. Softmax attention (`use_derf=False`) — the standard Jet-Long merge applies.
+- The remote view is realized by an **on-the-fly correction rotation** on the already-RoPE'd
+  KV cache (RoPE composes additively), so **the cache is never rewritten** and decode is cheap.
+- Enabled via config: `use_jetlong=true`, `jetlong_w0=2048`, `jetlong_w_pretrained=4096`,
+  `max_position_embeddings=10240`. Set `use_jetlong=false` to recover the exact base model.
+- The inclusion–exclusion / CuTe throughput kernel from the paper is **not** included (it targets
+  100K+ contexts on H100); at 86M params the bifocal attention is computed directly.
+
+### Measured (PG-19-style perplexity on held-out text, lower is better)
+
+| Context length | Base model | This Jet-Long model |
+|---|---|---|
+| ≤ 4,096 (in-window) | *(identical — Jet-Long is a no-op)* | *(identical)* |
+| **10,240** | **59.68** | **16.04** |
+
+Beyond the training window the base model's perplexity blows up while Jet-Long stays flat —
+and long-context generation stays grammatical where the base model degrades into word-salad.
+
+### Usage
+
+Jet-Long is **on by default** in this repo. Pass explicit `position_ids` so RoPE gets true
+absolute positions during cached decode:
+
+```python
+import torch
+from transformers import AutoModelForCausalLM
+m = AutoModelForCausalLM.from_pretrained("Quazim0t0/Escarda-86M-Base-JL", trust_remote_code=True)
+ids = ...              # up to ~10,240 tokens
+pos = torch.arange(ids.shape[1]).unsqueeze(0)
+out = m(input_ids=ids, position_ids=pos, use_cache=True)   # prefill, then decode step-by-step
+```
+
+*Method: Tang, Wang, Gu, Han, Cai — “Jet-Long: Efficient Long-Context Extension with Dynamic
+Bifocal RoPE”, arXiv:2607.07740. Applied here zero-shot to SpikeWhaleLM; no weights were retrained.*
+
+
+**Escarda-86M-Base** is a ~86M-parameter, from-scratch decoder-only language model — the
+**base** sibling of [Quazim0t0/Escarda-86M](https://huggingface.co/Quazim0t0/Escarda-86M)
+(the chat-tuned model). It shares the same `SpikeWhaleLM` architecture (Multi-head Latent
+Attention, an n-gram "engram" memory, hash-lookup layers, hyper-connections, an HRM
+refinement step, and JEPA / multi-token-prediction training objectives) and the same
+custom ChatML-aware tokenizer.
+
+This checkpoint is a JEPA-distilled base. It is best used as a **starting point for
+continued pretraining / fine-tuning** rather than as a chat assistant.
+
+> **Related models:** SFT / chat model → [Quazim0t0/Escarda-86M](https://huggingface.co/Quazim0t0/Escarda-86M)
+> · live demo → [Escarda-86M-Chat Space](https://huggingface.co/spaces/Quazim0t0/Escarda-86M-Chat)
+
+Trained using **Modal's credits** during the **Small Models, Big Adventures Hackathon**.
+
+---
+
+## Model summary
+
+| | |
+|---|---|
+| **Parameters** | ~85.7M (`tie_word_embeddings=True`) |
+| **Type** | Decoder-only LM (`SpikeWhaleLM`, `model_type: spike_whale`) |
+| **Hidden size / layers** | 640 / 16 |
+| **Attention** | 10 heads (`head_dim=64`), 1 KV head (MQA), MLA low-rank Q/O, decoupled RoPE(16)+NoPE(48), QK-norm |
+| **Context length** | 4096 tokens |
+| **Vocab** | 16,512 (custom `length-max` tokenizer) |
+| **License** | Apache-2.0 |
+
+For the full architecture description see the
+[chat model's card](https://huggingface.co/Quazim0t0/Escarda-86M#architecture).
+
+---
+
+<!-- ARCH_TOK_START -->
+## Architecture
+
+These models are built on **SpikeWhaleLM**, a custom ~86M-parameter decoder-only transformer
+(16 layers, hidden size 640, 4096-token context, 16,512 vocab, tied input/output embeddings).
+It combines several non-standard components:
+
+- **Multi-head Latent Attention (MLA + XSA)** — queries and the output projection are
+  LoRA-compressed (rank 128); each head splits into a decoupled RoPE part (dim 16) and a
+  position-agnostic NoPE part (dim 48); 10 query heads share a **single KV head**
+  (multi-query attention), with QK-norm for stable logits.
+- **Engram n-gram memory** — a gated associative memory that hashes local n-grams (up to
+  trigrams) into a learned 4,096-entry table and mixes the result back into the residual stream.
+- **Hash-lookup layers (×2)** — multi-head content-addressable features alongside the token
+  embeddings.
+- **Hyper-Connections** — learned, width-expanded residual connections mixed via
+  Sinkhorn-normalized routing, in place of the plain residual add.
+- **HRM refinement** — a Hierarchical Reasoning Model block that performs an extra latent
+  "think a bit more" refinement pass over the hidden states before the output head.
+- **Multi-Token Prediction (MTP)** — a DeepSeek-V3-style auxiliary training head predicting
+  more than one next token (no inference cost).
+- Feed-forward is **dense** (the block is MoE-capable, but MoE is disabled in this release).
+
+> **JEPA vs HRM.** The **Escarda** models are trained with **both HRM refinement and a JEPA (Joint-Embedding Predictive Architecture) auxiliary objective** (`use_hrm_refine=True`, `use_jepa=True`) — the JEPA term predicts future latent states during training to shape the model's representations. The sibling **Byrne** models drop JEPA and use **HRM refinement only**.
+
+<a href="https://hfviewer.com/Quazim0t0/Escarda-86M-Base?utm_source=huggingface&amp;utm_medium=embedded_model_card&amp;utm_campaign=Quazim0t0__Escarda-86M-Base_card&amp;utm_content=embedded_card_open_viewer&amp;from=embedded-model-card" target="_blank" rel="noopener">
+  <img
+    src="https://hfviewer.com/api/card.svg?source=Quazim0t0%2FEscarda-86M-Base&amp;granularity=auto&amp;v=20260516-title-pills-card"
+    alt="Architecture graph for Quazim0t0/Escarda-86M-Base. Open in hfviewer"
+    width="100%"
+  />
+</a>
+
+## Tokenizer
+
+These models use **`SpikeTokenizer`**, a custom **byte-level "length-max" (greedy
+longest-match)** tokenizer with a **16,512-token vocabulary** — not a standard BPE/HF
+tokenizer. Text is UTF-8 encoded, each byte mapped to a latin-1 character, then greedily
+matched against the vocab using the longest key that fits at each position. It is
+**ChatML-aware**, with atomic special tokens for framing and reasoning/tool markers
+(`<|im_start|>`, `<|im_end|>`, `<think>`/`</think>`, `<begin_solution>`/`<end_solution>`,
+tool-call markers) plus `<bos>`/`<eos>`/`<pad>`/`<unk>`. It ships as a `PreTrainedTokenizer`
+subclass (`spike_tokenizer.py`) and loads via
+`AutoTokenizer.from_pretrained(..., trust_remote_code=True)`.
+<!-- ARCH_TOK_END -->
+
+## Evaluation
+
+splits. byte_ppl is `exp(sum_NLL_nats / total_UTF8_bytes)` on WikiText-2 test (tokenizer-
+independent). BLiMP is fraction of minimal pairs with `logprob(good) > logprob(bad)`
+(12 paradigms × 150). Stderr is binomial `sqrt(p(1-p)/n)`.
+
+### Language modeling
+
+| Metric | Value |
+|---|---|
+| WikiText-2 byte_ppl ↓ | **2.2228** |
+| BLiMP acc ↑ | 0.7144 |
+
+### Multiple-choice suite
+
+| Task | acc | ± | acc_norm | ± |
+|---|---|---|---|---|
+| arc_easy | 0.3801 | 0.0100 | 0.3615 | 0.0099 |
+| arc_challenge | 0.1886 | 0.0114 | 0.2235 | 0.0122 |
+| hellaswag | 0.2759 | 0.0045 | 0.2832 | 0.0045 |
+| winogrande | 0.5162 | 0.0140 | — | — |
+| piqa | 0.5843 | 0.0115 | 0.5631 | 0.0116 |
+| openbookqa | 0.1300 | 0.0150 | 0.2500 | 0.0194 |
+| boolq | 0.5138 | 0.0087 | — | — |
+
+### ArithMark-2.0 ([AxiomicLabs](https://huggingface.co/datasets/AxiomicLabs/ArithMark-2.0))
+
+| Metric | Value |
+|---|---|
+| acc | 0.2536 ± 0.0087 |
+| acc_norm | 0.2348 ± 0.0085 |
+
+n = 2,500 · chance = 0.25.
+
+> **Note:** as a distilled base, this checkpoint has the **lowest byte-perplexity** of the
+> Escarda family but trades off downstream task accuracy — a good reminder that perplexity
+> alone is not a reliable capability ranking. For the strongest chat behaviour use
+> [Escarda-86M](https://huggingface.co/Quazim0t0/Escarda-86M); use **this** model when you
+> want a low-loss base to continue pretraining or fine-tune.
+
+---
+
+## Training & token budget
+
+- **Tokens:** ~20B (from-scratch pretraining of the SpikeWhale base, ~28k steps); this
+  checkpoint is a JEPA-distilled snapshot of that base.
+- **Token/param ratio:** ~233 tokens/param (20B / 85.7M) — roughly **11–12× the Chinchilla
+  ~20-tokens/param compute-optimal heuristic**, i.e. a deliberately **over-trained small
+  model** (the inference-efficient trade-off).
+
+Fitting the Chinchilla data term to this model's own pretraining loss curve gives:
+
+`L(D) ≈ 2.611 + 77,715 · D^(−0.537)`  (nats/token, R² = 0.92)
+
+From that fit:
+- **Compute-optimal tokens for this 86M size ≈ 4.3B** → the 20B run is **~4.6× past
+  compute-optimal**.
+- **Diminishing-returns knee ≈ 22.5B** tokens (where +1B tokens buys < 0.005 nats) — the
+  20B stopping point lands **right at the knee**, a well-judged budget.
+- The model is **parameter-bound, not data-bound** at 20B: the capacity term (~0.82 nats)
+  exceeds the data term (~0.54), so extra tokens help little. Doubling to 40B is projected
+  to lower loss only ~0.07 nats (~7% perplexity) with negligible downstream gain — the lever
+  for better quality is **more parameters, not more tokens**. (This is also why, as a
+  distilled base, it reaches the lowest perplexity of the family without the best downstream
+  scores — it is already at its data-term floor.)
+
+*Caveats: single-size fit (folds irreducible loss + capacity floor into one constant); the
+cosine-LR decay inflates the fitted exponent, so treat β as an upper bound; token counts are
+anchored to the ~20B figure and scale linearly if that differs.*
+
+---
+
+## Usage
+
+Custom architecture — load with `trust_remote_code=True` (the modeling code ships in this
+repo via `auto_map`):
+
+```python
+from transformers import AutoModelForCausalLM
+model = AutoModelForCausalLM.from_pretrained(
+    "Quazim0t0/Escarda-86M-Base", trust_remote_code=True)
+```
+
+The tokenizer is the custom `SpikeTokenizer` (`tokenizer.json`, `algorithm: length-max`);
+load it with the `spike_tokenizer.py` helper from the project rather than `AutoTokenizer`.
+
+## Acknowledgements
+
+Built with **Modal** credits during the **Small Models, Big Adventures Hackathon**, and
+released to the community as a base to build on.
+
+<!-- CITE_START -->
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{escarda86mbase,
+  title        = {Escarda-86M-Base: A ~86M-parameter SpikeWhaleLM},
+  author       = {Dean Byrne (Quazim0t0)},
+  year         = {2026},
+  howpublished = {HuggingFace, \url{https://huggingface.co/Quazim0t0/Escarda-86M-Base}},
+  note         = {Quazim0t0/Escarda-86M-Base}
+}
+```
+<!-- CITE_END -->
+
+
+## Escarda vs Byrne — vision family comparison
+
+The **Byrne** family uses HRM refinement. **Escarda** = Byrne **+ JEPA** (Joint-Embedding
+Predictive head) added *alongside* HRM in both the vision encoder and the LM trunk —
+auxiliary only, **zero inference cost**.
+
+**Vision encoder (DINOv2 teacher-alignment, n=1024 held-out):**
+
+| | Byrne-VE | Escarda-VE |
+|---|---|---|
+| Params | 39.34M | 39.60M (+JEPA head) |
+| CLS cosine | **0.776** | 0.771 |
+| PATCH cosine | **0.600** | 0.584 |
+| JEPA self-consistency | — | 0.040 |
+
+**Docling (same held-out doc images, atomic DocTags):** both emit well-formed DocTags;
+Byrne-Docling is marginally more complete on the hardest samples (closes `</formula>`,
+includes the `<code>` wrapper), consistent with its slightly higher teacher-alignment.
+Escarda-Docling is structurally on par and adds the JEPA representation-learning trait.
+
+**Pros/cons.** *Byrne (HRM):* higher teacher-alignment, all capacity on distillation
+fidelity; no self-supervised objective. *Escarda (HRM+JEPA):* self-supervised
+neighbour-prediction (richer spatial structure) at zero inference cost, trading ~1–3%
+teacher-alignment. Same size class.
+
+Family repos: [Byrne-VE](https://huggingface.co/Quazim0t0/Byrne-VE) ·
+[Escarda-VE](https://huggingface.co/Quazim0t0/Escarda-VE) ·
+[Byrne-Docling-131M](https://huggingface.co/Quazim0t0/Byrne-Docling-131M) ·
+[Escarda-Docling-126M](https://huggingface.co/Quazim0t0/Escarda-Docling-126M)
+
+## Update: engram repair (behavior-preserving)
+
+The n-gram Engram memory in the original weights was degenerate: with the frozen
+LSH compressor at init scale, every token hashed to bucket 0, so only one table
+row ever received gradient. This revision rescales the (frozen) compressor and
+broadcasts the learned bucket-0 vector across all table rows.
+
+**Outputs are bit-identical to the previous revision** (verified: max logit
+difference 0.0 across a prompt battery). The only change: the Engram's hash now
+spreads across the full table and every bucket is independently trainable — so
+if you distill or SFT on top of this base, the n-gram memory will actually learn
+instead of staying a constant bias.

config.json

@@ -0,0 +1,70 @@
+{
+  "architectures": [
+    "SpikeWhaleLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.SpikeWhaleConfig",
+    "AutoModel": "model_v2.SpikeWhaleLM",
+    "AutoModelForCausalLM": "model_v2.SpikeWhaleLM"
+  },
+  "attention_dropout": 0.0,
+  "bos_token_id": 2,
+  "dtype": "float32",
+  "engram_compress_dim": 32,
+  "engram_gate_init_bias": -1.0,
+  "engram_max_ngram": 3,
+  "engram_num_heads": 2,
+  "engram_table_size": 4096,
+  "eos_token_id": 3,
+  "hc_eps": 1e-06,
+  "hc_mult": 2,
+  "hc_sinkhorn_iters": 20,
+  "head_dim": 64,
+  "hidden_dropout": 0.0,
+  "hidden_size": 640,
+  "hrm_refine_dim": 128,
+  "hrm_refine_steps": 1,
+  "initializer_range": 0.02,
+  "jepa_horizon": 1,
+  "jepa_loss_weight": 0.1,
+  "jepa_pred_dim": 256,
+  "max_position_embeddings": 10240,
+  "model_type": "spike_whale",
+  "moe_aux_loss_coef": 0.01,
+  "moe_intermediate_size": 2000,
+  "moe_layers": [],
+  "mtp_loss_weight": 0.3,
+  "n_routed_experts": 6,
+  "n_shared_experts": 1,
+  "nope_head_dim": 48,
+  "norm_topk_prob": true,
+  "num_attention_heads": 10,
+  "num_experts_per_tok": 2,
+  "num_hash_layers": 2,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 1,
+  "num_nextn_predict_layers": 1,
+  "o_lora_rank": 128,
+  "q_lora_rank": 128,
+  "qk_rope_head_dim": 16,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000.0,
+  "routed_scaling_factor": 1.0,
+  "scoring_func": "sqrtsoftplus",
+  "tie_word_embeddings": true,
+  "transformers_version": "5.8.0",
+  "use_derf": false,
+  "use_engram": true,
+  "use_hrm_refine": true,
+  "use_hyper_connections": true,
+  "use_jepa": true,
+  "use_moe": false,
+  "use_qk_norm": true,
+  "use_value_embed": false,
+  "use_xsa": true,
+  "vocab_size": 16512,
+  "zloss_coef": 0.0001,
+  "use_jetlong": true,
+  "jetlong_w_pretrained": 4096,
+  "jetlong_w0": 2048
+}

config.py

@@ -0,0 +1,164 @@
+"""
+config.py -- SpikeWhale: combined config from SpikeTransformer (My Project) + NanoWhale (DeepSeek-V4).
+
+Features carried from My Project (not in NanoWhale):
+  - DERF attention: erf(alpha*score+bias)*gamma replaces softmax
+  - XSA (Exclusive Self-Attention): orthogonality correction removes self-echo from attn output
+  - Engram N-gram module: hash-table N-gram lookup with DERF gate injected into embeddings
+  - Three-tier optimizer: embed/table params trained at lower LR
+
+Features carried from NanoWhale (not in My Project):
+  - MLA (Multi-Head Latent Attention): low-rank Q projection + direct K,V (MQA)
+  - Partial RoPE: rotary embeddings on only qk_rope_head_dim dims of Q and K
+  - Low-rank grouped output projection (o_lora_rank)
+  - Hyper-Connections: hc_mult residual streams with learned routing between layers
+  - Shared expert in MoE (always-active expert alongside routed experts)
+  - sqrtsoftplus expert scoring (vs softmax in My Project)
+  - Hash-based routing for first num_hash_layers layers
+  - norm_topk_prob + routed_scaling_factor
+  - Multi-Token Prediction (MTP): extra heads predict k steps ahead
+  - torch.compile, FineWeb-Edu streaming, Trackio, YAML configs in train.py
+"""
+
+from transformers import PretrainedConfig
+
+
+class SpikeWhaleConfig(PretrainedConfig):
+    model_type = "spike_whale"
+
+    def __init__(
+        self,
+        # Standard
+        vocab_size: int = 16512,   # SpikeTokenizer: 16384 base + 128 padded special slots
+        hidden_size: int = 2048,
+        num_hidden_layers: int = 11,
+        max_position_embeddings: int = 4096,
+        rms_norm_eps: float = 1e-6,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = False,
+        hidden_dropout: float = 0.0,
+        bos_token_id: int = 0,
+        eos_token_id: int = 1,
+        # MLA Attention (NanoWhale)
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 1,   # 1 = MQA; >1 = GQA
+        q_lora_rank: int = 160,         # low-rank Q: hidden -> q_lora_rank -> num_heads*head_dim
+        head_dim: int = 96,             # total per-head dim = nope_head_dim + qk_rope_head_dim
+        qk_rope_head_dim: int = 32,     # RoPE applied only to these dims
+        o_lora_rank: int = 80,          # low-rank output: num_heads*head_dim -> o_lora_rank -> hidden
+        attention_dropout: float = 0.0,
+        rope_theta: float = 10000.0,
+        # DERF + XSA (My Project)
+        use_derf: bool = True,
+        use_xsa: bool = True,
+        # MoE (combined)
+        use_moe: bool = True,
+        moe_intermediate_size: int = 640,
+        n_routed_experts: int = 4,
+        n_shared_experts: int = 1,      # NanoWhale: always-active shared expert
+        num_experts_per_tok: int = 2,
+        norm_topk_prob: bool = True,    # NanoWhale: normalize top-k routing weights
+        scoring_func: str = "sqrtsoftplus",  # NanoWhale: sqrt(softplus(x)) vs softmax
+        routed_scaling_factor: float = 1.0,  # NanoWhale: scale routed expert weights
+        num_hash_layers: int = 2,       # NanoWhale: first N layers use hash routing
+        moe_aux_loss_coef: float = 0.01,
+        moe_layers: list = None,
+        # Hyper-Connections (NanoWhale)
+        use_hyper_connections: bool = True,
+        hc_mult: int = 4,               # number of parallel residual streams
+        hc_sinkhorn_iters: int = 20,
+        hc_eps: float = 1e-6,
+        # Multi-Token Prediction (NanoWhale)
+        num_nextn_predict_layers: int = 1,  # extra MTP heads (0 = disabled)
+        # Engram N-gram module (My Project)
+        use_engram: bool = True,
+        engram_compress_dim: int = 64,
+        engram_num_heads: int = 4,
+        engram_table_size: int = 8192,
+        engram_max_ngram: int = 3,
+        engram_gate_init_bias: float = -4.0,
+        # HRM-inspired iterative refinement (EXPERIMENTAL; off by default).
+        # Adds one small block that refines the final hidden state over N inner
+        # steps before the output norm. This is the "iterative refinement" part
+        # that the ARC-Prize ablation found carried most of HRM's benefit -- NOT
+        # the full two-timescale H/L hierarchy. Honestly labeled HRM-inspired.
+        use_hrm_refine: bool = False,
+        hrm_refine_steps: int = 3,       # inner refinement iterations
+        hrm_refine_dim: int = 256,       # bottleneck width of the refine MLP
+        # --- v2 additions ---
+        use_qk_norm: bool = True,        # per-head RMSNorm on Q,K before RoPE
+        zloss_coef: float = 1e-4,        # log^2(Z) penalty on lm_head logits (0=off)
+        mtp_loss_weight: float = 0.3,    # down-weight for MTP CE loss
+        use_value_embed: bool = False,   # per-layer value-embedding residual (zero-init)
+        # --- JEPA secondary prediction head (jepa_v2) ---
+        # Predicts the trunk's FUTURE hidden state in representation space
+        # (stop-gradient target), complementing MTP which predicts future
+        # TOKENS. Same bottleneck-MLP shape as the HRM refinement block.
+        use_jepa: bool = True,
+        jepa_horizon: int = 1,           # predict hidden state k=1..horizon ahead
+        jepa_pred_dim: int = 256,        # bottleneck width of the predictor MLP
+        jepa_loss_weight: float = 0.1,   # weight of the (1 - cosine) JEPA loss
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.initializer_range = initializer_range
+        self.hidden_dropout = hidden_dropout
+
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.q_lora_rank = q_lora_rank
+        self.head_dim = head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.nope_head_dim = head_dim - qk_rope_head_dim
+        self.o_lora_rank = o_lora_rank
+        self.attention_dropout = attention_dropout
+        self.rope_theta = rope_theta
+        self.use_derf = use_derf
+        self.use_xsa = use_xsa
+
+        self.use_moe = use_moe
+        self.moe_intermediate_size = moe_intermediate_size
+        self.n_routed_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.norm_topk_prob = norm_topk_prob
+        self.scoring_func = scoring_func
+        self.routed_scaling_factor = routed_scaling_factor
+        self.num_hash_layers = num_hash_layers
+        self.moe_aux_loss_coef = moe_aux_loss_coef
+        self.moe_layers = moe_layers if moe_layers is not None else list(range(num_hidden_layers))
+
+        self.use_hyper_connections = use_hyper_connections
+        self.hc_mult = hc_mult
+        self.hc_sinkhorn_iters = hc_sinkhorn_iters
+        self.hc_eps = hc_eps
+
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+
+        self.use_engram = use_engram
+        self.engram_compress_dim = engram_compress_dim
+        self.engram_num_heads = engram_num_heads
+        self.engram_table_size = engram_table_size
+        self.engram_max_ngram = engram_max_ngram
+        self.engram_gate_init_bias = engram_gate_init_bias
+        self.use_hrm_refine = use_hrm_refine
+        self.hrm_refine_steps = hrm_refine_steps
+        self.hrm_refine_dim = hrm_refine_dim
+        self.use_qk_norm = use_qk_norm
+        self.zloss_coef = zloss_coef
+        self.mtp_loss_weight = mtp_loss_weight
+        self.use_value_embed = use_value_embed
+        self.use_jepa = use_jepa
+        self.jepa_horizon = jepa_horizon
+        self.jepa_pred_dim = jepa_pred_dim
+        self.jepa_loss_weight = jepa_loss_weight

model.safetensors

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

model_v2.py

@@ -0,0 +1,1042 @@
+"""
+model_v2.py -- SpikeWhaleLM v2: optimized base architecture.
+
+Changes vs model.py (v1):
+
+  PERFORMANCE
+  - SparseMoEFFN: sort-based expert dispatch (one contiguous slice per expert,
+    index_add_ scatter-back) replaces per-expert boolean masking. Far fewer
+    kernel launches, torch.compile-friendly (no data-dependent boolean
+    indexing in the hot path).
+  - Shared experts fused into ONE ExpertFFN with n_shared * intermediate width
+    (mathematically equivalent to the averaged sum, 1 matmul set instead of N).
+
+  QUALITY / STABILITY
+  - QK-Norm: per-head RMSNorm on Q and K before RoPE (Gemma2/OLMo2-style).
+    Stabilizes attention logits, tolerates higher LR. (cfg.use_qk_norm, default ON)
+  - z-loss on lm_head logits: zloss_coef * mean(log^2 Z). Prevents logit drift.
+    (cfg.zloss_coef, default 1e-4; set 0 to disable)
+  - MTP heads REDESIGNED: instead of K independent full H x V matrices (which at
+    50M params dwarfed the model), each MTP head is now a small zero-init H x H
+    projection feeding the SHARED lm_head. Param cost per head: H^2 instead of
+    H*V. MTP loss is down-weighted by cfg.mtp_loss_weight (default 0.3).
+  - HC output: learned softmax mix over streams (HCOutputMix) instead of mean().
+  - Value-embedding residual (nanoGPT-speedrun style): per-layer learned gate
+    (zero-init => exact no-op at init) adds a projection of the token embedding
+    into each block's input. (cfg.use_value_embed, default OFF = opt-in)
+
+  All new config keys are read with getattr(cfg, key, default) so your existing
+  config.py works unmodified. NOTE: QK-Norm and HCOutputMix add parameters, so
+  v1 checkpoints need load_state_dict(strict=False) (new params keep init;
+  QK-Norm at init is NOT identity -- prefer training v2 from scratch, or set
+  use_qk_norm=False to stay v1-loadable).
+
+  XSA is kept byte-identical to v1 but read the note in MLADerfXSAAttention:
+  with num_kv_heads == 1 it removes the SAME rank-1 value subspace from every
+  head. A/B it at 50M before keeping it in the final base.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from torch.utils.checkpoint import checkpoint as gradient_checkpoint
+
+# Force the LOCAL jepa_v2/config.py (with the JEPA fields) even when this file
+# is imported from the project root, where the root config.py would win.
+import os as _os, sys as _sys
+_sys.path.insert(0, _os.path.dirname(_os.path.abspath(__file__)))
+try:
+    # Dotted import so HuggingFace's trust_remote_code loader fetches config.py
+    # as a relative dependency; falls back to flat import for local script use.
+    from .config import SpikeWhaleConfig
+except ImportError:
+    from config import SpikeWhaleConfig
+
+
+# ---------------------------------------------------------------------------
+# Primitives
+# ---------------------------------------------------------------------------
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    """RoPE for the rope partition of Q and K (qk_rope_head_dim dims only)."""
+
+    def __init__(self, dim: int, max_positions: int = 4096, theta: float = 10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        t = torch.arange(max_positions).float()
+        freqs = torch.outer(t, inv_freq)
+        self.register_buffer("cos_cache", freqs.cos(), persistent=False)
+        self.register_buffer("sin_cache", freqs.sin(), persistent=False)
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
+        cos = self.cos_cache[position_ids].unsqueeze(1)   # [B, 1, S, rope_dim//2]
+        sin = self.sin_cache[position_ids].unsqueeze(1)
+        d = cos.shape[-1]
+        x1, x2 = x[..., :d], x[..., d:]
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# Engram: N-gram hash lookup + DERF gate (unchanged from v1)
+# ---------------------------------------------------------------------------
+
+class TokenCompressor(nn.Module):
+    def __init__(self, embed_dim: int, compress_dim: int):
+        super().__init__()
+        self.proj = nn.Linear(embed_dim, compress_dim, bias=False)
+        nn.init.normal_(self.proj.weight, std=0.02)
+        # Frozen LSH-style projection: gradient never reaches it through the
+        # .long() hash cast, so a fixed random projection is correct (see v1).
+        self.proj.weight.requires_grad_(False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.proj(x)
+
+
+class MultiHeadHashLookup(nn.Module):
+    def __init__(self, num_heads: int, table_size: int,
+                 compress_dim: int, out_dim: int, max_ngram: int = 3):
+        super().__init__()
+        self.num_heads = num_heads
+        self.table_size = table_size
+        self.max_ngram = max_ngram
+        self.out_dim = out_dim
+
+        self.tables = nn.ModuleList([
+            nn.Embedding(table_size, out_dim) for _ in range(num_heads)
+        ])
+        for t in self.tables:
+            nn.init.normal_(t.weight, std=0.01)
+
+        for n in range(1, max_ngram + 1):
+            for k in range(n):
+                proj = torch.randn(num_heads, compress_dim)
+                proj = proj / (proj.norm(dim=1, keepdim=True) + 1e-8)
+                self.register_buffer(f"hash_proj_n{n}_p{k}", proj)
+
+    def forward(self, compressed: torch.Tensor) -> torch.Tensor:
+        B, S, _ = compressed.shape
+        device = compressed.device
+        out = torch.zeros(B, S, self.out_dim, device=device, dtype=compressed.dtype)
+        norm = torch.zeros(S, device=device)
+
+        for n in range(1, self.max_ngram + 1):
+            if S < n:
+                continue
+            valid_len = S - n + 1
+            start = n - 1
+
+            h = torch.zeros(B, valid_len, self.num_heads, device=device)
+            for k in range(n):
+                proj = getattr(self, f"hash_proj_n{n}_p{k}")
+                h = h + torch.matmul(compressed[:, k:k + valid_len, :].float(), proj.t())
+
+            idx = h.abs().long() % self.table_size
+
+            for head_idx, table in enumerate(self.tables):
+                out[:, start:, :] = out[:, start:, :] + table(idx[:, :, head_idx])
+
+            norm[start:] += self.num_heads
+
+        return (out / norm.view(1, -1, 1).clamp(min=1)).to(compressed.dtype)
+
+
+class DERFContextGate(nn.Module):
+    def __init__(self, obs_size: int, init_bias: float = -4.0):
+        super().__init__()
+        self.proj = nn.Linear(obs_size * 2, obs_size)
+        self.alpha = nn.Parameter(torch.ones(obs_size))
+        self.bias = nn.Parameter(torch.full((obs_size,), init_bias))
+        self.gamma = nn.Parameter(torch.ones(obs_size))
+
+    def forward(self, retrieved: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        logits = self.proj(torch.cat([retrieved, x], dim=-1))
+        gate = self.gamma * ((torch.erf(self.alpha * logits + self.bias) + 1.0) / 2.0)
+        return retrieved * gate
+
+
+class EngramModule(nn.Module):
+    def __init__(self, cfg: SpikeWhaleConfig):
+        super().__init__()
+        self.compressor = TokenCompressor(cfg.hidden_size, cfg.engram_compress_dim)
+        self.lookup = MultiHeadHashLookup(
+            cfg.engram_num_heads, cfg.engram_table_size,
+            cfg.engram_compress_dim, cfg.hidden_size, cfg.engram_max_ngram,
+        )
+        self.gate = DERFContextGate(cfg.hidden_size, cfg.engram_gate_init_bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        compressed = self.compressor(x.detach())
+        retrieved = self.lookup(compressed)
+        return self.gate(retrieved, x)
+
+
+# ---------------------------------------------------------------------------
+# Hyper-Connections
+# ---------------------------------------------------------------------------
+
+class HyperConnectionLayer(nn.Module):
+    """Simplified HC: softmax pre-mix / post-distribute over hc_mult streams.
+    Asymmetric init (v1 bugfix) so streams diverge and gradients flow."""
+    def __init__(self, hidden_size: int, hc_mult: int,
+                 sinkhorn_iters: int = 20, eps: float = 1e-6):
+        super().__init__()
+        self.hc_mult = hc_mult
+        self.pre_weight = nn.Parameter(
+            torch.linspace(0.5, -0.5, hc_mult) / max(hc_mult, 1)
+        )
+        self.post_weight = nn.Parameter(
+            torch.linspace(-0.5, 0.5, hc_mult) / max(hc_mult, 1)
+        )
+
+    def pre_op(self, copies: torch.Tensor) -> torch.Tensor:
+        w = F.softmax(self.pre_weight, dim=0)
+        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)
+
+    def post_op(self, copies: torch.Tensor, delta: torch.Tensor) -> torch.Tensor:
+        w = F.softmax(self.post_weight, dim=0)
+        return copies + delta.unsqueeze(1) * w.view(1, -1, 1, 1)
+
+
+class HCOutputMix(nn.Module):
+    """
+    NEW (v2): learned combination of the hc_mult streams at the model output,
+    replacing the v1 mean(dim=1). Mean forces the streams toward redundancy at
+    exactly the point where you want them specialized. Initialized uniform so
+    it starts identical to mean() -- a strict generalization, zero risk.
+    """
+    def __init__(self, hc_mult: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hc_mult))   # softmax(0)=uniform=mean
+
+    def forward(self, copies: torch.Tensor) -> torch.Tensor:
+        w = F.softmax(self.weight, dim=0)
+        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)
+
+
+# ---------------------------------------------------------------------------
+# MLA + (DERF) + XSA Attention, now with QK-Norm
+# ---------------------------------------------------------------------------
+
+class MLADerfXSAAttention(nn.Module):
+    """
+    v2 additions:
+      - QK-Norm (cfg.use_qk_norm, default True): per-head RMSNorm applied to Q
+        and K BEFORE the rope/nope split. Bounds attention logits, the standard
+        modern stability fix; composes cleanly with SDPA and partial RoPE.
+
+    XSA NOTE (unchanged mechanics, important caveat): with num_kv_heads == 1
+    (MQA) every query head shares the same value vector, so the self-projection
+    subtraction removes the SAME rank-1 value subspace from all heads -- much
+    more aggressive than per-head XSA. Ablate use_xsa on/off at 50M before
+    locking the base config.
+    """
+
+    def __init__(self, cfg: SpikeWhaleConfig):
+        super().__init__()
+        self.num_heads = cfg.num_attention_heads
+        self.num_kv_heads = cfg.num_key_value_heads
+        self.head_dim = cfg.head_dim
+        self.qk_rope_head_dim = cfg.qk_rope_head_dim
+        self.nope_head_dim = cfg.nope_head_dim
+        self.hidden_size = cfg.hidden_size
+        self.use_derf = cfg.use_derf
+        # Jet-Long dynamic bifocal RoPE (arXiv:2607.07740), inference-time only.
+        self.use_jetlong = getattr(cfg, 'use_jetlong', False)
+        self.jetlong_w_pretrained = int(getattr(cfg, 'jetlong_w_pretrained', 4096))
+        self.jetlong_w0 = int(getattr(cfg, 'jetlong_w0', 2048))
+        self.use_xsa = cfg.use_xsa
+        self.dropout_p = cfg.attention_dropout
+        self.kv_groups = self.num_heads // self.num_kv_heads
+        self.use_qk_norm = getattr(cfg, "use_qk_norm", True)
+
+        self.q_a_proj = nn.Linear(cfg.hidden_size, cfg.q_lora_rank, bias=False)
+        self.q_a_norm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps)
+        self.q_b_proj = nn.Linear(cfg.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+
+        self.k_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+
+        self.o_a_proj = nn.Linear(self.num_heads * self.head_dim, cfg.o_lora_rank, bias=False)
+        self.o_b_proj = nn.Linear(cfg.o_lora_rank, cfg.hidden_size, bias=False)
+
+        # QK-Norm: one RMSNorm over head_dim, shared across heads (Gemma-2 style).
+        if self.use_qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+
+        self.rope = RotaryEmbedding(
+            self.qk_rope_head_dim,
+            max_positions=cfg.max_position_embeddings,
+            theta=cfg.rope_theta,
+        )
+
+        if self.use_derf:
+            self.derf_alpha = nn.Parameter(torch.ones(self.num_heads))
+            self.derf_bias = nn.Parameter(torch.zeros(self.num_heads))
+            self.derf_gamma = nn.Parameter(torch.ones(self.num_heads))
+
+        for m in (self.q_a_proj, self.q_b_proj, self.k_proj,
+                  self.v_proj, self.o_a_proj, self.o_b_proj):
+            nn.init.normal_(m.weight, std=cfg.initializer_range)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        B, S, _ = x.shape
+
+        q = self.q_a_norm(self.q_a_proj(x))
+        q = self.q_b_proj(q).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+
+        k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        # QK-Norm before RoPE (v2). Cache stores the NORMALIZED k so prefill and
+        # incremental decode agree.
+        if self.use_qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+
+        q_nope = q[..., :self.nope_head_dim]
+        q_rope = q[..., self.nope_head_dim:]
+        k_nope = k[..., :self.nope_head_dim]
+        k_rope = k[..., self.nope_head_dim:]
+
+        q_rope = self.rope(q_rope, position_ids)
+        k_rope = self.rope(k_rope, position_ids)
+
+        q = torch.cat([q_nope, q_rope], dim=-1)
+        k = torch.cat([k_nope, k_rope], dim=-1)
+
+        if past_key_value is not None:
+            k = torch.cat([past_key_value[0], k], dim=2)
+            v = torch.cat([past_key_value[1], v], dim=2)
+        present = (k, v) if use_cache else None
+        N = k.shape[2]
+
+        if self.kv_groups > 1:
+            k = k.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
+                B, self.num_heads, N, self.head_dim)
+            v = v.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
+                B, self.num_heads, N, self.head_dim)
+
+        # --- Jet-Long dynamic bifocal RoPE (arXiv:2607.07740) ---
+        # No-op within the native window (G==1); position-aliased remote
+        # window beyond it. Correction rotation composes on the already-RoPE'd
+        # cache, so the KV cache (and TriAttention compression) is untouched.
+        y = None
+        if getattr(self, 'use_jetlong', False) and not self.use_derf:
+            _G = max(1, math.ceil(N / self.jetlong_w_pretrained))
+            if _G > 1:
+                y = self._jetlong_attend(q, k, v, position_ids, S, N, _G)
+        if y is not None:
+            pass
+        elif self.use_derf:
+            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+            if attention_mask is None and past_key_value is None:
+                is_masked = torch.triu(
+                    torch.ones(S, N, dtype=torch.bool, device=scores.device),
+                    diagonal=N - S + 1,
+                ).unsqueeze(0).unsqueeze(0)
+            else:
+                is_masked = (attention_mask < -1.0) if attention_mask is not None \
+                    else torch.zeros_like(scores, dtype=torch.bool)
+
+            safe_scores = scores.masked_fill(is_masked, -10000.0)
+
+            a = self.derf_alpha.view(1, -1, 1, 1)
+            b = self.derf_bias.view(1, -1, 1, 1)
+            g = self.derf_gamma.view(1, -1, 1, 1)
+
+            attn_weights = g * torch.erf(a * safe_scores + b)
+            attn_weights = (attn_weights + g) / 2.0
+            attn_weights = attn_weights.masked_fill(is_masked, 0.0)
+            attn_weights = attn_weights / (attn_weights.sum(dim=-1, keepdim=True) + 1e-8)
+
+            if self.dropout_p > 0 and self.training:
+                attn_weights = F.dropout(attn_weights, p=self.dropout_p)
+
+            y = torch.matmul(attn_weights, v)
+        else:
+            q = q.contiguous()
+            k = k.contiguous()
+            v = v.contiguous()
+            drop = self.dropout_p if self.training else 0.0
+            if past_key_value is None and attention_mask is None:
+                y = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=drop)
+            else:
+                if attention_mask is not None:
+                    is_masked = (attention_mask < -1.0)
+                else:
+                    is_masked = torch.triu(
+                        torch.ones(S, N, dtype=torch.bool, device=q.device),
+                        diagonal=N - S + 1,
+                    ).unsqueeze(0).unsqueeze(0)
+                y = F.scaled_dot_product_attention(
+                    q, k, v, attn_mask=~is_masked, dropout_p=drop)
+
+        if self.use_xsa:
+            past_len = N - S
+            v_self = v[:, :, past_len:past_len + S, :]
+            vn = v_self / (v_self.norm(dim=-1, keepdim=True) + 1e-8)
+            projection = (y * vn).sum(dim=-1, keepdim=True) * vn
+            y = y - projection
+
+        y = y.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim)
+        y = self.o_b_proj(self.o_a_proj(y))
+        return y, present
+
+    def _jetlong_attend(self, q, k, v, position_ids, S, N, G):
+        """Bifocal softmax attention (Jet-Long, arXiv:2607.07740).
+
+        q, k arrive already RoPE-rotated at their true absolute positions. We add
+        a *correction rotation* delta = floor(p/G) - p to realize the remote,
+        position-aliased view (RoPE composes additively, so this works directly on
+        the rotated -- and possibly TriAttention-compressed -- cache without
+        rewriting it), then merge the local and remote windows by query-key
+        distance. Handles prefill (S==N) and decode (S==1, N==cache+1) uniformly.
+        Only reached when use_jetlong and G>1, so behaviour within the native
+        window is byte-identical to the base model."""
+        import math as _m
+        B = q.shape[0]
+        nd = self.nope_head_dim
+        dev = q.device
+        if position_ids is not None:
+            q_pos = position_ids.reshape(B, S).to(dev).long()
+        else:
+            q_pos = torch.arange(N - S, N, device=dev).view(1, S).expand(B, S)
+        _kp = getattr(self, '_jl_key_pos', None)
+        if _kp is not None and _kp.numel() == N:
+            k_pos = _kp.to(dev).long().view(1, N).expand(B, N)
+        else:
+            k_pos = torch.arange(N, device=dev).view(1, N).expand(B, N)
+        dq = torch.div(q_pos, G, rounding_mode='floor') - q_pos          # [B,S]
+        dk = torch.div(k_pos, G, rounding_mode='floor') - k_pos          # [B,N]
+        inv_freq = self.rope.inv_freq.to(dev).float()                    # [rd/2]
+
+        def rot_by(x, delta):
+            d = x.shape[-1] // 2
+            ang = delta.float().unsqueeze(1).unsqueeze(-1) * inv_freq    # [B,1,T,d]
+            cos, sin = ang.cos(), ang.sin()
+            x1, x2 = x[..., :d], x[..., d:]
+            return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+        q_nope, q_rope = q[..., :nd], q[..., nd:]
+        k_nope, k_rope = k[..., :nd], k[..., nd:]
+        q_rope_rem = rot_by(q_rope, dq)
+        k_rope_rem = rot_by(k_rope, dk)
+
+        scale = 1.0 / _m.sqrt(self.head_dim)
+        s_nope = torch.matmul(q_nope, k_nope.transpose(-2, -1))
+        s_loc = (s_nope + torch.matmul(q_rope, k_rope.transpose(-2, -1))) * scale
+        s_rem = (s_nope + torch.matmul(q_rope_rem, k_rope_rem.transpose(-2, -1))) * scale
+
+        dist = q_pos.view(B, 1, S, 1) - k_pos.view(B, 1, 1, N)          # i - j
+        local_mask = (dist <= self.jetlong_w0) & (dist >= 0)
+        causal = dist >= 0
+        scores = torch.where(local_mask, s_loc, s_rem)
+        scores = scores.masked_fill(~causal, float('-inf'))
+        attn = torch.softmax(scores.float(), dim=-1).to(v.dtype)
+        return torch.matmul(attn, v)
+
+
+# ---------------------------------------------------------------------------
+# MoE FFN -- v2: sort-based dispatch + fused shared expert
+# ---------------------------------------------------------------------------
+
+class ExpertFFN(nn.Module):
+    """Single SwiGLU expert."""
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        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: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+def sqrtsoftplus(x: torch.Tensor) -> torch.Tensor:
+    return torch.sqrt(F.softplus(x) + 1e-8)
+
+
+class SparseMoEFFN(nn.Module):
+    """
+    v2 changes:
+      - FUSED shared expert: one ExpertFFN with width n_shared * intermediate,
+        scaled by 1/n_shared on output -- equivalent to v1's averaged Python
+        loop, one fused matmul set. (state-dict key changes: shared_expert.*)
+      - SORT-BASED dispatch for routed experts: flatten (token, slot) pairs,
+        argsort by expert id, run each expert on ONE contiguous slice, weighted
+        index_add_ back. No boolean masks, no nonzero(), no per-expert scatter.
+    Routing logic (hash routing, sqrtsoftplus, aux loss) is unchanged.
+    """
+    def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int = 0):
+        super().__init__()
+        self.n_routed_experts = cfg.n_routed_experts
+        self.n_shared_experts = cfg.n_shared_experts
+        self.num_experts_per_tok = cfg.num_experts_per_tok
+        self.norm_topk_prob = cfg.norm_topk_prob
+        self.scoring_func = cfg.scoring_func
+        self.routed_scaling_factor = cfg.routed_scaling_factor
+        self.use_hash_routing = layer_idx < cfg.num_hash_layers
+        self.aux_loss_coef = cfg.moe_aux_loss_coef
+
+        self.router = nn.Linear(cfg.hidden_size, cfg.n_routed_experts, bias=False)
+        self.experts = nn.ModuleList([
+            ExpertFFN(cfg.hidden_size, cfg.moe_intermediate_size)
+            for _ in range(cfg.n_routed_experts)
+        ])
+        # Fused shared expert (v2)
+        self.shared_expert = (
+            ExpertFFN(cfg.hidden_size,
+                      cfg.moe_intermediate_size * cfg.n_shared_experts)
+            if cfg.n_shared_experts > 0 else None
+        )
+
+        self._last_aux_loss: Optional[torch.Tensor] = None
+
+    def forward(self, x: torch.Tensor,
+                position_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, S, H = x.shape
+        x_flat = x.view(B * S, H)
+        T = B * S
+        K = self.num_experts_per_tok
+
+        # Shared expert: always active, single fused pass.
+        if self.shared_expert is not None:
+            shared_out = self.shared_expert(x_flat)
+            if self.n_shared_experts > 1:
+                shared_out = shared_out / self.n_shared_experts
+        else:
+            shared_out = None
+
+        # ---- Routing (unchanged logic) ----
+        if self.use_hash_routing:
+            if position_ids is not None:
+                base = (position_ids.reshape(T, 1) % self.n_routed_experts).long()
+            else:
+                base = (torch.arange(T, device=x.device) % self.n_routed_experts).unsqueeze(1)
+            offsets = torch.arange(K, device=x.device)
+            top_k_indices = (base + offsets.unsqueeze(0)) % self.n_routed_experts  # [T, K]
+            top_k_weights = torch.full((T, K), 1.0 / K, device=x.device, dtype=x_flat.dtype)
+            self._last_aux_loss = None
+        else:
+            router_logits = self.router(x_flat)
+            if self.scoring_func == "sqrtsoftplus":
+                routing_scores = sqrtsoftplus(router_logits)
+            else:
+                routing_scores = F.softmax(router_logits, dim=-1)
+
+            top_k_scores, top_k_indices = torch.topk(routing_scores, K, dim=-1)
+            if self.norm_topk_prob:
+                top_k_weights = top_k_scores / (top_k_scores.sum(dim=-1, keepdim=True) + 1e-8)
+            else:
+                top_k_weights = top_k_scores
+            top_k_weights = top_k_weights * self.routed_scaling_factor
+
+            softmax_probs = F.softmax(router_logits, dim=-1)
+            expert_mask = torch.zeros_like(softmax_probs)
+            expert_mask.scatter_(1, top_k_indices, 1.0)
+            f_e = expert_mask.mean(0)
+            p_e = softmax_probs.mean(0)
+            self._last_aux_loss = self.n_routed_experts * (f_e * p_e).sum() * self.aux_loss_coef
+
+        # ---- Sort-based dispatch (v2) ----
+        # Flatten the (token, slot) assignment: T*K rows total.
+        flat_expert = top_k_indices.reshape(-1)                       # [T*K]
+        flat_weight = top_k_weights.reshape(-1, 1)                    # [T*K, 1]
+        flat_token = torch.arange(T, device=x.device).repeat_interleave(K)  # [T*K]
+
+        order = torch.argsort(flat_expert, stable=True)               # group by expert
+        sorted_expert = flat_expert[order]
+        sorted_token = flat_token[order]
+        sorted_weight = flat_weight[order]
+
+        counts = torch.bincount(sorted_expert, minlength=self.n_routed_experts)
+        # boundaries per expert in the sorted order (CPU sync once per forward;
+        # unavoidable without grouped-GEMM, still vastly cheaper than v1's
+        # per-expert nonzero/masking)
+        counts_list = counts.tolist()
+
+        gathered = x_flat[sorted_token]                               # [T*K, H]
+        out_flat = torch.zeros_like(x_flat)
+
+        start = 0
+        for expert_idx, cnt in enumerate(counts_list):
+            if cnt == 0:
+                continue
+            end = start + cnt
+            seg = gathered[start:end]
+            seg_out = self.experts[expert_idx](seg) * sorted_weight[start:end]
+            out_flat.index_add_(0, sorted_token[start:end], seg_out.to(out_flat.dtype))
+            start = end
+
+        if shared_out is not None:
+            out_flat = out_flat + shared_out
+        return out_flat.view(B, S, H)
+
+    def get_aux_loss(self) -> Optional[torch.Tensor]:
+        return self._last_aux_loss
+
+
+class DenseFFN(nn.Module):
+    def __init__(self, cfg: SpikeWhaleConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.moe_intermediate_size, bias=False)
+        self.up_proj = nn.Linear(cfg.hidden_size, cfg.moe_intermediate_size, bias=False)
+        self.down_proj = nn.Linear(cfg.moe_intermediate_size, cfg.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor,
+                position_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+    def get_aux_loss(self) -> Optional[torch.Tensor]:
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Transformer block
+# ---------------------------------------------------------------------------
+
+class TransformerBlock(nn.Module):
+    def __init__(self, cfg: SpikeWhaleConfig, layer_idx: int):
+        super().__init__()
+        self.use_hc = cfg.use_hyper_connections
+        self.hidden_dropout = cfg.hidden_dropout
+        self.use_value_embed = getattr(cfg, "use_value_embed", False)
+
+        self.attn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.attn = MLADerfXSAAttention(cfg)
+        self.ffn_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+
+        if cfg.use_moe and layer_idx in cfg.moe_layers:
+            self.ffn = SparseMoEFFN(cfg, layer_idx)
+            self.is_moe = True
+        else:
+            self.ffn = DenseFFN(cfg)
+            self.is_moe = False
+
+        if self.use_hc:
+            self.hc_attn = HyperConnectionLayer(cfg.hidden_size, cfg.hc_mult,
+                                                cfg.hc_sinkhorn_iters, cfg.hc_eps)
+            self.hc_ffn = HyperConnectionLayer(cfg.hidden_size, cfg.hc_mult,
+                                               cfg.hc_sinkhorn_iters, cfg.hc_eps)
+
+        # NEW (v2, opt-in): value-embedding residual. Zero-init gate -> exact
+        # no-op at init; learns to mix raw token-embedding signal into each
+        # block's input (nanoGPT-speedrun "value embedding"/U-net skip family;
+        # consistent wins at the 50-500M scale).
+        if self.use_value_embed:
+            self.ve_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(
+        self,
+        x: torch.Tensor,           # [B, hc_mult, S, H] if HC else [B, S, H]
+        position_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Tuple] = None,
+        use_cache: bool = False,
+        token_embed: Optional[torch.Tensor] = None,   # [B, S, H] (value-embed)
+    ) -> Tuple[torch.Tensor, Optional[Tuple], Optional[torch.Tensor]]:
+
+        # --- Attention sub-layer ---
+        if self.use_hc:
+            h = self.hc_attn.pre_op(x)
+        else:
+            h = x
+
+        if self.use_value_embed and token_embed is not None:
+            h = h + torch.tanh(self.ve_gate) * token_embed
+
+        attn_out, present = self.attn(
+            self.attn_norm(h), position_ids, attention_mask, past_key_value, use_cache
+        )
+        attn_out = F.dropout(attn_out, p=self.hidden_dropout, training=self.training)
+
+        if self.use_hc:
+            x = self.hc_attn.post_op(x, attn_out)
+            h = self.hc_ffn.pre_op(x)
+        else:
+            h = h + attn_out
+
+        # --- FFN sub-layer ---
+        ffn_out = self.ffn(self.ffn_norm(h), position_ids)
+        ffn_out = F.dropout(ffn_out, p=self.hidden_dropout, training=self.training)
+
+        if self.use_hc:
+            x = self.hc_ffn.post_op(x, ffn_out)
+        else:
+            x = h + ffn_out
+
+        return x, present, self.ffn.get_aux_loss()
+
+
+# ---------------------------------------------------------------------------
+# HRM refinement
+# ---------------------------------------------------------------------------
+
+class HRMRefinementBlock(nn.Module):
+    def __init__(self, hidden_size: int, refine_dim: int, steps: int, eps: float = 1e-6):
+        super().__init__()
+        self.steps = steps
+        self.norm = RMSNorm(hidden_size, eps)
+        self.down = nn.Linear(hidden_size * 2, refine_dim, bias=False)
+        self.up = nn.Linear(refine_dim, hidden_size, bias=False)
+        self.gate = nn.Parameter(torch.zeros(steps))
+        nn.init.normal_(self.down.weight, std=0.02)
+        # FIX vs the original: up was ALSO zero-init, so update == 0 (kills the
+        # gate's gradient) and tanh(gate) == 0 (kills up's gradient) -- a saddle
+        # both gradients can never leave; the block stayed a no-op forever.
+        # The zero gate alone already makes init an exact no-op; up must be
+        # nonzero so the gate receives gradient.
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        anchor = x
+        h = x
+        for t in range(self.steps):
+            inp = torch.cat([self.norm(h), anchor], dim=-1)
+            update = self.up(F.silu(self.down(inp)))
+            h = h + torch.tanh(self.gate[t]) * update
+        return h
+
+
+# ---------------------------------------------------------------------------
+# JEPA secondary prediction head (jepa_v2)
+# ---------------------------------------------------------------------------
+
+class JEPAPredictorBlock(nn.Module):
+    """
+    JEPA-inspired representation-space prediction (I-JEPA / LLM-JEPA family):
+    from the trunk's hidden state at position t, predict the trunk's OWN hidden
+    state at position t+k. The target is stop-gradient (detached), the standard
+    JEPA asymmetry that prevents the trivial collapse where the trunk just
+    makes all hidden states identical.
+
+    This complements the MTP heads: MTP predicts future TOKENS through the
+    lm_head (output space); JEPA predicts the future REPRESENTATION directly
+    (embedding space), pressuring the trunk to encode where its own state is
+    going -- abstract next-step structure rather than surface vocabulary.
+
+    Deliberately shaped like HRMRefinementBlock: RMSNorm -> down (bottleneck)
+    -> SiLU -> up -> per-offset tanh-gated residual. gate starts at zero so the
+    predictor is exactly identity at init (zero-risk insertion); up is NORMAL
+    init (not zero) so the gate actually receives gradient -- see the
+    double-zero saddle note on HRMRefinementBlock above.
+    """
+    def __init__(self, hidden_size: int, pred_dim: int, horizon: int, eps: float = 1e-6):
+        super().__init__()
+        self.horizon = horizon
+        self.norm = RMSNorm(hidden_size, eps)
+        self.down = nn.Linear(hidden_size, pred_dim, bias=False)
+        self.up = nn.Linear(pred_dim, hidden_size, bias=False)
+        self.gate = nn.Parameter(torch.zeros(horizon))
+        nn.init.normal_(self.down.weight, std=0.02)
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, h: torch.Tensor, k: int) -> torch.Tensor:
+        """Predict the hidden state k steps ahead of each position in h."""
+        update = self.up(F.silu(self.down(self.norm(h))))
+        return h + torch.tanh(self.gate[k - 1]) * update
+
+
+# ---------------------------------------------------------------------------
+# Full model
+# ---------------------------------------------------------------------------
+
+class SpikeWhaleModel(nn.Module):
+    """Decoder stack without LM head."""
+
+    def __init__(self, cfg: SpikeWhaleConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        nn.init.normal_(self.embed_tokens.weight, std=cfg.initializer_range)
+
+        self.engram = EngramModule(cfg) if cfg.use_engram else None
+        self.layers = nn.ModuleList([
+            TransformerBlock(cfg, layer_idx=i)
+            for i in range(cfg.num_hidden_layers)
+        ])
+        self.norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps)
+        self.hc_out_mix = (
+            HCOutputMix(cfg.hc_mult) if cfg.use_hyper_connections else None
+        )
+        self.hrm_refine = (
+            HRMRefinementBlock(cfg.hidden_size, cfg.hrm_refine_dim, cfg.hrm_refine_steps,
+                               cfg.rms_norm_eps)
+            if getattr(cfg, "use_hrm_refine", False) else None
+        )
+        self.use_value_embed = getattr(cfg, "use_value_embed", False)
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[Tuple]] = None,
+        use_cache: bool = False,
+        engram_context_ids: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[List[Tuple]], torch.Tensor]:
+        B, S = input_ids.shape
+        device = input_ids.device
+
+        if position_ids is None:
+            past_len = past_key_values[0][0].shape[2] if past_key_values else 0
+            position_ids = torch.arange(
+                past_len, past_len + S, device=device
+            ).unsqueeze(0).expand(B, -1)
+
+        x = self.embed_tokens(input_ids)
+        token_embed = x if self.use_value_embed else None
+
+        if self.engram is not None:
+            if engram_context_ids is not None and engram_context_ids.numel() > 0:
+                # Cached decode: the n-gram hashes need the (max_ngram - 1)
+                # tokens BEFORE this window, which a KV cache does not carry.
+                # Prepend their embeddings, run the engram over the joined
+                # window, and keep only this window's positions.
+                ctx = self.embed_tokens(engram_context_ids)
+                n_ctx = ctx.shape[1]
+                x = x + self.engram(torch.cat([ctx, x], dim=1))[:, n_ctx:, :]
+            else:
+                x = x + self.engram(x)
+
+        if self.cfg.use_hyper_connections:
+            x = x.unsqueeze(1).expand(-1, self.cfg.hc_mult, -1, -1).clone()
+
+        present_key_values = [] if use_cache else None
+        total_aux_loss = torch.tensor(0.0, device=device)
+
+        # Gradient checkpointing is incompatible with use_cache (the cache from
+        # the discarded forward would be silently wrong on recompute).
+        assert not (self.gradient_checkpointing and self.training and use_cache), \
+            "use_cache=True is not supported with gradient checkpointing"
+
+        for layer_idx, layer in enumerate(self.layers):
+            pkv = past_key_values[layer_idx] if past_key_values else None
+
+            if self.gradient_checkpointing and self.training:
+                x, present, aux_loss = gradient_checkpoint(
+                    layer, x, position_ids, attention_mask, None, False, token_embed,
+                    use_reentrant=False,
+                )
+            else:
+                x, present, aux_loss = layer(
+                    x, position_ids, attention_mask, pkv, use_cache, token_embed)
+
+            if use_cache:
+                present_key_values.append(present)
+            if aux_loss is not None:
+                total_aux_loss = total_aux_loss + aux_loss
+
+        if self.cfg.use_hyper_connections:
+            x = self.hc_out_mix(x)          # v2: learned mix (init == mean)
+
+        if self.hrm_refine is not None:
+            x = self.hrm_refine(x)
+
+        x = self.norm(x)
+        return x, present_key_values, total_aux_loss
+
+
+class MTPHead(nn.Module):
+    """
+    v2 MTP head: small zero-init H x H projection feeding the SHARED lm_head.
+    Cost per head: H^2 params (e.g. 1M at H=1024) instead of H*V (e.g. 50M+).
+    Zero-init means at step 0 the head predicts exactly what lm_head predicts
+    for the residual path = 0, i.e. uniform-ish gradient pressure; the residual
+    form (x + proj(x)) keeps it anchored to the trunk representation.
+    """
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+        return hidden + self.proj(hidden)
+
+
+class SpikeWhaleLM(PreTrainedModel):
+    """
+    v2 loss = CE + zloss_coef * z-loss
+            + mtp_loss_weight * mean(MTP CE)
+            + jepa_loss_weight * mean(JEPA 1-cosine)   (jepa_v2)
+            + MoE aux loss
+    """
+    config_class = SpikeWhaleConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["TransformerBlock"]
+
+    def __init__(self, cfg: SpikeWhaleConfig):
+        super().__init__(cfg)
+        self.model = SpikeWhaleModel(cfg)
+        self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False)
+        nn.init.normal_(self.lm_head.weight, std=cfg.initializer_range)
+
+        self.zloss_coef = getattr(cfg, "zloss_coef", 1e-4)
+        self.mtp_loss_weight = getattr(cfg, "mtp_loss_weight", 0.3)
+
+        # v2 MTP: H x H residual projections sharing lm_head (see MTPHead).
+        self.mtp_heads = nn.ModuleList([
+            MTPHead(cfg.hidden_size)
+            for _ in range(cfg.num_nextn_predict_layers)
+        ]) if cfg.num_nextn_predict_layers > 0 else None
+
+        # JEPA secondary prediction head (representation-space, stop-grad target).
+        self.jepa_loss_weight = getattr(cfg, "jepa_loss_weight", 0.1)
+        self.jepa_horizon = getattr(cfg, "jepa_horizon", 1)
+        self.jepa = (
+            JEPAPredictorBlock(cfg.hidden_size,
+                               getattr(cfg, "jepa_pred_dim", 256),
+                               self.jepa_horizon, cfg.rms_norm_eps)
+            if getattr(cfg, "use_jepa", False) else None
+        )
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def tie_weights(self, **kwargs):
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+    def save_pretrained(self, *args, **kwargs):
+        tied = (
+            self.config.tie_word_embeddings
+            and self.lm_head.weight.data_ptr() == self.model.embed_tokens.weight.data_ptr()
+        )
+        if tied:
+            self.lm_head.weight = nn.Parameter(self.model.embed_tokens.weight.detach().clone())
+        try:
+            super().save_pretrained(*args, **kwargs)
+        finally:
+            if tied:
+                self.lm_head.weight = self.model.embed_tokens.weight
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, SpikeWhaleModel):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[Tuple]] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        engram_context_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        hidden, present_kvs, aux_loss = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            engram_context_ids=engram_context_ids,
+        )
+
+        logits = self.lm_head(hidden)
+        loss = None
+
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            flat_logits = shift_logits.view(-1, shift_logits.size(-1))
+            flat_labels = shift_labels.view(-1)
+            loss = F.cross_entropy(flat_logits, flat_labels, ignore_index=-100)
+
+            # z-loss (v2): penalize log^2 of the partition function on valid
+            # positions. Keeps logits from drifting; pairs well with Muon.
+            if self.zloss_coef > 0:
+                valid = flat_labels != -100
+                if valid.any():
+                    log_z = torch.logsumexp(flat_logits[valid].float(), dim=-1)
+                    loss = loss + self.zloss_coef * (log_z ** 2).mean()
+
+            # MTP (v2): residual H x H head -> shared lm_head, down-weighted.
+            if self.mtp_heads is not None and self.mtp_loss_weight > 0:
+                mtp_total = torch.tensor(0.0, device=loss.device)
+                n_active = 0
+                for k, head in enumerate(self.mtp_heads, start=1):
+                    offset = k + 1
+                    if hidden.size(1) > offset:
+                        mtp_hidden = head(hidden[..., :-offset, :])
+                        mtp_logits = self.lm_head(mtp_hidden)
+                        mtp_labels = labels[..., offset:].contiguous()
+                        mtp_total = mtp_total + F.cross_entropy(
+                            mtp_logits.reshape(-1, mtp_logits.size(-1)),
+                            mtp_labels.reshape(-1),
+                            ignore_index=-100,
+                        )
+                        n_active += 1
+                if n_active > 0:
+                    loss = loss + self.mtp_loss_weight * mtp_total / n_active
+
+            # JEPA (jepa_v2): predict the trunk's hidden state k steps ahead in
+            # representation space. Target is DETACHED (JEPA stop-gradient) and
+            # the loss is (1 - cosine), computed only on positions whose target
+            # carries a real label (skips padding / masked prompt tokens).
+            if self.jepa is not None and self.jepa_loss_weight > 0:
+                jepa_total = torch.tensor(0.0, device=loss.device)
+                n_jepa = 0
+                for k in range(1, self.jepa_horizon + 1):
+                    if hidden.size(1) <= k:
+                        break
+                    pred = self.jepa(hidden[..., :-k, :], k)
+                    target = hidden[..., k:, :].detach()
+                    valid = labels[..., k:] != -100
+                    if not valid.any():
+                        continue
+                    cos = F.cosine_similarity(pred.float(), target.float(), dim=-1)
+                    jepa_total = jepa_total + (1.0 - cos)[valid].mean()
+                    n_jepa += 1
+                if n_jepa > 0:
+                    loss = loss + self.jepa_loss_weight * jepa_total / n_jepa
+
+            loss = loss + aux_loss
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=present_kvs,
+        )
+
+    def count_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters())

spike_tokenizer.py

@@ -0,0 +1,117 @@
+"""
+spike_tokenizer.py -- HuggingFace-compatible wrapper for the custom
+byte-level "length-max" (greedy longest-match) tokenizer in tokenizer.json.
+
+The raw tokenizer.json is NOT a HuggingFace `tokenizers` file; it is a plain
+dict {vocab, vocab_size, max_token_len, algorithm:"length-max"}. This wrapper
+makes it loadable by AutoTokenizer.from_pretrained / save_pretrained and
+exposes encode/decode + the bos/eos/pad/unk ids the training scripts expect.
+
+Encoding scheme (verified): byte-level. Text is UTF-8 encoded, each byte mapped
+to its latin-1 character, then greedily matched against the vocab using the
+longest key that matches at each position (max key length = max_token_len).
+"""
+import json, os
+from typing import List, Optional
+from transformers import PreTrainedTokenizer
+
+
+class SpikeTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "tokenizer.json"}
+    model_input_names = ["input_ids"]
+
+    def __init__(self, vocab_file=None, **kwargs):
+        with open(vocab_file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        self._vocab = data["vocab"]                       # str -> id
+        self._ids_to_tokens = {i: t for t, i in self._vocab.items()}
+        self.max_token_len = int(data.get("max_token_len", 24))
+        # length-bucketed keys for fast greedy match (longest length first)
+        self._lengths = sorted({len(k) for k in self._vocab}, reverse=True)
+
+        # Appended special tokens (im_start / <think> / <begin_solution> / ...).
+        # They already live in self._vocab at their real ids; we hand them to the
+        # HF base class as `additional_special_tokens` so its AddedToken trie:
+        #   (1) splits them out ATOMICALLY before our byte-level greedy match
+        #       (verified: each maps back to its existing vocab id, no phantom id), and
+        #   (2) drops them on decode(skip_special_tokens=True).
+        # The set is stored in tokenizer.json under "special_tokens" so it
+        # survives save_pretrained/from_pretrained round-trips.
+        self._extra_specials = [
+            t for t in data.get("special_tokens", []) if t in self._vocab
+        ]
+        if self._extra_specials:
+            existing = list(kwargs.get("additional_special_tokens", []) or [])
+            merged = existing + [t for t in self._extra_specials if t not in existing]
+            kwargs["additional_special_tokens"] = merged
+
+        kwargs.setdefault("bos_token", "<bos>")
+        kwargs.setdefault("eos_token", "<eos>")
+        kwargs.setdefault("unk_token", "<unk>")
+        kwargs.setdefault("pad_token", "<pad>")
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._vocab)
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    # --- core byte-level greedy tokenization ---
+    def _tokenize(self, text: str) -> List[str]:
+        s = text.encode("utf-8").decode("latin-1")  # one char per byte
+        out, i, n = [], 0, len(s)
+        while i < n:
+            matched = None
+            hi = min(self.max_token_len, n - i)
+            for L in range(hi, 0, -1):
+                sub = s[i:i + L]
+                if sub in self._vocab:
+                    matched = sub
+                    break
+            if matched is None:           # single byte always exists in vocab
+                matched = s[i]
+            out.append(matched)
+            i += len(matched)
+        return out
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self._vocab.get(token, self._vocab["<unk>"])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self._ids_to_tokens.get(index, "<unk>")
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        # transformers 5.x hands the FULL token list here (special tokens
+        # included; skip_special_tokens is already applied upstream via
+        # convert_ids_to_tokens). So we can't just byte-decode everything: a
+        # special token like "<|im_start|>" is a literal marker, not latin-1
+        # bytes. Decode runs of ordinary byte-tokens together (needed so
+        # multi-byte UTF-8 sequences reassemble) and emit any special token
+        # inline as its literal string.
+        specials = {"<pad>", "<unk>", "<bos>", "<eos>", *self._extra_specials}
+        out, buf = [], []
+        for tok in tokens:
+            if tok in specials:
+                if buf:
+                    out.append("".join(buf).encode("latin-1").decode("utf-8", errors="replace"))
+                    buf = []
+                out.append(tok)
+            else:
+                buf.append(tok)
+        if buf:
+            out.append("".join(buf).encode("latin-1").decode("utf-8", errors="replace"))
+        return "".join(out)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None):
+        os.makedirs(save_directory, exist_ok=True)
+        fn = (filename_prefix + "-" if filename_prefix else "") + "tokenizer.json"
+        path = os.path.join(save_directory, fn)
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump({"vocab": self._vocab, "vocab_size": self.vocab_size,
+                       "max_token_len": self.max_token_len,
+                       "algorithm": "length-max",
+                       "special_tokens": list(self._extra_specials)},
+                      f, ensure_ascii=False)
+        return (path,)

tokenizer_config.json

@@ -0,0 +1,204 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<eos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16384": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16385": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16386": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16387": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16388": {
+      "content": "<begin_solution>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16389": {
+      "content": "<end_solution>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16390": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16391": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16392": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16393": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16394": {
+      "content": "<|system|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16395": {
+      "content": "<|user|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16396": {
+      "content": "<|assistant|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16397": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16398": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16399": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16400": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "backend": "custom",
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "extra_special_tokens": [
+    "</think>",
+    "</tool_call>",
+    "</tool_response>",
+    "<begin_solution>",
+    "<end_solution>",
+    "<think>",
+    "<tool_call>",
+    "<tool_response>",
+    "<|assistant|>",
+    "<|endoftext|>",
+    "<|fim_middle|>",
+    "<|fim_prefix|>",
+    "<|fim_suffix|>",
+    "<|im_end|>",
+    "<|im_start|>",
+    "<|system|>",
+    "<|user|>"
+  ],
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "tokenizer_class": "SpikeTokenizer",
+  "unk_token": "<unk>",
+  "auto_map": {
+    "AutoTokenizer": [
+      "spike_tokenizer.SpikeTokenizer",
+      null
+    ]
+  }
+}