Publish Modilify Mk2 Preview step 1250 schema25

.gitattributes

@@ -1,3 +1,2 @@
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
-assets/01-LOGO.jpg filter=lfs diff=lfs merge=lfs -text

NOTICE.md

@@ -3,12 +3,13 @@
 Copyright 2026 Modilify
 
 This distribution is derived from `google/diffusiongemma-26B-A4B-it`, published
-by Google DeepMind under Apache License 2.0. It retains the upstream multimodal
-encoder, vision tower, vision projection, tokenizer, processor assets, and base
-language-model parameters. Modilify added dual-timescale latent deliberation
-and an excess-entropy confidence-and-entropy commit policy. The inference graph
-restores the Gemma 4 vision tower on the encoder so text, image, and video
-prompts share one rolling-canvas decoder.
+by Google DeepMind under Apache License 2.0. This PyTorch distribution retains
+the text encoder/decoder backbone, tokenizer, and chat template. It migrates
+the step 1250 schema25 weights from the native MLX release, preserving unfused
+LoRA adapters and dual-timescale GDN2 trajectory memory. The runtime follows
+the Modilify-Mk reference implementation and its confidence-and-entropy prefix
+commit policy. It contains no vision tower, vision projection, or image/video
+inference path.
 
 The Modilify Open Model License 1.0 applies to Modilify's distribution and
 original contributions. It does not erase, narrow, or replace rights and notices

README.md

@@ -3,278 +3,60 @@ license: other
 license_name: modilify-open-model-license-1.0
 license_link: LICENSE
 library_name: transformers
-pipeline_tag: image-text-to-text
+pipeline_tag: text-generation
 tags:
   - diffusion
-  - multimodal
-  - image-text-to-text
   - mixture-of-experts
   - trust-remote-code
+  - safetensors
 ---
 
 ![LOGO](assets/01-LOGO.jpg)
 
-# Modilify Mk2 Preview · 26B-A5B
+# Modilify Mk2 Preview
 
-**Refine a whole canvas. Think in latent space. Carry memory forward.**
+PyTorch text inference for the **step 1250, schema25** checkpoint migrated
+from `Modilify-Mk2-preview-mlx`. The model uses a shared DiffusionGemma text
+encoder/decoder, a rolling 256-token canvas, and dual-timescale GDN2 memory.
+The release contains 32 Safetensors shards totaling **48.23 GiB**.
 
-Modilify Mk2 is a **26B-A5B multimodal block-diffusion model** that brings parallel token refinement, latent deliberation, and persistent trajectory memory into one generation loop. A rolling 256-token canvas gives the model room to revise upcoming text together. A dedicated latent Transformer turns the history of those revisions into context for the next step. As output advances, a learned memory writer carries information across canvas boundaries.
+Dense and expert LoRA adapters remain unfused. Model loading preserves the
+26 FP32 GDN2/norm parameters alongside BF16 weights. No vision tower is included.
 
-The central idea is simple: **give each answer a workspace, a memory, and a variable compute budget.** Mk2 can refine its internal state over multiple passes before committing text, and release multiple tokens together when its confidence-and-entropy policy permits. Deliberation happens in continuous hidden states, without requiring every internal update to become a visible reasoning token.
+## Inference
 
-Built on DiffusionGemma, Mk2 combines sparse expert routing with a dual-timescale latent architecture. Text, images, and sampled video frames feed the same generation path.
-
-## What makes Mk2 different
-
-| Architecture choice | What it enables |
-| --- | --- |
-| **Parallel block diffusion** | Revise a 256-token canvas jointly and commit a variable-length prefix, allowing multiple output tokens per denoising pass. |
-| **Trajectory-aware latent deliberation** | Condition the next revision on how hidden states have evolved, including their changes, acceleration, and residuals relative to the latest state. |
-| **Memory at two timescales** | Rebuild working state every pass while retaining persistent slots across rolling-window shifts within a generation. |
-| **Memory inside the decoder** | Feed working and persistent memory directly into full-attention decoder layers so both can influence token refinement. |
-| **Adaptive commitment** | Use confidence and entropy to decide how much text to release, with explicit budgets for continued refinement and forced progress. |
-| **Sparse multimodal foundation** | Select 8 of 128 experts per token and bring text, image, and video context into a shared decoder. |
-
-## Inside the generation loop
-
-### A canvas built for revision
-
-Mk2 maintains a rolling canvas of 256 candidate tokens. Each denoising pass updates the candidates using the encoded prompt, the current canvas, and latent context. Attention lets positions within the canvas inform one another before the output prefix is finalized. After a commit, the canvas shifts forward and opens space for new candidates.
-
-This creates two useful degrees of freedom: **how many times to refine** and **how many tokens to release**. A pass can commit a longer prefix when the policy permits, or spend additional computation refining an uncertain frontier.
-
-### Deliberation that reads its own trajectory
-
-A 4-layer latent Transformer, 2,816 dimensions wide, builds working context before each decoder denoising pass. It reads the current canvas alongside three complementary sources of state:
-
-- **Per-token trajectory history:** 16 recent frames represented through four views—hidden state, first difference, second difference, and residual from the latest state—projected to rank 1,024. These give the processor access to the direction and stability of recent revisions. History follows surviving canvas tokens; new positions start empty.
-- **Denoising tape:** 16 pooled probes per frame summarize canvas activity in a time-indexed ring. The tape preserves recent step-level context as token positions move through the window.
-- **Persistent memory:** 256 slots of 2,816 dimensions carry learned summaries across commits within the generation.
-
-The resulting working context enters the decoder through its self-conditioning bridge and working-memory bus. Full-attention layers also read a separate persistent-memory bus. **The refinement history becomes an input to the next refinement.**
-
-### Fast working state, lasting commit memory
-
-Working state is recomputed on every denoising pass. Persistent slots update only when tokens commit, using a Transformer writer with a separate gate for each slot. That writer draws on the committed region's trajectory, working state, and final decoder representation.
-
-This separates rapid revision from memory consolidation: the canvas can keep changing while persistent memory stays stable between commits. When the window advances, the memory slots remain available to later tokens.
-
-### Compute that follows the commit frontier
-
-Mk2 combines proposal confidence with an excess-entropy penalty and selects the longest prefix whose cumulative failure score stays below the configured budget. A tighter budget requires stronger evidence before normal commitment; a looser budget admits longer prefixes for the same scores. Stagnation handling and a pondering watchdog bound continued refinement.
-
-Temperature and commit budget are exposed directly at inference time, making the generation policy adjustable per request. These scores govern commitment; they are not calibrated guarantees of factual correctness.
-
-### One generation path for text, images, and video
-
-The Gemma 4 vision tower supplies visual features through the DiffusionGemma multimodal encoder. Text, image, and sampled video-frame inputs are encoded into the prefix KV cache that conditions the decoder. The rolling text canvas then uses the same latent deliberation and memory loop across all three input modalities.
-
-## Model Summary
-
-| | |
-| --- | ---: |
-| Architecture | Mixture-of-Experts block diffusion + dual-timescale latent Transformer |
-| Model Size | 26B-A5B |
-| Vision Encoder | Gemma 4 Vision |
-| Layers | 30 |
-| Number of Experts | 128 |
-| Selected Experts per Token | 8 |
-| Vocabulary Size | 262,144 |
-| Configured Context Length | 262,144 tokens |
-| Sliding Window | 1024 |
-| Canvas Length | 256 |
-| Latent Width | 2,816 |
-| Latent Transformer | 4 layers, 16 attention heads |
-| Persistent Memory | 256 slots × 2,816 dimensions |
-| Memory Writer | 2-layer commit-sequence Transformer + per-slot gated writer |
-| Trajectory History | 16 frames, 4 views, rank 1,024 |
-| Denoise Tape | 16 probes per frame |
-| Modality | Text, Image, Video |
-| Preview checkpoint | Training step 900 |
-| Adaptation tokens | ~12.4 million |
-
-## Preview release
-
-This first public preview includes **merged BF16 weights, inference code, processor, and tokenizer**. The text backbone and latent stack are exported from training step 900, after approximately 12.4 million adaptation tokens. The Gemma 4 vision tower is restored from DiffusionGemma.
-
-The release makes the architecture available for hands-on exploration and evaluation. Comprehensive benchmark results are not included; measured speed, reasoning quality, and multimodal reliability remain to be established for specific workloads.
-
-## Getting Started
-
-Transformers 5.14.1 is the minimum supported version.
-
-```shell
-pip install -U transformers torch accelerate
-```
-
-### Text generation
+Use Transformers **5.14.1**, PyTorch, Accelerate, and Safetensors. Run from
+this directory or replace `model_path` with its location. Inference needs
+additional memory beyond the weights.
 
 ```python
 import torch
-from transformers import AutoModelForMultimodalLM, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-model_id = "modilify/Modilify-Mk2-preview"
-processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForMultimodalLM.from_pretrained(
-    model_id,
+model_path = "."
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
     trust_remote_code=True,
     dtype=torch.bfloat16,
     device_map="auto",
-)
-
-messages = [{"role": "user", "content": "Explain why the sky is blue."}]
-inputs = processor.apply_chat_template(
-    messages,
+    attn_implementation="sdpa",
+).eval()
+inputs = tokenizer.apply_chat_template(
+    [{"role": "user", "content": "Explain why the sky is blue."}],
     tokenize=True,
     add_generation_prompt=True,
     enable_thinking=False,
     return_dict=True,
     return_tensors="pt",
 ).to(model.device)
-
-output = model.generate(
-    **inputs,
-    max_new_tokens=256,
-    denoise_temperature=0.8,
-    commit_failure_budget=0.2,
-)
-new_tokens = output.sequences[:, inputs["input_ids"].shape[1]:]
-print(processor.batch_decode(new_tokens, skip_special_tokens=False)[0])
-```
-
-### Image input
-
-```python
-from PIL import Image
-
-image = Image.open("example.jpg").convert("RGB")
-messages = [{
-    "role": "user",
-    "content": [
-        {"type": "image", "image": image},
-        {"type": "text", "text": "Describe the image and identify uncertainty."},
-    ],
-}]
-inputs = processor.apply_chat_template(
-    messages,
-    tokenize=True,
-    add_generation_prompt=True,
-    enable_thinking=True,
-    return_dict=True,
-    return_tensors="pt",
-).to(model.device)
-output = model.generate(**inputs, max_new_tokens=256)
-```
-
-### Video-frame input
-
-The processor represents video as a sampled sequence of frames. The following example uses PyAV to decode a short local clip and samples at most 32 RGB frames.
-
-```python
-import av
-from PIL import Image
-
-container = av.open("short_clip.mp4")
-decoded = [Image.fromarray(frame.to_rgb().to_ndarray()) for frame in container.decode(video=0)]
-stride = max(1, len(decoded) // 32)
-frames = decoded[::stride][:32]
-
-messages = [{
-    "role": "user",
-    "content": [
-        {"type": "video", "video": frames},
-        {"type": "text", "text": "Summarize the main visual events in order."},
-    ],
-}]
-inputs = processor.apply_chat_template(
-    messages,
-    tokenize=True,
-    add_generation_prompt=True,
-    enable_thinking=True,
-    return_dict=True,
-    return_tensors="pt",
-).to(model.device)
-output = model.generate(**inputs, max_new_tokens=256)
-```
-
-## Thinking mode
-
-Latent deliberation runs inside the generation loop regardless of the chat template's thinking flag. The flag controls the prompt's request for a textual thought channel:
-
-- `enable_thinking=True` inserts a system turn that contains `<|think|>` and still ends the prompt at `<|turn>model`.
-- `enable_thinking=False` does **not** inject an empty thought channel. The prompt ends at `<|turn>model`.
-
-The model may still open `<|channel>thought` on its own when thinking is disabled. The template flag does not guarantee suppression of generated thought-channel text. Applications should handle that channel explicitly before displaying an answer.
-
-## Configurable inference
-
-The two primary knobs are sampling temperature and the prefix failure budget. Both default to the values used in Mk2 training (`0.8` and `0.2`) and can be changed per call or on the config object.
-
-| Parameter | Default | Meaning |
-| --- | ---: | --- |
-| `denoise_temperature` | 0.8 | Sampling temperature for every canvas step |
-| `commit_failure_budget` | 0.2 | Cumulative prefix risk limit for normal commits |
-| `jump_failure_budget` | 2.0 | Cumulative risk limit for forced jumps |
-| `jump_on_no_progress_after` | 12 | Stagnation steps before a forced jump |
-| `max_ponder_steps` | 64 | Watchdog multiplier per requested token |
-| `min_trajectory_progress` | 0.005 | Minimum fused-risk improvement counted as progress |
-| `canvas_length` | 256 | Rolling diffusion canvas length |
-| `repetition_penalty` | 1.0 | Transformers-style repetition penalty |
-| `turn_end_token_id` | 106 | Gemma turn terminator |
-
-Call-site override:
-
-```python
-output = model.generate(
-    **inputs,
-    max_new_tokens=256,
-    denoise_temperature=0.4,
-    commit_failure_budget=0.05,
-)
-```
-
-Load-time override:
-
-```python
-from transformers import AutoConfig
-
-config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
-config.denoise_temperature = 0.4
-config.commit_failure_budget = 0.05
-model = AutoModelForMultimodalLM.from_pretrained(
-    model_id,
-    config=config,
-    trust_remote_code=True,
-    dtype=torch.bfloat16,
-    device_map="auto",
-)
-```
-
-A tighter commit budget allows fewer tokens for the same confidence-and-entropy scores; a looser budget allows more. Temperature changes the sampling distribution and also affects those scores, so its effect on throughput depends on the prompt and generation trajectory. Measure latency and output quality together when tuning these controls.
-
-Generation supports left-padded batches with independent stopping. Batch prompts of similar lengths together for the best throughput. Streaming and caller-supplied KV caches remain limited to batch size 1.
-
-## Evaluation status, limitations, and risks
-
-This preview does not include a complete accuracy, robustness, calibration, fairness, or safety evaluation. Architectural features describe how Mk2 generates; they do not establish benchmark superiority or fitness for a particular deployment. The configured context limit is not a validated long-context quality result.
-
-The model can hallucinate facts, citations, visual details, or temporal relationships; reproduce bias, unsafe content, personal information, or copyrighted material; and consume substantial time and memory during long iterative generation. Confidence-based commits control compute. They do not certify that a prefix is true. Visual performance can degrade with poor resolution, motion, occlusion, unusual aspect ratios, or domain shift.
-
-Evaluate the exact deployment on representative, adversarial, and out-of-distribution inputs. Use layered safeguards, monitoring, incident response, and qualified human review. Never delegate autonomous high-risk medical, legal, financial, employment, housing, education, critical-infrastructure, or safety decisions to the model.
-
-## License
-
-Released under the [Modilify Open Model License 1.0](LICENSE), subject to its responsible-use and derivative-impact terms. Upstream rights, attribution, Apache-2.0 text, and the impact-statement template are retained in [NOTICE.md](NOTICE.md).
-
-## Citation
-
-```bibtex
-@software{modilify_mk2_preview_2026,
-  title = {Modilify Mk2 Preview: 26B-A5B},
-  author = {Modilify},
-  year = {2026},
-  note = {A multimodal dual-timescale latent-deliberation derivative of DiffusionGemma}
-}
+output = model.generate(**inputs, max_new_tokens=128, seed=42)
+print(tokenizer.decode(output.sequences[0, inputs.input_ids.shape[1]:],
+                       skip_special_tokens=True))
 ```
 
-Also cite the upstream DiffusionGemma release as requested by Google DeepMind.
+For MPS, use `device_map={"": "mps"}`. Set `enable_thinking=True` in the chat
+template to request thinking. Generation uses temperature 0.8, top-k 40,
+min-p 0.05, target confidence 0.5, and failure budget 0.2. `max_denoising_steps`
+can bound generation. Static batches and continuous batching share the
+confidence-prefix commit policy; each request retains its own cache and RNG.

chat_template.jinja

