0.3.2a2: MLP-scope recompute, lookahead options, web UI fields, compiled document-reset fix

CHANGELOG.md

@@ -2,6 +2,32 @@
 
 ## [Unreleased]
 
+## [0.3.2a2] — 2026-11-09 (beta)
+
+Adds narrower recomputation, two lookahead options and web UI fields for every long-context switch, and fixes a crash found by training on the real corpus. Every new switch is off by default. `reports/long_context_20261109/real_corpus_ab.json` holds the measurements below.
+
+### Added
+- `recompute_scope="mlp"`: with `activation_checkpointing`, recompute only the MoE/FFN of each block (`"block"` recomputes the whole block, as before). At 32k on the real corpus it saves 1.4 GB (5%) and loses 11% speed, against 6.9 GB (25%) and 19% for whole blocks, so it suits runs that need only a little more memory.
+- `mla_upproj_recompute`: keep only the MLA latent and rebuild K and V per attention block. Exact, but it saved no memory (24.47 -> 24.29 GiB at 64k) and costs 6.7% speed, so it stays off.
+- `lookahead_input_detach` (the lookahead router reads a gradient-free hidden state) and `temporal_grad_span` (cut the temporal-locality gradient chain every n positions). Both are exact in the forward pass. On the en-zh data they were neutral (validation CE within 1e-3), so they stay off.
+- Web UI: every long-context field (`attention_block_size`, `ce_chunk_size`, `moe_chunk_tokens`, `recompute_scope`, `recompute_policy`, `mla_upproj_recompute`, `document_masking`, `reset_router_at_documents`, `reset_lookahead_at_documents`, `lookahead_input_detach`, `temporal_grad_span`) with labels and help in all four languages.
+
+### Fixed
+- `reset_router_at_documents` and `document_masking` raised "Attempting to eval an array without a primitive" inside the compiled LongRun step, because the document ids were built with NumPy from a traced array. `chronos.expert_clusters.segment_ids_mx` builds them in MLX, and a regression test runs the option through `mx.compile`.
+- `reset_lookahead_at_documents="false"` (the string the web UI sends) was read as true. Strings are now parsed; the empty string follows the router switch.
+
+### Measured on the real en-zh corpus (production 50M config, random init, M5 Max, bf16)
+`long` = `attention_block_size=1024`, `ce_chunk_size=1024`, `moe_chunk_tokens=8192`. Validation CE differs by at most 3e-3 between arms.
+
+| Context | Default: tokens/s, peak | long: tokens/s, peak |
+|---|---|---|
+| 8k | 8.9k, 40.6 GiB | 19.5k, 7.5 GiB |
+| 16k | 1.9k, 141.5 GiB (over this machine's 128 GB) | 15.2k, 13.8 GiB |
+| 32k | not run | 8.8k, 27.0 GiB |
+| 64k | not run | 3.7k, 54.6 GiB |
+
+Larger attention blocks speed up one MLA layer (0.62 s -> 0.53 s at 64k with blocks of 8192) but slow the whole model, because the sliding-window layers get slower (9.9 s -> 16.6 s per step); 1024 stays the recommendation.
+
 ## [0.3.2a1] — 2026-11-09 (beta)
 
 This release lets the MLX trainer run long contexts on one Mac. Every new switch is off by default, so existing configs, checkpoints and recipes behave as before.

chronos/__init__.py

@@ -1,5 +1,5 @@
 # Project Chronos: On-device low-latency lookahead dual-layer MoE inference
-__version__ = "0.3.2a1"
+__version__ = "0.3.2a2"
 
 # Register with HuggingFace AutoXXX factories so
 # AutoModelForCausalLM.from_pretrained("path/to/chronos") works out of the box.

chronos/config_schema.py

@@ -46,7 +46,7 @@
     "recommended_resident_experts", "offload_working_set_experts", "storage_format",
     "cluster_manifest_path", "training_paging", "optimizer_paging",
     "optimizer_page_dir", "page_size_mb", "activation_checkpointing", "ce_chunk_size", "attention_block_size",
-    "reset_router_at_documents", "document_masking", "recompute_policy", "moe_chunk_tokens",
+    "reset_router_at_documents", "document_masking", "recompute_policy", "recompute_scope", "lookahead_input_detach", "temporal_grad_span", "moe_chunk_tokens", "mla_upproj_recompute",
     "reset_lookahead_at_documents", "rope_train_scaling",
     "modalities", "auto_detect_modalities", "image_size", "image_patch_size",
     "max_image_tokens", "audio_sample_rate", "max_audio_seconds", "max_audio_tokens",

chronos/expert_clusters.py

@@ -226,6 +226,17 @@ def segment_ids(input_ids, chunk: int = 0, boundary_id: int = SEGMENT_BOUNDARY_I
     return np.cumsum(new, axis=1)
 
 
+def segment_ids_mx(input_ids, boundary_id: int = SEGMENT_BOUNDARY_ID):
+    """``segment_ids`` (chunk = 0) for an MLX array, traceable under ``mx.compile``.
+
+    Same rule: a separator token closes the document it ends, so the id of position t is the
+    number of separators strictly before t.  Returns an int32 ``[B, S]`` array."""
+    import mlx.core as mx
+    is_bnd = (input_ids == int(boundary_id)).astype(mx.int32)
+    before = mx.cumsum(is_bnd, axis=1) - is_bnd
+    return mx.stop_gradient(before)
+
+
 def segment_index(segments):
     """``(inverse [B*S] int32, n)``: dense batch-wide segment number of each position."""
     import numpy as np

chronos/mlx/attention.py

@@ -21,7 +21,7 @@
 from chronos.mlx.checkpointing import Dropout as TapeDropout   # replayed under recomputation
 
 from chronos.mlx.kv_quant import kv_read, kv_store
-from chronos.mlx.blockwise_attention import blockwise_attention, use_blockwise
+from chronos.mlx.blockwise_attention import blockwise_attention, blockwise_mla_attention, use_blockwise
 
 
 def _rope_freqs(
@@ -308,16 +308,22 @@ def __call__(
         if self._use_absorbed(S, T_full):
             out = self._absorbed(q_nope, q_rope, c_kv_full, k_rope, scale, allowed)
         else:
-            k_nope = self.k_nope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.nope_dim)
-            v = self.v_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.head_dim)
-            q = mx.concatenate([q_nope, q_rope], axis=-1).transpose(0, 2, 1, 3)   # [B, nH, S, d]
-            k = mx.concatenate([k_nope, k_rope], axis=-1).transpose(0, 2, 1, 3)   # [B, nKV, T, d]
             blk = use_blockwise(self._block_cfg, S, T_full, self.training) if mask is None else 0
-            if blk:
-                out = blockwise_attention(q, k, v.transpose(0, 2, 1, 3), scale=scale, window=None, block=blk,
-                                     segments=getattr(self, "_segments", None), past=T_full - S)
+            seg = getattr(self, "_segments", None)
+            q = mx.concatenate([q_nope, q_rope], axis=-1).transpose(0, 2, 1, 3)   # [B, nH, S, d]
+            if blk and self.training and getattr(self._block_cfg, "mla_upproj_recompute", False):
+                # Keep c_kv only: K and V are rebuilt per block (forward and backward).
+                out = blockwise_mla_attention(q, c_kv_full, k_rope, self.k_nope_proj.weight, self.v_proj.weight,
+                  scale=scale, block=blk, segments=seg, past=T_full - S)
             else:
-                out = self._attend(q, k, v.transpose(0, 2, 1, 3), scale, allowed)
+                k_nope = self.k_nope_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.nope_dim)
+                v = self.v_proj(c_kv_full).reshape(B, T_full, self.n_kv_heads, self.head_dim)
+                k = mx.concatenate([k_nope, k_rope], axis=-1).transpose(0, 2, 1, 3)   # [B, nKV, T, d]
+                if blk:
+                    out = blockwise_attention(q, k, v.transpose(0, 2, 1, 3), scale=scale, window=None, block=blk,
+                         segments=seg, past=T_full - S)
+                else:
+                    out = self._attend(q, k, v.transpose(0, 2, 1, 3), scale, allowed)
             out = out.transpose(0, 2, 1, 3)
         out = self.resid_dropout(self.o_proj(out.reshape(B, S, -1)))
         if return_shared_kv:

chronos/mlx/blockwise_attention.py

@@ -80,6 +80,51 @@ def blockwise_attention(q, k, v, *, scale, window=None, block=1024, segments=Non
     return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=2)
 
 
+def _mla_block(qb, cb, rb, w_k, w_v, mask, scale):
+    """One query block of MLA attention, K and V rebuilt from the latent slice.
+
+    ``cb [B, t, L]`` latent, ``rb [B, t, nKV, r]`` rotated rope keys,
+    ``w_k [nKV*nope, L]`` and ``w_v [nKV*d, L]`` the up-projections."""
+    B, t, _ = cb.shape
+    n_kv, r = rb.shape[2], rb.shape[3]
+    k_nope = (cb @ w_k.T).reshape(B, t, n_kv, -1)
+    v = (cb @ w_v.T).reshape(B, t, n_kv, -1)
+    k = mx.concatenate([k_nope, rb], axis=-1).transpose(0, 2, 1, 3)
+    return mx.fast.scaled_dot_product_attention(qb, k, v.transpose(0, 2, 1, 3), scale=scale, mask=mask)
+
+
+_mla_block_ckpt = mx.checkpoint(_mla_block)
+
+
+def blockwise_mla_attention(q, c_kv, k_rope, w_k, w_v, *, scale, block=1024, segments=None, past=0):
+    """Full-causal MLA attention that never keeps the expanded K and V.
+
+    ``q [B, nH, S, d]``; ``c_kv [B, T, L]``; ``k_rope [B, T, nKV, r]`` (already
+    rotated); ``w_k``/``w_v`` the ``k_nope_proj`` / ``v_proj`` weights.  Every
+    block runs under ``mx.checkpoint`` with the latent slice as its input, so
+    the forward pass keeps ``c_kv`` (L wide) instead of K and V (nKV*d wide
+    each), and the backward pass rebuilds them one block at a time.  The same
+    keys and softmax set as :func:`blockwise_attention`, so the output is
+    equal up to summation order.  Returns ``[B, nH, S, d]``."""
+    S, T = q.shape[2], c_kv.shape[1]
+    block = max(1, int(block))
+    seg = None if segments is None else mx.stop_gradient(segments)
+    outs, prev = [], None
+    for start in range(0, S, block):
+        end = min(S, start + block)
+        qa0, qa1 = past + start, past + end
+        k1 = min(T, qa1)
+        seg_q = None if seg is None else seg[:, qa0:qa1]
+        seg_k = None if seg is None else seg[:, :k1]
+        mask = _mask_for(qa0, qa1, 0, k1, None, seg_q, seg_k)
+        qb = q[:, :, start:end]
+        if prev is not None:
+            qb = mx.depends(qb, prev)
+        prev = _mla_block_ckpt(qb, c_kv[:, :k1], k_rope[:, :k1], w_k, w_v, mask, scale)
+        outs.append(prev)
+    return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=2)
+
+
 def use_blockwise(config, S, T, training):
     """Rows per attention block from ``config.attention_block_size`` (0 = off).
 

chronos/mlx/model.py

@@ -96,6 +96,7 @@ def __init__(self, layer_id: int, config):
             self.self_attn = SlidingWindowAttentionMLX(config)
             self.self_attn.window_size = None     # dense causal GQA (a 2**63 window overflowed int32)
 
+        object.__setattr__(self, "_cfg", config)
         if getattr(config, 'use_moe', True):
             self.mlp = ChronosMLXMOE(config)
         else:
@@ -121,10 +122,16 @@ def __call__(
                 h, new_cache = out
             x = x + h
         mlp_in = self.post_norm(x)
+        # recompute_scope="mlp": only the FFN/MoE is a recomputed region.
+        sub = (self.training and cache is None and getattr(self._cfg, "recompute_scope", "block") == "mlp"
+            and bool(getattr(self._cfg, "activation_checkpointing", False)))
         if isinstance(self.mlp, ChronosMLXMOE):
-            x = x + self.mlp(mlp_in, available_expert_mask=available_expert_mask)
+            if sub:
+                x = x + call_block(self.mlp, mlp_in, available_expert_mask=available_expert_mask)
+            else:
+                x = x + self.mlp(mlp_in, available_expert_mask=available_expert_mask)
         else:
-            x = x + self.mlp(mlp_in)
+            x = x + call_block(self.mlp, mlp_in) if sub else x + self.mlp(mlp_in)
         if return_shared_kv:
             return x, new_cache, shared
         return x, new_cache
@@ -232,7 +239,8 @@ def _forward(self, input_ids, cache, available_expert_masks, logits_dtype, depth
         def run(i, hidden, slot=None, **kw):
             avail = available_expert_masks[i] if available_expert_masks is not None else None
             c = cache[i if slot is None else slot] if "kv_override" not in kw else None
-            if recompute and c is None and _recompute_block(cfg, i, self.layers[i]):
+            if (recompute and c is None and getattr(cfg, "recompute_scope", "block") == "block"
+                and _recompute_block(cfg, i, self.layers[i])):
                 return call_block(self.layers[i], hidden, mask=mask, cache=None, available_expert_mask=avail, **kw)
             return self.layers[i](hidden, mask=mask, cache=c, available_expert_mask=avail, **kw)
 
@@ -257,7 +265,8 @@ def run(i, hidden, slot=None, **kw):
             for i in range(cfg.prelude_layers):
                 x, new_caches[i] = run(i, x)
             # Lookahead observes the post-prelude state (Torch parity).
-            lookahead_probs = self.lookahead_router(x)
+            lookahead_probs = self.lookahead_router(
+                mx.stop_gradient(x) if getattr(cfg, "lookahead_input_detach", False) else x)
             # MLX always returns a cache, so per_iteration (one slot per
             # iteration) never exits early; shared_first/first_only may.
             allow_exit = mode != "per_iteration"
@@ -321,7 +330,8 @@ def run(i, hidden, slot=None, **kw):
                 x, candidate = run(i, x)
                 new_caches.append(candidate)
                 if i == 0:
-                    lookahead_probs = self.lookahead_router(x)
+                    lookahead_probs = self.lookahead_router(
+                        mx.stop_gradient(x) if getattr(cfg, "lookahead_input_detach", False) else x)
 
         x = self.norm(x)
         if return_hidden:

chronos/mlx/training/trainer.py

@@ -162,10 +162,26 @@ def _group_mask(token_mask: mx.array | None, groups: int) -> mx.array | None:
     return mx.repeat(token_mask, groups, axis=0)
 
 
+_TEMPORAL_SPAN: list = [0]   # temporal_grad_span of the loss being built (0 = full chain)
+
+
+def _span_detach(prev: mx.array) -> mx.array:
+    """The older side of every adjacent pair, stop-gradient at span boundaries.
+
+    With ``temporal_grad_span`` = n the sequence is cut in blocks of n positions: inside a
+    block the older token keeps its gradient; the pair that straddles a block edge detaches
+    it, so no gradient chain runs across more than n steps."""
+    n = int(_TEMPORAL_SPAN[0] or 0)
+    if n <= 0 or prev.shape[1] <= 1:
+        return prev
+    edge = (mx.arange(prev.shape[1]) % n) == (n - 1)
+    return mx.where(edge[None, :, None], mx.stop_gradient(prev), prev)
+
+
 def _temporal_locality_loss(router_mean: mx.array, token_mask: mx.array | None = None) -> mx.array:
     if router_mean.shape[1] <= 1:
         return mx.zeros(())
-    diff = router_mean[:, 1:, :] - router_mean[:, :-1, :]
+    diff = router_mean[:, 1:, :] - _span_detach(router_mean[:, :-1, :])
     return _masked_mean((diff * diff).sum(axis=-1), _pair_mask(token_mask))
 
 
@@ -571,18 +587,20 @@ def _chronos_regularized_loss(model, base_loss, lookahead_probs, config, router_
     reset_look = getattr(config, "reset_lookahead_at_documents", None)
     reset_look = reset_router if reset_look is None else bool(reset_look)
     if input_ids is not None and (reset_router or reset_look):
-        from chronos.expert_clusters import segment_ids
-        docs = mx.array(segment_ids(input_ids))
+        from chronos.expert_clusters import segment_ids_mx
+        docs = segment_ids_mx(input_ids)
         segments = docs if reset_router else None
         look_segments = docs if reset_look else None
     _PAIR_SEGMENTS[0] = segments
     _LOOKAHEAD_SEGMENTS[0] = look_segments
+    _TEMPORAL_SPAN[0] = int(getattr(config, "temporal_grad_span", 0) or 0)
     try:
         return _regularized_terms(model, base_loss, lookahead_probs, config, router_ref, metrics,
                                   draft_kl, token_mask, input_ids)
     finally:
         _PAIR_SEGMENTS[0] = None
         _LOOKAHEAD_SEGMENTS[0] = None
+        _TEMPORAL_SPAN[0] = 0
 
 
 def _regularized_terms(model, base_loss, lookahead_probs, config, router_ref, metrics, draft_kl,
@@ -648,10 +666,10 @@ def lm_ce(model, input_ids, labels, *, aligned: bool = False, need_logits: bool
     from chronos.mlx.chunked_ce import ce_chunk_for, chunked_lm_loss, chunked_shifted_lm_loss
     chunk = 0 if need_logits else ce_chunk_for(model.config, int(input_ids.shape[0]) * int(input_ids.shape[1]))
     if getattr(model.config, "document_masking", False) and "segments" not in forward_kw:
-        from chronos.expert_clusters import segment_ids
+        from chronos.expert_clusters import segment_ids_mx
         # Varlen packing: queries attend only inside their own document; the
         # separator token closes the document it ends (segment_ids rule).
-        forward_kw["segments"] = mx.array(segment_ids(input_ids))
+        forward_kw["segments"] = segment_ids_mx(input_ids)
     if chunk:
         hidden, lookahead, _cache = model(input_ids, return_hidden=True, **forward_kw)
         loss_fn = chunked_lm_loss if aligned else chunked_shifted_lm_loss

chronos/model/config.py

@@ -111,6 +111,15 @@ def __init__(self, **kwargs):
         self.recompute_policy: str = str(kwargs.pop("recompute_policy", "all") or "all").lower()
         if self.recompute_policy not in ("all", "mla", "swa", "alternate"):
             raise ValueError(f"unknown recompute_policy {self.recompute_policy!r}")
+        # What activation_checkpointing recomputes: "block" (a whole block) or "mlp" (only the
+        # MoE/FFN of each block: most saved activations, none of the quadratic attention).
+        # MLX only; recompute_policy applies to scope=block.
+        self.recompute_scope: str = str(kwargs.pop("recompute_scope", "block") or "block").lower()
+        if self.recompute_scope not in ("block", "mlp"):
+            raise ValueError(f"unknown recompute_scope {self.recompute_scope!r}")
+        # MLA training with attention_block_size: keep only the latent c_kv and rebuild K and V
+        # per attention block, in both passes (no [B, T, nKV, d] K/V tensors).  Exact.
+        self.mla_upproj_recompute: bool = bool(kwargs.pop("mla_upproj_recompute", False))
         # MoE sparse dispatch in chunks of this many tokens, each recomputed in the
         # backward pass, so [N*K, I] expert activations never exist for a whole row
         # (0 = one dispatch, the historical path).
@@ -119,7 +128,17 @@ def __init__(self, **kwargs):
         # reset_router_at_documents; True/False decouples it from the temporal and
         # router-locality terms.
         _lrd = kwargs.pop("reset_lookahead_at_documents", None)
+        if isinstance(_lrd, str):      # web UI tri-state: "" follows the router switch
+            _lrd = None if _lrd.strip() == "" else _lrd.strip().lower() in ("1", "true", "yes", "on")
         self.reset_lookahead_at_documents = None if _lrd is None else bool(_lrd)
+        # Two-stage lookahead decoupling.  True: the LookaheadRouter reads a stop-gradient copy of
+        # the hidden state, so its losses train only the router and add no gradient path
+        # through the backbone (fine-tune the router on a frozen trunk).
+        self.lookahead_input_detach: bool = bool(kwargs.pop("lookahead_input_detach", False))
+        # > 0: the temporal / locality terms keep gradient only for the newer token of a pair
+        # every this-many positions (blocks of that length, the older side is detached), so
+        # no gradient chain crosses a whole row.  0 = full chain (historical).
+        self.temporal_grad_span: int = int(kwargs.pop("temporal_grad_span", 0) or 0)
         # Long-context RoPE for training: {"type": "yarn"|"linear"|"ntk", "factor": f,
         # "original_max_position_embeddings": n, ...}; applied after the base init.
         _rts = kwargs.pop("rope_train_scaling", None)