@@ -1,6 +1,6 @@
 {#
-  Template: Modilify Canonical Chat Template
-  Author: Modilify
+  Template: Google Gemma 4 Canonical Chat Template
+  Author: Google Gemma Engineering Team
   Published: 2026-07-09
   Context: Fixed tool-calling loops, turn closures, and thinking content-ordering.
 #}
@@ -268,7 +268,7 @@
 
             {%- set ns_tr_out = namespace(flag=false) -%}
             {%- if message.get('tool_responses') -%}
-                {#- Legacy: tool_responses embedded on the assistant message -#}
+                {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
                 {%- for tool_response in message.get('tool_responses') -%}
                     {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
                     {%- set ns_tr_out.flag = true -%}
@@ -384,4 +384,4 @@
     {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%}
         {{- '<|channel>thought\n' -}}
     {%- endif -%}
-{%- endif -%}
+{%- endif -%}

commit_policy.py

@@ -1,13 +1,11 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
-"""Confidence-and-entropy commit policy for inference."""
+"""Confidence-and-entropy prefix commitment."""
 
 from __future__ import annotations
 
 from collections.abc import Sequence
 from dataclasses import dataclass
-import math
 
+import math
 import torch
 
 from .latent_deliberation import (
@@ -20,12 +18,32 @@ JUMP_FAILURE_BUDGET = 2.0
 FUSED_EPS = 1e-6
 
 
+def commit_target_confidence_bias(
+    target_confidence: float | None,
+    failure_budget: float = 0.2,
+    budget_safety_ratio: float = 0.85,
+) -> float:
+    if target_confidence is None or target_confidence <= 0.0 or target_confidence >= 1.0:
+        return 0.0
+    target_failure = min(budget_safety_ratio * failure_budget, 1.0 - target_confidence)
+    target_failure = max(target_failure, 1.0e-4)
+    target_conf = 1.0 - target_failure
+    logit_c = math.log(target_conf / (1.0 - target_conf))
+    logit_p = math.log(target_confidence / (1.0 - target_confidence))
+    return max(logit_c - logit_p, 0.0)
+
+
 def fused_commit_confidence(
     proposal_confidence: torch.Tensor,
     token_entropy: torch.Tensor,
     *,
-    vocab_size: int = 256000,
     eps: float = FUSED_EPS,
+    entropy_weight: float = 1.0,
+    confidence_power: float = 1.0,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    target_confidence: float | None = None,
+    failure_budget: float = 0.2,
 ) -> torch.Tensor:
     """Fuse proposal confidence with token entropy into effective commit confidence.
 
@@ -34,11 +52,9 @@ def fused_commit_confidence(
         p   = clamp(proposal_confidence, eps, 1 - eps)
         h2  = -p * log(p) - (1 - p) * log(1 - p)     # binary entropy of p
         excess = max(token_entropy - h2, 0)
-        base_fused = sigmoid(logit(p) - excess)
-        fused = base_fused.square()
+        base_fused = sigmoid(logit(p) + bias - entropy_weight * excess)
+        fused = base_fused ** confidence_power
 
-    When the token entropy equals the binary entropy implied by p, the fused
-    confidence equals p^2. Entropy *above* h2 pulls fused below p^2.
     """
 
     p = proposal_confidence.float().nan_to_num(0.5).clamp(min=eps, max=1.0 - eps)
@@ -48,12 +64,23 @@ def fused_commit_confidence(
     h2 = -p * torch.log(p) - (1.0 - p) * torch.log1p(-p)
 
     excess = (H - h2).clamp(min=0.0)
+    k_eff = None
+    if top_k is not None and top_k > 0:
+        k_eff = torch.full_like(p, float(top_k))
+    if min_p is not None and min_p > 0:
+        thresh = torch.clamp_min(float(min_p) * p, 1.0e-6)
+        k_min_p = torch.clamp_min((1.0 - p) / thresh, 1.0)
+        k_eff = k_min_p if k_eff is None else torch.minimum(k_eff, k_min_p)
+    if k_eff is not None:
+        max_excess = (1.0 - p) * torch.log(k_eff)
+        excess = torch.minimum(excess, max_excess)
 
     # logit(p) = log(p / (1-p)) = log(p) - log1p(-p)
     logit_p = torch.log(p) - torch.log1p(-p)
 
-    base_fused = torch.sigmoid(logit_p - excess)
-    fused = base_fused.square()
+    bias = commit_target_confidence_bias(target_confidence, failure_budget=failure_budget)
+    base_fused = torch.sigmoid(logit_p + bias - float(entropy_weight) * excess)
+    fused = base_fused.pow(float(confidence_power))
     return fused.clamp(min=eps, max=1.0 - eps)
 
 
@@ -71,7 +98,7 @@ def fused_commit_failure_rate(
 
 @dataclass(frozen=True)
 class CommitPolicyDecision:
-    """One inference transition from proposal to committed prefix."""
+    """Proposal-to-commit transition."""
 
     normal_lengths: torch.LongTensor
     commit_lengths: torch.LongTensor
@@ -184,7 +211,6 @@ def select_commit_lengths(
     stagnation_threshold: int,
     min_progress: float,
     max_ponder_steps: int | None = None,
-    jump_failure_budget: float | None = None,
     valid_mask: torch.BoolTensor | None = None,
 ) -> CommitPolicyDecision:
     """Use normal sampled commits and a fixed-budget greedy JUMP.
@@ -242,8 +268,6 @@ def select_commit_lengths(
         stagnation_steps,
         commit_lengths=normal,
         active_rows=active_rows,
-        progress_scores=progress,
-        min_progress=min_progress,
     )
     jump_rows = normal.eq(0) & active_rows & should_force_trajectory_jump(
         next_stagnation,
@@ -256,9 +280,7 @@ def select_commit_lengths(
     jump_commit = bounded_prefix_failure_commit_lengths(
         greedy_token_ids,
         jump_failure_rate,
-        failure_budget=(
-            JUMP_FAILURE_BUDGET if jump_failure_budget is None else float(jump_failure_budget)
-        ),
+        failure_budget=JUMP_FAILURE_BUDGET,
         remaining_lengths=remaining_lengths,
         stop_token_id=stop_token_id,
         valid_mask=valid_mask,
@@ -291,15 +313,3 @@ def select_commit_lengths(
         ponder_steps=next_ponder,
         stagnation_steps=next_stagnation,
     )
-
-
-__all__ = [
-    "CommitPolicyDecision",
-    "JUMP_FAILURE_BUDGET",
-    "bounded_prefix_failure_commit_lengths",
-    "first_committed_token_lengths",
-    "fused_commit_confidence",
-    "fused_commit_failure_rate",
-    "prefix_failure_commit_lengths",
-    "select_commit_lengths",
-]

config.json

@@ -1,54 +1,27 @@
 {
-  "architectures": [
-    "ModilifyMk2ForBlockDiffusion"
-  ],
-  "auto_map": {
-    "AutoConfig": "configuration_modilify_mk2.ModilifyMk2Config",
-    "AutoModel": "modeling_modilify_mk2.ModilifyMk2Model",
-    "AutoModelForCausalLM": "modeling_modilify_mk2.ModilifyMk2ForBlockDiffusion",
-    "AutoModelForMultimodalLM": "modeling_modilify_mk2.ModilifyMk2ForBlockDiffusion"
-  },
-  "boi_token_id": 255999,
-  "bos_token_id": 2,
   "canvas_length": 256,
-  "channel_end_token_id": 101,
+  "commit_confidence_power": 1.0,
+  "commit_entropy_weight": 1.0,
   "commit_failure_budget": 0.2,
+  "commit_min_p": 0.05,
   "commit_sequence_dim": 1024,
-  "commit_sequence_layers": 2,
-  "denoise_temperature": 0.8,
-  "dtype": "bfloat16",
-  "eoi_token_id": 258882,
-  "eos_token_id": [
-    1,
-    106
-  ],
-  "experience_roles": 3,
-  "image_token_id": 258880,
+  "commit_target_confidence": 0.5,
+  "commit_top_k": 40,
+  "eos_token_id": 1,
   "initializer_range": 0.02,
-  "jump_failure_budget": 2.0,
-  "jump_on_no_progress_after": 12,
   "kv_cache_bucket_size": 128,
   "latent_dim": 2816,
-  "latent_dropout": 0.0,
   "latent_ffn_dim": 7168,
   "latent_history_kv_rank": 1024,
-  "latent_history_length": 16,
-  "latent_history_views": 4,
   "latent_local_attention_window": 128,
-  "latent_memory_slots": 256,
   "latent_num_heads": 16,
   "latent_num_layers": 4,
-  "latent_tape_probes": 16,
+  "latent_tape_probes": 4,
   "latent_working_last_block_global": true,
-  "max_ponder_steps": 64,
-  "memory_scheme": "dual_timescale_transformer_trajectory_memory",
-  "min_trajectory_progress": 0.005,
+  "memory_architecture": "compact_gdn2_v2",
   "model_type": "modilify_mk2",
-  "pad_token_id": 0,
   "persistent_memory_bus": true,
-  "persistent_memory_write": "commit_only_transformer",
-  "repetition_penalty": 1.0,
-  "state_schema_version": 23,
+  "state_schema_version": 25,
   "terminal_token_ids": [
     106,
     50
@@ -122,56 +95,40 @@
     "sliding_window": 1024,
     "tie_word_embeddings": true,
     "top_k_experts": 8,
-    "use_bidirectional_attention": "vision",
+    "use_bidirectional_attention": null,
     "vocab_size": 262144
   },
   "tie_word_embeddings": true,
   "transformers_version": "5.14.1",
   "turn_end_token_id": 106,
-  "vision_config": {
-    "_name_or_path": "",
-    "architectures": null,
-    "attention_bias": false,
-    "attention_dropout": 0.0,
-    "chunk_size_feed_forward": 0,
-    "default_output_length": 280,
-    "dtype": "bfloat16",
-    "global_head_dim": 72,
-    "head_dim": 72,
-    "hidden_activation": "gelu_pytorch_tanh",
-    "hidden_size": 1152,
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "initializer_range": 0.02,
-    "intermediate_size": 4304,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "max_position_embeddings": 131072,
-    "model_type": "gemma4_vision",
-    "num_attention_heads": 16,
-    "num_hidden_layers": 27,
-    "num_key_value_heads": 16,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "patch_size": 16,
-    "pooling_kernel_size": 3,
-    "position_embedding_size": 10240,
-    "problem_type": null,
-    "return_dict": true,
-    "rms_norm_eps": 1e-06,
-    "rope_parameters": {
-      "rope_theta": 100.0,
-      "rope_type": "default"
-    },
-    "standardize": true,
-    "use_clipped_linears": false
-  },
   "vocab_chunk_size": 32768,
   "working_memory_bus": true,
-  "writer_slot_gate": "per_slot"
+  "architectures": [
+    "ModilifyMk2ForBlockDiffusion"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_modilify_mk2.ModilifyMk2Config",
+    "AutoModel": "modeling_modilify_mk2.ModilifyMk2Model",
+    "AutoModelForCausalLM": "modeling_modilify_mk2.ModilifyMk2ForBlockDiffusion"
+  },
+  "dtype": "bfloat16",
+  "lora_config": {
+    "r": 16,
+    "alpha": 16,
+    "target_modules": [
+      "q_proj",
+      "k_proj",
+      "v_proj",
+      "o_proj",
+      "gate_proj",
+      "up_proj",
+      "down_proj",
+      "proj"
+    ],
+    "expert_r": 8,
+    "expert_alpha": 8
+  },
+  "precision_policy": "gdn2_small_fp32_v1",
+  "global_step": 1250,
+  "base_model": "google/diffusiongemma-26B-A4B-it"
 }

configuration_modilify_mk2.py

@@ -1,323 +1,99 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
-"""Inference configuration for Modilify Mk2."""
-
+"""Text-only schema25 inference configuration."""
 from __future__ import annotations
-
-import math
 from collections.abc import Sequence
 from typing import Any
+from transformers import PreTrainedConfig
+from transformers.models.diffusion_gemma import DiffusionGemmaTextConfig
 
-from transformers.models.diffusion_gemma import (
-    DiffusionGemmaConfig,
-    DiffusionGemmaTextConfig,
-)
-
-
-MEMORY_SCHEME = "dual_timescale_transformer_trajectory_memory"
-HISTORY_VIEWS = 4
-EXPERIENCE_ROLES = 3
-COMMIT_SEQUENCE_LAYERS = 2
 DENOISE_TEMPERATURE = 0.8
-COMMIT_FAILURE_BUDGET = 0.2
-JUMP_FAILURE_BUDGET = 2.0
-VOCAB_CHUNK_SIZE = 32_768
-
-_DROPPED_TRAINING_FIELDS = frozenset(
-    {
-        "training_scheme",
-        "training_bptt_steps",
-        "training_prefix_cache",
-        "token_loss_weight",
-        "jump_token_loss_weight",
-        "confidence_calibration_loss_weight",
-        "denoise_improvement_loss_weight",
-        "commit_throughput_softness",
-        "commit_throughput_target_tpd",
-        "commit_throughput_loss_weight",
-        "terminal_stop_loss_weight",
-        "terminal_stop_target_probability",
-        "latent_working_bus_unfreeze_steps",
-        "latent_persistent_bus_unfreeze_steps",
-        "virtual_commit_chunk_sizes",
-        "virtual_commit_chunk_probs",
-        "decoder_checkpoint_policy",
-        "fused_entropy_weight",
-        "schema_version",
-        "sampler_entropy_bound",
-        "loss_transition_steps",
-        "initial_jump_token_loss_weight",
-        "initial_confidence_calibration_loss_weight",
-        "initial_denoise_improvement_loss_weight",
-        "readiness_loss_weight",
-        "commit_risk_budget",
-        "router_load_balance_weight",
-        "router_z_loss_weight",
-        "latent_refinement_steps",
-        "latent_memory_bus_unfreeze_steps",
-    }
-)
-
 
 class ModilifyMk2TextConfig(DiffusionGemmaTextConfig):
-    """Text configuration for the Modilify Mk2 decoder."""
-
     model_type = "modilify_mk2_text"
-
-
-class ModilifyMk2Config(DiffusionGemmaConfig):
-    """Multimodal inference configuration for dual-timescale Modilify Mk2.
-
-    Args:
-        text_config: DiffusionGemma text configuration or its serialized form.
-        vision_config: Gemma 4 vision configuration or its serialized form.
-        denoise_temperature: Sampling temperature used at every denoising step.
-        commit_failure_budget: Maximum cumulative failure risk for normal commits.
-        jump_failure_budget: Maximum cumulative failure risk for forced jumps.
-        canvas_length: Rolling diffusion canvas length.
-        latent_dim: Width of the working trajectory state.
-        latent_ffn_dim: Feed-forward width inside the latent Transformer.
-        latent_memory_slots: Number of persistent latent memory slots.
-        latent_num_layers: Number of latent Transformer blocks.
-        latent_num_heads: Number of latent attention heads.
-        latent_local_attention_window: Local token-attention radius.
-        latent_dropout: Latent Transformer dropout probability.
-        latent_history_length: Packed per-token trajectory history length.
-        latent_tape_probes: Denoise-time tape probes per frame.
-        jump_on_no_progress_after: Stagnation steps before a forced jump.
-        max_ponder_steps: Maximum denoising iterations per requested token.
-        min_trajectory_progress: Minimum fused-risk improvement counted as progress.
-        turn_end_token_id: Native Gemma turn terminator.
-        repetition_penalty: Transformers-style repetition penalty. ``1.0`` disables
-            it.
-        kwargs: Standard DiffusionGemma configuration values.
-    """
-
+    vocab_size: int = 262_144
+    hidden_size: int = 2816
+    intermediate_size: int = 2112
+    num_hidden_layers: int = 30
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    head_dim: int = 256
+    max_position_embeddings: int = 262_144
+    sliding_window: int = 1024
+    use_bidirectional_attention: str | None = None
+    num_global_key_value_heads: int | None = 2
+    global_head_dim: int = 512
+    num_experts: int | None = 128
+    top_k_experts: int | None = 8
+    moe_intermediate_size: int | None = 704
+
+class ModilifyMk2Config(PreTrainedConfig):
     model_type = "modilify_mk2"
-    sub_configs = {
-        "text_config": ModilifyMk2TextConfig,
-        **{
-            key: value
-            for key, value in DiffusionGemmaConfig.sub_configs.items()
-            if key != "text_config"
-        },
-    }
+    sub_configs = {"text_config": ModilifyMk2TextConfig}
 
     def __init__(
-        self,
-        text_config: (
-            ModilifyMk2TextConfig
-            | DiffusionGemmaTextConfig
-            | dict[str, Any]
-            | None
-        ) = None,
-        vision_config: Any | dict[str, Any] | None = None,
-        *,
-        denoise_temperature: float = DENOISE_TEMPERATURE,
-        commit_failure_budget: float = COMMIT_FAILURE_BUDGET,
-        jump_failure_budget: float = JUMP_FAILURE_BUDGET,
+        self, text_config: ModilifyMk2TextConfig | dict[str, Any] | None = None, *,
         canvas_length: int = 256,
         initializer_range: float = 0.02,
-        memory_scheme: str = MEMORY_SCHEME,
+        tie_word_embeddings: bool = True,
+        state_schema_version: int = 25,
+        memory_architecture: str = "compact_gdn2_v2",
         latent_dim: int = 2816,
         latent_ffn_dim: int = 7168,
-        latent_memory_slots: int = 256,
         latent_num_layers: int = 4,
         latent_num_heads: int = 16,
         latent_local_attention_window: int = 128,
-        latent_dropout: float = 0.0,
-        latent_history_length: int = 16,
-        latent_tape_probes: int = 16,
-        latent_history_views: int = HISTORY_VIEWS,
-        latent_history_kv_rank: int | None = None,
+        latent_tape_probes: int = 4,
+        latent_history_kv_rank: int = 1024,
         latent_working_last_block_global: bool = True,
         working_memory_bus: bool = True,
         persistent_memory_bus: bool = True,
-        persistent_memory_write: str = "commit_only_transformer",
-        experience_roles: int = EXPERIENCE_ROLES,
-        commit_sequence_layers: int = COMMIT_SEQUENCE_LAYERS,
-        commit_sequence_dim: int | None = None,
-        writer_slot_gate: str = "per_slot",
+        commit_sequence_dim: int = 1024,
         kv_cache_bucket_size: int = 128,
+        vocab_chunk_size: int = 32768,
         turn_end_token_id: int = 106,
-        terminal_token_ids: Sequence[int] | None = None,
-        channel_end_token_id: int = 101,
-        vocab_chunk_size: int = VOCAB_CHUNK_SIZE,
-        jump_on_no_progress_after: int = 12,
-        max_ponder_steps: int = 64,
-        min_trajectory_progress: float = 0.005,
-        repetition_penalty: float = 1.0,
-        state_schema_version: int = 23,
+        terminal_token_ids: Sequence[int] = (106, 50),
+        eos_token_id: int = 1,
+        commit_failure_budget: float = 0.2,
+        commit_top_k: int | None = 40,
+        commit_min_p: float | None = 0.05,
+        commit_target_confidence: float | None = 0.5,
+        commit_entropy_weight: float = 1.0,
+        commit_confidence_power: float = 1.0,
         **kwargs: Any,
     ) -> None:
-        kwargs.pop("model_type", None)
-        for name in _DROPPED_TRAINING_FIELDS:
-            kwargs.pop(name, None)
-        if isinstance(text_config, DiffusionGemmaTextConfig):
-            text_payload = text_config.to_dict()
-            text_payload.pop("model_type", None)
-            if text_payload.get("use_bidirectional_attention") in (None, False):
-                text_payload["use_bidirectional_attention"] = "vision"
-            text_config = ModilifyMk2TextConfig(**text_payload)
+        if state_schema_version != 25 or memory_architecture != "compact_gdn2_v2":
+            raise ValueError("This release requires schema25 compact_gdn2_v2 weights.")
+        if text_config is None:
+            text_config = ModilifyMk2TextConfig()
         elif isinstance(text_config, dict):
-            text_payload = dict(text_config)
-            text_payload.pop("model_type", None)
-            if text_payload.get("use_bidirectional_attention") in (None, False):
-                text_payload["use_bidirectional_attention"] = "vision"
-            text_config = ModilifyMk2TextConfig(**text_payload)
-        elif text_config is None:
-            text_config = ModilifyMk2TextConfig(use_bidirectional_attention="vision")
-
-        self.denoise_temperature = float(denoise_temperature)
-        self.commit_failure_budget = float(commit_failure_budget)
-        self.jump_failure_budget = float(jump_failure_budget)
-        self.canvas_length = int(canvas_length)
-        self.memory_scheme = str(memory_scheme)
-        self.latent_dim = int(latent_dim)
-        self.latent_ffn_dim = int(latent_ffn_dim)
-        self.latent_memory_slots = int(latent_memory_slots)
-        self.latent_num_layers = int(latent_num_layers)
-        self.latent_num_heads = int(latent_num_heads)
-        self.latent_local_attention_window = int(latent_local_attention_window)
-        self.latent_dropout = float(latent_dropout)
-        self.latent_history_length = int(latent_history_length)
-        self.latent_tape_probes = int(latent_tape_probes)
-        self.latent_history_views = int(latent_history_views)
-        if latent_history_kv_rank is None:
-            rank = min(1024, self.latent_dim)
-            rank -= rank % max(self.latent_num_heads, 1)
-            if rank <= 0:
-                rank = self.latent_num_heads
-            self.latent_history_kv_rank = rank
-        else:
-            self.latent_history_kv_rank = int(latent_history_kv_rank)
-        self.latent_working_last_block_global = bool(latent_working_last_block_global)
-        self.working_memory_bus = bool(working_memory_bus)
-        self.persistent_memory_bus = bool(persistent_memory_bus)
-        self.persistent_memory_write = str(persistent_memory_write)
-        self.experience_roles = int(experience_roles)
-        self.commit_sequence_layers = int(commit_sequence_layers)
-        if commit_sequence_dim is None:
-            self.commit_sequence_dim = self.latent_history_kv_rank
-        else:
-            self.commit_sequence_dim = int(commit_sequence_dim)
-        self.writer_slot_gate = str(writer_slot_gate)
-        self.kv_cache_bucket_size = int(kv_cache_bucket_size)
-        self.turn_end_token_id = int(turn_end_token_id)
-        if terminal_token_ids is None:
-            self.terminal_token_ids = (int(self.turn_end_token_id),)
-        else:
-            self.terminal_token_ids = tuple(int(token_id) for token_id in terminal_token_ids)
-        self.channel_end_token_id = int(channel_end_token_id)
-        self.vocab_chunk_size = int(vocab_chunk_size)
-        self.jump_on_no_progress_after = int(jump_on_no_progress_after)
-        self.max_ponder_steps = int(max_ponder_steps)
-        self.min_trajectory_progress = float(min_trajectory_progress)
-        self.repetition_penalty = float(repetition_penalty)
-        self.state_schema_version = int(state_schema_version)
-        super().__init__(
-            text_config=text_config,
-            vision_config=vision_config,
-            initializer_range=initializer_range,
-            **kwargs,
-        )
-        self.model_type = type(self).model_type
-        if not hasattr(self, "eos_token_id"):
-            self.eos_token_id = self.text_config.eos_token_id
-        if not hasattr(self, "pad_token_id"):
-            self.pad_token_id = self.text_config.pad_token_id
-        if not hasattr(self, "bos_token_id"):
-            self.bos_token_id = self.text_config.bos_token_id
-        self._validate_modilify()
-
-    def _validate_modilify(self) -> None:
-        """Validate inference architecture and policy values."""
-
-        if self.memory_scheme != MEMORY_SCHEME:
-            raise ValueError(
-                f"`memory_scheme` must be {MEMORY_SCHEME!r}."
-            )
-        policy_values = (
-            self.denoise_temperature,
-            self.commit_failure_budget,
-            self.jump_failure_budget,
-            self.min_trajectory_progress,
-            self.repetition_penalty,
-        )
-        if any(not math.isfinite(value) for value in policy_values):
-            raise ValueError("Modilify Mk2 policy values must be finite.")
-        positive = (
-            self.denoise_temperature,
-            self.commit_failure_budget,
-            self.jump_failure_budget,
-            self.canvas_length,
-            self.latent_dim,
-            self.latent_ffn_dim,
-            self.latent_memory_slots,
-            self.latent_num_layers,
-            self.latent_num_heads,
-            self.latent_local_attention_window,
-            self.latent_history_length,
-            self.latent_tape_probes,
-            self.latent_history_kv_rank,
-            self.commit_sequence_layers,
-            self.commit_sequence_dim,
-            self.kv_cache_bucket_size,
-            self.vocab_chunk_size,
-            self.jump_on_no_progress_after,
-            self.max_ponder_steps,
-            self.repetition_penalty,
-        )
-        if any(value <= 0 for value in positive):
-            raise ValueError(
-                "Modilify Mk2 dimensions, budgets, intervals, and "
-                "`repetition_penalty` must be positive."
-            )
-        if self.latent_history_views != HISTORY_VIEWS:
-            raise ValueError(f"`latent_history_views` must be {HISTORY_VIEWS}.")
-        if self.experience_roles != EXPERIENCE_ROLES:
-            raise ValueError(f"`experience_roles` must be {EXPERIENCE_ROLES}.")
-        if self.commit_sequence_dim % self.latent_num_heads:
-            raise ValueError("`commit_sequence_dim` must be divisible by `latent_num_heads`.")
-        if self.commit_sequence_dim != self.latent_history_kv_rank:
-            raise ValueError("`commit_sequence_dim` must equal `latent_history_kv_rank`.")
-        if self.persistent_memory_write != "commit_only_transformer":
-            raise ValueError("`persistent_memory_write` must be commit_only_transformer.")
-        if self.writer_slot_gate != "per_slot":
-            raise ValueError("`writer_slot_gate` must be per_slot.")
-        if self.latent_history_kv_rank > self.latent_dim:
-            raise ValueError("`latent_history_kv_rank` must not exceed `latent_dim`.")
-        if self.latent_history_kv_rank % self.latent_num_heads:
-            raise ValueError("`latent_history_kv_rank` must be divisible by `latent_num_heads`.")
-        if not isinstance(self.channel_end_token_id, int) or self.channel_end_token_id < 0:
-            raise ValueError("`channel_end_token_id` must be a non-negative integer.")
-        if not self.terminal_token_ids:
-            raise ValueError("`terminal_token_ids` must not be empty.")
-        if any(
-            not isinstance(token_id, int) or token_id < 0
-            for token_id in self.terminal_token_ids
-        ):
-            raise ValueError("`terminal_token_ids` must be non-negative integers.")
-        if self.latent_dim % self.latent_num_heads:
-            raise ValueError("`latent_dim` must be divisible by `latent_num_heads`.")
-        if not 0.0 <= self.latent_dropout < 1.0:
-            raise ValueError("`latent_dropout` must be in [0, 1).")
-        if self.min_trajectory_progress < 0:
-            raise ValueError("`min_trajectory_progress` must be non-negative.")
-
-
-ModilifyMk2Config.register_for_auto_class()
-
-
-__all__ = [
-    "COMMIT_FAILURE_BUDGET",
-    "COMMIT_SEQUENCE_LAYERS",
-    "DENOISE_TEMPERATURE",
-    "EXPERIENCE_ROLES",
-    "HISTORY_VIEWS",
-    "JUMP_FAILURE_BUDGET",
-    "MEMORY_SCHEME",
-    "ModilifyMk2Config",
-    "ModilifyMk2TextConfig",
-    "VOCAB_CHUNK_SIZE",
-]
+            payload = dict(text_config)
+            payload.pop("model_type", None)
+            text_config = ModilifyMk2TextConfig(**payload)
+        text_config.use_bidirectional_attention = None
+        self.text_config = text_config
+        self.canvas_length = canvas_length
+        self.initializer_range = initializer_range
+        self.state_schema_version = state_schema_version
+        self.memory_architecture = memory_architecture
+        self.latent_dim = latent_dim
+        self.latent_ffn_dim = latent_ffn_dim
+        self.latent_num_layers = latent_num_layers
+        self.latent_num_heads = latent_num_heads
+        self.latent_local_attention_window = latent_local_attention_window
+        self.latent_tape_probes = latent_tape_probes
+        self.latent_history_kv_rank = latent_history_kv_rank
+        self.latent_working_last_block_global = latent_working_last_block_global
+        self.working_memory_bus = working_memory_bus
+        self.persistent_memory_bus = persistent_memory_bus
+        self.commit_sequence_dim = commit_sequence_dim
+        self.kv_cache_bucket_size = kv_cache_bucket_size
+        self.vocab_chunk_size = vocab_chunk_size
+        self.turn_end_token_id = turn_end_token_id
+        self.terminal_token_ids = terminal_token_ids
+        self.commit_failure_budget = commit_failure_budget
+        self.commit_top_k = commit_top_k
+        self.commit_min_p = commit_min_p
+        self.commit_target_confidence = commit_target_confidence
+        self.commit_entropy_weight = commit_entropy_weight
+        self.commit_confidence_power = commit_confidence_power
+        super().__init__(tie_word_embeddings=tie_word_embeddings,
+                         eos_token_id=eos_token_id, **kwargs)

continuous_batching.py

@@ -1,6 +1,4 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
-"""Continuous batching for Modilify Mk2 behind the Transformers public API shape.
+"""Continuous batching for ModilifyMk2 behind the Transformers public API shape.
 
 The upstream continuous runner is autoregressive: it persists every query in a
 paged cache and emits exactly one token per request and step.  ModilifyMk2 instead
@@ -47,21 +45,9 @@ from .generation_modilify_mk2 import (
     NoiseCanvasSampler,
     _add_repetition_history,
     _flatten_token_ids,
-    build_denoise_trace_event,
-    deterministic_episode_iteration_bound,
-)
-from .latent_deliberation import (
-    LatentDeliberationState,
-    TrajectoryHistory,
-    cat_latent_states,
-    cat_trajectory_history,
-    cat_trajectory_tape,
-    empty_trajectory_tape,
-    infer_commit_reason,
-    slice_latent_state,
-    slice_trajectory_history,
-    slice_trajectory_tape,
 )
+from .latent_deliberation import LatentDeliberationState, cat_latent_states, infer_commit_reason, slice_latent_state
+from .generation_modilify_mk2 import deterministic_episode_iteration_bound
 
 
 _TERMINAL_REASONS = frozenset(
@@ -100,16 +86,13 @@ def continuous_config_fingerprint(
 
 @dataclass
 class ModilifyMk2ContinuousGenerationOutput(GenerationOutput):
-    """Official ``GenerationOutput`` plus ModilifyMk2 request-local diagnostics."""
-
+    """Request output for the continuous generation runner."""
     stop_reason: str | None = None
     committed_tokens: int = 0
     denoise_steps: int = 0
     no_progress_steps: int = 0
     jump_count: int = 0
     forced_jump_bad_count: int = 0
-    heavy_forward_count: int = 0
-    latent_context_update_count: int = 0
     average_commit_len: float = 0.0
     tokens_per_forward: float = 0.0
     seed: int | None = None
@@ -121,8 +104,6 @@ class ModilifyMk2ContinuousGenerationOutput(GenerationOutput):
     is_stream_update: bool = False
     delta_tokens: list[int] = field(default_factory=list)
     state_shift_count: int = 0
-    latent_memory_norm: float = 0.0
-    state_retention_score: float = 0.0
 
     def is_finished(self) -> bool:
         """Treat failed/cancelled requests as terminal for every consumer API."""
@@ -133,7 +114,6 @@ class ModilifyMk2ContinuousGenerationOutput(GenerationOutput):
 @dataclass
 class ModilifyMk2RequestState:
     """All mutable state required to suspend and re-batch one request."""
-
     request_id: str
     prompt_ids: list[int]
     max_new_tokens: int
@@ -142,7 +122,6 @@ class ModilifyMk2RequestState:
     record_timestamps: bool
     seed: int
     max_denoising_steps: int | None
-    trace_callback: Callable[[dict[str, object]], None] | None = None
     created_time: float = field(default_factory=time.perf_counter)
     status: RequestStatus = RequestStatus.PENDING
     started_time: float = -1.0
@@ -178,10 +157,9 @@ def _slice_rolling_state(state: ModilifyMk2RollingState, row: int) -> ModilifyMk
         canvas=_clone_tensor_row(state.canvas, row),
         confidence=_clone_tensor_row(state.confidence, row),
         entropy=_clone_tensor_row(state.entropy, row),
-        age=_clone_tensor_row(state.age, row),
         latent_state=slice_latent_state(state.latent_state, selected),
-        history=slice_trajectory_history(state.history, selected),
-        tape=slice_trajectory_tape(state.tape, selected),
+
+        head=_clone_tensor_row(state.head, row),
     )
 
 
@@ -190,10 +168,9 @@ def _pack_rolling_states(states: Sequence[ModilifyMk2RollingState]) -> ModilifyM
         canvas=torch.cat([state.canvas for state in states], dim=0),
         confidence=torch.cat([state.confidence for state in states], dim=0),
         entropy=torch.cat([state.entropy for state in states], dim=0),
-        age=torch.cat([state.age for state in states], dim=0),
         latent_state=cat_latent_states([state.latent_state for state in states]),
-        history=cat_trajectory_history([state.history for state in states]),
-        tape=cat_trajectory_tape([state.tape for state in states]),
+
+        head=torch.cat([state.head for state in states], dim=0),
     )
 
 
@@ -320,9 +297,6 @@ class ModilifyMk2ContinuousBatchingManager:
         workload_hints: Any = None,
     ) -> None:
         del workload_hints
-        # Generation must not silently mutate the caller's train/eval mode.
-        # Inference mode below disables autograd without changing module-local
-        # dropout or other training flags.
         self.model = model
         self.generation_config = copy.deepcopy(
             generation_config or getattr(model, "generation_config", None)
@@ -348,6 +322,12 @@ class ModilifyMk2ContinuousBatchingManager:
         self.sampler: NoiseCanvasSampler = model._prepare_sampler(
             self.generation_config, model.config.canvas_length
         )
+        pad_token_id = self.generation_config.pad_token_id
+        if pad_token_id is None:
+            pad_token_id = getattr(model.config, "pad_token_id", None)
+        if isinstance(pad_token_id, (list, tuple)):
+            pad_token_id = pad_token_id[0]
+        self.pad_token_id = int(0 if pad_token_id is None else pad_token_id)
         self.run_id = uuid.uuid4().hex
         self.warmed_up = False
         self.destroyed = False
@@ -707,20 +687,12 @@ class ModilifyMk2ContinuousBatchingManager:
         ):
             raise ValueError("`input_ids` must be a non-empty list of integer token IDs.")
         seed = request_kwargs.pop("seed", None)
-        trace_callback = request_kwargs.pop("denoise_trace_callback", None)
         max_denoising_steps = request_kwargs.pop(
             "max_denoising_steps", self.generation_config.max_denoising_steps
         )
         if request_kwargs:
             unsupported = ", ".join(sorted(request_kwargs))
             raise ValueError(f"Unsupported per-request generation options: {unsupported}")
-        if trace_callback is not None and not callable(trace_callback):
-            raise TypeError("`denoise_trace_callback` must be callable.")
-        if trace_callback is not None and self.max_requests_per_batch > 1:
-            raise ValueError(
-                "ModilifyMk2 denoise tracing remains a batch-size-1 interface; "
-                "set `max_requests_per_batch=1`."
-            )
         limit = self.generation_config.max_new_tokens if max_new_tokens is None else max_new_tokens
         if not isinstance(limit, int) or limit <= 0:
             raise ValueError("`max_new_tokens` must be a positive integer.")
@@ -788,7 +760,6 @@ class ModilifyMk2ContinuousBatchingManager:
                 record_timestamps=bool(record_timestamps),
                 seed=resolved_seed & ((1 << 63) - 1),
                 max_denoising_steps=max_denoising_steps,
-                trace_callback=trace_callback,
             )
             state.reserved_blocks = math.ceil(
                 (len(state.prompt_ids) + state.max_new_tokens) / self.block_size
@@ -955,8 +926,6 @@ class ModilifyMk2ContinuousBatchingManager:
             ),
             jump_count=state.jumps,
             forced_jump_bad_count=state.forced_jump_tokens,
-            heavy_forward_count=state.denoise_steps,
-            latent_context_update_count=state.denoise_steps,
             average_commit_len=len(state.generated_tokens) / shifts,
             tokens_per_forward=len(state.generated_tokens) / steps,
             seed=state.seed,
@@ -974,16 +943,6 @@ class ModilifyMk2ContinuousBatchingManager:
                 state.last_delta_tokens if delta_tokens is None else delta_tokens
             ),
             state_shift_count=state.shifts,
-            latent_memory_norm=(
-                0.0
-                if state.rolling_state is None
-                else float(
-                    state.rolling_state.latent_state.memory_slots.float()
-                    .norm(dim=-1)
-                    .mean()
-                )
-            ),
-            state_retention_score=1.0 if state.shifts else 0.0,
         )
 
     def _finish(
@@ -1052,21 +1011,15 @@ class ModilifyMk2ContinuousBatchingManager:
         generator = torch.Generator(device=self.device)
         generator.manual_seed(state.seed)
         state.generator = generator
-        try:
-            canvas = self.sampler.initialize_canvas(
-                1, self.device, generators=[generator]
-            )
-        except TypeError:
-            canvas = self.sampler.initialize_canvas(1, self.device)
-        dtype = self.model.model.decoder.embed_tokens.weight.dtype
+        canvas = self.sampler.initialize_canvas(
+            1, self.device, generators=[generator]
+        )
+        canvas[:, state.max_new_tokens:] = self.pad_token_id
         canvas_length = int(self.model.config.canvas_length)
         latent = LatentDeliberationState.empty(
             batch_size=1,
             canvas_length=canvas_length,
-            latent_dim=self.model.config.latent_dim,
-            memory_slots=self.model.config.latent_memory_slots,
             device=self.device,
-            dtype=dtype,
         )
         state.rolling_state = ModilifyMk2RollingState(
             canvas=canvas,
@@ -1077,22 +1030,9 @@ class ModilifyMk2ContinuousBatchingManager:
                 device=self.device,
                 dtype=torch.float32,
             ),
-            age=torch.zeros(1, canvas_length, device=self.device, dtype=torch.int32),
             latent_state=latent,
-            history=TrajectoryHistory.empty(
-                batch_size=1,
-                canvas_length=canvas_length,
-                hidden_size=self.model.config.text_config.hidden_size,
-                history_length=self.model.config.latent_history_length,
-                device=self.device,
-                dtype=dtype,
-            ),
-            tape=empty_trajectory_tape(
-                batch_size=1,
-                config=self.model.config,
-                device=self.device,
-                dtype=dtype,
-            ),
+
+            head=torch.zeros(1, device=self.device, dtype=torch.long),
         )
         state.cache = self.cache_pool.prefill(state.prompt_ids)
         state.logical_length = len(state.prompt_ids)
@@ -1203,8 +1143,7 @@ class ModilifyMk2ContinuousBatchingManager:
                     stagnation_steps=rolling.latent_state.stagnation_steps[row : row + 1],
                     active_rows=torch.ones(1, device=self.device, dtype=torch.bool),
                     remaining_lengths=torch.tensor([remaining], device=self.device),
-                    failure_budget=self.generation_config.commit_failure_budget,
-                    jump_failure_budget=self.generation_config.jump_failure_budget,
+                    failure_budget=self.model.config.commit_failure_budget,
                     stop_token_id=state.eos_token_ids,
                     max_ponder_steps=self.generation_config.max_ponder_steps,
                     stagnation_threshold=self.generation_config.jump_on_no_progress_after,
@@ -1222,6 +1161,14 @@ class ModilifyMk2ContinuousBatchingManager:
 
     @torch.inference_mode()
     def _run_batch_step(self, states: Sequence[ModilifyMk2RequestState]) -> list[str]:
+        if len(states) > 1:
+            # The BF16 trunk changes reductions with batch shape. Routed
+            # experts amplify those differences into different commits, so
+            # run each active row with its independent-request numerics.
+            finished = []
+            for state in states:
+                finished.extend(self._run_batch_step([state]))
+            return finished
         started = time.perf_counter()
         rolling_states = [state.rolling_state for state in states]
         if any(state is None for state in rolling_states):
@@ -1261,15 +1208,12 @@ class ModilifyMk2ContinuousBatchingManager:
             decoder_input_ids=rolling.canvas,
             previous_confidence=rolling.confidence,
             previous_entropy=rolling.entropy,
-            token_age=rolling.age,
             latent_state=rolling.latent_state,
-            history=rolling.history,
-            tape=rolling.tape,
+
             decoder_position_ids=decoder_positions,
             decoder_read_cache=True,
             decoder_attention_mask=decoder_mask,
             compact_vocab=True,
-            denoise_temperature=self.generation_config.denoise_temperature,
             repetition_token_mask=repetition_history,
             repetition_penalty=self.repetition_penalty,
             sampling_generators=generators,
@@ -1295,46 +1239,56 @@ class ModilifyMk2ContinuousBatchingManager:
             output.next_latent_state,
             confidence=next_confidence.detach().float(),
             entropy=token_entropy.detach().float(),
-            age=rolling.age + 1,
-            token_changed=next_canvas.ne(rolling.canvas).detach().float(),
-            confidence_delta=next_confidence.detach().float() - rolling.confidence,
-            entropy_delta=token_entropy.detach().float() - rolling.entropy,
         )
-        live_mask = torch.ones(
-            rolling.canvas.shape, device=rolling.canvas.device, dtype=torch.bool
+        remaining_lengths = torch.tensor(
+            [state.max_new_tokens - len(state.generated_tokens) for state in states],
+            device=self.device,
+            dtype=torch.long,
+        )
+        live_mask = torch.arange(canvas_length, device=self.device)[None, :].lt(
+            remaining_lengths[:, None]
         )
-        tape_probes, tape_valid = self.model.latent_deliberation.encode_tape_frame(
-            output.heavy_hidden_state, live_mask
+        next_latent = self.model.latent_deliberation.observe_state(
+            next_latent, output.heavy_hidden_state, output.working_state,
+            live_mask, rolling.head,
         )
         next_state = ModilifyMk2RollingState(
             canvas=next_canvas,
             confidence=next_confidence,
             entropy=token_entropy,
-            age=rolling.age + 1,
             latent_state=next_latent,
-            history=rolling.history.append(
-                output.heavy_hidden_state,
-                next_confidence,
-                token_entropy,
-                next_canvas.ne(rolling.canvas).detach().float(),
-                live_mask=live_mask,
-            ),
-            tape=rolling.tape.append(tape_probes, tape_valid),
+
+            head=rolling.head,
         )
         normal_failure_rate = fused_commit_failure_rate(
             proposal_confidence,
             token_entropy,
-            vocab_size=self.model.config.text_config.vocab_size,
+            entropy_weight=self.model.config.commit_entropy_weight,
+            confidence_power=self.model.config.commit_confidence_power,
+            top_k=getattr(self.model.config, "commit_top_k", None),
+            min_p=getattr(self.model.config, "commit_min_p", None),
+            target_confidence=getattr(self.model.config, "commit_target_confidence", None),
+            failure_budget=float(self.model.config.commit_failure_budget),
         )
         jump_failure_rate = fused_commit_failure_rate(
             greedy_confidence,
             token_entropy,
-            vocab_size=self.model.config.text_config.vocab_size,
+            entropy_weight=self.model.config.commit_entropy_weight,
+            confidence_power=self.model.config.commit_confidence_power,
+            top_k=getattr(self.model.config, "commit_top_k", None),
+            min_p=getattr(self.model.config, "commit_min_p", None),
+            target_confidence=getattr(self.model.config, "commit_target_confidence", None),
+            failure_budget=float(self.model.config.commit_failure_budget),
         )
         previous_failure_rate = fused_commit_failure_rate(
             rolling.confidence,
             rolling.entropy,
-            vocab_size=self.model.config.text_config.vocab_size,
+            entropy_weight=self.model.config.commit_entropy_weight,
+            confidence_power=self.model.config.commit_confidence_power,
+            top_k=getattr(self.model.config, "commit_top_k", None),
+            min_p=getattr(self.model.config, "commit_min_p", None),
+            target_confidence=getattr(self.model.config, "commit_target_confidence", None),
+            failure_budget=float(self.model.config.commit_failure_budget),
         )
         (
             normal_commit,
@@ -1372,17 +1326,13 @@ class ModilifyMk2ContinuousBatchingManager:
                 stagnation_steps=next_stagnation,
             ),
         )
-        unshifted_trace_states = [
-            _slice_rolling_state(next_state, row) for row in range(batch_size)
-        ]
-        if output.history_projected is None or output.working_state is None:
+        if output.working_state is None:
             raise RuntimeError("Forward did not return working trajectory features.")
         next_state = self.model._write_committed_memory(
-            previous_history=rolling.history,
             next_state=next_state,
             working_state=output.working_state,
-            history_projected=output.history_projected,
             heavy_hidden=output.heavy_hidden_state,
+            commit_token_ids=commit_token_ids,
             commit_lengths=commit_lengths,
             prefix_lengths=logical_lengths,
             commit_reason=infer_commit_reason(
@@ -1399,6 +1349,8 @@ class ModilifyMk2ContinuousBatchingManager:
             commit_lengths,
             self.sampler,
             generators=generators,
+            remaining_lengths=remaining_lengths - commit_lengths,
+            pad_token_id=self.pad_token_id,
         )
         shifted_states = [
             _slice_rolling_state(shifted, row) for row in range(batch_size)
@@ -1462,31 +1414,6 @@ class ModilifyMk2ContinuousBatchingManager:
                 reason = "episode_watchdog"
 
             elapsed = time.perf_counter() - started
-            if state.trace_callback is not None:
-                trace = build_denoise_trace_event(
-                    denoise_step=state.denoise_steps,
-                    prefix_length=state.logical_length - commit_length,
-                    committed_before=before,
-                    committed_after=len(state.generated_tokens),
-                    no_progress_steps=int(next_stagnation[row]),
-                    policy_prefix_mask=policy_prefix_mask[row : row + 1],
-                    commit_length=commit_length,
-                    ponder_fallback=bool(jump_rows[row]),
-                    state=unshifted_trace_states[row],
-                    proposal=proposal[row : row + 1],
-                    committed_token_ids=commit_token_ids[row : row + 1, :commit_length],
-                    step_elapsed_seconds=elapsed,
-                    latent_residual_diagnostics=None,
-                )
-                trace["request_id"] = state.request_id
-                trace["batch_size"] = batch_size
-                try:
-                    state.trace_callback(trace)
-                except Exception as error:
-                    warnings.warn(
-                        f"Denoise trace callback failed for {state.request_id}: {error!r}",
-                        stacklevel=2,
-                    )
 
             if reason is not None:
                 self._finish(state, reason)
@@ -1691,7 +1618,6 @@ def generate_static_batch_with_logical_cache(
         return torch.tensor(
             [getattr(output, name) for output in ordered],
             device=input_ids.device,
-            dtype=dtype,
         )
 
     return ModilifyMk2GenerationOutput(
@@ -1705,22 +1631,6 @@ def generate_static_batch_with_logical_cache(
         no_progress_steps=tensor("no_progress_steps", dtype=torch.long),
         jump_count=tensor("jump_count", dtype=torch.long),
         forced_jump_bad_count=tensor("forced_jump_bad_count", dtype=torch.long),
-        heavy_forward_count=tensor("heavy_forward_count", dtype=torch.long),
-        latent_context_update_count=tensor(
-            "latent_context_update_count", dtype=torch.long
-        ),
         average_commit_len=tensor("average_commit_len", dtype=torch.float32),
         state_shift_count=tensor("state_shift_count", dtype=torch.long),
-        latent_memory_norm=tensor("latent_memory_norm", dtype=torch.float32),
-        state_retention_score=tensor("state_retention_score", dtype=torch.float32),
     )
-
-
-__all__ = [
-    "ModilifyMk2ContinuousBatchingManager",
-    "ModilifyMk2ContinuousGenerationOutput",
-    "ModilifyMk2LogicalCachePool",
-    "ModilifyMk2RequestState",
-    "continuous_config_fingerprint",
-    "generate_static_batch_with_logical_cache",
-]

gdn2_memory.py

@@ -0,0 +1,112 @@
+"""Reference GDN2 matrix memory for trajectory-time recurrence.
+
+The state is FP32.  Leading dimensions are independent rows or canvas cells;
+only the last three dimensions, [heads, key, value], belong to the rule.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class _GateProjection(nn.Module):
+    """GDN2 gates need channel control, not a second full-width content map."""
+
+    def __init__(self, source: int, target: int, rank: int) -> None:
+        super().__init__()
+        self.down = nn.Linear(source, rank, bias=False)
+        self.up = nn.Linear(rank, target, bias=False)
+
+    def forward(self, source: torch.Tensor) -> torch.Tensor:
+        return self.up(self.down(source))
+
+
+class GDN2Memory(nn.Module):
+    def __init__(self, input_dim: int, heads: int, key_dim: int, value_dim: int,
+                 *, observation_dim: int | None = None) -> None:
+        super().__init__()
+        if min(input_dim, heads, key_dim, value_dim) <= 0:
+            raise ValueError("GDN2 dimensions must be positive.")
+        self.heads, self.key_dim, self.value_dim = heads, key_dim, value_dim
+        observation_dim = input_dim if observation_dim is None else observation_dim
+        if observation_dim <= 0:
+            raise ValueError("GDN2 observation width must be positive.")
+        self.q_proj = nn.Linear(input_dim, heads * key_dim, bias=False)
+        self.k_proj = nn.Linear(observation_dim, heads * key_dim, bias=False)
+        self.v_proj = nn.Linear(observation_dim, heads * value_dim, bias=False)
+        self.f_proj = _GateProjection(observation_dim, heads * key_dim, min(observation_dim, key_dim))
+        self.b_proj = nn.Linear(observation_dim, heads * key_dim, bias=False)
+        self.w_proj = nn.Linear(observation_dim, heads * value_dim, bias=False)
+        self.g_proj = _GateProjection(input_dim, heads * value_dim, min(input_dim, value_dim))
+        self.o_proj = nn.Linear(heads * value_dim, input_dim, bias=False)
+        self.a_log = nn.Parameter(torch.zeros(heads))
+        self.dt_bias = nn.Parameter(torch.full((heads, key_dim), -6.906255))
+
+    def _projections(self, source: torch.Tensor):
+        source = F.rms_norm(source.float(), (source.shape[-1],), eps=1e-6).to(source.dtype)
+        shape = source.shape[:-1]
+        key_shape = (*shape, self.heads, self.key_dim)
+        value_shape = (*shape, self.heads, self.value_dim)
+        k = F.normalize(F.silu(self.k_proj(source).float()).reshape(key_shape), dim=-1)
+        v = F.silu(self.v_proj(source).float()).reshape(value_shape)
+        head_rate = self.a_log.float().exp().reshape(*((1,) * (source.ndim - 1)), self.heads, 1)
+        decay = torch.exp(-head_rate * F.softplus(
+            self.f_proj(source).float().reshape(key_shape) + self.dt_bias.float()))
+        erase = torch.sigmoid(self.b_proj(source).float().reshape(key_shape))
+        write = torch.sigmoid(self.w_proj(source).float().reshape(value_shape))
+        return k, v, decay, erase, write
+
+    def _query(self, source: torch.Tensor) -> torch.Tensor:
+        return F.normalize(F.silu(self.q_proj(source).float()).reshape(
+            *source.shape[:-1], self.heads, self.key_dim), dim=-1)
+
+    def _output(self, value: torch.Tensor, source: torch.Tensor) -> torch.Tensor:
+        gate = F.silu(self.g_proj(source).float().reshape(value.shape))
+        value = F.rms_norm(value, (self.value_dim,), eps=1e-6) * gate
+        return self.o_proj(value.flatten(-2).to(source.dtype))
+
+    def read(self, state: torch.Tensor, source: torch.Tensor) -> torch.Tensor:
+        if state.shape != (*source.shape[:-1], self.heads, self.key_dim, self.value_dim):
+            raise ValueError("GDN2 state and query leading dimensions differ.")
+        value = torch.einsum("...hk,...hkv->...hv", self._query(source), state.float())
+        return self._output(value, source)
+
+    def read_shared(self, state: torch.Tensor, source: torch.Tensor) -> torch.Tensor:
+        """Read one row matrix for all queries without a canvas-sized matrix product."""
+        if source.ndim != 3 or state.shape != (source.shape[0], self.heads, self.key_dim, self.value_dim):
+            raise ValueError("Shared GDN2 read requires [batch, canvas, width] queries.")
+        query = self._query(source).transpose(1, 2)
+        value = (query @ state.float()).transpose(1, 2)
+        return self._output(value, source)
+
+    @staticmethod
+    def _transition(state, k, v, decay, erase, write, valid):
+        decayed = state.float() * decay.unsqueeze(-1)
+        old = torch.einsum("...hk,...hkv->...hv", erase * k, decayed)
+        candidate = decayed + k.unsqueeze(-1) * (write * v - old).unsqueeze(-2)
+        return candidate if valid is None else torch.where(
+            valid[..., None, None, None], candidate, state.float())
+
+    def transition(self, state: torch.Tensor, source: torch.Tensor,
+                   valid: torch.Tensor | None = None) -> torch.Tensor:
+        if state.shape != (*source.shape[:-1], self.heads, self.key_dim, self.value_dim):
+            raise ValueError("GDN2 state and observation leading dimensions differ.")
+        k, v, decay, erase, write = self._projections(source)
+        if valid is not None:
+            if valid.shape != source.shape[:-1]:
+                raise ValueError("GDN2 valid mask must match observation rows.")
+        return self._transition(state, k, v, decay, erase, write, valid)
+
+    def write_sequence(self, state: torch.Tensor, source: torch.Tensor,
+                       valid: torch.Tensor) -> torch.Tensor:
+        """Project a commit packet once, then preserve ordered GDN2 recurrence."""
+        if source.ndim != 3 or valid.shape != source.shape[:2]:
+            raise ValueError("GDN2 sequence and mask must share [batch, length].")
+        if state.shape != (source.shape[0], self.heads, self.key_dim, self.value_dim):
+            raise ValueError("GDN2 sequence state shape differs.")
+        projected = self._projections(source)
+        for index in range(source.shape[1]):
+            state = self._transition(state, *(part[:, index] for part in projected), valid[:, index])
+        return state

gdn2_trajectory.py

@@ -0,0 +1,86 @@
+"""Dual-timescale matrix state with denoise and commit lifetimes.
+
+This module is independent of the decoder and commit policy.  It provides the
+reference state transition used when replacing the old history and slot paths.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from .gdn2_memory import GDN2Memory
+
+
+@dataclass(frozen=True)
+class GDN2TrajectoryState:
+    cells: torch.Tensor
+    row: torch.Tensor
+    persistent: torch.Tensor
+    seen: torch.Tensor
+
+
+    def shift(self, lengths: torch.Tensor) -> "GDN2TrajectoryState":
+        """Shift a non-ring canvas and zero its newly filled tail."""
+        batch, canvas = self.seen.shape
+        if lengths.shape != (batch,):
+            raise ValueError("Commit lengths must be per row.")
+        physical = torch.arange(canvas, device=self.seen.device)[None, :] + lengths[:, None]
+        kept = physical < canvas
+        selected = physical.clamp_max(canvas - 1)
+        cells = self.cells.gather(
+            1, selected[..., None, None, None].expand_as(self.cells)
+        )
+        seen = self.seen.gather(1, selected)
+        return GDN2TrajectoryState(
+            cells.masked_fill(~kept[..., None, None, None], 0.0),
+            self.row, self.persistent, seen & kept,
+        )
+
+
+class GDN2TrajectoryMemory(nn.Module):
+    def __init__(self, width: int, *, probes: int = 4,
+                 working_heads: int = 16, working_key: int = 64,
+                 working_value: int = 64, persistent_heads: int = 16,
+                 persistent_key: int = 128, persistent_value: int = 128,
+                 persistent_observation_dim: int | None = None) -> None:
+        super().__init__()
+        if probes <= 0:
+            raise ValueError("Probe count must be positive.")
+        self.probes = probes
+        self.cell = GDN2Memory(width, working_heads, working_key, working_value)
+        self.row = GDN2Memory(width, working_heads, working_key, working_value)
+        self.persistent = GDN2Memory(width, persistent_heads, persistent_key, persistent_value,
+                                    observation_dim=persistent_observation_dim)
+        self.probe_embed = nn.Embedding(probes, width)
+        nn.init.normal_(self.probe_embed.weight, std=0.02)
+
+    def read(self, state: GDN2TrajectoryState, query: torch.Tensor) -> torch.Tensor:
+        result = (self.cell.read(state.cells, query)
+                  + self.row.read_shared(state.row, query)
+                  + self.persistent.read_shared(state.persistent, query))
+        return torch.where(state.seen[..., None], result, torch.zeros_like(result))
+
+    def observe(self, state: GDN2TrajectoryState, observation: torch.Tensor,
+                live: torch.Tensor, head: torch.Tensor) -> GDN2TrajectoryState:
+        batch, canvas, width = observation.shape
+        if live.shape != (batch, canvas) or head.shape != (batch,):
+            raise ValueError("Observation mask and head have incorrect shapes.")
+        cells = self.cell.transition(state.cells, observation, live)
+        seen = state.seen | live
+        logical_idx = (head[:, None] + torch.arange(canvas, device=head.device)[None, :]) % canvas
+        logical = observation.gather(1, logical_idx[..., None].expand(-1, -1, width))
+        logical_live = live.gather(1, logical_idx)
+        row_state = state.row
+        for probe in range(self.probes):
+            lo = canvas * probe // self.probes
+            hi = canvas * (probe + 1) // self.probes
+            selected = logical_live[:, lo:hi]
+            count = selected.sum(dim=1, keepdim=True)
+            pooled = (logical[:, lo:hi].float() * selected[..., None]).sum(dim=1)
+            pooled = (pooled / count.clamp_min(1)).to(observation.dtype)
+            pooled = pooled + self.probe_embed.weight[probe].to(observation.dtype)
+            row_state = self.row.transition(row_state, pooled, count[:, 0] > 0)
+        return GDN2TrajectoryState(cells, row_state, state.persistent, seen)

generation_config.json

@@ -1,24 +1,10 @@
 {
-  "commit_failure_budget": 0.2,
-  "confidence_threshold": null,
-  "denoise_temperature": 0.8,
-  "eos_token_id": [
-    1,
-    106
-  ],
-  "jump_failure_budget": 2.0,
+  "eos_token_id": 1,
   "jump_on_no_progress_after": 12,
   "max_denoising_steps": null,
-  "max_new_tokens": 256,
   "max_ponder_steps": 64,
   "min_trajectory_progress": 0.005,
   "repetition_penalty": 1.0,
   "repetition_penalty_exclude_token_ids": [],
-  "return_dict_in_generate": true,
-  "sampler_config": null,
-  "stability_threshold": null,
-  "t_max": 0.8,
-  "t_min": 0.8,
-  "transformers_version": "5.14.1",
   "turn_end_token_id": 106
 }

generation_modilify_mk2.py

@@ -1,14 +1,11 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
-"""Rolling generation for latent-memory Modilify Mk2."""
+"""Rolling text generation for latent-memory ModilifyMk2."""
 
 from __future__ import annotations
 
-from collections.abc import Callable, Sequence
+from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass, replace
 import math
-import time
 from typing import Any
 
 import torch
@@ -22,36 +19,37 @@ from transformers.models.diffusion_gemma import (
     DiffusionGemmaGenerationMixin,
 )
 from .commit_policy import (
-    first_committed_token_lengths,
     fused_commit_failure_rate,
     select_commit_lengths,
 )
-from .configuration_modilify_mk2 import (
-    COMMIT_FAILURE_BUDGET,
-    DENOISE_TEMPERATURE,
-    JUMP_FAILURE_BUDGET,
-)
-from .latent_deliberation import (
-    LatentDeliberationState,
-    TrajectoryHistory,
-    TrajectoryTape,
-    choose_trajectory_history,
-    choose_trajectory_tape,
-    empty_trajectory_tape,
-    infer_commit_reason,
-)
+from .configuration_modilify_mk2 import DENOISE_TEMPERATURE
+from .latent_deliberation import LatentDeliberationState, infer_commit_reason
 
 
+def deterministic_episode_iteration_bound(
+    response_lengths: torch.LongTensor,
+    *,
+    max_ponder_steps: int,
+) -> int:
+    """Return a safe watchdog bound without assuming canvas-sized jumps.
+
+    Every active row must advance by at least one token no later than the
+    configured no-progress threshold.  Normal commits can also be only one
+    token long, so a bound based on the number of canvases is not valid.
+    """
+    if response_lengths.numel() == 0:
+        raise ValueError("`response_lengths` must be non-empty.")
+    if max_ponder_steps <= 0:
+        raise ValueError("`max_ponder_steps` must be positive.")
+    return max(1, int(response_lengths.max()) * max_ponder_steps)
+
 _PARENT_GENERATION_KEYS = frozenset({
     "max_new_tokens",
     "max_length",
     "return_dict_in_generate",
     "max_denoising_steps",
-    "sampler_config",
     "t_min",
     "t_max",
-    "stability_threshold",
-    "confidence_threshold",
     "cache_implementation",
     "cache_config",
     "disable_compile",
@@ -62,38 +60,8 @@ _PARENT_GENERATION_KEYS = frozenset({
     "_from_model_config",
     "transformers_version",
 })
-_IGNORED_GENERATION_KEYS = frozenset({
-    "compile_generation",
-    "sliding_denoise",
-    "one_token_per_denoise_step",
-    "adaptive_ponder_budget",
-    "force_commit_on_max_steps",
-    "ponder_budget_id",
-})
-_FORBIDDEN_COMMIT_FIELDS = (
-    "sampler_config",
-    "stability_threshold",
-    "confidence_threshold",
-    "one_token_per_denoise_step",
-)
-
-
-def _reject_legacy_commit_fields(fields: dict[str, object]) -> None:
-    configured = {
-        name: value
-        for name, value in fields.items()
-        if value not in (None, False)
-    }
-    if configured:
-        raise ValueError(
-            "ModilifyMk2 accepts only the confidence-prefix commit policy; "
-            f"unsupported generation fields: {sorted(configured)}"
-        )
-
-
 def _flatten_token_ids(*values: object) -> set[int]:
     """Normalize scalar and sequence token-ID configuration values."""
-
     token_ids: set[int] = set()
     for value in values:
         if value is None:
@@ -128,21 +96,7 @@ def _add_repetition_history(
 
 class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
     def __init__(self, **kwargs):
-        _reject_legacy_commit_fields({
-            name: kwargs.pop(name)
-            for name in _FORBIDDEN_COMMIT_FIELDS
-            if name in kwargs
-        })
         self.turn_end_token_id: int | None = kwargs.pop("turn_end_token_id", None)
-        self.denoise_temperature: float = float(
-            kwargs.pop("denoise_temperature", DENOISE_TEMPERATURE)
-        )
-        self.commit_failure_budget: float = float(
-            kwargs.pop("commit_failure_budget", COMMIT_FAILURE_BUDGET)
-        )
-        self.jump_failure_budget: float = float(
-            kwargs.pop("jump_failure_budget", JUMP_FAILURE_BUDGET)
-        )
         self.max_ponder_steps: int = kwargs.pop("max_ponder_steps", 64)
         self.jump_on_no_progress_after: int = kwargs.pop("jump_on_no_progress_after", 12)
         self.min_trajectory_progress: float = float(kwargs.pop("min_trajectory_progress", 0.005))
@@ -151,8 +105,6 @@ class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
         self.repetition_penalty_exclude_token_ids: list[int] = list(
             dict.fromkeys(int(token_id) for token_id in excluded_token_ids or ())
         )
-        for name in _IGNORED_GENERATION_KEYS:
-            kwargs.pop(name, None)
         kwargs.pop("t_min", None)
         kwargs.pop("t_max", None)
         parent_kwargs = {
@@ -166,26 +118,14 @@ class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
             confidence_threshold=None,
             **parent_kwargs,
         )
-        self.t_min = self.denoise_temperature
-        self.t_max = self.denoise_temperature
-        self.validate()
+        self.t_min = DENOISE_TEMPERATURE
+        self.t_max = DENOISE_TEMPERATURE
 
     def update(self, defaults_only=False, allow_custom_entries=False, **kwargs):
-        """Apply supported overrides, including temperature and failure budget."""
+        """Apply supported overrides while keeping temperature globally fixed."""
 
-        _reject_legacy_commit_fields({
-            name: kwargs.pop(name)
-            for name in _FORBIDDEN_COMMIT_FIELDS
-            if name in kwargs
-        })
         if "turn_end_token_id" in kwargs:
             self.turn_end_token_id = kwargs.pop("turn_end_token_id")
-        if "denoise_temperature" in kwargs:
-            self.denoise_temperature = float(kwargs.pop("denoise_temperature"))
-        if "commit_failure_budget" in kwargs:
-            self.commit_failure_budget = float(kwargs.pop("commit_failure_budget"))
-        if "jump_failure_budget" in kwargs:
-            self.jump_failure_budget = float(kwargs.pop("jump_failure_budget"))
         if "max_ponder_steps" in kwargs:
             self.max_ponder_steps = kwargs.pop("max_ponder_steps")
         if "jump_on_no_progress_after" in kwargs:
@@ -199,8 +139,6 @@ class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
             self.repetition_penalty_exclude_token_ids = list(
                 dict.fromkeys(int(token_id) for token_id in excluded_token_ids or ())
             )
-        for name in _IGNORED_GENERATION_KEYS:
-            kwargs.pop(name, None)
         kwargs.pop("t_min", None)
         kwargs.pop("t_max", None)
         unused = super().update(
@@ -211,8 +149,8 @@ class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
         self.sampler_config = None
         self.stability_threshold = None
         self.confidence_threshold = None
-        self.t_min = self.denoise_temperature
-        self.t_max = self.denoise_temperature
+        self.t_min = DENOISE_TEMPERATURE
+        self.t_max = DENOISE_TEMPERATURE
         return unused
 
     def validate(self, **kwargs):
@@ -234,12 +172,6 @@ class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
             raise ValueError("`jump_on_no_progress_after` must be a positive integer.")
         if not isinstance(self.min_trajectory_progress, (int, float)) or self.min_trajectory_progress < 0:
             raise ValueError("`min_trajectory_progress` must be a non-negative number.")
-        if not math.isfinite(self.denoise_temperature) or self.denoise_temperature <= 0:
-            raise ValueError("`denoise_temperature` must be positive.")
-        if not math.isfinite(self.commit_failure_budget) or self.commit_failure_budget <= 0:
-            raise ValueError("`commit_failure_budget` must be positive.")
-        if not math.isfinite(self.jump_failure_budget) or self.jump_failure_budget <= 0:
-            raise ValueError("`jump_failure_budget` must be positive.")
         if not math.isfinite(self.repetition_penalty) or self.repetition_penalty <= 0:
             raise ValueError("`repetition_penalty` must be a finite positive number.")
         if any(
@@ -252,33 +184,9 @@ class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
 
     @classmethod
     def from_model_config(cls, model_config):
-        """Build generation defaults from a model configuration."""
+        """Build the only generation field owned by the ModilifyMk2 model config."""
 
-        return cls(
-            turn_end_token_id=model_config.turn_end_token_id,
-            denoise_temperature=getattr(
-                model_config, "denoise_temperature", DENOISE_TEMPERATURE
-            ),
-            commit_failure_budget=getattr(
-                model_config, "commit_failure_budget", COMMIT_FAILURE_BUDGET
-            ),
-            jump_failure_budget=getattr(
-                model_config, "jump_failure_budget", JUMP_FAILURE_BUDGET
-            ),
-            max_ponder_steps=getattr(model_config, "max_ponder_steps", 64),
-            jump_on_no_progress_after=getattr(
-                model_config, "jump_on_no_progress_after", 12
-            ),
-            min_trajectory_progress=getattr(
-                model_config, "min_trajectory_progress", 0.005
-            ),
-            repetition_penalty=getattr(model_config, "repetition_penalty", 1.0),
-            eos_token_id=getattr(
-                model_config,
-                "eos_token_id",
-                model_config.text_config.eos_token_id,
-            ),
-        )
+        return cls(turn_end_token_id=model_config.turn_end_token_id)
 
     @staticmethod
     def _get_default_generation_params() -> dict[str, object]:
@@ -292,20 +200,6 @@ class ModilifyMk2GenerationConfig(DiffusionGemmaGenerationConfig):
         }
 
 
-def deterministic_episode_iteration_bound(
-    response_lengths: torch.LongTensor,
-    *,
-    max_ponder_steps: int,
-) -> int:
-    """Return a safe watchdog bound without assuming canvas-sized jumps."""
-
-    if response_lengths.numel() == 0:
-        raise ValueError("`response_lengths` must be non-empty.")
-    if max_ponder_steps <= 0:
-        raise ValueError("`max_ponder_steps` must be positive.")
-    return max(1, int(response_lengths.max()) * max_ponder_steps)
-
-
 @dataclass
 class ModilifyMk2GenerationOutput(ModelOutput):
     sequences: torch.LongTensor
@@ -318,135 +212,18 @@ class ModilifyMk2GenerationOutput(ModelOutput):
     no_progress_steps: int | torch.LongTensor | None = None
     jump_count: int | torch.LongTensor | None = None
     forced_jump_bad_count: int | torch.LongTensor | None = None
-    heavy_forward_count: int | torch.LongTensor | None = None
-    latent_context_update_count: int | torch.LongTensor | None = None
     average_commit_len: float | torch.FloatTensor | None = None
     state_shift_count: int | torch.LongTensor | None = None
-    latent_memory_norm: float | torch.FloatTensor | None = None
-    state_retention_score: float | torch.FloatTensor | None = None
-    logits: None = None
-    scores: None = None
-    hidden_states: None = None
 
 
 @dataclass
 class ModilifyMk2RollingState:
     """All real iterative state; no vocabulary-sized tensor is retained."""
-
     canvas: torch.LongTensor
     confidence: torch.FloatTensor
     entropy: torch.FloatTensor
-    age: torch.IntTensor
     latent_state: LatentDeliberationState
-    history: TrajectoryHistory
-    tape: TrajectoryTape
-
-
-def qualified_single_turn_end_lengths(
-    token_ids: torch.LongTensor,
-    qualified_mask: torch.BoolTensor,
-    turn_end_token_id: int,
-) -> torch.LongTensor:
-    if token_ids.ndim != 2 or token_ids.shape != qualified_mask.shape:
-        raise ValueError("Token IDs and qualification mask must share shape [batch, canvas].")
-    qualified_prefix = qualified_mask.long().cumprod(dim=1).sum(dim=-1)
-    clipped = first_committed_token_lengths(
-        token_ids,
-        qualified_prefix,
-        turn_end_token_id,
-    )
-    found = clipped.lt(qualified_prefix)
-    boundary_is_turn = token_ids.gather(
-        1,
-        clipped.clamp_min(1).sub(1)[:, None],
-    ).squeeze(1).eq(turn_end_token_id)
-    return torch.where(
-        found | boundary_is_turn,
-        clipped,
-        torch.zeros_like(clipped),
-    )
-
-
-def _prefix_length(prefix_mask: torch.BoolTensor) -> int:
-    common = prefix_mask.all(dim=0)
-    rejected = (~common).nonzero(as_tuple=False)
-    return common.shape[0] if rejected.numel() == 0 else int(rejected[0, 0])
-
-
-def _tensor_stats(value: torch.Tensor) -> dict[str, float]:
-    data = value.detach().float()
-    return {
-        "min": float(data.min()), "max": float(data.max()),
-        "mean": float(data.mean()), "norm": float(data.norm()),
-    }
-
-
-def _memory_diversity_stats(memory_slots: torch.Tensor) -> dict[str, float]:
-    """Trace-only collapse diagnostics; the pairwise term is tiny (32 slots)."""
-
-    memory = memory_slots.detach().float()
-    centered = memory - memory.mean(dim=1, keepdim=True)
-    maximum_difference = (
-        (memory[:, 1:] - memory[:, :-1]).abs().amax()
-        if memory.shape[1] > 1 else memory.new_zeros(())
-    )
-    normalized = torch.nn.functional.normalize(memory, dim=-1, eps=1.0e-6)
-    pairwise = normalized @ normalized.transpose(-1, -2)
-    off_diagonal = ~torch.eye(memory.shape[1], device=memory.device, dtype=torch.bool)
-    return {
-        "memory_slot_std": float(centered.square().mean().sqrt()),
-        "memory_slot_max_difference": float(maximum_difference),
-        "memory_slot_pairwise_cosine_mean": float(pairwise[:, off_diagonal].mean()),
-    }
-
-
-def build_denoise_trace_event(
-    *,
-    denoise_step: int,
-    prefix_length: int,
-    committed_before: int,
-    committed_after: int,
-    no_progress_steps: int,
-    policy_prefix_mask: torch.BoolTensor,
-    commit_length: int,
-    ponder_fallback: bool,
-    state: ModilifyMk2RollingState,
-    proposal: torch.LongTensor,
-    committed_token_ids: torch.LongTensor,
-    step_elapsed_seconds: float,
-    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None,
-) -> dict[str, object]:
-    return {
-        "event": "denoise_step",
-        "denoise_step": denoise_step,
-        "prefix_length": prefix_length,
-        "committed_before": committed_before,
-        "committed_after": committed_after,
-        "no_progress_steps": no_progress_steps,
-        "policy_prefix_count": int(policy_prefix_mask.sum()),
-        "policy_prefix_length": _prefix_length(policy_prefix_mask),
-        "commit_length": commit_length,
-        "ponder_fallback": bool(ponder_fallback),
-        "confidence": _tensor_stats(state.confidence),
-        "entropy": _tensor_stats(state.entropy),
-        "age": _tensor_stats(state.age),
-        "token_changed": _tensor_stats(state.latent_state.token_changed),
-        "confidence_delta": _tensor_stats(state.latent_state.confidence_delta),
-        "entropy_delta": _tensor_stats(state.latent_state.entropy_delta),
-        "ponder_steps": int(state.latent_state.ponder_steps[0]),
-        "stagnation_steps": int(state.latent_state.stagnation_steps[0]),
-        "history_fill": float(state.history.valid[0].float().mean()),
-        "memory_slots": _tensor_stats(state.latent_state.memory_slots),
-        "memory_diversity": _memory_diversity_stats(state.latent_state.memory_slots),
-        "latent_residual": {
-            name: float(value.detach().float())
-            for name, value in (latent_residual_diagnostics or {}).items()
-        },
-        "proposal_token_ids": proposal[0].detach().cpu().tolist(),
-        "draft_token_ids": state.canvas[0].detach().cpu().tolist(),
-        "committed_token_ids": committed_token_ids[0].detach().cpu().tolist(),
-        "step_elapsed_seconds": step_elapsed_seconds,
-    }
+    head: torch.LongTensor | None = None
 
 
 class NoiseCanvasSampler:
@@ -659,114 +436,48 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
             vocab_size=self.config.text_config.vocab_size,
         )
 
-    @staticmethod
-    def _shift_state(
-        state: ModilifyMk2RollingState,
-        commit_length: int,
-        sampler: NoiseCanvasSampler,
-        **_: object,
-    ) -> ModilifyMk2RollingState:
-        if commit_length == 0:
-            return state
-        canvas_length = state.canvas.shape[1]
-        if not 0 < commit_length <= canvas_length:
-            raise ValueError(f"`commit_length` must be in [1, {canvas_length}].")
-        tail = sampler.initialize_canvas(state.canvas.shape[0], state.canvas.device)[:, :commit_length]
-        canvas = torch.cat((state.canvas[:, commit_length:], tail), dim=1)
-
-        def shift(value: torch.Tensor | None, fill_value: float | int = 0) -> torch.Tensor | None:
-            if value is None:
-                return None
-            tail_state = torch.full(
-                (value.shape[0], commit_length, *value.shape[2:]),
-                fill_value, device=value.device, dtype=value.dtype,
-            )
-            return torch.cat((value[:, commit_length:], tail_state), dim=1)
-
-        if hasattr(sampler, "initial_entropy"):
-            unknown_entropy = float(sampler.initial_entropy)
-        elif hasattr(sampler, "vocab_size"):
-            unknown_entropy = math.log(sampler.vocab_size)
-        else:
-            unknown_entropy = float(state.entropy.max())
-
-        return ModilifyMk2RollingState(
-            canvas=canvas,
-            confidence=shift(state.confidence),
-            entropy=shift(state.entropy, unknown_entropy),
-            age=shift(state.age),
-            latent_state=state.latent_state.shift(
-                commit_length, entropy_fill_value=unknown_entropy
-            ),
-            history=state.history.shift(commit_length, entropy_fill_value=unknown_entropy),
-            tape=state.tape,
-        )
-
-    @staticmethod
-    def _merge_state_rows(
-        previous: ModilifyMk2RollingState,
-        updated: ModilifyMk2RollingState,
-        update_mask: torch.BoolTensor,
-    ) -> ModilifyMk2RollingState:
-        """Keep inactive batch rows bit-identical while active rows advance."""
-
-        def choose(old: torch.Tensor, new: torch.Tensor) -> torch.Tensor:
-            mask = update_mask.view(update_mask.shape[0], *([1] * (old.ndim - 1)))
-            return torch.where(mask, new, old)
-
-        old_latent = previous.latent_state
-        new_latent = updated.latent_state
-        latent = LatentDeliberationState(
-            memory_slots=choose(old_latent.memory_slots, new_latent.memory_slots),
-            confidence=choose(old_latent.confidence, new_latent.confidence),
-            entropy=choose(old_latent.entropy, new_latent.entropy),
-            age=choose(old_latent.age, new_latent.age),
-            token_changed=choose(old_latent.token_changed, new_latent.token_changed),
-            confidence_delta=choose(old_latent.confidence_delta, new_latent.confidence_delta),
-            entropy_delta=choose(old_latent.entropy_delta, new_latent.entropy_delta),
-            ponder_steps=choose(old_latent.ponder_steps, new_latent.ponder_steps),
-            stagnation_steps=choose(old_latent.stagnation_steps, new_latent.stagnation_steps),
-        )
-        return ModilifyMk2RollingState(
-            canvas=choose(previous.canvas, updated.canvas),
-            confidence=choose(previous.confidence, updated.confidence),
-            entropy=choose(previous.entropy, updated.entropy),
-            age=choose(previous.age, updated.age),
-            latent_state=latent,
-            history=choose_trajectory_history(
-                previous.history, updated.history, update_mask
-            ),
-            tape=choose_trajectory_tape(previous.tape, updated.tape, update_mask),
-        )
-
     def _write_committed_memory(
         self,
         *,
-        previous_history: TrajectoryHistory,
         next_state: ModilifyMk2RollingState,
         working_state: torch.Tensor,
-        history_projected: torch.Tensor,
         heavy_hidden: torch.Tensor,
+        commit_token_ids: torch.LongTensor,
         commit_lengths: torch.Tensor,
         prefix_lengths: torch.Tensor,
         commit_reason: torch.Tensor,
-        **kwargs: Any,
+        canvas_head: torch.Tensor | None = None,
     ) -> ModilifyMk2RollingState:
         if not bool(commit_lengths.gt(0).any()):
             return next_state
-        memory, _diagnostics = self.latent_deliberation.commit_write(
+        batch, canvas = commit_token_ids.shape
+        if working_state.shape[:2] != (batch, canvas):
+            raise ValueError("Committed token IDs must match the unshifted canvas.")
+        max_commit = min(int(commit_lengths.max()), canvas)
+        if canvas_head is None:
+            selected_ids = commit_token_ids[:, :max_commit]
+        else:
+            head = canvas_head.to(device=commit_token_ids.device, dtype=torch.long).view(batch, 1)
+            physical = (head + torch.arange(max_commit, device=commit_token_ids.device)) % canvas
+            selected_ids = commit_token_ids.gather(1, physical)
+        with torch.no_grad():
+            committed_token_embeddings = self.model.decoder.embed_tokens(selected_ids)
+        memory = self.latent_deliberation.commit_write(
             memory=next_state.latent_state.memory_slots,
             working_state=working_state,
-            history=previous_history,
-            history_projected=history_projected,
+
             heavy_hidden=heavy_hidden,
+            committed_token_embeddings=committed_token_embeddings,
             commit_lengths=commit_lengths,
-            prefix_lengths=prefix_lengths,
             commit_reason=commit_reason,
+            canvas_head=canvas_head,
         )
         return replace(
             next_state,
-            latent_state=replace(next_state.latent_state, memory_slots=memory),
+            latent_state=replace(
+                next_state.latent_state, memory_slots=memory,
+                gdn2=replace(next_state.latent_state.gdn2, persistent=memory),
+            ),
         )
 
     @staticmethod
@@ -775,6 +486,8 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
         commit_lengths: torch.LongTensor,
         sampler: NoiseCanvasSampler,
         generators: Sequence[torch.Generator] | None = None,
+        remaining_lengths: torch.LongTensor | None = None,
+        pad_token_id: int = 0,
     ) -> ModilifyMk2RollingState:
         """Shift every rolling row by its own committed prefix length."""
 
@@ -788,12 +501,9 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 # Seeded generation deliberately advances every active request
                 # once per denoise step, independent of the other active rows.
                 for generator in generators:
-                    try:
-                        sampler.initialize_canvas(
-                            1, state.canvas.device, generators=[generator]
-                        )
-                    except TypeError:
-                        sampler.initialize_canvas(1, state.canvas.device)
+                    sampler.initialize_canvas(
+                        1, state.canvas.device, generators=[generator]
+                    )
             return state
         positions = torch.arange(canvas_length, device=state.canvas.device)[None, :]
         source = positions + commit_lengths[:, None]
@@ -820,18 +530,21 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
             for row, (commit_length, generator) in enumerate(
                 zip(commit_lengths.detach().cpu().tolist(), generators, strict=True)
             ):
-                try:
-                    sampled = sampler.initialize_canvas(
-                        1,
-                        state.canvas.device,
-                        generators=[generator],
-                    )
-                except TypeError:
-                    # Preserve compatibility with deterministic test/custom
-                    # samplers written before per-request RNG was introduced.
-                    sampled = sampler.initialize_canvas(1, state.canvas.device)
+                sampled = sampler.initialize_canvas(
+                    1,
+                    state.canvas.device,
+                    generators=[generator],
+                )
                 tail[row] = sampled[0]
         canvas = torch.cat((state.canvas, tail), dim=1).gather(1, source)
+        if remaining_lengths is not None:
+            if remaining_lengths.shape != (batch_size,):
+                raise ValueError("Remaining lengths must have shape [batch].")
+            canvas = canvas.masked_fill(
+                (positions >= (canvas_length - commit_lengths)[:, None])
+                & (positions >= remaining_lengths[:, None]),
+                pad_token_id,
+            )
         unknown_entropy = float(sampler.initial_entropy)
         latent = state.latent_state
         committed = commit_lengths.gt(0)
@@ -839,25 +552,21 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
             memory_slots=latent.memory_slots.clone(),
             confidence=shift(latent.confidence),
             entropy=shift(latent.entropy, unknown_entropy),
-            age=shift(latent.age),
-            token_changed=shift(latent.token_changed),
-            confidence_delta=shift(latent.confidence_delta),
-            entropy_delta=shift(latent.entropy_delta),
             ponder_steps=torch.where(
                 committed, torch.zeros_like(latent.ponder_steps), latent.ponder_steps
             ),
             stagnation_steps=torch.where(
                 committed, torch.zeros_like(latent.stagnation_steps), latent.stagnation_steps
             ),
+            gdn2=latent.gdn2.shift(commit_lengths),
         )
         return ModilifyMk2RollingState(
             canvas=canvas,
             confidence=shift(state.confidence),
             entropy=shift(state.entropy, unknown_entropy),
-            age=shift(state.age),
             latent_state=shifted_latent,
-            history=state.history.shift(commit_lengths, entropy_fill_value=unknown_entropy),
-            tape=state.tape,
+
+            head=state.head,
         )
 
     @torch.inference_mode()
@@ -868,7 +577,6 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
         streamer: BaseStreamer | None = None,
         generation_config: ModilifyMk2GenerationConfig | None = None,
         logits_processor: LogitsProcessorList | None = None,
-        denoise_trace_callback: Callable[[dict[str, object]], None] | None = None,
         **kwargs,
     ) -> ModilifyMk2GenerationOutput:
         request_seeds = kwargs.pop("seeds", None)
@@ -912,8 +620,6 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
             sampling_generators = None
         if batch_size > 1 and streamer is not None:
             raise ValueError("ModilifyMk2 streamers currently support batch size 1 only.")
-        if batch_size > 1 and denoise_trace_callback is not None:
-            raise ValueError("ModilifyMk2 denoise tracing currently supports batch size 1 only.")
         if batch_size > 1 and past_key_values is not None:
             raise ValueError("Batched ModilifyMk2 generation requires a fresh KV cache.")
         if batch_size > 1:
@@ -976,6 +682,7 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
         _, max_new_tokens = self._prepare_generated_length(
             generation_config, cached_length + input_width
         )
+
         max_iterations = deterministic_episode_iteration_bound(
             torch.tensor([max_new_tokens]),
             max_ponder_steps=generation_config.max_ponder_steps,
@@ -1016,40 +723,21 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
             prompt_positions = input_mask.long().cumsum(dim=-1).sub(1).clamp_min(0).to(torch.int32)
         logical_lengths = cache_attention_mask.long().sum(dim=-1)
         if input_width:
-            encoder_keys = ("pixel_values", "mm_token_type_ids", "image_position_ids")
-            encoder_kwargs = {
-                key: model_kwargs.pop(key)
-                for key in encoder_keys
-                if key in model_kwargs
-            }
             past_key_values = self.model.encoder(
                 input_ids=input_ids,
                 attention_mask=cache_attention_mask,
                 past_key_values=past_key_values,
                 position_ids=prompt_positions,
-                **encoder_kwargs,
             ).past_key_values
 
         sampler = self._prepare_sampler(generation_config, canvas_length)
         latent = LatentDeliberationState.empty(
             batch_size=batch_size, canvas_length=canvas_length,
-            latent_dim=self.config.latent_dim, memory_slots=self.config.latent_memory_slots,
-            device=device, dtype=dtype,
-        )
-        history = TrajectoryHistory.empty(
-            batch_size=batch_size,
-            canvas_length=canvas_length,
-            hidden_size=self.config.text_config.hidden_size,
-            history_length=self.config.latent_history_length,
             device=device,
-            dtype=dtype,
         )
-        try:
-            initial_canvas = sampler.initialize_canvas(
-                batch_size, device, generators=sampling_generators
-            )
-        except TypeError:
-            initial_canvas = sampler.initialize_canvas(batch_size, device)
+        initial_canvas = sampler.initialize_canvas(
+            batch_size, device, generators=sampling_generators
+        )
         state = ModilifyMk2RollingState(
             canvas=initial_canvas,
             confidence=torch.zeros(
@@ -1059,17 +747,9 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 (batch_size, canvas_length), math.log(self.config.text_config.vocab_size),
                 device=device, dtype=torch.float32,
             ),
-            age=torch.zeros(
-                batch_size, canvas_length, device=device, dtype=torch.int32
-            ),
             latent_state=latent,
-            history=history,
-            tape=empty_trajectory_tape(
-                batch_size=batch_size,
-                config=self.config,
-                device=device,
-                dtype=dtype,
-            ),
+
+            head=torch.zeros(batch_size, device=device, dtype=torch.long),
         )
         turn_end = (
             self.config.turn_end_token_id
@@ -1090,6 +770,13 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
         if isinstance(pad_token_id, (list, tuple)):
             pad_token_id = pad_token_id[0]
         pad_token_id = int(0 if pad_token_id is None else pad_token_id)
+        state = replace(
+            state,
+            canvas=state.canvas.masked_fill(
+                torch.arange(canvas_length, device=device)[None, :] >= max_new_tokens,
+                pad_token_id,
+            ),
+        )
         excluded_repetition_token_ids = _flatten_token_ids(
             generation_config.repetition_penalty_exclude_token_ids,
             generation_config.pad_token_id,
@@ -1122,7 +809,6 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
         jumps = torch.zeros_like(committed)
         forced_jump_tokens = torch.zeros_like(committed)
         shifts = torch.zeros_like(committed)
-        retention_scores = torch.zeros(batch_size, dtype=torch.float32, device=device)
         stop_codes = torch.zeros_like(committed)
         active_rows = torch.ones(batch_size, dtype=torch.bool, device=device)
         canvas_positions = torch.arange(canvas_length, device=device)[None, :]
@@ -1130,8 +816,6 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
             streamer.put(input_ids.cpu())
 
         while bool(active_rows.any()):
-            started = time.perf_counter()
-            prefix_length = past_key_values.get_seq_length()
             decoder_positions = (
                 logical_lengths[:, None]
                 + torch.arange(canvas_length, device=device)[None, :]
@@ -1153,13 +837,11 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 input_ids=None, past_key_values=past_key_values,
                 decoder_input_ids=state.canvas,
                 previous_confidence=state.confidence, previous_entropy=state.entropy,
-                token_age=state.age, latent_state=state.latent_state,
-                history=state.history,
-                tape=state.tape,
+                latent_state=state.latent_state,
+
                 decoder_position_ids=decoder_positions, decoder_read_cache=True,
                 decoder_attention_mask=decoder_attention_mask,
                 compact_vocab=True,
-                denoise_temperature=generation_config.denoise_temperature,
                 repetition_token_mask=repetition_history,
                 repetition_penalty=repetition_penalty,
                 sampling_generators=sampling_generators,
@@ -1186,43 +868,48 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 output.next_latent_state,
                 confidence=next_confidence.detach().float(),
                 entropy=token_entropy.detach().float(),
-                age=state.age + 1,
-                token_changed=next_canvas.ne(state.canvas).detach().float(),
-                confidence_delta=next_confidence.detach().float() - state.confidence,
-                entropy_delta=token_entropy.detach().float() - state.entropy,
             )
             remaining = torch.tensor(
                 max_new_tokens, device=device, dtype=torch.long
             ).sub(committed)
             remaining_canvas = remaining[:, None].gt(canvas_positions)
-            tape_probes, tape_valid = self.latent_deliberation.encode_tape_frame(
-                output.heavy_hidden_state, remaining_canvas
+            next_latent = self.latent_deliberation.observe_state(
+                next_latent, output.heavy_hidden_state, output.working_state,
+                remaining_canvas, state.head,
             )
             next_state = ModilifyMk2RollingState(
                 canvas=next_canvas, confidence=next_confidence,
-                entropy=token_entropy, age=state.age + 1,
+                entropy=token_entropy,
                 latent_state=next_latent,
-                history=state.history.append(
-                    output.heavy_hidden_state,
-                    next_confidence,
-                    token_entropy,
-                    next_canvas.ne(state.canvas).detach().float(),
-                    live_mask=remaining_canvas,
-                ),
-                tape=state.tape.append(tape_probes, tape_valid),
+
+                head=state.head,
             )
-            next_state = self._merge_state_rows(state, next_state, active_rows)
             normal_failure_rate = fused_commit_failure_rate(
                 proposal_confidence, token_entropy,
-                vocab_size=self.config.text_config.vocab_size,
+                entropy_weight=self.config.commit_entropy_weight,
+                confidence_power=self.config.commit_confidence_power,
+                top_k=getattr(self.config, "commit_top_k", None),
+                min_p=getattr(self.config, "commit_min_p", None),
+                target_confidence=getattr(self.config, "commit_target_confidence", None),
+                failure_budget=float(self.config.commit_failure_budget),
             )
             jump_failure_rate = fused_commit_failure_rate(
                 greedy_confidence, token_entropy,
-                vocab_size=self.config.text_config.vocab_size,
+                entropy_weight=self.config.commit_entropy_weight,
+                confidence_power=self.config.commit_confidence_power,
+                top_k=getattr(self.config, "commit_top_k", None),
+                min_p=getattr(self.config, "commit_min_p", None),
+                target_confidence=getattr(self.config, "commit_target_confidence", None),
+                failure_budget=float(self.config.commit_failure_budget),
             )
             previous_failure_rate = fused_commit_failure_rate(
                 state.confidence, state.entropy,
-                vocab_size=self.config.text_config.vocab_size,
+                entropy_weight=self.config.commit_entropy_weight,
+                confidence_power=self.config.commit_confidence_power,
+                top_k=getattr(self.config, "commit_top_k", None),
+                min_p=getattr(self.config, "commit_min_p", None),
+                target_confidence=getattr(self.config, "commit_target_confidence", None),
+                failure_budget=float(self.config.commit_failure_budget),
             )
             policy_decision = select_commit_lengths(
                 sampled_token_ids=proposal,
@@ -1234,15 +921,12 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 stagnation_steps=state.latent_state.stagnation_steps,
                 active_rows=active_rows,
                 remaining_lengths=remaining,
-                failure_budget=generation_config.commit_failure_budget,
-                jump_failure_budget=generation_config.jump_failure_budget,
+                failure_budget=self.config.commit_failure_budget,
                 stop_token_id=stop_token_ids,
                 max_ponder_steps=generation_config.max_ponder_steps,
                 stagnation_threshold=generation_config.jump_on_no_progress_after,
                 min_progress=generation_config.min_trajectory_progress,
             )
-            normal_commit = policy_decision.normal_lengths
-            policy_prefix_mask = canvas_positions.lt(normal_commit[:, None])
             next_ponder = policy_decision.ponder_steps
             next_stagnation = policy_decision.stagnation_steps
             commit_lengths = policy_decision.commit_lengths
@@ -1319,21 +1003,17 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 ).past_key_values
                 if streamer is not None:
                     streamer.put(committed_block.cpu())
-            committed += commit_lengths
-            logical_lengths += commit_lengths
             committed_rows = commit_lengths.gt(0)
             shifts += committed_rows.long()
             if (
-                output.history_projected is None
-                or output.working_state is None
+                output.working_state is None
             ):
                 raise RuntimeError("Forward did not return working trajectory features.")
             next_state = self._write_committed_memory(
-                previous_history=state.history,
                 next_state=next_state,
                 working_state=output.working_state,
-                history_projected=output.history_projected,
                 heavy_hidden=output.heavy_hidden_state,
+                commit_token_ids=commit_token_ids,
                 commit_lengths=commit_lengths,
                 prefix_lengths=logical_lengths,
                 commit_reason=infer_commit_reason(
@@ -1348,9 +1028,12 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 commit_lengths,
                 sampler,
                 generators=sampling_generators,
+                remaining_lengths=remaining - commit_lengths,
+                pad_token_id=pad_token_id,
             )
-            retention_scores += committed_rows.float()
             state = shifted
+            committed += commit_lengths
+            logical_lengths += commit_lengths
 
             turn_hits = (
                 commit_token_ids.eq(turn_end) & commit_positions
@@ -1388,19 +1071,6 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
                 torch.full_like(stop_codes, 5),
                 stop_codes,
             )
-            if denoise_trace_callback is not None:
-                denoise_trace_callback(build_denoise_trace_event(
-                    denoise_step=int(denoise_steps[0]),
-                    prefix_length=prefix_length, committed_before=int(before[0]),
-                    committed_after=int(committed[0]),
-                    no_progress_steps=int(next_stagnation[0]),
-                    policy_prefix_mask=policy_prefix_mask,
-                    commit_length=int(commit_lengths[0]),
-                    ponder_fallback=bool(jump_rows[0]), state=next_state, proposal=proposal,
-                    committed_token_ids=commit_token_ids[:, :commit_width],
-                    step_elapsed_seconds=time.perf_counter() - started,
-                    latent_residual_diagnostics=output.latent_residual_diagnostics,
-                ))
             active_rows = stop_codes.eq(0)
 
         output_width = int(committed.max())
@@ -1419,8 +1089,6 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
         )
         tokens_per_forward = committed.float() / denoise_steps.clamp_min(1).float()
         average_commit_len = committed.float() / shifts.clamp_min(1).float()
-        latent_memory_norm = state.latent_state.memory_slots.float().norm(dim=-1).mean(dim=-1)
-        state_retention_score = retention_scores / shifts.clamp_min(1).float()
 
         def scalar_or_tensor(value: torch.Tensor, *, floating: bool = False):
             if batch_size > 1:
@@ -1439,17 +1107,6 @@ class ModilifyMk2GenerationMixin(DiffusionGemmaGenerationMixin):
             no_progress_steps=scalar_or_tensor(state.latent_state.stagnation_steps),
             jump_count=scalar_or_tensor(jumps),
             forced_jump_bad_count=scalar_or_tensor(forced_jump_tokens),
-            heavy_forward_count=scalar_or_tensor(denoise_steps),
-            latent_context_update_count=scalar_or_tensor(denoise_steps),
             average_commit_len=scalar_or_tensor(average_commit_len, floating=True),
             state_shift_count=scalar_or_tensor(shifts),
-            latent_memory_norm=scalar_or_tensor(latent_memory_norm, floating=True),
-            state_retention_score=scalar_or_tensor(state_retention_score, floating=True),
         )
-
-
-__all__ = [
-    "ModilifyMk2GenerationConfig", "ModilifyMk2GenerationMixin", "ModilifyMk2GenerationOutput",
-    "ModilifyMk2RollingState", "NoiseCanvasSampler",
-    "build_denoise_trace_event", "qualified_single_turn_end_lengths",
-]

latent_deliberation.py

@@ -1,45 +1,31 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
-"""Dual-timescale Transformer memory for Modilify Mk2 inference.
-
-Working trajectory state is recomputed every denoise step from the packed
-personal history. Persistent slots are commit-invariant and mutate only in
-TransformerCommitWriter.
-"""
+"""Schema25 Torch GDN2 trajectory state, processor, and memory readers."""
 
 from __future__ import annotations
 
 from collections.abc import Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+import weakref
 import math
 
 import torch
 from torch import nn
 from torch.nn import functional as F
 
+from .gdn2_trajectory import GDN2TrajectoryMemory, GDN2TrajectoryState
 
-_AGE_MAX = 4096
-_PONDER_MAX = 1024
-_STAGNATION_MAX = 1024
-_METADATA_HIDDEN = 64
-_FILM_RANK = 64
-_FOURIER_WAVES = 4
-_RETIREMENT_FRAMES = 4
-_KEY_META_DIM = 13
-_QUERY_META_DIM = 8 + 6
-_ROW_META_DIM = 2
-_HISTORY_VIEWS = 4
-_EXPERIENCE_ROLES = 3
-_EXPERIENCE_CANVAS_STRIPE = 8
-_GATE_BIAS = -3.0
 
 COMMIT_REASON_NONE = 0
 COMMIT_REASON_NORMAL = 1
 COMMIT_REASON_FORCED_JUMP = 2
 COMMIT_REASON_TERMINAL = 3
-COMMIT_REASON_FALLBACK = 4
-COMMIT_REASON_TRAINING_RANDOM = 5
-COMMIT_REASON_COUNT = 6
+
+
+def _sdpa_mask_value(dtype: torch.dtype) -> float:
+    """Additive SDPA mask that stays finite on MPS fp16/bf16."""
+
+    if dtype in (torch.float16, torch.bfloat16):
+        return -1.0e4
+    return -1.0e9
 
 
 def _fp32_scaled_dot_product_attention(
@@ -93,197 +79,15 @@ def _fp32_scaled_dot_product_attention(
     return output.to(dtype=output_dtype)
 
 
-@dataclass
-class TrajectoryHistory:
-    """Detached ring of recent canvas heavies. Canvas is dimension 2."""
-
-    hidden: torch.Tensor
-    confidence: torch.Tensor
-    entropy: torch.Tensor
-    token_changed: torch.Tensor
-    valid: torch.Tensor
-
-    @classmethod
-    def empty(
-        cls,
-        *,
-        batch_size: int,
-        canvas_length: int,
-        hidden_size: int,
-        history_length: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> "TrajectoryHistory":
-        return cls(
-            hidden=torch.zeros(
-                batch_size, history_length, canvas_length, hidden_size,
-                device=device, dtype=dtype,
-            ),
-            confidence=torch.zeros(
-                batch_size, history_length, canvas_length,
-                device=device, dtype=torch.float32,
-            ),
-            entropy=torch.zeros(
-                batch_size, history_length, canvas_length,
-                device=device, dtype=torch.float32,
-            ),
-            token_changed=torch.zeros(
-                batch_size, history_length, canvas_length,
-                device=device, dtype=torch.float32,
-            ),
-            valid=torch.zeros(
-                batch_size, history_length, canvas_length,
-                device=device, dtype=torch.bool,
-            ),
-        )
-
-    def detach(self) -> "TrajectoryHistory":
-        return TrajectoryHistory(
-            hidden=self.hidden.detach(),
-            confidence=self.confidence.detach(),
-            entropy=self.entropy.detach(),
-            token_changed=self.token_changed.detach(),
-            valid=self.valid.detach(),
-        )
-
-    def append(
-        self,
-        hidden: torch.Tensor,
-        confidence: torch.Tensor,
-        entropy: torch.Tensor,
-        token_changed: torch.Tensor,
-        live_mask: torch.Tensor | None = None,
-    ) -> "TrajectoryHistory":
-        # Truncated-BPTT observation. Replay uses temporal context / committed
-        # memory, not this ring; keeping frames live would retain every heavy
-        # decoder graph across the chunk.
-        frame = hidden.detach()
-        if live_mask is None:
-            newest_valid = torch.ones(
-                self.valid.shape[0],
-                self.valid.shape[2],
-                device=self.valid.device,
-                dtype=torch.bool,
-            )
-        else:
-            newest_valid = live_mask.to(device=self.valid.device, dtype=torch.bool)
-            if newest_valid.shape != self.valid[:, 0].shape:
-                raise ValueError("`live_mask` must have shape [batch, canvas].")
-        return TrajectoryHistory(
-            hidden=torch.cat((self.hidden[:, 1:], frame.unsqueeze(1)), dim=1),
-            confidence=torch.cat(
-                (self.confidence[:, 1:], confidence.detach().float().unsqueeze(1)),
-                dim=1,
-            ),
-            entropy=torch.cat(
-                (self.entropy[:, 1:], entropy.detach().float().unsqueeze(1)),
-                dim=1,
-            ),
-            token_changed=torch.cat(
-                (self.token_changed[:, 1:], token_changed.detach().float().unsqueeze(1)),
-                dim=1,
-            ),
-            valid=torch.cat((self.valid[:, 1:], newest_valid.unsqueeze(1)), dim=1),
-        )
-
-    def shift(
-        self,
-        commit_lengths: torch.Tensor | int,
-        *,
-        entropy_fill_value: float = 0.0,
-    ) -> "TrajectoryHistory":
-        batch, history_length, canvas_length, _hidden = self.hidden.shape
-        if isinstance(commit_lengths, int):
-            lengths = torch.full(
-                (batch,), commit_lengths, device=self.hidden.device, dtype=torch.long
-            )
-        else:
-            lengths = commit_lengths.to(device=self.hidden.device, dtype=torch.long)
-        if bool((lengths <= 0).all()):
-            return self.detach()
-        positions = torch.arange(canvas_length, device=self.hidden.device)[None, :]
-        source = positions + lengths[:, None]
-        retained = source.lt(canvas_length)
-
-        def shifted(tensor: torch.Tensor, fill_value: float | int = 0) -> torch.Tensor:
-            index = source.clamp_max(canvas_length - 1)
-            extra = tensor.ndim - 3
-            view = index.view(batch, 1, canvas_length, *([1] * extra)).expand_as(tensor)
-            gathered = tensor.gather(2, view)
-            fill = torch.as_tensor(fill_value, device=tensor.device, dtype=tensor.dtype)
-            mask = retained.view(batch, 1, canvas_length, *([1] * extra))
-            return torch.where(mask, gathered, fill)
-
-        return TrajectoryHistory(
-            hidden=shifted(self.hidden),
-            confidence=shifted(self.confidence),
-            entropy=shifted(self.entropy, entropy_fill_value),
-            token_changed=shifted(self.token_changed),
-            valid=shifted(self.valid, False),
-        )
-
-
-@dataclass
-class TrajectoryTape:
-    """Row-level denoise snapshots. Time axis does not follow canvas shift."""
-
-    probes: torch.Tensor
-    valid: torch.Tensor
-
-    @classmethod
-    def empty(
-        cls,
-        *,
-        batch_size: int,
-        tape_length: int,
-        num_probes: int,
-        probe_dim: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> "TrajectoryTape":
-        return cls(
-            probes=torch.zeros(
-                batch_size, tape_length, num_probes, probe_dim,
-                device=device, dtype=dtype,
-            ),
-            valid=torch.zeros(
-                batch_size, tape_length, device=device, dtype=torch.bool,
-            ),
-        )
-
-    def detach(self) -> "TrajectoryTape":
-        return TrajectoryTape(probes=self.probes.detach(), valid=self.valid.detach())
-
-    def append(self, probes: torch.Tensor, valid: torch.Tensor) -> "TrajectoryTape":
-        # Canvas snapshots are detached before pooling. Keep the pool graph so
-        # the next denoise's tape read can train the compressor; TBPTT still
-        # cuts at `detach()`.
-        flag = valid.to(device=self.valid.device, dtype=torch.bool)
-        if flag.ndim == 0:
-            flag = flag.expand(self.valid.shape[0])
-        if flag.shape != self.valid[:, 0].shape:
-            raise ValueError("Tape frame validity must have shape [batch].")
-        if probes.shape[0] != self.probes.shape[0] or probes.shape[-2:] != self.probes.shape[-2:]:
-            raise ValueError("Tape probes do not match the ring.")
-        return TrajectoryTape(
-            probes=torch.cat((self.probes[:, 1:], probes.unsqueeze(1)), dim=1),
-            valid=torch.cat((self.valid[:, 1:], flag.unsqueeze(1)), dim=1),
-        )
-
-
 @dataclass
 class LatentDeliberationState:
     """Persistent slots plus per-canvas trajectory clocks. No token latents."""
-
     memory_slots: torch.Tensor
     confidence: torch.Tensor
     entropy: torch.Tensor
-    age: torch.Tensor
-    token_changed: torch.Tensor
-    confidence_delta: torch.Tensor
-    entropy_delta: torch.Tensor
     ponder_steps: torch.Tensor
     stagnation_steps: torch.Tensor
+    gdn2: GDN2TrajectoryState
 
     @classmethod
     def empty(
@@ -291,84 +95,35 @@ class LatentDeliberationState:
         *,
         batch_size: int,
         canvas_length: int,
-        latent_dim: int,
-        memory_slots: int,
         device: torch.device,
-        dtype: torch.dtype,
     ) -> "LatentDeliberationState":
+        persistent = torch.zeros(batch_size, 16, 128, 128, device=device, dtype=torch.float32)
         return cls(
-            memory_slots=torch.zeros(
-                batch_size, memory_slots, latent_dim, device=device, dtype=dtype
-            ),
+            memory_slots=persistent,
             confidence=torch.zeros(
                 batch_size, canvas_length, device=device, dtype=torch.float32
             ),
             entropy=torch.zeros(
                 batch_size, canvas_length, device=device, dtype=torch.float32
             ),
-            age=torch.zeros(
-                batch_size, canvas_length, device=device, dtype=torch.int32
-            ),
-            token_changed=torch.zeros(
-                batch_size, canvas_length, device=device, dtype=torch.float32
-            ),
-            confidence_delta=torch.zeros(
-                batch_size, canvas_length, device=device, dtype=torch.float32
-            ),
-            entropy_delta=torch.zeros(
-                batch_size, canvas_length, device=device, dtype=torch.float32
-            ),
             ponder_steps=torch.zeros(batch_size, device=device, dtype=torch.int32),
             stagnation_steps=torch.zeros(batch_size, device=device, dtype=torch.int32),
-        )
-
-    def detach(self) -> "LatentDeliberationState":
-        return LatentDeliberationState(
-            memory_slots=self.memory_slots.detach(),
-            confidence=self.confidence.detach(),
-            entropy=self.entropy.detach(),
-            age=self.age.detach(),
-            token_changed=self.token_changed.detach(),
-            confidence_delta=self.confidence_delta.detach(),
-            entropy_delta=self.entropy_delta.detach(),
-            ponder_steps=self.ponder_steps.detach(),
-            stagnation_steps=self.stagnation_steps.detach(),
-        )
-
-    def shift(
-        self, committed: int, *, entropy_fill_value: float = 0.0
-    ) -> "LatentDeliberationState":
-        canvas_length = self.confidence.shape[1]
-        if not 0 <= committed <= canvas_length:
-            raise ValueError("`committed` must be in [0, canvas_length].")
-        if committed == 0:
-            return self.detach()
-
-        def shifted(tensor: torch.Tensor, fill_value: float | int = 0) -> torch.Tensor:
-            result = torch.full_like(tensor, fill_value)
-            if committed < canvas_length:
-                result[:, : canvas_length - committed] = tensor[:, committed:]
-            return result
-
-        return LatentDeliberationState(
-            memory_slots=self.memory_slots.clone(),
-            confidence=shifted(self.confidence),
-            entropy=shifted(self.entropy, entropy_fill_value),
-            age=shifted(self.age),
-            token_changed=shifted(self.token_changed),
-            confidence_delta=shifted(self.confidence_delta),
-            entropy_delta=shifted(self.entropy_delta),
-            ponder_steps=torch.zeros_like(self.ponder_steps),
-            stagnation_steps=torch.zeros_like(self.stagnation_steps),
+            gdn2=GDN2TrajectoryState(
+                cells=torch.zeros(batch_size, canvas_length, 16, 64, 64,
+                                  device=device, dtype=torch.float32),
+                row=torch.zeros(batch_size, 16, 64, 64,
+                                device=device, dtype=torch.float32),
+                persistent=persistent,
+                seen=torch.zeros(batch_size, canvas_length,
+                                 device=device, dtype=torch.bool),
+            ),
         )
 
 
 @dataclass
 class LatentProcessorOutput:
     context: torch.Tensor
-    working_state: torch.Tensor
     state: LatentDeliberationState
-    history_projected: torch.Tensor
 
 
 def advance_trajectory_clocks(
@@ -377,13 +132,9 @@ def advance_trajectory_clocks(
     *,
     commit_lengths: torch.LongTensor,
     active_rows: torch.BoolTensor,
-    progress_scores: torch.Tensor | None = None,
-    min_progress: float = 0.0,
 ) -> tuple[torch.IntTensor, torch.IntTensor]:
     """Advance useful-ponder and stagnation clocks for each row."""
 
-    if min_progress < 0:
-        raise ValueError("`min_progress` must be non-negative.")
     if not (
         ponder_steps.shape == stagnation_steps.shape == commit_lengths.shape
         == active_rows.shape
@@ -411,85 +162,12 @@ def should_force_trajectory_jump(
     ponder_steps: torch.Tensor | None = None,
     max_ponder_steps: int | None = None,
 ) -> torch.BoolTensor:
-    """Determine whether to force a trajectory JUMP for each row.
-
-    JUMP is triggered when stagnation_steps reaches or exceeds stagnation_threshold
-    (default 12) and the current single-step progress is not strictly greater than
-    min_progress (default 0.005).
-    """
-    if stagnation_threshold <= 0:
-        raise ValueError("`stagnation_threshold` must be positive.")
-    if min_progress < 0:
-        raise ValueError("`min_progress` must be non-negative.")
-
-    if progress_scores is None:
-        stagnation_jump = stagnation_steps.ge(stagnation_threshold)
-    else:
-        if progress_scores.shape != stagnation_steps.shape:
-            raise ValueError("`progress_scores` must share shape with `stagnation_steps`.")
-        stagnation_jump = stagnation_steps.ge(stagnation_threshold) & progress_scores.le(
-            float(min_progress)
-        )
-
+    jump = stagnation_steps.ge(stagnation_threshold)
+    if progress_scores is not None:
+        jump = jump & progress_scores.le(float(min_progress))
     if ponder_steps is not None and max_ponder_steps is not None and max_ponder_steps > 0:
-        ponder_jump = ponder_steps.ge(max_ponder_steps)
-        return (ponder_jump | stagnation_jump).to(torch.bool)
-
-    return stagnation_jump.to(torch.bool)
-
-
-def _logit01(values: torch.Tensor) -> torch.Tensor:
-    clipped = values.clamp(1.0e-6, 1.0 - 1.0e-6)
-    return torch.log(clipped) - torch.log1p(-clipped)
-
-
-def _renorm_confidence(confidence: torch.Tensor) -> torch.Tensor:
-    return _logit01(confidence).clamp(-8.0, 8.0) / 8.0
-
-
-def _canvas_fourier(canvas_length: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-    positions = torch.arange(canvas_length, device=device, dtype=dtype) / max(canvas_length, 1)
-    features = []
-    for wave in range(_FOURIER_WAVES):
-        angle = (2.0 ** wave) * math.pi * positions
-        features.append(torch.sin(angle))
-        features.append(torch.cos(angle))
-    return torch.stack(features, dim=-1)
-
-
-def _safe_cosine(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
-    left_n = F.normalize(left.float(), dim=-1, eps=1.0e-6)
-    right_n = F.normalize(right.float(), dim=-1, eps=1.0e-6)
-    return (left_n * right_n).sum(dim=-1).clamp(-1.0, 1.0)
-
-
-def _sdpa_mask_value(dtype: torch.dtype) -> float:
-    """Additive SDPA mask that stays finite on MPS fp16/bf16."""
-
-    if dtype in (torch.float16, torch.bfloat16):
-        return -1.0e4
-    return -1.0e9
-
-
-def _apply_rotary(payload: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-    dim = payload.shape[-1]
-    half = dim // 2
-    if half == 0:
-        return payload
-    device = payload.device
-    inv = torch.arange(half, device=device, dtype=torch.float32)
-    inv = 10000.0 ** (-inv / max(half, 1))
-    angle = positions.to(dtype=torch.float32).unsqueeze(-1) * inv
-    cos = angle.cos().to(dtype=payload.dtype)
-    sin = angle.sin().to(dtype=payload.dtype)
-    while cos.ndim < payload.ndim:
-        cos = cos.unsqueeze(1)
-        sin = sin.unsqueeze(1)
-    left, right = payload[..., :half], payload[..., half: half * 2]
-    rotated = torch.cat((left * cos - right * sin, left * sin + right * cos), dim=-1)
-    if dim > half * 2:
-        rotated = torch.cat((rotated, payload[..., half * 2 :]), dim=-1)
-    return rotated
+        jump = jump | ponder_steps.ge(max_ponder_steps)
+    return jump.to(torch.bool)
 
 
 class _RMSNorm(nn.Module):
@@ -500,7 +178,7 @@ class _RMSNorm(nn.Module):
 
     def forward(self, hidden: torch.Tensor) -> torch.Tensor:
         rms = hidden.float().square().mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
-        return (hidden.float() * rms).to(dtype=hidden.dtype) * self.weight.to(dtype=hidden.dtype)
+        return (hidden.float() * rms * self.weight.float()).to(dtype=hidden.dtype)
 
 
 class _SwiGLU(nn.Module):
@@ -514,216 +192,6 @@ class _SwiGLU(nn.Module):
         return self.down(F.silu(self.gate(hidden)) * self.up(hidden))
 
 
-class SharedHistoryProjector(nn.Module):
-    """2816 → rank projection, once per denoise (and once more on the commit tail)."""
-
-    def __init__(self, hidden_size: int, rank: int) -> None:
-        super().__init__()
-        self.norm = _RMSNorm(hidden_size)
-        self.proj = nn.Linear(hidden_size, rank, bias=False)
-
-    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
-        return self.proj(self.norm(hidden))
-
-
-class CanvasProbePool(nn.Module):
-    """Learned queries compress one canvas snapshot into P rank-space probes."""
-
-    def __init__(self, rank: int, num_probes: int, num_heads: int) -> None:
-        super().__init__()
-        if rank % num_heads:
-            raise ValueError("Tape rank must be divisible by heads.")
-        if num_probes <= 0:
-            raise ValueError("`num_probes` must be positive.")
-        self.rank = rank
-        self.num_probes = num_probes
-        self.num_heads = num_heads
-        self.head_dim = rank // num_heads
-        self.queries = nn.Parameter(torch.empty(num_probes, rank))
-        self.k_proj = nn.Linear(rank, rank, bias=False)
-        self.v_proj = nn.Linear(rank, rank, bias=False)
-        self.reset_parameters()
-
-    @torch.no_grad()
-    def reset_parameters(self) -> None:
-        nn.init.normal_(self.queries, mean=0.0, std=0.02)
-
-    def forward(
-        self,
-        projected: torch.Tensor,
-        live_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        batch, canvas, _rank = projected.shape
-        heads = self.num_heads
-        head_dim = self.head_dim
-        query = self.queries.to(dtype=projected.dtype).view(1, self.num_probes, heads, head_dim)
-        query = query.expand(batch, -1, -1, -1).permute(0, 2, 1, 3)
-        keys = self.k_proj(projected).view(batch, canvas, heads, head_dim).transpose(1, 2)
-        values = self.v_proj(projected).view(batch, canvas, heads, head_dim).transpose(1, 2)
-        if live_mask is None:
-            allowed = torch.ones(batch, canvas, device=projected.device, dtype=torch.bool)
-        else:
-            allowed = live_mask.to(device=projected.device, dtype=torch.bool)
-            if allowed.shape != (batch, canvas):
-                raise ValueError("`live_mask` must have shape [batch, canvas].")
-        has_live = allowed.any(dim=-1)
-        safe = allowed.clone()
-        safe[:, 0] = safe[:, 0] | ~has_live
-        # This small P×canvas pool is a poor place to trade stability for BF16:
-        # on MPS the first B16 tape frame can contain NaNs even with finite Q/K/V
-        # and at least one unmasked key per row.  Perform only the attention
-        # reduction in FP32; projections and stored probes retain model dtype.
-        additive = torch.zeros(
-            batch, 1, 1, canvas, device=projected.device, dtype=torch.float32
-        )
-        additive = additive.masked_fill(
-            ~safe.view(batch, 1, 1, canvas), _sdpa_mask_value(torch.float32)
-        )
-        context = _fp32_scaled_dot_product_attention(
-            query, keys, values, attn_mask=additive
-        )
-        probes = context.transpose(1, 2).reshape(batch, self.num_probes, self.rank)
-        return probes * has_live.to(dtype=probes.dtype).view(batch, 1, 1)
-
-
-class SharedPersistentKV(nn.Module):
-    """Frozen-M key/value projection shared across working-processor blocks."""
-
-    def __init__(self, dim: int, num_heads: int, kv_rank: int) -> None:
-        super().__init__()
-        if kv_rank % num_heads:
-            raise ValueError("`kv_rank` must be divisible by `num_heads`.")
-        self.num_heads = num_heads
-        self.kv_rank = kv_rank
-        self.head_dim = kv_rank // num_heads
-        self.address_norm = _RMSNorm(dim)
-        self.value_norm = _RMSNorm(dim)
-        self.k_proj = nn.Linear(dim, kv_rank, bias=False)
-        self.v_proj = nn.Linear(dim, kv_rank, bias=False)
-
-    def forward(
-        self, memory: torch.Tensor, slot_identity: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        batch, slots, _dim = memory.shape
-        keys = self.k_proj(self.address_norm(memory + slot_identity))
-        values = self.v_proj(self.value_norm(memory))
-        keys = keys.view(batch, slots, self.num_heads, self.head_dim).transpose(1, 2)
-        values = values.view(batch, slots, self.num_heads, self.head_dim).transpose(1, 2)
-        return keys, values
-
-
-class _HistoryAttention(nn.Module):
-    """Per-position attention over shared rank-space history keys."""
-
-    def __init__(self, dim: int, num_heads: int, kv_rank: int) -> None:
-        super().__init__()
-        if kv_rank % num_heads:
-            raise ValueError("`kv_rank` must be divisible by `num_heads`.")
-        self.num_heads = num_heads
-        self.kv_rank = kv_rank
-        self.head_dim = kv_rank // num_heads
-        self.q_proj = nn.Linear(dim, kv_rank, bias=False)
-        self.o_proj = nn.Linear(kv_rank, dim, bias=False)
-        self.q_norm = _RMSNorm(dim)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        keys: torch.Tensor,
-        values: torch.Tensor,
-        *,
-        attn_bias: torch.Tensor,
-        key_mask: torch.Tensor,
-    ) -> torch.Tensor:
-        batch, canvas, _dim = query.shape
-        heads = self.num_heads
-        head_dim = self.head_dim
-        query = self.q_proj(self.q_norm(query))
-        query = query.view(batch, canvas, heads, head_dim).transpose(1, 2)
-        if keys.ndim == 5:
-            expected = (batch, heads, canvas, keys.shape[-2], head_dim)
-            if keys.shape != expected or values.shape != expected:
-                raise ValueError("Preformatted history K/V dimensions do not match.")
-            slots = int(keys.shape[-2])
-        else:
-            slots = int(keys.shape[2])
-            keys = keys.view(batch, canvas, slots, heads, head_dim).permute(0, 3, 1, 2, 4)
-            values = values.view(batch, canvas, slots, heads, head_dim).permute(0, 3, 1, 2, 4)
-        query = query.reshape(batch * heads * canvas, 1, head_dim)
-        keys = keys.reshape(batch * heads * canvas, slots, head_dim)
-        values = values.reshape(batch * heads * canvas, slots, head_dim)
-        has_hist = key_mask.any(dim=-1)
-        safe_mask = key_mask.clone()
-        safe_mask[..., 0] = safe_mask[..., 0] | ~has_hist
-        mask = safe_mask.reshape(batch, 1, canvas, slots)
-        mask = mask.expand(-1, heads, -1, -1).reshape(batch * heads * canvas, 1, slots)
-        bias = attn_bias.reshape(batch * heads * canvas, 1, slots)
-        additive = bias.masked_fill(~mask, _sdpa_mask_value(query.dtype))
-        context = _fp32_scaled_dot_product_attention(
-            query, keys, values, attn_mask=additive
-        )
-        context = context.view(batch, heads, canvas, head_dim).transpose(1, 2).reshape(
-            batch, canvas, self.kv_rank
-        )
-        output = self.o_proj(context)
-        keep = has_hist.unsqueeze(-1)
-        return torch.where(keep, output, output.new_zeros(output.shape))
-
-
-class _QueryOutputAttention(nn.Module):
-    """Q/O attention against precomputed K/V."""
-
-    def __init__(self, dim: int, num_heads: int, kv_rank: int) -> None:
-        super().__init__()
-        if kv_rank % num_heads:
-            raise ValueError("`kv_rank` must be divisible by `num_heads`.")
-        self.num_heads = num_heads
-        self.kv_rank = kv_rank
-        self.head_dim = kv_rank // num_heads
-        self.q_proj = nn.Linear(dim, kv_rank, bias=False)
-        self.o_proj = nn.Linear(kv_rank, dim, bias=False)
-        self.q_norm = _RMSNorm(dim)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        keys: torch.Tensor,
-        values: torch.Tensor,
-        attn_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        batch, queries, _dim = query.shape
-        heads = self.num_heads
-        head_dim = self.head_dim
-        query = self.q_proj(self.q_norm(query)).view(batch, queries, heads, head_dim).transpose(1, 2)
-        mask = attn_mask
-        keep_rows = None
-        if mask is not None and mask.ndim == 2:
-            if mask.shape[0] == batch:
-                if mask.dtype == torch.bool:
-                    keep_rows = mask.any(dim=-1)
-                    safe = mask.clone()
-                    safe[:, 0] = safe[:, 0] | ~keep_rows
-                    mask = safe
-                mask = mask.view(batch, 1, 1, mask.shape[-1])
-            else:
-                mask = mask.view(1, 1, queries, keys.shape[-2])
-        elif mask is not None and mask.ndim == 3:
-            mask = mask.unsqueeze(1)
-        if mask is not None and mask.dtype == torch.bool:
-            additive = torch.zeros(
-                mask.shape, device=query.device, dtype=query.dtype
-            )
-            mask = additive.masked_fill(~mask, _sdpa_mask_value(query.dtype))
-        context = _fp32_scaled_dot_product_attention(
-            query, keys, values, attn_mask=mask
-        )
-        context = context.transpose(1, 2).reshape(batch, queries, self.kv_rank)
-        output = self.o_proj(context)
-        if keep_rows is not None:
-            output = output * keep_rows.to(dtype=output.dtype).view(batch, 1, 1)
-        return output
-
-
 class _RankAttention(nn.Module):
     """Sequence attention in a rank-``kv_rank`` subspace, then map back to ``dim``."""
 
@@ -767,86 +235,8 @@ class _RankAttention(nn.Module):
         return self.o_proj(context)
 
 
-class _ProcessorBlock(nn.Module):
-    """History-read canvas state, local or global mixing, read-only slot CA, SwiGLU."""
-
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        local_attention_window: int,
-        kv_rank: int,
-        ffn_dim: int,
-        *,
-        global_attention: bool,
-    ) -> None:
-        super().__init__()
-        self.history_attention = _HistoryAttention(dim, num_heads, kv_rank)
-        self.state_norm = _RMSNorm(dim)
-        self.local_attention = _RankAttention(dim, num_heads, kv_rank)
-        self.local_attention_window = local_attention_window
-        self.global_attention = global_attention
-        self.register_buffer("_local_attention_mask", torch.empty(0), persistent=False)
-        self.token_memory_attention = _QueryOutputAttention(dim, num_heads, kv_rank)
-        self.tape_attention = _QueryOutputAttention(dim, num_heads, kv_rank)
-        self.token_ff_norm = _RMSNorm(dim)
-        self.ff = _SwiGLU(dim, ffn_dim)
-
-    def _local_mask(self, canvas: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-        if (
-            self._local_attention_mask.shape != (canvas, canvas)
-            or self._local_attention_mask.device != device
-            or self._local_attention_mask.dtype != dtype
-        ):
-            positions = torch.arange(canvas, device=device)
-            allowed = (positions[:, None] - positions[None, :]).abs() < self.local_attention_window
-            mask = torch.zeros(canvas, canvas, device=device, dtype=dtype)
-            self._local_attention_mask = mask.masked_fill(~allowed, _sdpa_mask_value(dtype))
-        return self._local_attention_mask
-
-    def forward(
-        self,
-        canvas_state: torch.Tensor,
-        history_keys: torch.Tensor,
-        history_values: torch.Tensor,
-        attn_bias: torch.Tensor,
-        key_mask: torch.Tensor,
-        memory_keys: torch.Tensor,
-        memory_values: torch.Tensor,
-        history_query: torch.Tensor,
-        tape_keys: torch.Tensor | None = None,
-        tape_values: torch.Tensor | None = None,
-        tape_mask: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        canvas_state = canvas_state + self.history_attention(
-            history_query, history_keys, history_values,
-            attn_bias=attn_bias, key_mask=key_mask,
-        )
-        if tape_keys is not None and tape_values is not None:
-            canvas_state = canvas_state + self.tape_attention(
-                self.state_norm(canvas_state), tape_keys, tape_values, attn_mask=tape_mask
-            )
-        normalized = self.state_norm(canvas_state)
-        attn_mask = None if self.global_attention else self._local_mask(
-            canvas_state.shape[1], canvas_state.device, canvas_state.dtype
-        )
-        canvas_state = canvas_state + self.local_attention(
-            normalized, normalized, normalized, attn_mask=attn_mask
-        )
-        normalized = self.state_norm(canvas_state)
-        canvas_state = canvas_state + self.token_memory_attention(
-            normalized, memory_keys, memory_values
-        )
-        return canvas_state + self.ff(self.token_ff_norm(canvas_state))
-
-
 class DecoderMemoryBus(nn.Module):
-    """Sidecar readers on full-attention decoder layers.
-
-    ``alpha`` starts at 0 so the residual is zero. Scale with ``tanh(alpha)``
-    rather than a hard ``where(alpha == 0)`` so the gate stays differentiable.
-    ``o_proj`` is *not* zeroed: that pair was a dead-gradient product.
-    """
+    """Read memory through a per-head gated residual."""
 
     def __init__(
         self,
@@ -893,20 +283,7 @@ class DecoderMemoryBus(nn.Module):
         span = max(2 * max_relative_span - 1, 1)
         self.rel_bias = nn.Parameter(torch.zeros(max(num_heads, 1), span))
         self.max_relative_span = max_relative_span
-        self.enabled = False
-        self.freeze()
-
-    def freeze(self) -> None:
-        self.enabled = False
-        for parameter in self.parameters():
-            parameter.requires_grad_(False)
 
-    def unfreeze(self) -> None:
-        if self.num_readers <= 0:
-            return
-        self.enabled = True
-        for parameter in self.parameters():
-            parameter.requires_grad_(True)
 
     def prepare_kv(
         self,
@@ -915,8 +292,6 @@ class DecoderMemoryBus(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor] | None:
         if self.num_readers <= 0:
             return None
-        if self.training and not self.enabled:
-            return None
         if self.address_with_identity:
             if slot_identity is None:
                 raise ValueError("Persistent bus requires slot identity on keys.")
@@ -931,9 +306,23 @@ class DecoderMemoryBus(nn.Module):
         values = self.v_proj(mapped_values).view(batch, slots, heads, head_dim).transpose(1, 2)
         return keys, values
 
-    def _relative_mask(self, queries: int, keys: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor | None:
+    def _relative_mask(
+        self,
+        queries: int,
+        keys: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        positions: torch.Tensor | None = None,
+    ) -> torch.Tensor | None:
         if not self.relative_bias:
             return None
+        if positions is not None:
+            if positions.shape[1] != queries or queries != keys:
+                raise ValueError("Working memory positions must match both canvas axes.")
+            relative = (
+                positions[:, :, None] - positions[:, None, :] + (keys - 1)
+            ).clamp(0, self.rel_bias.shape[1] - 1)
+            return self.rel_bias[:, relative].permute(1, 0, 2, 3).to(dtype=dtype)
         q = torch.arange(queries, device=device)
         k = torch.arange(keys, device=device)
         rel = (q[:, None] - k[None, :] + (keys - 1)).clamp(0, self.rel_bias.shape[1] - 1)
@@ -945,18 +334,28 @@ class DecoderMemoryBus(nn.Module):
         reader_index: int,
         keys: torch.Tensor,
         values: torch.Tensor,
+        positions: torch.Tensor | None = None,
+        *, key_seen: torch.Tensor | None = None,
     ) -> torch.Tensor:
         batch, canvas, _dim = hidden.shape
         heads = self.num_heads
         head_dim = self.head_dim
         query = self.q_proj[reader_index](self.q_norm(hidden))
         query = query.view(batch, canvas, heads, head_dim).transpose(1, 2)
-        bias = self._relative_mask(canvas, keys.shape[2], hidden.device, query.dtype)
-        if bias is not None:
+        bias = self._relative_mask(
+            canvas, keys.shape[2], hidden.device, query.dtype, positions
+        )
+        if bias is not None and bias.ndim == 3:
             bias = bias.unsqueeze(0)
+        if key_seen is not None:
+            key_mask = torch.zeros((batch, 1, 1, keys.shape[2]), device=hidden.device, dtype=torch.float32)
+            key_mask = key_mask.masked_fill(~key_seen[:, None, None, :], -1e9)
+            bias = key_mask if bias is None else bias.float() + key_mask
         context = _fp32_scaled_dot_product_attention(
             query, keys, values, attn_mask=bias
         )
+        if key_seen is not None:
+            context = torch.where(key_seen.any(-1)[:, None, None, None], context, torch.zeros_like(context))
         scale = torch.tanh(self.alpha[reader_index]).to(dtype=hidden.dtype).view(1, heads, 1, 1)
         context = context * scale
         context = context.transpose(1, 2).reshape(batch, canvas, self.kv_rank)
@@ -968,1059 +367,6 @@ class DecoderMemoryBus(nn.Module):
         self.rel_bias.zero_()
 
 
-class _SequenceBlock(nn.Module):
-    def __init__(self, dim: int, num_heads: int, ffn_dim: int) -> None:
-        super().__init__()
-        self.attn = _RankAttention(dim, num_heads, dim)
-        self.norm = _RMSNorm(dim)
-        self.ff_norm = _RMSNorm(dim)
-        self.ff = _SwiGLU(dim, ffn_dim)
-
-    def forward(self, hidden: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor:
-        normalized = self.norm(hidden)
-        hidden = hidden + self.attn(normalized, normalized, normalized, attn_mask=attn_mask)
-        return hidden + self.ff(self.ff_norm(hidden))
-
-
-class ExperienceRoleEncoder(nn.Module):
-    """Three role queries over a committed token's packed rank-space trajectory."""
-
-    def __init__(self, rank: int, num_heads: int, hidden_size: int) -> None:
-        super().__init__()
-        self.rank = rank
-        self.num_heads = num_heads
-        self.head_dim = rank // num_heads
-        self.role_queries = nn.Parameter(torch.empty(_EXPERIENCE_ROLES, rank))
-        self.working_proj = nn.Linear(hidden_size, rank, bias=False)
-        self.q_proj = nn.Linear(rank, rank, bias=False)
-        self.k_proj = nn.Linear(rank, rank, bias=False)
-        self.v_proj = nn.Linear(rank, rank, bias=False)
-        self.o_proj = nn.Linear(rank, rank, bias=False)
-        self.reason_embed = nn.Embedding(COMMIT_REASON_COUNT, rank)
-        self.reset_parameters()
-
-    @torch.no_grad()
-    def reset_parameters(self) -> None:
-        nn.init.normal_(self.role_queries, mean=0.0, std=0.02)
-
-    def forward(
-        self,
-        *,
-        working_state: torch.Tensor,
-        history_keys: torch.Tensor,
-        history_values: torch.Tensor,
-        history_mask: torch.Tensor,
-        z_final: torch.Tensor,
-        commit_reason: torch.Tensor,
-    ) -> torch.Tensor:
-        batch, canvas, slots, rank = history_keys.shape
-        working = self.working_proj(working_state)
-        extra = torch.stack((z_final, working), dim=2)
-        keys = torch.cat((history_keys, extra), dim=2)
-        values = torch.cat((history_values, extra), dim=2)
-        extra_mask = torch.ones(batch, canvas, 2, device=history_mask.device, dtype=torch.bool)
-        mask = torch.cat((history_mask, extra_mask), dim=2)
-        roles = self.role_queries.to(dtype=keys.dtype).view(1, 1, _EXPERIENCE_ROLES, rank)
-        roles = roles.expand(batch, canvas, -1, -1)
-        reason = self.reason_embed(commit_reason.clamp(0, COMMIT_REASON_COUNT - 1))
-        roles = roles + reason.to(dtype=roles.dtype).view(batch, 1, 1, rank)
-        heads = self.num_heads
-        head_dim = self.head_dim
-        query = self.q_proj(roles).view(batch, canvas, _EXPERIENCE_ROLES, heads, head_dim)
-        query = query.permute(0, 3, 1, 2, 4).reshape(
-            batch * heads * canvas, _EXPERIENCE_ROLES, head_dim
-        )
-        key = self.k_proj(keys).view(batch, canvas, slots + 2, heads, head_dim)
-        key = key.permute(0, 3, 1, 2, 4).reshape(batch * heads * canvas, slots + 2, head_dim)
-        value = self.v_proj(values).view(batch, canvas, slots + 2, heads, head_dim)
-        value = value.permute(0, 3, 1, 2, 4).reshape(batch * heads * canvas, slots + 2, head_dim)
-        attn_mask = mask.view(batch, 1, canvas, 1, slots + 2)
-        attn_mask = attn_mask.expand(-1, heads, -1, _EXPERIENCE_ROLES, -1)
-        attn_mask = attn_mask.reshape(batch * heads * canvas, _EXPERIENCE_ROLES, slots + 2)
-        additive = torch.zeros(
-            query.shape[0],
-            query.shape[1],
-            key.shape[1],
-            device=query.device,
-            dtype=query.dtype,
-        )
-        additive = additive.masked_fill(~attn_mask, _sdpa_mask_value(query.dtype))
-        context = _fp32_scaled_dot_product_attention(
-            query, key, value, attn_mask=additive
-        )
-        context = context.view(batch, heads, canvas, _EXPERIENCE_ROLES, head_dim)
-        context = context.permute(0, 2, 3, 1, 4).reshape(batch, canvas, _EXPERIENCE_ROLES, rank)
-        return self.o_proj(context)
-
-
-class CommitSequenceTransformer(nn.Module):
-    """Bidirectional phrase-level mixer over 3L experience role tokens."""
-
-    def __init__(self, dim: int, num_heads: int, num_layers: int, ffn_dim: int) -> None:
-        super().__init__()
-        self.role_embed = nn.Embedding(_EXPERIENCE_ROLES, dim)
-        self.blocks = nn.ModuleList(
-            [_SequenceBlock(dim, num_heads, ffn_dim) for _ in range(num_layers)]
-        )
-        self.norm = _RMSNorm(dim)
-
-    def forward(
-        self,
-        tokens: torch.Tensor,
-        positions: torch.Tensor,
-        valid: torch.Tensor,
-    ) -> torch.Tensor:
-        batch, length, dim = tokens.shape
-        roles = torch.arange(_EXPERIENCE_ROLES, device=tokens.device).repeat(length // _EXPERIENCE_ROLES + 1)
-        roles = roles[:length]
-        hidden = tokens + self.role_embed(roles).to(dtype=tokens.dtype)
-        heads = self.blocks[0].attn.num_heads
-        head_dim = dim // heads
-        hidden = hidden.view(batch, length, heads, head_dim).transpose(1, 2)
-        hidden = _apply_rotary(hidden, positions)
-        hidden = hidden.transpose(1, 2).reshape(batch, length, dim)
-        keep_rows = valid.any(dim=-1)
-        safe = valid.clone()
-        if length > 0:
-            safe[:, 0] = safe[:, 0] | ~keep_rows
-        keep = safe.unsqueeze(1) & safe.unsqueeze(2)
-        attn_mask = torch.zeros(
-            batch, length, length, device=tokens.device, dtype=tokens.dtype
-        )
-        attn_mask = attn_mask.masked_fill(
-            ~keep, _sdpa_mask_value(tokens.dtype)
-        )
-        for block in self.blocks:
-            hidden = block(hidden, attn_mask)
-        hidden = self.norm(hidden)
-        return torch.where(valid.unsqueeze(-1), hidden, hidden.new_zeros(hidden.shape))
-
-
-class TransformerCommitWriter(nn.Module):
-    """Identity-init slot-gated persistent write."""
-
-    def __init__(
-        self,
-        dim: int,
-        rank: int,
-        num_heads: int,
-        ffn_dim: int,
-        experience_dim: int,
-    ) -> None:
-        super().__init__()
-        self.cross = _RankAttention(dim, num_heads, rank)
-        self.experience_up = (
-            nn.Identity()
-            if experience_dim == dim
-            else nn.Linear(experience_dim, dim, bias=False)
-        )
-        self.self_attn = _RankAttention(dim, num_heads, rank)
-        self.ff = _SwiGLU(dim, ffn_dim)
-        self.ff_norm = _RMSNorm(dim)
-        self.norm = _RMSNorm(dim)
-        self.gate = nn.Linear(dim * 2, 1, bias=True)
-        self.beta_write = nn.Parameter(torch.zeros(()))
-        self.gamma_ca = nn.Parameter(torch.tensor(0.1))
-        self.gamma_sa = nn.Parameter(torch.zeros(()))
-        self.gamma_ffn = nn.Parameter(torch.tensor(0.1))
-        self.reset_identity_parameters()
-
-    @torch.no_grad()
-    def reset_identity_parameters(self) -> None:
-        nn.init.zeros_(self.gate.weight)
-        nn.init.constant_(self.gate.bias, _GATE_BIAS)
-        self.beta_write.zero_()
-        self.gamma_sa.zero_()
-        self.gamma_ca.copy_(self.gamma_ca.new_tensor(0.1))
-        self.gamma_ffn.copy_(self.gamma_ffn.new_tensor(0.1))
-
-    def forward(
-        self,
-        memory: torch.Tensor,
-        experience: torch.Tensor,
-        experience_mask: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        mapped = self.experience_up(experience)
-        row_has_experience = experience_mask.any(dim=-1)
-        safe_mask = experience_mask
-        if experience_mask.shape[-1] > 0:
-            safe_mask = experience_mask.clone()
-            safe_mask[:, 0] = safe_mask[:, 0] | ~row_has_experience
-        keep = safe_mask.unsqueeze(1) & torch.ones(
-            memory.shape[0], memory.shape[1], 1, device=memory.device, dtype=torch.bool
-        )
-        attn_mask = torch.zeros(
-            memory.shape[0],
-            memory.shape[1],
-            mapped.shape[1],
-            device=memory.device,
-            dtype=memory.dtype,
-        )
-        attn_mask = attn_mask.masked_fill(~keep, _sdpa_mask_value(memory.dtype))
-        delta_ca = self.cross(self.norm(memory), mapped, mapped, attn_mask=attn_mask)
-        hidden = memory + self.gamma_ca.to(dtype=memory.dtype) * delta_ca
-        delta_sa = self.self_attn(self.norm(hidden), hidden, hidden)
-        hidden = hidden + self.gamma_sa.to(dtype=memory.dtype) * delta_sa
-        delta_ff = self.ff(self.ff_norm(hidden))
-        proposed = hidden + self.gamma_ffn.to(dtype=memory.dtype) * delta_ff
-        delta = proposed - memory
-        gate = torch.sigmoid(self.gate(torch.cat((memory, proposed), dim=-1)))
-        scale = torch.tanh(self.beta_write).to(dtype=memory.dtype)
-        written = memory + scale * gate * delta
-        written = torch.where(row_has_experience.view(memory.shape[0], 1, 1), written, memory)
-        return written, gate, delta
-
-
-class LatentDeliberationTransformer(nn.Module):
-    """Working trajectory processor plus commit-only persistent writer."""
-
-    def __init__(
-        self,
-        *,
-        hidden_size: int,
-        vocab_size: int,
-        latent_dim: int = 2816,
-        ffn_dim: int = 7168,
-        memory_slots: int = 256,
-        num_layers: int = 4,
-        num_heads: int = 16,
-        local_attention_window: int = 128,
-        dropout: float = 0.0,
-        history_length: int = 16,
-        tape_probes: int = 16,
-        history_kv_rank: int = 1024,
-        num_memory_readers: int = 0,
-        num_working_readers: int | None = None,
-        num_persistent_readers: int | None = None,
-        working_last_block_global: bool = True,
-        experience_roles: int = _EXPERIENCE_ROLES,
-        commit_sequence_layers: int = 2,
-        commit_sequence_dim: int | None = None,
-        writer_ffn_dim: int | None = None,
-        max_canvas_length: int = 256,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__()
-        del dropout, experience_roles
-        if latent_dim % num_heads:
-            raise ValueError("`latent_dim` must be divisible by `num_heads`.")
-        if local_attention_window <= 0:
-            raise ValueError("`local_attention_window` must be positive.")
-        if history_length <= 0:
-            raise ValueError("`history_length` must be positive.")
-        if ffn_dim <= 0:
-            raise ValueError("`ffn_dim` must be positive.")
-        if history_kv_rank % num_heads or history_kv_rank > latent_dim:
-            raise ValueError("Invalid history K/V rank.")
-        self.hidden_size = hidden_size
-        self.vocab_size = vocab_size
-        self.latent_dim = latent_dim
-        self.memory_slots = memory_slots
-        self.history_length = history_length
-        self.tape_probes = int(tape_probes)
-        if self.tape_probes <= 0:
-            raise ValueError("`tape_probes` must be positive.")
-        self.history_views = _HISTORY_VIEWS
-        self.log_vocab = math.log(max(vocab_size, 2))
-        packet_dim = int(commit_sequence_dim or history_kv_rank)
-        if packet_dim % num_heads:
-            raise ValueError("`commit_sequence_dim` must be divisible by `num_heads`.")
-        self.packet_dim = packet_dim
-        self.history_in = (
-            nn.Identity()
-            if hidden_size == latent_dim
-            else nn.Linear(hidden_size, latent_dim, bias=False)
-        )
-        self.query_in = (
-            nn.Identity()
-            if hidden_size == latent_dim
-            else nn.Linear(hidden_size, latent_dim, bias=False)
-        )
-        self.history_projector = SharedHistoryProjector(latent_dim, history_kv_rank)
-        self.tape_pool = CanvasProbePool(history_kv_rank, self.tape_probes, num_heads)
-        self.persistent_kv = SharedPersistentKV(latent_dim, num_heads, history_kv_rank)
-        self.bias_in = nn.Linear(
-            _KEY_META_DIM + _QUERY_META_DIM + _ROW_META_DIM, _METADATA_HIDDEN, bias=True
-        )
-        self.bias_out = nn.Linear(_METADATA_HIDDEN, num_heads, bias=True)
-        nn.init.zeros_(self.bias_out.weight)
-        nn.init.zeros_(self.bias_out.bias)
-        self.film_in = nn.Linear(_KEY_META_DIM, _FILM_RANK, bias=True)
-        self.film_out = nn.Linear(_FILM_RANK, 2 * history_kv_rank, bias=True)
-        nn.init.zeros_(self.film_out.weight)
-        nn.init.zeros_(self.film_out.bias)
-        self.blocks = nn.ModuleList(
-            [
-                _ProcessorBlock(
-                    latent_dim,
-                    num_heads,
-                    local_attention_window,
-                    history_kv_rank,
-                    ffn_dim,
-                    global_attention=bool(
-                        working_last_block_global and index == num_layers - 1
-                    ),
-                )
-                for index in range(num_layers)
-            ]
-        )
-        self.output_norm = _RMSNorm(latent_dim)
-        self.output_to_hidden = (
-            nn.Identity()
-            if hidden_size == latent_dim
-            else nn.Linear(latent_dim, hidden_size, bias=False)
-        )
-        self.memory_slot_identity = nn.Parameter(torch.empty(memory_slots, latent_dim))
-        bus_heads = num_heads if hidden_size % num_heads == 0 else 1
-        bus_rank = history_kv_rank if history_kv_rank % bus_heads == 0 else bus_heads
-        working_readers = num_memory_readers if num_working_readers is None else num_working_readers
-        persistent_readers = (
-            num_memory_readers if num_persistent_readers is None else num_persistent_readers
-        )
-        self.working_memory_bus = DecoderMemoryBus(
-            hidden_size=hidden_size,
-            num_heads=bus_heads,
-            num_readers=working_readers,
-            memory_dim=hidden_size,
-            kv_rank=bus_rank,
-            relative_bias=True,
-            address_with_identity=False,
-            max_relative_span=max_canvas_length,
-        )
-        self.persistent_memory_bus = DecoderMemoryBus(
-            hidden_size=hidden_size,
-            num_heads=bus_heads,
-            num_readers=persistent_readers,
-            memory_dim=latent_dim,
-            kv_rank=bus_rank,
-            relative_bias=False,
-            address_with_identity=True,
-            max_relative_span=max_canvas_length,
-        )
-        self.experience_encoder = ExperienceRoleEncoder(
-            packet_dim, num_heads, hidden_size
-        )
-        self.commit_sequence = CommitSequenceTransformer(
-            packet_dim,
-            num_heads,
-            commit_sequence_layers,
-            max(packet_dim * 2, packet_dim),
-        )
-        self.commit_writer = TransformerCommitWriter(
-            latent_dim,
-            history_kv_rank,
-            num_heads,
-            writer_ffn_dim or ffn_dim,
-            packet_dim,
-        )
-        self.reset_identity_parameters()
-
-    @property
-    def memory_bus(self) -> DecoderMemoryBus:
-        return self.persistent_memory_bus
-
-    @torch.no_grad()
-    def reset_memory_slot_identity(self) -> None:
-        workspace = torch.empty_like(self.memory_slot_identity, dtype=torch.float32)
-        if self.memory_slots <= self.latent_dim:
-            nn.init.orthogonal_(workspace)
-        else:
-            nn.init.normal_(workspace, mean=0.0, std=1.0)
-            workspace = F.normalize(workspace, dim=-1)
-        self.memory_slot_identity.copy_(workspace.to(dtype=self.memory_slot_identity.dtype))
-
-    @torch.no_grad()
-    def reset_identity_parameters(self) -> None:
-        """Initialize direct parameters after generic PreTrainedModel init."""
-
-        self.reset_memory_slot_identity()
-        # These direct Parameters are created on `meta` during low-memory
-        # from_pretrained loading. Generic initialization covers Linear,
-        # Embedding, and RMSNorm modules, but not standalone query tensors.
-        self.tape_pool.reset_parameters()
-        self.experience_encoder.reset_parameters()
-        nn.init.zeros_(self.bias_out.weight)
-        nn.init.zeros_(self.bias_out.bias)
-        nn.init.zeros_(self.film_out.weight)
-        nn.init.zeros_(self.film_out.bias)
-        self.commit_writer.reset_identity_parameters()
-        self.working_memory_bus.reset_identity_parameters()
-        self.persistent_memory_bus.reset_identity_parameters()
-
-    def scaled_memory_slot_identity(
-        self,
-        *,
-        batch_size: int,
-        device: torch.device,
-        dtype: torch.dtype,
-    ) -> torch.Tensor:
-        identity = F.normalize(self.memory_slot_identity.float(), dim=-1)
-        identity = identity * math.sqrt(self.latent_dim)
-        return identity.to(device=device, dtype=dtype).unsqueeze(0).expand(
-            batch_size, -1, -1
-        )
-
-    def project_context(self, canvas_state: torch.Tensor) -> torch.Tensor:
-        return self.output_to_hidden(self.output_norm(canvas_state))
-
-    def encode_tape_frame(
-        self,
-        hidden: torch.Tensor,
-        live_mask: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Compress a detached canvas snapshot into tape probes."""
-
-        projected = self.history_projector(self.history_in(hidden.detach()))
-        if not bool(torch.isfinite(projected.detach()).all()):
-            raise FloatingPointError("Non-finite trajectory tape projection.")
-        probes = self.tape_pool(projected, live_mask)
-        # Materialize the MPS FP32 reduction before the probes enter the
-        # recurrent ring.  Besides fail-fast validation, this is a required
-        # producer/consumer barrier for the next BF16 denoise on MPS.
-        if not bool(torch.isfinite(probes.detach()).all()):
-            raise FloatingPointError("Non-finite trajectory tape probes.")
-        if live_mask is None:
-            valid = torch.ones(hidden.shape[0], device=hidden.device, dtype=torch.bool)
-        else:
-            valid = live_mask.to(device=hidden.device, dtype=torch.bool).any(dim=-1)
-        return probes, valid
-
-    def _tape_keys(
-        self,
-        tape: TrajectoryTape,
-        dtype: torch.dtype,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[None, None, None]:
-        if not bool(tape.valid.any()):
-            return None, None, None
-        rank = self.tape_pool.rank
-        heads = self.tape_pool.num_heads
-        head_dim = self.tape_pool.head_dim
-        batch, tape_length, probes, probe_dim = tape.probes.shape
-        if probe_dim != rank:
-            raise ValueError("Tape probe width does not match the projector rank.")
-        keys = tape.probes.to(dtype=dtype).reshape(batch, tape_length * probes, heads, head_dim)
-        keys = keys.transpose(1, 2)
-        mask = tape.valid.unsqueeze(-1).expand(-1, -1, probes).reshape(batch, tape_length * probes)
-        return keys, keys, mask
-
-    def _geometry(self, history: TrajectoryHistory) -> dict[str, torch.Tensor]:
-        valid = history.valid
-        batch, history_length, canvas, hidden_size = history.hidden.shape
-        velocity = torch.zeros(
-            batch, history_length, canvas,
-            device=history.hidden.device, dtype=torch.float32,
-        )
-        acceleration = torch.zeros_like(velocity)
-        reversal = torch.zeros_like(velocity)
-        recurrence = torch.zeros_like(velocity)
-        osc = torch.zeros_like(velocity)
-        scale = math.sqrt(max(hidden_size, 1))
-
-        # Geometry is diagnostic conditioning over a detached history ring.
-        # Processing one time edge at a time keeps only three FP32 frames and
-        # two deltas live instead of materializing FP32 hidden/delta/accel for
-        # the complete [B, H, C, D] ring.  Each scalar uses the same FP32
-        # subtraction, norm, and cosine operations as the dense formulation.
-        if history_length > 1:
-            previous_previous: torch.Tensor | None = None
-            previous = history.hidden[:, 0].float()
-            previous_delta: torch.Tensor | None = None
-            for index in range(1, history_length):
-                current = history.hidden[:, index].float()
-                delta = current - previous
-                velocity[:, index] = delta.norm(dim=-1) / scale
-                if previous_previous is not None and previous_delta is not None:
-                    accel = current - 2.0 * previous + previous_previous
-                    acceleration[:, index] = accel.norm(dim=-1) / scale
-                    reversal[:, index] = -_safe_cosine(delta, previous_delta)
-                    recurrence[:, index] = _safe_cosine(current, previous_previous)
-                    osc[:, index] = (
-                        recurrence[:, index] - _safe_cosine(current, previous)
-                    )
-                previous_previous = previous
-                previous = current
-                previous_delta = delta
-        raw_mask = valid
-        delta_mask = valid.clone()
-        delta_mask[:, 0] = False
-        if valid.shape[1] > 1:
-            delta_mask[:, 1:] = valid[:, 1:] & valid[:, :-1]
-        accel_mask = valid.clone()
-        accel_mask[:, :2] = False
-        if valid.shape[1] > 2:
-            accel_mask[:, 2:] = valid[:, 2:] & valid[:, 1:-1] & valid[:, :-2]
-        return {
-            "velocity": velocity.masked_fill(~delta_mask, 0.0),
-            "acceleration": acceleration.masked_fill(~accel_mask, 0.0),
-            "reversal": reversal.masked_fill(~accel_mask, 0.0),
-            "recurrence": recurrence.masked_fill(~accel_mask, 0.0),
-            "oscillation": osc.masked_fill(~accel_mask, 0.0),
-            "raw_mask": raw_mask,
-            "delta_mask": delta_mask,
-            "accel_mask": accel_mask,
-        }
-
-    def _frame_metadata(
-        self,
-        history: TrajectoryHistory,
-        state: LatentDeliberationState,
-        dtype: torch.dtype,
-        geometry: dict[str, torch.Tensor],
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        batch, history_length, canvas, _dim = history.hidden.shape
-        log_vocab = self.log_vocab
-        confidence = _renorm_confidence(history.confidence.to(dtype=dtype))
-        entropy = (history.entropy.to(dtype=dtype) / log_vocab).clamp(0.0, 1.0)
-        if history_length == 1:
-            recency = torch.ones(
-                batch, history_length, canvas, device=history.hidden.device, dtype=dtype
-            )
-        else:
-            recency = torch.linspace(
-                0.0, 1.0, history_length, device=history.hidden.device, dtype=dtype
-            )
-            recency = recency.view(1, history_length, 1).expand(batch, history_length, canvas)
-        retirement = torch.zeros(
-            batch, history_length, canvas, device=history.hidden.device, dtype=dtype
-        )
-        oldest = min(_RETIREMENT_FRAMES, history_length)
-        if oldest:
-            scale = torch.linspace(
-                1.0, 1.0 / oldest, oldest, device=history.hidden.device, dtype=dtype
-            )
-            retirement[:, :oldest] = scale.view(1, oldest, 1)
-            retirement = retirement * history.valid.to(dtype=dtype)
-        changed = history.token_changed.to(dtype=dtype)
-        delta_c = torch.zeros_like(confidence)
-        delta_e = torch.zeros_like(entropy)
-        if history_length > 1:
-            delta_c[:, 1:] = (confidence[:, 1:] - confidence[:, :-1]).clamp(-1.0, 1.0)
-            delta_e[:, 1:] = ((history.entropy[:, 1:] - history.entropy[:, :-1]) / log_vocab).clamp(
-                -1.0, 1.0
-            ).to(dtype=dtype)
-        age = (
-            state.age.to(dtype=dtype).clamp_max(_AGE_MAX).log1p()
-            / math.log1p(_AGE_MAX)
-        )
-        age_frames = torch.zeros_like(confidence)
-        age_frames[:, -1] = age
-        key_meta = torch.stack(
-            (
-                confidence, entropy, age_frames, changed, delta_c, delta_e,
-                recency, retirement,
-                geometry["velocity"].to(dtype=dtype),
-                geometry["acceleration"].to(dtype=dtype),
-                geometry["reversal"].to(dtype=dtype),
-                geometry["recurrence"].to(dtype=dtype),
-                geometry["oscillation"].to(dtype=dtype),
-            ),
-            dim=-1,
-        )
-        query_scalars = torch.stack(
-            (
-                _renorm_confidence(state.confidence.to(dtype=dtype)),
-                (state.entropy.to(dtype=dtype) / log_vocab).clamp(0.0, 1.0),
-                age,
-                state.token_changed.to(dtype=dtype),
-                state.confidence_delta.to(dtype=dtype).clamp(-1.0, 1.0),
-                (state.entropy_delta.to(dtype=dtype) / log_vocab).clamp(-1.0, 1.0),
-            ),
-            dim=-1,
-        )
-        fourier = _canvas_fourier(canvas, history.hidden.device, dtype).unsqueeze(0).expand(
-            batch, -1, -1
-        )
-        query_meta = torch.cat((fourier, query_scalars), dim=-1)
-        ponder = (
-            state.ponder_steps.to(dtype=dtype).clamp_max(_PONDER_MAX).log1p()
-            / math.log1p(_PONDER_MAX)
-        )
-        stagnation = (
-            state.stagnation_steps.to(dtype=dtype).clamp_max(_STAGNATION_MAX).log1p()
-            / math.log1p(_STAGNATION_MAX)
-        )
-        row_meta = torch.stack((ponder, stagnation), dim=-1)
-        return key_meta, query_meta, row_meta, history.valid
-
-    def _history_views(
-        self,
-        history: TrajectoryHistory,
-        key_meta: torch.Tensor,
-        geometry: dict[str, torch.Tensor],
-        dtype: torch.dtype,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        mapped = self.history_in(history.hidden.to(dtype=dtype))
-        projected = self.history_projector(mapped)
-        return self._views_from_projected(
-            projected, history, key_meta, geometry, dtype,
-            preformat_heads=True,
-        )
-
-    def _views_from_projected(
-        self,
-        projected: torch.Tensor,
-        history: TrajectoryHistory,
-        key_meta: torch.Tensor,
-        geometry: dict[str, torch.Tensor],
-        dtype: torch.dtype,
-        *,
-        preformat_heads: bool = False,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        valid = history.valid
-        raw_mask = geometry["raw_mask"]
-        delta_mask = geometry["delta_mask"]
-        accel_mask = geometry["accel_mask"]
-        latest_index = (
-            valid.to(torch.int64) * (
-                torch.arange(valid.shape[1], device=valid.device).view(1, -1, 1) + 1
-            )
-        ).amax(dim=1) - 1
-        has_latest = latest_index.ge(0)
-        latest_index = latest_index.clamp_min(0)
-        gather = latest_index.view(projected.shape[0], 1, projected.shape[2], 1).expand(
-            -1, 1, -1, projected.shape[-1]
-        )
-        z_latest = projected.gather(1, gather).squeeze(1)
-        residual = projected - z_latest.unsqueeze(1)
-        residual_mask = raw_mask & has_latest.unsqueeze(1)
-        if projected.shape[1] > 1:
-            delta = torch.cat(
-                (torch.zeros_like(projected[:, :1]), projected[:, 1:] - projected[:, :-1]),
-                dim=1,
-            )
-        else:
-            delta = torch.zeros_like(projected)
-        accel = torch.zeros_like(projected)
-        if projected.shape[1] > 2:
-            accel[:, 2:] = projected[:, 2:] - 2.0 * projected[:, 1:-1] + projected[:, :-2]
-        views = torch.stack((projected, delta, accel, residual), dim=2)
-        view_mask = torch.stack((raw_mask, delta_mask, accel_mask, residual_mask), dim=2)
-        film = self.film_out(F.silu(self.film_in(key_meta.to(dtype=dtype))))
-        scale, shift = film.chunk(2, dim=-1)
-        numeric_mask = view_mask.unsqueeze(-1).to(dtype=views.dtype)
-        # The unmasked `views` tensor is dead here.  Mask it in-place and use
-        # it as the key storage instead of allocating a second full copy.
-        # Applying the same mask again after the FiLM shift preserves invalid
-        # entries as exact zeros and leaves every valid entry unchanged.
-        views.mul_(numeric_mask)
-        values = views * (1.0 + scale.unsqueeze(2))
-        values.add_(shift.unsqueeze(2)).mul_(numeric_mask)
-        keys = views
-        batch, history_length, views_n, canvas, dim = keys.shape
-        if preformat_heads:
-            heads = self.blocks[0].history_attention.num_heads
-            head_dim = dim // heads
-            keys = keys.view(
-                batch, history_length, views_n, canvas, heads, head_dim
-            ).permute(0, 4, 3, 1, 2, 5).reshape(
-                batch, heads, canvas, history_length * views_n, head_dim
-            )
-            values = values.view(
-                batch, history_length, views_n, canvas, heads, head_dim
-            ).permute(0, 4, 3, 1, 2, 5).reshape(
-                batch, heads, canvas, history_length * views_n, head_dim
-            )
-        else:
-            keys = keys.permute(0, 3, 1, 2, 4).reshape(
-                batch, canvas, history_length * views_n, dim
-            )
-            values = values.permute(0, 3, 1, 2, 4).reshape(
-                batch, canvas, history_length * views_n, dim
-            )
-        key_mask = view_mask.permute(0, 3, 1, 2).reshape(batch, canvas, history_length * views_n)
-        return keys, values, key_mask, projected
-
-    def _attention_bias(
-        self,
-        key_meta: torch.Tensor,
-        query_meta: torch.Tensor,
-        row_meta: torch.Tensor,
-        view_mask: torch.Tensor,
-        num_heads: int,
-        dtype: torch.dtype,
-    ) -> torch.Tensor:
-        batch, history_length, canvas, _meta = key_meta.shape
-        views = _HISTORY_VIEWS
-        query = query_meta[:, None, :, :].expand(-1, history_length, -1, -1)
-        row = row_meta[:, None, None, :].expand(-1, history_length, canvas, -1)
-        packed = torch.cat((key_meta.to(dtype=dtype), query, row), dim=-1)
-        bias = self.bias_out(F.silu(self.bias_in(packed)))
-        bias = bias.permute(0, 3, 2, 1).unsqueeze(-1).expand(-1, -1, -1, -1, views)
-        return bias.reshape(batch, num_heads, canvas, history_length * views).to(dtype=dtype)
-
-    def forward(
-        self,
-        *,
-        token_embeddings: torch.Tensor,
-        confidence: torch.Tensor,
-        entropy: torch.Tensor,
-        state: LatentDeliberationState,
-        history: TrajectoryHistory,
-        tape: TrajectoryTape | None = None,
-    ) -> LatentProcessorOutput:
-        if token_embeddings.ndim != 3:
-            raise ValueError("`token_embeddings` must have shape [batch, canvas, hidden].")
-        batch_size, canvas_length, hidden_size = token_embeddings.shape
-        if hidden_size != self.hidden_size:
-            raise ValueError("Unexpected hidden size for latent deliberation.")
-        if state.memory_slots.shape != (batch_size, self.memory_slots, self.latent_dim):
-            raise ValueError("State memory slots do not match this module.")
-        if history.hidden.shape[:3] != (batch_size, self.history_length, canvas_length):
-            raise ValueError("Trajectory history does not match the current canvas.")
-        if state.age.dtype is not torch.int32:
-            raise TypeError("Latent deliberation ages must use int32.")
-        if tape is None:
-            tape = TrajectoryTape.empty(
-                batch_size=batch_size,
-                tape_length=self.history_length,
-                num_probes=self.tape_probes,
-                probe_dim=self.tape_pool.rank,
-                device=token_embeddings.device,
-                dtype=token_embeddings.dtype,
-            )
-        if tape.probes.shape[:2] != (batch_size, self.history_length):
-            raise ValueError("Trajectory tape does not match the current batch.")
-
-        dtype = token_embeddings.dtype
-        query = self.query_in(token_embeddings)
-        geometry = self._geometry(history)
-        key_meta, query_meta, row_meta, _valid = self._frame_metadata(
-            history, state, dtype, geometry
-        )
-        keys, values, key_mask, projected = self._history_views(
-            history, key_meta, geometry, dtype
-        )
-        num_heads = self.blocks[0].history_attention.num_heads
-        attn_bias = self._attention_bias(
-            key_meta, query_meta, row_meta, key_mask, num_heads, dtype
-        )
-        tape_keys, tape_values, tape_mask = self._tape_keys(tape, dtype)
-        # Working state accumulates from zero. Empty history and zero memory
-        # therefore produce a zero self-conditioning residual.
-        canvas_state = torch.zeros_like(query)
-        slot_identity = self.scaled_memory_slot_identity(
-            batch_size=batch_size, device=state.memory_slots.device, dtype=query.dtype
-        )
-        memory_keys, memory_values = self.persistent_kv(state.memory_slots, slot_identity)
-        for block in self.blocks:
-            canvas_state = block(
-                canvas_state, keys, values, attn_bias, key_mask,
-                memory_keys, memory_values, query,
-                tape_keys, tape_values, tape_mask,
-            )
-        context = self.project_context(canvas_state)
-        next_state = LatentDeliberationState(
-            memory_slots=state.memory_slots,
-            confidence=confidence.to(dtype=torch.float32),
-            entropy=entropy.to(dtype=torch.float32),
-            age=state.age,
-            token_changed=state.token_changed,
-            confidence_delta=state.confidence_delta,
-            entropy_delta=state.entropy_delta,
-            ponder_steps=state.ponder_steps,
-            stagnation_steps=state.stagnation_steps,
-        )
-        return LatentProcessorOutput(
-            context=context,
-            working_state=context,
-            state=next_state,
-            history_projected=projected,
-        )
-
-    def _slice_commit_canvas(
-        self,
-        *,
-        working_state: torch.Tensor,
-        history: TrajectoryHistory,
-        history_projected: torch.Tensor,
-        heavy_hidden: torch.Tensor,
-        max_commit: int,
-    ) -> tuple[torch.Tensor, TrajectoryHistory, torch.Tensor, torch.Tensor]:
-        return (
-            working_state[:, :max_commit],
-            TrajectoryHistory(
-                hidden=history.hidden[:, :, :max_commit],
-                confidence=history.confidence[:, :, :max_commit],
-                entropy=history.entropy[:, :, :max_commit],
-                token_changed=history.token_changed[:, :, :max_commit],
-                valid=history.valid[:, :, :max_commit],
-            ),
-            history_projected[:, :, :max_commit],
-            heavy_hidden[:, :max_commit],
-        )
-
-    def _experience_encoder_step(
-        self,
-        working_state: torch.Tensor,
-        history_keys: torch.Tensor,
-        history_values: torch.Tensor,
-        history_mask: torch.Tensor,
-        z_final: torch.Tensor,
-        commit_reason: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.experience_encoder(
-            working_state=working_state,
-            history_keys=history_keys,
-            history_values=history_values,
-            history_mask=history_mask,
-            z_final=z_final,
-            commit_reason=commit_reason,
-        )
-
-    def _encode_experience_roles(
-        self,
-        *,
-        working_state: torch.Tensor,
-        history_keys: torch.Tensor,
-        history_values: torch.Tensor,
-        history_mask: torch.Tensor,
-        z_final: torch.Tensor,
-        commit_reason: torch.Tensor,
-    ) -> torch.Tensor:
-        canvas = int(working_state.shape[1])
-        stripe = min(_EXPERIENCE_CANVAS_STRIPE, canvas)
-        parts: list[torch.Tensor] = []
-        for start in range(0, canvas, stripe):
-            stop = min(start + stripe, canvas)
-            parts.append(
-                self._experience_encoder_step(
-                    working_state[:, start:stop],
-                    history_keys[:, start:stop],
-                    history_values[:, start:stop],
-                    history_mask[:, start:stop],
-                    z_final[:, start:stop],
-                    commit_reason,
-                )
-            )
-        return parts[0] if len(parts) == 1 else torch.cat(parts, dim=1)
-
-    def _pack_experience(
-        self,
-        *,
-        working_state: torch.Tensor,
-        history: TrajectoryHistory,
-        history_projected: torch.Tensor,
-        heavy_hidden: torch.Tensor,
-        commit_lengths: torch.Tensor,
-        prefix_lengths: torch.Tensor,
-        commit_reason: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        batch, canvas, _hidden = working_state.shape
-        max_commit = min(int(commit_lengths.max().clamp_min(0)), canvas)
-        if max_commit <= 0:
-            empty = working_state.new_zeros(batch, 0, self.packet_dim)
-            return empty, empty.new_zeros(batch, 0, dtype=torch.bool), empty.new_zeros(batch, 0)
-        working_state, history, history_projected, heavy_hidden = self._slice_commit_canvas(
-            working_state=working_state,
-            history=history,
-            history_projected=history_projected,
-            heavy_hidden=heavy_hidden,
-            max_commit=max_commit,
-        )
-        dtype = working_state.dtype
-        geometry = self._geometry(history)
-        key_meta, _query_meta, _row_meta, _valid = self._frame_metadata(
-            history,
-            LatentDeliberationState.empty(
-                batch_size=batch,
-                canvas_length=max_commit,
-                latent_dim=self.latent_dim,
-                memory_slots=self.memory_slots,
-                device=working_state.device,
-                dtype=dtype,
-            ),
-            dtype,
-            geometry,
-        )
-        keys, values, key_mask, _projected = self._views_from_projected(
-            history_projected.to(dtype=dtype), history, key_meta, geometry, dtype
-        )
-        # Heavy is a TBPTT observation here, same as history frames: CE already
-        # backpropagated through this decoder stack. The writer stays in the
-        # temporal graph via `working_state` and the committed memory output.
-        z_final = self.history_projector(
-            self.history_in(heavy_hidden.detach().to(dtype=dtype))
-        )
-        roles = self._encode_experience_roles(
-            working_state=working_state,
-            history_keys=keys,
-            history_values=values,
-            history_mask=key_mask,
-            z_final=z_final,
-            commit_reason=commit_reason,
-        )
-        tokens = roles.reshape(batch, max_commit * _EXPERIENCE_ROLES, -1)
-        token_valid = torch.arange(max_commit, device=working_state.device)[None, :] < commit_lengths[:, None]
-        valid = token_valid.unsqueeze(-1).expand(-1, -1, _EXPERIENCE_ROLES).reshape(
-            batch, max_commit * _EXPERIENCE_ROLES
-        )
-        abs_pos = prefix_lengths[:, None] + torch.arange(
-            max_commit, device=working_state.device
-        )[None, :]
-        pos = abs_pos.unsqueeze(-1).expand(-1, -1, _EXPERIENCE_ROLES).reshape(
-            batch, max_commit * _EXPERIENCE_ROLES
-        )
-        mixed = self.commit_sequence(tokens, pos, valid)
-        return mixed, valid, pos
-
-    def commit_write(
-        self,
-        *,
-        memory: torch.Tensor,
-        working_state: torch.Tensor,
-        history: TrajectoryHistory,
-        history_projected: torch.Tensor,
-        heavy_hidden: torch.Tensor,
-        commit_lengths: torch.Tensor,
-        prefix_lengths: torch.Tensor | None = None,
-        commit_reason: torch.Tensor | None = None,
-        **kwargs: Any,
-    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
-        batch = memory.shape[0]
-        canvas = working_state.shape[1]
-        lengths = commit_lengths.to(device=memory.device, dtype=torch.long)
-        if prefix_lengths is None:
-            prefixes = torch.zeros(batch, device=memory.device, dtype=torch.long)
-        else:
-            prefixes = prefix_lengths.to(device=memory.device, dtype=torch.long)
-        if commit_reason is None:
-            reasons = torch.full(
-                (batch,), COMMIT_REASON_NORMAL, device=memory.device, dtype=torch.long
-            )
-        else:
-            reasons = commit_reason.to(device=memory.device, dtype=torch.long)
-        if bool((lengths <= 0).all()):
-            zero = memory.new_zeros(batch, self.memory_slots, 1)
-            return memory, {
-                "gate_mean": zero.mean(),
-                "gate_max": zero.amax(),
-                "gate_gt_01": zero.new_zeros(()),
-                "gate_gt_05": zero.new_zeros(()),
-                "delta_norm_mean": memory.new_zeros(()),
-            }
-        experience, experience_mask, _pos = self._pack_experience(
-            working_state=working_state,
-            history=history,
-            history_projected=history_projected,
-            heavy_hidden=heavy_hidden,
-            commit_lengths=lengths,
-            prefix_lengths=prefixes,
-            commit_reason=reasons,
-        )
-        max_commit = int(lengths.max())
-        role_stop = max_commit * _EXPERIENCE_ROLES
-        chunk_tokens = experience[:, :role_stop]
-        chunk_mask = experience_mask[:, :role_stop]
-        if chunk_tokens.shape[1] > 0:
-            written, last_gate, last_delta = self.commit_writer(
-                memory, chunk_tokens, chunk_mask
-            )
-        else:
-            written = memory
-            last_gate = memory.new_zeros(batch, self.memory_slots, 1)
-            last_delta = memory.new_zeros(memory.shape)
-        gate = last_gate.detach()
-        delta_norm = last_delta.detach().float().norm(dim=-1)
-        diagnostics = {
-            "gate_mean": gate.mean(),
-            "gate_max": gate.amax(),
-            "gate_gt_01": gate.gt(0.1).float().sum(),
-            "gate_gt_05": gate.gt(0.5).float().sum(),
-            "delta_norm_mean": delta_norm.mean(),
-        }
-        unchanged = lengths.le(0).view(batch, 1, 1)
-        written = torch.where(unchanged, memory, written)
-        return written, diagnostics
-
-
-def slice_trajectory_history(
-    history: TrajectoryHistory, rows: slice | torch.Tensor
-) -> TrajectoryHistory:
-    return TrajectoryHistory(
-        hidden=history.hidden[rows],
-        confidence=history.confidence[rows],
-        entropy=history.entropy[rows],
-        token_changed=history.token_changed[rows],
-        valid=history.valid[rows],
-    )
-
-
-def cat_trajectory_history(
-    histories: Sequence[TrajectoryHistory],
-) -> TrajectoryHistory:
-    return TrajectoryHistory(
-        hidden=torch.cat([history.hidden for history in histories], dim=0),
-        confidence=torch.cat([history.confidence for history in histories], dim=0),
-        entropy=torch.cat([history.entropy for history in histories], dim=0),
-        token_changed=torch.cat([history.token_changed for history in histories], dim=0),
-        valid=torch.cat([history.valid for history in histories], dim=0),
-    )
-
-
-def choose_trajectory_history(
-    previous: TrajectoryHistory,
-    updated: TrajectoryHistory,
-    update_mask: torch.Tensor,
-) -> TrajectoryHistory:
-    def choose(old: torch.Tensor, new: torch.Tensor) -> torch.Tensor:
-        mask = update_mask.view(update_mask.shape[0], *([1] * (old.ndim - 1)))
-        return torch.where(mask, new, old)
-
-    return TrajectoryHistory(
-        hidden=choose(previous.hidden, updated.hidden),
-        confidence=choose(previous.confidence, updated.confidence),
-        entropy=choose(previous.entropy, updated.entropy),
-        token_changed=choose(previous.token_changed, updated.token_changed),
-        valid=choose(previous.valid, updated.valid),
-    )
-
-
-def slice_trajectory_tape(
-    tape: TrajectoryTape, rows: slice | torch.Tensor
-) -> TrajectoryTape:
-    return TrajectoryTape(probes=tape.probes[rows], valid=tape.valid[rows])
-
-
-def cat_trajectory_tape(tapes: Sequence[TrajectoryTape]) -> TrajectoryTape:
-    return TrajectoryTape(
-        probes=torch.cat([tape.probes for tape in tapes], dim=0),
-        valid=torch.cat([tape.valid for tape in tapes], dim=0),
-    )
-
-
-def choose_trajectory_tape(
-    previous: TrajectoryTape,
-    updated: TrajectoryTape,
-    update_mask: torch.Tensor,
-) -> TrajectoryTape:
-    def choose(old: torch.Tensor, new: torch.Tensor) -> torch.Tensor:
-        mask = update_mask.view(update_mask.shape[0], *([1] * (old.ndim - 1)))
-        return torch.where(mask, new, old)
-
-    return TrajectoryTape(
-        probes=choose(previous.probes, updated.probes),
-        valid=choose(previous.valid, updated.valid),
-    )
-
-
-def empty_trajectory_tape(
-    *,
-    batch_size: int,
-    config: object,
-    device: torch.device,
-    dtype: torch.dtype,
-) -> TrajectoryTape:
-    rank = int(getattr(config, "latent_history_kv_rank"))
-    return TrajectoryTape.empty(
-        batch_size=batch_size,
-        tape_length=int(getattr(config, "latent_history_length")),
-        num_probes=int(getattr(config, "latent_tape_probes", 16)),
-        probe_dim=rank,
-        device=device,
-        dtype=dtype,
-    )
-
-
 def slice_latent_state(
     state: LatentDeliberationState, rows: slice | torch.Tensor
 ) -> LatentDeliberationState:
@@ -2028,12 +374,12 @@ def slice_latent_state(
         memory_slots=state.memory_slots[rows],
         confidence=state.confidence[rows],
         entropy=state.entropy[rows],
-        age=state.age[rows],
-        token_changed=state.token_changed[rows],
-        confidence_delta=state.confidence_delta[rows],
-        entropy_delta=state.entropy_delta[rows],
         ponder_steps=state.ponder_steps[rows],
         stagnation_steps=state.stagnation_steps[rows],
+        gdn2=GDN2TrajectoryState(
+            cells=state.gdn2.cells[rows], row=state.gdn2.row[rows],
+            persistent=state.gdn2.persistent[rows], seen=state.gdn2.seen[rows],
+        ),
     )
 
 
@@ -2045,12 +391,14 @@ def cat_latent_states(states: Sequence[LatentDeliberationState]) -> LatentDelibe
         memory_slots=cat("memory_slots"),
         confidence=cat("confidence"),
         entropy=cat("entropy"),
-        age=cat("age"),
-        token_changed=cat("token_changed"),
-        confidence_delta=cat("confidence_delta"),
-        entropy_delta=cat("entropy_delta"),
         ponder_steps=cat("ponder_steps"),
         stagnation_steps=cat("stagnation_steps"),
+        gdn2=GDN2TrajectoryState(
+            cells=torch.cat([state.gdn2.cells for state in states], dim=0),
+            row=torch.cat([state.gdn2.row for state in states], dim=0),
+            persistent=torch.cat([state.gdn2.persistent for state in states], dim=0),
+            seen=torch.cat([state.gdn2.seen for state in states], dim=0),
+        ),
     )
 
 
@@ -2060,7 +408,6 @@ def infer_commit_reason(
     jump_rows: torch.Tensor | None = None,
     commit_token_ids: torch.Tensor | None = None,
     terminal_token_ids: Sequence[int] = (),
-    training_random: bool = False,
 ) -> torch.Tensor:
     """Return per-row commit-reason codes. No hard skip; writer sees the label."""
 
@@ -2072,7 +419,7 @@ def infer_commit_reason(
     )
     committed = commit_lengths.gt(0)
     default = (
-        COMMIT_REASON_TRAINING_RANDOM if training_random else COMMIT_REASON_NORMAL
+        COMMIT_REASON_NORMAL
     )
     reasons = torch.where(committed, torch.full_like(reasons, default), reasons)
     if jump_rows is not None:
@@ -2097,38 +444,181 @@ def infer_commit_reason(
     return reasons
 
 
-def memory_bus_parameter_names(module: nn.Module) -> list[str]:
-    return [
-        name for name, _parameter in module.named_parameters()
-        if "working_memory_bus." in name or "persistent_memory_bus." in name
-        or "memory_bus." in name
-    ]
-
-
-__all__ = [
-    "COMMIT_REASON_FALLBACK",
-    "COMMIT_REASON_FORCED_JUMP",
-    "COMMIT_REASON_NONE",
-    "COMMIT_REASON_NORMAL",
-    "COMMIT_REASON_TERMINAL",
-    "COMMIT_REASON_TRAINING_RANDOM",
-    "DecoderMemoryBus",
-    "LatentDeliberationState",
-    "LatentDeliberationTransformer",
-    "LatentProcessorOutput",
-    "TrajectoryHistory",
-    "TrajectoryTape",
-    "advance_trajectory_clocks",
-    "cat_latent_states",
-    "cat_trajectory_history",
-    "cat_trajectory_tape",
-    "choose_trajectory_history",
-    "choose_trajectory_tape",
-    "empty_trajectory_tape",
-    "infer_commit_reason",
-    "memory_bus_parameter_names",
-    "should_force_trajectory_jump",
-    "slice_latent_state",
-    "slice_trajectory_history",
-    "slice_trajectory_tape",
-]
+class _CanvasBlock(nn.Module):
+    def __init__(self, width: int, heads: int, rank: int, ffn: int,
+                 window: int, global_attention: bool) -> None:
+        super().__init__()
+        self.norm = _RMSNorm(width)
+        self.attn = _RankAttention(width, heads, rank)
+        self.ff_norm = _RMSNorm(width)
+        self.ff = _SwiGLU(width, ffn)
+        self.window = window
+        self.global_attention = global_attention
+
+    def forward(self, hidden: torch.Tensor, seen: torch.Tensor,
+                offsets: torch.Tensor) -> torch.Tensor:
+        allowed = seen[:, None, :].expand(-1, hidden.shape[1], -1)
+        if not self.global_attention:
+            allowed = allowed & ((offsets[:, :, None] - offsets[:, None, :]).abs() < self.window)
+        additive = torch.zeros(allowed.shape, device=hidden.device, dtype=torch.float32)
+        additive = additive.masked_fill(~allowed, -1e9)
+        normed = self.norm(hidden)
+        hidden = hidden + self.attn(normed, normed, normed, attn_mask=additive)
+        hidden = hidden + self.ff(self.ff_norm(hidden))
+        return torch.where(seen[..., None], hidden, torch.zeros_like(hidden))
+
+
+class _WorkingBus(DecoderMemoryBus):
+    def prepare_kv(self, memory: tuple[torch.Tensor, torch.Tensor],
+                   slot_identity: torch.Tensor | None = None):
+        del slot_identity
+        working, seen = memory
+        pair = super().prepare_kv(working)
+        return None if pair is None else (*pair, seen)
+
+    def read(self, hidden: torch.Tensor, reader_index: int,
+             keys: torch.Tensor, values: torch.Tensor, seen: torch.Tensor,
+             positions: torch.Tensor | None = None) -> torch.Tensor:
+        written = super().read(hidden, reader_index, keys, values, positions, key_seen=seen)
+        return torch.where(seen[..., None], written, hidden)
+
+
+class _PersistentBus(nn.Module):
+    def __init__(self, memory: nn.Module, readers: int) -> None:
+        super().__init__()
+        object.__setattr__(self, "_memory_ref", weakref.ref(memory))
+        self.num_readers = readers
+        self.alpha = nn.Parameter(torch.zeros(max(readers, 1), 1))
+
+
+    def reset_identity_parameters(self) -> None:
+        with torch.no_grad():
+            self.alpha.zero_()
+
+    def prepare_kv(self, memory: torch.Tensor,
+                   seen: torch.Tensor | None = None):
+        if self.num_readers <= 0:
+            return None
+        return memory, seen
+
+    def read(self, hidden: torch.Tensor, reader_index: int,
+             memory: torch.Tensor, seen: torch.Tensor | None) -> torch.Tensor:
+        delta = self._memory_ref().read_shared(memory, hidden)
+        if seen is not None:
+            delta = delta * seen[..., None].to(delta.dtype)
+        return hidden + torch.tanh(self.alpha[reader_index]).to(hidden.dtype) * delta
+
+
+class LatentDeliberationTransformer(nn.Module):
+    def __init__(self, *, hidden_size: int,
+                 latent_dim: int = 2816, ffn_dim: int = 7168,
+                  num_layers: int = 4,
+                 num_heads: int = 16, local_attention_window: int = 128,
+                  tape_probes: int = 4,
+
+                 history_kv_rank: int = 1024, num_memory_readers: int = 0,
+                 num_working_readers: int | None = None,
+                 num_persistent_readers: int | None = None,
+                 working_last_block_global: bool = True,
+                 commit_sequence_dim: int | None = None,
+                 max_canvas_length: int = 256) -> None:
+        super().__init__()
+        if hidden_size != latent_dim:
+            raise ValueError("Schema25 GDN2 requires hidden_size == latent_dim.")
+        self.hidden_size = hidden_size
+        self.latent_dim = latent_dim
+        self.tape_probes = tape_probes
+        self.packet_dim = int(commit_sequence_dim or history_kv_rank)
+        if self.packet_dim % num_heads:
+            raise ValueError("Commit packet rank must be divisible by attention heads.")
+        self.trajectory = GDN2TrajectoryMemory(
+            hidden_size, probes=tape_probes, persistent_observation_dim=self.packet_dim,
+        )
+        self.blocks = nn.ModuleList([
+            _CanvasBlock(hidden_size, num_heads, history_kv_rank, ffn_dim,
+                         local_attention_window,
+                         bool(working_last_block_global and i == num_layers - 1))
+            for i in range(num_layers)
+        ])
+        self.output_norm = _RMSNorm(hidden_size)
+        readers = num_memory_readers if num_working_readers is None else num_working_readers
+        persistent_readers = (
+            num_memory_readers if num_persistent_readers is None else num_persistent_readers
+        )
+        bus_rank = history_kv_rank if history_kv_rank % num_heads == 0 else num_heads
+        self.working_memory_bus = _WorkingBus(
+            hidden_size, num_heads, readers, hidden_size, bus_rank,
+            relative_bias=True, max_relative_span=max_canvas_length,
+        )
+        self.persistent_memory_bus = _PersistentBus(
+            self.trajectory.persistent, persistent_readers,
+        )
+        self.experience_in = nn.Linear(hidden_size * 3, self.packet_dim, bias=False)
+        self.reason_embed = nn.Embedding(6, self.packet_dim)
+
+    def reset_identity_parameters(self) -> None:
+        self.working_memory_bus.reset_identity_parameters()
+        self.persistent_memory_bus.reset_identity_parameters()
+
+
+    def forward(self, *, token_embeddings: torch.Tensor, confidence: torch.Tensor,
+                entropy: torch.Tensor, state: LatentDeliberationState,
+                canvas_head: torch.Tensor | None = None) -> LatentProcessorOutput:
+        batch, canvas, width = token_embeddings.shape
+        if state.memory_slots.shape != state.gdn2.persistent.shape:
+            raise ValueError("Persistent GDN2 state shape differs from memory slots.")
+        memory = replace(state.gdn2, persistent=state.memory_slots)
+        hidden = self.trajectory.read(memory, token_embeddings)
+        offsets = torch.arange(canvas, device=token_embeddings.device)[None, :].expand(batch, -1)
+        if canvas_head is not None:
+            offsets = (offsets - canvas_head[:, None]) % canvas
+        for block in self.blocks:
+            hidden = block(hidden, memory.seen, offsets)
+        hidden = self.output_norm(hidden) * memory.seen[..., None].to(hidden.dtype)
+        next_state = replace(state, confidence=confidence.float(), entropy=entropy.float(),
+                             gdn2=memory)
+        return LatentProcessorOutput(hidden, next_state)
+
+    def observe_state(self, state: LatentDeliberationState,
+                      heavy: torch.Tensor, working: torch.Tensor,
+                      live: torch.Tensor, head: torch.Tensor) -> LatentDeliberationState:
+        source = heavy.detach() + working
+        updated = self.trajectory.observe(state.gdn2, source, live, head)
+        return replace(state, gdn2=updated)
+
+
+    def commit_write(self, *, memory: torch.Tensor, working_state: torch.Tensor,
+                     heavy_hidden: torch.Tensor,
+                     committed_token_embeddings: torch.Tensor,
+                     commit_lengths: torch.Tensor,
+                     commit_reason: torch.Tensor | None = None,
+                     canvas_head: torch.Tensor | None = None):
+        batch, canvas, width = working_state.shape
+        count = int(commit_lengths.max().item())
+        if count <= 0:
+            return memory
+        index = torch.arange(count, device=working_state.device)[None, :].expand(batch, -1)
+        if canvas_head is not None:
+            index = (index + canvas_head[:, None]) % canvas
+        selected_working = working_state.gather(
+            1, index[..., None].expand(-1, -1, width)
+        )
+        selected_heavy = heavy_hidden.detach().gather(
+            1, index[..., None].expand(-1, -1, width)
+        )
+        if committed_token_embeddings.shape != selected_working.shape:
+            raise ValueError("Committed embeddings do not match the prefix.")
+        # Canvas processing already supplies bidirectional spatial context.
+        # Separate normalized role channels feed the ordered GDN2 writer directly.
+        # Unit-floor normalization keeps a zero Working state at zero without
+        # amplifying its derivative by 1/sqrt(eps) on the first denoise.
+        roles = tuple(F.rms_norm(value.float(), (width,), eps=1.0).to(value.dtype)
+                      for value in (selected_heavy, selected_working,
+                                    committed_token_embeddings.detach()))
+        packet = self.experience_in(torch.cat(roles, dim=-1))
+        reason = torch.zeros(batch, device=packet.device, dtype=torch.long) if commit_reason is None else commit_reason.long()
+        reason = torch.where((reason == 1) | (reason == 2), 5, reason).clamp(0, 5)
+        packet = packet + self.reason_embed(reason)[:, None]
+        valid = torch.arange(count, device=packet.device)[None, :] < commit_lengths[:, None]
+        written = self.trajectory.persistent.write_sequence(memory, packet, valid)
+        return written

lora.py

@@ -0,0 +1,96 @@
+"""Unfused dense and expert inference adapters."""
+from __future__ import annotations
+from dataclasses import dataclass
+import math
+import torch
+from torch import nn
+from transformers.models.diffusion_gemma.modeling_diffusion_gemma import DiffusionGemmaTextExperts
+from .mps_ops import attach_expert_lora
+
+@dataclass
+class LoRAConfig:
+    r: int
+    alpha: int
+    target_modules: tuple[str, ...]
+    expert_r: int = 8
+    expert_alpha: int = 8
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, base: nn.Linear, config: LoRAConfig):
+        super().__init__()
+        if config.r <= 0:
+            raise ValueError("LoRA rank must be positive.")
+        self.base = base
+        self.scaling = config.alpha / config.r
+        self.lora_a = nn.Parameter(
+            torch.empty(config.r, base.in_features, device=base.weight.device, dtype=base.weight.dtype)
+        )
+        self.lora_b = nn.Parameter(
+            torch.zeros(base.out_features, config.r, device=base.weight.device, dtype=base.weight.dtype)
+        )
+        nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5))
+        self.base.requires_grad_(False)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        update = torch.nn.functional.linear(
+            torch.nn.functional.linear(inputs, self.lora_a), self.lora_b,
+        )
+        return self.base(inputs) + self.scaling * update
+
+def get_parent_module(root: nn.Module, module_name: str) -> tuple[nn.Module, str]:
+    parent = root
+    parts = module_name.split(".")
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    return parent, parts[-1]
+
+def inject_lora(model: nn.Module, config: LoRAConfig) -> int:
+    targets = set(config.target_modules)
+    replacements = []
+    for name, module in model.named_modules():
+        if not isinstance(module, nn.Linear):
+            continue
+        if ".router." in name:
+            continue
+        in_trunk = ".layers." in name
+        if name.rsplit(".", 1)[-1] not in targets or not in_trunk:
+            continue
+        replacements.append((name, module))
+    for name, module in replacements:
+        parent, child = get_parent_module(model, name)
+        setattr(parent, child, LoRALinear(module, config))
+    expert_modules = [
+        module for module in model.modules()
+        if isinstance(module, DiffusionGemmaTextExperts)
+    ]
+    for module in expert_modules:
+        attach_expert_lora(
+            module,
+            rank=config.expert_r,
+            alpha=config.expert_alpha,
+        )
+    share_tied_lora_parameters(model)
+    if not replacements and not expert_modules:
+        raise RuntimeError("No LoRA target modules were found.")
+    return len(replacements) + len(expert_modules)
+
+def share_tied_lora_parameters(model: nn.Module) -> None:
+    """Use one adapter for the encoder/decoder pair whose base weights are tied."""
+    trunk = getattr(model, "model", model)
+    encoder = trunk.encoder.language_model
+    decoder = trunk.decoder
+    expert_names = ("lora_gate_up_a", "lora_gate_up_b", "lora_down_a", "lora_down_b")
+    for name, decoder_module in decoder.named_modules():
+        if not name:
+            continue
+        try:
+            encoder_module = encoder.get_submodule(name)
+        except AttributeError:
+            continue
+        if isinstance(decoder_module, LoRALinear) and isinstance(encoder_module, LoRALinear):
+            encoder_module.lora_a = decoder_module.lora_a
+            encoder_module.lora_b = decoder_module.lora_b
+        for parameter_name in expert_names:
+            if hasattr(decoder_module, parameter_name) and hasattr(encoder_module, parameter_name):
+                setattr(encoder_module, parameter_name, getattr(decoder_module, parameter_name))

model-00001-of-00032.safetensors

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

model-00002-of-00032.safetensors

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

model-00003-of-00032.safetensors

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

model-00004-of-00032.safetensors

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

model-00005-of-00032.safetensors

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

model-00006-of-00032.safetensors

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

model-00007-of-00032.safetensors

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

model-00008-of-00032.safetensors

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

model-00009-of-00032.safetensors

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

model-00010-of-00032.safetensors

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

model-00011-of-00032.safetensors

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

model-00012-of-00032.safetensors

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

model-00013-of-00032.safetensors

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

model-00014-of-00032.safetensors

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

model-00015-of-00032.safetensors

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

model-00016-of-00032.safetensors

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

model-00017-of-00032.safetensors

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

model-00018-of-00032.safetensors

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

model-00019-of-00032.safetensors

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

model-00020-of-00032.safetensors

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

model-00021-of-00032.safetensors

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

model-00022-of-00032.safetensors

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

model-00023-of-00032.safetensors

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

model-00024-of-00032.safetensors

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

model-00025-of-00032.safetensors

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

model-00026-of-00032.safetensors

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

model-00027-of-00032.safetensors

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

model-00028-of-00032.safetensors

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

model-00029-of-00032.safetensors

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

model-00030-of-00032.safetensors

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

model-00031-of-00032.safetensors

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

model-00032-of-00032.safetensors

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

modeling_modilify_mk2.py

@@ -1,11 +1,9 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
-"""Multimodal Modilify Mk2 model with recurrent latent deliberation."""
+"""Text-only ModilifyMk2 model with recurrent latent deliberation."""
 
 from __future__ import annotations
 
 from collections.abc import Sequence
-from dataclasses import dataclass, replace
+from dataclasses import dataclass
 from typing import Any
 
 import torch
@@ -13,7 +11,6 @@ from torch import nn
 from torch.nn import functional as F
 from transformers.cache_utils import Cache
 from transformers.modeling_outputs import BaseModelOutputWithPast
-from transformers.utils import ModelOutput
 
 from transformers.masking_utils import (
     ALL_MASK_ATTENTION_FUNCTIONS,
@@ -21,7 +18,7 @@ from transformers.masking_utils import (
 )
 from transformers.models.diffusion_gemma import (
     DiffusionGemmaDecoderModel,
-    DiffusionGemmaEncoderModel,
+    DiffusionGemmaEncoderTextModel,
     DiffusionGemmaPreTrainedModel,
 )
 from transformers.models.diffusion_gemma.modeling_diffusion_gemma import (
@@ -29,61 +26,58 @@ from transformers.models.diffusion_gemma.modeling_diffusion_gemma import (
     DiffusionGemmaTextRouter,
 )
 
-from .mps_ops import mps_segmented_experts_forward as _mps_segmented_experts_forward  # noqa: F401
-from .configuration_modilify_mk2 import ModilifyMk2Config
+from .mps_ops import register_mps_backends
+from .lora import LoRAConfig, inject_lora
+from .configuration_modilify_mk2 import ModilifyMk2Config, DENOISE_TEMPERATURE
 from .generation_modilify_mk2 import ModilifyMk2GenerationConfig, ModilifyMk2GenerationMixin
-from .latent_deliberation import (
-    LatentDeliberationState,
-    LatentDeliberationTransformer,
-    LatentProcessorOutput,
-    TrajectoryHistory,
-    TrajectoryTape,
-    empty_trajectory_tape,
-)
+from .latent_deliberation import LatentDeliberationState, LatentProcessorOutput, _sdpa_mask_value
+from .latent_deliberation import LatentDeliberationTransformer
 from .vocab_ops import chunked_vocab_statistics
 
 
-@dataclass
-class ModilifyMk2DecoderOutput(BaseModelOutputWithPast):
-    token_embeddings: torch.FloatTensor | None = None
-    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
+PRECISION_POLICY = "gdn2_small_fp32_v1"
+
+
+def small_parameter_fp32(name: str) -> bool:
+    return name.startswith("latent_deliberation.") and (
+        name.endswith((".dt_bias", ".a_log"))
+        or (name.endswith(".weight") and "norm" in name.rsplit(".", 2)[-2])
+    )
+
+
+register_mps_backends()
 
 
 @dataclass
 class ModilifyMk2ModelOutput(BaseModelOutputWithPast):
-    token_embeddings: torch.FloatTensor | None = None
     encoder_last_hidden_state: torch.FloatTensor | None = None
-    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
 
 
 @dataclass
-class ModilifyMk2BlockDiffusionOutput(ModelOutput):
-    """Inference output used by the rolling diffusion generator."""
-
-    logits: torch.FloatTensor | None = None
-    heavy_hidden_state: torch.FloatTensor | None = None
-    next_latent_state: LatentDeliberationState | None = None
+class ModilifyMk2BlockDiffusionOutput:
+    logits: torch.FloatTensor | None
+    heavy_hidden_state: torch.FloatTensor
+    next_latent_state: LatentDeliberationState
     past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
     encoder_last_hidden_state: torch.FloatTensor | None = None
-    temporal_context: torch.FloatTensor | None = None
-    history_projected: torch.FloatTensor | None = None
     working_state: torch.FloatTensor | None = None
-    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
     proposal: torch.LongTensor | None = None
     proposal_confidence: torch.FloatTensor | None = None
     token_entropy: torch.FloatTensor | None = None
     greedy_proposal: torch.LongTensor | None = None
     greedy_confidence: torch.FloatTensor | None = None
 
+    def to_tuple(self) -> tuple[Any, ...]:
+        return tuple(value for value in (
+            self.logits, self.past_key_values, self.hidden_states,
+            self.attentions, self.encoder_last_hidden_state,
+            self.heavy_hidden_state, self.next_latent_state,
+        ) if value is not None)
 
-class ModilifyMk2RMSNorm(DiffusionGemmaRMSNorm):
-    """Official RMSNorm parameters, pre-fusion same-dtype forward."""
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        normed_output = self._norm(hidden_states)
-        if self.with_scale:
-            normed_output = normed_output * self.weight.to(dtype=normed_output.dtype)
-        return normed_output.type_as(hidden_states)
+    def __getitem__(self, key: str | int) -> Any:
+        return getattr(self, key) if isinstance(key, str) else self.to_tuple()[key]
 
 
 class ModilifyMk2TextRouter(DiffusionGemmaTextRouter):
@@ -91,7 +85,7 @@ class ModilifyMk2TextRouter(DiffusionGemmaTextRouter):
 
     def __init__(self, config: Any) -> None:
         super().__init__(config)
-        self.norm = ModilifyMk2RMSNorm(self.hidden_size, eps=self.eps, with_scale=False)
+        self.norm = DiffusionGemmaRMSNorm(self.hidden_size, eps=self.eps, with_scale=False)
 
     def forward(
         self, hidden_states: torch.Tensor
@@ -114,18 +108,26 @@ class ModilifyMk2TextRouter(DiffusionGemmaTextRouter):
         return router_probabilities, top_k_weights, top_k_index
 
 
+def _finite_decoder_attention_mask(
+    mask: torch.Tensor | None,
+    *,
+    dtype: torch.dtype,
+) -> torch.Tensor | None:
+    """Keep boolean masks and finite values for fully blocked query rows."""
+
+    del dtype
+    if mask is None:
+        return None
+    if mask.dtype == torch.bool:
+        return mask | ~mask.any(dim=-1, keepdim=True)
+    return mask.clamp_min(_sdpa_mask_value(mask.dtype))
+
+
 def install_modilify_mk2_trunk_semantics(module: nn.Module) -> None:
     """Swap official leaf modules on this instance. Never patch Transformers classes."""
 
     for name, child in list(module.named_children()):
-        if type(child) is DiffusionGemmaRMSNorm:
-            dim = int(child.weight.shape[0]) if child.with_scale else 1
-            replacement = ModilifyMk2RMSNorm(
-                dim, eps=child.eps, with_scale=child.with_scale
-            )
-            replacement.load_state_dict(child.state_dict())
-            setattr(module, name, replacement)
-        elif type(child) is DiffusionGemmaTextRouter:
+        if type(child) is DiffusionGemmaTextRouter:
             replacement = ModilifyMk2TextRouter(child.config)
             replacement.load_state_dict(child.state_dict())
             setattr(module, name, replacement)
@@ -133,11 +135,26 @@ def install_modilify_mk2_trunk_semantics(module: nn.Module) -> None:
             install_modilify_mk2_trunk_semantics(child)
 
 
-class ModilifyMk2EncoderModel(DiffusionGemmaEncoderModel):
-    """Unmodified Transformers DiffusionGemma multimodal encoder."""
+class ModilifyMk2EncoderModel(DiffusionGemmaPreTrainedModel):
+    """Checkpoint-compatible wrapper around the text encoder trunk."""
 
     config_class = ModilifyMk2Config
 
+    def __init__(self, config: ModilifyMk2Config):
+        super().__init__(config)
+        self.language_model = DiffusionGemmaEncoderTextModel(config.text_config)
+        install_modilify_mk2_trunk_semantics(self)
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.language_model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.language_model.embed_tokens = value
+
+    def forward(self, **kwargs: Any) -> BaseModelOutputWithPast:
+        return self.language_model(**kwargs)
+
 
 class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
     """Diffusion decoder accepting compact self-conditioning embeddings."""
@@ -171,7 +188,12 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
         if isinstance(decoder_attention_mask, dict) and all(
             mask.ndim == 4 for mask in decoder_attention_mask.values()
         ):
-            return decoder_attention_mask
+            return {
+                name: _finite_decoder_attention_mask(
+                    mask, dtype=inputs_embeds.dtype
+                )
+                for name, mask in decoder_attention_mask.items()
+            }
 
         text_config = config.get_text_config() if hasattr(config, "get_text_config") else config
         q_length = inputs_embeds.shape[1]
@@ -193,23 +215,24 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
                 max_length = sliding_layer.get_max_length() + additional_kv_length
                 if kv_length >= max_length:
                     kv_length = max_length
-            mask_mapping[layer_pattern] = ALL_MASK_ATTENTION_FUNCTIONS[
-                config._attn_implementation
-            ](
-                batch_size=inputs_embeds.shape[0],
-                q_length=q_length,
-                kv_length=kv_length,
-                q_offset=q_offset,
-                kv_offset=kv_offset,
-                mask_function=bidirectional_mask_function,
-                attention_mask=decoder_attention_mask,
-                allow_is_causal_skip=False,
-                allow_is_bidirectional_skip=True,
-                local_size=getattr(text_config, "sliding_window", None),
+            mask_mapping[layer_pattern] = _finite_decoder_attention_mask(
+                ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation](
+                    batch_size=inputs_embeds.shape[0],
+                    q_length=q_length,
+                    kv_length=kv_length,
+                    q_offset=q_offset,
+                    kv_offset=kv_offset,
+                    mask_function=bidirectional_mask_function,
+                    attention_mask=decoder_attention_mask,
+                    allow_is_causal_skip=False,
+                    allow_is_bidirectional_skip=True,
+                    local_size=getattr(text_config, "sliding_window", None),
+                    dtype=inputs_embeds.dtype,
+                    config=text_config,
+                    use_vmap=False,
+                    device=inputs_embeds.device,
+                ),
                 dtype=inputs_embeds.dtype,
-                config=text_config,
-                use_vmap=False,
-                device=inputs_embeds.device,
             )
         return mask_mapping
 
@@ -217,9 +240,7 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
         self,
         token_embeddings: torch.Tensor,
         latent_context: torch.Tensor | None,
-        *,
-        collect_diagnostics: bool = True,
-    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    ) -> torch.Tensor:
         """Map latent context through the frozen native self-conditioning bridge."""
 
         if latent_context is None:
@@ -235,10 +256,6 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
             * mapper.up_proj(normalized_context)
         )
         mapped_fp32 = mapped_context.float()
-        # Use the energy directly in the cap denominator.  Computing
-        # sqrt(E[x²]) and immediately squaring it again has an undefined
-        # backward at the identity-init point x=0 (0/0 in d(sqrt)/dx), which
-        # poisoned the detached temporal-context VJP on every denoise.
         mapped_token_energy = mapped_fp32.square().mean(dim=-1, keepdim=True)
         token_rms_per_token = (
             token_embeddings.float().square().mean(dim=-1, keepdim=True).sqrt()
@@ -249,21 +266,7 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
         )
         mapped_context = (mapped_fp32 * soft_cap_scale).to(dtype=mapped_context.dtype)
         combined = token_embeddings + mapped_context
-        diagnostics: dict[str, torch.Tensor] = {}
-        if collect_diagnostics:
-            token_rms = token_embeddings.detach().float().square().mean().sqrt()
-            mapped_rms = mapped_context.detach().float().square().mean().sqrt()
-            diagnostics = {
-                "token_embedding_rms": token_rms,
-                "latent_context_rms": context.detach().float().square().mean().sqrt(),
-                "mapped_sc_rms": mapped_rms,
-                "actual_residual_rms": mapped_rms,
-                "latent_to_embedding_rms_ratio": (
-                    mapped_rms / token_rms.clamp_min(1.0e-12)
-                ),
-                "merged_input_rms": combined.detach().float().square().mean().sqrt(),
-            }
-        return mapper.post_norm(combined), diagnostics
+        return mapper.post_norm(combined)
 
     def _run_stack(
         self,
@@ -288,10 +291,8 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
             )
         working_bus = kwargs.pop("working_bus", None)
         working_state = kwargs.pop("working_state", None)
+        working_canvas_head = kwargs.pop("working_canvas_head", None)
         persistent_bus = kwargs.pop("persistent_bus", None)
-        memory_bus = kwargs.pop("memory_bus", None)
-        if memory_bus is None:
-            memory_bus = persistent_bus
         memory_slots = kwargs.pop("memory_slots", None)
         slot_identity = kwargs.pop("slot_identity", None)
         cache = past_key_values
@@ -304,8 +305,8 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
         if working_bus is not None and working_state is not None:
             working_kv = working_bus.prepare_kv(working_state)
         memory_kv = None
-        if memory_bus is not None and memory_slots is not None:
-            memory_kv = memory_bus.prepare_kv(memory_slots, slot_identity)
+        if persistent_bus is not None and memory_slots is not None:
+            memory_kv = persistent_bus.prepare_kv(memory_slots, slot_identity)
         working_reader = 0
         reader_index = 0
         for index in range(self.text_config.num_hidden_layers):
@@ -325,14 +326,17 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
                     and working_reader < working_bus.num_readers
                 ):
                     hidden = working_bus.read(
-                        hidden, working_reader, working_kv[0], working_kv[1]
+                        hidden, working_reader, *working_kv,
+                        positions=(
+                            decoder_position_ids if working_canvas_head is not None else None
+                        ),
                     )
                     working_reader += 1
                 if (
                     memory_kv is not None
-                    and reader_index < memory_bus.num_readers
+                    and reader_index < persistent_bus.num_readers
                 ):
-                    hidden = memory_bus.read(
+                    hidden = persistent_bus.read(
                         hidden, reader_index, memory_kv[0], memory_kv[1]
                     )
                     reader_index += 1
@@ -346,15 +350,13 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
         temporal_context_embeddings: torch.FloatTensor | None = None,
         decoder_attention_mask: torch.Tensor | dict | None = None,
         decoder_position_ids: torch.LongTensor | None = None,
-        collect_latent_diagnostics: bool = True,
-        memory_bus: Any | None = None,
         memory_slots: torch.Tensor | None = None,
         working_bus: Any | None = None,
         working_state: torch.Tensor | None = None,
         persistent_bus: Any | None = None,
         slot_identity: torch.Tensor | None = None,
         **kwargs: Any,
-    ) -> ModilifyMk2DecoderOutput:
+    ) -> BaseModelOutputWithPast:
         if "use_cache" in kwargs:
             raise ValueError("The diffusion decoder always reads the supplied cache.")
         if decoder_token_embeddings is None:
@@ -364,22 +366,15 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
             expected = (*decoder_input_ids.shape, self.text_config.hidden_size)
             if token_embeddings.shape != expected:
                 raise ValueError("Precomputed decoder embeddings have the wrong shape.")
-        context_embeddings = (
-            torch.zeros_like(token_embeddings)
-            if temporal_context_embeddings is None
-            else temporal_context_embeddings.to(token_embeddings)
-        )
-        inputs_embeds, diagnostics = self.merge_latent_context(
+        inputs_embeds = self.merge_latent_context(
             token_embeddings,
-            context_embeddings if temporal_context_embeddings is not None else None,
-            collect_diagnostics=collect_latent_diagnostics,
+            temporal_context_embeddings,
         )
         hidden = self._run_stack(
             inputs_embeds,
             past_key_values=past_key_values,
             decoder_attention_mask=decoder_attention_mask,
             decoder_position_ids=decoder_position_ids,
-            memory_bus=memory_bus,
             memory_slots=memory_slots,
             working_bus=working_bus,
             working_state=working_state,
@@ -387,11 +382,9 @@ class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
             slot_identity=slot_identity,
             **kwargs,
         )
-        return ModilifyMk2DecoderOutput(
+        return BaseModelOutputWithPast(
             last_hidden_state=hidden,
             past_key_values=past_key_values,
-            token_embeddings=token_embeddings,
-            latent_residual_diagnostics=diagnostics,
         )
 
 
@@ -412,6 +405,16 @@ class ModilifyMk2Model(DiffusionGemmaPreTrainedModel):
         self.encoder = ModilifyMk2EncoderModel(config)
         self.decoder = ModilifyMk2DecoderModel(config)
         install_modilify_mk2_trunk_semantics(self)
+        if getattr(config, "lora_config", None):
+            config.text_config._experts_implementation = "mps_segmented"
+            inject_lora(self, LoRAConfig(**config.lora_config))
+            self._tied_weights_keys = {
+                **self._tied_weights_keys,
+                r"encoder.language_model.layers\.(?:[^.]+\.)*lora_[ab]$":
+                    r"decoder.layers\.(?:[^.]+\.)*lora_[ab]$",
+                r"encoder.language_model.layers\.(?:[^.]+\.)*lora_(?:gate_up|down)_[ab]$":
+                    r"decoder.layers\.(?:[^.]+\.)*lora_(?:gate_up|down)_[ab]$",
+            }
         self.post_init()
 
     def get_encoder(self):
@@ -440,8 +443,6 @@ class ModilifyMk2Model(DiffusionGemmaPreTrainedModel):
         decoder_attention_mask: torch.Tensor | dict | None = None,
         decoder_position_ids: torch.LongTensor | None = None,
         return_encoder_outputs: bool = True,
-        collect_latent_diagnostics: bool = True,
-        memory_bus: Any | None = None,
         memory_slots: torch.Tensor | None = None,
         working_bus: Any | None = None,
         working_state: torch.Tensor | None = None,
@@ -450,15 +451,13 @@ class ModilifyMk2Model(DiffusionGemmaPreTrainedModel):
         **kwargs: Any,
     ) -> ModilifyMk2ModelOutput:
         encoder_hidden = None
-        encoder_keys = ("pixel_values", "mm_token_type_ids", "image_position_ids", "inputs_embeds")
-        encoder_kwargs = {key: kwargs.pop(key) for key in encoder_keys if key in kwargs}
         if input_ids is not None:
             encoded = self.encoder(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
                 past_key_values=past_key_values,
                 position_ids=position_ids,
-                **encoder_kwargs,
+                **kwargs,
             )
             past_key_values = encoded.past_key_values
             if return_encoder_outputs:
@@ -472,8 +471,6 @@ class ModilifyMk2Model(DiffusionGemmaPreTrainedModel):
             temporal_context_embeddings=temporal_context_embeddings,
             decoder_attention_mask=decoder_attention_mask,
             decoder_position_ids=decoder_position_ids,
-            collect_latent_diagnostics=collect_latent_diagnostics,
-            memory_bus=memory_bus,
             memory_slots=memory_slots,
             working_bus=working_bus,
             working_state=working_state,
@@ -484,9 +481,7 @@ class ModilifyMk2Model(DiffusionGemmaPreTrainedModel):
         return ModilifyMk2ModelOutput(
             last_hidden_state=decoded.last_hidden_state,
             past_key_values=past_key_values,
-            token_embeddings=decoded.token_embeddings,
             encoder_last_hidden_state=encoder_hidden,
-            latent_residual_diagnostics=decoded.latent_residual_diagnostics,
         )
 
 
@@ -494,14 +489,13 @@ class ModilifyMk2ForBlockDiffusion(DiffusionGemmaPreTrainedModel, ModilifyMk2Gen
     config_class = ModilifyMk2Config
     _tied_weights_keys = {"lm_head.weight": "model.decoder.embed_tokens.weight"}
     generation_config_class = ModilifyMk2GenerationConfig
+    supports_gradient_checkpointing = False
 
     @torch.no_grad()
     def _init_weights(self, module: nn.Module) -> None:
         super()._init_weights(module)
         if isinstance(module, LatentDeliberationTransformer):
             module.reset_identity_parameters()
-            module.working_memory_bus.freeze()
-            module.persistent_memory_bus.freeze()
 
     def __init__(self, config: ModilifyMk2Config):
         super().__init__(config)
@@ -509,15 +503,11 @@ class ModilifyMk2ForBlockDiffusion(DiffusionGemmaPreTrainedModel, ModilifyMk2Gen
         layer_types = tuple(getattr(config.text_config, "layer_types", None) or ())
         self.latent_deliberation = LatentDeliberationTransformer(
             hidden_size=config.text_config.hidden_size,
-            vocab_size=config.text_config.vocab_size,
             latent_dim=config.latent_dim,
             ffn_dim=config.latent_ffn_dim,
-            memory_slots=config.latent_memory_slots,
             num_layers=config.latent_num_layers,
             num_heads=config.latent_num_heads,
             local_attention_window=config.latent_local_attention_window,
-            dropout=config.latent_dropout,
-            history_length=config.latent_history_length,
             tape_probes=config.latent_tape_probes,
             history_kv_rank=config.latent_history_kv_rank,
             num_memory_readers=sum(layer_type == "full_attention" for layer_type in layer_types),
@@ -530,13 +520,17 @@ class ModilifyMk2ForBlockDiffusion(DiffusionGemmaPreTrainedModel, ModilifyMk2Gen
                 if config.persistent_memory_bus else 0
             ),
             working_last_block_global=config.latent_working_last_block_global,
-            experience_roles=config.experience_roles,
-            commit_sequence_layers=config.commit_sequence_layers,
             commit_sequence_dim=config.commit_sequence_dim,
             max_canvas_length=config.canvas_length,
         )
         self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
         self.final_logit_softcapping = config.text_config.final_logit_softcapping
+        policy = getattr(config, "precision_policy", PRECISION_POLICY)
+        if policy != PRECISION_POLICY:
+            raise ValueError(f"Unsupported export precision policy: {policy}")
+        self._keep_in_fp32_modules_strict = {
+            name for name, _ in self.named_parameters() if small_parameter_fp32(name)
+        }
         self.post_init()
         install_modilify_mk2_trunk_semantics(self)
 
@@ -544,192 +538,101 @@ class ModilifyMk2ForBlockDiffusion(DiffusionGemmaPreTrainedModel, ModilifyMk2Gen
         logits = self.lm_head(hidden)
         return torch.tanh(logits / self.final_logit_softcapping) * self.final_logit_softcapping
 
+
     def _prepare_latent_context(
         self,
         decoder_input_ids: torch.LongTensor,
         *,
-        history: TrajectoryHistory | None,
-        tape: TrajectoryTape | None,
+
         confidence: torch.Tensor | None,
         entropy: torch.Tensor | None,
-        age: torch.Tensor | None,
         latent_state: LatentDeliberationState | None,
-    ) -> tuple[torch.Tensor, LatentDeliberationState, torch.Tensor, torch.Tensor]:
+        canvas_head: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, LatentDeliberationState, torch.Tensor]:
         batch, canvas = decoder_input_ids.shape
-        dtype = self.model.decoder.embed_tokens.weight.dtype
         if latent_state is None:
             latent_state = LatentDeliberationState.empty(
                 batch_size=batch, canvas_length=canvas,
-                latent_dim=self.config.latent_dim, memory_slots=self.config.latent_memory_slots,
-                device=decoder_input_ids.device, dtype=dtype,
-            )
-        if history is None:
-            history = TrajectoryHistory.empty(
-                batch_size=batch,
-                canvas_length=canvas,
-                hidden_size=self.config.text_config.hidden_size,
-                history_length=self.config.latent_history_length,
                 device=decoder_input_ids.device,
-                dtype=dtype,
-            )
-        if tape is None:
-            tape = empty_trajectory_tape(
-                batch_size=batch,
-                config=self.config,
-                device=decoder_input_ids.device,
-                dtype=dtype,
             )
         confidence = latent_state.confidence if confidence is None else confidence.squeeze(-1).float()
         entropy = latent_state.entropy if entropy is None else entropy.squeeze(-1).float()
-        if age is not None:
-            latent_state = replace(
-                latent_state,
-                age=age.to(device=decoder_input_ids.device, dtype=torch.int32),
-            )
         token_embeddings = self.model.decoder.embed_tokens(decoder_input_ids)
         processed: LatentProcessorOutput = self.latent_deliberation(
             token_embeddings=token_embeddings,
             confidence=confidence,
             entropy=entropy,
             state=latent_state,
-            history=history,
-            tape=tape,
+
+            canvas_head=canvas_head,
         )
+        latent_context = processed.context
+        next_state = processed.state
         return (
-            processed.context,
-            processed.state,
+            latent_context,
+            next_state,
             token_embeddings,
-            processed.history_projected,
         )
 
     def forward(
-        self,
-        *,
-        input_ids: torch.LongTensor | None = None,
+        self, *, input_ids: torch.LongTensor | None = None,
         attention_mask: torch.Tensor | dict | None = None,
         past_key_values: Cache | None = None,
         position_ids: torch.LongTensor | None = None,
         decoder_input_ids: torch.LongTensor,
         previous_confidence: torch.FloatTensor | None = None,
         previous_entropy: torch.FloatTensor | None = None,
-        token_age: torch.Tensor | None = None,
         latent_state: LatentDeliberationState | None = None,
-        history: TrajectoryHistory | None = None,
-        tape: TrajectoryTape | None = None,
-        history_hidden_state: TrajectoryHistory | torch.FloatTensor | None = None,
+        canvas_head: torch.Tensor | None = None,
         decoder_attention_mask: torch.Tensor | dict | None = None,
         decoder_position_ids: torch.LongTensor | None = None,
         return_encoder_outputs: bool = True,
         compact_vocab: bool = False,
-        denoise_temperature: float | None = None,
         repetition_token_mask: torch.BoolTensor | None = None,
         repetition_penalty: float = 1.0,
         sampling_generators: Sequence[torch.Generator] | None = None,
-        collect_latent_diagnostics: bool = True,
         **kwargs: Any,
     ) -> ModilifyMk2BlockDiffusionOutput:
-        if history is None and isinstance(history_hidden_state, TrajectoryHistory):
-            history = history_hidden_state
-        (
-            latent_context,
-            next_state,
-            decoder_token_embeddings,
-            history_projected,
-        ) = self._prepare_latent_context(
-            decoder_input_ids,
-            history=history,
-            tape=tape,
-            confidence=previous_confidence,
-            entropy=previous_entropy,
-            age=token_age,
-            latent_state=latent_state,
+        latent_context, next_state, token_embeddings = self._prepare_latent_context(
+            decoder_input_ids, confidence=previous_confidence, entropy=previous_entropy,
+             latent_state=latent_state, canvas_head=canvas_head,
         )
         working_bus = self.latent_deliberation.working_memory_bus
         persistent_bus = self.latent_deliberation.persistent_memory_bus
-        working_readers_active = working_bus.num_readers > 0
-        memory_readers_active = persistent_bus.num_readers > 0
         outputs = self.model(
             input_ids=input_ids, attention_mask=attention_mask,
             past_key_values=past_key_values, position_ids=position_ids,
-            decoder_input_ids=decoder_input_ids,
-            decoder_token_embeddings=decoder_token_embeddings,
+            decoder_input_ids=decoder_input_ids, decoder_token_embeddings=token_embeddings,
             temporal_context_embeddings=latent_context,
             decoder_attention_mask=decoder_attention_mask,
             decoder_position_ids=decoder_position_ids,
             return_encoder_outputs=return_encoder_outputs,
-            collect_latent_diagnostics=collect_latent_diagnostics,
-            working_bus=working_bus if working_readers_active else None,
-            working_state=latent_context if working_readers_active else None,
-            persistent_bus=persistent_bus if memory_readers_active else None,
-            memory_slots=next_state.memory_slots if memory_readers_active else None,
-            slot_identity=(
-                self.latent_deliberation.scaled_memory_slot_identity(
-                    batch_size=decoder_input_ids.shape[0],
-                    device=decoder_input_ids.device,
-                    dtype=latent_context.dtype,
-                )
-                if memory_readers_active else None
-            ),
-            **kwargs,
-        )
-        temperature = (
-            self.config.denoise_temperature
-            if denoise_temperature is None
-            else float(denoise_temperature)
+            working_bus=working_bus,
+            working_state=(latent_context, next_state.gdn2.seen),
+        working_canvas_head=canvas_head,
+        persistent_bus=persistent_bus, memory_slots=next_state.memory_slots,
+        slot_identity=next_state.gdn2.seen,
+        **kwargs,
         )
-        proposal = proposal_confidence = token_entropy = None
-        greedy_proposal = greedy_confidence = None
-        logits = None
+        proposal = proposal_confidence = token_entropy = greedy_proposal = greedy_confidence = None
         if compact_vocab:
-            (
-                proposal,
-                proposal_confidence,
-                token_entropy,
-                greedy_proposal,
-                greedy_confidence,
-            ) = chunked_vocab_statistics(
-                outputs.last_hidden_state.detach(),
-                self.lm_head.weight.detach(),
-                softcap=self.final_logit_softcapping,
-                temperature=temperature,
+            logits = None
+            proposal, proposal_confidence, token_entropy, greedy_proposal, greedy_confidence = chunked_vocab_statistics(
+                outputs.last_hidden_state, self.lm_head.weight,
+                softcap=self.final_logit_softcapping, temperature=DENOISE_TEMPERATURE,
                 chunk_size=self.config.vocab_chunk_size,
-                repetition_token_mask=repetition_token_mask,
-                repetition_penalty=repetition_penalty,
+                repetition_token_mask=repetition_token_mask, repetition_penalty=repetition_penalty,
                 sampling_generators=sampling_generators,
+                top_k=self.config.commit_top_k, min_p=self.config.commit_min_p,
             )
         else:
             logits = self._finalize_logits(outputs.last_hidden_state)
         return ModilifyMk2BlockDiffusionOutput(
-            logits=logits,
-            past_key_values=outputs.past_key_values,
+            logits=logits, past_key_values=outputs.past_key_values,
             encoder_last_hidden_state=outputs.encoder_last_hidden_state,
-            heavy_hidden_state=outputs.last_hidden_state,
-            next_latent_state=next_state,
-            temporal_context=latent_context,
-            history_projected=history_projected,
-            working_state=latent_context,
-            latent_residual_diagnostics=outputs.latent_residual_diagnostics,
-            proposal=proposal,
-            proposal_confidence=proposal_confidence,
-            token_entropy=token_entropy,
-            greedy_proposal=greedy_proposal,
+            heavy_hidden_state=outputs.last_hidden_state, next_latent_state=next_state,
+             working_state=latent_context,
+            proposal=proposal, proposal_confidence=proposal_confidence,
+            token_entropy=token_entropy, greedy_proposal=greedy_proposal,
             greedy_confidence=greedy_confidence,
         )
-
-
-ModilifyMk2Model.register_for_auto_class("AutoModel")
-ModilifyMk2ForBlockDiffusion.register_for_auto_class("AutoModelForCausalLM")
-ModilifyMk2ForBlockDiffusion.register_for_auto_class("AutoModelForMultimodalLM")
-
-
-__all__ = [
-    "ModilifyMk2BlockDiffusionOutput",
-    "ModilifyMk2Config",
-    "ModilifyMk2DecoderModel",
-    "ModilifyMk2EncoderModel",
-    "ModilifyMk2ForBlockDiffusion",
-    "ModilifyMk2Model",
-    "ModilifyMk2RMSNorm",
-    "ModilifyMk2TextRouter",
-    "install_modilify_mk2_trunk_semantics",
-]

mps_ops.py

@@ -1,5 +1,3 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
 """MPS-specific kernels that preserve the model's mathematical operations."""
 
 from __future__ import annotations
@@ -8,10 +6,75 @@ import math
 
 import torch
 from torch import nn
-from torch.nn import functional as F
 from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS, _grouped_linear
 
 
+def attach_expert_lora(
+    experts: nn.Module,
+    *,
+    rank: int,
+    alpha: int,
+) -> None:
+    """Attach packed LoRA parameters to one expert module instance."""
+    if rank <= 0:
+        raise ValueError("Expert LoRA rank must be positive.")
+    if alpha <= 0:
+        raise ValueError("Expert LoRA alpha must be positive.")
+    parameter_names = (
+        "lora_gate_up_a",
+        "lora_gate_up_b",
+        "lora_down_a",
+        "lora_down_b",
+    )
+    existing = tuple(hasattr(experts, name) for name in parameter_names)
+    if any(existing):
+        if not all(existing):
+            raise RuntimeError("Expert LoRA parameters are only partially initialized.")
+        return
+
+    experts.lora_scaling = float(alpha) / rank
+    experts.lora_gate_up_a = nn.Parameter(
+        torch.empty(
+            experts.num_experts,
+            rank,
+            experts.hidden_dim,
+            device=experts.gate_up_proj.device,
+            dtype=experts.gate_up_proj.dtype,
+        )
+    )
+    experts.lora_gate_up_b = nn.Parameter(
+        torch.zeros(
+            experts.num_experts,
+            2 * experts.intermediate_dim,
+            rank,
+            device=experts.gate_up_proj.device,
+            dtype=experts.gate_up_proj.dtype,
+        )
+    )
+    experts.lora_down_a = nn.Parameter(
+        torch.empty(
+            experts.num_experts,
+            rank,
+            experts.intermediate_dim,
+            device=experts.down_proj.device,
+            dtype=experts.down_proj.dtype,
+        )
+    )
+    experts.lora_down_b = nn.Parameter(
+        torch.zeros(
+            experts.num_experts,
+            experts.hidden_dim,
+            rank,
+            device=experts.down_proj.device,
+            dtype=experts.down_proj.dtype,
+        )
+    )
+    nn.init.kaiming_uniform_(experts.lora_gate_up_a, a=math.sqrt(5))
+    nn.init.kaiming_uniform_(experts.lora_down_a, a=math.sqrt(5))
+    experts.gate_up_proj.requires_grad_(False)
+    experts.down_proj.requires_grad_(False)
+
+
 def lora_aware_eager_experts_forward(
     self: nn.Module,
     hidden_states: torch.Tensor,
@@ -24,7 +87,6 @@ def lora_aware_eager_experts_forward(
         expert_mask = expert_mask.permute(2, 1, 0)
         expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
 
-    lora_dropout = getattr(self, "lora_dropout", nn.Identity())
     lora_scaling = getattr(self, "lora_scaling", 1.0)
 
     for expert_idx in expert_hit:
@@ -36,7 +98,7 @@ def lora_aware_eager_experts_forward(
         gate_up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx])
         if hasattr(self, "lora_gate_up_a"):
             update = nn.functional.linear(
-                nn.functional.linear(lora_dropout(current_state), self.lora_gate_up_a[expert_idx]),
+                nn.functional.linear(current_state, self.lora_gate_up_a[expert_idx]),
                 self.lora_gate_up_b[expert_idx],
             )
             gate_up = gate_up + lora_scaling * update
@@ -45,7 +107,7 @@ def lora_aware_eager_experts_forward(
         expert_output = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
         if hasattr(self, "lora_down_a"):
             update = nn.functional.linear(
-                nn.functional.linear(lora_dropout(current_hidden_states), self.lora_down_a[expert_idx]),
+                nn.functional.linear(current_hidden_states, self.lora_down_a[expert_idx]),
                 self.lora_down_b[expert_idx],
             )
             expert_output = expert_output + lora_scaling * update
@@ -55,120 +117,15 @@ def lora_aware_eager_experts_forward(
     return final_hidden_states
 
 
-def grouped_expert_offsets(
-    sorted_expert_ids: torch.Tensor,
-    num_experts: int,
-) -> torch.Tensor:
-    """Build grouped-mm offsets without MPS ``histc`` host synchronization."""
-    counts = torch.bincount(sorted_expert_ids, minlength=num_experts)[:num_experts]
-    return torch.cumsum(counts, dim=0, dtype=torch.int32)
-
-
-def mps_grouped_mm_experts_forward(
-    self: torch.nn.Module,
-    hidden_states: torch.Tensor,
-    top_k_index: torch.Tensor,
-    top_k_weights: torch.Tensor,
-) -> torch.Tensor:
-    """Transformers grouped-mm MoE with histogram-free MPS routing offsets."""
-    num_top_k = top_k_index.size(-1)
-    num_tokens = hidden_states.size(0)
-    hidden_dim = hidden_states.size(-1)
-    sample_weights = top_k_weights.reshape(-1)
-    expert_ids = top_k_index.reshape(-1)
-
-    expert_ids_g, perm = torch.sort(expert_ids)
-    selected_hidden_states_g = hidden_states[perm // num_top_k]
-    sample_weights_g = sample_weights[perm]
-    offsets = grouped_expert_offsets(expert_ids_g, self.num_experts)
-
-    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
-    expert_ids_g.clamp_(max=self.num_experts - 1)
-    selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
-
-    gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj
-    gate_up_biases = (
-        self.gate_up_proj_bias[expert_ids_g]
-        if self.has_gate and self.has_bias
-        else self.up_proj_bias[expert_ids_g]
-        if self.has_bias
-        else None
-    )
-    projected = _grouped_linear(
-        selected_hidden_states_g,
-        gate_up_weights,
-        offsets,
-        bias=gate_up_biases,
-        is_transposed=self.is_transposed,
-    )
-    projected = self._apply_gate(projected) if self.has_gate else self.act_fn(projected)
-
-    down_biases = self.down_proj_bias[expert_ids_g] if self.has_bias else None
-    projected = _grouped_linear(
-        projected,
-        self.down_proj,
-        offsets,
-        bias=down_biases,
-        is_transposed=self.is_transposed,
-    )
-    weighted = projected * sample_weights_g.unsqueeze(-1)
-    weighted.masked_fill_(sentinel_mask, 0.0)
-
-    # Scatter directly back to the original top-k order.  Constructing the
-    # inverse permutation and then gathering performs the same permutation
-    # with an extra index tensor and an extra read pass over `weighted`.
-    reordered = torch.empty_like(weighted)
-    reordered.index_copy_(0, perm, weighted)
-    weighted = reordered
-    return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)
-
-
-ALL_EXPERTS_FUNCTIONS.register("mps_grouped_mm", mps_grouped_mm_experts_forward)
-
-
-class _IndependentExpertLoRA(torch.autograd.Function):
-    """Capacity-padded reference path for independent expert LoRA."""
-
-    @staticmethod
-    def forward(
-        ctx,
-        inputs: torch.Tensor,
-        weight_a: torch.Tensor,
-        weight_b: torch.Tensor,
-        expert_ids: torch.LongTensor,
-        counts: torch.LongTensor,
-        capacity: int,
-    ) -> torch.Tensor:
-        offsets = counts.cumsum(0)
-        starts = offsets - counts
-        slots = torch.arange(inputs.shape[0], device=inputs.device) - starts[expert_ids]
-        padded = inputs.new_zeros((weight_a.shape[0], capacity, inputs.shape[-1]))
-        padded[expert_ids, slots] = inputs
-        low_rank = torch.bmm(padded, weight_a.transpose(1, 2))
-        updates = torch.bmm(low_rank, weight_b.transpose(1, 2))
-        ctx.save_for_backward(
-            padded,
-            low_rank,
-            weight_a,
-            weight_b,
-            expert_ids,
-            slots,
-        )
-        return updates[expert_ids, slots]
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor):
-        padded, low_rank, weight_a, weight_b, expert_ids, slots = ctx.saved_tensors
-        grad_padded = grad_output.new_zeros(
-            (weight_a.shape[0], padded.shape[1], grad_output.shape[-1])
-        )
-        grad_padded[expert_ids, slots] = grad_output
-        grad_weight_b = torch.bmm(grad_padded.transpose(1, 2), low_rank)
-        grad_low_rank = torch.bmm(grad_padded, weight_b)
-        grad_weight_a = torch.bmm(grad_low_rank.transpose(1, 2), padded)
-        grad_inputs_padded = torch.bmm(grad_low_rank, weight_a)
-        grad_inputs = grad_inputs_padded[expert_ids, slots]
-        return grad_inputs, grad_weight_a, grad_weight_b, None, None, None
+def _padded_expert_lora(inputs: torch.Tensor, weight_a: torch.Tensor, weight_b: torch.Tensor, expert_ids: torch.LongTensor, counts: torch.LongTensor, capacity: int) -> torch.Tensor:
+    offsets = counts.cumsum(0)
+    starts = offsets - counts
+    slots = torch.arange(inputs.shape[0], device=inputs.device) - starts[expert_ids]
+    padded = inputs.new_zeros((weight_a.shape[0], capacity, inputs.shape[-1]))
+    padded[expert_ids, slots] = inputs
+    low_rank = torch.bmm(padded, weight_a.transpose(1, 2))
+    updates = torch.bmm(low_rank, weight_b.transpose(1, 2))
+    return updates[expert_ids, slots]
 
 
 _INDEPENDENT_EXPERT_LORA_METAL_SOURCE = r"""
@@ -238,116 +195,6 @@ kernel void expert_lora_forward_b(
   output[index] = bfloat(value);
 }
 
-kernel void expert_lora_backward_low_rank(
-    device const bfloat* grad_output [[buffer(0)]],
-    device const bfloat* weight_b [[buffer(1)]],
-    device const long* expert_ids [[buffer(2)]],
-    device bfloat* grad_low_rank [[buffer(3)]],
-    constant uint& row_count [[buffer(4)]],
-    constant uint& output_dim [[buffer(5)]],
-    uint row [[threadgroup_position_in_grid]],
-    uint lane [[thread_index_in_threadgroup]]) {
-  if (row >= row_count) return;
-  const uint expert = uint(expert_ids[row]);
-  threadgroup float partial[kThreads * kRank];
-  for (uint rank = 0; rank < kRank; ++rank) {
-    float value = 0.0f;
-    const uint grad_base = row * output_dim;
-    for (uint column = lane; column < output_dim; column += kThreads) {
-      const uint b_index =
-          (expert * output_dim + column) * kRank + rank;
-      value += float(grad_output[grad_base + column])
-          * float(weight_b[b_index]);
-    }
-    partial[lane * kRank + rank] = value;
-  }
-  threadgroup_barrier(mem_flags::mem_threadgroup);
-  for (uint stride = kThreads / 2; stride > 0; stride >>= 1) {
-    if (lane < stride) {
-      for (uint rank = 0; rank < kRank; ++rank) {
-        partial[lane * kRank + rank] +=
-            partial[(lane + stride) * kRank + rank];
-      }
-    }
-    threadgroup_barrier(mem_flags::mem_threadgroup);
-  }
-  if (lane == 0) {
-    for (uint rank = 0; rank < kRank; ++rank) {
-      grad_low_rank[row * kRank + rank] = bfloat(partial[rank]);
-    }
-  }
-}
-
-kernel void expert_lora_backward_inputs(
-    device const bfloat* grad_low_rank [[buffer(0)]],
-    device const bfloat* weight_a [[buffer(1)]],
-    device const long* expert_ids [[buffer(2)]],
-    device bfloat* grad_inputs [[buffer(3)]],
-    constant uint& row_count [[buffer(4)]],
-    constant uint& input_dim [[buffer(5)]],
-    uint index [[thread_position_in_grid]]) {
-  const uint total = row_count * input_dim;
-  if (index >= total) return;
-  const uint row = index / input_dim;
-  const uint column = index - row * input_dim;
-  const uint expert = uint(expert_ids[row]);
-  float value = 0.0f;
-  for (uint rank = 0; rank < kRank; ++rank) {
-    const uint a_index =
-        (expert * kRank + rank) * input_dim + column;
-    value += float(grad_low_rank[row * kRank + rank])
-        * float(weight_a[a_index]);
-  }
-  grad_inputs[index] = bfloat(value);
-}
-
-kernel void expert_lora_backward_a(
-    device const bfloat* grad_low_rank [[buffer(0)]],
-    device const bfloat* inputs [[buffer(1)]],
-    device const long* offsets [[buffer(2)]],
-    device bfloat* grad_weight_a [[buffer(3)]],
-    constant uint& expert_count [[buffer(4)]],
-    constant uint& input_dim [[buffer(5)]],
-    uint index [[thread_position_in_grid]]) {
-  const uint total = expert_count * kRank * input_dim;
-  if (index >= total) return;
-  const uint column = index % input_dim;
-  const uint rank_expert = index / input_dim;
-  const uint rank = rank_expert % kRank;
-  const uint expert = rank_expert / kRank;
-  const uint start = uint(offsets[expert]);
-  const uint stop = uint(offsets[expert + 1]);
-  float value = 0.0f;
-  for (uint row = start; row < stop; ++row) {
-    value += float(grad_low_rank[row * kRank + rank])
-        * float(inputs[row * input_dim + column]);
-  }
-  grad_weight_a[index] = bfloat(value);
-}
-
-kernel void expert_lora_backward_b(
-    device const bfloat* grad_output [[buffer(0)]],
-    device const bfloat* low_rank [[buffer(1)]],
-    device const long* offsets [[buffer(2)]],
-    device bfloat* grad_weight_b [[buffer(3)]],
-    constant uint& expert_count [[buffer(4)]],
-    constant uint& output_dim [[buffer(5)]],
-    uint index [[thread_position_in_grid]]) {
-  const uint total = expert_count * output_dim * kRank;
-  if (index >= total) return;
-  const uint rank = index % kRank;
-  const uint output_expert = index / kRank;
-  const uint column = output_expert % output_dim;
-  const uint expert = output_expert / output_dim;
-  const uint start = uint(offsets[expert]);
-  const uint stop = uint(offsets[expert + 1]);
-  float value = 0.0f;
-  for (uint row = start; row < stop; ++row) {
-    value += float(grad_output[row * output_dim + column])
-        * float(low_rank[row * kRank + rank]);
-  }
-  grad_weight_b[index] = bfloat(value);
-}
 """
 
 
@@ -363,118 +210,20 @@ def _get_independent_expert_lora_metal_library():
     return _independent_expert_lora_metal_library
 
 
-class _MetalIndependentExpertLoRA(torch.autograd.Function):
-    """Independent rank-8 expert A/B with direct BF16 Metal forward/backward."""
-
-    @staticmethod
-    def forward(
-        ctx,
-        inputs: torch.Tensor,
-        weight_a: torch.Tensor,
-        weight_b: torch.Tensor,
-        expert_ids: torch.LongTensor,
-        counts: torch.LongTensor,
-    ) -> torch.Tensor:
-        inputs = inputs.contiguous()
-        weight_a = weight_a.contiguous()
-        weight_b = weight_b.contiguous()
-        expert_ids = expert_ids.contiguous()
-        expert_count, rank, input_dim = weight_a.shape
-        row_count = inputs.shape[0]
-        output_dim = weight_b.shape[1]
-        if rank != 8:
-            raise ValueError("The direct Metal expert LoRA kernel requires rank 8.")
-        offsets = torch.cat(
-            (
-                counts.new_zeros(1, dtype=torch.int64),
-                counts.cumsum(0, dtype=torch.int64),
-            )
-        ).contiguous()
-        low_rank = inputs.new_empty((row_count, rank))
-        output = inputs.new_empty((row_count, output_dim))
-        library = _get_independent_expert_lora_metal_library()
-        library.expert_lora_forward_a(
-            inputs,
-            weight_a,
-            expert_ids,
-            low_rank,
-            row_count,
-            input_dim,
-            threads=(row_count * 256, 1, 1),
-            group_size=(256, 1, 1),
-        )
-        library.expert_lora_forward_b(
-            low_rank,
-            weight_b,
-            expert_ids,
-            output,
-            row_count,
-            output_dim,
-            threads=(row_count * output_dim, 1, 1),
-            group_size=(256, 1, 1),
-        )
-        ctx.save_for_backward(
-            inputs,
-            low_rank,
-            weight_a,
-            weight_b,
-            expert_ids,
-            offsets,
-        )
-        ctx.dimensions = (expert_count, row_count, input_dim, output_dim)
-        return output
-
-    @staticmethod
-    def backward(ctx, grad_output: torch.Tensor):
-        inputs, low_rank, weight_a, weight_b, expert_ids, offsets = ctx.saved_tensors
-        expert_count, row_count, input_dim, output_dim = ctx.dimensions
-        grad_output = grad_output.contiguous()
-        grad_low_rank = low_rank.new_empty(low_rank.shape)
-        grad_inputs = inputs.new_empty(inputs.shape)
-        grad_weight_a = weight_a.new_empty(weight_a.shape)
-        grad_weight_b = weight_b.new_empty(weight_b.shape)
-        library = _get_independent_expert_lora_metal_library()
-        library.expert_lora_backward_low_rank(
-            grad_output,
-            weight_b,
-            expert_ids,
-            grad_low_rank,
-            row_count,
-            output_dim,
-            threads=(row_count * 256, 1, 1),
-            group_size=(256, 1, 1),
-        )
-        library.expert_lora_backward_inputs(
-            grad_low_rank,
-            weight_a,
-            expert_ids,
-            grad_inputs,
-            row_count,
-            input_dim,
-            threads=(row_count * input_dim, 1, 1),
-            group_size=(256, 1, 1),
-        )
-        library.expert_lora_backward_a(
-            grad_low_rank,
-            inputs,
-            offsets,
-            grad_weight_a,
-            expert_count,
-            input_dim,
-            threads=(expert_count * 8 * input_dim, 1, 1),
-            group_size=(256, 1, 1),
-        )
-        library.expert_lora_backward_b(
-            grad_output,
-            low_rank,
-            offsets,
-            grad_weight_b,
-            expert_count,
-            output_dim,
-            threads=(expert_count * output_dim * 8, 1, 1),
-            group_size=(256, 1, 1),
-        )
-        return grad_inputs, grad_weight_a, grad_weight_b, None, None
+def _metal_expert_lora(inputs: torch.Tensor, weight_a: torch.Tensor, weight_b: torch.Tensor, expert_ids: torch.LongTensor) -> torch.Tensor:
+    inputs = inputs.contiguous()
+    weight_a = weight_a.contiguous()
+    weight_b = weight_b.contiguous()
+    expert_ids = expert_ids.contiguous()
+    expert_count, rank, input_dim = weight_a.shape
+    row_count = inputs.shape[0]
+    output_dim = weight_b.shape[1]
+    low_rank = inputs.new_empty((row_count, rank))
+    output = inputs.new_empty((row_count, output_dim))
+    library = _get_independent_expert_lora_metal_library()
+    library.expert_lora_forward_a(inputs, weight_a, expert_ids, low_rank, row_count, input_dim, threads=(row_count * 256, 1, 1), group_size=(256, 1, 1))
+    library.expert_lora_forward_b(low_rank, weight_b, expert_ids, output, row_count, output_dim, threads=(row_count * output_dim, 1, 1), group_size=(256, 1, 1))
+    return output
 
 
 def _independent_expert_lora(
@@ -496,14 +245,13 @@ def _independent_expert_lora(
         and weight_a.shape[1] == 8
         and hasattr(torch.mps, "compile_shader")
     ):
-        return _MetalIndependentExpertLoRA.apply(
+        return _metal_expert_lora(
             inputs,
             weight_a,
             weight_b,
             expert_ids,
-            counts,
         )
-    return _IndependentExpertLoRA.apply(
+    return _padded_expert_lora(
         inputs,
         weight_a,
         weight_b,
@@ -553,16 +301,7 @@ def mps_segmented_experts_forward(
         bias=gate_up_bias,
         is_transposed=self.is_transposed,
     )
-    if hasattr(self, "lora_gate_up_a") and self.training:
-        gate_up_update = _independent_expert_lora(
-            self.lora_dropout(selected_hidden_g),
-            self.lora_gate_up_a,
-            self.lora_gate_up_b,
-            expert_ids_g,
-            counts,
-        )
-        gate_up_g.add_(gate_up_update, alpha=self.lora_scaling)
-    elif hasattr(self, "lora_gate_up_a"):
+    if hasattr(self, "lora_gate_up_a"):
         gate_up_update = _independent_expert_lora(
             selected_hidden_g,
             self.lora_gate_up_a,
@@ -570,7 +309,7 @@ def mps_segmented_experts_forward(
             expert_ids_g,
             counts,
         )
-        gate_up_g.add_(gate_up_update, alpha=self.lora_scaling)
+        gate_up_g = gate_up_g + gate_up_update * self.lora_scaling
     activated_g = (
         self._apply_gate(gate_up_g)
         if self.has_gate
@@ -584,16 +323,7 @@ def mps_segmented_experts_forward(
         bias=down_bias,
         is_transposed=self.is_transposed,
     )
-    if hasattr(self, "lora_down_a") and self.training:
-        down_update = _independent_expert_lora(
-            self.lora_dropout(activated_g),
-            self.lora_down_a,
-            self.lora_down_b,
-            expert_ids_g,
-            counts,
-        )
-        projected_g.add_(down_update, alpha=self.lora_scaling)
-    elif hasattr(self, "lora_down_a"):
+    if hasattr(self, "lora_down_a"):
         down_update = _independent_expert_lora(
             activated_g,
             self.lora_down_a,
@@ -601,18 +331,12 @@ def mps_segmented_experts_forward(
             expert_ids_g,
             counts,
         )
-        projected_g.add_(down_update, alpha=self.lora_scaling)
+        projected_g = projected_g + down_update * self.lora_scaling
     weighted_g = projected_g * sample_weights_g[:, None]
     weighted = torch.empty_like(weighted_g)
     weighted.index_copy_(0, permutation, weighted_g)
     return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)
 
 
-ALL_EXPERTS_FUNCTIONS.register("mps_segmented", mps_segmented_experts_forward)
-
-
-__all__ = [
-    "grouped_expert_offsets",
-    "mps_grouped_mm_experts_forward",
-    "mps_segmented_experts_forward",
-]
+def register_mps_backends() -> None:
+    ALL_EXPERTS_FUNCTIONS.register("mps_segmented", mps_segmented_experts_forward)

vocab_ops.py

@@ -1,16 +1,10 @@
-# Copyright 2026 Modilify
-# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
-"""Exact memory-bounded vocabulary projection and sampling."""
-
+"""Memory-bounded vocabulary projection and exact inference sampling."""
 from __future__ import annotations
-
 import math
 from collections.abc import Sequence
-
 import torch
 from torch.nn import functional as F
 
-
 def _stable_max(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     """Return ``max`` indices that stay in-range on MPS NaN/-inf rows."""
 
@@ -20,322 +14,6 @@ def _stable_max(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         return best, index
     return best, index.clamp(0, width - 1)
 
-
-class _ChunkedVocabStatistics(torch.autograd.Function):
-    """One projection pass for CE and exact Gumbel-max categorical sampling."""
-
-    @staticmethod
-    def forward(
-        ctx,
-        hidden: torch.Tensor,
-        weight: torch.Tensor,
-        labels: torch.LongTensor,
-        softcap: float,
-        temperature: float,
-        chunk_size: int,
-    ) -> tuple[
-        torch.Tensor,
-        torch.Tensor,
-        torch.LongTensor,
-        torch.Tensor,
-        torch.Tensor,
-        torch.LongTensor,
-        torch.Tensor,
-    ]:
-        if hidden.ndim != 2 or weight.ndim != 2:
-            raise ValueError("Chunked vocabulary tensors must be matrices.")
-        if hidden.shape[1] != weight.shape[1]:
-            raise ValueError("Hidden and vocabulary projection dimensions differ.")
-        if temperature <= 0 or chunk_size <= 0:
-            raise ValueError("Temperature and vocabulary chunk size must be positive.")
-        labels = labels.to(device=hidden.device, dtype=torch.long)
-        if labels.shape != (hidden.shape[0],):
-            raise ValueError("Labels must contain one value per hidden row.")
-
-        rows = hidden.shape[0]
-        row_indices = torch.arange(rows, device=hidden.device)
-        valid = labels.ge(0)
-        raw_log_z = torch.full(
-            (rows,), -torch.inf, device=hidden.device, dtype=torch.float32
-        )
-        sample_log_z = torch.full_like(raw_log_z, -torch.inf)
-        gold_score = torch.zeros_like(raw_log_z)
-        best_gumbel = torch.full_like(raw_log_z, -torch.inf)
-        selected_score = torch.zeros_like(raw_log_z)
-        selected = torch.zeros(rows, device=hidden.device, dtype=torch.long)
-        greedy_score = torch.full_like(raw_log_z, -torch.inf)
-        greedy = torch.zeros(rows, device=hidden.device, dtype=torch.long)
-        moment_max = torch.full_like(raw_log_z, -torch.inf)
-        moment_sum = torch.zeros_like(raw_log_z)
-        moment_weighted = torch.zeros_like(raw_log_z)
-        vocab_size = weight.shape[0]
-
-        with torch.no_grad():
-            for start in range(0, vocab_size, int(chunk_size)):
-                stop = min(start + int(chunk_size), vocab_size)
-                raw = F.linear(hidden, weight[start:stop])
-                scores = (
-                    torch.tanh(raw.float() / float(softcap))
-                    * float(softcap)
-                )
-                raw_log_z = torch.logaddexp(
-                    raw_log_z,
-                    torch.logsumexp(scores, dim=-1),
-                )
-                sample_scores = scores / float(temperature)
-                sample_log_z = torch.logaddexp(
-                    sample_log_z,
-                    torch.logsumexp(sample_scores, dim=-1),
-                )
-
-                in_chunk = valid & labels.ge(start) & labels.lt(stop)
-                local_gold = (labels - start).clamp(0, stop - start - 1)
-                gold_score = torch.where(
-                    in_chunk,
-                    scores[row_indices, local_gold],
-                    gold_score,
-                )
-
-                uniform = torch.rand(
-                    sample_scores.shape,
-                    device=sample_scores.device,
-                    dtype=torch.float32,
-                ).clamp_(
-                    min=torch.finfo(torch.float32).tiny,
-                    max=1.0 - torch.finfo(torch.float32).eps,
-                )
-                gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
-                chunk_best, chunk_index = _stable_max(gumbel_scores)
-                replace_best = chunk_best.gt(best_gumbel)
-                candidate_score = sample_scores.gather(
-                    1, chunk_index[:, None]
-                ).squeeze(-1)
-                best_gumbel = torch.maximum(best_gumbel, chunk_best)
-                selected = torch.where(
-                    replace_best,
-                    chunk_index + start,
-                    selected,
-                )
-                selected_score = torch.where(
-                    replace_best,
-                    candidate_score,
-                    selected_score,
-                )
-
-                chunk_max, chunk_argmax = _stable_max(sample_scores)
-                replace_greedy = chunk_max.gt(greedy_score)
-                greedy_score = torch.maximum(greedy_score, chunk_max)
-                greedy = torch.where(
-                    replace_greedy,
-                    chunk_argmax + start,
-                    greedy,
-                )
-                shifted = torch.exp(sample_scores - chunk_max[:, None])
-                chunk_sum = shifted.sum(dim=-1)
-                chunk_weighted = (shifted * sample_scores).sum(dim=-1)
-                merged_max = torch.maximum(moment_max, chunk_max)
-                previous_scale = torch.exp(moment_max - merged_max)
-                chunk_scale = torch.exp(chunk_max - merged_max)
-                moment_sum = (
-                    moment_sum * previous_scale + chunk_sum * chunk_scale
-                )
-                moment_weighted = (
-                    moment_weighted * previous_scale
-                    + chunk_weighted * chunk_scale
-                )
-                moment_max = merged_max
-
-        raw_nll = torch.where(
-            valid,
-            raw_log_z - gold_score,
-            torch.zeros_like(raw_log_z),
-        )
-        temperature_nll = torch.where(
-            valid,
-            sample_log_z - gold_score / float(temperature),
-            torch.zeros_like(sample_log_z),
-        )
-        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
-        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
-        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
-            torch.finfo(torch.float32).tiny
-        )
-
-        ctx.save_for_backward(
-            hidden,
-            weight,
-            labels,
-            selected,
-            raw_log_z,
-            sample_log_z,
-            confidence,
-            entropy,
-        )
-        ctx.softcap = float(softcap)
-        ctx.temperature = float(temperature)
-        ctx.chunk_size = int(chunk_size)
-        ctx.mark_non_differentiable(selected, greedy, greedy_confidence)
-        return (
-            raw_nll,
-            temperature_nll,
-            selected,
-            confidence,
-            entropy,
-            greedy,
-            greedy_confidence,
-        )
-
-    @staticmethod
-    def backward(
-        ctx,
-        grad_raw_nll: torch.Tensor | None,
-        grad_temperature_nll: torch.Tensor | None,
-        grad_selected: torch.Tensor | None,
-        grad_confidence: torch.Tensor | None,
-        grad_entropy: torch.Tensor | None,
-        grad_greedy: torch.Tensor | None,
-        grad_greedy_confidence: torch.Tensor | None,
-    ):
-        del grad_selected, grad_greedy, grad_greedy_confidence
-        (
-            hidden,
-            weight,
-            labels,
-            selected,
-            raw_log_z,
-            sample_log_z,
-            confidence,
-            entropy,
-        ) = ctx.saved_tensors
-        grad_raw_nll = (
-            torch.zeros_like(raw_log_z)
-            if grad_raw_nll is None
-            else grad_raw_nll.float()
-        )
-        grad_temperature_nll = (
-            torch.zeros_like(sample_log_z)
-            if grad_temperature_nll is None
-            else grad_temperature_nll.float()
-        )
-        grad_confidence = (
-            torch.zeros_like(confidence)
-            if grad_confidence is None
-            else grad_confidence.float()
-        )
-        valid = labels.ge(0)
-        row_indices = torch.arange(hidden.shape[0], device=hidden.device)
-        grad_hidden = torch.zeros_like(hidden, dtype=torch.float32)
-        confidence_scale = (
-            grad_confidence * confidence / ctx.temperature
-        )
-        for start in range(0, weight.shape[0], ctx.chunk_size):
-            stop = min(start + ctx.chunk_size, weight.shape[0])
-            raw = F.linear(hidden, weight[start:stop])
-            scores = torch.tanh(raw.float() / ctx.softcap) * ctx.softcap
-            raw_probability = torch.exp(scores - raw_log_z[:, None])
-            sample_probability = torch.exp(
-                scores / ctx.temperature - sample_log_z[:, None]
-            )
-            score_gradient = (
-                raw_probability * (grad_raw_nll * valid)[:, None]
-                + sample_probability
-                * (grad_temperature_nll * valid)[:, None]
-                / ctx.temperature
-                - sample_probability * confidence_scale[:, None]
-            )
-            if grad_entropy is not None:
-                grad_entropy_f = grad_entropy.float()
-                sample_scores = scores / ctx.temperature
-                entropy_grad = -(
-                    grad_entropy_f / ctx.temperature
-                )[:, None] * sample_probability * (
-                    sample_scores - sample_log_z[:, None] + entropy[:, None]
-                )
-                score_gradient = score_gradient + entropy_grad
-
-            gold_in_chunk = valid & labels.ge(start) & labels.lt(stop)
-            gold_local = (labels - start).clamp(0, stop - start - 1)
-            score_gradient[row_indices, gold_local] -= (
-                grad_raw_nll * gold_in_chunk
-                + grad_temperature_nll
-                * gold_in_chunk
-                / ctx.temperature
-            )
-            selected_in_chunk = selected.ge(start) & selected.lt(stop)
-            selected_local = (selected - start).clamp(0, stop - start - 1)
-            score_gradient[row_indices, selected_local] += (
-                confidence_scale * selected_in_chunk
-            )
-            score_gradient *= 1.0 - (scores / ctx.softcap).square()
-            # The B16 shape feeds 4,096 rows through a 32,768-wide vocabulary
-            # chunk.  MPS BF16 GEMM can return NaNs for this backward-only
-            # projection even when every score gradient is finite.  Accumulate
-            # the exact VJP in FP32; casting the finished hidden gradient back
-            # to the model dtype happens only once below.
-            chunk_gradient = F.linear(
-                score_gradient,
-                weight[start:stop].float().transpose(0, 1),
-            )
-            grad_hidden.add_(chunk_gradient)
-
-        return (
-            grad_hidden.to(hidden.dtype),
-            None,
-            None,
-            None,
-            None,
-            None,
-        )
-
-
-def chunked_vocab_training_statistics(
-    hidden: torch.Tensor,
-    weight: torch.Tensor,
-    labels: torch.LongTensor,
-    *,
-    softcap: float,
-    temperature: float,
-    chunk_size: int,
-) -> tuple[
-    torch.Tensor,
-    torch.Tensor,
-    torch.LongTensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.LongTensor,
-    torch.Tensor,
-]:
-    """Return exact losses plus sampled and greedy proposal statistics."""
-
-    shape = labels.shape
-    outputs = _ChunkedVocabStatistics.apply(
-        hidden.reshape(-1, hidden.shape[-1]),
-        weight,
-        labels.reshape(-1),
-        float(softcap),
-        float(temperature),
-        int(chunk_size),
-    )
-    (
-        raw_nll,
-        temperature_nll,
-        proposal,
-        confidence,
-        entropy,
-        greedy,
-        greedy_confidence,
-    ) = outputs
-    return (
-        raw_nll.view(shape),
-        temperature_nll.view(shape),
-        proposal.view(shape),
-        confidence.view(shape),
-        entropy.view(shape),
-        greedy.view(shape),
-        greedy_confidence.view(shape),
-    )
-
-
 @torch.no_grad()
 def chunked_vocab_statistics(
     hidden: torch.Tensor,
@@ -347,6 +25,8 @@ def chunked_vocab_statistics(
     repetition_token_mask: torch.BoolTensor | None = None,
     repetition_penalty: float = 1.0,
     sampling_generators: Sequence[torch.Generator] | None = None,
+    top_k: int | None = 40,
+    min_p: float | None = 0.05,
 ) -> tuple[
     torch.LongTensor,
     torch.Tensor,
@@ -359,71 +39,77 @@ def chunked_vocab_statistics(
     if not math.isfinite(repetition_penalty) or repetition_penalty <= 0:
         raise ValueError("`repetition_penalty` must be a finite positive number.")
 
-    # Preserve the original custom-autograd path exactly when disabled. In
-    # particular, this keeps the same chunk-local RNG draws and default output.
-    if repetition_penalty != 1.0 or sampling_generators is not None:
-        if hidden.ndim < 2 or weight.ndim != 2:
-            raise ValueError("Chunked vocabulary tensors must have matrix features.")
-        if hidden.shape[-1] != weight.shape[1]:
-            raise ValueError("Hidden and vocabulary projection dimensions differ.")
-        if temperature <= 0 or chunk_size <= 0:
-            raise ValueError("Temperature and vocabulary chunk size must be positive.")
-        expected_mask_shape = (hidden.shape[0], weight.shape[0])
-        if repetition_penalty != 1.0:
-            if repetition_token_mask is None or repetition_token_mask.shape != expected_mask_shape:
-                raise ValueError(
-                    "`repetition_token_mask` must have shape [batch, vocabulary]."
-                )
-            if repetition_token_mask.dtype != torch.bool:
-                raise ValueError("`repetition_token_mask` must be a boolean tensor.")
-        if sampling_generators is not None and len(sampling_generators) != hidden.shape[0]:
+    if hidden.ndim < 2 or weight.ndim != 2:
+        raise ValueError("Chunked vocabulary tensors must have matrix features.")
+    if hidden.shape[-1] != weight.shape[1]:
+        raise ValueError("Hidden and vocabulary projection dimensions differ.")
+    if temperature <= 0 or chunk_size <= 0:
+        raise ValueError("Temperature and vocabulary chunk size must be positive.")
+    expected_mask_shape = (hidden.shape[0], weight.shape[0])
+    if repetition_penalty != 1.0:
+        if repetition_token_mask is None or repetition_token_mask.shape != expected_mask_shape:
             raise ValueError(
-                "`sampling_generators` must contain one generator per batch row."
+                "`repetition_token_mask` must have shape [batch, vocabulary]."
             )
-
-        output_shape = hidden.shape[:-1]
-        flat_hidden = hidden.reshape(-1, hidden.shape[-1])
-        rows_per_batch = math.prod(hidden.shape[1:-1])
-        batch_indices = torch.arange(
-            flat_hidden.shape[0], device=hidden.device
-        ).div(rows_per_batch, rounding_mode="floor")
-        sample_log_z = torch.full(
-            (flat_hidden.shape[0],),
-            -torch.inf,
-            device=hidden.device,
-            dtype=torch.float32,
-        )
-        best_gumbel = torch.full_like(sample_log_z, -torch.inf)
-        selected_score = torch.zeros_like(sample_log_z)
-        selected = torch.zeros(
-            flat_hidden.shape[0], device=hidden.device, dtype=torch.long
+        if repetition_token_mask.dtype != torch.bool:
+            raise ValueError("`repetition_token_mask` must be a boolean tensor.")
+    if sampling_generators is not None and len(sampling_generators) != hidden.shape[0]:
+        raise ValueError(
+            "`sampling_generators` must contain one generator per batch row."
         )
-        greedy_score = torch.full_like(sample_log_z, -torch.inf)
-        greedy = torch.zeros_like(selected)
-        moment_max = torch.full_like(sample_log_z, -torch.inf)
-        moment_sum = torch.zeros_like(sample_log_z)
-        moment_weighted = torch.zeros_like(sample_log_z)
 
-        for start in range(0, weight.shape[0], int(chunk_size)):
-            stop = min(start + int(chunk_size), weight.shape[0])
-            raw = F.linear(flat_hidden, weight[start:stop])
-            scores = torch.tanh(raw.float() / float(softcap)) * float(softcap)
-            if repetition_penalty != 1.0:
-                assert repetition_token_mask is not None
-                seen = repetition_token_mask[:, start:stop].index_select(
-                    0, batch_indices
-                )
-                penalized = torch.where(
-                    scores < 0,
-                    scores * float(repetition_penalty),
-                    scores / float(repetition_penalty),
-                )
-                scores = torch.where(seen, penalized, scores)
-            sample_scores = scores / float(temperature)
-            sample_log_z = torch.logaddexp(
-                sample_log_z, torch.logsumexp(sample_scores, dim=-1)
+    output_shape = hidden.shape[:-1]
+    flat_hidden = hidden.reshape(-1, hidden.shape[-1])
+    rows_per_batch = math.prod(hidden.shape[1:-1])
+    batch_indices = torch.arange(
+        flat_hidden.shape[0], device=hidden.device
+    ).div(rows_per_batch, rounding_mode="floor")
+    sample_log_z = torch.full(
+        (flat_hidden.shape[0],),
+        -torch.inf,
+        device=hidden.device,
+        dtype=torch.float32,
+    )
+    best_gumbel = torch.full_like(sample_log_z, -torch.inf)
+    selected_score = torch.zeros_like(sample_log_z)
+    selected = torch.zeros(
+        flat_hidden.shape[0], device=hidden.device, dtype=torch.long
+    )
+    greedy_score = torch.full_like(sample_log_z, -torch.inf)
+    greedy = torch.zeros_like(selected)
+    moment_max = torch.full_like(sample_log_z, -torch.inf)
+    moment_sum = torch.zeros_like(sample_log_z)
+    moment_weighted = torch.zeros_like(sample_log_z)
+
+    use_constrained = top_k is not None and top_k > 0
+    chunk_cand_scores = []
+    chunk_cand_tokens = []
+
+    for start in range(0, weight.shape[0], int(chunk_size)):
+        stop = min(start + int(chunk_size), weight.shape[0])
+        raw = F.linear(flat_hidden, weight[start:stop])
+        scores = torch.tanh(raw.float() / float(softcap)) * float(softcap)
+        if repetition_penalty != 1.0:
+            seen = repetition_token_mask[:, start:stop].index_select(
+                0, batch_indices
+            )
+            penalized = torch.where(
+                scores < 0,
+                scores * float(repetition_penalty),
+                scores / float(repetition_penalty),
             )
+            scores = torch.where(seen, penalized, scores)
+        sample_scores = scores / float(temperature)
+        sample_log_z = torch.logaddexp(
+            sample_log_z, torch.logsumexp(sample_scores, dim=-1)
+        )
 
+        if use_constrained:
+            k = min(int(top_k), stop - start)
+            c_score, c_idx = torch.topk(sample_scores, k=k, dim=-1)
+            chunk_cand_scores.append(c_score)
+            chunk_cand_tokens.append(c_idx + start)
+        else:
             if sampling_generators is None:
                 uniform = torch.rand(
                     sample_scores.shape,
@@ -431,8 +117,6 @@ def chunked_vocab_statistics(
                     dtype=torch.float32,
                 )
             else:
-                # Drawing each request from its own generator makes sampling
-                # invariant to dynamic admission, slot changes, and batch order.
                 per_request_shape = (
                     rows_per_batch,
                     sample_scores.shape[-1],
@@ -465,61 +149,75 @@ def chunked_vocab_statistics(
                 replace_best, candidate_score, selected_score
             )
 
-            chunk_max, chunk_argmax = _stable_max(sample_scores)
-            replace_greedy = chunk_max.gt(greedy_score)
-            greedy_score = torch.maximum(greedy_score, chunk_max)
-            greedy = torch.where(replace_greedy, chunk_argmax + start, greedy)
-            shifted = torch.exp(sample_scores - chunk_max[:, None])
-            chunk_sum = shifted.sum(dim=-1)
-            chunk_weighted = (shifted * sample_scores).sum(dim=-1)
-            merged_max = torch.maximum(moment_max, chunk_max)
-            previous_scale = torch.exp(moment_max - merged_max)
-            chunk_scale = torch.exp(chunk_max - merged_max)
-            moment_sum = moment_sum * previous_scale + chunk_sum * chunk_scale
-            moment_weighted = (
-                moment_weighted * previous_scale + chunk_weighted * chunk_scale
-            )
-            moment_max = merged_max
-
-        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
-        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
-        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
-            torch.finfo(torch.float32).tiny
+        chunk_max, chunk_argmax = _stable_max(sample_scores)
+        replace_greedy = chunk_max.gt(greedy_score)
+        greedy_score = torch.maximum(greedy_score, chunk_max)
+        greedy = torch.where(replace_greedy, chunk_argmax + start, greedy)
+        shifted = torch.exp(sample_scores - chunk_max[:, None])
+        chunk_sum = shifted.sum(dim=-1)
+        chunk_weighted = (shifted * sample_scores).sum(dim=-1)
+        merged_max = torch.maximum(moment_max, chunk_max)
+        previous_scale = torch.exp(moment_max - merged_max)
+        chunk_scale = torch.exp(chunk_max - merged_max)
+        moment_sum = moment_sum * previous_scale + chunk_sum * chunk_scale
+        moment_weighted = (
+            moment_weighted * previous_scale + chunk_weighted * chunk_scale
         )
-        return (
-            selected.view(output_shape),
-            confidence.view(output_shape),
-            entropy.view(output_shape),
-            greedy.view(output_shape),
-            greedy_confidence.view(output_shape),
+        moment_max = merged_max
+
+    if use_constrained:
+        all_scores = torch.cat(chunk_cand_scores, dim=-1)
+        all_tokens = torch.cat(chunk_cand_tokens, dim=-1)
+        global_k = min(int(top_k), all_scores.shape[-1])
+        cand_scores, global_idx = torch.topk(all_scores, k=global_k, dim=-1)
+        cand_tokens = all_tokens.gather(dim=-1, index=global_idx)
+        if min_p is not None and min_p > 0:
+            thresh = greedy_score + math.log(float(min_p))
+            valid_mask = cand_scores >= thresh[:, None]
+            eligible = cand_scores.masked_fill(~valid_mask, -float("inf"))
+        else:
+            eligible = cand_scores
+        if sampling_generators is None:
+            uniform = torch.rand(
+                eligible.shape,
+                device=eligible.device,
+                dtype=torch.float32,
+            )
+        else:
+            per_request_shape = (
+                rows_per_batch,
+                eligible.shape[-1],
+            )
+            uniform = torch.cat(
+                [
+                    torch.rand(
+                        per_request_shape,
+                        device=eligible.device,
+                        dtype=torch.float32,
+                        generator=generator,
+                    )
+                    for generator in sampling_generators
+                ],
+                dim=0,
+            )
+        uniform = uniform.clamp_(
+            min=torch.finfo(torch.float32).tiny,
+            max=1.0 - torch.finfo(torch.float32).eps,
         )
-
-    labels = torch.full(
-        hidden.shape[:-1],
-        -100,
-        device=hidden.device,
-        dtype=torch.long,
+        gumbel = eligible - torch.log(-torch.log(uniform))
+        _, chosen_idx = _stable_max(gumbel)
+        selected = cand_tokens.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
+        selected_score = cand_scores.gather(dim=-1, index=chosen_idx[:, None]).squeeze(-1)
+
+    confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
+    greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
+    entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
+        torch.finfo(torch.float32).tiny
     )
-    (
-        _,
-        _,
-        proposal,
-        confidence,
-        entropy,
-        greedy,
-        greedy_confidence,
-    ) = chunked_vocab_training_statistics(
-        hidden,
-        weight,
-        labels,
-        softcap=softcap,
-        temperature=temperature,
-        chunk_size=chunk_size,
+    return (
+        selected.view(output_shape),
+        confidence.view(output_shape),
+        entropy.view(output_shape),
+        greedy.view(output_shape),
+        greedy_confidence.view(output_shape),
     )
-    return proposal, confidence, entropy, greedy, greedy_confidence
-
-
-__all__ = [
-    "chunked_vocab_statistics",
-    "chunked_vocab_training_statistics",
-]